mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-15 22:56:49 +08:00
27 lines
885 B
Python
27 lines
885 B
Python
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
@triton.jit
|
|
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
|
|
pid = tl.program_id(axis=0)
|
|
block_start = pid * BLOCK_SIZE
|
|
offsets = block_start + tl.arange(0, BLOCK_SIZE)
|
|
mask = offsets < n_elements
|
|
x = tl.load(x_ptr + offsets, mask=mask)
|
|
y = tl.load(y_ptr + offsets, mask=mask)
|
|
output = x + y
|
|
tl.store(output_ptr + offsets, output, mask=mask)
|
|
|
|
def add(x: torch.Tensor, y: torch.Tensor):
|
|
output = torch.empty_like(x)
|
|
n_elements = output.numel()
|
|
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
|
|
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
|
|
return output
|
|
|
|
a = torch.rand(3, device="cuda")
|
|
b = a + a
|
|
b_compiled = add(a, a)
|
|
print(b_compiled - b)
|
|
print("If you see tensor([0., 0., 0.], device='cuda:0'), then it works") |