[Pallas TPU] Add pack_elementwise and unpack_elementwise primitives.

PiperOrigin-RevId: 828697560

Author: megrez-yliu

jax/_src/pallas/mosaic/lowering.py

@@ -3580,6 +3580,45 @@ def _stochastic_round_lowering_rule(
   return tpu.stochastic_convert(out_type, x, random_bits)
 
 
+def _check_elementwise_packing_dtypes(unpacked_dtype, packed_dtype):
+  if unpacked_dtype == jnp.float32 and packed_dtype == jnp.bfloat16:
+    return
+  if unpacked_dtype == jnp.int32 and packed_dtype in [
+      jnp.int16, jnp.int8, jnp.int4
+    ]:
+    return
+  raise ValueError(
+      f"Unsupported elementwise packing: {unpacked_dtype} -> {packed_dtype}. "
+      "Only f32 <-> bf16 and i32 <-> i16/i8/i4 are supported."
+  )
+
+
+@register_lowering_rule(tpu_primitives.pack_elementwise_p)
+def _pack_elementwise_lowering_rule(
+    ctx: LoweringRuleContext, *xs, packed_dtype
+):
+  in_aval = ctx.avals_in[0]
+  _check_elementwise_packing_dtypes(in_aval.dtype, packed_dtype)
+  packed_ir_type = _dtype_to_ir_type(packed_dtype)
+  out_type = ir.VectorType.get(
+      in_aval.shape, _dtype_to_ir_type(jnp.uint32)
+  )
+  return tpu.pack_elementwise(out_type, xs, target_type=packed_ir_type)
+
+
+@register_lowering_rule(tpu_primitives.unpack_elementwise_p)
+def _unpack_elementwise_lowering_rule(
+    ctx: LoweringRuleContext, x, index, packed_dtype, unpacked_dtype
+):
+  in_aval = ctx.avals_in[0]
+  _check_elementwise_packing_dtypes(unpacked_dtype, packed_dtype)
+  out_type = ir.VectorType.get(
+      in_aval.shape, _dtype_to_ir_type(unpacked_dtype)
+  )
+  return tpu.unpack_elementwise(
+      out_type, x, source_type=_dtype_to_ir_type(packed_dtype), index=index)
+
+
 @register_lowering_rule(tpu_primitives.bitcast_p)
 def _bitcast_lowering_rule(ctx: LoweringRuleContext, x, *, ty):
   del ty

jax/_src/pallas/mosaic/primitives.py

@@ -969,6 +969,60 @@ def _stochastic_round_abstract_eval(x, random_bits, *, target_dtype):
     )
   return jax_core.ShapedArray(x.shape, target_dtype)
 
+def _get_elementwise_packing_factor(unpacked_dtype, packed_dtype):
+  unpacked_bitwidth = dtypes.bit_width(unpacked_dtype)
+  packed_bitwidth = dtypes.bit_width(packed_dtype)
+  if unpacked_bitwidth % packed_bitwidth != 0:
+    raise ValueError(
+        "Unpacked bitwidth must be a multiple of packed bitwidth, got "
+        f"{unpacked_bitwidth} and {packed_bitwidth}"
+    )
+  return unpacked_bitwidth // packed_bitwidth
+
+pack_elementwise_p = jax_core.Primitive("pack_elementwise")
+
+
+def pack_elementwise(xs, *, packed_dtype):
+  return pack_elementwise_p.bind(*xs, packed_dtype=packed_dtype)
+
+
+@pack_elementwise_p.def_abstract_eval
+def _pack_elementwise_abstract_eval(*xs, packed_dtype):
+  if not xs:
+    raise ValueError("At least one source is required")
+  first = xs[0]
+  if not all(x.shape == first.shape for x in xs):
+    raise ValueError("All sources must have the same shape")
+  if not all(x.dtype == first.dtype for x in xs):
+    raise ValueError("All sources must have the same dtype")
+  packing_factor = _get_elementwise_packing_factor(first.dtype, packed_dtype)
+  if len(xs) != packing_factor:
+    raise ValueError(
+        "The number of sources must match the packing factor "
+        f"({packing_factor}), got {len(xs)}"
+    )
+  return jax_core.ShapedArray(first.shape, jnp.uint32)
+
+
+unpack_elementwise_p = jax_core.Primitive("unpack_elementwise")
+
+
+def unpack_elementwise(x, *, index, packed_dtype, unpacked_dtype):
+  return unpack_elementwise_p.bind(
+      x, index=index, packed_dtype=packed_dtype, unpacked_dtype=unpacked_dtype
+  )
+
+
+@unpack_elementwise_p.def_abstract_eval
+def _unpack_elementwise_abstract_eval(x, *, index, packed_dtype, unpacked_dtype):
+  if x.dtype != jnp.uint32:
+    raise ValueError(f"Source must be uint32, got {x.dtype}")
+  packing_factor = _get_elementwise_packing_factor(unpacked_dtype, packed_dtype)
+  if index < 0 or index >= packing_factor:
+    raise ValueError(
+        f"Index {index} is out of bounds for packing factor {packing_factor}")
+  return jax_core.ShapedArray(x.shape, unpacked_dtype)
+
 
 def with_memory_space_constraint(
     x: jax.Array, memory_space: Any

jax/experimental/pallas/tpu.py

@@ -45,13 +45,15 @@
 from jax._src.pallas.mosaic.primitives import load as load
 from jax._src.pallas.mosaic.primitives import make_async_copy as make_async_copy
 from jax._src.pallas.mosaic.primitives import make_async_remote_copy as make_async_remote_copy
+from jax._src.pallas.mosaic.primitives import pack_elementwise as pack_elementwise
 from jax._src.pallas.mosaic.primitives import prng_random_bits as prng_random_bits
 from jax._src.pallas.mosaic.primitives import prng_seed as prng_seed
 from jax._src.pallas.mosaic.primitives import repeat as repeat
 from jax._src.pallas.mosaic.primitives import roll as roll
 from jax._src.pallas.mosaic.primitives import stochastic_round as stochastic_round
 from jax._src.pallas.mosaic.primitives import store as store
 from jax._src.pallas.mosaic.primitives import touch as touch
+from jax._src.pallas.mosaic.primitives import unpack_elementwise as unpack_elementwise
 from jax._src.pallas.mosaic.primitives import with_memory_space_constraint as with_memory_space_constraint
 from jax._src.pallas.mosaic.random import sample_block as sample_block
 from jax._src.pallas.mosaic.random import stateful_bernoulli as stateful_bernoulli

tests/pallas/tpu_ops_test.py

@@ -716,7 +716,6 @@ def kernel(x_ref, b_ref, o_ref):
 
     result = pl.pallas_call(
         kernel,
-        in_specs=[pl.BlockSpec(), pl.BlockSpec()],
         out_shape=jax.ShapeDtypeStruct(x.shape, target_dtype),
     )(x, bits)
 
@@ -730,5 +729,108 @@ def kernel(x_ref, b_ref, o_ref):
     )
     self.assertTrue(jnp.all(is_correct))
 
