Generalize codegen_atomic_add to codegen_atomic for all atomic ops

Refactor the atomic codegen path so all atomic operations (add, and, or,
xor, max, min, xchg) route through a generic codegen_atomic(op, ...)
method on IndexingStrategy. This enables tensor_descriptor-based TMA
atomics for all supported reduction ops (add, and, max, min, or, xor),
with automatic fallback to pointer for unsupported ops (xchg, cas),
return-value-consuming calls, and non-relaxed memory semantics.

Author: Ethan Che

helion/_compiler/indexing_strategy.py

@@ -4,6 +4,7 @@
 import collections
 import dataclasses
 from typing import TYPE_CHECKING
+from typing import ClassVar
 from typing import NamedTuple
 
 import sympy
@@ -166,8 +167,9 @@ def codegen_store(
     ) -> ast.AST:
         raise NotImplementedError
 
-    def codegen_atomic_add(
+    def codegen_atomic(
         self,
+        op: str,
         state: CodegenState,
         fake_tensor: torch.Tensor,
         subscript: list[object],
@@ -315,8 +317,9 @@ def codegen_store(
             mask=indexing.mask_expr,
         )
 
-    def codegen_atomic_add(
+    def codegen_atomic(
         self,
+        op: str,
         state: CodegenState,
         fake_tensor: torch.Tensor,
         subscript: list[object],
@@ -326,7 +329,7 @@ def codegen_atomic_add(
         indexing = SubscriptIndexing.create(state, fake_tensor, subscript)
         name = state.device_function.tensor_arg(fake_tensor).name
         return expr_from_string(
-            f"tl.atomic_add({name} + {{offset}}, {{value}}, mask={{mask}}, sem={{sem}})",
+            f"tl.{op}({name} + {{offset}}, {{value}}, mask={{mask}}, sem={{sem}})",
             offset=indexing.index_expr,
             value=value,
             mask=indexing.mask_expr,
@@ -394,7 +397,6 @@ def codegen_store(
         )
 
 
-
 class TensorDescriptorIndexingStrategy(IndexingStrategy):
     """Use TensorDescriptor to load/store from tensors"""
 
@@ -573,23 +575,37 @@ def codegen_store(
             value=store_value,
         )
 
-    def codegen_atomic_add(
+    # Ops supported by TMA cp.reduce.async.bulk.tensor via Triton descriptor API
+    _TMA_ATOMIC_OPS: ClassVar[set[str]] = {
+        "atomic_add",
+        "atomic_and",
+        "atomic_max",
+        "atomic_min",
+        "atomic_or",
+        "atomic_xor",
+    }
+
+    def codegen_atomic(
         self,
+        op: str,
         state: CodegenState,
         fake_tensor: torch.Tensor,
         subscript: list[object],
         value: ast.AST,
         sem: ast.AST,
     ) -> ast.AST:
-        fallback = PointerIndexingStrategy().codegen_atomic_add
-        # Descriptor atomic_add returns void; fall back if the return value is used
+        fallback = PointerIndexingStrategy().codegen_atomic
+        # Only certain ops are supported by TMA reduce
+        if op not in self._TMA_ATOMIC_OPS:
+            return fallback(op, state, fake_tensor, subscript, value, sem)
+        # Descriptor atomics return void; fall back if the return value is used
         if state.fx_node is not None and len(state.fx_node.users) > 0:
-            return fallback(state, fake_tensor, subscript, value, sem)
-        # Descriptor atomic_add has no sem parameter; fall back for non-relaxed
+            return fallback(op, state, fake_tensor, subscript, value, sem)
+        # Descriptor atomics have no sem parameter; fall back for non-relaxed
         if isinstance(sem, ast.Constant) and sem.value != "relaxed":
-            return fallback(state, fake_tensor, subscript, value, sem)
+            return fallback(op, state, fake_tensor, subscript, value, sem)
         if not self.is_supported(state, fake_tensor, subscript):
-            return fallback(state, fake_tensor, subscript, value, sem)
+            return fallback(op, state, fake_tensor, subscript, value, sem)
         indexing = BlockedSubscriptIndexing.create(state, fake_tensor, subscript)
         desc_arg = indexing.tensor_descriptor_arg(state)
         atomic_value = indexing.reshape_store(state, value)
@@ -601,7 +617,7 @@ def codegen_atomic_add(
             )
 
         return expr_from_string(
-            f"{indexing.tensor_descriptor(state)}.atomic_add({indexing.offsets_str_permuted(state)}, {{value}})",
+            f"{indexing.tensor_descriptor(state)}.{op}({indexing.offsets_str_permuted(state)}, {{value}})",
             value=atomic_value,
         )
 

helion/language/atomic_ops.py

@@ -59,8 +59,9 @@ def _prepare_mem_args(
 
 
 def _codegen_common(
-    tl_func: str, state: CodegenState, value_exprs: list[ast.AST]
+    op: str, state: CodegenState, value_exprs: list[ast.AST]
 ) -> ast.AST:
+    """Route any single-value atomic op through the atomic_indexing strategy."""
     target = state.proxy_arg(0)
     index = state.proxy_arg(1)
     sem = expr_from_string(repr(state.proxy_arg(len(state.ast_args) - 1)))
@@ -70,23 +71,13 @@ def _codegen_common(
 
     host_function = HostFunction.current()
     if target not in host_function.tensor_to_origin:
-        raise exc.AtomicOnDeviceTensor(tl_func)
+        raise exc.AtomicOnDeviceTensor(op)
 
-    indices = SubscriptIndexing.create(state, target, index)
-    name = state.device_function.tensor_arg(target).name
-
-    placeholder_names = [f"v{i}" for i in range(len(value_exprs))]
-    values_section = (
-        ", " + ", ".join([f"{{{n}}}" for n in placeholder_names]) if value_exprs else ""
-    )
-    placeholders = dict(zip(placeholder_names, value_exprs, strict=False))
-    return expr_from_string(
-        f"tl.{tl_func}({name} + {{offset}}{values_section}, mask={{mask}}, sem={{sem}})",
-        offset=indices.index_expr,
-        mask=indices.mask_expr,
-        sem=sem,
-        **placeholders,
-    )
+    device_fn = state.device_function
+    indexing_idx = device_fn.atomic_op_index
+    device_fn.atomic_op_index += 1
+    strategy = device_fn.get_atomic_indexing_strategy(indexing_idx)
+    return strategy.codegen_atomic(op, state, target, index, value_exprs[0], sem)
 
 
 def _cute_pointer_expr(
@@ -590,23 +581,8 @@ def apply(t: torch.Tensor, idx_tuple: tuple, v: object) -> None:
 
 @_decorators.codegen(atomic_add, "triton")
 def _(state: CodegenState) -> ast.AST:
-    target = state.proxy_arg(0)
-    index = state.proxy_arg(1)
-    value_expr = _to_ast_values([state.ast_args[2]])[0]
-    sem = expr_from_string(repr(state.proxy_arg(len(state.ast_args) - 1)))
-
-    assert isinstance(target, torch.Tensor)
-    assert isinstance(index, list)
-
-    host_function = HostFunction.current()
-    if target not in host_function.tensor_to_origin:
-        raise exc.AtomicOnDeviceTensor("atomic_add")
-
-    device_fn = state.device_function
-    indexing_idx = device_fn.atomic_op_index
-    device_fn.atomic_op_index += 1
-    strategy = device_fn.get_atomic_indexing_strategy(indexing_idx)
-    return strategy.codegen_atomic_add(state, target, index, value_expr, sem)
+    value_expr = state.ast_args[2]
+    return _codegen_common("atomic_add", state, _to_ast_values([value_expr]))
 
 
 @_decorators.codegen(atomic_add, "cute")

test/test_atomic_ops.py

@@ -151,6 +151,58 @@ def atomic_cas_kernel(
     return x
 
 
+# 2D kernels for tensor descriptor atomic tests (TD requires ndim >= 2 + static_shapes)
+
+
+@helion.kernel(static_shapes=True)
+def atomic_add_2d_td_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    for i, j in hl.tile([x.size(0), x.size(1)]):
+        hl.atomic_add(x, [i, j], y[i, j])
+    return x
+
+
+@helion.kernel(static_shapes=True)
+def atomic_and_2d_td_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    for i, j in hl.tile([x.size(0), x.size(1)]):
+        hl.atomic_and(x, [i, j], y[i, j])
+    return x
+
+
+@helion.kernel(static_shapes=True)
+def atomic_or_2d_td_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    for i, j in hl.tile([x.size(0), x.size(1)]):
+        hl.atomic_or(x, [i, j], y[i, j])
+    return x
+
+
+@helion.kernel(static_shapes=True)
+def atomic_xor_2d_td_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    for i, j in hl.tile([x.size(0), x.size(1)]):
+        hl.atomic_xor(x, [i, j], y[i, j])
+    return x
+
+
+@helion.kernel(static_shapes=True)
+def atomic_max_2d_td_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    for i, j in hl.tile([x.size(0), x.size(1)]):
+        hl.atomic_max(x, [i, j], y[i, j])
+    return x
+
+
+@helion.kernel(static_shapes=True)
+def atomic_min_2d_td_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    for i, j in hl.tile([x.size(0), x.size(1)]):
+        hl.atomic_min(x, [i, j], y[i, j])
+    return x
+
+
+@helion.kernel(static_shapes=True)
+def atomic_xchg_2d_td_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    for i, j in hl.tile([x.size(0), x.size(1)]):
+        hl.atomic_xchg(x, [i, j], y[i, j])
+    return x
+
+
 @onlyBackends(["triton", "cute", "pallas"])
 class TestAtomicOperations(RefEagerTestBase, TestCase):
     def test_basic_atomic_add(self):
@@ -425,30 +477,8 @@ def test_atomic_cas(self):
 
     @onlyBackends("triton")
     @skipIfRocm("Tensor descriptor not supported on ROCm")
-    def test_atomic_add_tensor_descriptor(self):
-        """Test that atomic_add with tensor_descriptor indexing generates desc.atomic_add."""
-
-        @helion.kernel(
-            config=helion.Config(
-                block_sizes=[64, 64],
-                indexing="tensor_descriptor",
-                atomic_indexing="tensor_descriptor",
-            ),
-            static_shapes=True,
-        )
-        def atomic_add_td_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-            for i, j in hl.tile([x.size(0), x.size(1)]):
-                hl.atomic_add(x, [i, j], y[i, j])
-            return x
-
-        M, N = 128, 64
-        x = torch.zeros(M, N, device=DEVICE, dtype=torch.float32)
-        y = torch.ones(M, N, device=DEVICE, dtype=torch.float32)
-        code, result = code_and_output(atomic_add_td_kernel, (x, y))
-        expected = torch.ones(M, N, device=DEVICE, dtype=torch.float32)
-        torch.testing.assert_close(result, expected)
-        self.assertIn("desc.atomic_add(", code)
-        self.assertNotIn("tl.atomic_add", code)
+    def test_atomic_td_fallbacks(self):
+        """Test that tensor_descriptor atomics fall back to pointer when needed."""
 
         # Return value consumed: should fall back to pointer
         @helion.kernel(
@@ -466,14 +496,14 @@ def atomic_add_td_prev_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
                 out[i, j] = prev
             return out
 
-        x2 = torch.zeros(M, N, device=DEVICE, dtype=torch.float32)
-        y2 = torch.ones(M, N, device=DEVICE, dtype=torch.float32)
-        code2, result2 = code_and_output(atomic_add_td_prev_kernel, (x2, y2))
-        # prev should be zeros (the old values before adding ones)
-        expected2 = torch.zeros(M, N, device=DEVICE, dtype=torch.float32)
-        torch.testing.assert_close(result2, expected2)
-        self.assertIn("tl.atomic_add", code2)
-        self.assertNotIn("desc.atomic_add(", code2)
+        M, N = 128, 64
+        x = torch.zeros(M, N, device=DEVICE, dtype=torch.float32)
+        y = torch.ones(M, N, device=DEVICE, dtype=torch.float32)
+        code, result = code_and_output(atomic_add_td_prev_kernel, (x, y))
+        expected = torch.zeros(M, N, device=DEVICE, dtype=torch.float32)
+        torch.testing.assert_close(result, expected)
+        self.assertIn("tl.atomic_add", code)
+        self.assertNotIn("desc.atomic_add(", code)
 
         # Non-relaxed sem: should fall back to pointer
         @helion.kernel(
@@ -491,13 +521,13 @@ def atomic_add_td_release_kernel(
                 hl.atomic_add(x, [i, j], y[i, j], sem="release")
             return x
 
-        x3 = torch.zeros(M, N, device=DEVICE, dtype=torch.float32)
-        y3 = torch.ones(M, N, device=DEVICE, dtype=torch.float32)
-        code3, result3 = code_and_output(atomic_add_td_release_kernel, (x3, y3))
-        expected3 = torch.ones(M, N, device=DEVICE, dtype=torch.float32)
-        torch.testing.assert_close(result3, expected3)
-        self.assertIn("tl.atomic_add", code3)
-        self.assertNotIn("desc.atomic_add(", code3)
+        x2 = torch.zeros(M, N, device=DEVICE, dtype=torch.float32)
+        y2 = torch.ones(M, N, device=DEVICE, dtype=torch.float32)
+        code2, result2 = code_and_output(atomic_add_td_release_kernel, (x2, y2))
+        expected2 = torch.ones(M, N, device=DEVICE, dtype=torch.float32)
+        torch.testing.assert_close(result2, expected2)
+        self.assertIn("tl.atomic_add", code2)
+        self.assertNotIn("desc.atomic_add(", code2)
 
     @onlyBackends("triton")
     @skipIfRocm("Tensor descriptor not supported on ROCm")
@@ -536,6 +566,81 @@ def two_atomic_adds(
         self.assertNotIn("out1_desc", code)
         self.assertNotIn("tl.atomic_add(out2", code)
 
+    @onlyBackends("triton")
+    @skipIfRocm("Tensor descriptor not supported on ROCm")
+    def test_atomic_ops_tensor_descriptor(self):
+        """Test all TMA-supported atomic ops generate desc.atomic_{op} codegen."""
+        M, N = 128, 64
+        td_config = {
+            "block_sizes": [64, 64],
+            "indexing": "tensor_descriptor",
+            "atomic_indexing": "tensor_descriptor",
+        }
+        # (op_name, kernel, x, y, expected)
+        cases = [
+            (
+                "add",
+                atomic_add_2d_td_kernel,
+                torch.zeros(M, N, device=DEVICE, dtype=torch.float32),
+                torch.ones(M, N, device=DEVICE, dtype=torch.float32),
+                torch.ones(M, N, device=DEVICE, dtype=torch.float32),
+            ),
+            (
+                "and",
+                atomic_and_2d_td_kernel,
+                torch.full((M, N), 0b1111, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 0b1010, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 0b1010, device=DEVICE, dtype=torch.int32),
+            ),
+            (
+                "or",
+                atomic_or_2d_td_kernel,
+                torch.zeros(M, N, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 0b1010, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 0b1010, device=DEVICE, dtype=torch.int32),
+            ),
+            (
+                "xor",
+                atomic_xor_2d_td_kernel,
+                torch.full((M, N), 0b1010, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 0b1100, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 0b0110, device=DEVICE, dtype=torch.int32),
+            ),
+            (
+                "max",
+                atomic_max_2d_td_kernel,
+                torch.ones(M, N, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 5, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 5, device=DEVICE, dtype=torch.int32),
+            ),
+            (
+                "min",
+                atomic_min_2d_td_kernel,
+                torch.full((M, N), 10, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 3, device=DEVICE, dtype=torch.int32),
+                torch.full((M, N), 3, device=DEVICE, dtype=torch.int32),
+            ),
+        ]
+        for op_name, kernel, x, y, expected in cases:
+            with self.subTest(op=op_name):
+                code, result = code_and_output(kernel, (x, y), **td_config)
+                torch.testing.assert_close(result, expected)
+                self.assertIn(f"desc.atomic_{op_name}(", code)
+                self.assertNotIn(f"tl.atomic_{op_name}", code)
+
+        # xchg is NOT a TMA reduction op — should fall back to pointer
+        with self.subTest(op="xchg_fallback"):
+            x = torch.zeros(M, N, device=DEVICE, dtype=torch.int32)
+            y = torch.ones(M, N, device=DEVICE, dtype=torch.int32)
+            code, result = code_and_output(
+                atomic_xchg_2d_td_kernel, (x, y), **td_config
+            )
+            torch.testing.assert_close(
+                result, torch.ones(M, N, device=DEVICE, dtype=torch.int32)
+            )
+            self.assertIn("tl.atomic_xchg", code)
+            self.assertNotIn("desc.atomic_xchg", code)
+
 
 if __name__ == "__main__":
     unittest.main()