implement launch with cpp

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

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,