Repository: google/jax
Issue: remat_call results in unnecessary CSE foiling when used without partial_eval.

With remat, we want the forward pass to be lowered unmodified, and for the backwards to include a forwards + backwards. However, XLA does common subexpression elimination (CSE), which threatens this - so the `remat_call` primitive lowers with a CSE foiling gadget.

However, foiling CSE has a cost, e.g. missed fusions and other optimizations, so we don't want to do it unless necessary. So, at the moment, the partial-eval-of-remat machinery removes the `remat_call` from the forward pass jaxpr, but keeps it in the backwards pass.

Lost in this logic is that sometimes we just... don't partial eval the `remat_call`, so it sticks around into the final jaxpr and is lowered with CSE-foiling even though it's not needed.

```python
import functools
import jax
import jax.numpy as jnp

@jax.remat
def f(x):
  return jnp.sin(x)

print('Without partial_eval.')
print(jax.make_jaxpr(f)(jnp.zeros(4)))
print('\\nWith partial_eval.')
print(jax.make_jaxpr(functools.partial(jax.vjp, f))(jnp.zeros(4)))
```

```
Without partial_eval.
{ lambda  ; a.
  let b = remat_call[ call_jaxpr={ lambda  ; a.
                                   let b = sin a
                                   in (b,) }
                      concrete=False
                      name=f ] a
  in (b,) }
With partial_eval.
{ lambda  ; a.
  let b = sin a
  in (b, a) }
```

The conceptual error here is that one primitive is taking care of two concepts at once:
- Indicating that something is meant to be rematerialized.
- Indicating that we should foil XLA CSE around a certain block. 

One solution: separate this into two calls, one is `remat_call` and the other is `atomic_call`. `remat_call` is inert when lowering; `atomic_call` isn't. partial-evaling `remat_call` produces a n`atomic_call` in the backwards pass.

Changed files:
jax/interpreters/xla.py, jax/interpreters/partial_eval.py, tests/api_test.py, jax/_src/api.py

diff --git a/jax/_src/api.py b/jax/_src/api.py
index a21ea6a81..d19145ffa 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2408,7 +2408,8 @@ def eval_shape(fun: Callable, *args, **kwargs):
   return tree_unflatten(out_tree(), out)
 
 
