Allow DCE side-effecting custom calls

PiperOrigin-RevId: 829639173

Author: sharadmv

jax/_src/tpu_custom_call.py

@@ -133,6 +133,15 @@ def to_json(self) -> bytes:
     ).encode("ascii")
 
 
+class TpuSideEffectType(enum.Enum):
+  # No side effects, can be deduplicated / removed if unused.
+  PURE = "pure"
+  # Cannot be deduplicated, but can be removed if unused.
+  DATAFLOW_SIDE_EFFECTING = "dataflow_side_effecting"
+  # Cannot be deduplicated or removed.
+  SIDE_EFFECTING = "side_effecting"
+
+
 @dataclasses.dataclass(frozen=True)
 class CustomCallBackendConfig:
   """Represents an unserialized backend config for custom calls."""
@@ -304,7 +313,7 @@ def _tpu_custom_call_lowering(
     ctx: mlir.LoweringRuleContext,
     *in_nodes,  # pylint: disable=missing-function-docstring
     config: CustomCallBackendConfig,
-    has_side_effects: bool,
+    has_side_effects: TpuSideEffectType,
     kernel_name: str | None,
     out_avals: Any,
     input_output_aliases: tuple[tuple[int, int], ...],
@@ -340,24 +349,27 @@ def _tpu_custom_call_lowering(
   # information.
   if kernel_name is not None:
     extra_attributes = dict(kernel_name=ir.StringAttr.get(kernel_name))
-  has_side_effects = has_side_effects if has_side_effects is not None else False
   call = mlir.custom_call(
       "tpu_custom_call",
       result_types=result_types,
       operands=in_nodes,
       backend_config=config.to_json(),
       api_version=1,
-      has_side_effect=has_side_effects,
+      has_side_effect=has_side_effects != TpuSideEffectType.PURE,
       operand_output_aliases=dict(input_output_aliases),
       operand_layouts=_avals_to_layouts(ctx.avals_in),
       result_layouts=_avals_to_layouts(ctx.avals_out),
       result_shapes=result_shapes,
       extra_attributes=extra_attributes,
   )
+  metadata_dict = {}
   if metadata is not None:
-    call.attributes["mhlo.frontend_attributes"] = ir.DictAttr.get(
-        dict(kernel_metadata=ir.StringAttr.get(json.dumps(metadata)))
-    )
+    metadata_dict["kernel_metadata"] = ir.StringAttr.get(json.dumps(metadata))
+  assert isinstance(has_side_effects, TpuSideEffectType)
+  if has_side_effects == TpuSideEffectType.DATAFLOW_SIDE_EFFECTING:
+    metadata_dict["xla_allow_dce_side_effecting_op"] = ir.StringAttr.get("true")
+  if metadata_dict:
+    call.attributes["mhlo.frontend_attributes"] = ir.DictAttr.get(metadata_dict)
   return call.results
 
 
@@ -643,7 +655,7 @@ def lower_module_to_custom_call(
     input_output_aliases: tuple[tuple[int, int], ...],
     internal_scratch_in_bytes: int | None,
     collective_id: int | None,
-    has_side_effects: bool,
+    has_side_effects: bool | TpuSideEffectType,
     serialization_format: int | None,
     output_memory_spaces: tuple[MemorySpace | None, ...] | None,
     disable_bounds_checks: bool = False,
@@ -653,6 +665,12 @@ def lower_module_to_custom_call(
     allow_collective_id_without_custom_barrier: bool = False,
     shape_invariant_numerics: bool = False,
 ) -> Sequence[ir.Value]:
+  if isinstance(has_side_effects, bool):
+    has_side_effects = (
+        TpuSideEffectType.PURE
+        if not has_side_effects
+        else TpuSideEffectType.DATAFLOW_SIDE_EFFECTING
+    )
   config = _lower_to_custom_call_config(
       module,
       vmem_limit_bytes=vmem_limit_bytes,
@@ -694,7 +712,7 @@ def as_tpu_kernel(
     input_output_aliases: tuple[tuple[int, int], ...] = (),
     internal_scratch_in_bytes: int | None = None,
     collective_id: int | None = None,
-    has_side_effects: bool = False,
+    has_side_effects: TpuSideEffectType = TpuSideEffectType.PURE,
     serialization_format: int | None = 1,
     output_memory_spaces: tuple[MemorySpace | None, ...] | None = None,
     disable_bounds_checks: bool = False,
@@ -738,7 +756,7 @@ def lowered_as_tpu_kernel(
     needs_hlo_passes: bool = False,
     needs_layout_passes: bool = False,
     has_communication: bool = False,
-    has_side_effects: bool = False,
+    has_side_effects: bool | TpuSideEffectType = False,
     has_custom_barrier: bool = False,
     kernel_name: str | None = None,
     vmem_limit_bytes: int | None = None,
@@ -755,6 +773,12 @@ def lowered_as_tpu_kernel(
   lowered_module_asm = lowered_module.operation.get_asm(
       binary=True, enable_debug_info=True
   )
+  if isinstance(has_side_effects, bool):
+    has_side_effects = (
+        TpuSideEffectType.PURE
+        if not has_side_effects
+        else TpuSideEffectType.DATAFLOW_SIDE_EFFECTING
+    )
   config = _lowered_to_custom_call_config(
       lowered_module_asm,
       vmem_limit_bytes=vmem_limit_bytes,
@@ -784,7 +808,7 @@ def lowered_as_tpu_kernel(
 
 def _as_jax_callable(
     config: CustomCallBackendConfig,
-    has_side_effects: bool,
+    has_side_effects: TpuSideEffectType,
     out_type: Any,
     *,
     kernel_name: str | None,