[Pallas/Mosaic GPU] Allow scalar loads in warp-level lowering.

PiperOrigin-RevId: 835174212

Author: bchetioui

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -1536,6 +1536,9 @@ def _ndindexer_indices(
 
 
 @register_lowering_rule(sp.get_p, mgpu.LoweringSemantics.Lane)
+@register_lowering_rule(
+    sp.get_p, mgpu.LoweringSemantics.Lane, gpu_core.PrimitiveSemantics.Warp
+)
 def _get_lowering_rule(
     ctx: LoweringRuleContext, x_ref, *leaves, tree, optimized=True
 ):
@@ -1544,6 +1547,11 @@ def _get_lowering_rule(
         "Loads from TMEM are asynchronous operations and cannot be performed"
         " using the usual syntax. Please use plgpu.async_load_tmem instead."
     )
+  if (
+      ctx.avals_out[0].shape
+      and ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp
+  ):
+    raise ValueError("Can only load scalars in warp-level code.")
   if not isinstance(x_ref, ir.Value) and ir.MemRefType.isinstance(x_ref):
     raise TypeError(f"Can only load from references (got {x_ref}).")
   dtype = ctx.avals_out[0].dtype

jax/_src/pallas/mosaic_gpu/primitives.py

@@ -2938,6 +2938,11 @@ def _load_abstract_eval(src, *avals_flat, tree, optimized):
 lowering.register_lowering_rule(load_p, mgpu.LoweringSemantics.Lane)(
     lowering._get_lowering_rule
 )
+lowering.register_lowering_rule(
+    load_p, mgpu.LoweringSemantics.Lane, gpu_core.PrimitiveSemantics.Warp
+)(
+    lowering._get_lowering_rule
+)
 lowering.register_lowering_rule(load_p, mgpu.LoweringSemantics.Warpgroup)(
     lowering._get_lowering_rule_wg
 )

tests/pallas/mosaic_gpu_test.py

@@ -2592,6 +2592,33 @@ def _():
                           jnp.ones((128,), jnp.int32) * 3), axis=0)
     np.testing.assert_array_equal(result, expected)
 
+  def test_scalar_load(self):
+    warp_mesh = plgpu.WarpMesh(axis_name="warp")
+    @functools.partial(self.kernel,
+                       out_shape=jax.ShapeDtypeStruct((), jnp.int32))
+    def kernel(x_ref, y_ref):
+      @pl.core_map(warp_mesh)
+      def _():
+        warp_id = lax.axis_index("warp")
+        @pl.when(warp_id == 1)
+        def _():
+          y_ref[...] = x_ref[...]
+    np.testing.assert_array_equal(kernel(4), 4)
+
+  def test_non_scalar_load_raises(self):
+    warp_mesh = plgpu.WarpMesh(axis_name="warp")
+    @functools.partial(self.kernel,
+                       out_shape=jax.ShapeDtypeStruct((2,), jnp.int32))
+    def kernel(x_ref, y_ref):
+      @pl.core_map(warp_mesh)
+      def _():
+        warp_id = lax.axis_index("warp")
+        @pl.when(warp_id == 1)
+        def _():
+          y_ref[...] = x_ref[...]
+    with self.assertRaisesRegex(ValueError, "Can only load scalars",):
+      kernel(jnp.ones((2,), jnp.int32))
+
   @parameterized.parameters(
     lax.add, lax.sub, lax.mul, lax.div, lax.rem, lax.bitwise_and,
     lax.bitwise_or, lax.bitwise_xor, lax.max, lax.min,