minminsun commented on a change in pull request #4234: Auto TensorCore CodeGen URL: https://github.com/apache/incubator-tvm/pull/4234#discussion_r343661828
########## File path: tutorials/autotvm/tune_tensor_core_matmul.py ########## @@ -0,0 +1,219 @@ +import logging +import sys + +import numpy as np +import tvm + +from tvm import autotvm + + +def matmul_nn(A, B, L, dtype='float16', layout='NN'): + k = tvm.reduce_axis((0, L), name='k') + if dtype == 'float16': + out_type = 'float' + elif dtype == 'int8': + out_type = 'int' + if (layout == 'NN'): + return tvm.compute((N, M), lambda i, j: tvm.sum((A[i, k] * B[k, j]).astype(out_type), axis=k)) + if (layout == 'NT'): + return tvm.compute((N, M), lambda i, j: tvm.sum((A[k, i] * B[k, j]).astype(out_type), axis=k)) + if (layout == 'TN'): + return tvm.compute((N, M), lambda i, j: tvm.sum((A[i, k] * B[j, k]).astype(out_type), axis=k)) + if (layout == 'TT'): + return tvm.compute((N, M), lambda i, j: tvm.sum((A[k, i] * B[j, k]).astype(out_type), axis=k)) + [email protected] +def test_gemm_nn(N, L, M, dtype, layout): + if (layout == "NN"): Review comment: Done. Added more comments to the final formal tutorials. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services
