[mosaic_gpu] Added support for using cf.assert in Mosaic GPU kernels

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 `cf.assert` support 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/BUILD

@@ -82,6 +82,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:gpu_dialect",
         "//jaxlib/mlir:ir",

jaxlib/mlir/BUILD.bazel

@@ -65,7 +65,7 @@ symlink_inputs(
     name = "func_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@llvm-project//mlir/python:FuncPyFiles"],
+        "dialects": ["@llvm-project//mlir/python:FuncPyFiles"],
     }},
     deps = [
         ":core",
@@ -78,7 +78,7 @@ symlink_inputs(
     name = "vector_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@llvm-project//mlir/python:VectorOpsPyFiles"],
+        "dialects": ["@llvm-project//mlir/python:VectorOpsPyFiles"],
     }},
     deps = [
         ":core",
@@ -91,7 +91,7 @@ symlink_inputs(
     name = "math_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@llvm-project//mlir/python:MathOpsPyFiles"],
+        "dialects": ["@llvm-project//mlir/python:MathOpsPyFiles"],
     }},
     deps = [
         ":core",
@@ -104,7 +104,7 @@ symlink_inputs(
     name = "arithmetic_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@llvm-project//mlir/python:ArithOpsPyFiles"],
+        "dialects": ["@llvm-project//mlir/python:ArithOpsPyFiles"],
     }},
     deps = [
         ":core",
@@ -117,7 +117,20 @@ symlink_inputs(
     name = "memref_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@llvm-project//mlir/python:MemRefOpsPyFiles"],
+        "dialects": ["@llvm-project//mlir/python:MemRefOpsPyFiles"],
+    }},
+    deps = [
+        ":core",
+        ":ir",
+        ":mlir",
+    ],
+)
+
+symlink_inputs(
+    name = "control_flow_dialect",
+    rule = py_library,
+    symlinked_inputs = {"srcs": {
+        "dialects": ["@llvm-project//mlir/python:ControlFlowOpsPyFiles"],
     }},
     deps = [
         ":core",
@@ -130,7 +143,7 @@ symlink_inputs(
     name = "scf_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@llvm-project//mlir/python:SCFPyFiles"],
+        "dialects": ["@llvm-project//mlir/python:SCFPyFiles"],
     }},
     deps = [
         ":core",
@@ -143,7 +156,7 @@ symlink_inputs(
     name = "builtin_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@llvm-project//mlir/python:BuiltinOpsPyFiles"],
+        "dialects": ["@llvm-project//mlir/python:BuiltinOpsPyFiles"],
     }},
     deps = [
         ":core",
@@ -157,7 +170,7 @@ symlink_inputs(
     name = "chlo_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@stablehlo//:chlo_ops_py_files"],
+        "dialects": ["@stablehlo//:chlo_ops_py_files"],
     }},
     deps = [
         ":core",
@@ -171,7 +184,7 @@ symlink_inputs(
     name = "sparse_tensor_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@llvm-project//mlir/python:SparseTensorOpsPyFiles"],
+        "dialects": ["@llvm-project//mlir/python:SparseTensorOpsPyFiles"],
     }},
     deps = [
         ":core",
@@ -186,7 +199,7 @@ symlink_inputs(
     name = "mhlo_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@xla//xla/mlir_hlo:MhloOpsPyFiles"],
+        "dialects": ["@xla//xla/mlir_hlo:MhloOpsPyFiles"],
     }},
     deps = [
         ":core",
@@ -228,7 +241,7 @@ symlink_inputs(
     name = "stablehlo_dialect",
     rule = py_library,
     symlinked_inputs = {"srcs": {
-      "dialects": ["@stablehlo//:stablehlo_ops_py_files"],
+        "dialects": ["@stablehlo//:stablehlo_ops_py_files"],
     }},
     deps = [
         ":core",

jaxlib/mlir/_mlir_libs/BUILD.bazel

@@ -208,6 +208,7 @@ nanobind_pywrap_extension(
     deps = [
         "//jaxlib/mosaic/gpu:mlir_capi",
         "@llvm-project//mlir:CAPIArith",
+        "@llvm-project//mlir:CAPICF",
         "@llvm-project//mlir:CAPIGPU",
         "@llvm-project//mlir:CAPIIR",
         "@llvm-project//mlir:CAPILLVM",
@@ -297,4 +298,4 @@ nanobind_pywrap_extension(
         "@nanobind",
         "@stablehlo//:stablehlo_capi",
     ],
-)
+)

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);

jaxlib/tools/build_wheel.py

@@ -272,6 +272,7 @@ def prepare_wheel(wheel_sources_path: pathlib.Path, *, cpu, wheel_sources):
           f"{source_file_prefix}jaxlib/mlir/dialects/_arith_enum_gen.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/_arith_ops_gen.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/_builtin_ops_gen.py",
+          f"{source_file_prefix}jaxlib/mlir/dialects/_cf_ops_gen.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/_chlo_ops_gen.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/_func_ops_gen.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/_math_ops_gen.py",
@@ -296,6 +297,7 @@ def prepare_wheel(wheel_sources_path: pathlib.Path, *, cpu, wheel_sources):
           f"{source_file_prefix}jaxlib/mlir/dialects/_llvm_ops_gen.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/arith.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/builtin.py",
+          f"{source_file_prefix}jaxlib/mlir/dialects/cf.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/chlo.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/func.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/math.py",

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):