[mgpu_pallas] Better support for unsigned integers and floats in iota.

PiperOrigin-RevId: 701307324

Author: cperivol

jax/_src/pallas/mosaic_gpu/primitives.py

@@ -706,17 +706,33 @@ def _broadcasted_iota_abstract_eval(dtype, shape, dimension, layout):
 @lowering.register_lowering_rule(broadcasted_iota_p)
 def _broadcasted_iota_lowering(ctx: lowering.LoweringRuleContext, dtype, shape, dimension, layout):
   del ctx
-  undef = llvm_dialect.mlir_undef(mlir.dtype_to_ir_type(dtype))
+  # Unsigned integers (as opposed to signless) cause MLIR verification
+  # errors so we only use signless like Mosaic GPU does.
+  #
+  # TODO(cperivol): use mgpu.utils.dtype_to_ir_type() instead.
+  mlir_dtype = (
+      ir.IntegerType.get_signless(dtype.itemsize * 8)
+      if jnp.issubdtype(dtype, jnp.integer)
+      else mlir.dtype_to_ir_type(dtype)
+  )
+  undef = llvm_dialect.mlir_undef(mlir_dtype)
   is_signed = (
       jnp.issubdtype(dtype, jnp.signedinteger)
       if jnp.issubdtype(dtype, jnp.integer)
       else None
   )
-  mlir_dtype = mlir.dtype_to_ir_type(dtype)
+
+  i32 = ir.IntegerType.get_signless(32)
+  def _cast(x):
+    if ir.FloatType.isinstance(mlir_dtype):
+      x = arith_dialect.index_cast(i32, x)
+      return arith_dialect.uitofp(mlir_dtype, x)
+    else:
+      return arith_dialect.index_cast(mlir_dtype, x)
   return mgpu.FragmentedArray.splat(
       undef, shape, layout.value, is_signed=is_signed
   ).foreach(
-      lambda _, idx: arith_dialect.index_cast(mlir_dtype, idx[dimension]), create_array=True, is_signed=is_signed
+      lambda _, idx: _cast(idx[dimension]), create_array=True, is_signed=is_signed
   )
 
 

tests/pallas/mosaic_gpu_test.py

@@ -241,8 +241,9 @@ def kernel(x_ref, o_ref):
     # are never written to.
     np.testing.assert_array_equal(kernel(x)[:, :16], y[:, :16])
 
-  def test_iota(self):
-    dtype, dimension = jnp.int8, 1
+  @parameterized.parameters(jnp.float32, jnp.int32, jnp.uint32)
+  def test_iota(self, dtype):
+    dimension = 1
     @functools.partial(
         pl.pallas_call,
         out_shape=jax.ShapeDtypeStruct((128, 128), dtype),