[mosaic_gpu] Added debug_assert

The XLA GPU runtime does not yet handle device assertions well and will hang
if the assert is triggered. However, the assertion output still appears in
stderr, so I think having `debug_assert` is still useful.

PiperOrigin-RevId: 740717697

Author: superbobry

jax/BUILD

@@ -866,6 +866,7 @@ py_library_providing_imports_info(
         "//jax/_src/lib",
         "//jaxlib/mlir:arithmetic_dialect",
         "//jaxlib/mlir:builtin_dialect",
+        "//jaxlib/mlir:control_flow_dialect",
         "//jaxlib/mlir:func_dialect",
         "//jaxlib/mlir:gpu_dialect",
         "//jaxlib/mlir:ir",

jax/_src/lib/BUILD

@@ -52,6 +52,7 @@ py_library_providing_imports_info(
         "//jaxlib/mlir:arithmetic_dialect",
         "//jaxlib/mlir:builtin_dialect",
         "//jaxlib/mlir:chlo_dialect",
+        "//jaxlib/mlir:control_flow_dialect",
         "//jaxlib/mlir:func_dialect",
         "//jaxlib/mlir:ir",
         "//jaxlib/mlir:math_dialect",

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

@@ -55,3 +55,9 @@
 
 # Alias that is set up to abstract away the transition from MHLO to StableHLO.
 from jaxlib.mlir.dialects import stablehlo as hlo
+
+from jax._src import lib
+if lib.version >= (0, 6, 1):
+  from jaxlib.mlir.dialects import cf
+else:
+  cf = None  # type: ignore[no-redef]

jaxlib/mlir/_mlir_libs/register_jax_dialects.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include <nanobind/nanobind.h>
 
 #include "mlir-c/Dialect/Arith.h"  // IWYU pragma: keep
+#include "mlir-c/Dialect/ControlFlow.h"
 #include "mlir-c/Dialect/Func.h"  // IWYU pragma: keep
 #include "mlir-c/Dialect/GPU.h"  // IWYU pragma: keep
 #include "mlir-c/Dialect/LLVM.h"  // IWYU pragma: keep
@@ -50,6 +51,7 @@ NB_MODULE(register_jax_dialects, m) {
     REGISTER_DIALECT(scf);
     REGISTER_DIALECT(vector);
     // For Mosaic GPU
+    REGISTER_DIALECT(cf);
     REGISTER_DIALECT(gpu);
     REGISTER_DIALECT(nvgpu);
     REGISTER_DIALECT(nvvm);

jaxlib/mosaic/gpu/BUILD

@@ -155,6 +155,7 @@ cc_library(
         "@llvm-project//mlir:ArithTransforms",
         "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
         "@llvm-project//mlir:ComplexToLLVM",
+        "@llvm-project//mlir:ControlFlowDialect",
         "@llvm-project//mlir:ControlFlowToLLVM",
         "@llvm-project//mlir:ConversionPasses",
         "@llvm-project//mlir:ExecutionEngine",

jaxlib/mosaic/gpu/custom_call.cc

@@ -60,6 +60,7 @@ limitations under the License.
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
@@ -228,12 +229,12 @@ mlir::LogicalResult RunPasses(mlir::OpPassManager&& passes,
 
 void InitContext(mlir::MLIRContext* context) {
   mlir::DialectRegistry registry;
-  registry.insert<mlir::arith::ArithDialect, mlir::func::FuncDialect,
-                  mlir::math::MathDialect, mlir::memref::MemRefDialect,
-                  mlir::scf::SCFDialect, mlir::vector::VectorDialect,
-                  mlir::gpu::GPUDialect, mlir::nvgpu::NVGPUDialect,
-                  mlir::NVVM::NVVMDialect, mlir::LLVM::LLVMDialect,
-                  mosaic_gpu::MosaicGPUDialect>();
+  registry.insert<mlir::arith::ArithDialect, mlir::cf::ControlFlowDialect,
+                  mlir::func::FuncDialect, mlir::math::MathDialect,
+                  mlir::memref::MemRefDialect, mlir::scf::SCFDialect,
+                  mlir::vector::VectorDialect, mlir::gpu::GPUDialect,
+                  mlir::nvgpu::NVGPUDialect, mlir::NVVM::NVVMDialect,
+                  mlir::LLVM::LLVMDialect, mosaic_gpu::MosaicGPUDialect>();
   mlir::registerConvertNVVMToLLVMInterface(registry);
   mlir::registerConvertComplexToLLVMInterface(registry);
   mlir::registerConvertMemRefToLLVMInterface(registry);

tests/mosaic/gpu_test.py

@@ -32,6 +32,7 @@
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir import passmanager
 from jax._src.lib.mlir.dialects import arith
+from jax._src.lib.mlir.dialects import cf
 from jax._src.lib.mlir.dialects import scf
 from jax._src.lib.mlir.dialects import vector
 from jax.experimental.mosaic.gpu import dialect as mgpu_dialect  # pylint: disable=g-importing-member
@@ -237,7 +238,6 @@ def capture_stdout(self):
       mosaic_gpu_lib._mosaic_gpu_ext._sync_all_devices()
 
 
-
 class Sm90ATestCase(TestCase, jtu.CudaArchSpecificTest):
 
   def setUp(self):
@@ -3320,6 +3320,34 @@ def test_parse_indices_oob(self, indices):
     with self.assertRaisesRegex(IndexError, "out of bounds"):
       utils.parse_indices(indices, (2, 3, 4))
 
+  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
+  def test_assert(self):
+    if cf is None:
+      self.skipTest("``cf`` is not available")
+
+    def kernel(ctx: mgpu.LaunchContext, x_ref, out, scratch) -> None:
+      del ctx, out  # Unused.
+      # TODO(b/408271232): Use a False condition once the bug is fixed.
+      x = mgpu.FragmentedArray.load_strided(x_ref)
+      cond = x.reduce_sum(*scratch) != 42.0
+      cf.assert_(cond.registers.item(), "OOOPS")
+
+    f = mgpu.as_gpu_kernel(
+        kernel,
+        grid=(1, 1, 1),
+        block=(128, 1, 1),
+        in_shape=(jax.ShapeDtypeStruct((128,), jnp.float32),),
+        out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
+        smem_scratch_shape=(jax.ShapeDtypeStruct((4,), jnp.float32),),
+    )
+
+    with jtu.set_env(MOSAIC_GPU_DUMP_SASS="1"), self.capture_stdout() as sass:
+      f(jnp.ones((128,), jnp.float32))
+
+    # SASS doesn't seem to include the assertion message, so we are just
+    # checking that __assertfail appears in the symbol table for the kernel.
+    self.assertIn("__assertfail", sass())
+
 
 class SerializationTest(absltest.TestCase):