[PYTHON] Deleted 01-vector-add.py: it is an unnecessary duplicate of

01-vector-add.ipynb
This commit is contained in:
Philippe Tillet
2021-03-04 02:06:57 -05:00
parent 62835a0979
commit 3ecf834a69

View File

@@ -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)