Conflict free shared memory permutation in tensorir

LeiWang1999 · November 19, 2022, 5:33am

@vinx13 @masahi Thanks alot, and quite interesting, I think I need to rollback tvm and do test again. And I implemented cutlass dp4a permutation with tvm which can even be faster than cutlass about 3~4% in nn layout, so I believe we can do the same thing with tensorcore.

github.com

LeiWang1999/tvm_gpu_gemm/blob/master/tensorirscript_dp4a/sota_nn.py

"""
Problem definition:
    1.We read a int8 matrix A of shape (M, K) from global memory to shared memory.
    2.We need to do permutation on A to make it suitable for dp4a conflict free access.
    3.So we first need to read A from global memory to local memory.
    4.Then we need to do permutation on A in local memory.
    5.Finally we need to read A from local memory to shared memory.
Solution:
    In this python code, we use tensorir transform layout to do permutation.
    Take a Gemm example, and the size of Gemm is a Wrap tile of nvidia cutlass, which is 128x128x16.
Result:
    average time cost of 1 runs = 136.129 ms, 64616.1 GFLOPS. Sota Implementation!
"""

import tvm
from tvm import te
import numpy as np
import tvm.testing
from tvm.script import tir as T
import os

This file has been truncated. show original