Add missing dependency and remove workardound with the path

ROCm · Nov 19, 2024 · b63c225 · b63c225
1 parent a8ce7dc
commit b63c225
Show file tree

Hide file tree

Showing 6 changed files with 91 additions and 64 deletions.
diff --git a/.bazelrc b/.bazelrc
@@ -283,7 +283,6 @@ build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain
 build:rocm --define=using_rocm_hipcc=true
 build:rocm --define=tensorflow_mkldnn_contraction_kernel=0
 build:rocm --repo_env TF_NEED_ROCM=1
-test:rocm --test_env=ROCBLAS_TENSILE_LIBPATH=/opt/rocm/lib/rocblas/library
 
 build:sycl --crosstool_top=@local_config_sycl//crosstool:toolchain
 build:sycl --define=using_sycl=true

diff --git a/build_tools/rocm/run_xla.sh b/build_tools/rocm/run_xla.sh
@@ -55,19 +55,25 @@ UNSUPPORTED_GPU_TAGS="$(echo -requires-gpu-sm{60,70,80,86,89,90}{,-only})"
 TAGS_FILTER="${TAGS_FILTER},${UNSUPPORTED_GPU_TAGS// /,}"
 
 bazel \
+    --output_base="/home/atheodor/projects/xla/bazel-build-dir" \
     test \
+    --sandbox_debug \
+    --verbose_failures \
     --config=rocm \
     --build_tag_filters=${TAGS_FILTER} \
     --test_tag_filters=${TAGS_FILTER} \
     --test_timeout=920,2400,7200,9600 \
     --test_sharding_strategy=disabled \
     --test_output=errors \
-    --flaky_test_attempts=3 \
+    --flaky_test_attempts=1 \
     --keep_going \
     --local_test_jobs=${N_TEST_JOBS} \
     --test_env=TF_TESTS_PER_GPU=$TF_TESTS_PER_GPU \
     --test_env=TF_GPU_COUNT=$TF_GPU_COUNT \
     --action_env=XLA_FLAGS=--xla_gpu_force_compilation_parallelism=16 \
     --action_env=XLA_FLAGS=--xla_gpu_enable_llvm_module_compilation_parallelism=true \
     --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute \
