Add InvalidDeviceForLoop exception type

stack-info: PR: #205, branch: jansel/stack/57

Author: jansel

helion/_compiler/device_ir.py

@@ -507,7 +507,8 @@ def visit_For(self, node: ast.For) -> None:
         assert not node.orelse
         assert isinstance(node.iter, ExtendedAST)
         iter_type = node.iter._type_info
-        assert isinstance(iter_type, IterType)
+        if not isinstance(iter_type, IterType):
+            raise exc.InvalidDeviceForLoop(iter_type)
         inner_type: TypeInfo = iter_type.inner
         if node._loop_type == LoopType.GRID:
             self._assign(node.target, inner_type.proxy())

helion/exc.py

@@ -222,6 +222,10 @@ class UndefinedVariable(BaseError):
     message = "{} is not defined."
 
 
+class InvalidDeviceForLoop(BaseError):
+    message = "For loops on device must use `hl.tile` or `hl.grid`, got {0!s}."
+
+
 class StarredArgsNotSupportedOnDevice(BaseError):
     message = "*/** args are not supported inside the `hl.tile` or `hl.grid` loop."
 

test/test_errors.py

@@ -109,3 +109,18 @@ def fn(x: torch.Tensor) -> torch.Tensor:
             r"Expected ndim=1, but got ndim=2.*You have too many indices",
         ):
             code_and_output(fn, (torch.randn(8, device=DEVICE),))
+
+    def test_invalid_device_for_loop(self):
+        """Test that InvalidDeviceForLoop is raised for invalid for loops on device."""
+
+        @helion.kernel()
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            batch = x.size(0)
+            out = x.new_empty(batch)
+            for tile_batch in hl.tile(batch):
+                for i in range(10):
+                    out[tile_batch] = x[tile_batch] + i
+            return out
+
+        with self.assertRaises(helion.exc.InvalidDeviceForLoop):
+            code_and_output(fn, (torch.randn(8, device=DEVICE),))