a9464f4993
This PR does the following things: - Code refactoring on Load and Store op codegen, rewrite with same logic and share much code - Support the vectorized load/store
46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
import torch
|
|
|
|
import triton
|
|
import triton.language as tl
|
|
import triton.runtime as runtime
|
|
|
|
# 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()
|
|
cubin = triton.compile(empty_kernel,
|
|
"*fp32,i32,i32",
|
|
device=device,
|
|
constants={"BLOCK": 256},
|
|
output="cubin")
|
|
|
|
print('cubin size:', len(cubin))
|
|
assert len(cubin) > 0
|
|
|
|
|
|
def test_empty_kernel_launch():
|
|
device = torch.cuda.current_device()
|
|
binary = runtime.build_kernel(empty_kernel, "*fp32,i32,i32",
|
|
constants={"BLOCK": 256},
|
|
num_warps=4,
|
|
num_stages=3)
|
|
grid = lambda META: (
|
|
triton.cdiv(1024, META['BLOCK']) * triton.cdiv(1024, META['BLOCK']),
|
|
)
|
|
|
|
A = torch.zeros([1024], device="cuda")
|
|
runtime.launch_kernel(kernel=binary,
|
|
grid=grid,
|
|
device=device,
|
|
X=A,
|
|
stride_xm=256,
|
|
BLOCK=tl.constexpr(256))
|