Raise better error when hl.atomic_* is used on device tensor (#658)

Author: yf225

helion/exc.py

@@ -370,6 +370,13 @@ class CannotModifyHostVariableOnDevice(BaseError):
     message = "Cannot modify host variable '{0}' inside `hl.tile` or `hl.grid` loop without subscript assignment. Use '{0}[tile] = ...' instead."
 
 
+class AtomicOnDeviceTensor(BaseError):
+    message = (
+        "hl.{0}() target must be host-allocated tensor (i.e. allocated outside of hl.tile or hl.grid loop). "
+        "Tensors created inside device loops do not have an addressable pointer for atomics."
+    )
+
+
 class CannotReadDeviceVariableOnHost(BaseError):
     message = "Cannot read variable '{0}' defined inside `hl.tile` or `hl.grid` loop from host code."
 

helion/language/atomic_ops.py

@@ -10,6 +10,7 @@
 
 from .. import exc
 from .._compiler.ast_extension import expr_from_string
+from .._compiler.host_function import HostFunction
 from .._compiler.indexing_strategy import SubscriptIndexing
 from . import _decorators
 
@@ -64,6 +65,10 @@ def _codegen_common(
     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(tl_func)
+
     indices = SubscriptIndexing.create(state, target, index)
     name = state.device_function.tensor_arg(target).name
 

test/test_atomic_ops.py

@@ -276,6 +276,22 @@ def test_atomic_add_w_tile_attr(self):
         torch.testing.assert_close(result, expected)
         self.assertExpectedJournal(code)
 
+    @skipIfRefEager("Error only raises in normal mode")
+    def test_atomic_add_device_tensor_error(self):
+        @helion.kernel(static_shapes=True)
+        def kernel(x: torch.Tensor) -> torch.Tensor:
+            for tile in hl.tile(x.size(0), block_size=128):
+                device_tensor = hl.zeros([tile], dtype=x.dtype)
+                hl.atomic_add(device_tensor, [tile], x[tile])
+            return x
+
+        x = torch.ones(256, device=DEVICE, dtype=torch.float32)
+        with self.assertRaisesRegex(
+            helion.exc.AtomicOnDeviceTensor,
+            r"hl\.atomic_add\(\)",
+        ):
+            kernel(x)
+
     # New tests for other atomics (correctness only; no journal asserts)
     def test_atomic_and(self):
         x0 = torch.full((8,), 0b1111, device=DEVICE, dtype=torch.int32)