[Mosaic GPU][NFC] Create undefined values in test with ub.poison.

`ub.PoisonOp` and `llvm.UndefOp` serve a similar purpose in tests. However, `llvm.UndefOp` is more restrictive as it can only creates an undefined value of the specified **LLVM IR** dialect type.

PiperOrigin-RevId: 818583066

Author: allanrenucci

jax/_src/lib/mlir/dialects/__init__.py

@@ -56,8 +56,10 @@
 # TODO(joelwee): Remove this once jaxlib 0.8 is the minimum.
 try:
   from jaxlib.mlir.dialects import mpmd
+  from jaxlib.mlir.dialects import ub
 except ImportError:
   mpmd: Any = None  # type: ignore[no-redef]
+  ub: Any = None  # type: ignore[no-redef]
 from jaxlib.mlir.dialects import sdy
 
 # Alias that is set up to abstract away the transition from MHLO to StableHLO.

jaxlib/BUILD

@@ -86,6 +86,7 @@ pytype_strict_library(
         "//jaxlib/mlir:sdy_dialect",
         "//jaxlib/mlir:sparse_tensor_dialect",
         "//jaxlib/mlir:stablehlo_dialect",
+        "//jaxlib/mlir:ub_dialect",
         "//jaxlib/mlir:vector_dialect",
         "//jaxlib/mlir/_mlir_libs:_jax_mlir_ext",
         "//jaxlib/mosaic",

jaxlib/mlir/BUILD.bazel

@@ -87,6 +87,19 @@ symlink_inputs(
     ],
 )
 
+symlink_inputs(
+    name = "ub_dialect",
+    rule = py_library,
+    symlinked_inputs = {"srcs": {
+        "dialects": ["@llvm-project//mlir/python:UbPyFiles"],
+    }},
+    deps = [
+        ":core",
+        ":ir",
+        ":mlir",
+    ],
+)
+
 symlink_inputs(
     name = "math_dialect",
     rule = py_library,

jaxlib/tools/build_wheel.py

@@ -311,6 +311,7 @@ def prepare_wheel(wheel_sources_path: pathlib.Path, *, cpu, wheel_sources):
           f"{source_file_prefix}jaxlib/mlir/dialects/sdy.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/sparse_tensor.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/stablehlo.py",
+          f"{source_file_prefix}jaxlib/mlir/dialects/ub.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/vector.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/nvgpu.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/nvvm.py",

tests/mosaic/gpu_dialect_test.py

@@ -24,12 +24,12 @@
 from jax._src.interpreters import mlir as mlir_interpreter
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import arith
-from jax._src.lib.mlir.dialects import builtin
 from jax._src.lib.mlir.dialects import gpu
 from jax._src.lib.mlir.dialects import llvm
 from jax._src.lib.mlir.dialects import memref
 from jax._src.lib.mlir.dialects import nvvm
 from jax._src.lib.mlir.dialects import scf
+from jax._src.lib.mlir.dialects import ub
 from jax._src.lib.mlir.dialects import vector
 from jax.experimental.mosaic import gpu as mgpu
 from jax.experimental.mosaic.gpu import dialect_lowering as lowering
@@ -86,8 +86,7 @@ def workgroup_ptr_ty() -> ir.Type:
 
 def undefs(*tys: ir.Type) -> list[ir.Value]:
   """Returns a list of undefined values of the given types."""
-  # TODO(allanrenucci): Use `ub.poison` once Python bindings are available.
-  return [builtin.unrealized_conversion_cast([ty], []) for ty in tys]
+  return [ub.poison(ty) for ty in tys]
 
 
 class MosaicGpuTest(parameterized.TestCase):