[Pallas TPU] Allow closed over scalars in core_map code

This allows doing things like dynamic indexing of Refs using just
regular scalars from outside the kernel.

PiperOrigin-RevId: 842429684

Author: sharadmv

jax/_src/pallas/core.py

@@ -1535,8 +1535,13 @@ def default_mesh_discharge_rule(
     scratch_shapes,
 ):
   """Discharges a ``core_map`` over a mesh to a ``pallas_call``."""
-  del out_avals  # Unused.
   default_memory_space = memory_space
+  if not all(
+      isinstance(aval, state.AbstractRef) for aval in (in_avals + out_avals)
+  ):
+    raise ValueError(
+        "default_mesh_discharge_rule only supports Ref inputs/outputs."
+    )
 
   def body(*args):
     # Due to aliasing, ``args`` contains aliased inputs and outputs so we
@@ -1605,15 +1610,24 @@ def _core_map_discharge_rule(in_avals, out_avals, *args_flat, jaxpr, debug_info,
         for var in jaxpr.constvars
         if not isinstance(aval := var.aval, state.AbstractRef)
     ]
-    if consts_avals:
+    is_scalar_const_aval = [
+        isinstance(aval, jax_core.ShapedArray) and not aval.shape
+        for aval in consts_avals
+    ]
+    if not all(is_scalar_const_aval):
       ctx = jax_core.JaxprPpContext()
