[BYOC] CUTLASS integration

[Laurawly]

As discussed in RFC: https://discuss.tvm.apache.org/t/rfc-byoc-nvidia-cutlass-integration/9147, this PR is a CUTLASS integration of GEMM to TVM. It also includes a profiler to search for best params in CUTLASS.
@masahi Please take over this PR, Thanks!

[masahi]

Thanks @Laurawly, I'll work on this as my top priority!

[junrushao]

Really cool work @Laurawly!

[masahi]

One question for @comaniac: Currently, this integration is implemented with C source backend. But one big use case for cutlass BYOC is for dynamic workloads, for which we cannot do well at the moment. I think C source backend is not compatible with dynamic inputs, is that right? On the other hand, with json codegen/runtime, I can see a way to support dynamic inputs.

However, since cutlass is a template library, the only way it could work with the json codegen/runtime is to instantiate all possible templates that we care about ahead of time, and build them together with libtvm.so. I wonder if that is feasible. cc @Laurawly
One benefit of the above approach is that we don't need to commit the cutlass code generator into our repo, which is a big part of this PR.

[comaniac]

IMHO, CUTLASS doesn't naturally benefit dynamic workloads due to the exact reason you mentioned. We internally use CUTLASS for training and it works well because we JIT kernels with known shapes in runtime.

In the case of CUTLASS with BYOC in TVM for inference, my impression is we could leverage high performance kernel templates while 1) keeping the binary self-contained, 2) fusing ops, and 3) having lightweight tuning (e.g., ~10 trials similar to CUDNN). On the other hand, dynamic workloads are still challenging, and hopefully our ongoing efforts of dynamic kernel tuning and generation could be landed soon to make it happen.

[masahi]

hmm, are you concerned about slow performance due to lack of tuning, or the integration problem that I brought up? I wouldn't worry about the former because cublas is fast without tuning. So I expect cutlass to perform equally well.

Since our cublas offload supports dynamic inputs, not supporting dynamic inputs for cutlass would be a big bummer imo. So I want to discuss the integration problem first, and investigate performance issues later.

[zhiics]

I think even C source code should be able to handle dynamic shapes as well since it only expects tensors at the runtime, or I might forget something here. But in general, json format is more recommended as it is easier to debug and maintain. It is also more friendly in handling constants.

[masahi]

I think even C source code should be able to handle dynamic shapes as well since it only expects tensors at the runtime

Exactly, I've looked at the code a bit and it seems currently we generate an API that takes DLTensors as inputs, but it passes only raw pointers to the generated backend calls. It is easy to change things a bit to pass shape information as well.

tvm/src/relay/backend/contrib/codegen_c/codegen_c.h

Lines 176 to 185 in 3cb838d

	code_stream_ << func_name << "_(";
	for (size_t i = 0; i < args.size(); i++) {
	const auto& dtype_str = GetDtypeString(args[i]);
	code_stream_ << "(" << dtype_str << "*)(arg" << i << "->data),\n";
	PrintIndents();
	}
	for (size_t i = 0; i < outs.size() - 1; i++) {
	code_stream_ << "(" << outs[i].dtype << "*)(out" << i << "->data),\n";
	PrintIndents();
	}

I'm glad to find at least one path that supports all use cases. For the json codegen/runtime path, I'm not sure how to integrate cutlass's JIT codegen + compile approach with it.

[comaniac]

Ah sorry I didn't make it clear. The interface of C source codegen does deal with dynamic workloads, because it takes the raw pointer which could be any size in run time.

What I meant was how to generate CUTLASS kernels that are able to perform well with all shapes. In @Laurawly's post, they generate lots of kernels to cover possible shapes, which result in 7GB binary. I assume they also generate a run time dispatching logic (also in the generated C source code) to determine which kernel should be used given the known shape in run time. Obviously, the binary size will definitely be an issue for this solution.

For JSON codegen/runtime, it would be similar to TensorRT: We simply dump a JSON graph in codegen without doing anything else. Meanwhile, we have a custom runtime that JITs/catches CUTLASS kernels based on known shapes. This results in a much smaller binary, but the first execution (or an execution with new shapes) may take several seconds or even a minute to JIT all kernels.

[hwu36]

I assume they also generate a run time dispatching logic (also in the generated C source code) to determine which kernel should be used given the known shape in run time.

At this moment, CUTLASS does not provide heuristics to choose which kernel to use based on the runtime information. The user of CUTLASS needs to make the decision to pick kernels.

[comaniac]

Otherwise LGTM.

[comaniac]

I'm wondering whether we should use "profile" here. My impression to "profile" is more like an analysis rather than transformation. For example, profile_executor sounds getting a report of detail execution latencies. Maybe "tune_cutlass_kernels" (similar to AutoTVM/Ansor) or "search_cutlass_kernels" (similar to CuDNN) might be better? Would like to hear to others as well.

[masahi]

ok I'll change to tune_cutlass_kernels. profile is what cutlass uses in their repo (e.g. cutlass_profiler), but yeah, might be confusing to TVM users.

My only concern with tune is, people might get a wrong impression that tune_cutlass_kernels is an optional thing, since in Ansor / AutoTVM tuning is optional. We can add "default kernels" later to avoid having tuning as a hard requirement. We need default kernels for dynamic workload anyway.

[comaniac]

That's a fair concern and I agree with your plan.

[masahi]

Thanks @Laurawly @comaniac @junrushao1994 @zhiics @hwu36 this is merged!

[hwu36]

This is so cool. Thank you everyone.