Fix TensorDescriptor handling in _find_device (#35)

Author: yf225

helion/_compat.py

@@ -23,14 +23,20 @@ def _supports_tensor_descriptor() -> bool:
     if major < 9:
         return False
     try:
-        return get_triton_tensor_descriptor_import_path() is not None
+        return get_triton_tensor_descriptor_class() is not None
     except ImportError:
         return False
 
 
 @functools.cache
-def get_triton_tensor_descriptor_import_path() -> str:
-    """Attempt to import TensorDescriptor object from known Triton modules."""
+def get_triton_tensor_descriptor_class_import_path() -> str:
+    cls = get_triton_tensor_descriptor_class()
+    return f"from {cls.__module__} import {cls.__qualname__}"
+
+
+@functools.cache
+def get_triton_tensor_descriptor_class() -> type[object]:
+    """Attempt to import TensorDescriptor class from known Triton modules."""
     possible_modules = [
         "triton.tools.experimental_descriptor",
         "triton.tools.tensor_descriptor",
@@ -39,10 +45,12 @@ def get_triton_tensor_descriptor_import_path() -> str:
         try:
             module = importlib.import_module(module_name)
             if hasattr(module, "TensorDescriptor"):
-                return f"from {module_name} import TensorDescriptor"
+                return module.TensorDescriptor
         except ImportError:
             continue
-    raise ImportError("TensorDescriptor not found in any of the known Triton modules.")
+    raise ImportError(
+        "TensorDescriptor class not found in any of the known Triton modules."
+    )
 
 
 @functools.cache

helion/_compiler/output_header.py

@@ -4,7 +4,7 @@
 
 from .. import exc
 from .ast_read_writes import ReadWrites
-from helion._compat import get_triton_tensor_descriptor_import_path
+from helion._compat import get_triton_tensor_descriptor_class_import_path
 from helion._compat import supports_tensor_descriptor
 
 if TYPE_CHECKING:
@@ -24,7 +24,9 @@
 }
 
 if supports_tensor_descriptor():
-    library_imports["TensorDescriptor"] = get_triton_tensor_descriptor_import_path()
+    library_imports["TensorDescriptor"] = (
+        get_triton_tensor_descriptor_class_import_path()
+    )
 
 disallowed_names: dict[str, None] = dict.fromkeys(
     [

helion/runtime/kernel.py

@@ -13,6 +13,8 @@
 from torch._inductor.codecache import PyCodeCache
 
 from .. import exc
+from .._compat import get_triton_tensor_descriptor_class
+from .._compat import supports_tensor_descriptor
 from .._compiler.ast_extension import unparse
 from .._compiler.compile_environment import CompileEnvironment
 from .._compiler.generate_ast import generate_ast
@@ -503,6 +505,10 @@ def _find_device(args: tuple[object, ...]) -> torch.device:
     for arg in args:
         if isinstance(arg, torch.Tensor):
             return arg.device
+        if supports_tensor_descriptor() and isinstance(
+            arg, get_triton_tensor_descriptor_class()
+        ):
+            return arg.base.device  # pyre-ignore[16]
         if isinstance(arg, (tuple, list)):
             for item in arg:
                 try:

test/test_matmul.py

@@ -8,7 +8,7 @@
 
 import helion
 from helion import Config
-from helion._compat import get_triton_tensor_descriptor_import_path
+from helion._compat import get_triton_tensor_descriptor_class_import_path
 from helion._compat import supports_tensor_descriptor
 from helion._testing import DEVICE
 from helion._testing import code_and_output
@@ -355,7 +355,7 @@ def test_matmul_tensor_descriptor(self):
 import torch
 import triton
 import triton.language as tl
-{get_triton_tensor_descriptor_import_path()}
+{get_triton_tensor_descriptor_class_import_path()}
 
 @triton.jit
 def _matmul_kernel(x_desc, y_desc, out_desc, m, n, k, _BLOCK_SIZE_0: tl.constexpr, _BLOCK_SIZE_1: tl.constexpr, _BLOCK_SIZE_2: tl.constexpr):