Repository: jax-ml/jax
Pull request #24067 (2024-10-02 09:24:21)

Only thread a discharged ref value through a cond when it changes in some branch
Only thread a discharged ref value through a cond when it changes in some branch

Otherwise, we can simply pass it in as an argument, but we can avoid updating it
since it will always remain constant. Both programs have equivalent semantics,
but this one can be optimized better since it makes it more apparent that the
cond does not actually modify a ref.



Patch:
diff --git a/jax/_src/lax/control_flow/conditionals.py b/jax/_src/lax/control_flow/conditionals.py
index 4385939b8..e654ce953 100644
--- a/jax/_src/lax/control_flow/conditionals.py
+++ b/jax/_src/lax/control_flow/conditionals.py
@@ -858,13 +858,39 @@ mlir.register_lowering(cond_p, _cond_lowering)
 def _cond_state_discharge_rule(should_discharge, in_avals, out_avals, index, *args, branches):
   assert not should_discharge[0], "Can't discharge the index."
   discharged_branches = tuple(
+      discharge_state(branch.jaxpr, (), should_discharge=should_discharge[1:])[0]
+      for branch in branches
+  )
+  # Don't thread the ref values through the cond if they never change.
+  forwarded_outvars = None
+  for branch in discharged_branches:
+    invar_pos = {v: i for i, v in enumerate(branch.invars)}
+    branch_forwarding = [
+        invar_pos.get(v, None) if isinstance(v, core.Var) else None
+        for v in branch.outvars[len(out_avals) :]
+    ]
+    if forwarded_outvars is None:
+      forwarded_outvars = branch_forwarding
+    else:
+      forwarded_outvars = [
+          i if i == j else None
+          for i, j in zip(forwarded_outvars, branch_forwarding)
+      ]
+  assert forwarded_outvars is not None
+  all_outvars_fwd = [None] * len(out_avals) + forwarded_outvars
+  new_branches = tuple(
       core.ClosedJaxpr(
-          discharge_state(branch.jaxpr, (),
-                          should_discharge=should_discharge[1:])[0], ())
-      for branch in branches)
-  out_vals = cond_p.bind(index, *args, branches=discharged_branches)
-  out_vals, out_ref_vals = util.split_list(
-      out_vals, [len(out_avals)])
+          branch.replace(outvars=[v for v, fwd in zip(branch.outvars, all_outvars_fwd)
+                                  if fwd is None]), ())
+      for branch in discharged_branches
+  )
+  out_vals_no_fwd = cond_p.bind(index, *args, branches=new_branches)
+  out_vals, out_ref_vals_no_fwd = util.split_list(out_vals_no_fwd, [len(out_avals)])
+  # Insert forwarded values into reference outputs
+  ref_val_no_fwd_iter = iter(out_ref_vals_no_fwd)
+  out_ref_vals = [next(ref_val_no_fwd_iter) if fwd is None else args[fwd]
+                  for fwd in forwarded_outvars]
+  # Map reference outputs back to their invars
   ref_val_iter = iter(out_ref_vals)
   new_invals = []
   for should, aval in zip(should_discharge, in_avals):


Test patch:
diff --git a/tests/pallas/mosaic_gpu_test.py b/tests/pallas/mosaic_gpu_test.py
index 8c4a31488..0204cb1ad 100644
--- a/tests/pallas/mosaic_gpu_test.py
+++ b/tests/pallas/mosaic_gpu_test.py
@@ -538,12 +538,10 @@ class PallasCallTest(PallasTest):
       assert o_ref.shape == acc_ref.shape == (tile_m, tile_n)
       plgpu.wgmma(acc_ref, a_ref, b_ref)
       plgpu.wgmma_wait(0)  # TODO(apaszke): Delay the pipeline to avoid memory races
-      # TODO(apaszke): Only store in the last step. It doesn't work because we
-      # don't have partial discharge for control flow.
-      # is_last_step = pl.program_id(2) == grid_k - 1
-      # @pl.when(is_last_step)
-      # def _epilogue():
-      o_ref[...] = acc_ref[...].astype(dtype)
+      is_last_step = pl.program_id(2) == grid_k - 1
+      @pl.when(is_last_step)
+      def _epilogue():
+        o_ref[...] = acc_ref[...].astype(dtype)
 
     key1, key2 = jax.random.split(jax.random.key(42), 2)
     a = jax.random.uniform(key1, shape=(m, k), dtype=dtype)
diff --git a/tests/state_test.py b/tests/state_test.py
index c13bb2876..0d6cddfc8 100644
--- a/tests/state_test.py
+++ b/tests/state_test.py
@@ -1114,6 +1114,25 @@ class StateControlFlowTest(jtu.JaxTestCase):
     expected_false = 2.
     self.assertAllClose(out_false, expected_false)
 
+  def test_cond_readonly_refs(self):
+    def f(pred):
+      def body(refs):
+        x_ref, y_ref, z_ref = refs
+        def true_fun():
+          y_ref[()] = x_ref[()]
+        def false_fun():
+          y_ref[()] = x_ref[()] + z_ref[()]
+        lax.cond(pred, true_fun, false_fun)
+      return run_state(body)((1., 0., 2.))
+    jaxpr = jax.make_jaxpr(f)(True).jaxpr
+    [run_state_eqn] = jaxpr.eqns
+    *_, cond_eqn = discharge_state(run_state_eqn.params["jaxpr"], ())[0].eqns
+    self.assertIs(cond_eqn.primitive, lax.cond_p)
+    self.assertLen(cond_eqn.invars, 4)  # pred + 3x ref values
+    self.assertLen(cond_eqn.outvars, 1)  # only the updated ref value
+    self.assertAllClose(jax.jit(f)(True), (1., 1., 2.))
+    self.assertAllClose(jax.jit(f)(False), (1., 3., 2.))
+
   def test_simple_cond_using_multiple_refs_with_interleaved_consts(self):
     def f(pred):
       def body(refs):