support typedvalue

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

python/triton_tvm_ffi/__init__.py

@@ -1,3 +1,11 @@
-from . import utils
-
-__all__ = ["utils"]
+# tvm-ffi-stubgen(begin): export/_ffi_api
+# fmt: off
+# isort: off
+from ._ffi_api import *  # noqa: F403
+from ._ffi_api import __all__ as _ffi_api__all__
+if "__all__" not in globals():
+    __all__ = []
+__all__.extend(_ffi_api__all__)
+# isort: on
+# fmt: on
+# tvm-ffi-stubgen(end)

python/triton_tvm_ffi/_ffi_api.py

@@ -0,0 +1,54 @@
+# tvm-ffi-stubgen(begin): import-section
+# fmt: off
+# isort: off
+from __future__ import annotations
+from tvm_ffi import Object as _ffi_Object, init_ffi_api as _FFI_INIT_FUNC, register_object as _FFI_REG_OBJ
+from tvm_ffi.libinfo import load_lib_module as _FFI_LOAD_LIB
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from tvm_ffi import Object
+    from typing import Any
+# isort: on
+# fmt: on
+# tvm-ffi-stubgen(end)
+
+# tvm-ffi-stubgen(import-object): tvm_ffi.libinfo.load_lib_module;False;_FFI_LOAD_LIB
+LIB = _FFI_LOAD_LIB("triton_tvm_ffi", "triton_tvm_ffi")
+# tvm-ffi-stubgen(begin): global/triton_tvm_ffi
+# fmt: off
+_FFI_INIT_FUNC("triton_tvm_ffi", __name__)
+if TYPE_CHECKING:
+    def get_device_properties(_0: int, /) -> Mapping[str, int]: ...
+    def launch(*args: Any) -> Any: ...
+    def load_binary(_0: str, _1: bytes, _2: int, _3: int, /) -> tuple[int, int, int, int, int]: ...
+    def string_to_type(_0: str, /) -> int | None: ...
+    def type_to_string(_0: int, /) -> str: ...
+# fmt: on
+# tvm-ffi-stubgen(end)
+
+
+# tvm-ffi-stubgen(import-object): tvm_ffi.register_object;False;_FFI_REG_OBJ
+# tvm-ffi-stubgen(import-object): ffi.Object;False;_ffi_Object
+@_FFI_REG_OBJ("triton_tvm_ffi.TypedValue")
+class TypedValue(_ffi_Object):
+    # tvm-ffi-stubgen(begin): object/triton_tvm_ffi.TypedValue
+    # fmt: off
+    if TYPE_CHECKING:
+        @staticmethod
+        def __c_ffi_init__(_0: int, _1: Any, /) -> Object: ...
+    # fmt: on
+    # tvm-ffi-stubgen(end)
+
+
+__all__ = [
+    # tvm-ffi-stubgen(begin): __all__
+    "LIB",
+    "TypedValue",
+    "get_device_properties",
+    "launch",
+    "load_binary",
+    "string_to_type",
+    "type_to_string",
+    # tvm-ffi-stubgen(end)
+]

python/triton_tvm_ffi/driver.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 from ctypes import c_void_p
-from typing import List, Mapping, Tuple, Type
+from typing import Any, List, Mapping, Optional, Tuple, Type
 from triton.backends.nvidia.driver import CudaDriver
-from .utils import get_device_properties, launch, load_binary
+from . import TypedValue, get_device_properties, launch, load_binary, string_to_type
 
 
 class TVMFFIUtils(object):
@@ -53,7 +53,7 @@ class TVMLauncher(object):
     def __init__(self, src: List[bool], metadata, *args, **kwargs) -> TVMLauncher:
         super().__init__(*args, **kwargs)
 
-        self.mask: List[bool] = [annotation != "constexpr" for annotation in src.signature.values()]
+        self.signature = src.signature.values()
         self.num_ctas = getattr(metadata, "num_ctas", 1)
         self.launch = launch
         self.global_scratch_size = metadata.global_scratch_size
@@ -63,12 +63,6 @@ class TVMLauncher(object):
         self.launch_cooperative_grid = metadata.launch_cooperative_grid
         self.launch_pdl = metadata.launch_pdl
 
-        # We assume the global Triton allocator is not enabled: `_allocator` must be a NullAllocator.
-        # This module depends on NullAllocator behavior; ensure no other code replaces the allocator.
-        from triton.runtime._allocation import _allocator, NullAllocator
-
-        assert isinstance(_allocator.get(), NullAllocator)
-
     def __call__(
         self,
         gridX,
@@ -102,8 +96,15 @@ class TVMLauncher(object):
         )
         assert not self.launch_cooperative_grid
         assert not self.launch_pdl
-        assert len(self.mask) == len(args)
-        args = [arg for arg, m in zip(args, self.mask) if m]
+        assert len(self.signature) == len(args)
+
+        def canonicalize(arg: Any, sig: str) -> TypedValue:
+            ty: Optional[int] = string_to_type(sig)
+            assert ty is not None, sig
+            return TypedValue(ty, arg)
+
+        args = [canonicalize(arg, sig) for arg, sig in zip(args, self.signature)]
+
         return launch(
             gridX,
             gridY,

python/triton_tvm_ffi/utils/__init__.py

@@ -1,11 +0,0 @@
-# tvm-ffi-stubgen(begin): export/_ffi_api
-# fmt: off
-# isort: off
-from ._ffi_api import *  # noqa: F403
-from ._ffi_api import __all__ as _ffi_api__all__
-if "__all__" not in globals():
-    __all__ = []
-__all__.extend(_ffi_api__all__)
-# isort: on
-# fmt: on
-# tvm-ffi-stubgen(end)

python/triton_tvm_ffi/utils/_ffi_api.py

@@ -1,33 +0,0 @@
-# tvm-ffi-stubgen(begin): import-section
-# fmt: off
-# isort: off
-from __future__ import annotations
-from tvm_ffi import init_ffi_api as _FFI_INIT_FUNC
-from tvm_ffi.libinfo import load_lib_module as _FFI_LOAD_LIB
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from collections.abc import Mapping
-    from typing import Any
-# isort: on
-# fmt: on
-# tvm-ffi-stubgen(end)
-# tvm-ffi-stubgen(import-object): tvm_ffi.libinfo.load_lib_module;False;_FFI_LOAD_LIB
-LIB = _FFI_LOAD_LIB("triton_tvm_ffi", "utils")
-# tvm-ffi-stubgen(begin): global/triton_tvm_ffi.utils
-# fmt: off
-_FFI_INIT_FUNC("triton_tvm_ffi.utils", __name__)
-if TYPE_CHECKING:
-    def get_device_properties(_0: int, /) -> Mapping[str, int]: ...
-    def launch(*args: Any) -> Any: ...
-    def load_binary(_0: str, _1: bytes, _2: int, _3: int, /) -> tuple[int, int, int, int, int]: ...
-# fmt: on
-# tvm-ffi-stubgen(end)
-
-__all__ = [
-    # tvm-ffi-stubgen(begin): __all__
-    "LIB",
-    "get_device_properties",
-    "launch",
-    "load_binary",
-    # tvm-ffi-stubgen(end)
-]