[Mosaic GPU] Implement the skeleton of a lowering pass for the Mosaic GPU dialect.

Also add a lowering rule for `mosaic_gpu.initialize_barrier`.

PiperOrigin-RevId: 694276698

Author: bchetioui

jax/experimental/mosaic/gpu/__init__.py

@@ -27,6 +27,12 @@
     Union as Union,
     as_gpu_kernel as as_gpu_kernel,
 )
+
+if dialect is not None:
+  from .dialect_lowering import lower_mgpu_dialect
+else:
+  lower_mgpu_dialect = None
+
 from .fragmented_array import (
     FragmentedArray as FragmentedArray,
     FragmentedLayout as FragmentedLayout,

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -0,0 +1,125 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Lowering rules and pass for the MLIR Mosaic GPU dialect."""
+
+from collections.abc import Callable
+import functools
+import operator
+from typing import Sequence, Type
+
+from jax._src.interpreters import mlir as mlir_interpreter
+from jax._src.lib import mosaic_gpu_dialect as mgpu
+
+from jaxlib.mlir import ir
+from jaxlib.mlir.dialects import gpu
+from jaxlib.mlir.dialects import llvm
+from jaxlib.mlir.dialects import memref
+from jaxlib.mlir.dialects import nvvm
+from .utils import c, memref_ptr, single_thread_predicate
+
+
+MlirLoweringRule = Callable[[ir.Operation | ir.OpView], Sequence[ir.Value]]
+
+
+_lowerings: dict[str, MlirLoweringRule] = {}
+
+
+# TODO(bchetioui): Remove this when minimum jaxlib version >= 0.4.36.
+# Jaxlib doesn't contain Mosaic GPU dialect bindings.
+InitializeBarrierOp = mgpu.InitializeBarrierOp if mgpu is not None else None
+
+def _register_lowering(
+    op: str | Type[ir.OpView]
+) -> Callable[[MlirLoweringRule], MlirLoweringRule]:
+  def wrapper(f):
+    op_name = op if isinstance(op, str) else op.OPERATION_NAME  # pytype: disable=attribute-error
+    _lowerings[op_name] = f
+    return f
+
+  return wrapper
+
+
+def _lowered_barrier_type() -> ir.Type:
+  return ir.IntegerType.get_signless(64)
+
+
+def _gpu_address_space_to_nvptx(address_space: gpu.AddressSpace) -> int:
+  match address_space:
+    case gpu.AddressSpace.Global:
+      return 1
+    case gpu.AddressSpace.Workgroup:
+      return 3
+    case _:
+      raise NotImplementedError(f"address_space not supported: {address_space}")
+
+
+@_register_lowering(InitializeBarrierOp)
+def _initialize_barrier_op_lowering_rule(
+    initialize_barrier_op: InitializeBarrierOp) -> Sequence[ir.Value]:
+
+  shape = initialize_barrier_op.barriers_ref.type.shape
+  num_barriers = functools.reduce(operator.mul, shape, 1)
+
+  i32 = ir.IntegerType.get_signless(32)
+  workgroup_nvptx_address_space = _gpu_address_space_to_nvptx(
+      gpu.AddressSpace.Workgroup)
+  ptr_ty = ir.Type.parse(f"!llvm.ptr<{workgroup_nvptx_address_space}>")
+
+  lowered_barrier_type = _lowered_barrier_type()
+  lowered_barrier_ref = memref.alloca(
+      ir.MemRefType.get(shape, lowered_barrier_type), [], [])
+  barrier_ref_address = memref_ptr(
+      lowered_barrier_ref, memory_space=workgroup_nvptx_address_space)
+
+  predicate = single_thread_predicate(per_block=True)
+  for i in range(num_barriers):
+    nvvm.mbarrier_init_shared(
+        llvm.getelementptr(ptr_ty, barrier_ref_address, [], [i],
+                           lowered_barrier_type),
+        c(initialize_barrier_op.arrival_count.value, i32),
+        predicate=predicate
+    )
+  return barrier_ref_address,
+
+
+def lower_mgpu_dialect(module: ir.Module):
+  module.context.append_dialect_registry(mlir_interpreter.upstream_dialects)
+  module.context.load_all_available_dialects()
+
+  lowered_operations: set[ir.Operation | ir.OpView] = set()
+
+  def _lower_op(op: ir.OpView):
+    if op.name not in _lowerings:
+      return
+    lowering_rule = _lowerings[op.name]
+    new_results = lowering_rule(op)
+    for old, new in zip(op.results, new_results):
+      old.replace_all_uses_with(new)
+    lowered_operations.add(op)
+
+  def _traverse_and_lower_op(op: ir.OpView):
+    for region in op.operation.regions:
+      for block in region:
+        for block_op in list(block):
+          with ir.InsertionPoint(block_op):
+            _traverse_and_lower_op(block_op)
+    _lower_op(op)
+
+  with ir.InsertionPoint(module.body):
+    for op in module.body:
+      _traverse_and_lower_op(op)
+
+  for lowered_op in lowered_operations:
+    lowered_op.erase()

tests/mosaic/gpu_dialect_test.py

@@ -14,12 +14,19 @@
 # ==============================================================================
 """(Deviceless) tests for the Mosaic GPU MLIR dialect."""
 
+from typing import Callable
+
 from absl.testing import parameterized
 from jax._src import config
 from jax._src import test_util as jtu
+from jax._src.interpreters import mlir as mlir_interpreter
 from jax._src.lib.mlir import ir
-from jax.experimental.mosaic.gpu import dialect as mgpu  # pylint: disable=g-importing-member
+from jax._src.lib.mlir.dialects import arith
+from jax._src.lib.mlir.dialects import nvvm
+from jax._src.lib.mlir.dialects import scf
 
+from jax.experimental.mosaic.gpu import dialect as mgpu  # pylint: disable=g-importing-member
+from jax.experimental.mosaic.gpu import lower_mgpu_dialect  # pylint: disable=g-importing-member,g-multiple-import
 
 _cext = mgpu._cext if mgpu is not None else None
 
@@ -29,10 +36,35 @@
 
 def _make_ir_context():
   context = ir.Context()
+  context.append_dialect_registry(mlir_interpreter.upstream_dialects)
+  context.load_all_available_dialects()
   mgpu.register_dialect(context)
   return context
 
 
+def walk_operations(op: ir.OpView, callback):
+  for region in op.operation.regions:
+    for block in region:
+      for block_op in block:
+        walk_operations(block_op, callback)
+  callback(op)
+
+
+def find_if(module: ir.Module,
+            predicate: Callable[[ir.OpView], bool]) -> list[ir.OpView]:
+  result = []
+  def callback(op: ir.OpView):
+    if predicate(op):
+      result.append(op)
+  for op in module.body.operations:
+    walk_operations(op, callback)
+  return result
+
+
+def is_mosaic_gpu_op(op: ir.OpView) -> bool:
+  return op.name.startswith("mosaic_gpu.")
+
+
 class DialectTest(parameterized.TestCase):
 
   def setUp(self):
@@ -72,5 +104,57 @@ def test_initialize_barrier_op_with_a_positive_arrival_count_passes(self):
                           mgpu.InitializeBarrierOp)
 
 
+class DialectLoweringTest(DialectTest):
+
+  def test_lowering_removes_mosaic_gpu_ops(self):
+    with ir.InsertionPoint(self.module.body):
+      mgpu.initialize_barrier(
+          ir.MemRefType.get((1, 2), ir.Type.parse("!mosaic_gpu.barrier")),
+          arrival_count=1)
+    lower_mgpu_dialect(self.module)
+
+    self.assertEmpty(
+        list(filter(is_mosaic_gpu_op, self.module.body.operations)))
+
+  def test_lowering_traverses_regions_correctly(self):
+    with ir.InsertionPoint(self.module.body):
+      bool_type = ir.IntegerType.get_signless(1)
+      cst_true = arith.constant(bool_type, ir.IntegerAttr.get(bool_type, 1))
+      if_op = scf.IfOp(cst_true)
+      with ir.InsertionPoint(if_op.then_block):
+        mgpu.initialize_barrier(
+            ir.MemRefType.get((1, 2), ir.Type.parse("!mosaic_gpu.barrier")),
+            arrival_count=1)
+        scf.yield_([])
+    lower_mgpu_dialect(self.module)
+
+    self.assertEmpty(
+        list(filter(is_mosaic_gpu_op, if_op.then_block.operations)))
+
+  def test_initialize_barrier_op_lowering_rule(self):
+    shape = (3, 4)
+    num_shape_elements = shape[0] * shape[1]
+    arrival_count = 1337
+
+    with ir.InsertionPoint(self.module.body):
+      mgpu.initialize_barrier(
+          ir.MemRefType.get(shape, ir.Type.parse("!mosaic_gpu.barrier")),
+          arrival_count=arrival_count)
+    lower_mgpu_dialect(self.module)
+
+    all_mbarrier_init_shared_ops = find_if(
+        self.module,
+        lambda op: op.name == nvvm.MBarrierInitSharedOp.OPERATION_NAME)
+
+    # One nvvm.mbarrier_init_shared is issued per barrier.
+    self.assertLen(all_mbarrier_init_shared_ops, num_shape_elements)
+
+    # Each barrier has its count equal to the arrival count.
+    for op in all_mbarrier_init_shared_ops:
+      count = op.count.owner.opview
+      self.assertIsInstance(count, arith.ConstantOp)
+      self.assertEqual(count.literal_value, arrival_count)
+
+
 if __name__ == "__main__":
   parameterized.absltest.main(testLoader=jtu.JaxTestLoader())