How to create onnx module codegen?

[smallcoscat]

Refer to TVM tutorial Bring Your Own Codegen To TVM, which details how to create a self-defined c source module codegen.

However, ONNX is not a C source module, we should define an ONNX module node for codegen additionally. The following is the steps to create ONNX module codegen, and also applies to other external module codegen.

1. Create ONNXModuleCodegen to traverse relay IRModule and convert Relay op to ONNX op. (see codegen.cc)

   class ONNXModuleCodegen {
   public:
   ONNXModuleCodegen(){}
   runtime::Module CreateONNXModule(const ObjectRef& ref){
   
   auto mod = Downcast<IRModule>(ref);
   String codes = (*to_onnx_)(mod);
   /* Here, use String instead of std::string, because some byte info
    * will lost while calling PackedFunc object.
    */
   const auto* pf = runtime::Registry::Get("runtime.ONNXModuleCreate");
   CHECK(pf != nullptr) << "Cannot find onnx module to create the external runtime module";
   return (*pf)(codes, "onnx");
   }
   private:
   /*!
   * \brief The python function to convert relay module to onnx module.
   * \return byte array -> String
   */
   const PackedFunc* to_onnx_ = runtime::Registry::Get("tvm.relay.converter.to_onnx");
   };

   Register a global function, "relay.ext.onnx", whose body is a wrapper function to create onnx module.

   runtime::Module ONNXCompiler(const ObjectRef& ref) {
   ONNXModuleCodegen onnx;
   return onnx.CreateONNXModule(ref);
   }
   TVM_REGISTER_GLOBAL("relay.ext.onnx").set_body_typed(ONNXCompiler);

   Instead of writing op conversions in C++, use _registerfunc to register a global function "tvm.relay.converter.to_onnx", and write op conversions in Python to convert relay module to onnx module. (see converter/onnx.py)

   @tvm.register_func("tvm.relay.converter.to_onnx")
   def convert_to_onnx(model):
   ...
   opset = onnx.defs.onnx_opset_version() # get the supported opset version
   data = ""
   global_vars = model.get_global_vars()
   for i, global_var in enumerate(global_vars):
       func = model[global_var]
       sub_model = tvm.IRModule().from_expr(func.body)
       sub_model = fuse_ops(sub_model)
       func = sub_model["main"]
       graph_name = global_var.name_hint
       # Traverse the Relay function and record the nodes.
       sub_onnx_model = ONNXGenerator({}, opset, graph_name, "").to_onnx(func)
       bytes_data = get_onnx_bytes(sub_onnx_model)
       data += graph_name +"<"+ str(bytes_data)[2:-1]+">"; 
   return data

2. Create ONNXModuleNode which is subclass of ModuleNode to create a specific runtime module. (see source_module.cc)

   class ONNXModuleNode : public runtime::ModuleNode {
   public:
   ONNXModuleNode(std::string code,
                  std::string fmt)
     : code_(code), fmt_(fmt) {}
   const char* type_key() const {
   return "onnx";
   }
   PackedFunc GetFunction(
     const std::string& name,
     const ObjectPtr<Object>& sptr_to_self) final {
   LOG(FATAL) << "ONNX Source module cannot execute, to get executable module"
              << " build TVM with \'" << fmt_ << "\' runtime support";
   return PackedFunc();
   }
   
   std::string GetSource(const std::string& format) final {
   return code_;
   }
   ...
   void SaveToFile(const std::string& file_name,
                 const std::string& format) final {
   std::string fmt = GetFileFormat(file_name, format);
   std::string folder;
   size_t pos = file_name.find_last_of("\\/");
   if(pos!=std::string::npos){
       folder = file_name.substr(0,pos+1);
   }else{
       folder = file_name+"/";
   }
   auto datas = Split(code_,'>');
   if (fmt == "onnx") {
     CHECK_NE(code_.size(), 0);
     std::stringstream ss;
     for(auto data : datas){
         auto split_data = Split(data,'<');
         ss<<folder<<split_data[0].c_str()<<"."<<fmt;
         SaveBinaryToFile(ss.str(), ConvertEscape(split_data[1]));
         ss.str("");
         ss.clear();
   
     }
   } else {
     CHECK_EQ(fmt, fmt_)
         << "Can only save to format=" << fmt_;
   }
   }
   ...
   runtime::Module ONNXModuleCreate(String code, std::string fmt) {
   /* Here, use String instead of std::string, because some byte info
    * will lost while calling PackedFunc object.
    */
   auto n = make_object<ONNXModuleNode>(code, fmt);
   return runtime::Module(n);
   }
   TVM_REGISTER_GLOBAL("runtime.ONNXModuleCreate")
   .set_body_typed(ONNXModuleCreate);

3. Create a cmake file for ONNX, and include cmake in CMakeLists.txt. When user wants to use onnx module codegen, set _USE_ONNXCODEGEN "ON", and build TVM source code. (see ONNX.cmake)

   if(USE_ONNX_CODEGEN STREQUAL "ON")
   file(GLOB ONNX_RELAY_CONTRIB_SRC src/relay/backend/contrib/onnx/codegen.cc)
   list(APPEND COMPILER_SRCS ${ONNX_RELAY_CONTRIB_SRC})
   message(STATUS "Build with ONNX codegen.")
   endif()

[AugustineYang]

Hi. Really appreciate your work 'cause we are trying to do the exactly same thing. But I'm wondering if the project was modified after this issue, 'cause I cannot find the function converter_to_onnx() in the onnx.py. And other functions mentioned in this issue seem different with the ones in the latest files of the branch relay2onnx.