Merge pull request #21447 from gnecula:lower_for_export

PiperOrigin-RevId: 638186743

Author: jax authors

jax/_src/api.py

@@ -1846,7 +1846,7 @@ def lower(*args, **kwargs) -> stages.Lowered:
         fun, in_axes, out_axes, static_broadcasted_tuple, donate_tuple,
         devices, backend, axis_size, args, kwargs)
     abstract_args = list(map(shaped_abstractify, p.flat_args))
-    computation = pxla.lower_parallel_callable(
+    computation, closed_jaxpr = pxla.lower_parallel_callable(
         p.flat_fun, backend, axis_name,
         axis_size=p.local_axis_size, global_axis_size=p.global_axis_size,
         devices=p.devices,
@@ -1858,7 +1858,8 @@ def lower(*args, **kwargs) -> stages.Lowered:
         avals=abstract_args,
         lowering_parameters=lowering_parameters)
     return stages.Lowered.from_flat_info(
-        computation, p.in_tree, abstract_args, donate_tuple, p.out_tree())
+        computation, p.in_tree, abstract_args, donate_tuple, p.out_tree(),
+        fun_name=p.flat_fun.__name__, jaxpr=closed_jaxpr)
 
   return lower
 

jax/_src/interpreters/pxla.py

@@ -556,7 +556,7 @@ def parallel_callable(fun: lu.WrappedFun,
                       donated_invars: Sequence[bool],
                       is_explicit_global_axis_size: bool,
                       *avals):
