[CUTLASS] Support more kernels: int8, tf32, and 3xtf32

[masahi]

@comaniac @Laurawly @hwu36

The change is mostly boilerplate. Supports int8, tf32, and 3xtf32 for both gemm and conv2d. TF32 and 3xtf32 are meant for training and not expected to be fast on Geforce cards (see NVIDIA/cutlass#390). But I found that

• FasterRCNN runs reasonably fast on RTX 3070 with TF32 (about 28 msec on (1, 3, 512, 512) input).
• MaskRCNN with 3xtf32 results match with PyTorch fp32 reference (atol = rtol = 1e-5).

For int8, I tested on quantized resnet50. But for some reason, the performance is worse than fp16. I haven't investigated deeply yet. In general, int8 models available for benchmarking are quite limited.

[comaniac]

Sorry for the late review. LGTM.

[comaniac]

Thanks @masahi

[hwu36]

@masahi, cutlass profiler misses many 256x64, 64x256 tile sizes for turing/volta kernels. It needs some change like

--- a/tools/library/scripts/generator.py
+++ b/tools/library/scripts/generator.py
@@ -724,6 +724,8 @@ def GenerateSM70_TensorOp_884(manifest, cuda_version):
       TileDescription([256, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([128, 256, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc),
       TileDescription([128, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc),
+      TileDescription([256,  64, 32], 2, [4, 1, 1], math_inst, min_cc, max_cc),
+      TileDescription([ 64, 256, 32], 2, [1, 4, 1], math_inst, min_cc, max_cc),
       TileDescription([ 64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([128,  64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([ 64,  64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc),
@@ -941,6 +943,8 @@ def GenerateSM75_TensorOp_1688(manifest, cuda_version):
       TileDescription([256, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([128, 256, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc),
       TileDescription([128, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc),
+      TileDescription([ 64, 256, 32], 2, [1, 4, 1], math_inst, min_cc, max_cc),
+      TileDescription([256,  64, 32], 2, [4, 1, 1], math_inst, min_cc, max_cc),
       TileDescription([ 64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([128,  64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([ 64,  64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc),
@@ -1218,6 +1222,8 @@ def GenerateSM75_TensorOp_8832_TN(manifest, cuda_version):
       TileDescription([256, 128, 128], 2, [4, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([128, 256, 128], 2, [2, 4, 1], math_inst, min_cc, max_cc),
       TileDescription([128, 128, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc),
+      TileDescription([256,  64, 128], 2, [4, 1, 1], math_inst, min_cc, max_cc),
+      TileDescription([ 64, 256, 128], 2, [1, 4, 1], math_inst, min_cc, max_cc),
       TileDescription([ 64, 128, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([128,  64, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([ 64,  64, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc),
@@ -1355,6 +1361,8 @@ def GenerateSM75_TensorOp_88128(manifest, cuda_version):
       TileDescription([256, 128, 512], 2, [4, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([128, 256, 512], 2, [2, 4, 1], math_inst, min_cc, max_cc),
       TileDescription([128, 128, 512], 2, [2, 2, 1], math_inst, min_cc, max_cc),
+      TileDescription([ 64, 256, 512], 2, [1, 4, 1], math_inst, min_cc, max_cc),
+      TileDescription([256,  64, 512], 2, [4, 1, 1], math_inst, min_cc, max_cc),
       TileDescription([ 64, 128, 512], 2, [2, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([128,  64, 512], 2, [2, 2, 1], math_inst, min_cc, max_cc),
       TileDescription([ 64,  64, 512], 2, [2, 2, 1], math_inst, min_cc, max_cc),

Also, cutlass works the best with cuda 11.3+.