I tried to do to the TensorRT integration what #11631 did to the CUTLASS integration, viz:
- Make sure all compilation options are passed in Target instances. This helps Collage.
- Use a custom pass invoked via RelayToTIRTargetHooks instead of the relay.ext.$toolchain mechanism.
This helps use decouple external codegen from lowering.
This PR collects the prep for that change:
- TensorRT uses the JSONSerializer visitor to encode each partition function. Previously, when the
visitor encountered a Constant it simply generated and recorded a name for the constant. Then,
completely separately, and via a callback in TECompiler, the function is visited again in the
same order and with the same name generation convention by a ConstantUpdater to actually collect the
bindings, which are then encoded into a ConstLoaderModule to be made available at runtime.
However if all TensorRT compilation is to be done by a stand-alone pass there's no TECompiler callback
hackery available. So I've added a "const_name_to_ndarray" attribute to the IRModule of type
Map<String, runtime::NDArray> so that named constants can be accumulated throughout compilation by
any pass which needs to do so. Then the Graph, AOT and VM executors are all updated to merge those
constants into the final runtime artifact
(Compare with "Constants", the equivalent attribute for extracting TIR AllocateConsts.)
- The TensorRT tests use the create_executor interface but it wasn't quite ready for the
new more general form of passing list-of-targets.
- I want TensorRT compilation to work out of the box without the need for any special targets if
all the default options should apply. Go back and make the CUTLASS integration I did follow the
same convention.
- To test this I also switched the 'demo' "ccompiler" external codegen target to IRModule-at-a-time
style. This means we can test most of external codegen machinery in one place without depending on
any target which may not be enabled in CI (eg TensorRT):
- Target instances are plumbed correctly so compile-time options are available.
- External modules are conveyed to the final export library.
- Constant bindings are conveyed to the metadata module.
* [BYOC] InlineCompilerFunctions helper pass
The TensorRT BYOC integration needs to 'undo' partitionings in some situations. Add an
InlineCompilerFunctions pass to make that robust. In particular, it must undo both the
'partitioning' (ie separating out the "Compiler" function) and any 'compositing' (ie separating
out small sub-graphs as "Composite" functions).
Fix misspelled nn.bias_add while there.
Note that the current implementation is broken but untested in CI. I have all the tests
fixed in a follow-up PR.
* - Lints
* - Only AOT compilation paths ensure "executor" is provided as a Target attribute.
In #11474 I got ready to switch CUTLASS from function-at-a-time to IRModule-at-a-time compilation.
However my approach didn't handle dynamic shape functions, so I adjust it here.
The idea is still that such passes will leave behind
calls to 'extern' functions. However, converting those
calls to 'call_lowered' form in
MarkCompilerFunctionsAsExtern is too soon since only
the TECompiler knows how to capture all the attributes
necessary to support dynamic shape functions.
So stop doing that in MarkCompilerFunctionsAsExtern and
instead support this case properly in the TECompiler.
While there try to chip away at the chronic lack of structure in te_compiler.cc. Every little bit helps.
Add a basic unit test.
* [BYOC] Two helper passes for external codegen using RelayToTIR custom pass machinery
(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).
For reasons explained in the above thread I'm moving CUTLASS to be IRModule-at-a-time external codegen
using a custom RelayToTIR pass instead of the traditional function-at-a-time external codegen using
a relay.ext.cutlass registered function. This means some of the rewriing done on-the-fly by LowerTEPass now
needs to be done by the custom pass directly. This PR supplies two passes which ease that burden:
- Before starting the CUTLASS-specific processing, make sure all "Compiler" attributed functions have
unique global definitions (ie are outlined). Though functions start in this form after BYOC partitioning,
under Graph and AOT compilation flows those functions are then inlined to pass through the 'codegen' keyhole
which assumes the whole model is just one self-contained main function. This pass will undo that. (I gave up
trying to just remove the inlining in the first place.)
- After the CUTLASS-specific processing the now compiled "Compiler" attributed functions need to marked as
'extern'. The te_compiler.cc uses the "ExternalSymbol" attribute for that, but since a) the symbol name
is never needed, on the presense of the attribute is significant downstream and b) "ExternalSymbol" is
easy to confuse with "global_symbol", I just replaced "ExternalSymbol" with "Extern" with an Integer(1)
(cf "Primitive").
The outlining pass is a little more general than necessary because it (will also) be used by Collage to
rewrite the IRModule into optimally partitioned form while making maximal reuse of partition functions.
Hence the abstract GlobalSymbolCache.
* - Andrew's comments
Mark Shields