mirror of
https://github.com/sgjzfzzf/triton-tvm-ffi.git
synced 2026-07-01 08:51:56 +08:00
59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
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, load_binary
|
|
|
|
|
|
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'
|
|
)
|
|
|
|
|
|
class TVMFFIDriver(CudaDriver):
|
|
def __init__(self, *args, **kwargs) -> TVMFFIDriver:
|
|
super().__init__(*args, **kwargs)
|
|
self.utils: TVMFFIUtils = TVMFFIUtils()
|
|
|
|
|
|
del CudaDriver
|