Yes, you are right. Some key words are confusing.
I guess the module can be seen in some way as a hash map with func name as key and function implement as value.
The following code can be found in graph_runtime_codegen.cc.
LoweredOutput Codegen(relay::Function func) {
auto pf = GetPackedFunc("relay.backend.GraphPlanMemory");
storage_device_map_ = (*pf)(func);
// First we convert all the parameters into input nodes.
for (auto param : func->params) {
auto node_ptr = GraphInputNode::make_node_ptr(param->name_hint(), GraphAttrs());
var_map_[param.get()] = AddNode(node_ptr, param);
}
heads_ = VisitExpr(func->body);
std::ostringstream os;
dmlc::JSONWriter writer(&os);
GetJSON(&writer);
LoweredOutput ret;
ret.graph_json = os.str();
ret.params = params_;
for (auto& kv : lowered_funcs_) {
if (ret.lowered_funcs.count(kv.first) == 0) {
ret.lowered_funcs.Set(kv.first, IRModule::Empty());
}
auto& mod = ret.lowered_funcs[kv.first];
mod->Update(kv.second);
ret.lowered_funcs.Set(kv.first, mod);
}
ret.external_mods = compile_engine_->LowerExternalFunctions();
return ret;
}
struct LoweredOutput {
std::string graph_json;
Map<std::string, IRModule> lowered_funcs;
Array<tvm::runtime::Module> external_mods;
std::unordered_map<std::string, tvm::runtime::NDArray> params;
};
And the IRModule shown above may have something to do with parallelism. And IRModule is something just like a lower function but not a module.
struct GraphCodegen {
public:
GraphCodegen() {
auto pf = GetPackedFunc("relay.build_module._GraphRuntimeCodegen");
mod = (*pf)();
}
...
In build_module.cc, the lowered_funcs have been transformed into module.
class BuildModule(object):
"""Build an IR module to run on TVM graph runtime. This class is used
to expose the `RelayBuildModule` APIs implemented in C++.
"""
def __init__(self):
self.mod = _build_module._BuildModule()
...
def build(self, mod, target=None, target_host=None, params=None):
target = _update_target(target)
# Setup the params.
if params:
self._set_params(params)
# Build the IR module
self._build(mod, target, target_host)
# Get artifacts
graph_json = self.get_json()
mod = self.get_module()
params = self.get_params()
And when we invoke relay.build we would get graph_json, mod, params. We use graph_json, mod and ctx to create the graph runtime.
Maybe it is serial. Just copy from one device to another device and switch runtime.