make mlir arg and result names work with pmap

This is a follow-up on #15080 to restore (and indeed fix!) how pmap builds a jaxpr with debug info (i.e. parameter names and result paths). The difference with the machinery in #15080 is just to deal with pmap being final-style (i.e. build the jaxpr at the last second, well after pytrees have been flattened away and transformations have been applied), whereas the machinery for pjit in imagine, plumbing for the former is a bit more long-range and subtle. The main idea here is that we need to annotate and maintain debug info on the lu.WrappedFun instance, which we first form at the api.py level, then pass through transformations (which can either update or drop debug info), then finally hand off to the impl rule to be traced to a jaxpr. It makes sense as an annotation, parallel with the in_type annotation used for dynamic shapes, because the debug info has to be updated as transformations are applied, since they might e.g. add tangent inputs and outputs. In more detail: with an initial-style higher-orer primitive (like pjit), a jaxpr is formed immediately. Transformations, like autodiff, are jaxpr-to-jaxpr, and so those transformations (like ad.jvp_jaxpr) need to return a new jaxpr either with updated debug info or no debug info at all. (The initial implementation in #15080 doesn't provide updated debug info in any of those jaxpr-to-jaxpr transformation functions, so the debug info is only applied to the jaxpr and then lowered to MLIR when the pjit as at the top level.) For final-style, like pmap here, instead of transformations being jaxpr-to-jaxpr, they're WrappedFun-to-WrappedFun. And so, analogously, transformations, like ad.JVPTrace.process_map, would need to produce a WrappedFun with updated debug info or no debug info at all. (ALso analogously to #15080, this PR only implements enough for the debug info to be preserved for top-level pmaps.) This PR doens't yet delete the trace-time debug info in partial_eval.py. But that'll happen too!
jax-ml · joglekara · Mar 21, 2023 · Mar 21, 2023 · Mar 24, 2023 · Mar 24, 2023
commit c5ba4d3daf3aa2b45383733afbab70c39903ceab
diff --git a/jax/_src/api.py b/jax/_src/api.py
@@ -59,7 +59,7 @@
     rebase_donate_argnums, _ensure_index, _ensure_index_tuple,
     shaped_abstractify, _ensure_str_tuple, argnames_partial_except,
     validate_argnames, validate_argnums, check_callable, resolve_argnums,
-    FLAGS)
+    debug_info, result_paths, debug_info_final, FLAGS)
 from jax._src.lax import lax as lax_internal
 from jax._src.lib import jax_jit
 from jax._src.lib import xla_client as xc
@@ -1634,6 +1634,8 @@ def _prepare_pmap(fun, in_axes, out_axes, static_broadcasted_tuple,
   if in_devices is not None and len(in_devices) == 0:
     raise ValueError("'devices' argument to pmap must be non-empty, or None.")
 
+  dbg = debug_info('pmap', fun, args, kwargs, static_broadcasted_tuple, ())
+
   f = lu.wrap_init(fun)
   if static_broadcasted_tuple:
     if max(static_broadcasted_tuple) >= len(args):
@@ -1671,7 +1673,9 @@ def _prepare_pmap(fun, in_axes, out_axes, static_broadcasted_tuple,
       kws=True))
   local_axis_size = _mapped_axis_size(fun, in_tree, args, in_axes_flat, "pmap")
 
+  f, res_paths = result_paths(f)
   flat_fun, out_tree = flatten_fun(f, in_tree)
+  flat_fun = debug_info_final(flat_fun, dbg, res_paths)
 
   if any(out_axis is None for out_axis in tree_flatten(out_axes)):
     raise NotImplementedError("None out_axes in pmap are not supported yet")

diff --git a/jax/_src/api_util.py b/jax/_src/api_util.py
@@ -15,8 +15,8 @@
 import inspect
 import operator
 from functools import partial
-from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Optional,
-                    Sequence, Set, Tuple, Union,)
+from typing import (Any, Callable, Dict, Iterable, List, Optional, Sequence,
+                    Set, Tuple, Union)
 import warnings
 
 import numpy as np
@@ -30,7 +30,9 @@
     treedef_children, generate_key_paths, keystr)
 from jax._src.tree_util import _replace_nones
 from jax._src import linear_util as lu
