Add helion.exc.CannotModifyHostVariableOnDevice and helion.exc.CannotReadDeviceVariableOnHost (#290)

Author: jansel

docs/api/exceptions.md

@@ -167,6 +167,14 @@ These exceptions occur when Helion language functions are used incorrectly with
 .. autoclass:: UndefinedVariable
 
    Raised when referencing undefined variables.
+
+.. autoclass:: CannotModifyHostVariableOnDevice
+
+   Raised when modifying host variables inside device loops without subscript assignment.
+
+.. autoclass:: CannotReadDeviceVariableOnHost
+
+   Raised when attempting to read variables defined inside device loops from host context.
 ```
 
 ## Type and Inference Errors

helion/_compiler/type_propagation.py

@@ -123,6 +123,12 @@ def get(self, name: str) -> TypeInfo:
             return self.variables[name]
         return self.parent.get(name)
 
+    def maybe_get(self, name: str) -> TypeInfo | None:
+        try:
+            return self.get(name)
+        except exc.UndefinedVariable:
+            return None
+
     def set(self, name: str, type_info: TypeInfo) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         self.variables[name] = type_info
 
@@ -1545,6 +1551,13 @@ def _compare(self, op: ast.cmpop, left: TypeInfo, right: TypeInfo) -> TypeInfo:
 
     def _assign(self, lhs: ast.AST, rhs: TypeInfo) -> None:
         if isinstance(lhs, ast.Name):
+            # Check if we're trying to modify a host variable inside a device loop
+            if (
+                (existing_type := self.scope.maybe_get(lhs.id)) is not None
+                and existing_type.origin.is_host()
+                and rhs.origin.is_device()
+            ):
+                raise exc.CannotModifyHostVariableOnDevice(lhs.id) from None
             return self.scope.set(lhs.id, rhs)
         if isinstance(lhs, ast.Starred):
             try:
@@ -1679,7 +1692,10 @@ def visit_Dict(self, node: ast.Dict) -> TypeInfo:
         return DictType(element_types=element_types, origin=self.origin())
 
     def visit_Name(self, node: ast.Name) -> TypeInfo:
-        return self.scope.get(node.id)
+        result = self.scope.get(node.id)
+        if self.device_loop_depth == 0 and result.origin.is_device():
+            raise exc.CannotReadDeviceVariableOnHost(node.id)
+        return result
 
     visit_Starred: _VisitMethod = generic_visit  # pyright: ignore[reportAssignmentType,reportIncompatibleMethodOverride]
 
@@ -1996,6 +2012,8 @@ def visit_Try(self, node: ast.Try) -> TypeInfo:
     def _not_on_device_statement(self, node: ast.AST) -> TypeInfo:
         if self.device_loop_depth:
             raise exc.NotAllowedOnDevice(type(node).__name__)
+        for child_node in ast.iter_child_nodes(node):
+            self.visit(child_node)
         return NoType(origin=self.origin())
 
     visit_ExceptHandler: _VisitMethod = _not_on_device_statement  # pyright: ignore[reportAssignmentType,reportIncompatibleMethodOverride]

helion/exc.py

@@ -304,3 +304,11 @@ class TypeInferenceError(BaseError):
 
 class NotAllowedInHelperFunction(BaseError):
     message = "This operation is not allowed inside helper functions. It requires kernel context."
+
+
+class CannotModifyHostVariableOnDevice(BaseError):
+    message = "Cannot modify host variable '{0}' inside `hl.tile` or `hl.grid` loop without subscript assignment. Use '{0}[tile] = ...' instead."
+
+
+class CannotReadDeviceVariableOnHost(BaseError):
+    message = "Cannot read variable '{0}' defined inside `hl.tile` or `hl.grid` loop from host code."

helion/language/loops.py

@@ -193,8 +193,8 @@ def tile_info_example(x: torch.Tensor) -> torch.Tensor:
         * tile(begin, end, block_size) iterates begin to end-1, fixed block_size
         * tile(end, block_size=block_size) iterates 0 to end-1, fixed block_size
 
-        Block sizes can be registered for autotuning explicitly with :func:`~helion.language.register_block_size`.
-        And passed in to as ``block_size`` argument if one needs two loops to use the same block size.  Passing
+        Block sizes can be registered for autotuning explicitly with :func:`~helion.language.register_block_size`
+        and passed as the ``block_size`` argument if one needs two loops to use the same block size.  Passing
         ``block_size=None`` is equivalent to calling register_block_size.
 
         Use ``tile`` in most cases. Use ``grid`` when you need explicit control over the launch grid.

test/data/all_ast_nodes.py

@@ -135,8 +135,8 @@ def all_ast_nodes(x, y):
     i = 3
     t = 0
 
-    with contextlib.nullcontext():
-        e3 = 1
+    # with contextlib.nullcontext():
+    #    e3 = 1
 
     global global0
 

test/test_errors.py

@@ -125,6 +125,50 @@ def fn(x: torch.Tensor) -> torch.Tensor:
         with self.assertRaises(helion.exc.InvalidDeviceForLoop):
             code_and_output(fn, (torch.randn(8, device=DEVICE),))
 
+    def test_return_inside_grid_loop(self):
+        """Test that return statement inside hl.grid loop raises proper error."""
+
+        @helion.kernel()
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            batch = x.size(0)
+            out = x.new_empty(batch)
+            for tile_batch in hl.grid(batch):
+                if x[tile_batch] > 0:
+                    return out  # This should not be allowed
+                out[tile_batch] = x[tile_batch] * 2
+            return out
+
+        with self.assertRaises(helion.exc.NotAllowedOnDevice):
+            code_and_output(fn, (torch.randn(8, device=DEVICE),))
+
+    def test_assign_without_subscript1(self):
+        """Test that modifying host variables inside device loops raises proper error."""
+
+        @helion.kernel()
+        def bad_fn(x: torch.Tensor) -> torch.Tensor:
+            batch = x.size(0)
+            result = torch.empty_like(x)
+            for tile_batch in hl.tile(batch):
+                # shouldn't be able to modify host variables on device
+                result = x[tile_batch] * 2
+            return result
+
+        with self.assertRaises(helion.exc.CannotModifyHostVariableOnDevice):
+            code_and_output(bad_fn, (torch.randn(8, device=DEVICE),))
+
+    def test_assign_without_subscript2(self):
+        """Test that reading device variables from host context raises proper error."""
+
+        @helion.kernel()
+        def bad_fn(x: torch.Tensor) -> torch.Tensor:
+            batch = x.size(0)
+            for tile_batch in hl.tile(batch):
+                result = x[tile_batch] * 2
+            return result  # shouldn't be able to read device variable here
+
+        with self.assertRaises(helion.exc.CannotReadDeviceVariableOnHost):
+            code_and_output(bad_fn, (torch.randn(8, device=DEVICE),))
+
 
 if __name__ == "__main__":
     unittest.main()

test/test_reductions.expected

@@ -108,6 +108,7 @@ def reduce_kernel(x: torch.Tensor, fn: Callable[[torch.Tensor], torch.Tensor], o
         # UnaryOp: LiteralType(-1) DeviceOrigin(location=<SourceLocation test_reductions.py:55>)
         # Constant: LiteralType(1) DeviceOrigin(location=<SourceLocation test_reductions.py:55>)
         out[tile_n] = fn(x[tile_n, :], dim=-1)
+    # Name: TensorType([x_size0], torch.float32) SourceOrigin(location=<SourceLocation test_reductions.py:49>)
     return out
 
 def root_graph_0():

test/test_type_propagation.expected

@@ -33,6 +33,7 @@ def add(x, y):
         # Name: TensorType([y_size0, x_size1], torch.int32) GetItemOrigin(value=SourceOrigin(location=<SourceLocation basic_kernels.py:10>), key=1)
         # Name: SequenceType([TileIndexType(0), TileIndexType(1)]) SourceOrigin(location=<SourceLocation basic_kernels.py:12>)
         out[tile] = x[tile] + y[tile]
+    # Name: TensorType([y_size0, x_size1], torch.int32) SourceOrigin(location=<SourceLocation basic_kernels.py:11>)
     return out
 
 def root_graph_0():
@@ -383,8 +384,6 @@ def all_ast_nodes(x, y):
     i = 3
     # Constant: LiteralType(0) SourceOrigin(location=<SourceLocation all_ast_nodes.py:136>)
     t = 0
-    with contextlib.nullcontext():
-        e3 = 1
     global global0
     # Call: TensorType([y_size0, x_size1], torch.int32) SourceOrigin(location=<SourceLocation all_ast_nodes.py:143>)
     # Attribute: CallableType(_VariableFunctionsClass.empty_like) AttributeOrigin(value=GlobalOrigin(name='torch'), key='empty_like')
@@ -437,6 +436,7 @@ def all_ast_nodes(x, y):
     # Name: SymIntType(u16) SourceOrigin(location=<SourceLocation all_ast_nodes.py:149>)
     # Name: TensorType([y_size0, x_size1], torch.int32) SourceOrigin(location=<SourceLocation all_ast_nodes.py:150>)
     combined = [v, z]
+    # Name: TensorType([y_size0, x_size1], torch.int32) SourceOrigin(location=<SourceLocation all_ast_nodes.py:143>)
     return out
 
 def root_graph_0():
@@ -488,6 +488,7 @@ def hl_full_usage(x: torch.Tensor):
         # Name: SequenceType([TileIndexType(0), TileIndexType(1)]) SourceOrigin(location=<SourceLocation basic_kernels.py:39>)
         # Name: TensorType([block_size_0, block_size_1], torch.int32) DeviceOrigin(location=<SourceLocation basic_kernels.py:42>)
         out[tile] = tmp
+    # Name: TensorType([x_size0, x_size1], torch.int32) SourceOrigin(location=<SourceLocation basic_kernels.py:38>)
     return out
 
 def root_graph_0():
@@ -545,6 +546,7 @@ def hl_zeros_usage(x: torch.Tensor):
         # Name: SequenceType([TileIndexType(0), TileIndexType(1)]) SourceOrigin(location=<SourceLocation basic_kernels.py:28>)
         # Name: TensorType([block_size_0, block_size_1], torch.int32) DeviceOrigin(location=<SourceLocation basic_kernels.py:31>)
         out[tile] = tmp
+    # Name: TensorType([x_size0, x_size1], torch.int32) SourceOrigin(location=<SourceLocation basic_kernels.py:27>)
     return out
 
 def root_graph_0():
@@ -640,6 +642,7 @@ def matmul(x: torch.Tensor, y: torch.Tensor):
         # Name: TileIndexType(1) SourceOrigin(location=<SourceLocation matmul.py:19>)
         # Name: TensorType([block_size_0, block_size_1], torch.float32) DeviceOrigin(location=<SourceLocation matmul.py:22>)
         out[tile_m, tile_n] = acc
+    # Name: TensorType([512, 512], torch.float32) SourceOrigin(location=<SourceLocation matmul.py:16>)
     return out
 
 def for_loop_0(arg0_1: "f32[u0, u1]"):
@@ -697,6 +700,7 @@ def fn(x):
         # Name: TensorType([x_size0, x_size1], torch.int32) ArgumentOrigin(name='x')
         # Name: SequenceType([TileIndexType(0), TileIndexType(1)]) SourceOrigin(location=<SourceLocation test_type_propagation.py:77>)
         out[tile] = x[tile].sin()
+    # Name: TensorType([x_size0, x_size1], torch.int32) SourceOrigin(location=<SourceLocation test_type_propagation.py:76>)
     return out
 
 def root_graph_0():
@@ -747,6 +751,7 @@ def pointwise_device_loop(x: torch.Tensor):
             # Name: TileIndexType(1) DeviceOrigin(location=<SourceLocation basic_kernels.py:52>)
             # Constant: LiteralType(1) DeviceOrigin(location=<SourceLocation basic_kernels.py:53>)
             out[tile_n, tile_m] = torch.sigmoid(x[tile_n, tile_m] + 1)
+    # Name: TensorType([x_size0, x_size1], torch.int32) SourceOrigin(location=<SourceLocation basic_kernels.py:49>)
     return out
 
 def for_loop_0():
@@ -805,6 +810,7 @@ def torch_ops_pointwise(x, y):
         # Name: TensorType([y_size0], torch.int32) ArgumentOrigin(name='y')
         # Name: SequenceType([TileIndexType(0)]) SourceOrigin(location=<SourceLocation basic_kernels.py:20>)
         out[tile] = torch.sigmoid(torch.add(torch.sin(x[tile]), torch.cos(y[tile])))
+    # Name: TensorType([x_size0], torch.int32) SourceOrigin(location=<SourceLocation basic_kernels.py:19>)
     return out
 
 def root_graph_0():