[CUTLASS] Support conv2d activation fusion

[masahi]

@comaniac @Laurawly @hwu36 @manishucsd @junrushao1994 @vinx13

This adds conv2d + activation fusion support for bias_add, relu, and fp32 sigmoid. I have pending PRs at the cutlass repo NVIDIA/cutlass#378, NVIDIA/cutlass#379 which will enable fp16 sigmoid, silu, and hardswish fusion as well.

End to end results

All code, models, and nvprof dump etc are available at https://github.com/masahi/tvm-cutlass-eval

All numbers in milli sec, using fp16 models with fp16 accumulation running on tensorcore, measured on RTX 3070

Model name	Input size	CUTLASS (no fusion)	CUTLASS with fusion	cuDNN	AutoTVM TensorCore	TensorRT fp16
resnet50	(8, 3, 224, 224)	3.47	3.16	4.08	4.14	2.53
efficientnet_v2	(8, 3, 224, 224)	8.18	8.26 (7.80 if use fast_math)	14.0	13.2	5.25
DETR-R50	(8, 3, 800, 750)	61.06	57.19	68.4	80.5	NA
deeplabv3_mobilenet_v3_large	(8, 3, 512, 512)	10.3	9.16	17.5	15.9	19.2 (?)
YOLOv5l	(8, 3, 512, 512)	33.6	25.4 (24.2 if use fast_math)	34.8	N/A	N/A

Observations

• First of all, the cutlass results above are still suboptimal because residual blocks are not completely fused due to a limitation in cutlass discussed in Epilogue with mutiple sources NVIDIA/cutlass#347. So if we look at the partitioned subgraph of resnet50, for example, the main function looks like this (see the whole model at https://github.com/masahi/tvm-cutlass-eval/blob/master/resnet50/resnet50_partition.txt):

  ...
  %4 = @tvmgen_default_cutlass_main_3(%3, meta[relay.Constant][2] /* ty=Tensor[(64, 1, 1, 64), float16] */, meta[relay.Constant][3] /* ty=Tensor[(64), float16] */) /* ty=Tensor[(8, 56, 56, 64), float16] */;
  %5 = @tvmgen_default_cutlass_main_6(%4, meta[relay.Constant][4] /* ty=Tensor[(64, 3, 3, 64), float16] */, meta[relay.Constant][5] /* ty=Tensor[(64), float16] */) /* ty=Tensor[(8, 56, 56, 64), float16] */;
  %6 = @tvmgen_default_cutlass_main_9(%5, meta[relay.Constant][6] /* ty=Tensor[(256, 1, 1, 64), float16] */, meta[relay.Constant][7] /* ty=Tensor[(256), float16] */) /* ty=Tensor[(8, 56, 56, 256), float16] */;
  %7 = @tvmgen_default_cutlass_main_12(%3, meta[relay.Constant][8] /* ty=Tensor[(256, 1, 1, 64), float16] */, meta[relay.Constant][9] /* ty=Tensor[(256), float16] */) /* ty=Tensor[(8, 56, 56, 256), float16] */;
  %8 = add(%6, %7) /* ty=Tensor[(8, 56, 56, 256), float16] */;
  %9 = nn.relu(%8) /* ty=Tensor[(8, 56, 56, 256), float16] */;
  %10 = @tvmgen_default_cutlass_main_15(%9, meta[relay.Constant][10] /* ty=Tensor[(64, 1, 1, 256), float16] */, meta[relay.Constant][11] /* ty=Tensor[(64), float16] */) /* ty=Tensor[(8, 56, 56, 64), float16] */;
  %11 = @tvmgen_default_cutlass_main_18(%10, meta[relay.Constant][12] /* ty=Tensor[(64, 3, 3, 64), float16] */, meta[relay.Constant][13] /* ty=Tensor[(64), float16] */) /* ty=Tensor[(8, 56, 56, 64), float16] */;
  %12 = @tvmgen_default_cutlass_main_21(%11, meta[relay.Constant][14] /* ty=Tensor[(256, 1, 1, 64), float16] */, meta[relay.Constant][15] /* ty=Tensor[(256), float16] */) /* ty=Tensor[(8, 56, 56, 256), float16] */;
  %13 = add(%12, %9) /* ty=Tensor[(8, 56, 56, 256), float16] */;
  %14 = nn.relu(%13) /* ty=Tensor[(8, 56, 56, 256), float16] */;
  ...

The intermediate adds are the element-wise addition in the residual block, which can in principle be fused with the preceding conv2d. nvprof output shows that these unfused ops are taking more than 20% of e2e time. If we fuse them, I believe we could approach TRT-level performance.

• Despite its suboptimality, enabling activation fusion still results in a good speedup.
• Cutlass results are always better than cuDNN ones, even if fusion is disabled. The difference is especially large for deeplabv3 and efficientnetv2, both of which use a lot of depthwise conv2d. The nvprof output shows that cuDNN is spending a lot of time doing layout transform (see for example). Maybe our use of cuDNN is not optimal in terms of API usage and kernel selections. nvprof dumps for cuDNN-runs are available in the repository, if someone wants to compare against cutlass nvprof dumps.
• Fusing sigmoid activation into cutlass conv2d results in worse runtime (see efficientnet_v2 row), which I found odd since TVM should be generating essentially the same sigmoid kernel in the unfused case. The fast-math option, discussed in Support half precision sigmoid activation NVIDIA/cutlass#378 (comment), makes it a lot better at the cost of slight accuracy loss.
• efficientnetv2 and deeplabv3 use a lot of depthwise conv2d, which is not currently offloaded to cutlass. They are actually the bottleneck in the results above, which can be observed by looking at nvprof dumps. Using AutoTVM for them should help bring down e2e time further.

Known issues

• Kernel profiling time is extremely slow. For DETR-R50, I had to wait like 20 min. More discussion in [CUTLASS] Add conv2d profiler #9737 (comment)
• Trying to fuse silu activation in YOLOv5l results in a strange type inference error during MergeComposite. So the YOLOv5l result above doesn't use silu fusion. (Fixed a bug in type relation)
• I also tried MaskRCNN, but currently it results in a strange output (no detection) when cutlass BYOC is used. The same issue also applies to cuDNN offload, but TVM-native kernels (target = "cuda") work correctly.

[masahi]

@comaniac @Laurawly Ready to merge?

[comaniac]

LGTM. Thanks.

[comaniac]

Thanks @masahi