gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[EXAMPLES][TUTORIAL] Changed to new triton.kernel API

Browse Source
This commit is contained in:
Philippe Tillet
2020-07-08 13:39:19 -04:00
committed by Philippe Tillet
parent c33d6d15f5
commit 4ccd78f1a6
1 changed files with 21 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
python/examples/tutorials/mat_mul.py
View File

@@ -23,9 +23,7 @@ class _dot(torch.autograd.Function):
for(int k=K; k>0; k-=TK) { for(int k=K; k>0; k-=TK) {
TYPE a[TM, TK] = *pa; TYPE a[TM, TK] = *pa;
TYPE b[TK, TN] = *pb; TYPE b[TK, TN] = *pb;
c += a @ b; c += a @ b;
pa = pa + TK * 1; pa = pa + TK * 1;
pb = pb + TK * ldb; pb = pb + TK * ldb;
} }
@@ -41,31 +39,34 @@ class _dot(torch.autograd.Function):
return c return c
kernel = dict()
@staticmethod @staticmethod
def _call(a, b): def _call(a, b):
# shapes
M, K = a.shape M, K = a.shape
K, N = b.shape K, N = b.shape
# leading dimension
lda = K lda = K
ldb = N ldb = N
ldc = N ldc = N
dtype = a.dtype dtype = a.dtype
# create kernel if necessary
c = triton.empty([M,N], dtype=dtype) if dtype not in _dot.kernel:
grid = lambda opt: [triton.cdiv(M, opt.d('TM')), triton.cdiv(N, opt.d('TN'))]
defines = { defines = {
'TYPE' : dtype, 'TYPE' : dtype,
'TM' : [32,64,128], 'TM' : [64, 128],
'TN' : [32,64,128], 'TN' : [64, 128],
'TK' : [8], 'TK' : [8, 16],
} }
_dot.kernel[dtype] = triton.kernel(_dot.src, num_warps=[2, 4], defines=defines)
_dot.kernel = triton.kernel(_dot.src, defines=defines) kernel = _dot.kernel[dtype]
_dot.kernel(a, b, c, M, N, K, lda, ldb, ldc, # allocate output
grid=grid, num_warps=4, defines=defines) c = triton.empty([M,N], dtype=dtype)
# enqueue
grid = lambda opt: [triton.cdiv(M, opt.d('TM')),
triton.cdiv(N, opt.d('TN'))]
kernel(a, b, c, M, N, K, lda, ldb, ldc, grid=grid)
return c return c