* [BYOC] Make CUTLASS BYOC integration 'Collage friendly'
(See https://discuss.tvm.apache.org/t/byoc-supporting-cutlass-byoc-with-collage/12796/6 for
context, which in turn is part of Collage (https://github.com/apache/tvm-rfcs/blob/main/rfcs/0062-collage.md).
Currently CUTLASS has four entry points:
- The usual 'partition_for_cutlass' partitioning function, using the
standard pattern table and pass machinery (see cutlass/build.py).
- A 'tune_cutlass_kernels' function which augments CUTLASS partition
functions with the results of building and running test kernels (see cutlass/build.py).
- A 'relay.ext.cutlass' external codegen function which inspects the
turning results and generates a CSourceModule for each partitions
(see cutlass/codegen.cc).
- A 'build_cutlass_kernels_vm' function which runs 'export_library' with
all the nvcc compiler options needed to build all the CSourceModules
(see cutlass/bild.py).
For Collage we'd like CUTLASS to have only two entry points: 'partition_for_cutlass',
and 'relay.ext.cutlass' or equivalent. This makes the CUTLASS external codegen integration
composable with other integrations, which in turn helps Collage avoid having to understand any
external codegen APIs other than the global pattern table and the custom compilation function/pass.
Collage also tends to end up requiring multiple partitions for the same backend since it is
more aggressive at mixing-and-matching smaller sub-graphs between backends. Thus we'd also like
to make sure all tuning, generated code and compilation overhead is shared between all such CUTLASS
partitions.
So, in this PR:
- We add all the CUTLASS-specific tuning and compilation options as new Target
attributes for the 'external codegen' "cutlass" TargetKind (cutlass/target.cc).
The user now has one place to provide those settings, and we've already done the
legwork to plumb the target instance.
- We replace 'relay.ext.cutlass' with a 'RelayToTIR' custom pass hook
'CompileForCutlass' (see cutlass/codegen.cc). This pass obviously can see all
the CUTLASS partitions in the IRModule, so we can now share tuning results
between them all and can be sure to generate a single CSourceModule. The pass can
also invoke the compiler to yield a StaticModule, which we've also already done the
legwork to support. In this way all CUTLASS-specific steps are handled at once.
- For convenience we supply 'finalize_modules' and 'finalize_modules_vm' which
invoke nvcc for final linking (using export_library as usual). However, there's now
nothing CUTLASS specific in those helpers other than their overriding of the 'compiler' to
be nvcc.
- test_cutlass.py is updated to use the new API.
Though this is a breaking change for existing users of the CUTLASS integration the
change is pretty minor, as shown in test_cutlass.py.
* - Masa's comments
* - Remove unnecessary save.
Mark Shields
dfc8e95604