22ec22c257
This PR merges the new runtime back into the `triton-mlir` branch. This adds caching and just-in-time compilation functionality to the triton-mlir project, and paves the way for re-using tests from the master branch.
33 lines
806 B
Python
33 lines
806 B
Python
import torch
|
|
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
# trigger the torch.device implicitly to ensure cuda context initialization
|
|
torch.zeros([10], device=torch.device('cuda'))
|
|
|
|
|
|
@triton.jit
|
|
def empty_kernel(X, stride_xm, BLOCK: tl.constexpr):
|
|
pass
|
|
|
|
|
|
def test_empty_kernel_cubin_compile():
|
|
|
|
device = torch.cuda.current_device()
|
|
kernel = triton.compile(empty_kernel,
|
|
"*fp32,i32,i32",
|
|
device=device,
|
|
constants={"BLOCK": 256})
|
|
|
|
assert len(kernel.asm["cubin"]) > 0
|
|
|
|
|
|
def test_empty_kernel_launch():
|
|
grid = lambda META: (
|
|
triton.cdiv(1024, META['BLOCK']) * triton.cdiv(1024, META['BLOCK']),
|
|
)
|
|
|
|
A = torch.zeros([1024], device="cuda")
|
|
empty_kernel[grid](X=A, stride_xm=256, BLOCK=256)
|