Add a primitive to foil DCE

Author: dougalm

jax/_src/checkify.py

@@ -842,7 +842,7 @@ def new_body_f(*c_consts_and_vals):
     c_consts, vals = split_list(c_consts_and_vals, [c_consts_num])
     out = body_f(*vals)
     # This checks if the next cond application will error
-    _ = cond_f(*c_consts, *out)
+    lax.dce_sink(cond_f(*c_consts, *out))
     return out
   new_body_f_ = lu.wrap_init(new_body_f, debug_info=body_jaxpr.jaxpr.debug_info)
   c_consts_avals = cond_jaxpr.in_avals[:c_consts_num]

jax/_src/lax/__init__.py

@@ -62,7 +62,9 @@
   convert_element_type as convert_element_type,
   convert_element_type_p as convert_element_type_p,
   copy_p as copy_p,
+  dce_sink_p as dce_sink_p,
   cos as cos,
+  dce_sink as dce_sink,
   cos_p as cos_p,
   cosh as cosh,
   cosh_p as cosh_p,

jax/_src/lax/lax.py

@@ -36,6 +36,7 @@
 from jax._src import core
 from jax._src import dispatch
 from jax._src import dtypes
+from jax._src import effects
 from jax._src import linear_util as lu
 from jax._src import pjit
 from jax._src import pretty_printer as pp
@@ -8457,6 +8458,26 @@ 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
+# so the computation producing it won't be eliminated.
+def dce_sink(val):
+  tree_util.tree_map(dce_sink_p.bind, val)
+
+class NoDCEEffect(effects.Effect):
+  pass
+no_dce_effect = NoDCEEffect()
+effects.control_flow_allowed_effects.add_type(NoDCEEffect)
+effects.lowerable_effects.add_type(NoDCEEffect)
+
+dce_sink_p = core.Primitive('dce_sink')
+dce_sink_p.def_impl(lambda _: [])
+dce_sink_p.multiple_results = True
+dce_sink_p.def_effectful_abstract_eval(lambda _: ([], {no_dce_effect}))
+mlir.register_lowering(dce_sink_p, lambda ctx, _: [])
+ad.deflinear(dce_sink_p, lambda _: [])
+pe.def_trivial_padding(dce_sink_p)
+batching.defvectorized(dce_sink_p)
+
 def rng_bit_generator(key, shape, dtype=np.uint32,
                       algorithm=RandomAlgorithm.RNG_DEFAULT,
                       *, out_sharding=None):