Add helion.exc.DeviceTensorSubscriptAssignmentNotAllowed (#292)

Author: jansel

docs/api/exceptions.md

@@ -175,6 +175,10 @@ These exceptions occur when Helion language functions are used incorrectly with
 .. autoclass:: CannotReadDeviceVariableOnHost
 
    Raised when attempting to read variables defined inside device loops from host context.
+
+.. autoclass:: DeviceTensorSubscriptAssignmentNotAllowed
+
+   Raised when attempting to assign to subscript of device tensor.
 ```
 
 ## Type and Inference Errors

helion/_compiler/device_ir.py

@@ -758,7 +758,14 @@ def visit_Assign(self, node: ast.Assign) -> None:
             raise exc.NonTensorSubscriptAssign(lhs_type, rhs_type)
         assert isinstance(target.value, ExtendedAST)
         target_origin = target.value._type_info.origin  # pyright: ignore[reportOptionalMemberAccess]
-        assert target_origin.is_host()
+        if not target_origin.is_host():
+            # Get the variable name for the error message
+            var_name = (
+                target.value.id
+                if isinstance(target.value, ast.Name)
+                else str(target.value)
+            )
+            raise exc.DeviceTensorSubscriptAssignmentNotAllowed(var_name)
         val = self.visit(node.value)
         self._assign_subscript(target, val)
 

helion/exc.py

@@ -312,3 +312,7 @@ class CannotModifyHostVariableOnDevice(BaseError):
 
 class CannotReadDeviceVariableOnHost(BaseError):
     message = "Cannot read variable '{0}' defined inside `hl.tile` or `hl.grid` loop from host code."
+
+
+class DeviceTensorSubscriptAssignmentNotAllowed(BaseError):
+    message = "Cannot assign to subscript of device tensor '{0}'."

test/test_errors.py

@@ -169,6 +169,35 @@ def bad_fn(x: torch.Tensor) -> torch.Tensor:
         with self.assertRaises(helion.exc.CannotReadDeviceVariableOnHost):
             code_and_output(bad_fn, (torch.randn(8, device=DEVICE),))
 
+    def test_device_tensor_subscript(self):
+        @helion.kernel()
+        def bad_fn(x: torch.Tensor) -> torch.Tensor:
+            batch = x.size(0)
+            result = torch.empty_like(x)
+            for i in hl.tile(batch):
+                tmp = x[i] * 2
+                tmp[0] = 1  # This should not be allowed
+                result[i] = tmp
+            return result
+
+        with self.assertRaises(helion.exc.DeviceTensorSubscriptAssignmentNotAllowed):
+            code_and_output(bad_fn, (torch.randn(8, device=DEVICE),))
+
+    def test_closure_fn(self):
+        @helion.kernel()
+        def bad_fn(x: torch.Tensor) -> torch.Tensor:
+            def closure_fn():
+                pass
+
+            batch = x.size(0)
+            result = torch.empty_like(x)
+            for i in hl.tile(batch):
+                result[i] = x[i] * 2
+            return result
+
+        with self.assertRaises(helion.exc.StatementNotSupported):
+            code_and_output(bad_fn, (torch.randn(8, device=DEVICE),))
+
 
 if __name__ == "__main__":
     unittest.main()