-from jax._src.util import safe_map, WrapKwArgs, Hashable, Unhashable
+from jax._src.linear_util import TracingDebugInfo
+from jax._src.util import (safe_map, WrapKwArgs, Hashable, HashableFunction,
+                           Unhashable)
 from jax._src.config import flags, bool_env
 from jax._src import traceback_util
 traceback_util.register_exclusion(__file__)
@@ -585,22 +587,14 @@ def api_hook(fun, tag: str):
   return fun
 
 
-class TracingDebugInfo(NamedTuple):
-  # Packages up trace/staging-time debug info about a func and its parameters,
-  # formed just before staging to a jaxpr and read in trace-time error messages.
-  # TODO(mattjj): delete partial_eval.DebugInfo, replace all uses with this cls
-  traced_for: str             # e.g. 'jit', 'scan', etc
-  func_src_info: str          # e.g. f'{fun.__name__} at {filename}:{lineno}'
-  arg_names: Tuple[str, ...]  # e.g. ('args[0]', ... )
-
 def debug_info(traced_for: str, fun: Callable, args: Tuple[Any],
                kwargs: Dict[str, Any], static_argnums: Tuple[int, ...],
                static_argnames: Tuple[str, ...]) -> Optional[TracingDebugInfo]:
   """Try to build trace-time debug info for fun when applied to args/kwargs."""
   src = fun_sourceinfo(fun)
   arg_names = _arg_names(fun, args, kwargs, static_argnums, static_argnames)
   if src is None or arg_names is None: return None
-  return TracingDebugInfo(traced_for, src, arg_names)
+  return TracingDebugInfo(traced_for, src, arg_names, None)
 
 # TODO(mattjj): make this function internal to this module
 def fun_sourceinfo(fun: Callable) -> Optional[str]:
@@ -635,13 +629,23 @@ def result_paths(*args, **kwargs):
   yield ans, [keystr(path) for path, _ in generate_key_paths(ans)]
 
 def jaxpr_debug_info(jaxpr: core.Jaxpr, trace_debug: Optional[TracingDebugInfo],
-                     result_paths: Optional[Tuple[Optional[str], ...]]
+                     result_paths: Optional[Tuple[Optional[str], ...]] = None,
                      ) -> core.Jaxpr:
   """Add debug info to jaxpr, given trace-time debug info and result paths."""
-  if trace_debug is not None and result_paths is not None:
-    debug_info = core.JaxprDebugInfo(
-        trace_debug.traced_for, trace_debug.func_src_info,
-        trace_debug.arg_names, result_paths)
-  else:
-    debug_info = None
-  return jaxpr.replace(debug_info=debug_info) if debug_info else jaxpr
+  if trace_debug is None:
+    return jaxpr
+  assert (result_paths is not None) ^ (trace_debug.result_paths is not None)
+  if result_paths is None:
+    result_paths = trace_debug.result_paths()  # type: ignore
+  debug_info = core.JaxprDebugInfo(
+      trace_debug.traced_for, trace_debug.func_src_info,
+      trace_debug.arg_names, result_paths)
+  return jaxpr.replace(debug_info=debug_info)
+
+def debug_info_final(f: lu.WrappedFun, dbg: Optional[TracingDebugInfo],
+                     res_paths: Callable[[], Tuple[str, ...]]) -> lu.WrappedFun:
+  "Attach trace-time debug info and result paths lazy thunk to an lu.WrappedFun"
+  if dbg is None: return f
+  assert dbg.result_paths is None
+  res_paths_ = HashableFunction(res_paths, closure=())
+  return lu.add_debug_info(f, dbg._replace(result_paths=res_paths_))
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
@@ -1116,6 +1116,7 @@ def stage_parallel_callable(
                                    event=dispatch.JAXPR_TRACE_EVENT):
       jaxpr, out_sharded_avals, consts = pe.trace_to_jaxpr_final(
           fun, global_sharded_avals, pe.debug_info_final(fun, "pmap"))
