Merge pull request jax-ml#25710 from apaszke:mgpu_dialect_fix

PiperOrigin-RevId: 711430610

Author: Google-ML-Automation

jaxlib/mlir/_mlir_libs/BUILD.bazel

@@ -158,9 +158,10 @@ py_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi",
+        ":jaxlib_mlir_capi_shared_library",
+        "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi_headers",
         "@llvm-project//mlir:CAPIIRHeaders",
-        "@llvm-project//mlir:MLIRBindingsPythonHeadersAndDeps",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeadersAndDeps",
         "@nanobind",
     ],
 )
@@ -380,6 +381,7 @@ cc_library(
     name = "jaxlib_mlir_capi_objects",
     deps = [
         "//jaxlib/mosaic:tpu_dialect_capi_objects",
+        "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi_objects",
         "@llvm-project//mlir:CAPIArithObjects",
         "@llvm-project//mlir:CAPIGPUObjects",
         "@llvm-project//mlir:CAPIIRObjects",

jaxlib/mosaic/dialect/gpu/BUILD

@@ -215,3 +215,26 @@ cc_library(
         "@llvm-project//mlir:CAPIIR",
     ],
 )
+
+# Header-only target, used when using the C API from a separate shared library.
+cc_library(
+    name = "gpu_dialect_capi_headers",
+    hdrs = DIALECT_CAPI_HEADERS,
+    deps = [
+        ":mosaic_gpu_inc_gen",
+        "@llvm-project//mlir:CAPIIRHeaders",
+    ],
+)
+
+# Alwayslink target, used when exporting the C API from a shared library.
+cc_library(
+    name = "gpu_dialect_capi_objects",
+    srcs = DIALECT_CAPI_SOURCES,
+    hdrs = DIALECT_CAPI_HEADERS,
+    deps = [
+        ":mosaic_gpu",
+        ":mosaic_gpu_inc_gen",
+        "@llvm-project//mlir:CAPIIRObjects",
+    ],
+    alwayslink = True,
+)

jaxlib/mosaic/python/mosaic_gpu.py

@@ -33,4 +33,5 @@
   from mlir.dialects._ods_common import _cext  # type: ignore[import-not-found]
 
 
-_cext.globals.append_dialect_search_prefix("jax.jaxlib.mosaic.python")
+# Add the parent module to the search prefix
+_cext.globals.append_dialect_search_prefix(__name__[:__name__.rfind(".")])