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