[CUTLASS] Support batch_matmul

[masahi]

Adds support for offloading batch_matmul via GemmBatched cutlass kernel. Also supports dynamic shape.
I didn't add a profiler specifically for GemmBatched kernel - I piggy back on the dense profiler because cutlass GemmBatched uses the same kernel for normal gemm parallelized over the batch dimension (Grid Z dim).

This allows me to test cutlass byoc on Huggingface BERT-large end to end. I also have a number for TensorRT measured using their BERT demo, but TensorRT one is using google's own implementation of BERT.

This is the current result comparing cutlass offload, autotvm native and TensorRT, all using tensor core on rtx3070. Input size is (8, 128) and numbers in milliseconds.

CUTLASS	AutoTVM	TensorRT (google's implementation)
23.0551	23.6729	14.0

Here is the detailed nvprof output from cutlass and autotvm:

• cutlass https://gist.github.com/masahi/01c064529f859afefd0cc34ef138aa08
• autotvm https://gist.github.com/masahi/45ac7c45b637c2f3e4c35f8db11e9c88

As you can see, for now activation fusion is enabled only for GeLU. There are other activations or elemwise ops that can be fused in principle, such as

• https://gist.github.com/masahi/45ac7c45b637c2f3e4c35f8db11e9c88#file-bert_large_gelu_fusion_autotvm-txt-L30-L32
• https://gist.github.com/masahi/45ac7c45b637c2f3e4c35f8db11e9c88#file-bert_large_gelu_fusion_autotvm-txt-L26

Together, they account for about 15% of total execution time, which is a shame. To fuse them, we need to overcome all of following blockers:

• Support ND input dense to remove reshape between dense and activations ([Topi] Allow relay.nn.dense support arbitrary number dimensions #8412)
• Support multiple "source" tensors in cutlass epilogue (Epilogue with mutiple sources NVIDIA/cutlass#347)
• Enable fusion of cast op into activation (Allowing "source" (bias tensor) and output tensor to have different data type NVIDIA/cutlass#352)

The nvprof output also shows that softmax is a huge bottleneck, accounting for 24% of e2e time. Apparently TVM's cuda softmax is 2x slower than cudnn, here is the result if we use cudnn's softmax:

CUTLASS	AutoTVM	TensorRT (google's implementation)
20.2517	21.1869	14.0

cc @comaniac @Laurawly @zhiics @hwu36

[comaniac]

I like the detail benchmarking. LGTM.

[masahi]

Thanks @comaniac

[hwu36]

batch gemm is nice to have, but it is a bit yesterday now. :)

cutlass group gemm is released with 2.8 and it is the way to go.