put typedvalues initialization into cpp

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

python/triton_tvm_ffi/_ffi_api.py

@@ -6,6 +6,7 @@ from tvm_ffi import Object as _ffi_Object, init_ffi_api as _FFI_INIT_FUNC, regis
 from tvm_ffi.libinfo import load_lib_module as _FFI_LOAD_LIB
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
+    from collections.abc import Sequence
     from tvm_ffi import Object
     from typing import Any
 # isort: on
@@ -33,6 +34,10 @@ class TypedValue(_ffi_Object):
     if TYPE_CHECKING:
         @staticmethod
         def __c_ffi_init__(_0: int, _1: Any, /) -> Object: ...
+        @staticmethod
+        def make_typed_value(_0: str, _1: Any, /) -> TypedValue | None: ...
+        @staticmethod
+        def make_typed_values(_0: Sequence[str], _1: Sequence[Any], /) -> Sequence[TypedValue]: ...
     # fmt: on
     # tvm-ffi-stubgen(end)
 

python/triton_tvm_ffi/driver.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, List, Optional, Type
+from typing import Any, List, Optional, Sequence, Type
 from triton.backends.nvidia.driver import CudaDriver
 from triton.runtime import _allocation
 from . import TypedValue, utils, string_to_type
@@ -10,7 +10,7 @@ class TVMLauncher(object):
     def __init__(self, src, metadata, *args, **kwargs) -> TVMLauncher:
         super().__init__(*args, **kwargs)
 
-        self.signature: List[str] = src.signature.values()
+        self.signature: List[str] = [*src.signature.values()]
         self.num_ctas: int = getattr(metadata, "num_ctas", 1)
         self.launch = utils.launch
         self.global_scratch_size: int = metadata.global_scratch_size
@@ -51,14 +51,8 @@ class TVMLauncher(object):
         )
         assert not self.launch_cooperative_grid
         assert not self.launch_pdl
-        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)]
+        args: Sequence[TypedValue] = TypedValue.make_typed_values(self.signature, args)
 
         return self.launch(
             gridX,
@@ -74,7 +68,7 @@ class TVMLauncher(object):
             self.launch_pdl,
             global_scratch,
             profile_scratch,
-            *args,
+            args,
         )
 
 

python/triton_tvm_ffi/utils.py

@@ -5,7 +5,8 @@ from __future__ import annotations
 from tvm_ffi import init_ffi_api as _FFI_INIT_FUNC
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from collections.abc import Mapping
+    from collections.abc import Mapping, Sequence
+    from tvm_ffi import Object
     from typing import Any
 # isort: on
 # fmt: on
@@ -19,7 +20,7 @@ if TYPE_CHECKING:
     def cuOccupancyMaxActiveClusters(*args: Any) -> Any: ...
     def fill_tma_descriptor(*args: Any) -> Any: ...
     def get_device_properties(_0: int, /) -> Mapping[str, int]: ...
-    def launch(*args: Any) -> Any: ...
+    def launch(_0: int, _1: int, _2: int, _3: int, _4: int, _5: tuple[int, int, int], _6: Object, _7: Object, _8: Object, _9: bool, _10: bool, _11: Object, _12: Object, _13: Sequence[Any], /) -> None: ...
     def load_binary(_0: str, _1: bytes, _2: int, _3: int, /) -> tuple[int, int, int, int, int]: ...
     def set_printf_fifo_size(*args: Any) -> Any: ...
 # fmt: on