junrushao1994 commented on a change in pull request #8802:
URL: https://github.com/apache/tvm/pull/8802#discussion_r692664180
##########
File path: src/relay/backend/te_compiler.cc
##########
@@ -875,6 +878,121 @@ LoweredModule LowerTE(const IRModule& module, TargetMap
targets, DeviceMap devic
return lowered_module;
}
+IRModule LoweredModuleToIRModule(LoweredModule mod) {
+ Map<GlobalVar, BaseFunc> unified_funcs;
+ Map<GlobalTypeVar, TypeData> unified_type_defs;
+
+ // copy main module funcs to unified funcs (what target do we need to
annotate with here?)
+ for (const auto& kv : mod.main_module->functions) {
+ const GlobalVar& var = kv.first;
+ const BaseFunc& func = kv.second;
+ ICHECK(!func->IsInstance<tir::PrimFuncNode>());
+ unified_funcs.Set(var, func);
+ }
+
+ // copy the type definitions for the main module
+ for (const auto& kv : mod.main_module->type_definitions) {
+ const GlobalTypeVar& ty_var = kv.first;
+ const TypeData& ty_data = kv.second;
+ unified_type_defs.Set(ty_var, ty_data);
+ }
+ // Move functions in per target IRModule into unified module
+ // Also move the type definitions
+ for (const auto& kv : mod.per_target_module) {
+ const String target = kv.first;
+ const IRModule target_module = kv.second;
+ // Move the per module functions, and annotate the funcs with their target
+ for (const auto& kv : target_module->functions) {
+ const GlobalVar& var = kv.first;
+ const BaseFunc& func = kv.second;
+ ICHECK(func->IsInstance<tir::PrimFuncNode>())
+ << "We expect the target_module to contain only PrimFuncs at this
point, but got "
+ << func->GetTypeKey();
+ tir::PrimFunc primFunc =
WithAttr(Downcast<tir::PrimFunc>(std::move(func)), attr::kTarget,
+ runtime::String(target));
Review comment:
Sure I don't have strong opinion. Let me know if there is anything I
could help with :-)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
