[Pallas Fuser] Add output_fusion_mask support

Currently, the fusion API assumes by default that all of the outputs of a @fuse-decorated function are computed jointly in one big output fusion.

For example, in the following snippet
```python
@fuse
def f(x, y):
  z1, z2 = fusable_f(x, y)
  return g(z1, z2)
```

it assumes that `g` is a single function that operates on z1 and z2 jointly. However, in practice, the fusable may want two separate output fusions:
```python
@fuse
def f(x, y):
  z1, z2 = fusable_f(x, y)
  return g1(z1), g2(z2)
```

This is a special case of the general function but the fusable may not be materializing z1 and z2 at the same time so may not be able to compute this efficiently with a single function g.

By decorating a fusable with an output fusion prefix (in the above example `(True, True)`), the fusable will now be given a pair of functions `g1` and `g2` if the output fusion is "separable". For example, we'd error for the following example:

```python
@fuse
def f(x, y):
  z1, z2 = fusable_f(x, y)
  return z1 + z2
```

because z1 and z2 interact with each other in the output fusion.

The rationale for providing a PyTree prefix (as opposed to a more general mechanism) is that the fusable can group its outputs into subtrees that it can identify with the output prefix. This does restrict the types of output groups that are possible (outputs must be part of the same shared subtree, as opposed to arbitrarily scattered throughput the output pytree), but this is an okay restriction because the fusable author is responsible for the grouping and can always construct it that way.

PiperOrigin-RevId: 744784770

Author: sharadmv

jax/_src/pallas/fuser/BUILD

@@ -99,6 +99,7 @@ pytype_strict_library(
         "//jax:core",
         "//jax:partial_eval",
         "//jax:tree_util",
+        "//jax:util",
     ],
 )
 

