update to remove previously hard-coded rocprofiler-sdk path (#369)

* update to remove previously hard-coded rocprofiler-sdk path and add skip_rocprofiler_sdk to avoid loading `rocprofiler-sdk`

Author: cj401-amd

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",
 }
 
 

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

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*"]),

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:

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",