Fix bug with unit sized dims and block_sizes

stack-info: PR: #932, branch: jansel/stack/191

Author: jansel

helion/_compiler/compile_environment.py

@@ -582,9 +582,11 @@ def from_config(
 @dataclasses.dataclass
 class LoopSpecBlockSizeSource(BlockSizeSource):
     def from_config(self, config: Config, block_size_info: BlockSizeInfo) -> int:
-        index = CompileEnvironment.current().config_spec.block_sizes.block_id_to_index(
-            block_size_info.block_id
-        )
+        env = CompileEnvironment.current()
+        size = block_size_info.size
+        if isinstance(size, (int, torch.SymInt)) and env.known_equal(size, 1):
+            return 1
+        index = env.config_spec.block_sizes.block_id_to_index(block_size_info.block_id)
         return config.block_sizes[index]
 
 

test/test_matmul.py

@@ -272,6 +272,32 @@ def matmul_split_k(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         torch.testing.assert_close(result, expected, atol=1e-1, rtol=1e-2)
         self.assertExpectedJournal(code)
 
+    @skipIfRefEager("config_spec is not supported in ref eager mode")
+    def test_matmul_config_reuse_with_unit_dim(self):
+        torch.manual_seed(0)
+        big_args = (
+            torch.randn([64, 64], device=DEVICE, dtype=torch.float32),
+            torch.randn([64, 64], device=DEVICE, dtype=torch.float32),
+        )
+        big_bound = matmul_with_addmm.bind(big_args)
+        big_spec = big_bound.config_spec
+        self.assertEqual(len(big_spec.block_sizes), 3)
+        big_config = big_spec.default_config()
+
+        small_args = (
+            torch.randn([1, 64], device=DEVICE, dtype=torch.float32),
+            torch.randn([64, 64], device=DEVICE, dtype=torch.float32),
+        )
+        small_bound = matmul_with_addmm.bind(small_args)
+        small_spec = small_bound.config_spec
+        self.assertEqual(len(small_spec.block_sizes), 3)
+
+        # Previously raised when reusing configs tuned on larger shapes.
+        small_bound.set_config(big_config)
+        result = small_bound(*small_args)
+        expected = small_args[0] @ small_args[1]
+        torch.testing.assert_close(result, expected, atol=1e-1, rtol=1e-2)
+
     def test_matmul_packed_rhs(self):
         @helion.kernel(static_shapes=False)
         def matmul_with_packed_b(