junrushao1994 commented on a change in pull request #8802: URL: https://github.com/apache/tvm/pull/8802#discussion_r692565459
########## 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: Hey the `target`-to-string behavior is going to be deprecated in the future, per this RFC: https://discuss.tvm.apache.org/t/rfc-tvm-target-specification/6844. Is it possible to just annotate the target object onto the function? Thanks a lot! -- 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: commits-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org