Fix block size variable handling and atomic operations with symints (#177)

Author: jansel

helion/language/_tracing_ops.py

@@ -39,9 +39,10 @@ def _(state: CodegenState) -> ast.AST:
     val = state.fx_node.meta["val"]
     assert isinstance(val, (torch.SymInt, torch.SymFloat, torch.SymBool)), val
     if (block_idx := CompileEnvironment.current().get_block_id(val)) is not None:
-        if state.device_function.block_size_var(block_idx) is None:
-            # this should be unused
-            return expr_from_string("block_size_var_optimized_away")
+        block_size_var = state.device_function.block_size_var(block_idx)
+        if block_size_var is None:
+            return expr_from_string("1")
+        return expr_from_string(block_size_var)
     return state.codegen.lift(
         expr_from_string(state.sympy_expr(val._sympy_())),
         dce=True,

helion/language/memory_ops.py

@@ -4,9 +4,11 @@
 from typing import TYPE_CHECKING
 
 import torch
+from torch._inductor.codegen.simd import constant_repr
 from torch.fx import has_side_effect
 
 from .. import exc
+from .._compiler.ast_extension import expr_from_string
 from .._compiler.indexing_strategy import SubscriptIndexing
 from . import _decorators
 
@@ -170,26 +172,19 @@ def _(
 
 @_decorators.codegen(atomic_add)
 def _(state: CodegenState) -> ast.AST:
-    import ast
-
-    from .._compiler.ast_extension import expr_from_string
-
     target = state.proxy_arg(0)
     index = state.proxy_arg(1)
-    value = state.proxy_arg(2)
-    sem = expr_from_string(f"'{state.proxy_arg(3)}'")
+    sem = expr_from_string(repr(state.proxy_arg(3)))
 
     assert isinstance(target, torch.Tensor)
-    assert isinstance(index, (list))
+    assert isinstance(index, list)
 
     indices = SubscriptIndexing.create(state, target, index)
     name = state.device_function.tensor_arg(target).name
 
-    value_expr = (
-        state.ast_args[2]
-        if isinstance(value, torch.Tensor)
-        else ast.Constant(value=value)
-    )
+    value_expr = state.ast_args[2]
+    if isinstance(value_expr, (int, float, bool)):
+        value_expr = expr_from_string(constant_repr(value_expr))
     assert isinstance(value_expr, ast.AST)
     return expr_from_string(
         f"tl.atomic_add({name} + offset, value, mask=mask, sem=sem)",

helion/language/tile_ops.py

@@ -124,17 +124,8 @@ def tile_block_size(tile: Tile) -> int:
 
 @_decorators.register_fake(tile_block_size)
 def _(tile: torch.SymInt) -> torch.SymInt:
-    assert isinstance(tile, torch.SymInt)
-    return CompileEnvironment.current().create_unbacked_symint()
-
-
-@_decorators.codegen(tile_block_size)
-def _(state: CodegenState) -> ast.AST:
-    index = _get_tile_index(state)
-    block_size_var = state.device_function.block_size_var(index)
+    return tile
 
-    if block_size_var is not None:
-        return expr_from_string(block_size_var)
 
-    # Final fallback for grid tiles with block_size=1
-    return expr_from_string("1")
+# since we return tile above, no codegen is needed for this function.
+# codegen is handled in _get_symnode()

test/test_indexing.py

@@ -279,6 +279,21 @@ def test_block_size_access(x: torch.Tensor) -> torch.Tensor:
         expected = torch.full_like(x, 1, dtype=torch.int32)
         torch.testing.assert_close(result, expected)
 
+    def test_atomic_add_symint(self):
+        @helion.kernel(config={"block_size": 32})
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            for tile in hl.tile(x.size(0)):
+                hl.atomic_add(x, [tile], tile.block_size + 1)
+            return x
+
+        x = torch.zeros([200], device=DEVICE)
+        expected = x + 33
+        code, result = code_and_output(
+            fn,
+            (x,),
+        )
+        torch.testing.assert_close(result, expected)
+
 
 if __name__ == "__main__":
     unittest.main()