Add tl._experimental_make_tensor_descriptor support (#322)

Author: oulgen

helion/_compat.py

@@ -22,7 +22,17 @@ def _supports_tensor_descriptor() -> bool:
     major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
     if major < 9:
         return False
-    return hasattr(triton.language, "make_tensor_descriptor")
+    return hasattr(triton.language, "make_tensor_descriptor") or hasattr(
+        triton.language, "_experimental_make_tensor_descriptor"
+    )
+
+
+@functools.cache
+def get_tensor_descriptor_fn_name() -> str:
+    if hasattr(triton.language, "make_tensor_descriptor"):
+        return "tl.make_tensor_descriptor"
+    assert hasattr(triton.language, "_experimental_make_tensor_descriptor")
+    return "tl._experimental_make_tensor_descriptor"
 
 
 @functools.cache

helion/_compiler/device_function.py

@@ -17,6 +17,7 @@
 from torch._inductor.codegen.triton import texpr
 from torch.fx.graph import _Namespace
 
+from .._compat import get_tensor_descriptor_fn_name
 from .ast_extension import ExtendedAST
 from .ast_extension import create
 from .ast_extension import create_arg
@@ -347,8 +348,9 @@ def tensor_descriptor_arg(
             sizes = ", ".join([arg.name for arg in size_args])
             strides = ", ".join([arg.name for arg in stride_args])
 
+            tensor_descriptor_fn_name = get_tensor_descriptor_fn_name()
             descriptor_stmt = statement_from_string(
-                f"{desc_name} = tl.make_tensor_descriptor({tensor_arg.name}, [{sizes}], [{strides}], [{block_size_expr}])"
+                f"{desc_name} = {tensor_descriptor_fn_name}({tensor_arg.name}, [{sizes}], [{strides}], [{block_size_expr}])"
             )
             self.preamble.append(descriptor_stmt)
 

helion/_testing.py

@@ -16,6 +16,7 @@
 from triton.testing import do_bench
 
 from .runtime.config import Config
+from helion._compat import get_tensor_descriptor_fn_name
 
 if TYPE_CHECKING:
     import types
@@ -220,6 +221,12 @@ def normalize_id(test_id: str) -> str:
         assert match, f"Test ID '{test_id}' does not match expected format"
         return match.group(1)
 
+    @staticmethod
+    def normalize_tensor_descriptors(code: str) -> str:
+        return code.replace(
+            get_tensor_descriptor_fn_name(), "tl.make_tensor_descriptor"
+        )
+
     def lookup(self, test_id: str, value: str) -> tuple[str, str]:
         test_id = self.normalize_id(test_id)
         if self._current_id != test_id:
@@ -234,6 +241,7 @@ def lookup(self, test_id: str, value: str) -> tuple[str, str]:
             expected_values.append("")
             expected = ""
 
+        value = self.normalize_tensor_descriptors(value)
         value = value.strip()
         if value != expected and os.environ.get("EXPECTTEST_ACCEPT", "0") not in {
             "0",

test/test_indexing.py

@@ -5,6 +5,7 @@
 import torch
 
 import helion
+from helion._compat import get_tensor_descriptor_fn_name
 from helion._compat import supports_tensor_descriptor
 from helion._testing import DEVICE
 from helion._testing import TestCase
@@ -366,6 +367,10 @@ def test_broadcasting_block_ptr_indexing(self):
         self.assertExpectedJournal(code)
 
     @unittest.skipIf(not supports_tensor_descriptor(), "TensorDescriptor not supported")
+    @unittest.skipIf(
+        get_tensor_descriptor_fn_name() == "tl._experimental_make_tensor_descriptor",
+        "LLVM ERROR: Illegal shared layout",
+    )
     def test_broadcasting_tensor_descriptor_indexing(self):
         x = torch.randn([16, 24, 32], device=DEVICE)
         bias1 = torch.randn([1, 24, 32], device=DEVICE)

test/test_persistent_kernels.py

@@ -5,6 +5,7 @@
 import torch
 
 import helion
+from helion._compat import get_tensor_descriptor_fn_name
 from helion._compat import supports_tensor_descriptor
 from helion._testing import DEVICE
 from helion._testing import TestCase
@@ -999,8 +1000,8 @@ def tensor_descriptor_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         torch.testing.assert_close(result_interleaved, expected)
 
         # Verify tensor descriptor features in code
-        self.assertIn("tl.make_tensor_descriptor", code_blocked)
-        self.assertIn("tl.make_tensor_descriptor", code_interleaved)
+        self.assertIn(get_tensor_descriptor_fn_name(), code_blocked)
+        self.assertIn(get_tensor_descriptor_fn_name(), code_interleaved)
 
         # Verify persistent kernel features
         self.assertIn("for virtual_pid in tl.range", code_blocked)

test/test_tensor_descriptor.py

@@ -5,6 +5,7 @@
 import torch
 
 import helion
+from helion._compat import get_tensor_descriptor_fn_name
 from helion._compat import supports_tensor_descriptor
 from helion._testing import DEVICE
 from helion._testing import TestCase
@@ -41,15 +42,15 @@ def kernel_with_permutation(x: torch.Tensor) -> torch.Tensor:
             kernel_with_permutation,
             (x,),
             indexing="tensor_descriptor",
-            block_sizes=[4, 8],
+            block_sizes=[8, 8],
         )
 
         # Check that the result is correct
         expected = x + 1.0
         torch.testing.assert_close(result, expected)
 
         # Check that the generated code contains permutation calls
-        self.assertIn("tl.make_tensor_descriptor", code)
+        self.assertIn(get_tensor_descriptor_fn_name(), code)
         # The tensor descriptor should be created with permuted dimensions
         # (sizes and strides should be reordered so stride==1 dim is last)
 
@@ -77,15 +78,15 @@ def kernel_no_permutation(x: torch.Tensor) -> torch.Tensor:
             kernel_no_permutation,
             (x,),
             indexing="tensor_descriptor",
-            block_sizes=[4, 8],
+            block_sizes=[8, 8],
         )
 
         # Check that the result is correct
         expected = x * 2.0
         torch.testing.assert_close(result, expected)
 
         # Check that the generated code contains tensor descriptor
-        self.assertIn("tl.make_tensor_descriptor", code)
+        self.assertIn(get_tensor_descriptor_fn_name(), code)
         # Should not contain permute calls since no permutation needed
         self.assertNotIn("tl.permute", code)
 
@@ -121,7 +122,7 @@ def kernel_3d_permutation(x: torch.Tensor) -> torch.Tensor:
         torch.testing.assert_close(result, expected)
 
         # Should contain both tensor descriptor and permute operations
-        self.assertIn("tl.make_tensor_descriptor", code)
+        self.assertIn(get_tensor_descriptor_fn_name(), code)
         self.assertIn("tl.permute", code)
 
     @unittest.skipUnless(
@@ -149,15 +150,15 @@ def kernel_transpose_case(x: torch.Tensor) -> torch.Tensor:
             kernel_transpose_case,
             (x,),
             indexing="tensor_descriptor",
-            block_sizes=[4, 8],
+            block_sizes=[8, 8],
         )
 
         # Check correctness
         expected = x * x
         torch.testing.assert_close(result, expected)
 
         # Should handle the permutation properly
-        self.assertIn("tl.make_tensor_descriptor", code)
+        self.assertIn(get_tensor_descriptor_fn_name(), code)
         self.assertIn("tl.permute", code)
 
     @unittest.skipUnless(
@@ -183,14 +184,14 @@ def kernel_different_blocks(x: torch.Tensor) -> torch.Tensor:
             kernel_different_blocks,
             (x,),
             indexing="tensor_descriptor",
-            block_sizes=[4, 8],
+            block_sizes=[8, 8],
         )
 
         expected = x + 5.0
         torch.testing.assert_close(result, expected)
 
         # Should contain permutation and tensor descriptor
-        self.assertIn("tl.make_tensor_descriptor", code)
+        self.assertIn(get_tensor_descriptor_fn_name(), code)
         self.assertIn("tl.permute", code)
 
         # The block sizes should also be permuted in the tensor descriptor
@@ -223,14 +224,14 @@ def kernel_store_permutation(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             kernel_store_permutation,
             (x, y),
             indexing="tensor_descriptor",
-            block_sizes=[4, 8],
+            block_sizes=[8, 8],
         )
 
         expected = x * 3.0
         torch.testing.assert_close(result, expected)
 
         # Should have permutation for both load and store
-        self.assertIn("tl.make_tensor_descriptor", code)
+        self.assertIn(get_tensor_descriptor_fn_name(), code)
         self.assertIn("tl.permute", code)
 
     @unittest.skipUnless(
@@ -301,7 +302,7 @@ def kernel_small_block(x: torch.Tensor) -> torch.Tensor:
 
         # Should fall back to block_ptr or pointer indexing instead of tensor descriptor
         # If our fix works, this should NOT contain tensor descriptor
-        self.assertNotIn("tl.make_tensor_descriptor", code)
+        self.assertNotIn(get_tensor_descriptor_fn_name(), code)
 
         # But should still work correctly
         expected = x + 1.0