mirror of
https://github.com/sgjzfzzf/triton-tvm-ffi.git
synced 2026-07-01 08:51:56 +08:00
133 lines
4.5 KiB
Python
133 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
from ctypes import c_void_p
|
|
from typing import List, Mapping, Tuple, Type
|
|
from triton.backends.nvidia.driver import CudaDriver
|
|
from .utils import get_device_properties, launch, 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 TVMLauncher(object):
|
|
def __init__(self, src: List[bool], metadata, *args, **kwargs) -> TVMLauncher:
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.mask: List[bool] = [annotation != "constexpr" for annotation in 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
|
|
|
|
# We assume the global Triton allocator is not enabled: `_allocator` must be a NullAllocator.
|
|
# This module depends on NullAllocator behavior; ensure no other code replaces the allocator.
|
|
from triton.runtime._allocation import _allocator, NullAllocator
|
|
|
|
assert isinstance(_allocator.get(), NullAllocator)
|
|
|
|
def __call__(
|
|
self,
|
|
gridX,
|
|
gridY,
|
|
gridZ,
|
|
stream,
|
|
function,
|
|
kernel_metadata,
|
|
launch_metadata,
|
|
launch_enter_hook,
|
|
launch_exit_hook,
|
|
*args,
|
|
):
|
|
from triton.runtime import _allocation
|
|
|
|
def allocate_scratch(size, align, allocator):
|
|
if size > 0:
|
|
grid_size = gridX * gridY * gridZ
|
|
alloc_size = grid_size * self.num_ctas * size
|
|
alloc_fn = allocator.get()
|
|
return alloc_fn(alloc_size, align, stream)
|
|
return None
|
|
|
|
global_scratch = allocate_scratch(
|
|
self.global_scratch_size, self.global_scratch_align, _allocation._allocator
|
|
)
|
|
profile_scratch = allocate_scratch(
|
|
self.profile_scratch_size,
|
|
self.profile_scratch_align,
|
|
_allocation._profile_allocator,
|
|
)
|
|
assert not self.launch_cooperative_grid
|
|
assert not self.launch_pdl
|
|
assert len(self.mask) == len(args)
|
|
args = [arg for arg, m in zip(args, self.mask) if m]
|
|
return launch(
|
|
gridX,
|
|
gridY,
|
|
gridZ,
|
|
stream,
|
|
function,
|
|
kernel_metadata,
|
|
launch_metadata,
|
|
launch_enter_hook,
|
|
launch_exit_hook,
|
|
self.launch_cooperative_grid,
|
|
self.launch_pdl,
|
|
global_scratch,
|
|
profile_scratch,
|
|
*args,
|
|
)
|
|
|
|
|
|
class TVMFFIDriver(CudaDriver):
|
|
def __init__(self, *args, **kwargs) -> TVMFFIDriver:
|
|
super().__init__(*args, **kwargs)
|
|
self.utils: TVMFFIUtils = TVMFFIUtils()
|
|
self.launcher_cls: Type[TVMLauncher] = TVMLauncher
|
|
|
|
|
|
del CudaDriver
|