[Error Message] Update block config size length mismatch (#139)

Author: drisspg

helion/_compiler/type_propagation.py

@@ -479,7 +479,11 @@ def _device_indexing_size(self, key: TypeInfo) -> list[int | torch.SymInt]:
             else:
                 raise exc.InvalidIndexingType(k)
         if inputs_consumed != self.fake_value.ndim:
-            raise exc.RankMismatch(self.fake_value.ndim, inputs_consumed)
+            raise exc.RankMismatch(
+                self.fake_value.ndim,
+                inputs_consumed,
+                f"tensor shape: {tuple(self.fake_value.shape)}",
+            )
         return output_sizes
 
     def propagate_setitem(
@@ -488,11 +492,16 @@ def propagate_setitem(
         if origin.is_host():
             warning(exc.TensorOperationInWrapper)
         else:
-            lhs_rank = len(self._device_indexing_size(key))
+            lhs_shape = self._device_indexing_size(key)
+            lhs_rank = len(lhs_shape)
             if isinstance(value, TensorType):
                 rhs_rank = value.fake_value.ndim
                 if lhs_rank != rhs_rank:
-                    raise exc.RankMismatch(lhs_rank, rhs_rank)
+                    raise exc.RankMismatch(
+                        lhs_rank,
+                        rhs_rank,
+                        f"LHS shape: {tuple(lhs_shape)}, RHS shape: {tuple(value.fake_value.shape)}",
+                    )
             elif isinstance(value, UnknownType):
                 raise exc.TypePropagationError(value)
             else:

helion/autotuner/block_id_sequence.py

@@ -159,7 +159,9 @@ def _normalize(
                     values.append(spec._fill_missing())
             except NotImplementedError:
                 raise InvalidConfig(
-                    f"Not enough values for config[{name!r}], expected {size}, got {len(values)}"
+                    f"Not enough values for config[{name!r}]: expected {size} block sizes "
+                    f"(one for each tiled dimension), got {len(values)}. "
+                    f"Did you forget to specify block sizes for all your hl.tile() dimensions?"
                 ) from None
         for i, spec in enumerate(self._data):
             values[i] = spec._normalize(f"config[{name}][{i}]", values[i])

helion/exc.py

@@ -82,7 +82,26 @@ class NestedGridLoop(BaseError):
 
 
 class RankMismatch(BaseError):
-    message = "Expected ndim={0}, but got ndim={1}"
+    message = "Expected ndim={expected_ndim}, but got ndim={actual_ndim}{shape_part}. You have {direction}."
+
+    def __init__(
+        self, expected_ndim: int, actual_ndim: int, shape_info: str = ""
+    ) -> None:
+        if actual_ndim > expected_ndim:
+            direction = "too many indices"
+        elif actual_ndim < expected_ndim:
+            direction = "too few indices"
+        else:
+            direction = "indices that don't match expected structure"
+
+        shape_part = f" ({shape_info})" if shape_info else ""
+
+        super().__init__(
+            expected_ndim=expected_ndim,
+            actual_ndim=actual_ndim,
+            shape_part=shape_part,
+            direction=direction,
+        )
 
 
 class InvalidIndexingType(BaseError):

test/test_errors.py

@@ -35,3 +35,77 @@ def fn(x: torch.Tensor) -> torch.Tensor:
 
         with self.assertRaises(helion.exc.OverpackedTile):
             code_and_output(fn, (torch.randn(100, 100, device=DEVICE),))
+
+    def test_invalid_config_insufficient_block_sizes(self):
+        """Test that InvalidConfig shows helpful message for missing block sizes."""
+
+        @helion.kernel(config=helion.Config(block_sizes=[32, 64]))
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            batch, seq_len, hidden = x.size()
+            out = torch.empty_like(x)
+            for tile_batch, tile_seq, tile_hidden in hl.tile([batch, seq_len, hidden]):
+                out[tile_batch, tile_seq, tile_hidden] = x[
+                    tile_batch, tile_seq, tile_hidden
+                ]
+            return out
+
+        with self.assertRaisesRegex(
+            helion.exc.InvalidConfig,
+            r"Not enough values for config.*expected 3 block sizes.*got 2.*"
+            r"Did you forget to specify block sizes for all your hl\.tile\(\) dimensions\?",
+        ):
+            code_and_output(
+                fn,
+                (torch.randn(4, 8, 16, device=DEVICE),),
+            )
+
+    def test_rank_mismatch_assignment(self):
+        """Test that RankMismatch shows tensor shapes in assignment errors."""
+
+        @helion.kernel()
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            batch, seq_len = x.size()
+            out = x.new_empty(batch, seq_len)
+            for tile_batch, tile_seq in hl.tile([batch, seq_len]):
+                scalar_val = x[tile_batch, 0].sum()  # Creates 0D tensor
+                out[tile_batch, tile_seq] = scalar_val  # 0D -> 2D assignment
+            return out
+
+        with self.assertRaisesRegex(
+            helion.exc.RankMismatch,
+            r"Expected ndim=2, but got ndim=0.*You have too few indices",
+        ):
+            code_and_output(fn, (torch.randn(4, 8, device=DEVICE),))
+
+    def test_rank_mismatch_indexing(self):
+        """Test that RankMismatch shows tensor shapes in indexing errors."""
+
+        @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]):
+                scalar_val = x[tile_batch].sum()  # 1d index for 2d tensor
+                out = scalar_val
+            return out
+
+        with self.assertRaisesRegex(
+            helion.exc.RankMismatch,
+            r"Expected ndim=2, but got ndim=1.*You have too few indices",
+        ):
+            code_and_output(fn, (torch.randn(4, 8, device=DEVICE),))
+
+    def test_rank_mismatch_indexing_too_many(self):
+        @helion.kernel()
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            batch = x.size(0)
+            fill = x.new_empty(batch, batch)
+            for tile_batch in hl.tile(batch):
+                fill = x[tile_batch, tile_batch]  # 2d index for 1d tensor
+            return fill
+
+        with self.assertRaisesRegex(
+            helion.exc.RankMismatch,
+            r"Expected ndim=1, but got ndim=2.*You have too many indices",
+        ):
+            code_and_output(fn, (torch.randn(8, device=DEVICE),))