-      pp_const_avals = ", ".join(
-          jax_core.pp_aval(aval, ctx) for aval in consts_avals
+      non_scalar_const_avals = [
+          aval
+          for aval, is_scalar in zip(consts_avals, is_scalar_const_aval)
+          if not is_scalar
+      ]
+      non_scalar_const_pp_avals = ", ".join(
+          jax_core.pp_aval(aval, ctx) for aval in non_scalar_const_avals
       )
       raise ValueError(
           "The kernel function in core_map"
-          f" {debug_info.func_src_info} captures constants"
-          f" [{pp_const_avals}]. You should pass them as inputs."
+          f" {debug_info.func_src_info} captures non-scalar constants"
+          f" [{non_scalar_const_pp_avals}]. You should pass them as inputs."
       )
   return _core_map_mesh_rules[type(mesh)](
       in_avals, out_avals, *args_flat, jaxpr=jaxpr, mesh=mesh, **kwargs

jax/_src/pallas/mosaic/core.py

@@ -26,9 +26,11 @@
 import jax.numpy as jnp
 from jax.extend import backend as jex_backend
 from jax._src import core as jax_core
+from jax._src import linear_util as lu
 from jax._src import state
 from jax._src import util
 from jax._src.frozen_dict import FrozenDict
+from jax._src.interpreters import partial_eval as pe
 from jax._src.pallas import core as pallas_core
 import numpy as np
 
@@ -336,6 +338,49 @@ def _tensorcore_mesh_discharge_rule(
           "TensorCoreMesh does not support VMEM inputs/outputs when there are"
           " >1 cores. Use HBM or ANY instead."
       )
+  def allowed_aval(aval):
+    if isinstance(aval, state.AbstractRef):
+      return True
+    if isinstance(aval, jax_core.ShapedArray):
+      # Only scalars are allowed.
+      return not aval.shape
+    return False
+  assert all(allowed_aval(v.aval) for v in jaxpr.constvars + jaxpr.invars)
+
+  is_scalar_const = [
+      isinstance(v.aval, jax_core.ShapedArray) and not v.aval.shape
+      for v in jaxpr.constvars
+  ]
+  if any(is_scalar_const):
+    # Rewrite body jaxpr to take in scalar values as Refs.
+    def new_body(*args):
+      args = [
+          a[0] if is_scalar else a
+          for a, is_scalar in zip(args, is_scalar_const)
+      ]
+      return jax_core.eval_jaxpr(jaxpr, args)
+    # TODO(sharadmv): Remove this once Mosaic support passing scalars as values.
+    new_trace_avals = [
+        state.AbstractRef(  # pylint: disable=g-long-ternary
+            jax_core.ShapedArray((1,), v.aval.dtype),
+            memory_space=MemorySpace.SMEM,
+        )
+        if is_scalar
+        else v.aval
+        for v, is_scalar in zip(jaxpr.constvars, is_scalar_const)
+    ]
+    new_jaxpr, _, _ = pe.trace_to_jaxpr_dynamic(
+        lu.wrap_init(
+            new_body, debug_info=jaxpr.debug_info.with_unknown_names()
+        ),
+        new_trace_avals,
+    )
+    jaxpr = new_jaxpr.replace(invars=[], constvars=new_jaxpr.invars)
+    args = tuple(
+        a[None] if is_scalar else a
+        for a, is_scalar in zip(args, is_scalar_const)
+    )
+    in_avals, out_avals = util.split_list(new_trace_avals, [len(in_avals)])
   return pallas_core.default_mesh_discharge_rule(
       in_avals,
       out_avals,

jax/_src/pallas/mosaic/sc_core.py

@@ -219,6 +219,11 @@ def _scalar_subcore_mesh_discharge_rule(
     compiler_params = tpu_core.CompilerParams()
   if compiler_params.dimension_semantics is not None:
     raise ValueError("ScalarSubcoreMesh does not support dimension_semantics=")
+  sa_avals = [a for a in in_avals if isinstance(a, jax_core.ShapedArray)]
+  if sa_avals:
+    raise NotImplementedError(
+        f"Cannot close over values in core_map: {sa_avals}"
+    )
   return pallas_core.default_mesh_discharge_rule(
       in_avals,
       out_avals,

jax/_src/pallas/mosaic_gpu/core.py

@@ -1368,6 +1368,11 @@ def _gpu_mesh_discharge_rule(
     )
   if not compiler_params:
     compiler_params = CompilerParams()
+  sa_avals = [a for a in in_avals if isinstance(a, jax_core.ShapedArray)]
+  if sa_avals:
+    raise NotImplementedError(
+        f"Cannot close over values in core_map: {sa_avals}"
+    )
   return pallas_core.default_mesh_discharge_rule(
       in_avals,
       out_avals,

tests/pallas/tpu_pallas_state_test.py

@@ -324,9 +324,38 @@ def kernel(x_ref, out_ref, tmp_ref):
       return kernel(x)
 
     x = jnp.arange(8 * 128, dtype=jnp.int32).reshape((8, 128))
-    with self.assertRaisesRegex(Exception, "core_map .* captures constants"):
+    with self.assertRaisesRegex(
+        Exception, "core_map .* captures non-scalar constants"
+    ):
       f(x)
 
+  def test_capture_scalar(self):
+    @jax.jit
+    def f(x, i):
+      @pl.kernel(out_shape=jax.ShapeDtypeStruct(x.shape[1:], jnp.int32),
+                 mesh=pltpu.create_tensorcore_mesh("x", num_cores=1))
+      def kernel(x_ref, out_ref):
+        pltpu.sync_copy(x_ref.at[i], out_ref)
+      return kernel(x)
+
+    x = jnp.arange(4 * 8 * 128, dtype=jnp.int32).reshape((4, 8, 128))
+    for i in range(x.shape[0]):
+      out = f(x, i)
+      np.testing.assert_array_equal(out, x[i])
+
+    @jax.jit
+    def g(x, i):
+      @pl.kernel(out_shape=jax.ShapeDtypeStruct((2, *x.shape[1:]), jnp.int32),
+                 mesh=pltpu.create_tensorcore_mesh("x", num_cores=1))
+      def kernel(x_ref, out_ref):
+        pltpu.sync_copy(x_ref.at[pl.ds(i, 2)], out_ref)
+      return kernel(x)
+
+    x = jnp.arange(4 * 8 * 128, dtype=jnp.int32).reshape((4, 8, 128))
+    for i in range(3):
+      out = g(x, i)
+      np.testing.assert_array_equal(out, x[i:i+2])
+
   def test_kernel_helper_with_scratch(self):
     mesh = pltpu.create_tensorcore_mesh("x")
     def body(x_ref, o_ref, scratch_ref):