+  def _pack_unpack_elementwise_test_data(
+      self, shape, unpacked_dtype, packed_dtype):
+    """Generates data for test_pack_elementwise and test_unpack_elementwise."""
+    bitwidth = dtypes.bit_width(packed_dtype)
+    num_sources = 32 // bitwidth
+    if unpacked_dtype == jnp.int32:
+      stacked_sources = jax.random.randint(
+          jax.random.key(0),
+          (num_sources, *shape),
+          minval=-1000,
+          maxval=1000,
+          dtype=unpacked_dtype,
+      )
+    else:
+      stacked_sources = jax.random.uniform(
+          jax.random.key(0), (num_sources, *shape), dtype=unpacked_dtype
+      )
+    stacked_results = (
+        stacked_sources.astype(packed_dtype)
+        .view(getattr(jnp, f"uint{bitwidth}"))
+        .astype(jnp.uint32)
+    )
+    shifts = jnp.arange(num_sources, dtype=jnp.uint32) * bitwidth
+    shifts = jnp.expand_dims(shifts, axis=tuple(range(1, stacked_results.ndim)))
+    packed_data = jnp.bitwise_or.reduce(stacked_results << shifts, axis=0)
+    return stacked_sources, packed_data
+
+  @parameterized.product(
+      config=[
+          (jnp.float32, jnp.bfloat16),
+          (jnp.int32, jnp.int16),
+          (jnp.int32, jnp.int8),
+          (jnp.int32, jnp.int4),
+      ],
+      shape=[(8, 128), (2, 15, 300)],
+  )
+  def test_pack_elementwise(self, config, shape):
+    unpacked_dtype, packed_dtype = config
+    if not jtu.is_device_tpu_at_least(version=5):
+      self.skipTest("Requires TPU v5+")
+    if not jtu.if_cloud_tpu_at_least(2025, 11, 7):
+      self.skipTest("Test requires libtpu from 2025/11/7 or later")
+
+    bitwidth = dtypes.bit_width(packed_dtype)
+    num_sources = 32 // bitwidth
+
+    def kernel(xs_ref, o_ref):
+      xs = [xs_ref[i] for i in range(num_sources)]
+      o_ref[...] = pltpu.pack_elementwise(xs, packed_dtype=packed_dtype)
+
+    stacked_sources, expected = self._pack_unpack_elementwise_test_data(
+        shape, unpacked_dtype, packed_dtype
+    )
+
+    result = self.pallas_call(
+        kernel,
+        out_shape=jax.ShapeDtypeStruct(shape, jnp.uint32),
+    )(stacked_sources)
+
+    np.testing.assert_array_equal(result, expected)
+
+  @parameterized.product(
+      config=[
+          (jnp.float32, jnp.bfloat16),
+          (jnp.int32, jnp.int16),
+          (jnp.int32, jnp.int8),
+          (jnp.int32, jnp.int4),
+      ],
+      index=[0, 1, 3],
+      shape=[(8, 128), (2, 15, 300)],
+  )
+  def test_unpack_elementwise(self, config, index, shape):
+    unpacked_dtype, packed_dtype = config
+    if not jtu.is_device_tpu_at_least(version=5):
+      self.skipTest("Requires TPU v5+")
+    if not jtu.if_cloud_tpu_at_least(2025, 11, 7):
+      self.skipTest("Test requires libtpu from 2025/11/7 or later")
+
+    bitwidth = dtypes.bit_width(packed_dtype)
+    packing_factor = 32 // bitwidth
+
+    if index >= packing_factor:
+      self.skipTest(
+          f"Index {index} out of bounds for packing factor {packing_factor}")
+
+    def kernel(x_ref, o_ref):
+      o_ref[...] = pltpu.unpack_elementwise(
+          x_ref[...], index=index,
+          packed_dtype=packed_dtype, unpacked_dtype=unpacked_dtype
+      )
+
+    sources, packed = self._pack_unpack_elementwise_test_data(
+        shape, unpacked_dtype, packed_dtype
+    )
+    expected = sources[index].astype(packed_dtype).astype(unpacked_dtype)
+
+    result = self.pallas_call(
+        kernel,
+        out_shape=jax.ShapeDtypeStruct(shape, unpacked_dtype),
+    )(packed)
+
+    np.testing.assert_array_equal(result, expected)
+
 if __name__ == "__main__":
   absltest.main()