[External CodeGen] Status of Annotating composite functions?

I was playing around with the new external codegen & composite functions and noticed this flow fails inside of AnnotateTarget. Wondering if it is because of this TODO

mod = tvm.IRModule.from_expr(f)
mod = MergeComposite(pattern_table)(mod)
mod = AnnotateTarget("my_external_compiler")(mod)
mod = PartitionGraph()(mod)
TVMError: Check failed: ref->template IsInstance<typename SubRef: :ContainerType>(): Downcast from relay.Function to relay.Op failed.

I’m interested in offloading some subgraphs to my external codegen, lets say Conv2d+bias+relu. Is this the right mechanism to use? I can annotate individual ops for external codegen if I skip MergeComposite, but after this I’m not sure how to merge the annotated ops.

Thanks

@zhiics @comaniac @mbaret

You should be able to use the op-based annotation now. The PR we merged today provides a pass to combine supported ops in a single subgraph: https://github.com/apache/incubator-tvm/pull/5134

AnnotateTarget doesn’t support composite functions at the moment. I intend to send a PR to resolve this very soon (hopefully this week). You can use MergeCompilerRegions if you like, but this will only be applicable if you also support conv2d, bias and relu individually as well as merged.

Thank you @comaniac and @mbaret. I will see if MergeCompilerRegions is something that will work for us, otherwise will keep an eye out for your PR.

Status update! I’ve put on the following two PRs which hopefully will allow for composite function annotation: 5261, 5262. Feel free to take a look.

1 Like

@adb I had a old PR https://github.com/apache/incubator-tvm/pull/4741 which demonstrates conv + bias + relu fusion in the “hard” way (before composite was introduced). I’ll send a new one ASAP after PRs by @mbaret are merged.

Hi @masahi thanks for bringing this to my attention. Looks like this PR could work for us too. As a first pass we hope to target the most common fusion patterns as in your PR.

Thanks @mbaret! Will be trying out some patterns today.

@adb The PR is up https://github.com/apache/incubator-tvm/pull/5272

1 Like