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 []
@@ -641,7 +641,7 @@ def async_remote_copy(src_ref, dst_ref, send_sem, recv_sem, device_id,
 
 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()),
@@ -675,7 +675,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 []
@@ -691,7 +691,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,
@@ -1128,7 +1128,7 @@ def semaphore_signal(
   )
 
 
-@semaphore_signal_p.def_abstract_eval
+@semaphore_signal_p.def_effectful_abstract_eval2
 def _semaphore_signal_abstract_eval(
     *avals,
     args_tree,
@@ -1218,7 +1218,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