Improve error message for overpacked tiles (#126)

Author: jansel

helion/_compiler/type_propagation.py

@@ -474,6 +474,8 @@ def _device_indexing_size(self, key: TypeInfo) -> list[int | torch.SymInt]:
             elif isinstance(k, TensorType) and k.fake_value.ndim == 1:
                 inputs_consumed += 1
                 output_sizes.append(k.fake_value.size(0))
+            elif k.contains_type(TileIndexType):
+                raise exc.OverpackedTile(k)
             else:
                 raise exc.InvalidIndexingType(k)
         if inputs_consumed != self.fake_value.ndim:

helion/exc.py

@@ -125,6 +125,13 @@ class FailedToUnpackTile(BaseError):
     )
 
 
+class OverpackedTile(BaseError):
+    message = (
+        "Got a tile wrapped inside a container when indexing a tensor: {0!s}\n"
+        "Did you mix up `hl.tile([x])` and `hl.tile(x)`?"
+    )
+
+
 class AssignmentMultipleTargets(NotAllowedOnDevice):
     message = "Assignment with multiple targets (a=b=1) is not allowed inside the `hl.tile` or `hl.grid` loop."
 

test/test_errors.py

@@ -23,3 +23,15 @@ def sum_kernel(x: torch.Tensor) -> torch.Tensor:
 
         with self.assertRaises(helion.exc.FailedToUnpackTile):
             code_and_output(sum_kernel, (torch.randn(2, 3, 4, device=DEVICE),))
+
+    def test_tile_overpacking(self):
+        @helion.kernel()
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            batch = x.size(0)
+            out = x.new_empty(batch)
+            for tile_wrapped_in_tuple in hl.tile([batch]):
+                out[tile_wrapped_in_tuple] = x[tile_wrapped_in_tuple, :].sum(1)
+            return out
+
+        with self.assertRaises(helion.exc.OverpackedTile):
+            code_and_output(fn, (torch.randn(100, 100, device=DEVICE),))