-    -- //xla/...
+    -- //xla/tests:dot_operation_test_gpu_amd_any \
+    #-- //xla/...
+    #-- //xla/service/gpu/tests:element_wise_row_vectorization_test_gpu_amd_any \
+
diff --git a/third_party/tsl/third_party/gpus/rocm/BUILD.tpl b/third_party/tsl/third_party/gpus/rocm/BUILD.tpl
@@ -38,6 +38,44 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
+cc_library(
+    name = "rocm",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":hip",
+        ":hipblas",
+        ":hipblaslt",
+        ":hiprand",
+        ":hipsolver",
+        ":hipsparse",
+        ":hsa-rocr",
+        ":miopen",
+        ":rocblas",
+        ":rocm_config",
+        ":rocsolver",
+        ":roctracer",
+    ] + select_threshold(
+        above_or_eq = [":hipfft"],
+        below = [":rocfft"],
+        threshold = 40100,
+        value = rocm_version_number(),
+    ),
+)
+
+cc_library(
+    name = "hsa-rocr",
+    hdrs = glob(["%{rocm_root}/include/hsa/**"]),
+    data = glob(["%{rocm_root}/lib/libhsa-runtime*.so*"]),
+    include_prefix = "rocm",
+    includes = [
+        "%{rocm_root}/include",
+    ],
+    linkstatic = 1,
+    strip_include_prefix = "%{rocm_root}",
+    visibility = ["//visibility:public"],
+    deps = [":rocm_config"],
+)
+
 cc_library(
     name = "hip",
     srcs = glob(["%{rocm_root}/lib/libamdhip*.so*"]),
@@ -55,16 +93,16 @@ cc_library(
 
 cc_library(
     name = "rocblas",
-    srcs = glob(["%{rocm_root}/lib/librocblas*.so*"]),
     hdrs = glob(["%{rocm_root}/include/rocblas/**"]),
+    copts = ["-rpath local_config_rocm/rocm/rocm_dis/lib"],
     data = glob([
+        "%{rocm_root}/lib/librocblas*.so*",
         "%{rocm_root}/lib/rocblas/**",
     ]),
     include_prefix = "rocm",
     includes = [
         "%{rocm_root}/include",
     ],
-    linkstatic = 1,
     strip_include_prefix = "%{rocm_root}",
     visibility = ["//visibility:public"],
     deps = [":rocm_config"],
@@ -140,28 +178,6 @@ cc_library(
     deps = [":rocm_config"],
 )
 
-cc_library(
-    name = "rocm",
-    visibility = ["//visibility:public"],
-    deps = [
-        ":hip",
-        ":hipblas",
-        ":hiprand",
-        ":hipsolver",
-        ":hipsparse",
-        ":miopen",
-        ":rocblas",
-        ":rocm_config",
-        ":rocsolver",
-        ":roctracer",
-    ] + select_threshold(
-        above_or_eq = [":hipfft"],
-        below = [":rocfft"],
-        threshold = 40100,
-        value = rocm_version_number(),
-    ),
-)
-
 filegroup(
     name = "rocm_bin",
     srcs = glob(["%{rocm_root}/bin/**/*"]),
@@ -252,9 +268,26 @@ cc_library(
 
 cc_library(
     name = "hipblas",
-    srcs = glob(["%{rocm_root}/lib/libhipblas*.so*"]),
+    srcs = glob(["%{rocm_root}/lib/libhipblas.so*"]),
     hdrs = glob(["%{rocm_root}/include/hipblas/**"]),
-    data = glob(["%{rocm_root}/lib/libhipblas*.so*"]),
+    data = glob(["%{rocm_root}/lib/libhipblas.so*"]),
+    include_prefix = "rocm",
+    includes = [
+        "%{rocm_root}/include/",
+    ],
+    strip_include_prefix = "%{rocm_root}",
+    visibility = ["//visibility:public"],
+    deps = [":rocm_config"],
+)
+
+cc_library(
+    name = "hipblaslt",
+    hdrs = glob(["%{rocm_root}/include/hipblaslt/**"]),
+    copts = ["-rpath local_config_rocm/rocm/rocm_dis/lib"],
+    data = glob([
+        "%{rocm_root}/lib/hipblaslt/**",
+        "%{rocm_root}/lib/libhipblaslt.so*",
+    ]),
     include_prefix = "rocm",
     includes = [
         "%{rocm_root}/include/",

diff --git a/third_party/tsl/third_party/gpus/rocm/rocm_redist_ubuntu_20_04.bzl b/third_party/tsl/third_party/gpus/rocm/rocm_redist_ubuntu_20_04.bzl
@@ -17,6 +17,14 @@ rocm_redist_ubuntu_20_04 = {
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblas-dev6.2.0/hipblas-dev6.2.0_2.2.0.60200-66~20.04_amd64.deb",
                 sha256 = "ab3ee54b33eba013fbf3d9aefe64b54e1918b9fb72790ca0b57fb391cb662cf0",
             ),
+            struct(
+                url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblaslt6.2.0/hipblaslt6.2.0_0.8.0.60200-66~20.04_amd64.deb",
+                sha256 = "a9fb63b30c3383b37479253c59c85493b75646037e016f68b136d12c5e6a4adb",
+            ),
+            struct(
+                url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblaslt-dev6.2.0/hipblaslt-dev6.2.0_0.8.0.60200-66~20.04_amd64.deb",
+                sha256 = "5e2ca28ea10ff16eb75bde5c69aef6dbe3566d8fc49c3d2ace0ae7a777c1ffae",
+            ),
             struct(
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipcc6.2.0/hipcc6.2.0_1.1.1.60200-66~20.04_amd64.deb",
                 sha256 = "a68123c046b8c913705262014463a8a30768167a1b68a78d8455deaf85a802d7",

diff --git a/third_party/tsl/third_party/gpus/rocm/rocm_redist_ubuntu_22_04.bzl b/third_party/tsl/third_party/gpus/rocm/rocm_redist_ubuntu_22_04.bzl
@@ -17,6 +17,14 @@ rocm_redist_ubuntu_22_04 = {
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblas-dev6.2.0/hipblas-dev6.2.0_2.2.0.60200-66~22.04_amd64.deb",
                 sha256 = "885cf3f3a52ebde9caadf6348a6cda28fd15e3bc52bab0c90b587d72b29ff7ef",
             ),
+            struct(
+                url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblaslt6.2.0/hipblaslt6.2.0_0.8.0.60200-66~22.04_amd64.deb",
+                sha256 = "af3bea7cda7c1af147c3baae1cb3a8846ff571fe713c3a83d0924810bee734fe",
+            ),
+            struct(
+                url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblaslt-dev6.2.0/hipblaslt-dev6.2.0_0.8.0.60200-66~22.04_amd64.deb",
+                sha256 = "d9d2c80228ebfe74ebb98fc74fda57be498ab84ca40358355c616dfd38efded2",
+            ),
             struct(
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipcc6.2.0/hipcc6.2.0_1.1.1.60200-66~22.04_amd64.deb",
                 sha256 = "468026fa8eb70121f0c545557a926ddc41228cef9457b4a00d8fc3a36b04310f",

diff --git a/third_party/tsl/third_party/gpus/rocm/rocm_redist_ubuntu_24_04.bzl b/third_party/tsl/third_party/gpus/rocm/rocm_redist_ubuntu_24_04.bzl
@@ -2,177 +2,150 @@ rocm_redist_ubuntu_24_04 = {
     "6.2.0": {
         "archives": [
             struct(
-                name = "hip-runtime-amd-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hip-runtime-amd6.2.0/hip-runtime-amd6.2.0_6.2.41133.60200-66~24.04_amd64.deb",
                 sha256 = "4831c09058882d531d10ead6a13ada15e7e9cbb182bd56ab62f2404207488f3e",
             ),
             struct(
-                name = "hip-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hip-dev6.2.0/hip-dev6.2.0_6.2.41133.60200-66~24.04_amd64.deb",
                 sha256 = "5e6601ada30432ee0dab0473585bdf1fa7c398f0c655538d48eba9c44e6dc77a",
             ),
             struct(
-                name = "hipblas",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblas6.2.0/hipblas6.2.0_2.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "7ff8f6308c744c71008959b17ab6338de1c6fd3e4581dd94271e6eca9fdc4c13",
             ),
             struct(
-                name = "hipblas-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblas-dev6.2.0/hipblas-dev6.2.0_2.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "e9f71e71db600d72dcb2b61e64b965b6c60d47bd4bb699e8abec85edb260b819",
             ),
             struct(
-                name = "hipcc",
+                url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblaslt6.2.0/hipblaslt6.2.0_0.8.0.60200-66~24.04_amd64.deb",
+                sha256 = "e5dfd8ba9e49f919a96c102d3a652e8ef0c4d1a63b3f3909c856d40b1745e2a9",
+            ),
+            struct(
+                url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipblaslt-dev6.2.0/hipblaslt-dev6.2.0_0.8.0.60200-66~24.04_amd64.deb",
+                sha256 = "639bd47010035ee6719425510be33d2f54483004a909dfa4c64f853d7394a22f",
+            ),
+            struct(
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipcc6.2.0/hipcc6.2.0_1.1.1.60200-66~24.04_amd64.deb",
                 sha256 = "c2782a98633e4400f46ba732605e56b2821366db60ec06d88db0615e4d1acf3c",
             ),
             struct(
-                name = "hipcub-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipcub-dev6.2.0/hipcub-dev6.2.0_3.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "48fec4d06aef3159db4117125b728242a1eeb480ea3d55d3901d945d4b883694",
             ),
             struct(
-                name = "hipfft",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipfft6.2.0/hipfft6.2.0_1.0.14.60200-66~24.04_amd64.deb",
                 sha256 = "8dd73cdbd4f0563f4a0481304771e4cbcac5905eea1f2d8ef41f922cdf9aba85",
             ),
             struct(
-                name = "hipfft-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipfft-dev6.2.0/hipfft-dev6.2.0_1.0.14.60200-66~24.04_amd64.deb",
                 sha256 = "e3c0a4ebda8d3aacd44b19c6872f23222513be0a5c04f793605088d9183f1be4",
             ),
             struct(
-                name = "hipsolver",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipsolver6.2.0/hipsolver6.2.0_2.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "adbba9ffcf8b5e4202efbe45924d87520bf4100ec5464bd0ba3beb61cb535c6c",
             ),
             struct(
-                name = "hipsolver-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipsolver-dev6.2.0/hipsolver-dev6.2.0_2.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "01d3dd6195111808b40a5837d3e51d8c27c4700b4bd8bb2d901e39d0474fd98a",
             ),
             struct(
-                name = "hipsparse",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipsparse6.2.0/hipsparse6.2.0_3.1.1.60200-66~24.04_amd64.deb",
                 sha256 = "2ba33a96388cd3edd7b5b8b261fe99cbd569894f4d7db291fc0dd0ff5d7c67ce",
             ),
             struct(
-                name = "hipsparse-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hipsparse-dev6.2.0/hipsparse-dev6.2.0_3.1.1.60200-66~24.04_amd64.deb",
                 sha256 = "6a767f493a722e2d4260a9bc23cf9db66fd275a094b395c768e305f60d6b4fe9",
             ),
             struct(
-                name = "hiprand",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hiprand6.2.0/hiprand6.2.0_2.11.0.60200-66~24.04_amd64.deb",
                 sha256 = "82f182134b415080ba4a12fd7993b6099ee9b9e549c72bfebee24c8486704078",
             ),
             struct(
-                name = "hiprand-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hiprand-dev6.2.0/hiprand-dev6.2.0_2.11.0.60200-66~24.04_amd64.deb",
                 sha256 = "011d5c28f45cd9d756e0cf6ea6a3d37eabd98a3381ffd961c772ab92a37e4ee8",
             ),
             struct(
-                name = "hsarocr",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hsa-rocr6.2.0/hsa-rocr6.2.0_1.14.0.60200-66~24.04_amd64.deb",
                 sha256 = "fa04f707debb75087ea2bf5e327602034eaa3a6900421f2cf32ad5f5f1c887b9",
             ),
             struct(
-                name = "miopen",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/m/miopen-hip6.2.0/miopen-hip6.2.0_3.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "2dbf6d126d0de6930e0cd94d0e525e07d3019d90bd7256f3151a7f1fbc2250af",
             ),
             struct(
-                name = "miopen-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/m/miopen-hip-dev/miopen-hip-dev_3.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "df5fdd2218e4d380b133ba402f3734fbe0589d9cdd8618a101b71b968909b4ba",
             ),
             struct(
-                name = "rccl",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rccl6.2.0/rccl6.2.0_2.20.5.60200-66~24.04_amd64.deb",
                 sha256 = "4d7efa4ee6aa2bf69b0aab449cc1d01c25ca65814e1b3cb07f6b59fa8b1608b8",
             ),
             struct(
-                name = "rccl-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rccl-dev6.2.0/rccl-dev6.2.0_2.20.5.60200-66~24.04_amd64.deb",
                 sha256 = "4ab4f880344e04d61b6fa746be5c4bdc2841409fb6987ee61e39c6420b4eca42",
             ),
             struct(
-                name = "rocblas",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocblas6.2.0/rocblas6.2.0_4.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "521c87ce396c6ce10076cc641b6035451fd68ddb36a684c5a9c9538dfc831ade",
             ),
             struct(
-                name = "rocblas-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocblas-dev/rocblas-dev_4.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "00f135ce2ae47c35085ef06248ff7d5ce8c12fd0d5b82e7bd77b1dbc0ce7058e",
             ),
             struct(
-                name = "rocfft-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocfft-dev6.2.0/rocfft-dev6.2.0_1.0.28.60200-66~24.04_amd64.deb",
                 sha256 = "eb3904263b396d46799eeea1081d8e8d1a551a890432a803364db2d013849f92",
             ),
             struct(
-                name = "rocm_core",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocm-core/rocm-core_6.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "af5fcbe8dc2b6cbec30e2d39d30736e8a47a0b9d0ca2be7f179f2947f9c98245",
             ),
             struct(
-                name = "rocm_hip_libraries",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocm-hip-libraries/rocm-hip-libraries_6.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "228f07a3caefc41f6efd5345eb9d3630f1db769f9b4abd1313cbcb32d077ce53",
             ),
             struct(
-                name = "rocm_hip_dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/h/hip-dev/hip-dev_6.2.41133.60200-66~24.04_amd64.deb",
                 sha256 = "cda72054d2011dbb062e75386766d928fd8905c15c88685c3ef87fc963bd88ad",
             ),
             struct(
-                name = "rocm-device-libs",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocm-device-libs6.2.0/rocm-device-libs6.2.0_1.0.0.60200-66~24.04_amd64.deb",
                 sha256 = "298544f717dfb236b9257b19a0ab81abaaa770128976d4abfdea546cd32d8b02",
             ),
             struct(
-                name = "rocminfo",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocminfo6.2.0/rocminfo6.2.0_1.0.0.60200-66~24.04_amd64.deb",
                 sha256 = "8e78ed8e480b55a496153b150acb22bab39c3bb8cf1e62f9aff7eaf75a3a3a92",
             ),
             struct(
-                name = "rocm-llvm",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocm-llvm6.2.0/rocm-llvm6.2.0_18.0.0.24292.60200-66~24.04_amd64.deb",
                 sha256 = "72c388eae7c0f54151b46fbd8fa6e26f1ca81e2b8b415c43411a156b3f25b6e7",
             ),
             struct(
-                name = "rocm-llvm-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocm-llvm-dev/rocm-llvm-dev_18.0.0.24292.60200-66~24.04_amd64.deb",
                 sha256 = "3e85a859c5dafa82a9a57dda096d566b821217bacfac995f7cc45ed460b68999",
             ),
             struct(
-                name = "rocprim-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocprim-dev6.2.0/rocprim-dev6.2.0_3.2.0.60200-66~24.04_amd64.deb",
                 sha256 = "6c832e2feb0885fbe481245825c76a466921b294f530eb0d0da70a44cfe6e608",
             ),
             struct(
-                name = "rocrand-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocrand-dev/rocrand-dev_3.1.0.60200-66~24.04_amd64.deb",
                 sha256 = "2a2a95185ce0e54df226474b2f5cfcdc9e5ede5a6d88a8a70c2635ea2237abba",
             ),
             struct(
-                name = "roctracer",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/roctracer6.2.0/roctracer6.2.0_4.1.60200.60200-66~24.04_amd64.deb",
                 sha256 = "2f2fb6f8d06ace89131934c833b0ea359335a4b45aeec1559b293d7bc14b1d1d",
             ),
             struct(
-                name = "roctracer-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/roctracer-dev6.2.0/roctracer-dev6.2.0_4.1.60200.60200-66~24.04_amd64.deb",
                 sha256 = "c6c781ee87c459aed32e943b389137f98ecd402fb83a3d1c98de9a76abadc3a3",
             ),
             struct(
-                name = "rocsolver",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocsolver6.2.0/rocsolver6.2.0_3.26.0.60200-66~24.04_amd64.deb",
                 sha256 = "5e4b3e38556f0826e5322971635a49a72283d60862ccc4d28efd11c8fb955b47",
             ),
             struct(
-                name = "rocsolver-dev",
                 url = "https://repo.radeon.com/rocm/apt/6.2/pool/main/r/rocsolver-dev6.2.0/rocsolver-dev6.2.0_3.26.0.60200-66~24.04_amd64.deb",
                 sha256 = "5bb6ae92a25f33488f2ee5f123ac4f67ad130e18e4949161715451509be3b89d",
             ),