[PYTHON] Deleted 01-vector-add.py: it is an unnecessary duplicate of
01-vector-add.ipynb
This commit is contained in:
@@ -1,76 +0,0 @@
|
||||
import torch
|
||||
import triton
|
||||
|
||||
# source-code for Triton compute kernel
|
||||
# here we just copy-paste the above code without the extensive comments.
|
||||
# you may prefer to store it in a .c file and load it from there instead.
|
||||
_src = """
|
||||
__global__ void add(float* z, float* x, float* y, int N){
|
||||
// program id
|
||||
int pid = get_program_id(0);
|
||||
// create arrays of pointers
|
||||
int offset[BLOCK] = pid * BLOCK + 0 ... BLOCK;
|
||||
float* pz[BLOCK] = z + offset;
|
||||
float* px[BLOCK] = x + offset;
|
||||
float* py[BLOCK] = y + offset;
|
||||
// bounds checking
|
||||
bool check[BLOCK] = offset < N;
|
||||
// write-back
|
||||
*?(check)pz = *?(check)px + *?(check)py;
|
||||
}
|
||||
"""
|
||||
# This function returns a callable `triton.kernel` object
|
||||
# created from the above source code.
|
||||
# For portability, we maintain a cache of kernels for different `torch.device`
|
||||
# We compile the kernel with -DBLOCK=1024
|
||||
_kernels = dict()
|
||||
|
||||
def make_add_kernel(device):
|
||||
if device not in _kernels:
|
||||
defines = {'BLOCK': 1024}
|
||||
autotune_vals = [({'BLOCK': '1024'}, 4), ({'BLOCK': '2048'}, 4)]
|
||||
autotune_key = ["N"]
|
||||
_kernels[device] = triton.kernel(_src, device=device, defines=defines, autotune_vals=autotune_vals,
|
||||
autotune_key=autotune_key)
|
||||
return _kernels[device]
|
||||
|
||||
# This is a standard torch custom autograd Function
|
||||
# The only difference is that we can now use the above kernel
|
||||
# in the `forward` and `backward` functions.`
|
||||
class _add(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, x, y):
|
||||
# constraints of the op
|
||||
assert x.dtype == torch.float32
|
||||
# *allocate output*
|
||||
z = torch.empty_like(x)
|
||||
# *create launch grid*:
|
||||
# this is a function which takes compilation parameters `opt`
|
||||
# as input and returns a tuple of int (i.e., launch grid) for the kernel.
|
||||
# triton.cdiv is a shortcut for ceil division:
|
||||
# triton.cdiv(a, b) = (a + b - 1) // b
|
||||
grid = lambda opt: (triton.cdiv(z.shape[0], opt.BLOCK), )
|
||||
# *launch kernel*:
|
||||
# pointer to the data of torch tensors can be retrieved with
|
||||
# the `.data_ptr()` method
|
||||
kernel = make_add_kernel(z.device)
|
||||
kernel(z.data_ptr(), x.data_ptr(), y.data_ptr(), z.shape[0], grid=grid)
|
||||
return z
|
||||
|
||||
# Just like we standard PyTorch ops
|
||||
# We use the `.apply` method to create a
|
||||
# callable object for our function
|
||||
add = _add.apply
|
||||
|
||||
torch.manual_seed(0)
|
||||
x = torch.rand(32, device='cuda')
|
||||
y = torch.rand(32, device='cuda')
|
||||
za = x + y
|
||||
zb = add(x, y)
|
||||
print(za)
|
||||
print(zb)
|
||||
print(f'The maximum difference between torch and triton is ' f'{torch.max(torch.abs(za - zb))}')
|
||||
|
||||
th_ms = triton.testing.do_bench(lambda: x + y)
|
||||
tr_ms = triton.testing.do_bench(lambda: add(x, y))
|
||||
print(th_ms, tr_ms)
|
Reference in New Issue
Block a user