forked from ccf-ai-infra/Intro-ops
50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
from __future__ import annotations
|
|
|
|
import ctypes
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from .ctypes_bindings import CFunctions, Descriptor, check_status
|
|
|
|
|
|
@dataclass
|
|
class PreparedOp:
|
|
funcs: CFunctions
|
|
descriptor: Descriptor
|
|
workspace: torch.Tensor | None
|
|
runner_args: tuple[Any, ...]
|
|
stream_tensor: torch.Tensor | None = None
|
|
|
|
def run(self) -> None:
|
|
stream = torch.cuda.current_stream(device=self.stream_tensor.device if self.stream_tensor is not None else None)
|
|
workspace_ptr = None if self.workspace is None else ctypes.c_void_p(self.workspace.data_ptr())
|
|
workspace_size = 0 if self.workspace is None else self.workspace.numel()
|
|
status = self.funcs.execute(
|
|
self.descriptor,
|
|
workspace_ptr,
|
|
workspace_size,
|
|
*self.runner_args,
|
|
ctypes.c_void_p(stream.cuda_stream),
|
|
)
|
|
check_status(status)
|
|
|
|
def destroy(self) -> None:
|
|
if self.descriptor:
|
|
check_status(self.funcs.destroy(self.descriptor))
|
|
self.descriptor = Descriptor()
|
|
|
|
def __enter__(self) -> "PreparedOp":
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
self.destroy()
|
|
|
|
def __del__(self) -> None:
|
|
try:
|
|
self.destroy()
|
|
except Exception:
|
|
pass
|
|
|