diff --git a/build_tools/lint/tags.py b/build_tools/lint/tags.py index 35a722e38d6f7..96e509cf02575 100644 --- a/build_tools/lint/tags.py +++ b/build_tools/lint/tags.py @@ -94,6 +94,7 @@ "multi_gpu_h100": ( "Used by `xla_test` to signal that multiple H100s are needed." ), + "skip_rocprofiler_sdk": "used to skip rocmtracer test as it calls rocprofiler-sdk via rocprofiler_force_configure", } diff --git a/build_tools/rocm/run_xla.sh b/build_tools/rocm/run_xla.sh index 548d18bcf55dd..99db81cd84df2 100755 --- a/build_tools/rocm/run_xla.sh +++ b/build_tools/rocm/run_xla.sh @@ -69,7 +69,7 @@ BAZEL_DISK_CACHE_DIR="/tf/disk_cache/rocm-jaxlib-v0.6.0" mkdir -p ${BAZEL_DISK_CACHE_DIR} SCRIPT_DIR=$(realpath $(dirname $0)) -TAG_FILTERS=$($SCRIPT_DIR/rocm_tag_filters.sh),gpu,-multigpu,-multi_gpu_h100,requires-gpu-amd +TAG_FILTERS=$($SCRIPT_DIR/rocm_tag_filters.sh),gpu,-multigpu,-multi_gpu_h100,requires-gpu-amd,-skip_rocprofiler_sdk SANITIZER_ARGS=() if [[ $1 == "asan" ]]; then diff --git a/third_party/gpus/rocm/BUILD.tpl b/third_party/gpus/rocm/BUILD.tpl index e2143439b5799..f6a60c453c707 100644 --- a/third_party/gpus/rocm/BUILD.tpl +++ b/third_party/gpus/rocm/BUILD.tpl @@ -347,6 +347,19 @@ cc_library( deps = [":rocm_config"], ) +cc_library( + name = "rocprofiler-sdk", + srcs = glob(["%{rocm_root}/lib/librocprofiler-sdk*.so*"]), + hdrs = glob(["%{rocm_root}/include/rocprofiler-sdk/**"]), + include_prefix = "rocm", + includes = [ + "%{rocm_root}/include/", + ], + strip_include_prefix = "%{rocm_root}", + visibility = ["//visibility:public"], + deps = [":rocm_config"], +) + cc_library( name = "rocsolver", srcs = glob(["%{rocm_root}/lib/librocsolver*.so*"]), diff --git a/third_party/gpus/rocm_configure.bzl b/third_party/gpus/rocm_configure.bzl index 82e6db92734f5..8b6c833bcf3f8 100644 --- a/third_party/gpus/rocm_configure.bzl +++ b/third_party/gpus/rocm_configure.bzl @@ -341,6 +341,7 @@ def _find_libs(repository_ctx, rocm_config, miopen_path, rccl_path, bash_bin): ("rocsolver", rocm_config.rocm_toolkit_path), ("hipfft", rocm_config.rocm_toolkit_path), ("rocrand", rocm_config.rocm_toolkit_path), + ("rocprofiler-sdk", rocm_config.rocm_toolkit_path), ] ] if int(rocm_config.rocm_version_number) >= 40500: diff --git a/xla/backends/profiler/gpu/BUILD b/xla/backends/profiler/gpu/BUILD index 4fedad1de8da3..4bbff8eac368b 100644 --- a/xla/backends/profiler/gpu/BUILD +++ b/xla/backends/profiler/gpu/BUILD @@ -246,12 +246,6 @@ cc_library( "manual", "rocm-only", ], - linkopts = [ - "-Wl,--no-as-needed", - "-L/opt/rocm/lib", - "-Wl,-rpath,/opt/rocm/lib", - "-lhsa-runtime64", - ], deps = [ "//xla/tsl/profiler/backends/cpu:annotation_stack", "//xla/tsl/profiler/utils:time_utils", @@ -266,6 +260,7 @@ cc_library( "@tsl//tsl/platform:errors", "@tsl//tsl/platform:logging", "@tsl//tsl/platform:macros", + "@local_config_rocm//rocm:rocprofiler-sdk", ], visibility = ["//visibility:public"], ) @@ -274,13 +269,6 @@ cc_library( name = "rocm_collector", srcs = ["rocm_collector.cc"], hdrs = ["rocm_collector.h"], - # copybara:uncomment compatible_with = ["//buildenv/target:non_prod"], - linkopts = select({ - "//conditions:default": [ - "-L/opt/rocm/lib", # search path for all ROCm shared objects - "-lrocprofiler-sdk", # the library that owns the missing symbols - ], - }), tags = [ "gpu", "rocm-only", @@ -315,6 +303,7 @@ cc_library( "@tsl//tsl/platform:types", "@tsl//tsl/profiler/lib:profiler_factory", "@tsl//tsl/profiler/lib:profiler_interface", + "@local_config_rocm//rocm:rocprofiler-sdk", ], ) @@ -401,6 +390,7 @@ xla_test( tags = [ "gpu", "rocm-only", + "skip_rocprofiler_sdk", # due to rocprofiler-sdk's rocprofiler_force_configure ] + if_google([ # Optional: only run internally if ROCm config is enabled "manual",