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

python/examples/tutorials/mat_mul.py

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