40 lines
843 B
Python
40 lines
843 B
Python
import torch
|
|
import triton
|
|
|
|
class _add(torch.autograd.Function):
|
|
src = """
|
|
__global__ void add(float* z, float* x, float* y, int N) {
|
|
|
|
int pid = get_program_id(0);
|
|
|
|
int offset[TILE] = pid * TILE + 0 ... TILE;
|
|
float* pz[TILE] = z + offset;
|
|
float* px[TILE] = x + offset;
|
|
float* py[TILE] = y + offset;
|
|
|
|
bool check[TILE] = offset < N;
|
|
|
|
*pz = *px + *py;
|
|
}
|
|
"""
|
|
|
|
kernel = triton.kernel(src, defines={'TILE': 1024}, num_warps=[4])
|
|
|
|
@staticmethod
|
|
def forward(ctx, x, y):
|
|
z = torch.empty_like(x).cuda()
|
|
N = x.numel()
|
|
grid = lambda opt: (triton.cdiv(N, opt.d('TILE')),)
|
|
_add.kernel(z,x,y, N, grid=grid)
|
|
return z
|
|
|
|
add = _add.apply
|
|
|
|
# test
|
|
torch.manual_seed(0)
|
|
x = torch.rand(900).cuda()
|
|
y = torch.rand(900).cuda()
|
|
za = x + y
|
|
zb = add(x, y)
|
|
print(torch.allclose(za,zb))
|