put utils in C++ side

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

python/triton_tvm_ffi/driver.py

@@ -1,67 +1,23 @@
 from __future__ import annotations
 
-from ctypes import c_void_p
-from typing import Any, List, Mapping, Optional, Tuple, Type
+from typing import Any, List, Optional, Type
 from triton.backends.nvidia.driver import CudaDriver
-from . import TypedValue, get_device_properties, launch, load_binary, string_to_type
-
-
-class TVMFFIUtils(object):
-    def __new__(cls: Type[TVMFFIUtils]) -> TVMFFIUtils:
-        if not hasattr(cls, "instance"):
-            cls.instance = super().__new__(cls)
-        return cls.instance
-
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-
-    def load_binary(
-        self, name: str, data: bytes, shared: int, device: int
-    ) -> Tuple[c_void_p, c_void_p, int, int, int]:
-        return load_binary(name, data, shared, device)
-
-    def get_device_properties(self, device_id: int) -> Mapping[str, int]:
-        return get_device_properties(device_id)
-
-    def cuOccupancyMaxActiveClusters(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"cuOccupancyMaxActiveClusters isn\'t expected to be invoked"'
-        )
-
-    def set_printf_fifo_size(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"set_printf_fifo_size" isn\'t expected to be invoked'
-        )
-
-    def fill_tma_descriptor(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"fill_tma_descriptor" hasn\'t been supported for Hopper'
-        )
-
-    def launch(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"launch" is introduced in triton after commit d2b3925410689155e0f6028e8554bba972989348, which is still not supported yed'
-        )
-
-    def build_signature_metadata(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"launch" is introduced in triton after commit d2b3925410689155e0f6028e8554bba972989348, which is still not supported yed'
-        )
+from . import TypedValue, utils, string_to_type
 
 
 class TVMLauncher(object):
     def __init__(self, src: List[bool], metadata, *args, **kwargs) -> TVMLauncher:
         super().__init__(*args, **kwargs)
 
-        self.signature = src.signature.values()
-        self.num_ctas = getattr(metadata, "num_ctas", 1)
-        self.launch = launch
-        self.global_scratch_size = metadata.global_scratch_size
-        self.global_scratch_align = metadata.global_scratch_align
-        self.profile_scratch_size = metadata.profile_scratch_size
-        self.profile_scratch_align = metadata.profile_scratch_align
-        self.launch_cooperative_grid = metadata.launch_cooperative_grid
-        self.launch_pdl = metadata.launch_pdl
+        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
 
     def __call__(
         self,
@@ -105,7 +61,7 @@ class TVMLauncher(object):
 
         args = [canonicalize(arg, sig) for arg, sig in zip(args, self.signature)]
 
-        return launch(
+        return self.launch(
             gridX,
             gridY,
             gridZ,
@@ -126,7 +82,7 @@ class TVMLauncher(object):
 class TVMFFIDriver(CudaDriver):
     def __init__(self, *args, **kwargs) -> TVMFFIDriver:
         super().__init__(*args, **kwargs)
-        self.utils: TVMFFIUtils = TVMFFIUtils()
+        self.utils = utils
         self.launcher_cls: Type[TVMLauncher] = TVMLauncher