implement launch with cpp

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

python/triton_tvm_ffi/_ffi_api.py

@@ -27,17 +27,16 @@ if TYPE_CHECKING:
 
 # 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
+@_FFI_REG_OBJ("triton_tvm_ffi.TVMFFILauncherImpl")
+class TVMFFILauncherImpl(_ffi_Object):
+    """FFI binding for `triton_tvm_ffi.TVMFFILauncherImpl`."""
+
+    # tvm-ffi-stubgen(begin): object/triton_tvm_ffi.TVMFFILauncherImpl
     # fmt: off
     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]: ...
+        def __c_ffi_init__(_0: Sequence[int], _1: bool, _2: bool, /) -> Object: ...
+        def launch(self, _1: int, _2: int, _3: int, _4: int, _5: int, _6: tuple[int, int, int], _7: Object, _8: Object, _9: Object, _10: Object, _11: Object, _12: Sequence[Any], /) -> None: ...
     # fmt: on
     # tvm-ffi-stubgen(end)
 
@@ -45,7 +44,7 @@ class TypedValue(_ffi_Object):
 __all__ = [
     # tvm-ffi-stubgen(begin): __all__
     "LIB",
-    "TypedValue",
+    "TVMFFILauncherImpl",
     "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 typing import Any, List, Optional, Sequence, Type
+from typing import Any, Callable, Final, List, Sequence, Type, Union
 from triton.backends.nvidia.driver import CudaDriver
 from triton.runtime import _allocation
-from . import TypedValue, utils, string_to_type
+from . import TVMFFILauncherImpl, utils, string_to_type
 
 
 class TVMLauncher(object):
@@ -11,14 +11,34 @@ class TVMLauncher(object):
         super().__init__(*args, **kwargs)
 
         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
-        self.global_scratch_align: int = metadata.global_scratch_align
-        self.profile_scratch_size: int = metadata.profile_scratch_size
-        self.profile_scratch_align: int = metadata.profile_scratch_align
-        self.launch_cooperative_grid: bool = metadata.launch_cooperative_grid
-        self.launch_pdl: bool = metadata.launch_pdl
+        self.num_ctas: Final[int] = getattr(metadata, "num_ctas", 1)
+        self.global_scratch_size: Final[int] = metadata.global_scratch_size
+        self.global_scratch_align: Final[int] = metadata.global_scratch_align
+        self.profile_scratch_size: Final[int] = metadata.profile_scratch_size
+        self.profile_scratch_align: Final[int] = metadata.profile_scratch_align
+        self.launch_cooperative_grid: Final[bool] = metadata.launch_cooperative_grid
+        self.launch_pdl: Final[bool] = metadata.launch_pdl
+        self.impl: TVMFFILauncherImpl = TVMFFILauncherImpl(
+            [string_to_type(t) for t in self.signature],
+            self.launch_cooperative_grid,
+            self.launch_pdl,
+        )
+        self.launch: Callable[
+            [
+                int,
+                int,
+                int,
+                int,
+                int,
+                tuple[int, int, int],
+                object,
+                object,
+                object,
+                object,
+                object,
+                Sequence[Union[Any]],
+            ]
+        ] = self.impl.launch
 
     def __call__(
         self,
@@ -52,9 +72,9 @@ class TVMLauncher(object):
         assert not self.launch_cooperative_grid
         assert not self.launch_pdl
 
-        args: Sequence[TypedValue] = TypedValue.make_typed_values(self.signature, args)
+        # args: Sequence[TypedValue] = TypedValue.make_typed_values(self.signature, args)
 
-        return self.launch(
+        return self.impl.launch(
             gridX,
             gridY,
             gridZ,
@@ -64,8 +84,6 @@ class TVMLauncher(object):
             launch_metadata,
             launch_enter_hook,
             launch_exit_hook,
-            self.launch_cooperative_grid,
-            self.launch_pdl,
             global_scratch,
             profile_scratch,
             args,

python/triton_tvm_ffi/utils.py

@@ -5,8 +5,7 @@ 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, Sequence
-    from tvm_ffi import Object
+    from collections.abc import Mapping
     from typing import Any
 # isort: on
 # fmt: on
@@ -20,7 +19,6 @@ 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(_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
@@ -32,7 +30,6 @@ __all__ = [
     "cuOccupancyMaxActiveClusters",
     "fill_tma_descriptor",
     "get_device_properties",
-    "launch",
     "load_binary",
     "set_printf_fifo_size",
     # tvm-ffi-stubgen(end)