+  jaxpr = api_util.jaxpr_debug_info(jaxpr, fun.debug_info)
   jaxpr = dispatch.apply_outfeed_rewriter(jaxpr)
 
   assert len(out_sharded_avals) == len(pci.out_axes), (
@@ -1264,7 +1265,6 @@ def lower_parallel_callable(
       raise ValueError("Ordered effects not supported in `pmap`.")
     unordered_effects = list(
         effects.ordered_effects.filter_not_in(closed_jaxpr.effects))
-    arg_names, result_names = _debug_names(jaxpr.debug_info)
     lowering_result = mlir.lower_jaxpr_to_module(
         module_name,
         closed_jaxpr,
@@ -1278,7 +1278,8 @@ def lower_parallel_callable(
         replicated_args=replicated_args,
         arg_shardings=_shardings_to_mlir_shardings(parts.arg_parts),
         result_shardings=_shardings_to_mlir_shardings(parts.out_parts),
-        arg_names=arg_names, result_names=result_names)
+        arg_names=jaxpr.debug_info and jaxpr.debug_info.arg_names,
+        result_names=jaxpr.debug_info and jaxpr.debug_info.result_paths)
     module, keepalive, host_callbacks = (
         lowering_result.module, lowering_result.keepalive,
         lowering_result.host_callbacks)
@@ -2564,7 +2565,6 @@ def lower_sharding_computation(
   unordered_effects = list(
       effects.ordered_effects.filter_not_in(closed_jaxpr.effects))
   ordered_effects = list(effects.ordered_effects.filter_in(closed_jaxpr.effects))
-  arg_names, result_names = _debug_names(jaxpr.debug_info)
   lowering_result = mlir.lower_jaxpr_to_module(
       module_name,
       closed_jaxpr,
@@ -2579,7 +2579,8 @@ def lower_sharding_computation(
       replicated_args=replicated_args,
       arg_shardings=in_op_shardings,
       result_shardings=out_op_shardings,
-      arg_names=arg_names, result_names=result_names)
+      arg_names=jaxpr.debug_info and jaxpr.debug_info.arg_names,
+      result_names=jaxpr.debug_info and jaxpr.debug_info.result_paths)
 
   module, keepalive, host_callbacks = (
       lowering_result.module, lowering_result.keepalive,
@@ -2750,7 +2751,6 @@ def lower_mesh_computation(
       closed_jaxpr.effects))
     ordered_effects = list(effects.ordered_effects.filter_in(
       closed_jaxpr.effects))
-    arg_names, result_names = _debug_names(jaxpr.debug_info)
     lowering_result = mlir.lower_jaxpr_to_module(
         module_name,
         closed_jaxpr,
@@ -2764,8 +2764,8 @@ def lower_mesh_computation(
         replicated_args=replicated_args,
         arg_shardings=in_partitions,
         result_shardings=out_partitions,
-        arg_names=arg_names,
-        result_names=result_names)
+        arg_names=jaxpr.debug_info and jaxpr.debug_info.arg_names,
+        result_names=jaxpr.debug_info and jaxpr.debug_info.result_paths)
     module, keepalive, host_callbacks = (
         lowering_result.module, lowering_result.keepalive,
         lowering_result.host_callbacks)
@@ -2791,12 +2791,6 @@ def lower_mesh_computation(
       device_assignment=list(mesh.devices.flat),
       committed=True)
 
-def _debug_names(
-    dbg: Optional[core.JaxprDebugInfo]
-) -> Union[Tuple[None, None],
-           Tuple[Sequence[Optional[str]], Sequence[Optional[str]]]]:
-  return (None, None) if dbg is None else (dbg.arg_names, dbg.result_paths)
-
 class MeshComputation(stages.XlaLowering):
   _hlo: Optional[ir.Module]
   _executable: Optional[MeshExecutable]

diff --git a/jax/_src/linear_util.py b/jax/_src/linear_util.py
@@ -64,7 +64,7 @@ def trans1(static_arg, *dynamic_args, **kwargs):
 from __future__ import annotations
 
 from functools import partial
-from typing import Any, Tuple, Callable
+from typing import Any, Tuple, Callable, Optional, NamedTuple
 import weakref
 
 from jax.tree_util import tree_map
@@ -124,14 +124,15 @@ class WrappedFun:
     params: extra parameters to pass as keyword arguments to `f`, along with the
       transformed keyword arguments.
   """
-  __slots__ = ("f", "transforms", "stores", "params", "in_type")
+  __slots__ = ("f", "transforms", "stores", "params", "in_type", "debug_info")
 
-  def __init__(self, f, transforms, stores, params, in_type):
+  def __init__(self, f, transforms, stores, params, in_type, debug_info):
     self.f = f
     self.transforms = transforms
     self.stores = stores
     self.params = params
     self.in_type = in_type
+    self.debug_info = debug_info
 
   @property
   def __name__(self):
@@ -140,7 +141,7 @@ def __name__(self):
   def wrap(self, gen, gen_static_args, out_store) -> WrappedFun:
     """Add another transform and its store."""
     return WrappedFun(self.f, ((gen, gen_static_args),) + self.transforms,
-                      (out_store,) + self.stores, self.params, None)
+                      (out_store,) + self.stores, self.params, None, None)
 
   def populate_stores(self, stores):
     """Copy the values from the `stores` into `self.stores`."""
@@ -199,11 +200,13 @@ def transform_to_str(x):
     return "Wrapped function:\n" + '\n'.join(transformation_stack) + '\nCore: ' + fun_name(self.f) + '\n'
 
   def __hash__(self):
-    return hash((self.f, self.transforms, self.params, self.in_type))
+    return hash((self.f, self.transforms, self.params, self.in_type,
+                 self.debug_info))
 
   def __eq__(self, other):
     return (self.f == other.f and self.transforms == other.transforms and
-            self.params == other.params and self.in_type == other.in_type)
+            self.params == other.params and self.in_type == other.in_type and
+            self.debug_info == other.debug_info)
 
 @curry
 def transformation(gen, fun: WrappedFun, *gen_static_args) -> WrappedFun:
@@ -231,15 +234,15 @@ def fun_name(f):
 def wrap_init(f, params=None) -> WrappedFun:
   """Wraps function `f` as a `WrappedFun`, suitable for transformation."""
   params = () if params is None else tuple(sorted(params.items()))
-  return WrappedFun(f, (), (), params, None)
+  return WrappedFun(f, (), (), params, None, None)
 
 
-def annotate(f: WrappedFun, in_type: core.InputType) -> WrappedFun:
+def annotate(f: WrappedFun, in_type: Optional[core.InputType]) -> WrappedFun:
   assert f.in_type is None
   if in_type is None:
     return f
   _check_input_type(in_type)
-  return WrappedFun(f.f, f.transforms, f.stores, f.params, in_type)
+  return WrappedFun(f.f, f.transforms, f.stores, f.params, in_type, f.debug_info)
 
 def _check_input_type(in_type: core.InputType) -> None:
   # Check that in_type is syntactically well-formed
@@ -271,6 +274,24 @@ def valid_size(d) -> bool:
   assert all(provided)
 
 
+class TracingDebugInfo(NamedTuple):
+  # Packages up trace/staging-time debug info about a func and its parameters,
+  # formed just before staging to a jaxpr and read in trace-time error messages.
+  # TODO(mattjj): delete partial_eval.DebugInfo, replace all uses with this cls
+  traced_for: str             # e.g. 'jit', 'scan', etc
+  func_src_info: str          # e.g. f'{fun.__name__} at {filename}:{lineno}'
+  arg_names: Tuple[str, ...]  # e.g. ('args[0]', ... )
+  result_paths: Optional[Callable[[], Tuple[str, ...]]]
+
+def add_debug_info(f: WrappedFun, debug_info: Optional[TracingDebugInfo]
+                   ) -> WrappedFun:
+  """Produce a new WrappedFun with debug_info attached."""
+  assert f.debug_info is None
+  if debug_info is None:
+    return f
+  return WrappedFun(f.f, f.transforms, f.stores, f.params, f.in_type, debug_info)
+
+
 def cache(call: Callable):
   """Memoization decorator for functions taking a WrappedFun as first argument.
 

diff --git a/tests/pmap_test.py b/tests/pmap_test.py
@@ -2036,7 +2036,6 @@ def test_remat_of_pmap_policy(self, remat):
     self.assertEqual(jaxpr_text.count(' cos '), 2)
 
   def test_pmap_lower_arg_info(self):
-    raise SkipTest("arg info not plumbed to pmap yet")  # TODO(mattjj)
     def f(x, y, *args, **kwargs):
       return y['hi'] + args[1] + sum(kwargs.values())
 
@@ -2052,7 +2051,6 @@ def f(x, y, *args, **kwargs):
     self.assertIn("kwargs['w']", mhlo_str)
 
   def test_pmap_lower_result_info(self):
-    raise SkipTest("arg info not plumbed to pmap yet")  # TODO(mattjj)
     def f(x, y, z):
       return {'a': x, 'b': [y]}