diff --git a/runtimes/rocm/libpjrt_rocm.BUILD.bazel b/runtimes/rocm/libpjrt_rocm.BUILD.bazel index 4683699..1625740 100644 --- a/runtimes/rocm/libpjrt_rocm.BUILD.bazel +++ b/runtimes/rocm/libpjrt_rocm.BUILD.bazel @@ -65,7 +65,6 @@ copy_to_directory( "@rocblas//:runfiles", "@rocm-core", "@rocm-device-libs//:runfiles", - "@rocm-llvm//:lld", "@rocm-smi-lib//:rocm_smi", "@rocprofiler-register", "@rocfft", @@ -79,12 +78,6 @@ copy_to_directory( "@libdrm-amdgpu1", "@libtinfo6", "@zlib1g", - - # lld dependencies - "@libxml2", - "@libicu70//:libicuuc70", - "@libicu70//:libicudata70", - "@liblzma5", ] + select({ ":_hipblaslt": ["@hipblaslt//:runfiles"], "//conditions:default": [], @@ -97,10 +90,6 @@ copy_to_directory( "libelf1": "lib", "hipblaslt": "lib", "rocblas": "lib", - "libxml2": "lib", - "libicuuc70": "lib", - "liblzma5": "lib", - "lld": "llvm/bin", }, add_directory_to_runfiles = True, include_external_repositories = ["**"], diff --git a/runtimes/rocm/packages.lock.json b/runtimes/rocm/packages.lock.json index c8bd357..ebc49d9 100755 --- a/runtimes/rocm/packages.lock.json +++ b/runtimes/rocm/packages.lock.json @@ -4315,69 +4315,6 @@ "https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocfft/rocfft_1.0.31.60304-76~22.04_amd64.deb" ], "version": "1.0.31.60304-76~22.04" - }, - { - "arch": "amd64", - "dependencies": [ - { - "key": "zlib1g_1-1.2.11.dfsg-2ubuntu9.2_amd64", - "name": "zlib1g", - "version": "1:1.2.11.dfsg-2ubuntu9.2" - }, - { - "key": "libc6_2.35-0ubuntu3.10_amd64", - "name": "libc6", - "version": "2.35-0ubuntu3.10" - }, - { - "key": "libcrypt1_1-4.4.27-1_amd64", - "name": "libcrypt1", - "version": "1:4.4.27-1" - }, - { - "key": "libgcc-s1_12.3.0-1ubuntu1_22.04_amd64", - "name": "libgcc-s1", - "version": "12.3.0-1ubuntu1~22.04" - }, - { - "key": "gcc-12-base_12.3.0-1ubuntu1_22.04_amd64", - "name": "gcc-12-base", - "version": "12.3.0-1ubuntu1~22.04" - }, - { - "key": "liblzma5_5.2.5-2ubuntu1_amd64", - "name": "liblzma5", - "version": "5.2.5-2ubuntu1" - }, - { - "key": "libicu70_70.1-2_amd64", - "name": "libicu70", - "version": "70.1-2" - }, - { - "key": "libstdc-p--p-6_12.3.0-1ubuntu1_22.04_amd64", - "name": "libstdc++6", - "version": "12.3.0-1ubuntu1~22.04" - } - ], - "key": "libxml2_2.9.13-p-dfsg-1ubuntu0.7_amd64", - "name": "libxml2", - "sha256": "b53f47fa4ecdbab12b05250bb13b01b0bd42936323e6484b038298e45e2d5c10", - "urls": [ - "https://snapshot.ubuntu.com/ubuntu/20250711T030400Z/pool/main/libx/libxml2/libxml2_2.9.13+dfsg-1ubuntu0.7_amd64.deb" - ], - "version": "2.9.13+dfsg-1ubuntu0.7" - }, - { - "arch": "amd64", - "dependencies": [], - "key": "libicu70_70.1-2_amd64", - "name": "libicu70", - "sha256": "58a154f6307289813da2276f900498ef536ae7c0522d2cf31a3c3c5cf62dfd9a", - "urls": [ - "https://snapshot.ubuntu.com/ubuntu/20250711T030400Z/pool/main/i/icu/libicu70_70.1-2_amd64.deb" - ], - "version": "70.1-2" } ], "version": 1 diff --git a/runtimes/rocm/packages.yaml b/runtimes/rocm/packages.yaml index 2b23f76..a9afb53 100644 --- a/runtimes/rocm/packages.yaml +++ b/runtimes/rocm/packages.yaml @@ -35,6 +35,3 @@ packages: - "hipblaslt" # - "hipblaslt-dev" - "hip-runtime-amd" - # - "rocm-llvm" - # rocm-llvm > ld.ldd missing dependency - - "libxml2" diff --git a/runtimes/rocm/rocm.bzl b/runtimes/rocm/rocm.bzl index dea799c..14646d7 100644 --- a/runtimes/rocm/rocm.bzl +++ b/runtimes/rocm/rocm.bzl @@ -30,24 +30,6 @@ _UBUNTU_PACKAGES = { ]), "libtinfo6": packages.filegroup(name = "libtinfo6", srcs = ["lib/x86_64-linux-gnu/libtinfo.so.6"]), "zlib1g": packages.filegroup(name = "zlib1g", srcs = ["lib/x86_64-linux-gnu/libz.so.1"]), - "liblzma5": packages.filegroup(name = "liblzma5", srcs = ["lib/x86_64-linux-gnu/liblzma.so.5"]), - "libxml2": "\n".join([ - packages.load_("@zml//bazel:patchelf.bzl", "patchelf"), - packages.patchelf( - name = "libxml2", - shared_library = "usr/lib/x86_64-linux-gnu/libxml2.so.2", - set_rpath = '$ORIGIN', - ), - ]), - "libicu70": "\n".join([ - packages.load_("@zml//bazel:patchelf.bzl", "patchelf"), - packages.patchelf( - name = "libicuuc70", - shared_library = "usr/lib/x86_64-linux-gnu/libicuuc.so.70", - set_rpath = '$ORIGIN', - ), - packages.filegroup(name = "libicudata70", srcs = ["usr/lib/x86_64-linux-gnu/libicudata.so.70"]) - ]), } _ROCM_PACKAGES = { @@ -130,15 +112,6 @@ _ROCM_PACKAGES = { packages.filegroup(name = "hiprtc", srcs = ["lib/libhiprtc.so.6"]), ]), "hipsolver": packages.filegroup(name = "hipsolver", srcs = ["lib/libhipsolver.so.0"]), - "rocm-llvm": "\n".join([ - packages.load_("@zml//bazel:patchelf.bzl", "patchelf"), - packages.patchelf( - name = "lld", - #TODO: Rename attr to elf_file or file ? - shared_library = "llvm/bin/ld.lld", - set_rpath = '$ORIGIN/../../lib', - ), - ]), } def _rocm_impl(mctx): diff --git a/zml/module.zig b/zml/module.zig index c6824b4..8ab1aed 100644 --- a/zml/module.zig +++ b/zml/module.zig @@ -898,6 +898,9 @@ fn compileModuleToPjrtExecutable(arena: std.mem.Allocator, platform: Platform, m // Disable Triton GEMM on ROCM. For some reason it's much, much slower when // enabled on CDNA and it's used on RDNA. Disable it altogether. setFlag(&options, "xla_gpu_enable_triton_gemm", false); + // Use lld from libllvm instead of invoking the ld.lld binary. + // This saves us from having to sandbox it. + setFlag(&options, "xla_gpu_use_inprocess_lld", true); }, else => {}, }