[NFC] Add name argument to pallas_call primitive.

This change introduces an optional `name` string argument to the `pallas_call` primitive and its associated lowering and transformation rules. The `name` is currently passed through but not yet used by the backend-specific lowering implementations for TPU, GPU, and Triton.

A follow up CL will use it to extend the name stack on TPU

PiperOrigin-RevId: 825262144

Author: sharadmv

jax/_src/pallas/core.py

@@ -1391,6 +1391,7 @@ def core_map(
     interpret: Whether to run the function in interpret mode.
     debug: Whether or not to out helpful debugging information.
     cost_estimate: The cost estimate of the function.
+    name: The (optional) name of the kernel.
     metadata: Optional dictionary of information about the kernel that will be
       serialized as JSON in the HLO. Can be used for debugging and analysis.
   """

jax/_src/pallas/hlo_interpreter.py

@@ -352,8 +352,9 @@ def pallas_call_hlo_interpret(
     cost_estimate: CostEstimate,
     out_avals: tuple[jax_core.AbstractValue, ...],
     metadata: frozen_dict.FrozenDict[str, str] | None,
+    name: str | None,
 ):
-  del mesh, compiler_params, cost_estimate, out_avals, metadata
+  del mesh, compiler_params, cost_estimate, out_avals, metadata, name
   debug_info = jaxpr.debug_info
   # If we're in interpret mode, we *scan* over the grid and eval the
   # discharged jaxpr.

jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py

@@ -1959,8 +1959,9 @@ def interpret_pallas_call(
     out_avals: tuple[jax_core.AbstractValue, ...],
     interpret_params: InterpretParams,
     metadata: frozen_dict.FrozenDict[str, str] | None,
+    name: str | None,
 ):
-  del debug, cost_estimate, out_avals
+  del debug, cost_estimate, out_avals, name
   del metadata  # TODO(sharadmv): Add metadata to HLO.
 
   if isinstance(mesh, mosaic_core.TensorCoreMesh):

jax/_src/pallas/mosaic/pallas_call_registration.py

@@ -372,9 +372,10 @@ def pallas_call_tpu_lowering_rule(
     cost_estimate: pallas_core.CostEstimate | None,
     out_avals: tuple[jax_core.AbstractValue, ...],
     metadata: frozen_dict.FrozenDict[str, str] | None,
+    name: str | None,
 ):
   """Lowers a pallas_call to a Mosaic TPU custom call."""
-  del interpret  # Unused.
+  del interpret, name  # Unused.
 
   debug_info = jaxpr.debug_info
   if debug:

jax/_src/pallas/mosaic_gpu/pallas_call_registration.py

@@ -49,8 +49,9 @@ def pallas_call_lowering(
     cost_estimate: pallas_core.CostEstimate | None,
     out_avals: tuple[jax_core.AbstractValue, ...],
     metadata: frozen_dict.FrozenDict[str, str] | None,
+    name: str | None,
 ):
-  del metadata  # TODO(sharadmv): Add metadata to HLO.
+  del metadata, name  # TODO(sharadmv): Add metadata to HLO.
   debug_info = jaxpr.debug_info
   del interpret, out_avals
   if grid_mapping.num_dynamic_grid_bounds:

jax/_src/pallas/pallas_call.py

@@ -152,6 +152,7 @@ def _pallas_call_to_lojax(
     out_avals: tuple[jax_core.AbstractValue, ...],
     backend: Backend | None,
     metadata: FrozenDict[str, str] | None,
+    name: str | None,
 ):
   if any(jax_core.get_aval(x).has_qdd for x in hi_args):
     raise NotImplementedError("pallas_call does not support QDD for inputs")
@@ -221,6 +222,7 @@ def _pallas_call_to_lojax(
       interpret=interpret,
       input_output_aliases=tuple(new_input_output_aliases),
       out_avals=tuple(lo_out_avals),
+      name=name,
   )
   return pe.raise_lo_outs(out_avals, lo_outs)
 pallas_call_p.to_lojax = _pallas_call_to_lojax  # type: ignore
@@ -241,6 +243,7 @@ def _pallas_call_jvp_rule(
     out_avals: tuple[jax_core.AbstractValue, ...],
     backend: Backend | None,
     metadata: FrozenDict[str, str] | None,
+    name: str | None,
 ):
   debug_info = jaxpr.debug_info
   if grid_mapping.num_dynamic_grid_bounds:
@@ -308,6 +311,7 @@ def _pallas_call_jvp_rule(
       out_avals=(*out_avals, *out_avals),
       backend=backend,
       metadata=metadata,
+      name=name,
   )
   out_primals, out_tangents = split_list(out_flat, [len(out_flat) // 2])
   return out_primals, out_tangents
@@ -457,6 +461,7 @@ def _batch_with_explicit_loop(
     out_avals: tuple[jax_core.AbstractValue, ...],
     backend: Backend | None,
     metadata: FrozenDict[str, str] | None,
+    name: str | None,
 ):
   """Batch the pallas_call by calling it in loop over the batch size.
 
@@ -526,6 +531,7 @@ def body(batch_index: jax.Array, state: list[jax.Array]) -> list[jax.Array]:
         out_avals=out_avals,
         backend=backend,
         metadata=metadata,
+        name=name,
     )
     for i, batch_out_array in enumerate(batch_out):
       state[i] = jax.lax.dynamic_update_index_in_dim(
@@ -557,6 +563,7 @@ def _pallas_call_batching_rule(
     out_avals: tuple[jax_core.AbstractValue, ...],
     backend: Backend | None,
     metadata: FrozenDict[str, str] | None = None,
+    name: str | None = None,
 ):
   if mesh is not None:
     raise NotImplementedError(
@@ -596,6 +603,7 @@ def get_size(i, x, d):
         out_avals=out_avals,
         backend=backend,
         metadata=metadata,
+        name=name,
     )
     return [jnp.expand_dims(x, 0) for x in out], (0,) * len(out)
 
@@ -631,6 +639,7 @@ def get_size(i, x, d):
         out_avals=out_avals,
         backend=backend,
         metadata=metadata,
+        name=name,
     )
   else:
     pass  # No dynamic grid dimensions
