mirror of
https://github.com/sgjzfzzf/triton-tvm-ffi.git
synced 2026-07-01 08:51:56 +08:00
remove CudaUtils dependencies
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
@@ -1,8 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Mapping, Type
|
||||
from ctypes import c_void_p
|
||||
from typing import Mapping, Tuple, Type
|
||||
from triton.backends.nvidia.driver import CudaDriver
|
||||
from .utils import get_device_properties
|
||||
from .utils import get_device_properties, load_binary
|
||||
|
||||
|
||||
class TVMFFIUtils(object):
|
||||
@@ -13,12 +14,11 @@ class TVMFFIUtils(object):
|
||||
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
from triton.backends.nvidia.driver import CudaUtils
|
||||
|
||||
self._utils: CudaUtils = CudaUtils()
|
||||
|
||||
def load_binary(self, *args, **kwargs):
|
||||
return self._utils.load_binary(*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)
|
||||
@@ -34,13 +34,19 @@ class TVMFFIUtils(object):
|
||||
)
|
||||
|
||||
def fill_tma_descriptor(self, *args, **kwargs):
|
||||
return self._utils.fill_tma_descriptor(*args, **kwargs)
|
||||
raise NotImplementedError(
|
||||
'"fill_tma_descriptor" hasn\'t been supported for Hopper'
|
||||
)
|
||||
|
||||
def launch(self, *args, **kwargs):
|
||||
return self._utils.launch(*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):
|
||||
return self._utils.build_signature_metadata(*args, **kwargs)
|
||||
raise NotImplementedError(
|
||||
'"launch" is introduced in triton after commit d2b3925410689155e0f6028e8554bba972989348, which is still not supported yed'
|
||||
)
|
||||
|
||||
|
||||
class TVMFFIDriver(CudaDriver):
|
||||
|
||||
Reference in New Issue
Block a user