diff --git a/tensorflow/core/profiler/internal/gpu/BUILD b/tensorflow/core/profiler/internal/gpu/BUILD index c6fe4d77031..670080573b2 100644 --- a/tensorflow/core/profiler/internal/gpu/BUILD +++ b/tensorflow/core/profiler/internal/gpu/BUILD @@ -15,6 +15,10 @@ load( "//tensorflow/core/platform:build_config_root.bzl", "tf_cuda_tests_tags", ) +load( + "//tensorflow/stream_executor:build_defs.bzl", + "tf_additional_cupti_deps", +) package( default_visibility = ["//tensorflow:internal"], @@ -91,9 +95,7 @@ tf_cuda_library( deps = [ "//tensorflow/core:lib", "//tensorflow/core:platform_base", - ] + if_cuda_is_configured_compat([ - "//tensorflow/stream_executor/cuda:cupti_stub", - ]), + ] + tf_additional_cupti_deps(), ) tf_cuda_library( @@ -104,9 +106,7 @@ tf_cuda_library( visibility = ["//visibility:public"], deps = [ ":cupti_interface", - ] + if_cuda_is_configured_compat([ - "//tensorflow/stream_executor/cuda:cupti_stub", - ]), + ] + tf_additional_cupti_deps(), ) tf_cuda_library( diff --git a/tensorflow/stream_executor/build_defs.bzl b/tensorflow/stream_executor/build_defs.bzl index 3cb24f8468f..830f526dbd7 100644 --- a/tensorflow/stream_executor/build_defs.bzl +++ b/tensorflow/stream_executor/build_defs.bzl @@ -10,6 +10,9 @@ def tf_additional_cuda_platform_deps(): def tf_additional_cuda_driver_deps(): return [":cuda_stub"] +def tf_additional_cupti_deps(): + return ["//tensorflow/stream_executor/cuda:cupti_stub"] + def tf_additional_cudnn_plugin_deps(): return [] diff --git a/tensorflow/stream_executor/cuda/cupti_stub.cc b/tensorflow/stream_executor/cuda/cupti_stub.cc index 130c3f96e44..feea09e002a 100644 --- a/tensorflow/stream_executor/cuda/cupti_stub.cc +++ b/tensorflow/stream_executor/cuda/cupti_stub.cc @@ -23,12 +23,16 @@ limitations under the License. namespace { // Returns DSO handle or null if loading the DSO fails. void* GetDsoHandle() { +#if defined(PLATFORM_GOOGLE) && (CUDA_VERSION > 10000) + return nullptr; +#else static auto handle = []() -> void* { auto handle_or = stream_executor::internal::DsoLoader::GetCuptiDsoHandle(); if (!handle_or.ok()) return nullptr; return handle_or.ValueOrDie(); }(); return handle; +#endif } template