Target codegen vs relax BYOC signatures

As from my understanting, relax BYOC codegen for foo is registered under relax.ext.foo and has signature:

Array<runtime::Module> FooCodegen(Array<Function> functions, Map<String, ObjectRef> options, Map<Constant, String> /*unused*/);

It gets called by the RunCodegen pass on functions with the codegen annotation set to foo.

Now I am also following the target/device integration documentation as I would like to be able to build model for my target for my new device by simply passing the target to tvm.build. The documentation states that I need to register the target codegen under target.build.foo with the signature:

runtime::Module GeneratorFooCode(IRModule mod, Target target);

These two are different, and I was hoping I could just reuse relax.ext.foo for target.build.foo but I am not sure how to get to a single runtime::Module from an Array<runtime::Module>. Also I am not sure whether this is the correct way for BYOC for new device to provide their code generation as this seems to leads to duplicate code in my example.

Am I missing something?

1 Like