-def checkpoint(fun: Callable, concrete: bool = False) -> Callable:
+def checkpoint(fun: Callable, concrete: bool = False, prevent_cse: bool = True,
+               ) -> Callable:
   """Make ``fun`` recompute internal linearization points when differentiated.
 
   The :func:`jax.checkpoint` decorator, aliased to ``jax.remat``, provides a
@@ -2446,6 +2447,14 @@ def checkpoint(fun: Callable, concrete: bool = False) -> Callable:
       control flow is optional, and disabled by default, because in some
       edge-case compositions with :func:`jax.jit` it can lead to some extra
       computation.
+    prevent_cse: Optional, boolean indicating whether to prevent common
+      subexpression elimination (CSE) optimizations in the HLO generated from
+      differentiation. This CSE prevention has costs because it can foil other
+      optimizations, and because it can incur high overheads on some backends,
+      especially GPU. The default is True because otherwise, under a ``jit`` or
+      ``pmap``, CSE can defeat the purpose of this decorator. But in some
+      settings, like when used inside a ``scan``, this CSE prevention mechanism
+      is unnecessary, in which case ``prevent_cse`` can be set to False.
 
   Returns:
     A function (callable) with the same input/output behavior as ``fun`` but
@@ -2496,7 +2505,8 @@ def checkpoint(fun: Callable, concrete: bool = False) -> Callable:
     args_flat, in_tree = tree_flatten((args, kwargs))
     flat_fun, out_tree = flatten_fun(lu.wrap_init(fun), in_tree)
     out_flat = pe.remat_call(flat_fun, *args_flat, name=flat_fun.__name__,
-                             concrete=concrete)
+                             concrete=concrete, prevent_cse=prevent_cse,
+                             differentiated=False)
     return tree_unflatten(out_tree(), out_flat)
   return fun_remat
 remat = checkpoint
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index 2a0455709..df3ef6796 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -814,12 +814,12 @@ def _remat_partial_eval(trace, _, f, tracers, params):
 
   # dce jaxpr outputs
   new_jaxpr = _dce_jaxpr(closed_jaxpr, out_unknowns, drop_outputs=True).jaxpr
-  new_params = dict(params, call_jaxpr=new_jaxpr)
+  new_params = dict(params, call_jaxpr=new_jaxpr, differentiated=True)
 
   # set up eqn for unknown outputs
   in_tracers = (*const_tracers, *env_tracers, *instantiated_tracers)
-  eqn = new_eqn_recipe(in_tracers, unknown_output_tracers, remat_call_p, new_params,
-                       source_info_util.current())
+  eqn = new_eqn_recipe(in_tracers, unknown_output_tracers, remat_call_p,
+                       new_params, source_info_util.current())
   for t in unknown_output_tracers: t.recipe = eqn
   return _zip_knowns(known_output_tracers, unknown_output_tracers, out_unknowns)
 call_partial_eval_rules[remat_call_p] = _remat_partial_eval
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 3d9b23057..b476c3043 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -1464,14 +1464,18 @@ def _remat_using_while(
 
 def _remat_translation_rule(c, axis_env, in_nodes,
                             name_stack, backend, name, call_jaxpr,
-                            device=None, concrete=None):
+                            prevent_cse, differentiated, concrete, device=None):
   del device, concrete  # Unused.
-  if backend == "gpu":
-    return _remat_using_while(
-        c, axis_env, in_nodes, name_stack, backend, name, call_jaxpr)
+  if differentiated and prevent_cse:
+    if backend == "gpu":
+      return _remat_using_while(
+          c, axis_env, in_nodes, name_stack, backend, name, call_jaxpr)
+    else:
+      return _remat_using_cond(
+          c, axis_env, in_nodes, name_stack, backend, name, call_jaxpr)
   else:
-    return _remat_using_cond(
-        c, axis_env, in_nodes, name_stack, backend, name, call_jaxpr)
+    outs = jaxpr_subcomp(c, call_jaxpr, backend, axis_env, (), "", *in_nodes)
+    return xops.Tuple(c, outs)
 
 call_translations[pe.remat_call_p] = _remat_translation_rule  # type: ignore
 
diff --git a/tests/api_test.py b/tests/api_test.py
index 989814686..7604b27b1 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -3214,6 +3214,40 @@ class RematTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(core.UnexpectedTracerError, "global state"):
       api.jit(f)()
 
+  def test_no_cse_widget_on_primals(self):
+    @api.remat
+    def g(x):
+      return lax.sin(lax.sin(x)), 3.
+
+    def f(x):
+      x, _ = g(x)
+      return x
+
+    c = api.xla_computation(f)(2.)
+    self.assertNotIn('while', c.as_hlo_text())
+    self.assertNotIn('conditional', c.as_hlo_text())
+
+    c = api.xla_computation(grad(f))(2.)
+    text = c.as_hlo_text()
+    self.assertTrue('while' in text or 'conditional' in text)
+
+  def test_no_cse_widget_with_prevent_cse_false(self):
+    @partial(api.remat, prevent_cse=False)
+    def g(x):
+      return lax.sin(lax.sin(x)), 3.
+
+    def f(x):
+      x, _ = g(x)
+      return x
+
+    c = api.xla_computation(f)(2.)
+    self.assertNotIn('while', c.as_hlo_text())
+    self.assertNotIn('conditional', c.as_hlo_text())
+
+    c = api.xla_computation(grad(f))(2.)
+    self.assertNotIn('while', c.as_hlo_text())
+    self.assertNotIn('conditional', c.as_hlo_text())
+
 
 class JaxprTest(jtu.JaxTestCase):