We propose a solution for TensorCore CodeGen with significant transparency, 
flexibility and usability. In this solution, the algorithm description and 
schedule of TensorCore CodeGen is no different than that of a normal CUDA 
CodeGen. All the information needed by wmma API, such as 
matrix_a/matrix_b/accumulator, row_major/col_major, warp tile size and so on, 
is automatically derived from the AST. Of course, not every algorithm and 
schedule is suitable for TensorCore computation. This solution will do the 
check and fall back to normal CUDA CodeGen for those that are not qualified for 
TensorCore CodeGen.

In this solution, 3 IRVisitors and 1 IRMutator are added. 
* IRVisitors: BodyVisitor, MMAMatcher and BufferAnalyser. 
* IRMutator: TensorCoreIRMutator.

BodyVisitor, which is called by ScheduleAnalyser, visits the body stmt of 
original ComputeOp to get the access indices of input matrices if it is 
recognized as matrix multiply. ScheduleAnalyser compares the access indices 
with the axis/reduce_axis of ComputeOp to figure out whether an input matrix is 
matrix_a or matrix_b, row_major or col_major.

MMAMatcher does the pattern matching on AST stmt. The pattern it tries to find 
out is as following:
 If matched, the a, b, c will be recorded as fragment registers, which are 
important inputs to the next visitor.

BufferAnalyser, the last visitor, will get all of the rest information needed 
for TensorCoreIRMutator, like strides of src/dst buffer for wmma load/store 
matrix operation, warp tile size for fragment allocation as well as checking 
whether the schedule is qualified for TensorCore, loops that need to be scaled 
after normal load/store and compute operation replaced by TensorCore 
operations, etc..

TensorCoreIRMutator mutates the AST stmt for TensorCore CodeGen. The subtree 
matched by MMAMatcher will be replaced with “mma_sync” extern call. Load/Store 
of fragments are replaced with “load/store_matrix_sync” extern call, with the 
thread index getting unified within a warp. Thread index unification, i.e. 
changing the index of every thread to the same as the first thread of the warp, 
is done by ThreadIdxMutator on the subtree.

The TensorCore IR Passes are applied before StorageFlatten because they need 
stride/shape and index of specific dimensions before they got flattened into 
one. Before StorageFlatten, “Allocation” is represented by Realize IR Node, 
which has no new_expr member as Allocate IR Node has. So we added it to Realize 
IR Node to carry the expr for fragment allocation and pass to Allocate IR Node. 
We noticed the comment of deprecating new_expr when merging with the latest TVM 
codebase. We would like to ask for a reconsideration of this decision, because 
it is really useful for some non-standard buffer allocations.

This solution is evaluated on a sample schedule of Matmul, which is based on 
AutoTVM. It supports fp16 and int8 data type, and three kinds of data layouts: 

On some model layers, we have already achieved better performance than 
### FP16 on V100, CUDA 9.0, Driver 396.44

* NMT Online Service (In-house Model)

| M, N, K | CUBLAS TensorCore | TVM TensorCore |  Speed Up |
| -------- | -------- | -------- | -------- |
| 512, 64, 512     | 9.05us     | 7.34us    | 1.23X |
| 512, 32, 512     | 8.30us     | 6.84us    | 1.21X |
| 512, 16, 512     | 7.88us     | 6.60us    | 1.19X |

* MobileNet (Public Model)

 TVM TensorCore | SpeedUp |
| -------- | -------- | -------- | -------- |
| 56 56 64 128 1 1 0 0 1 1 |  8.5220us | 6.9320us | 1.23X |
| 28 28 128 256 1 1 0 0 1 1 | 10.787us |  8.3490us | 1.29X |
| 28 28 256 256 1 1 0 0 1 1 | 15.188us |14.136us | 1.07X |

### Int8 on T4, CUDA10.1, Driver 418.39
* NMT Online Service (In-house Model)

| M, N, K |  CUBLAS TensorCore | TVM TensorCore |  Speed Up |
| -------- | -------- | -------- | -------- |
| 512, 64, 512     | 23.163us     | 22.603us    | 1.025X |
| 512, 32, 512     | 22.551us     | 14.263us    | 1.58X |
| 512, 16, 512     | 22.510us     |  11.015us    | 2.04X |

There are also many shapes on which CUBLAS/CUDNN is much better. The 
performance tuning is still on-going. 

-- Minmin Sun, Lanbo Li, Chenfan Jia and Jun Yang of Alibaba PAI team

You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:

Reply via email to