[FRONTEND] Expose end-to-end compile to python frontend (#58)

This commit is contained in:
Yan Chunwei
2022-08-18 01:42:48 +08:00
committed by GitHub
parent 95bbac41e7
commit b1673caaf6
15 changed files with 228 additions and 165 deletions

View File

@@ -0,0 +1,23 @@
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'))
def test_empty_kernel_cubin_compile():
@triton.jit
def kernel(X, stride_xm, stride_xn, BLOCK: tl.constexpr):
pass
device = torch.cuda.current_device()
cubin = triton.compile(kernel,
"*fp32,i32,i32",
device=device,
constants={"BLOCK": 256},
output="cubin")
print('cubin size:', len(cubin))
assert len(cubin) > 0