set auto_dce=False by default, enable in direct linearize

small other fixes

Author: mattjj

jax/_src/api.py

@@ -2396,6 +2396,7 @@ def make_jaxpr(
   { lambda ; a:f32[]. let
       b:f32[] = cos a
       c:f32[] = sin a
+      _:f32[] = sin b
       d:f32[] = cos b
       e:f32[] = mul 1.0:f32[] d
       f:f32[] = neg e

jax/_src/core.py

@@ -606,7 +606,7 @@ def def_effectful_abstract_eval(self, effectful_abstract_eval):
     return effectful_abstract_eval
 
   def def_effectful_abstract_eval2(self, abstract_eval):
-    self.abstract_eval = _generic_abstract_eval(abstract_eval)
+    self.abstract_eval = _generic_effectful_abstract_eval(abstract_eval)
     return abstract_eval
 
   def def_bind_with_trace(self, bind_with_trace):
@@ -639,8 +639,9 @@ class GenericEffect(Effect):
 generic_effect_set = {generic_effect}
 effects.lowerable_effects.add_type(GenericEffect)
 effects.control_flow_allowed_effects.add_type(GenericEffect)
+effects.custom_derivatives_allowed_effects.add_type(GenericEffect)
 
-def _generic_abstract_eval(abstract_eval):
+def _generic_effectful_abstract_eval(abstract_eval):
   def abstract_eval_(*args, **kwargs):
     return abstract_eval(*args, **kwargs), generic_effect_set
   return abstract_eval_

jax/_src/interpreters/ad.py

@@ -93,7 +93,7 @@ def linearize_subtrace(_f: Callable, _store: lu.Store, _tag: core.TraceTag,
                        *primals, **params):
   source_info = source_info_util.current()
   with core.take_current_trace() as parent_trace:
-    tangent_trace = pe.DynamicJaxprTrace(debug_info)
+    tangent_trace = pe.DynamicJaxprTrace(debug_info, auto_dce=True)
     tangent_trace.tag = _tag
     linearize_trace = LinearizeTrace(parent_trace, tangent_trace, tag=_tag)
     tracers = [LinearizeTracer(linearize_trace, p,
@@ -167,7 +167,7 @@ def _linearize_jaxpr(
 ) -> tuple[core.ClosedJaxpr, int, Sequence[bool], Sequence[int | None], core.ClosedJaxpr]:
   dbg = jaxpr.jaxpr.debug_info
   primal_trace = pe.DynamicJaxprTrace(dbg)
-  tangent_trace = pe.DynamicJaxprTrace(dbg)
+  tangent_trace = pe.DynamicJaxprTrace(dbg, auto_dce=True)
   lin_trace = LinearizeTrace(primal_trace, tangent_trace)
   tangent_trace.tag = lin_trace.tag
 
@@ -227,7 +227,7 @@ def direct_linearize(traceable: lu.WrappedFun,
                      primals, kwargs, *, has_aux=False, tag=None):
   with core.take_current_trace() as parent_trace:
     source_info = source_info_util.current()
-    tangent_trace = pe.DynamicJaxprTrace(traceable.debug_info)
+    tangent_trace = pe.DynamicJaxprTrace(traceable.debug_info, auto_dce=True)
     tangents = [tangent_trace.new_arg(get_aval(p).to_tangent_aval(), source_info) for p in primals]
     tangents = [Zero.from_primal_value(t) if dtype(t) == float0 else t for t in tangents]
     linearize_trace = LinearizeTrace(parent_trace, tangent_trace, tag=tag)

jax/_src/interpreters/partial_eval.py

@@ -1937,7 +1937,7 @@ class DynamicJaxprTrace(core.Trace):
   __slots__ = ("frame", "tag", "parent_trace")
 
   def __init__(self, debug_info: core.DebugInfo, parent_trace=None, lower=False,
-               auto_dce=True):
+               auto_dce=False):
     super().__init__()
     self.requires_low = lower
     self.frame = JaxprStackFrame(debug_info, auto_dce)
@@ -2316,7 +2316,7 @@ def trace_to_jaxpr_dynamic(
     *,
     keep_inputs: list[bool] | None = None,
     lower: bool = False,
-    auto_dce: bool = True,
+    auto_dce: bool = False,
 ) -> tuple[Jaxpr, list[AbstractValue], list[Any]]:
   keep_inputs = [True] * len(in_avals) if keep_inputs is None else keep_inputs
   parent_trace = core.trace_ctx.trace

jax/_src/lax/lax.py

@@ -8458,7 +8458,7 @@ def _propagate_mem_kind_copy(in_mem_kind):
   return in_mem_kind
 pxla.memory_kind_propagate_rule[copy_p] = _propagate_mem_kind_copy
 
-# the dce_sink_p primitive marks a value as "used" from the perspective of DCE
+# The dce_sink_p primitive marks a value as "used" from the perspective of DCE
 # so the computation producing it won't be eliminated.
 def dce_sink(val):
   tree_util.tree_map(dce_sink_p.bind, val)

jax/_src/pallas/pallas_call.py

@@ -99,6 +99,8 @@ def _pallas_call_abstract_eval(
     # Report effects that will be introduced when running/lowering
     # mosaic_tpu_interpret.mosaic_tpu_interpret.interpret_pallas_call .
     effs = mosaic_tpu_interpret.get_interpret_effects()
+  elif getattr(params.get('compiler_params', None), 'has_side_effects', False):
+    effs = jax_core.generic_effect_set
   else:
     effs = jax_core.no_effects
 

tests/api_test.py

@@ -5122,18 +5122,46 @@ def g():
                           lambda: (c, jnp.sum(d), d))
     self.assertLen(g().consts, 2)
 
+  # TODO(mattjj,dougalm): this test was flakey on CI; figure out how to enable?
+  # @jtu.run_on_devices('cpu')
+  # def test_implicit_dce_linearize(self):
+  #   def foo(x):
+  #     const = np.zeros((300,))
+  #     x * const
+  #     r = weakref.ref(const)
+  #     del const
+  #     assert r() is None, "oops, the constant wasn't DCE'd"
+  #     return x
+  #   with config.use_direct_linearize(True):
+  #     _ = jax.grad(foo)(3.)
 
-  def test_implicit_dce(self):
-    @api.jit
+  @jtu.run_on_devices('cpu')
+  def test_implicit_dce_linearize_jaxpr(self):
     def foo(x):
       const = np.zeros((300,))
+      x * const
       r = weakref.ref(const)
-      jnp.sin(const) + const
       del const
       assert r() is None, "oops, the constant wasn't DCE'd"
-      return x + x
+      return x
 
-    foo(1.0)
+    with config.use_direct_linearize(True):
+      _, f_vjp = jax.vjp(foo, 3.)
+
+    self.assertNotIn('mul', str(f_vjp))
+
+  # TODO(mattjj,dougalm): re-enable when we set auto_dce=True by default
+  # @jtu.run_on_devices('cpu')
+  # def test_implicit_dce(self):
+  #   @api.jit
+  #   def foo(x):
+  #     const = np.zeros((300,))
+  #     r = weakref.ref(const)
+  #     jnp.sin(const) + const
+  #     del const
+  #     assert r() is None, "oops, the constant wasn't DCE'd"
+  #     return x + x
+  #   foo(1.0)
 
 class RematTest(jtu.JaxTestCase):
 
@@ -6208,7 +6236,7 @@ def test_vjp_caching(self):
     with jtu.count_pjit_cpp_cache_miss() as count:  # noqa: F841
       for _ in range(20):
         f_vjp(1.)[0].block_until_ready()
-    self.assertEqual(count(), 1)  # backward_pass on bwd
+    self.assertLessEqual(count(), 2)
 
   def test_vjp_caching_static_argnums(self):
     identity = jax.remat(lambda x, y: jax.jit(lambda x: 2 * x if y else x)(x),
@@ -6217,7 +6245,7 @@ def test_vjp_caching_static_argnums(self):
     with jtu.count_jit_and_pmap_lowerings() as count:  # noqa: F841
       for _ in range(20):
         f_vjp(1.)[0].block_until_ready()
-    self.assertEqual(count(), 1)  # backward_pass on bwd
+    self.assertLessEqual(count(), 2)
 
   def test_fwd_caching(self):
     # see above test also
@@ -7044,7 +7072,7 @@ def body(c, _):
       return out
     jaxpr = api.make_jaxpr(f)([1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]).jaxpr
     self.assertLen(jaxpr.eqns, 1)
-    self.assertLen(jaxpr.eqns[0].params['jaxpr'].jaxpr.eqns, 4)
+    self.assertLen(jaxpr.eqns[0].params['jaxpr'].jaxpr.eqns, 5)
 
     # If we use the value at index 8 only, all the hidden sequence must be kept
     # and no eqns can be pruned.

tests/checkify_test.py

@@ -729,7 +729,7 @@ def sin_bwd(x2, g):
     err, y = checkify.checkify(jax.grad(sin),
                                errors=checkify.float_checks)(jnp.inf)
     self.assertIsNotNone(err.get())
-    self.assertStartsWith(err.get(), "nan generated by primitive: cos")
+    self.assertStartsWith(err.get(), "nan generated by primitive: sin")
 
   def test_scan_consts(self):
     def f(xs):

tests/pallas/mosaic_gpu_test.py

@@ -1320,6 +1320,7 @@ def cond(acc):
 
       def body(acc):
         del acc  # Unused.
+        o_ref[...] = o_ref[...]  # side-effect to prevent DCE
 
         # We deliberately do a cast here to trigger a layout mismatch.
         return plgpu.layout_cast(

tests/pjit_test.py

@@ -7854,7 +7854,7 @@ def f(x, y, a, b):
         f_bar(x, y, a, b)  # doesn't crash
 
         grad_jaxpr = f_bar.trace(x, y, a, b).jaxpr
-        reshard_eqn = grad_jaxpr.eqns[2].params['jaxpr'].eqns[0]
+        reshard_eqn = grad_jaxpr.eqns[4].params['jaxpr'].eqns[0]
         self.assertEqual(reshard_eqn.params['dst_sharding'].spec.reduced,
                         frozenset('y'))
         self.assertEqual(reshard_eqn.params['dst_sharding'].spec.unreduced,