[Runtime] Serialization/Deserialization of runtime module

[sunggg]

This PR implements serialization and deserialization of such runtime module by registering .set_creator and .set_repr_bytes.
If runtime module is binary serializable (see #14406), its roundtrip is possible.
When the runtime module is not binary serializable, such as csource_module, it will raise an error.

cc. @tqchen

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @areusch _{See #10317 for details}

_{Generated by tvm-bot}

[tqchen]

This is not sufficient, we will need to collect the import tree, as in normal export, so not only module itself get exported, but also the imported modules

[sunggg]

Thanks for the catch! This is addressed by using https://github.com/apache/tvm/blob/main/src/runtime/dso_library.cc#L152 which also collects import tree.

[tqchen]

Thanks @sunggg. The main comment is to include import tree information so we can recover modules that imports each other(e.g. a host module that imports cuda).

We can try to refactor and bring out logic in https://github.com/apache/tvm/blob/main/src/target/codegen.cc#L63 to achieve that goal. Also we need to check if the modules are all serializable and report error accordingly if they are not

[sunggg]

Thanks for the valuable feedback, @tqchen!
I found the existing support for serialization/deserialization of runtime module (including its import tree) that leverages python-side export_library mechanism.

To check if the modules are all exportable, it leverages the check introduced by #14406.
The example of the error message is as below:

AssertionError: Module RelayGraphRuntimeCodegenModule should be either dso exportable or binary serializable.

[sunggg]

@tvm-bot rerun

[sunggg]

Thank you, @tqchen for the suggestion. By reflecting your feedback, the latest commit cleans-up the APIs as follows:

// support full roundtrip of `runtime::Module` if the root and its all submodules are binary serializable
// These are used for structural hash of `runtime::Module` and thus also used for `SaveJson` and `LoadJson`
// Please note that `SerializeModuleToBytes` supports `include_dso` to keep the existing usage (e.g., `PackImportsToC`). 
// When `include_dso` is disabled and it encounters `DSOModules`, error will be raised.  
std::string SerializeModuleToBytes(const runtime::Module& mod, bool include_dso = true);
runtime::Module DeserializeModuleFromBytes(std::string blob);

// supports export/import of `runtime::Module` if the root and all submodules are either binary serializable or dso exportable
// this is currently used for pytorch integration: https://github.com/apache/tvm/pull/12232
std::string ExportModuleToBase64(tvm::runtime::Module module);
runtime::Module ImportModuleFromBase64(std::string base64str);

Whenever each API faces module/submodule that is not expected, the error will be triggered.

Also, SHashReduce and SEqualReduce are defined to enable structural equality check.
Thus, new test-cases verify roundtrippability via structural equality check.

[tqchen]

Thanks @sunggg some followup comments

[tqchen]

Some final nits and we are good to go

[tqchen]

Move these two functions to callers, since they are not necessarily the optimal way of decoding b64. The b64Stream is the preferred way. So avoid making them in support(and gradually we move the contrib impl over, or keep them self-contained to a sub module scope) is better

[sunggg]

Ah, I see. I moved them to contrib side (src/contrib/torch/tvm_module_wrapper/RuntimeModuleWrapperTVM.cc) since it is the only user currently.

[junrushao]

Hey just wanted to chime in a little bit - are we on the right track to get it in? Seems an important PR

[tqchen]

Thanks for the ping, this is ready and merged