Fix failing tests on main (#244)

There was an upstream triton change to min_dot_size.

Author: jansel

helion/_compat.py

@@ -32,9 +32,16 @@ def torch_dtype_to_tl(torch_dtype: torch.dtype) -> object:
     return getattr(tl, name_str)
 
 
-@functools.cache
 def min_dot_size(
     device: torch.device, lhs: torch.dtype, rhs: torch.dtype
+) -> tuple[int, int, int]:
+    # call private func we can patch in testing
+    return _min_dot_size(device, lhs, rhs)
+
+
+@functools.cache
+def _min_dot_size(
+    device: torch.device, lhs: torch.dtype, rhs: torch.dtype
 ) -> tuple[int, int, int]:
     if device.type != "cuda":
         # TODO(jansel): support non-cuda properly

helion/_compiler/inductor_lowering.py

@@ -786,12 +786,12 @@ def apply_dot_requirements(
     lshape = lproxy.size()
     rshape = rproxy.size()
     # use last two dimensions for dot (supports 2D and batched 3D tensors)
-    n, k = lshape[-2], lshape[-1]
-    k2, m = rshape[-2], rshape[-1]
+    m, k = lshape[-2], lshape[-1]
+    k2, n = rshape[-2], rshape[-1]
     assert k == k2, f"Mismatched k dimensions for dot: {k} vs {k2}"
     a, b, c = min_dot_size(lproxy.device, lproxy.dtype, rproxy.dtype)
     env = CompileEnvironment.current()
-    for shape, min_size in [(n, a), (k, b), (m, c)]:
+    for shape, min_size in [(m, a), (n, b), (k, c)]:
         block_idx = CompileEnvironment.current().get_block_id(shape)
         if block_idx is not None:
             env.block_sizes[block_idx].update_min_block(min_size, allow_flattened=True)

test/test_autotuner.py

@@ -35,6 +35,7 @@ def setUp(self):
         random.seed(112)
 
     @patch.object(_compat, "_supports_tensor_descriptor", lambda: True)
+    @patch.object(_compat, "_min_dot_size", lambda *args: (16, 16, 16))
     @patch.object(loops, "_supports_warp_specialize", lambda: True)
     def test_config_fragment0(self):
         args = (

test/test_grid.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import unittest
+from unittest.mock import patch
 
 from expecttest import TestCase
 import torch
 
 import helion
+from helion import _compat
 from helion._testing import DEVICE
 from helion._testing import code_and_output
 import helion.language as hl
@@ -30,6 +32,7 @@ def grid_2d_pytorch(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
 
 class TestGrid(TestCase):
+    @patch.object(_compat, "_min_dot_size", lambda *args: (16, 16, 16))
     def test_grid_1d(self):
         @helion.kernel(static_shapes=True)
         def grid_1d(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: