Make cuda-runtime-wrappers only depend on LLVM Support headers, not the library. This allows loading other shared libraries which include LLVM Support (i.e. the async runtime) along with the cuda-runtime-wrappers.

See https://reviews.llvm.org/D95613 for how this works for the cmake build.

PiperOrigin-RevId: 355427339
Change-Id: If56db748b39026dedea3b3ab4e81e538c0436aae

third_party/mlir/BUILD

@@ -3,6 +3,7 @@
 
 load("@org_tensorflow//third_party/mlir:tblgen.bzl", "gentbl")
 load("@org_tensorflow//third_party/mlir:linalggen.bzl", "genlinalg")
+load("@org_tensorflow//third_party/mlir:build_defs.bzl", "cc_headers_only")
 
 package(
     default_visibility = [":friends"],
@@ -3925,17 +3926,33 @@ cc_binary(
     ],
 )
 
+# This target provides the headers from LLVM's Support target without any of
+# the symbols. In particular, it does not contain the static registration code
+# which may be executed by at most one shared library loaded by ORCJit. Direct
+# dependencies need to avoid requiring symbols from LLVMSupport by adding
+# copts = ["-DLLVM_DISABLE_ABI_BREAKING_CHECKS_ENFORCING=1"].
+#
+# Bazel links the dependencies' object files instead of the archives, which
+# means that symbols are linked in even if none are used. The LLVM cmake build
+# on the other hand links archives (or shared libraries, depending on
+# BUILD_SHARED_LIBS), skipping them if none of the symbols are used.
+# See also https://reviews.llvm.org/D95613.
+cc_headers_only(
+    name = "LLVMSupportHeaders",
+    src = "@llvm-project//llvm:Support",
+)
+
 cc_binary(
     name = "tools/libcuda-runtime-wrappers.so",
     srcs = ["tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp"],
+    # Prevent needing EnableABIBreakingChecks symbol from LLVMSupport.
+    copts = ["-DLLVM_DISABLE_ABI_BREAKING_CHECKS_ENFORCING=1"],
     linkshared = True,
     deps = [
+        ":LLVMSupportHeaders",
         ":mlir_c_runner_utils",
         "//third_party/gpus/cuda:cuda_headers",
         "//third_party/gpus/cuda:libcuda",
-        # Note: for headers only, the archive itself should not be linked.
-        # See https://reviews.llvm.org/D95613 for details.
-        "@llvm-project//llvm:Support",
     ],
 )
 

third_party/mlir/build_defs.bzl

@@ -0,0 +1,14 @@
+def _cc_headers_only_impl(ctx):
+    return CcInfo(compilation_context = ctx.attr.src[CcInfo].compilation_context)
+
+cc_headers_only = rule(
+    implementation = _cc_headers_only_impl,
+    attrs = {
+        "src": attr.label(
+            mandatory = True,
+            providers = [CcInfo],
+        ),
+    },
+    doc = "Provides the headers from 'src' without linking anything.",
+    provides = [CcInfo],
+)