[Pallas][Easy] Terser printing of GridMapping unless debug is set.

PiperOrigin-RevId: 768932310

Author: Google-ML-Automation

jax/_src/pallas/core.py

@@ -499,6 +499,7 @@ def to_block_mapping(
       index_map_tree: tree_util.PyTreeDef,
       grid: GridMappingGrid,
       mapped_dims: tuple[int, ...],
+      debug: bool = False,
   ) -> BlockMapping:
     if self.index_map is None:
       index_map_func = default_index_map(len(array_aval.shape))
@@ -539,11 +540,15 @@ def to_block_mapping(
 
     fake_index_map_args, fake_index_map_kwargs = \
         index_map_tree.unflatten([False] * index_map_tree.num_leaves)
-    debug = api_util.debug_info("pallas_call index_map",
-                                index_map_func, fake_index_map_args,
-                                fake_index_map_kwargs)
+    debug_info = api_util.debug_info(
+        "pallas_call index_map",
+        index_map_func,
+        fake_index_map_args,
+        fake_index_map_kwargs,
+    )
     flat_index_map_fun, index_map_out_tree_thunk = api_util.flatten_fun(
-        lu.wrap_init(index_map_func, debug_info=debug), index_map_tree)
+        lu.wrap_init(index_map_func, debug_info=debug_info), index_map_tree
+    )
     with tracing_grid_env(grid, mapped_dims):
       jaxpr, out_avals, consts, () = pe.trace_to_jaxpr_dynamic(
           flat_index_map_fun, index_map_avals
@@ -553,7 +558,7 @@ def to_block_mapping(
 
     if len(unflat_avals) != len(block_shape):
       raise ValueError(
-          f"Index map function {debug.func_src_info} for "
+          f"Index map function {debug_info.func_src_info} for "
           f"{origin} must return "
           f"{len(block_shape)} values to match {block_shape=}. "
           f"Currently returning {len(unflat_avals)} values:"
@@ -581,14 +586,14 @@ def to_block_mapping(
     for i, ov in enumerate(out_avals):
       if ov.shape or ov.dtype not in [jnp.int32, jnp.int64]:
         raise ValueError(
-            f"Index map function {debug.func_src_info} for "
+            f"Index map function {debug_info.func_src_info} for "
             f"{origin} must return integer scalars. Output[{i}] has type "
             f"{ov}."
         )
 
     if consts:
       raise ValueError(
-          f"Index map function {debug.func_src_info} for "
+          f"Index map function {debug_info.func_src_info} for "
           f"{origin} must not capture constants: {consts}"
       )
 
@@ -604,6 +609,7 @@ def to_block_mapping(
         ),
         origin=origin,
         pipeline_mode=self.pipeline_mode,
+        debug=debug,
     )
     mapping.check_invariants()
     return mapping
@@ -645,6 +651,7 @@ class BlockMapping:
   origin: OriginStr
   transforms: Sequence[MemoryRefTransform] = ()
   pipeline_mode: Buffered | None = None
+  debug: bool = False
 
   def check_invariants(self) -> None:
     if not config.enable_checks.value: return
@@ -716,6 +723,24 @@ def has_trivial_window(self):
         return False
     return True
 
+  def __repr__(self):
+    if self.debug:
+      return (
+          f"BlockMapping(block_shape={self.block_shape}, "
+          f"transformed_block_aval={self.transformed_block_aval}, "
+          f"index_map_jaxpr={self.index_map_jaxpr}, "
+          f"index_map_out_tree={self.index_map_out_tree}, "
+          f"array_shape_dtype={self.array_shape_dtype}, "
+          f"origin={self.origin}, "
+          f"transforms={self.transforms}, "
+          f"pipeline_mode={self.pipeline_mode}, "
+          f"debug={self.debug})"
+      )
+    return f"BlockMapping(block_shape={self.block_shape})"
+
+  def __str__(self):
+    return self.__repr__()
+
 
 @contextlib.contextmanager
 def tracing_grid_env(grid: GridMappingGrid, mapped_dims: tuple[int, ...]):
@@ -780,6 +805,8 @@ class GridMapping:
   num_scratch_operands: int
   get_grid_indices: Callable | None = None
   local_grid_env: Callable | None = None
+  # Primarily dictates how much debugging information is printed.
+  debug: bool = False
 
   def check_invariants(self) -> None:
     if not config.enable_checks.value: return
@@ -903,6 +930,29 @@ def out_shapes(self) -> Iterable[jax.ShapeDtypeStruct]:
     return tuple(
         bm.array_shape_dtype for bm in self.block_mappings_output)
 
+  def __repr__(self):
+    if self.debug:
+      return (
+          f"GridMapping(grid={self.grid}, grid_names={self.grid_names}, "
+          f"block_mappings={self.block_mappings}, "
+          f"index_map_tree={self.index_map_tree}, "
+          f"index_map_avals={self.index_map_avals}, "
+          f"vmapped_dims={self.vmapped_dims}, "
+          f"num_index_operands={self.num_index_operands}, "
+          f"num_inputs={self.num_inputs}, "
+          f"num_outputs={self.num_outputs}, "
+          f"num_scratch_operands={self.num_scratch_operands}, "
+          f"get_grid_indices={self.get_grid_indices}, "
+          f"local_grid_env={self.local_grid_env}, "
+          f"debug={self.debug})"
+      )
+    return (
+        f"GridMapping(grid={self.grid}, block_mappings={self.block_mappings})"
+    )
+
+  def __str__(self):
+    return self.__repr__()
+
 
 def _is_valid_grid_dim(dim: int | jax.Array) -> bool:
   if isinstance(dim, jax.Array):
@@ -938,6 +988,7 @@ def _convert_block_spec_to_block_mapping(
     index_map_tree: tree_util.PyTreeDef,
     grid: GridMappingGrid,
     mapped_dims: tuple[int, ...],
+    debug: bool = False,
 ) -> BlockMapping:
   if block_spec is no_block_spec:
     block_spec = BlockSpec(None, None)
@@ -948,8 +999,10 @@ def _convert_block_spec_to_block_mapping(
       index_map_tree=index_map_tree,
       grid=grid,
       mapped_dims=mapped_dims,
+      debug=debug,
   )
 
+
 index_map_grid_aval = jax_core.ShapedArray((), jnp.int32)
 
 
@@ -1023,8 +1076,8 @@ def get_grid_mapping(
     out_avals: Sequence[jax_core.AbstractValue],
     out_tree: tree_util.PyTreeDef,
     out_origins: Sequence[OriginStr],
-) -> tuple[tuple[jax_core.AbstractValue, ...],
-           GridMapping]:
+    debug: bool = False,
+) -> tuple[tuple[jax_core.AbstractValue, ...], GridMapping]:
   if dynamic_shapes_export_enabled():
     dim_check : Any = jax_core.is_dim
   else:
@@ -1090,6 +1143,7 @@ def get_grid_mapping(
           index_map_tree=index_map_tree,
           grid=grid_mapping_grid,  # type: ignore[arg-type]
           mapped_dims=(),
+          debug=debug,
       ),
       flat_in_specs,
       in_origins[num_flat_scalar_prefetch:],
@@ -1112,6 +1166,7 @@ def get_grid_mapping(
           index_map_tree=index_map_tree,
           grid=grid_mapping_grid,  # type: ignore[arg-type]
           mapped_dims=(),
+          debug=debug,
       ),
       flat_out_specs,
       out_origins,
@@ -1128,6 +1183,7 @@ def get_grid_mapping(
       num_inputs=len(flat_in_specs),
       num_outputs=len(flat_out_specs),
       num_scratch_operands=num_flat_scratch_operands,
+      debug=debug,
   )
   grid_mapping.check_invariants()
   in_ref_avals = [bm.ref_aval for bm in in_block_mappings]

jax/_src/pallas/mosaic_gpu/core.py

@@ -808,6 +808,7 @@ def to_block_mapping(
       index_map_tree: tree_util.PyTreeDef,
       grid: pallas_core.GridMappingGrid,
       mapped_dims: tuple[int, ...],
+      debug: bool = False,
   ) -> pallas_core.BlockMapping:
     bm = super().to_block_mapping(
         origin,
@@ -816,6 +817,7 @@ def to_block_mapping(
         index_map_tree=index_map_tree,
         grid=grid,
         mapped_dims=mapped_dims,
+        debug=debug,
     )
     block_inner_aval = bm.block_aval.inner_aval
     for t in self.transforms:

jax/_src/pallas/pallas_call.py

@@ -1131,10 +1131,10 @@ def _ensure_2d_error_shape(arg):
   retrace_in_avals = [*shaped_scalar_avals, *error_memref_aval, *input_aval,
                       *error_memref_aval, *output_aval, *scratch_aval]
   jaxpr_flat_avals, jaxpr_in_tree = tree_util.tree_flatten(retrace_in_avals)
-  debug = api_util.debug_info("checkify_pallas", checked_kernel_fn,
+  debug_info = api_util.debug_info("checkify_pallas", checked_kernel_fn,
                               retrace_in_avals, {})
   wrapped_kernel_with_err, out_tree_thunk = api_util.flatten_fun_nokwargs(
-      lu.wrap_init(checked_kernel_fn, debug_info=debug), jaxpr_in_tree)
+      lu.wrap_init(checked_kernel_fn, debug_info=debug_info), jaxpr_in_tree)
 
   with pallas_core.tracing_grid_env(grid_mapping.grid, ()):
     final_jaxpr, _, _, () = pe.trace_to_jaxpr_dynamic(
@@ -1146,13 +1146,18 @@ def _ensure_2d_error_shape(arg):
   error_paths, _ = unzip2(tree_util.tree_flatten_with_path(error_block_specs)[0])
   error_origins = tuple(f"errors[{tree_util.keystr(p)}" for p in error_paths)
   error_block_mappings = map(
-        partial(
-            pallas_core._convert_block_spec_to_block_mapping,
-            index_map_avals=grid_mapping.index_map_avals,
-            index_map_tree=grid_mapping.index_map_tree,
-            grid=grid_mapping.grid,
-            mapped_dims=grid_mapping.vmapped_dims),
-        error_block_specs, error_origins, shaped_err_avals)
+      partial(
+          pallas_core._convert_block_spec_to_block_mapping,
+          index_map_avals=grid_mapping.index_map_avals,
+          index_map_tree=grid_mapping.index_map_tree,
+          grid=grid_mapping.grid,
+          mapped_dims=grid_mapping.vmapped_dims,
+          debug=True,
+      ),
+      error_block_specs,
+      error_origins,
+      shaped_err_avals,
+  )
   input_block_mappings, output_block_mappings = split_list(
       grid_mapping.block_mappings, [num_kernel_inputs,])
   grid_mapping_with_error = grid_mapping.replace(
@@ -1396,7 +1401,9 @@ def _pallas_call_state_discharge_rule(
           index_map_tree=grid_mapping.index_map_tree,
           grid=grid_mapping.grid,
           mapped_dims=grid_mapping.mapped_dims,
-          ) for ref_aval, block_spec in zip(ref_avals, ref_block_specs)
+          debug=debug,
+      )
+      for ref_aval, block_spec in zip(ref_avals, ref_block_specs)
   ]
   in_block_mappings, out_block_mappings = split_list(
       grid_mapping.block_mappings, [grid_mapping.num_inputs]
@@ -1665,8 +1672,14 @@ def wrapped(*args):
     # TODO(necula): check that input_output_aliases is well-formed: no duplicates, etc.
     kernel_args, grid_mapping = pallas_core.get_grid_mapping(
         grid_spec,
-        flat_in_avals, in_tree, in_origins,
-        flat_out_avals, out_tree, out_origins)
+        flat_in_avals,
+        in_tree,
+        in_origins,
+        flat_out_avals,
+        out_tree,
+        out_origins,
+        debug,
+    )
     flat_kernel_args, kernel_in_tree = tree_util.tree_flatten(kernel_args)
     flat_kernel_avals = tuple(
         x.ref if isinstance(x, state_types.TransformedRef) else x