From 980f1b17fb557048aeadf424b4b5978dbb03556f Mon Sep 17 00:00:00 2001
From: Tarry Singh <tarry.singh@deepkapha.com>
Date: Mon, 11 Mar 2024 10:15:22 +0000
Subject: [PATCH] Ensure all runtime plugins have correct SONAME values, fixing
 issues with prebuilt PJRT plugins.

---
 runtimes/cpu/cpu.bzl                   | 20 +++++++++++++++-----
 runtimes/rocm/libpjrt_rocm.BUILD.bazel |  1 +
 runtimes/tpu/libpjrt_tpu.BUILD.bazel   | 12 +++---------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/runtimes/cpu/cpu.bzl b/runtimes/cpu/cpu.bzl
index 31dfafe..bc5648d 100644
--- a/runtimes/cpu/cpu.bzl
+++ b/runtimes/cpu/cpu.bzl
@@ -1,24 +1,34 @@
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
-_BUILD = """\
+_BUILD_LINUX = """\
+load("@zml//bazel:cc_import.bzl", "cc_import")
 cc_import(
     name = "libpjrt_cpu",
-    shared_library = "libpjrt_cpu.{ext}",
-    visibility = ["//visibility:public"],
+    shared_library = "libpjrt_cpu.so",
+    soname = "libpjrt_cpu.so",
+    visibility = ["@zml//runtimes/cpu:__subpackages__"],
+)
+"""
+
+_BUILD_DARWIN = """\
+cc_import(
+    name = "libpjrt_cpu",
+    shared_library = "libpjrt_cpu.dylib",
+    visibility = ["@zml//runtimes/cpu:__subpackages__"],
 )
 """
 
 def _cpu_pjrt_plugin_impl(mctx):
     http_archive(
         name = "libpjrt_cpu_linux_amd64",
-        build_file_content = _BUILD.format(ext = "so"),
+        build_file_content = _BUILD_LINUX,
         sha256 = "0f2cb204015e062df5d1cbd39d8c01c076ab2b004d0f4f37f6d5e120d3cd7087",
         url = "https://github.com/zml/pjrt-artifacts/releases/download/v5.0.0/pjrt-cpu_linux-amd64.tar.gz",
     )
 
     http_archive(
         name = "libpjrt_cpu_darwin_arm64",
-        build_file_content = _BUILD.format(ext = "dylib"),
+        build_file_content = _BUILD_DARWIN,
         sha256 = "2ddb66a93c8a913e3bc8f291e01df59aa297592cc91e05aab2dd4813884098cb",
         url = "https://github.com/zml/pjrt-artifacts/releases/download/v5.0.0/pjrt-cpu_darwin-arm64.tar.gz",
     )
diff --git a/runtimes/rocm/libpjrt_rocm.BUILD.bazel b/runtimes/rocm/libpjrt_rocm.BUILD.bazel
index f19ab7f..506255e 100644
--- a/runtimes/rocm/libpjrt_rocm.BUILD.bazel
+++ b/runtimes/rocm/libpjrt_rocm.BUILD.bazel
@@ -65,6 +65,7 @@ cc_import(
         "dlopen": "zmlxrocm_dlopen",
     },
     shared_library = "libpjrt_rocm.so",
+    soname = "libpjrt_rocm.so",
     visibility = ["//visibility:public"],
     deps = [
         "@comgr//:amd_comgr",
diff --git a/runtimes/tpu/libpjrt_tpu.BUILD.bazel b/runtimes/tpu/libpjrt_tpu.BUILD.bazel
index b4385aa..2fad364 100644
--- a/runtimes/tpu/libpjrt_tpu.BUILD.bazel
+++ b/runtimes/tpu/libpjrt_tpu.BUILD.bazel
@@ -1,14 +1,8 @@
-load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
-
-copy_file(
-    name = "libpjrt_tpu_so",
-    src = "libtpu/libtpu.so",
-    out = "libpjrt_tpu.so",
-    allow_symlink = True,
-)
+load("@zml//bazel:cc_import.bzl", "cc_import")
 
 cc_import(
     name = "libpjrt_tpu",
-    shared_library = ":libpjrt_tpu_so",
+    shared_library = "libtpu/libtpu.so",
+    soname = "libpjrt_tpu.so",
     visibility = ["@zml//runtimes/tpu:__subpackages__"],
 )