Repository: jax-ml/jax
Pull request #25907 (2025-01-15 16:47:50)

[MosaicGPU] Move `gpu_address_space_to_nvptx` inside `utils.py` and use it.
[MosaicGPU] Move `gpu_address_space_to_nvptx` inside `utils.py` and use it.



Patch:
diff --git a/jax/experimental/mosaic/gpu/__init__.py b/jax/experimental/mosaic/gpu/__init__.py
index 086ea025b..be42e59e5 100644
--- a/jax/experimental/mosaic/gpu/__init__.py
+++ b/jax/experimental/mosaic/gpu/__init__.py
@@ -34,14 +34,12 @@ from .launch_context import (
 
 if dialect is not None:
   from .dialect_lowering import (
-      gpu_address_space_to_nvptx as gpu_address_space_to_nvptx,
       lower_mgpu_dialect as lower_mgpu_dialect,
   )
   from .layout_inference import (
       infer_layout as infer_layout,
   )
 else:
-  gpu_address_space_to_nvptx = None
   lower_mgpu_dialect = None
   infer_layout = None
 
diff --git a/jax/experimental/mosaic/gpu/core.py b/jax/experimental/mosaic/gpu/core.py
index d4ea4998c..ab70c19b9 100644
--- a/jax/experimental/mosaic/gpu/core.py
+++ b/jax/experimental/mosaic/gpu/core.py
@@ -183,9 +183,7 @@ def _construct_smem_reftree(
     def get_barrier_ptr(num_barriers: int) -> ir.Value:
       nonlocal dynamic_smem_offset
       workgroup_nvptx_address_space = (
-          dialect_lowering.gpu_address_space_to_nvptx(
-              gpu.AddressSpace.Workgroup
-          )
+          utils.gpu_address_space_to_nvptx(gpu.AddressSpace.Workgroup)
       )
       smem_base_ptr = utils.memref_ptr(
           dynamic_smem, memory_space=workgroup_nvptx_address_space
diff --git a/jax/experimental/mosaic/gpu/dialect_lowering.py b/jax/experimental/mosaic/gpu/dialect_lowering.py
index 5f7b321e1..642c768b2 100644
--- a/jax/experimental/mosaic/gpu/dialect_lowering.py
+++ b/jax/experimental/mosaic/gpu/dialect_lowering.py
@@ -33,7 +33,7 @@ import numpy as np
 from .fragmented_array import FragmentedArray, WGStridedFragLayout
 from .launch_context import LaunchContext
 from .layouts import from_strided_fragmented_layout_attr, has_any_layout_set, is_strided_fragmented_layout, should_have_layout, to_strided_fragmented_layout_attr
-from .utils import BarrierRef, c, ptr_as_memref, single_thread_predicate
+from .utils import BarrierRef, c, gpu_address_space_to_nvptx, ptr_as_memref, single_thread_predicate
 
 # mypy: ignore-errors
 
@@ -130,16 +130,6 @@ 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(
     _: LaunchContext,
diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py
index 038bf2d44..ed0c17eb4 100644
--- a/jax/experimental/mosaic/gpu/utils.py
+++ b/jax/experimental/mosaic/gpu/utils.py
@@ -48,6 +48,21 @@ MBARRIER_BYTES = 8
 # pylint: disable=line-too-long, wildcard-import, missing-function-docstring, bad-continuation, g-bad-todo, protected-access, g-explicit-length-test, missing-class-docstring, g-doc-return-or-yield, g-inconsistent-quotes
 
 
+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}")
+
+
+WORKGROUP_NVPTX_ADDRESS_SPACE = gpu_address_space_to_nvptx(
+    gpu.AddressSpace.Workgroup
+)
+
+
 def ptr_as_memref(ptr, memref_ty: ir.MemRefType, ptr_memory_space: int | None = None):
   i64 = ir.IntegerType.get_signless(64)
   rank = len(memref_ty.shape)
@@ -674,7 +689,7 @@ class BarrierRef:
       raise NotImplementedError("Only up to 32 barriers per group supported")
     i32 = ir.IntegerType.get_signless(32)
     i64 = ir.IntegerType.get_signless(64)
-    ptr = ir.Type.parse("!llvm.ptr<3>")
+    ptr = ir.Type.parse(f"!llvm.ptr<{WORKGROUP_NVPTX_ADDRESS_SPACE}>")
     phases = memref.alloca(ir.MemRefType.get((), i32), [], [])
     memref.store(c(0, i32), phases, [])
     with single_thread(per_block=True):
@@ -742,7 +757,7 @@ class BarrierRef:
     nvvm.mbarrier_arrive_expect_tx_shared(self.get_ptr(), bytes, predicate=predicate)
 
   def get_ptr(self):
-    ptr = ir.Type.parse("!llvm.ptr<3>")
+    ptr = ir.Type.parse(f"!llvm.ptr<{WORKGROUP_NVPTX_ADDRESS_SPACE}>")
     i64 = ir.IntegerType.get_signless(64)
     DYNAMIC32 = -2147483648
     return llvm.getelementptr(
@@ -754,7 +769,7 @@ class BarrierRef:
     return ptr_as_memref(
         self.base_address,
         ir.MemRefType.get(shape, ir.Type.parse("!mosaic_gpu.barrier")),
-        ptr_memory_space=3,
+        ptr_memory_space=WORKGROUP_NVPTX_ADDRESS_SPACE,
     )
 
   @classmethod
@@ -770,7 +785,9 @@ class BarrierRef:
       )
 
     return cls(
-        base_address=memref_ptr(barrier, memory_space=3),
+        base_address=memref_ptr(
+            barrier, memory_space=WORKGROUP_NVPTX_ADDRESS_SPACE
+        ),
         offset=c(0, ir.IntegerType.get_signless(64)),
         phases=None,
         num_barriers=(1 if memref_type.rank == 0 else memref_type.shape[0]),


Test patch:
diff --git a/tests/mosaic/gpu_dialect_test.py b/tests/mosaic/gpu_dialect_test.py
index 036e0e4d2..d3b10bf8a 100644
--- a/tests/mosaic/gpu_dialect_test.py
+++ b/tests/mosaic/gpu_dialect_test.py
@@ -30,6 +30,7 @@ from jax._src.lib.mlir.dialects import nvvm
 from jax._src.lib.mlir.dialects import scf
 from jax._src.lib.mlir.dialects import vector
 import jax.experimental.mosaic.gpu as mgpu
+import jax.experimental.mosaic.gpu.utils as mgpu_utils
 
 _cext = mgpu.dialect._cext if mgpu.dialect is not None else None
 
@@ -72,7 +73,7 @@ def is_mosaic_gpu_op(op: ir.OpView) -> bool:
 
 
 def workgroup_ptr_ty() -> ir.Type:
-  workgroup_nvptx_address_space = mgpu.gpu_address_space_to_nvptx(
+  workgroup_nvptx_address_space = mgpu_utils.gpu_address_space_to_nvptx(
       gpu.AddressSpace.Workgroup
   )
   return ir.Type.parse(f"!llvm.ptr<{workgroup_nvptx_address_space}>")