@@ -667,6 +676,7 @@ def get_size(i, x, d):
           out_avals=out_avals,
           backend=backend,
           metadata=metadata,
+          name=name,
       )
 
   if not dims:
@@ -1048,6 +1058,7 @@ def index_rewrite_kernel(*indexer_args):
       out_avals=batched_out_avals,
       backend=backend,
       metadata=metadata,
+      name=name,
   )
   return out, (0,) * len(out)
 
@@ -1523,6 +1534,7 @@ def _pallas_call_state_discharge_rule(
     out_avals: tuple[jax_core.AbstractValue, ...],
     backend: Backend | None,
     metadata: FrozenDict[str, str] | None,
+    name: str | None,
 ):
   del avals_out
   assert all(isinstance(v.aval, state.AbstractRef) for v in jaxpr.constvars)
@@ -1629,6 +1641,7 @@ def _rewritten_body(*args):
       out_avals=new_out_avals,
       backend=backend,
       metadata=metadata,
+      name=name,
   )
   refs_out, rest = split_list(out_flat, [num_refs])
   updated_vals_in = refs_out + [None] * len(rest_in_avals)
@@ -1909,6 +1922,7 @@ def wrapped(*args):
         cost_estimate=cost_estimate,
         backend=backend,
         metadata=FrozenDict(metadata) if metadata is not None else None,
+        name=name,
     )
     out = tree_util.tree_unflatten(out_tree, out_flat)
     return out

jax/_src/pallas/triton/pallas_call_registration.py

@@ -58,8 +58,9 @@ def pallas_call_lowering(
     cost_estimate: pallas_core.CostEstimate | None,
     out_avals: tuple[jax_core.AbstractValue, ...],
     metadata: frozen_dict.FrozenDict[str, str] | None,
+    name: str | None,
 ):
-  del interpret, out_avals, cost_estimate
+  del interpret, out_avals, cost_estimate, name
   debug_info = jaxpr.debug_info
   if grid_mapping.num_dynamic_grid_bounds:
     raise NotImplementedError(