jax/_src/pallas/fuser/fusable.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Fusable primitive."""
+from typing import Any
 
 import jax
 from jax._src import api_util
@@ -40,32 +41,38 @@ def _make_trivial_fusion(x: jax.Array) -> fusion_lib.Fusion:
   )
 
 
-def fusable(f):
-  def wrapper(*args):
-    def wrapped(*args):
-      in_fusions = tree_util.tree_map(_make_trivial_fusion, args)
-      return f(*in_fusions, None)
-
-    flat_args, in_tree = tree_util.tree_flatten(args)
-    debug_info = api_util.debug_info('fusable', wrapped, args, {})
-    flat_fun, out_tree_thunk = api_util.flatten_fun_nokwargs(
-        lu.wrap_init(wrapped, debug_info=debug_info), in_tree
-    )
-    flat_avals = [_get_aval(x) for x in flat_args]
-    jaxpr, _, consts, _ = pe.trace_to_jaxpr_dynamic(flat_fun, flat_avals)
-    out_tree = out_tree_thunk()
-    out = fusable_p.bind(
-        *consts,
-        *flat_args,
-        jaxpr=jaxpr,
-        num_consts=len(consts),
-        in_tree=in_tree,
-        out_tree=out_tree,
-        func=f,
-    )
-    return tree_util.tree_unflatten(out_tree, out)
-
-  return wrapper
+def fusable(f=None, *, output_fusion_prefix: Any = True):
+  def decorator(f):
+    def wrapper(*args):
+      def wrapped(*args):
+        in_fusions = tree_util.tree_map(_make_trivial_fusion, args)
+        return f(*in_fusions, None)
+
+      flat_args, in_tree = tree_util.tree_flatten(args)
+      debug_info = api_util.debug_info('fusable', wrapped, args, {})
+      flat_fun, out_tree_thunk = api_util.flatten_fun_nokwargs(
+          lu.wrap_init(wrapped, debug_info=debug_info), in_tree
+      )
+      flat_avals = [_get_aval(x) for x in flat_args]
+      jaxpr, _, consts, _ = pe.trace_to_jaxpr_dynamic(flat_fun, flat_avals)
+      out_tree = out_tree_thunk()
+      out = fusable_p.bind(
+          *consts,
+          *flat_args,
+          jaxpr=jaxpr,
+          num_consts=len(consts),
+          in_tree=in_tree,
+          out_tree=out_tree,
+          func=f,
+          output_fusion_prefix=output_fusion_prefix,
+      )
+      return tree_util.tree_unflatten(out_tree, out)
+
+    return wrapper
+
+  if f is not None:
+    return decorator(f)
+  return decorator
 
 
 @fusable_p.def_impl

jax/_src/pallas/fuser/jaxpr_fusion.py

@@ -14,15 +14,15 @@
 
 """Fuses a function."""
 
+from collections.abc import Sequence
+import functools
 from typing import Any
-
 import jax
 from jax._src import api_util
 from jax._src import core as jax_core
 from jax._src import linear_util as lu
 from jax._src import tree_util
 from jax._src.interpreters import partial_eval as pe
-
 from jax._src.pallas.fuser import fusable_dtype
 from jax._src.pallas.fuser import fusion as fusion_lib
 from jax._src.pallas.fuser.fusable import fusable_p
@@ -73,9 +73,9 @@ def wrapper(*args, **kwargs):
 _fusable: dict[jax_core.Primitive, Any] = {}
 
 
-def construct_fusion(
+def _construct_fusion_jaxpr(
     candidate_values, jaxpr: jax_core.Jaxpr, outvars, *invars, **kwargs
-) -> fusion_lib.Fusion:
+):
   flat_outvars, out_tree = tree_util.tree_flatten(outvars)
   flat_invars, in_tree = tree_util.tree_flatten((invars, kwargs))
   new_jaxpr_no_dce = jaxpr.replace(
@@ -94,12 +94,6 @@ def construct_fusion(
       c for used, c in zip(used_consts, candidate_values, strict=True) if used
   )
   kernel_in_tree = tree_util.tree_structure((invars, kwargs))
-
-  def _fn(*args, **kwargs):
-    flat_args, _ = tree_util.tree_flatten((args, kwargs))
-    out_flat = jax_core.eval_jaxpr(new_jaxpr, new_values, *flat_args)
-    return tree_util.tree_unflatten(out_tree, out_flat)
-
   flat_in_type = [
       jax.ShapeDtypeStruct(x.aval.shape, x.aval.dtype) for x in flat_invars
   ]
@@ -108,9 +102,158 @@ def _fn(*args, **kwargs):
       out_tree,
       [jax.ShapeDtypeStruct(x.aval.shape, x.aval.dtype) for x in flat_outvars],
   )
+  return new_jaxpr, new_values, in_type, out_type, out_tree
+
+
+def construct_fusion(
+    candidate_values, jaxpr: jax_core.Jaxpr, outvars, *invars, **kwargs
+) -> fusion_lib.Fusion:
+  new_jaxpr, new_values, in_type, out_type, out_tree = _construct_fusion_jaxpr(
+      candidate_values, jaxpr, outvars, *invars, **kwargs
+  )
+
+  def _fn(*args, **kwargs):
+    flat_args, _ = tree_util.tree_flatten((args, kwargs))
+    out_flat = jax_core.eval_jaxpr(new_jaxpr, new_values, *flat_args)
+    return tree_util.tree_unflatten(out_tree, out_flat)
+
   return fusion_lib.Fusion(_fn, in_type, out_type)
 
 
+def _find_downstream(
+    jaxpr: jax_core.Jaxpr, in_used: Sequence[bool]
+) -> tuple[bool, ...]:
+  # TODO(sharadmv): We use partial_eval to query downstream dependencies which
+  # is not an officially sanctioned way to do so, since PE is really used for
+  # AD. In the future, we should have a special Jaxpr API that queries this.
+  _, _, out_used, *_ = pe.partial_eval_jaxpr_custom(
+      jaxpr,
+      in_unknowns=in_used,
+      in_inst=in_used,
+      ensure_out_unknowns=False,
+      ensure_out_inst=False,
+      saveable=lambda *_, **__: False,
+  )
+  return tuple(out_used)
+
+
+def _construct_output_permutation(
+    used: list[tuple[bool, ...]],
+) -> list[int]:
+  order = []
+  for u in used:
+    true_vals = [i for i in range(len(u)) if u[i]]
+    order.extend(true_vals)
+  return [order.index(i) for i in range(len(order))]
+
+
+def _construct_output_fusions(
+    candidate_values,
+    jaxpr,
+    out_tree,
+    fusion_eqn_index,
+    fusion_eqn_outvars,  # Flat list of vars output by the fusable eqn
+    fusion_eqn_out_tree,  # Tree structure of the fusable eqn outputs
+    output_fusion_prefix,  # Pytree defining output groups
+):
+  # 1. Create jaxpr_out: represents computation *after* the fusable
+  #    Inputs: fusion_eqn_outvars
+  #    Outputs: jaxpr.outvars
+  jaxpr_out, all_values, _, _, _ = _construct_fusion_jaxpr(
+      candidate_values,
+      jaxpr.replace(
+          eqns=jaxpr.eqns[:fusion_eqn_index]
+          + jaxpr.eqns[fusion_eqn_index + 1 :]
+      ),
+      tree_util.tree_unflatten(out_tree, jaxpr.outvars),  # Original outputs
+      tree_util.tree_unflatten(
+          fusion_eqn_out_tree, fusion_eqn_outvars
+      ),  # Fusable outputs as inputs
+  )
+
+  # 2. Group fusable outputs based on the mask
+  unflat_fusable_outvars = jax.tree.unflatten(
+      fusion_eqn_out_tree, fusion_eqn_outvars
+  )
+  partial_flat = jax.tree.structure(output_fusion_prefix).flatten_up_to(
+      unflat_fusable_outvars
+  )
+
+  # 3. Calculate dependencies and check disjointness
+  downstream_outputs_used_masks = []  # List of bool tuples, one per group
+  already_used_final_outputs = set()  # Indices of final outputs already claimed
+  for outvars_group in partial_flat:
+    # Identify vars in this group
+    used_fusable_outvars = set(jax.tree.leaves(outvars_group))
+    # Create mask for jaxpr_out inputs corresponding to this group
+    in_used_mask = [
+        True if v in used_fusable_outvars else False for v in jaxpr_out.invars
+    ]
+    # Trace dependencies through jaxpr_out to find which final outputs are affected
+    downstream_used_mask = _find_downstream(
+        jaxpr_out, in_used_mask
+    )  # Mask for jaxpr_out.outvars (== jaxpr.outvars)
+
+    # Check for overlap in final output usage across groups
+    for i, used in enumerate(downstream_used_mask):
+      if used:
+        if i in already_used_final_outputs:
+          raise ValueError(
+              "Outputs must be disjoint in order to use separate output fusions"
+          )
+        already_used_final_outputs.add(i)
+    downstream_outputs_used_masks.append(downstream_used_mask)
+
+  # 4. Construct output permutation needed to restore original output order
+  output_permutation = _construct_output_permutation(
+      downstream_outputs_used_masks
+  )
+
+  # Construct fusions for each group by DCEing the jaxpr_out
+  output_fusions = []
+  for i, outvars_group in enumerate(partial_flat):
+    flat_group_vars, _ = tree_util.tree_flatten(outvars_group)
+    downstream_used_mask = downstream_outputs_used_masks[i]
+
+    used_jaxpr_invars = [False] * len(all_values) + [
+        v in flat_group_vars for v in jaxpr_out.invars
+    ]
+    jaxpr_out_for_group, used_consts, _ = pe.dce_jaxpr_consts(
+        jaxpr_out, downstream_used_mask, instantiate=used_jaxpr_invars
+    )
+    values_for_jaxpr = tuple(
+        c for used, c in zip(used_consts, all_values, strict=True) if used
+    )
+
+    def _fn(jaxpr, vals, *args, **kwargs):
+      flat_args, _ = tree_util.tree_flatten((args, kwargs))
+      out_flat = jax_core.eval_jaxpr(jaxpr, vals, *flat_args)
+      return tuple(out_flat)
+
+    fn = functools.partial(_fn, jaxpr_out_for_group, values_for_jaxpr)
+    in_type = jax.tree.map(
+        lambda v: jax.ShapeDtypeStruct(v.aval.shape, v.aval.dtype),  # pytype: disable=attribute-error
+        outvars_group,
+    )
+    out_type = tuple(
+        jax.ShapeDtypeStruct(v.aval.shape, v.aval.dtype)  # pytype: disable=attribute-error
+        for v in jaxpr_out_for_group.outvars
+    )
+    fusion = fusion_lib.Fusion(
+        fn,
+        (in_type, {}),
+        out_type,
+    )
+    output_fusions.append(fusion)
+
+  return (
+      tree_util.tree_unflatten(
+          tree_util.tree_structure(output_fusion_prefix), output_fusions
+      ),
+      output_permutation,
+  )
+
+
 def fuse_jaxpr(
     jaxpr: jax_core.Jaxpr, out_tree: tree_util.PyTreeDef, consts, *args
 ):
@@ -125,6 +268,15 @@ def fuse_jaxpr(
     raise ValueError("No fusable eqn found")
   fusion_eqn = jaxpr.eqns[fusion_eqn_index]
 
+  # Now let's check if we need to do any fusion at all, e.g. do the outputs of
+  # the jaxpr have any dependence on the fusion at all? We can DCE the jaxpr
+  # with all the inputs and outputs to check if there is a dependence.
+  dced_jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars),
+                               instantiate=True)
+  if not any(eqn.primitive is fusable_p for eqn in dced_jaxpr.eqns):
+    # Short circuit if there is nothing to fuse.
+    return jax_core.eval_jaxpr(dced_jaxpr, consts, *args)
+
   candidate_values = [*consts, *args]
 
   # Construct fusions for non-constant inputs to the fusable.
@@ -141,21 +293,20 @@ def fuse_jaxpr(
   in_fusions = tree_util.tree_unflatten(
       fusion_eqn.params["in_tree"], in_fusions_flat
   )
-  out_fusion = construct_fusion(
+  output_fusions, output_permutation = _construct_output_fusions(
       candidate_values,
-      jaxpr.replace(
-          eqns=jaxpr.eqns[:fusion_eqn_index]
-          + jaxpr.eqns[fusion_eqn_index + 1 :]
-      ),
-      tree_util.tree_unflatten(out_tree, jaxpr.outvars),
-      tree_util.tree_unflatten(
-          fusion_eqn.params["out_tree"], fusion_eqn.outvars
-      ),
+      jaxpr,
+      out_tree,
+      fusion_eqn_index,
+      fusion_eqn.outvars,
+      fusion_eqn.params["out_tree"],
+      fusion_eqn.params["output_fusion_prefix"],
   )
-  # Run the fusable.
-  out = fusion_eqn.params["func"](*in_fusions, out_fusion)
-
-  # Now return the flattened output (the fuse_jaxpr caller should unflatten).
-  out_flat = tree_util.tree_leaves(out)
-  assert len(out_flat) == len(jaxpr.outvars)
-  return out_flat
+  out = fusion_eqn.params["func"](*in_fusions, output_fusions)
+  flat_out = jax.tree.leaves(out)
+  permuted_out = [flat_out[i] for i in output_permutation]
+  assert len(permuted_out) == len(jaxpr.outvars), (
+      len(permuted_out),
+      len(jaxpr.outvars),
+  )
+  return permuted_out

tests/pallas/BUILD

@@ -680,6 +680,27 @@ jax_multiplatform_test(
     ] + py_deps("absl/testing") + py_deps("numpy"),
 )
 
+jax_multiplatform_test(
+    name = "fusion_test",
+    srcs = [
+        "fusion_test.py",
+    ],
+    disable_configs = [
+        "cpu",
+        "cpu_shardy",
+    ],
+    enable_backends = ["cpu"],
+    tags = [
+        "noasan",
+        "nomsan",
+        "notsan",
+    ],
+    deps = [
+        "//jax:pallas",
+        "//jax:pallas_fuser",
+    ] + py_deps("absl/testing") + py_deps("numpy"),
+)
+
 jax_multiplatform_test(
     name = "tpu_fusable_matmul_test",
     srcs = ["tpu_fusable_matmul_test.py"],