Add partial_eval_custom rule for for_loop

jax/_src/lax/control_flow/for_loop.py

@@ -33,7 +33,7 @@
 from jax._src import source_info_util
 from jax._src import state
 from jax._src.util import (partition_list, merge_lists, safe_map, safe_zip,
-                           split_list)
+                           split_list, split_dict)
 import jax.numpy as jnp
 import numpy as np
 
@@ -319,6 +319,7 @@ def _for_partial_eval(trace: pe.JaxprTrace, *tracers: pe.JaxprTracer,
                       jaxpr: core.Jaxpr, nsteps: int, reverse: bool,
                       which_linear: Tuple[bool, ...]) -> List[pe.JaxprTracer]:
   num_inputs = len(tracers)
+  assert num_inputs == len(jaxpr.invars) - 1
   in_unknowns = [not t.pval.is_known() for t in tracers]
   # We first need to run a fixpoint to determine which of the `Ref`s are unknown
   # after running the for loop. We want to use the jaxpr to determine which
@@ -446,6 +447,135 @@ def _for_partial_eval(trace: pe.JaxprTrace, *tracers: pe.JaxprTracer,
   return merge_lists(in_unknowns, known_outputs, unknown_outputs)
 pe.custom_partial_eval_rules[for_p] = _for_partial_eval
 
+def _for_partial_eval_custom(saveable, in_unknowns, in_inst, eqn):
+  jaxpr, nsteps, reverse, which_linear = split_dict(
+      eqn.params, ["jaxpr", "nsteps", "reverse", "which_linear"])
+  num_inputs = len(eqn.invars)
+  # We first need to run a fixpoint to determine which of the `Ref`s are unknown
+  # after running the for loop. However, the jaxpr has no outputs. Instead, we
+  # discharge the body and run the fixpoint with the discharged jaxpr. We can do
+  # this because the outputs of the discharged jaxpr are one-to-one with the
+  # inputs.
+  discharged_jaxpr, discharged_consts = discharge_state(jaxpr, ())
+  discharged_jaxpr = discharged_jaxpr.replace(
+      invars=discharged_jaxpr.constvars + discharged_jaxpr.invars,
+      constvars=[])
+  in_unknowns, in_inst = list(in_unknowns), list(in_inst)
+  for _ in range(num_inputs):
+    jaxpr_in_unknowns = [False] * len(discharged_consts) + [False, *in_unknowns]
+    _, _, out_unknowns, inst_out, _, = pe.partial_eval_jaxpr_custom(
+        discharged_jaxpr, jaxpr_in_unknowns, True,
+          ensure_out_unknowns=in_unknowns, ensure_out_inst=True,
+          saveable=saveable)
+    out_unknowns = list(out_unknowns)
+    if out_unknowns == in_unknowns:
+      break
+    in_unknowns = map(operator.or_, in_unknowns, out_unknowns)
+  else:
+    raise Exception("Invalid fixpoint")
+  del out_unknowns # Redundant since it's the same as `in_unknowns`
+  new_inst = [x for x, inst in zip(eqn.invars, in_inst)
+              if type(x) is core.Var and not inst]
+  in_inst = [True] * len(eqn.invars)
+
+  # We use `partial_eval_jaxpr_custom` here because it won't remove effectful
+  # primitives like `get`/`set`.
+  jaxpr_known_resout, jaxpr_staged_resin_, _, _, num_res = \
+        pe.partial_eval_jaxpr_custom(jaxpr, [False, *in_unknowns],
+            [True, *in_inst], [], [], saveable)
+
+  # `partial_eval_jaxpr_custom` will give us jaxprs that have hybrid `Ref` and
+  # non-Ref input/outputs. However, we'd like to bind these jaxprs to a
+  # `for`, which expects only `Ref` inputs and no output. We need to convert
+  # both of these jaxprs into ones that are compatible with `for`.
+  # TODO(sharadmv,mattjj): implement "passthrough" optimization.
+  # TODO(sharadmv,mattjj): rematerialize loop-dependent values instead of
+  # passing the loop index as a residual
+
+  # `jaxpr_known_resout` is a jaxpr that maps from all the input `Refs`
+  # to output residual values (none of them should be `Ref`s). We'll need to
+  # convert the output residual values into `Ref`s that are initially empty
+  # `Ref`s that are written to at the end of the jaxpr.
+
+  # # Loop-invariant residual optimization
+  # Here we are interested in finding out which of the residuals are *not*
+  # dependent on the loop index. If a residual is not dependent on the loop
+  # index, we don't need add an extra loop dimension we're reading from when we
+  # convert it from an output into a write.
+
+  # In order to detect which residuals are loop-invariant, we need to run a
+  # fixpoint. This is because the residual could be dependent on a `Ref` that
+  # changes each iteration of the loop so we need to first detect which `Ref`s
+  # are loop-varying. We can do this by discharging the state from the jaxpr and
+  # running partial_eval with initially only the loop-index being loop-varying.
+  # The fixpoint will eventually propagate the loop-varying-ness over the
+  # inputs/outputs and we will converge.
+  loop_var_res = [False] * len(jaxpr_known_resout.outvars)
+  loop_var_refs = [False] * (len(jaxpr_known_resout.invars) - 1)
+  discharged_jaxpr_known_resout = core.ClosedJaxpr(
+      *discharge_state(jaxpr_known_resout, ()))
+  for _ in range(len(discharged_jaxpr_known_resout.jaxpr.invars)):
+    (_, _, loop_var_outputs, _) = pe.partial_eval_jaxpr_nounits(
+          discharged_jaxpr_known_resout, [True] + loop_var_refs, False)
+    loop_var_res, loop_var_refs_ = split_list(
+        loop_var_outputs, [len(loop_var_res)])
+    if loop_var_refs == loop_var_refs_:
+      break
+    loop_var_refs = map(operator.or_, loop_var_refs, loop_var_refs_)
+  # Now that the fixpoint is complete, we know which residuals are
+  # loop-invariant.
+  loop_invar_res = map(operator.not_, loop_var_res)
+
+  jaxpr_known, res_avals = _convert_outputs_to_writes(nsteps,
+                                                      jaxpr_known_resout,
+                                                      loop_invar_res)
+
+  known_invars, _ = partition_list(in_unknowns, eqn.invars)
+  known_outvars, _ = partition_list(in_unknowns, eqn.outvars)
+  newvar = core.gensym()
+  resvars = map(newvar, res_avals)
+
+  @lu.wrap_init
+  def known(*known_vals):
+    empty_res = map(ad_util.zeros_like_aval, res_avals)
+    jaxpr_known_args = [*known_vals, *empty_res]
+    jaxpr_known_which_linear = (False,) * len(jaxpr_known_args)
+    return for_p.bind(*jaxpr_known_args, jaxpr=jaxpr_known, nsteps=nsteps,
+                      reverse=reverse, which_linear=jaxpr_known_which_linear)
+  call_jaxpr_, _, call_jaxpr_consts = pe.trace_to_jaxpr_dynamic(
+      known, [v.aval for v in known_invars])
+  call_jaxpr = core.ClosedJaxpr(call_jaxpr_, call_jaxpr_consts)
+  eqn_known = pe.new_jaxpr_eqn(known_invars, [*known_outvars, *resvars],
+                               core.closed_call_p, dict(call_jaxpr=call_jaxpr),
+                               call_jaxpr.effects, eqn.source_info)
+
+  jaxpr_staged = _convert_inputs_to_reads(nsteps, len(res_avals),
+                                          jaxpr_staged_resin_,
+                                          loop_invar_res)
+  which_linear_unknown = (False,) * num_res + tuple(which_linear)
+  params_staged = dict(eqn.params, jaxpr=jaxpr_staged, reverse=reverse,
+                                   nsteps=nsteps,
+                                   which_linear=which_linear_unknown)
+
+  @lu.wrap_init
+  def staged(*res_and_refs):
+    out_flat = for_p.bind(*res_and_refs, **params_staged)
+    _, ans = split_list(out_flat, [num_res])
+    _, ans = partition_list(inst_out, ans)
+    return ans
+  call_jaxpr_, _, call_jaxpr_consts = pe.trace_to_jaxpr_dynamic(
+      staged, [v.aval for v in [*resvars, *eqn.invars]])
+  assert len(jaxpr_staged.invars) - 1 == len(call_jaxpr_.invars)
+  call_jaxpr = core.ClosedJaxpr(call_jaxpr_, call_jaxpr_consts)
+  _, outvars = partition_list(inst_out, eqn.outvars)
+  eqn_staged = pe.new_jaxpr_eqn([*resvars, *eqn.invars], outvars,
+                               core.closed_call_p, dict(call_jaxpr=call_jaxpr),
+                               call_jaxpr.effects, eqn.source_info)
+  new_vars = [*new_inst, *resvars]
+  return eqn_known, eqn_staged, in_unknowns, inst_out, new_vars
+
+pe.partial_eval_jaxpr_custom_rules[for_p] = _for_partial_eval_custom
+
 def _convert_outputs_to_writes(
     nsteps: int, jaxpr: core.Jaxpr, loop_invar_res: Sequence[bool]
     ) -> Tuple[core.Jaxpr, List[core.ShapedArray]]:

jax/_src/state/primitives.py

@@ -328,10 +328,10 @@ def _state_partial_eval_custom(prim, saveable, unks_in, inst_in, eqn):
   if any(unks_in):
     res = [v for v, inst in zip(eqn.invars, inst_in) if not inst]
     return None, eqn, [True] * len(eqn.outvars), [True] * len(eqn.outvars), res
-  elif saveable(get_p, *[var.aval for var in eqn.invars], **eqn.params):
+  elif saveable(prim, *[var.aval for var in eqn.invars], **eqn.params):
     return eqn, None, [False] * len(eqn.outvars), [False] * len(eqn.outvars), []
   res = [v for v, inst in zip(eqn.invars, inst_in) if not inst]
-  return eqn, eqn, [False] * len(eqn.outvars), [True] * len(eqn.outvars), []
+  return eqn, eqn, [False] * len(eqn.outvars), [True] * len(eqn.outvars), res
 
 pe.partial_eval_jaxpr_custom_rules[get_p] = partial(_state_partial_eval_custom,
                                                     get_p)

jax/interpreters/partial_eval.py

@@ -1240,9 +1240,45 @@ def call_partial_eval_custom_rule(
   new_inst = [x for x, inst in zip(eqn.invars, inst_in)
               if type(x) is Var and not inst]
   return eqn_known, eqn_staged, unks_out, inst_out, new_inst + residuals
+
+def closed_call_partial_eval_custom_rule(
+    jaxpr_param_name: str,
+    saveable: Callable[..., bool], unks_in: List[bool], inst_in: List[bool],
+    eqn: JaxprEqn, *, res_aval: ResAvalUpdater = _default_res_aval_updater,
+  ) -> Tuple[JaxprEqn, JaxprEqn, Sequence[bool], Sequence[bool], List[Var]]:
+  # TODO(sharadmv,mattjj): dedup this rule with call_partial_eval_custom_rule.
+  closed_jaxpr = eqn.params[jaxpr_param_name]
+  jaxpr = convert_constvars_jaxpr(closed_jaxpr.jaxpr)
+  unks_in = [False] * len(closed_jaxpr.consts) + list(unks_in)
+  inst_in = [False] * len(closed_jaxpr.consts) + list(inst_in)
+  jaxpr_known, jaxpr_staged, unks_out, inst_out, num_res = \
+      partial_eval_jaxpr_custom(jaxpr, unks_in, inst_in, False, False, saveable)
+  ins_known, _ = partition_list(unks_in, eqn.invars)
+  out_binders_known, _ = partition_list(unks_out, eqn.outvars)
+  _, ins_staged = partition_list(inst_in, eqn.invars)
+  _, out_binders_staged = partition_list(inst_out, eqn.outvars)
+  newvar = core.gensym([jaxpr_known, jaxpr_staged])
+  params_known = {**eqn.params, jaxpr_param_name: core.ClosedJaxpr(jaxpr_known,
+    ())}
+  params_staged = {**eqn.params, jaxpr_param_name:
+      core.ClosedJaxpr(jaxpr_staged, ())}
+  residuals = [newvar(res_aval(params_known, var.aval))
+               for var in jaxpr_staged.invars[:num_res]]
+  eqn_known = new_jaxpr_eqn(ins_known, [*out_binders_known, *residuals],
+                            eqn.primitive, params_known, jaxpr_known.effects, eqn.source_info)
+  eqn_staged = new_jaxpr_eqn([*residuals, *ins_staged], out_binders_staged,
+                             eqn.primitive, params_staged,
+                             jaxpr_staged.effects, eqn.source_info)
+  assert len(eqn_staged.invars) == len(jaxpr_staged.invars)
+  new_inst = [x for x, inst in zip(eqn.invars, inst_in)
+              if type(x) is Var and not inst]
+  return eqn_known, eqn_staged, unks_out, inst_out, new_inst + residuals
+
 partial_eval_jaxpr_custom_rules[core.call_p] = \
     partial(call_partial_eval_custom_rule, 'call_jaxpr',
             lambda _, __, ___, ____, _____, x, y: (x, y))
+partial_eval_jaxpr_custom_rules[core.closed_call_p] = \
+    partial(closed_call_partial_eval_custom_rule, 'call_jaxpr')
 partial_eval_jaxpr_custom_rules[core.named_call_p] = \
     partial(call_partial_eval_custom_rule, 'call_jaxpr',
             lambda _, __, ___, ____, _____, x, y: (x, y))