mark some pallas/mosaic primitives as effectful to avoid DCE

Author: dougalm

jax/_src/core.py

@@ -605,6 +605,10 @@ def def_effectful_abstract_eval(self, effectful_abstract_eval):
     self.abstract_eval = effectful_abstract_eval
     return effectful_abstract_eval
 
+  def def_effectful_abstract_eval2(self, abstract_eval):
+    self.abstract_eval = _generic_abstract_eval(abstract_eval)
+    return abstract_eval
+
   def def_bind_with_trace(self, bind_with_trace):
     self.bind_with_trace = bind_with_trace
     return bind_with_trace
@@ -629,6 +633,17 @@ def abstract_eval_(*args, **kwargs):
     return abstract_eval(*args, **kwargs), no_effects
   return abstract_eval_
 
+class GenericEffect(Effect):
+  pass
+generic_effect = GenericEffect()
+generic_effect_set = {generic_effect}
+effects.lowerable_effects.add_type(GenericEffect)
+
+def _generic_abstract_eval(abstract_eval):
+  def abstract_eval_(*args, **kwargs):
+    return abstract_eval(*args, **kwargs), generic_effect_set
+  return abstract_eval_
+
 # -------------------- lifting --------------------
 
 # TODO(mattjj): replace this approach with a primitive-keyed table of rules

jax/_src/pallas/mosaic/primitives.py

@@ -485,7 +485,7 @@ def do_discharge_src_sem(src_sem=src_sem):
 dma_wait_p = jax_core.Primitive('dma_wait')
 dma_wait_p.multiple_results = True
 
-@dma_wait_p.def_abstract_eval
+@dma_wait_p.def_effectful_abstract_eval2
 def _dma_wait_abstract_eval(*args, tree, device_id_type):
   del args, tree, device_id_type
   return []
@@ -658,7 +658,7 @@ def async_remote_copy(
 
 get_barrier_semaphore_p = jax_core.Primitive('get_barrier_semaphore')
 
-@get_barrier_semaphore_p.def_abstract_eval
+@get_barrier_semaphore_p.def_effectful_abstract_eval2
 def _get_barrier_semaphore_abstract_eval():
   return state.AbstractRef(
       jax_core.ShapedArray((), pl_core.BarrierSemaphore()),
@@ -692,7 +692,7 @@ def get_barrier_semaphore():
 delay_p.multiple_results = True
 
 
-@delay_p.def_abstract_eval
+@delay_p.def_effectful_abstract_eval2
 def _delay_abstract_eval(nanos):
   del nanos
   return []
@@ -708,7 +708,7 @@ def delay(nanos):
 prng_seed_p.multiple_results = True
 
 
-@prng_seed_p.def_abstract_eval
+@prng_seed_p.def_effectful_abstract_eval2
 def _prng_seed_abstract_eval(*_):
   return []
 

jax/_src/pallas/primitives.py

@@ -1083,7 +1083,7 @@ def semaphore_read(sem_or_view):
   flat_args, args_tree = tree_util.tree_flatten(args)
   return semaphore_read_p.bind(*flat_args, args_tree=args_tree)
 
-@semaphore_read_p.def_abstract_eval
+@semaphore_read_p.def_effectful_abstract_eval2
 def _semaphore_read_abstract_eval(
     *avals,
     args_tree,
@@ -1155,7 +1155,7 @@ def _semaphore_signal_abstract_eval(
   check_sem_avals(sem_aval, sem_transforms_avals, "signal")
   if value_aval.dtype != jnp.dtype("int32"):
     raise ValueError("Must signal an int32 value.")
-  effs = set()
+  effs = {jax_core.generic_effect}
   if device_id_avals is not None:
     device_id_flat_avals = tree_util.tree_leaves(device_id_avals)
     for aval in device_id_flat_avals:
@@ -1230,7 +1230,7 @@ def semaphore_wait(sem_or_view, dec: int | jax.Array = 1):
   flat_args, args_tree = tree_util.tree_flatten(args)
   semaphore_wait_p.bind(*flat_args, args_tree=args_tree)
 
-@semaphore_wait_p.def_abstract_eval
+@semaphore_wait_p.def_effectful_abstract_eval2
 def _semaphore_wait_abstract_eval(*avals, args_tree):
   sem_aval, sem_transforms_avals, value_aval = tree_util.tree_unflatten(
       args_tree, avals