-  pmap_computation = lower_parallel_callable(
+  pmap_computation, _ = lower_parallel_callable(
       fun, backend_name, axis_name, axis_size, global_axis_size, devices, name,
       in_axes, out_axes_thunk, donated_invars,
       is_explicit_global_axis_size, avals,
@@ -679,7 +679,7 @@ def lower_parallel_callable(
     is_explicit_global_axis_size: bool,
     avals: Sequence[core.AbstractValue],
     *,
-    lowering_parameters: mlir.LoweringParameters) -> PmapComputation:
+    lowering_parameters: mlir.LoweringParameters) -> tuple[PmapComputation, core.ClosedJaxpr]:
   # Determine global_axis_size for use in AxisEnv.
   # TODO(mattjj,skyewm): revive this check (inner_pmap always False now)
   # if xb.process_count() > 1 and global_axis_size is None and inner_pmap:
@@ -761,6 +761,7 @@ def lower_parallel_callable(
   tuple_args = dispatch.should_tuple_args(len(shards.global_sharded_avals),
                                           backend.platform)
   module_name = f"pmap_{fun.__name__}"
+  platforms = lowering_parameters.platforms or (backend.platform,)
   with maybe_extend_axis_env(axis_name, global_axis_size, None):
     ordered_effects = list(
         effects.ordered_effects.filter_in(closed_jaxpr.effects))
@@ -776,7 +777,7 @@ def lower_parallel_callable(
           closed_jaxpr,
           ordered_effects=ordered_effects,
           backend_or_name=backend,
-          platforms=lowering_parameters.platforms or (backend.platform,),
+          platforms=platforms,
           axis_context=sharding_impls.ReplicaAxisContext(axis_env),
           name_stack=name_stack,
           donated_args=donated_invars,
@@ -787,14 +788,16 @@ def lower_parallel_callable(
           result_names=jaxpr.debug_info and jaxpr.debug_info.result_paths,
           num_replicas=replicas.num_global_replicas,
           lowering_parameters=lowering_parameters)
-  return PmapComputation(lowering_result.module, pci=pci, replicas=replicas,
+  return PmapComputation(lowering_result.module,
+                         platforms=platforms,
+                         pci=pci, replicas=replicas,
                          shards=shards, tuple_args=tuple_args,
                          unordered_effects=unordered_effects,
                          ordered_effects=ordered_effects,
                          keepalive=lowering_result.keepalive,
                          host_callbacks=lowering_result.host_callbacks,
                          jaxpr_debug_info=closed_jaxpr.jaxpr.debug_info,
-                         shape_poly_state=lowering_result.shape_poly_state)
+                         shape_poly_state=lowering_result.shape_poly_state), closed_jaxpr
 
 
 def _pmap_unmap_shaped_array(
@@ -907,10 +910,13 @@ def from_hlo(hlo: ir.Module,
                host_callbacks: list[Any],
                keepalive: Any,
                jaxpr_debug_info: core.JaxprDebugInfo,
+               platforms: Sequence[str],
                shape_poly_state: mlir.ShapePolyLoweringState | None = None,
                compiler_options=None):
+    del platforms
     if shape_poly_state is not None and shape_poly_state.uses_dim_vars:
       hlo = mlir.refine_polymorphic_shapes(hlo)
+
     devices = pci.devices
     if devices is None:
       if shards.num_global_shards > xb.device_count(pci.backend):
@@ -1941,7 +1947,6 @@ def _cached_lowering_to_hlo(closed_jaxpr, api_name, fun_name, backend,
         "The following ordered effects are not supported for "
         f"more than 1 device: {unsupported_effects}")
   ordered_effects = list(effects.ordered_effects.filter_in(closed_jaxpr.effects))
-
   with dispatch.log_elapsed_time(
         "Finished jaxpr to MLIR module conversion {fun_name} in {elapsed_time} sec",
         fun_name=str(name_stack), event=dispatch.JAXPR_TO_MLIR_MODULE_EVENT):
@@ -2141,6 +2146,7 @@ def lower_sharding_computation(
            for js, source_info in util.stable_unique(jaxpr_sharding))),
       devices_from_context)
 
+  platforms = lowering_parameters.platforms or (backend.platform,)
   # TODO(yashkatariya): Enable this when offload APIs are stable.
   # transfer_mem_kind_in_jaxpr = list(jaxpr_transfer_mem_kinds(jaxpr))
 
@@ -2204,6 +2210,7 @@ def lower_sharding_computation(
       kept_var_idx=kept_var_idx,
       mut=mut,
       backend=backend,
+      platforms=platforms,
       device_assignment=da_object,
       committed=committed,
       in_layouts=in_layouts,
@@ -2244,6 +2251,7 @@ def lower_mesh_computation(
     lowering_parameters: mlir.LoweringParameters) -> MeshComputation:
   assert not mesh.empty
   backend = xb.get_device_backend(mesh.devices.flat[0])
+  platforms = lowering_parameters.platforms or (backend.platform,)
   name_stack = source_info_util.new_name_stack(wrap_name(fun_name, api_name))
 
   global_axis_sizes = mesh.shape
@@ -2352,7 +2360,7 @@ def lower_mesh_computation(
           closed_jaxpr,
           ordered_effects=ordered_effects,
           backend_or_name=backend,
-          platforms=lowering_parameters.platforms or (backend.platform,),
+          platforms=platforms,
           axis_context=axis_ctx,
           name_stack=name_stack,
           donated_args=donated_invars,
@@ -2382,6 +2390,7 @@ def lower_mesh_computation(
       keepalive=lowering_result.keepalive,
       kept_var_idx=set(range(len(global_in_avals))),
       backend=backend,
+      platforms=platforms,
       device_assignment=_create_da_object(tuple(mesh.devices.flat)),
       committed=True,
       in_layouts=(None,) * len(global_in_avals),
@@ -2394,10 +2403,15 @@ class MeshComputation(stages.XlaLowering):
   _executable: MeshExecutable | None
 
   def __init__(self, name: str, hlo: ir.Module,
-               donated_invars: Sequence[bool], **compile_args):
+               donated_invars: Sequence[bool],
+               # TODO(necula): fix this when internal clients stop using this
+               # constructor directly.
+               platforms: Sequence[str] | None = None,
+               **compile_args):
     self._name = name
     self._hlo = hlo
     self._donated_invars = donated_invars
+    self._platforms = platforms
     self.compile_args = compile_args
     self._executable = None
 

jax/_src/maps.py

@@ -617,7 +617,7 @@ def lower(*args, **kwargs):
         '_experimental_lowering_platform', mlir.LoweringParameters())
     fun_flat, args_flat, params, in_tree, out_tree = infer_params(*args)
     avals_flat = [shaped_abstractify(arg) for arg in args_flat]
-    computation = make_xmap_callable(
+    computation, jaxpr = make_xmap_callable(
         fun_flat, params['name'], params['in_axes'], params['out_axes_thunk'],
         params['donated_invars'], params['global_axis_sizes'], params['axis_resources'],
         params['resource_env'], params['backend'], params['spmd_in_axes'],
@@ -628,7 +628,7 @@ def lower(*args, **kwargs):
     in_avals = in_tree.unflatten(avals_flat)
     return stages.Lowered.from_flat_info(
         computation, in_tree, in_avals, donate_argnums, out_tree(),
-        no_kwargs=True)
+        no_kwargs=True, fun_name=params['name'], jaxpr=jaxpr)
 
   fun_mapped.lower = lower
   return type_cast(stages.Wrapped, fun_mapped)
@@ -637,11 +637,12 @@ def xmap_impl(fun: lu.WrappedFun, *args, name, in_axes, out_axes_thunk, donated_
               global_axis_sizes, axis_resources, resource_env, backend,
               spmd_in_axes, spmd_out_axes_thunk):
   in_avals = [core.raise_to_shaped(core.get_aval(arg)) for arg in args]
-  xmap_callable = make_xmap_callable(
+  computation, _ = make_xmap_callable(
       fun, name, in_axes, out_axes_thunk, donated_invars, global_axis_sizes,
       axis_resources, resource_env, backend,
       spmd_in_axes, spmd_out_axes_thunk,
-      mlir.LoweringParameters(), *in_avals).compile().unsafe_call
+      mlir.LoweringParameters(), *in_avals)
+  xmap_callable = computation.compile().unsafe_call
   distributed_debug_log(("Running xmapped function", name),
                         ("python function", fun.f),
                         ("mesh", resource_env.physical_mesh),
@@ -708,15 +709,15 @@ def make_xmap_callable(fun: lu.WrappedFun,
         in_shardings, out_shardings, donated_invars,
         use_spmd_lowering, in_avals,
         tiling_method=tiling_method,
-        lowering_parameters=lowering_parameters)
+        lowering_parameters=lowering_parameters), jaxpr
   else:
     jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(f, in_avals)
     return pxla.lower_sharding_computation(
         core.ClosedJaxpr(jaxpr, consts), 'jit', name,
         (UNSPECIFIED,) * len(in_avals), (UNSPECIFIED,) * len(out_avals),
         (None,) * len(in_avals), (None,) * len(out_avals),
         donated_invars, keep_unused=True, inline=False,
-        devices_from_context=None, lowering_parameters=lowering_parameters)
+        devices_from_context=None, lowering_parameters=lowering_parameters), jaxpr
 
 
 class EvaluationPlan(NamedTuple):

jax/_src/pjit.py

@@ -469,7 +469,7 @@ def lower(*args, **kwargs):
     donate_argnums = tuple(i for i, d in enumerate(donated_invars) if d)
     return stages.Lowered.from_flat_info(
         lowering, in_tree, flat_global_in_avals, donate_argnums,
-        out_tree)
+        out_tree, fun_name=params["name"], jaxpr=params["jaxpr"])
 
   @api_boundary
   def eval_shape(*args, **kwargs):

jax/_src/stages.py

@@ -601,23 +601,30 @@ class Lowered(Stage):
   querying properties of lowered computations across JAX's various
   lowering paths (:func:`~jax.jit`, :func:`~jax.pmap`, etc.).
   """
-  __slots__ = ["args_info", "out_tree", "_lowering", "_no_kwargs"]
-
+  __slots__ = ["_lowering", "args_info", "out_tree", "_no_kwargs", "_fun_name", "_jaxpr"]
+  _lowering: XlaLowering
   args_info: Any                # PyTree of ArgInfo
   out_tree: tree_util.PyTreeDef
-  _lowering: XlaLowering
   _no_kwargs: bool
+  _fun_name: str
+  _jaxpr: core.ClosedJaxpr | None  # Can be None when this class is constructed
+                                   # outside of JAX core.
 
   def __init__(
       self,
       lowering: XlaLowering,
       args_info,  # PyTree of ArgInfo
       out_tree: tree_util.PyTreeDef,
-      no_kwargs: bool = False):
+      no_kwargs: bool = False,
+      fun_name: str = "unknown",
+      jaxpr: core.ClosedJaxpr | None = None):
+
     self._lowering = lowering
-    self._no_kwargs = no_kwargs
     self.args_info = args_info
     self.out_tree = out_tree
+    self._no_kwargs = no_kwargs
+    self._fun_name = fun_name
+    self._jaxpr = jaxpr
 
   @classmethod
   def from_flat_info(cls,
@@ -626,7 +633,9 @@ def from_flat_info(cls,
                      in_avals,
                      donate_argnums: tuple[int, ...],
                      out_tree: tree_util.PyTreeDef,
-                     no_kwargs: bool = False):
+                     no_kwargs: bool = False,
+                     fun_name: str = "unknown",
+                     jaxpr: core.ClosedJaxpr | None = None):
     """Initialize from flat info (``in_avals`` etc.) and an input PyTreeDef.
 
     Args:
@@ -635,12 +644,15 @@ def from_flat_info(cls,
       no_kwargs: If ``True`` the transformation, and the
         ``Compiled`` returned from this object will not support keyword
         arguments (an error will be raised if some are provided).
+      fun_name: the name of the lowered function.
+      jaxpr: the Jaxpr of the lowered function. The value `None` is for
+        backwards compatibility, and is used only outside JAX core.
     """
     return cls(
         lowering,
         make_args_info(in_tree, in_avals, donate_argnums),
         out_tree,
-        no_kwargs=no_kwargs)
+        no_kwargs=no_kwargs, fun_name=fun_name, jaxpr=jaxpr)
 
   def compile(
       self, compiler_options: CompilerOptions | None = None) -> Compiled: