[翻译]在TVM中加入自己的Codegen — tvm 0.7.dev1文档

随着深度学习工作负载所针对的硬件设备的数量不断增加,用户在各种设备上实现高性能所需的知识也在不断增加。为了使数据科学家不必担心开发新模型时的性能,硬件后端提供程序要么为MKLDNN或cuDNN等库提供许多常用的深度学习运算符,要么提供诸如TensorRT的框架以使用户以某种方式描述他们的模型实现高性能。但是,用户尝试在新的库或设备上工作时必须学习新的编程界面。结果,对于1)让所有用户和硬件后端提供者站在同一页面上,对统一编程接口的需求变得越来越重要。在此开发人员指南中,我们演示了作为硬件后端提供程序的您如何轻松实现自己的代码生成并将其注册为Relay后端编译器以支持您的硬件设备/库。本指南根据您需要的不同图形表示形式涵盖两种类型的代码生成器:1.您要生成C代码。如果您的硬件已经具有经过优化的C / C ++库,例如从Intel CBLAS / MKL到CPU和从NVIDIA CUBLAS到GPU,那么这就是您所需要的。幸运的是,C源代码模块与TVM运行时模块完全兼容,这意味着生成的代码可以由具有适当编译标志的任何C / C ++编译器进行编译,因此您唯一的任务就是实现一个为子图生成C代码的代码生成器和一个C源模块以集成到TVM运行时模块中。在下一节中,我们将演示如何为您的硬件实现C代码生成器。2.您想生成任何其他图形表示。您的硬件可能需要其他形式的图形表示形式,例如JSON。在这种情况下,您不仅需要实现代码生成,还需要实现自定义的TVM运行时模块,以使TVM运行时知道应如何执行此图形表示。如果您已经为硬件配备了完整的图形执行引擎,例如用于GPU的TensorRT,则可以考虑采用这种解决方案。在完成代码生成和运行时之后,您可以让客户使用您的自定义标签来注释他们的模型以使用它们。最终用户注释和启动特定代码生成的教程在此处(TBA)。## 实施C代码 在这一部分中,我们演示了如何使用预实现的运算符函数来实现生成C代码的代码生成。为简化起见,我们的示例代码生成器不依赖于第三方库。相反,我们在C中手动实现了两个宏:#define CSOURCEBINARY_OP_1D(p_ID, pOP, pDIM1) \extern “C” void pID(float a, float b, float out) { \for (int64t i = 0; i < p_DIM1; ++i) { \out[i] = a[i] pOP b[i]; } }
#define CSOURCEBINARY_OP_2D(p_ID, pOP, pDIM1, pDIM2) \extern “C” void pID(float
a, float b, float out) { \for (int64t i = 0; i < p_DIM1; ++i) { \for (int64t j = 0; j < p_DIM2; ++j) { \int64t k = i * p_DIM2 + j; \out[k] = a[k] pOP b[k]; } } }使用这两个宏,我们可以为一维和二维张量生成二进制运算符。例如,给定一个子图如下。假设所有输入都是二维张量,其形状为(10,10)。ccompiler_input0|add<—c_compiler_input1|subtract<—c_compiler_input2|multiply<—c_compiler_input3|out我们的目标是生成以下可编译代码以执行子图:#include#include#include#include#include#include
#define GCC_BINARY_OP_1D(p_ID
, pOP, pDIM1) \extern “C” void pID(float a, float b, float out) { \for (int64t i = 0; i < p_DIM1; ++i) { \out[i] = a[i] pOP b[i]; } }
#define GCCBINARY_OP_2D(p_ID, pOP, pDIM1, pDIM2) \extern “C” void pID(float
a, float b, float out) { \for (int64t i = 0; i < p_DIM1; ++i) { \for (int64t j = 0; j < p_DIM2; ++j) { \int64t k = i * p_DIM2 + j; \out[k] = a[k] pOP b[k]; } } }
// Note 1GCCBINARY_OP_2D(gcc_0_0,*,10,10);GCC_BINARY_OP_2D(gcc_0_1,-,10,10);GCC_BINARY_OP_2D(gcc_0_2,+,10,10);
// Note 2extern”C”voidgcc_0
(floatgcc_input0,floatgccinput1,floatgcc_input2,floatgcc_input3,floatout){floatbuf_0=(float)malloc(4100);floatbuf_1=(float)malloc(4100);gcc_0_2(gcc_input0,gcc_input1,buf_0);gcc_0_1(buf_0,gcc_input2,buf_1);gcc_0_0(buf_1,gcc_input3,out);free(buf_0);free(buf_1);}
// Note 3extern”C”intgcc_0_wrapper(DLTensor
arg0,DLTensorarg1,DLTensorarg2,DLTensorarg3,DLTensorout){gcc_0
(static_cast(arg0->data),static_cast(arg1->data),static_cast(arg2->data),static_cast(arg3->data),static_cast(out->data));return0;}TVM_DLL_EXPORT_TYPED_FUNC(gcc_0,gcc_0_wrapper);在这里,我们突出显示上面代码中标记的注释:

  • 注意1是子图中三个节点的功能实现。
  • 注意2是通过分配中间缓冲区并调用相应函数来执行子图的函数。
  • 注3是TVM运行时兼容的包装函数。它接受输入张量和一个输出张量的列表(最后一个参数),将它们转换为正确的数据类型,并调用注释2中描述的子图函数。此外,它TVM_DLL_EXPORT_TYPED_FUNC是一个TVM宏,它生成另一个函数gcc_0并使该函数统一通过将所有张量打包到来进行论证TVMArgs。结果,TVM运行时可以直接调用gcc_0以执行子图,而无需付出额外的努力。生成上述代码后,TVM可以将其与图形的其余部分一起编译,并导出单个库以进行部署。 在本节的其余部分,我们将逐步实现一个codegen以生成上述代码。您自己的代码生成器必须位于src/relay/backend/contrib/<your-codegen-name>/。在我们的示例中,我们将代码生成命名为“ codegenc”,并将其放在。随时检查此文件以获取完整的实现。具体来说,我们将在此文件中实现两个类,这是它们之间的关系:subgraphsubgraphTVMbackend——————————————->CSourceCodegen——————->CodegenC^|^|||||————————————————————————————————generatedCsourceruntimemodulegeneratedCcode当TVM后端在中继图中找到函数(子图)时,使用已注册的编译器标记进行注释(ccompiler在此示例中),TVM后端将调用CSourceCodegen并传递该子图。CSourceCodegen的成员函数CreateCSourceModule将1)为子图生成C代码,以及2)将生成的C代码包装到C源运行时模块以供TVM后端编译和部署。特别是,C代码生成对于CodegenC该类是透明的,因为它提供了许多有用的实用程序来简化代码生成的实现。以下各节将以自底向上的顺序实现这两个类。### 实施CodegenC 在中src/relay/backend/contrib/codegen_c/codegen.cc,我们首先在以下名称空间下创建一个codegen类骨架tvm.relay.contrib:#include#include#include#include#include
    #include#include
    #include”codegen_c.h”
    namespacetvm{namespacerelay{namespacecontrib{
    classCodegenC:publicExprVisitor,publicCodegenCBase{public:explicitCodegenC(conststd::string&id){this->ext_func_id
    =id;}
    voidVisitExpr(constVarNode*node){;}voidVisitExpr(constCallNodecall)final{;}std::stringJIT(){;}
    private:/
    ! \brief The function id that represents a C source function. /std::stringextfunc_id=””;/! \brief The index of a wrapped C function. /intfunc_idx=0;/! \brief The index of allocated buffers. /intbufidx=0;/! \brief The arguments of a C compiler compatible function. /std::vectorextfunc_args;/! \brief The statements of a C compiler compatible function. /std::vectorext_func_body;/! \brief The declaration statements of a C compiler compatible function. /std::vectorfuncdecl;/! \brief The declaration statements of buffers. /std::vectorbufdecl;/! \brief The name and index pairs for output. /std::vector>out;}的CodegenC类继承两个类:ExprVisitor提供能力横动子图,并收集所需的信息并生成子图的功能,例如`gcc_0;CodegenCBase提供了生成包装函数的功能和实用程序,例如gcc0`上面的示例。可以看出,我们只需要在此codegen类中实现三个功能即可使其工作。#### 运营商代码生成 我们首先实施。遍历子图时,此函数访问所有调用节点。每个呼叫节点都包含一个我们要卸载到您的硬件的运算符。结果,我们需要按照拓扑顺序使用正确的运算符生成相应的C代码。我们按以下步骤逐步实现此功能。VisitExpr(constCallNodecall)1.生成函数声明结果示例:GCCBINARY_OP_2D(gcc_0_0,,10,10);如上所示,要生成函数声明,我们需要1)函数名称(例如gcc_0_0),2)运算符的类型(例如`)和3)输入张量形状(例如)。幸运的是,可以从以下位置轻松获取此信息:(10,10)CallNode`std::ostringstreammacro_stream;std::ostringstreamdecl_stream;std::ostringstreambuf_stream;
    // Generate a unique function name you like.std::stringfunc_name=ext_func_id
    +”“+std::to_string(func_idx++);
    // Make function declaration string.macro_stream<<”CSOURCE_BINARY_OP
    “<args.size()<<”D(“<// Check the operator type.if(IsOp(call,”add”)){macrostream<<”+”;}elseif(IsOp(call,”subtract”)){macro_stream<<”-“;}elseif(IsOp(call,”multiply”)){macro_stream<<”*”;}else{LOG(FATAL)<<”Unrecognized op”;}
    // Extract the input tensor shape.autoin_shape=GetShape(call->args[0]->checked_type());for(size_ti=0;i<in_shape.size();++i){macro_stream<<”, “<<in_shape[i];}macro_stream<<”);”;func_decl
    .pushback(macro_stream.str());可以看出,我们将生成的代码推送给类成员变量`func_decl。这意味着在完成遍历整个子图之后,我们已经收集了所有必需的函数声明,而我们唯一需要做的就是让它们由GCC进行编译。其余的实现也遵循此概念。VisitExpr_(constCallNode*call)**2.生成函数调用**结果示例:gcc_0_0(buf_1,gcc_input3,out);生成函数声明后,我们需要生成具有正确输入和输出的函数调用。要知道调用此函数时应放置哪些输入或缓冲区,我们必须访问其参数:boolfirst=true;decl_stream<<func_name<<"(";for(size_ti=0;i<call->args.size();++i){VisitExpr(call->args[i]);// Note 1for(autoout:out_){if(!first){decl_stream<<", ";}first=false;decl_stream<<out.first;}}// Note 2同样,我们要突出显示以上代码中的注释:**注意1**:VisitExpr(call->args[i])是递归调用,以访问当前函数的参数。参数可以是另一个节点的输出或输入张量。在示例实现中,我们确保每个节点out`在离开访问者之前都更新一个类变量。这是一个例子:arg_nodearg_node<-Visitarg(Note1)arg_node|||curr_node<-Processcurr_nodecurr_node<-Put”buf_0”asaninputbuffer
    (a)out
    ={}(b)out={}(c)out={(“buf0”,20)}我们可以在上图中看到,类变量`out在访问参数节点之前为空,并填充了输出缓冲区名称和大小argnode。结果,当我们完成访问参数节点时,通过查看,我们知道应该放置适当的输入缓冲区out。您将out`在本节末尾和下一节中找到我们的更新方式。注意2:您可能会注意到,在此步骤中我们没有关闭函数调用字符串。当前的函数调用字符串如下所示:。这是因为我们没有将最后一个参数(即输出)放入此调用。函数调用的输出可以是分配的临时缓冲区,也可以是子图输出张量。为了简化起见,在此示例中,我们为每个调用节点分配一个输出缓冲区(下一步),并将结果从最后一个缓冲区复制到输出张量。gcc_0_0(buf_1,gcc_input33.生成输出缓冲区结果示例:floatbuf_0=(float)malloc(4*100);如上一步所述,除了子图输入和输出张量之外,我们可能还需要缓冲区来保留中间结果。为了生成缓冲区,我们提取形状信息以确定缓冲区的类型和大小:// This example only supports single output.autotype_node=call->checked_type().as();CHECK(type_node!=nullptr&&runtime::TypeMatch(type_node->dtype,kDLFloat,32))<<”Only support single output tensor with float type”;
    // Generate a unique buffer name.std::stringout=”buf
    “+std::tostring(buf_idx++);
    // Extract the shape to be the buffer size.autooutshape=GetShape(call->checked_type());intout_size=1;for(size_ti=0;i// Make the buffer allocation and push to the buffer declarations.buf_stream<<”float “<<out<<” = (float)std::malloc(4 * “<<out_size<<”);”;buf_decl.pushback(buf_stream.str());分配输出缓冲区后,我们现在可以关闭函数调用字符串,并将生成的函数调用推送到类变量ext_func_body。decl_stream<<”, “<<out<<”);”;ext_func_body.push_back(decl_stream.str());4.更新输出缓冲区为了让接受当前调用节点的输出作为其输入的下一个节点知道它应该使用哪个缓冲区,我们需要`out在离开此访问函数之前更新class变量:out_.clear();out_.push_back({out,out_size});恭喜你!我们已经完成了本课中最困难的功能。在接下来的两节中,我们只需要组成此函数中的一些次要缺失部分。#### 输入变量的代码生成[](https://docs.tvm.ai/dev/relay_bring_your_own_codegen.html#code-generation-for-input-variables) 回想一下,我们通过访问调用节点的参数来收集输入缓冲区信息(上一节的第二步),并处理了其参数是另一个调用节点的情况(第四步)。在本节中,我们VarNode以示例为例演示如何处理其他节点。VarNode表示模型中的输入张量。它唯一的,但重要的信息是出了名的提示(如dataweight等)。在访问时VarNode,我们只需更新类变量out`以传递名称提示,以便后代调用节点可以生成正确的函数调用。voidVisitExpr(constVarNodenode){extfunc_args.pushback(node->name_hint());out.clear();out_.push_back({node->name_hint(),0});}请注意,在此示例中,我们假设要卸载的子图仅具有调用节点和变量节点。如果子图包含其他类型的节点,例如TupleNode,则还需要访问它们并绕过输出缓冲区信息。#### 发出代码 该codegen类的最后一部分是一个JIT函数,该函数为子图发出C函数,并将我们刚生成的C代码用作函数体。请记住,除了在上一节中生成的子图函数外,我们还需要一个包装器函数,该函数具有统一的参数,TVM运行时可以调用和传递数据。幸运的是,我们继承的基类已经提供了实现JitImpl来生成函数。例如,我们可以调用JitImpl如下:JitImpl(“gcc_0”/ Subgraph symbol (ID) /,{“gcc_input0”,”gcc_input1”,”gcc_input2”,”gcc_input3”}/ Input arguments /,{“float buf_0 = (float)malloc(4 20)”,…}/ Buffer allocations /,{“gcc_0_2(gcc_input0, gcc_input1, buf_0);”}/ Function body /,{“out”}/ Output /);上面的调用将生成三个函数(一个来自TVM包装器宏):
  1. 子图函数gcc_0_(在函数名称的末尾还有一个下划线),其中包含我们生成的所有C代码以执行子图。
  2. 包装函数gcc_0__wrapper_带有DLTensor参数列表,这些参数将数据转换为正确的类型并调用gcc_0_
  3. gcc_0具有TVM统一函数参数的TVM运行时兼容函数可解压缩TVM打包的张量并调用gcc_0__wrapper_。 因此,实现过程中唯一需要做的JIT就是将我们生成的所有子图函数代码传递给JitImpl:std::stringJIT(){// Write function macrosfor(autodecl:funcdecl){codestream<// Use GenCFunc to generate the C code and wrap it as a C source module.runtime::ModuleCreateCSourceModule(constNodeRef&ref)override{;}
    private:std::ostringstreamcodestream;};#### 实施GenCFunc GenCFunc只需使用CodegenC我们刚刚实现的遍历Relay函数(子图)并获得生成的C代码即可。内置函数在Relay函数中GetExtSymbol检索唯一的符号名称(例如gcc_0),我们必须将其用作C函数名称,因为该符号将用于DSO运行时查找。voidGenCFunc(constFunction&func){CHECK(func.defined())<<”Input error: expect a Relay function.”;
    // Record the external symbol for runtime lookup.autosid=GetExtSymbol(func);
    CodeGenCbuilder(sid);builder.VisitExpr(func->body);codestream<\n";code_stream_<<"#include \n";code_stream_<<"#include \n";code_stream_<<"#include \n";code_stream_<<"#include \n";code_stream_<<"#include \n";code_stream_<<"#include \n";
    // Append some common macro for operator definition.constcharoperatormacro=R”op_macro(#define CSOURCE_BINARY_OP_1D(p_ID, pOP, pDIM1) \extern “C” void pID(float a, float b, float out) { \for (int64t i = 0; i < p_DIM1; ++i) { \out[i] = a[i] pOP b[i]; } }
    #define CSOURCEBINARY_OP_2D(p_ID, pOP, pDIM1, pDIM2) \extern “C” void pID(float a, float b, float out) { \for (int64t i = 0; i < p_DIM1; ++i) { \for (int64t j = 0; j < p_DIM2; ++j) { \int64_t k = i pDIM2 + j; \out[k] = a[k] pOP b[k]; } } })opmacro”;
    code_stream
    <// Generate C code for the subgraph.if(ref->IsInstance()){GenCFunc(Downcast(ref));}elseif(ref->IsInstance()){relay::Modulemod=Downcast(ref);for(constauto&it:mod->functions){GenCFunc(Downcast(it.second));}}else{LOG(FATAL)<<”The input ref is expected to be a Relay function or module”<<”\n”;}
    // Create a CSourceModuleconstautopf=runtime::Registry::Get(“module.csource_module_create”);CHECK(pf!=nullptr)<<”Cannot find csource module to create the external runtime module”;return(pf)(codestream.str(),”cc”);}### 注册您的Codegen 最后一步是将您的代码生成器注册到TVM后端。我们首先实现一个简单的函数来调用我们的代码生成器并生成一个运行时模块。runtime::ModuleCCompiler(constNodeRef&ref){CSourceCodegencsource;returncsource.CreateCSourceModule(ref);}最后,我们将此功能注册到TVM后端:TVMREGISTER_GLOBAL(“relay.ext.ccompiler”).set_body_typed(CCompiler);这里ccompiler是一个定制的标签,让TVM知道这是它应该使用生成当子标注有卸载子图的代码生成ccompiler。最后,一个好的做法是设置CMake配置标志,使其仅为客户提供编译器。我们首先创建一个cmake文件cmake/modules/contrib/CODEGENC.cmake:if(USE_CODEGENC)file(GLOBCSOURCE_RELAY_CONTRIB_SRCsrc/relay/backend/contrib/codegen_c/codegen.cc)list(APPENDCOMPILER_SRCS${CSOURCE_RELAY_CONTRIB_SRC})endif(USE_CODEGENC)这样,用户可以在配置TVM时使用config.cmake以下命令配置是否包括编译器:set(USE_CODEGENCON)## 为您的表示实现代码生成 尽管我们已经演示了如何实现C代码源,但是您的硬件可能需要其他形式的图形表示形式,例如JSON。在这种情况下,您可以修改CodegenC我们已经实现的类以生成自己的图形表示,并实现定制的运行时模块,以使TVM运行时知道应如何执行该图形表示。为了简化,我们在本指南中定义了一个名为“ ExampleJSON”的图形表示。ExampleJSON并不意味着真正的JSON,而仅仅是没有控制流的图形的简单表示。例如,假设我们有一个名为的子图subgraph_0:input0|add<—input1|subtract<—input2|multiply<—input3|out然后,该子图的ExampleJON如下所示:subgraph_0input 0 10 10input 1 10 10input 2 10 10input 3 10 10add 4 inputs: 0 1 shape: 10 10sub 5 inputs: 4 2 shape: 10 10mul 6 inputs: 5 3 shape: 10 10的input关键字声明它的ID和形状的输入张量;其他语句则以语法描述计算。inputs:[inputID]shape:[shape]在本节中,我们的目标是实现以下定制的TVM运行时模块以执行ExampleJSON图。runtime::ModuleExampleJsonCompiler(constNodeRef&ref){ExampleJsonCodeGencodegen(ref);std::stringcode=codegen.gen();// Note 1constautopf=runtime::Registry::Get(“module.examplejson_module_create”);// Note 2CHECK(pf!=nullptr)<<”Cannot find ExampleJson module to create the external runtime module”;return(pf)(code);}TVM_REGISTER_GLOBAL(“relay.ext.examplejsoncompiler”).set_body_typed(ExampleJsonCompiler);注意1:我们稍后将实现自定义代码生成,以通过生成子图来生成ExampleJSON代码字符串。注2:此行获得指向用于创建定制运行时模块的函数的指针。您会看到它采用了我们刚刚生成的ExampleJSON格式的子图代码,并初始化了运行时模块。在以下各节中,我们将介绍1)如何实现ExampleJsonCodeGen和2)如何实现和注册examplejson_module_create。### 实现ExampleJsonCodeGen 类似于C代码生成器,我们还衍生ExampleJsonCodeGenExprVisitor了利用访问者模式进行子图遍历的方法。另一方面,我们不需要继承,CodegenCBase因为我们不需要TVM C ++包装器。codegen类的实现如下:#include#include#include#include#include
    #include#include
    namespacetvm{namespacerelay{namespacecontrib{
    classExampleJsonCodeGen:publicExprVisitor{public:explicitExampleJsonCodeGen();
    // Note 1voidVisitExpr
    (constVarNodenode){/ Skip in this example. /}voidVisitExpr_(constCallNodecall)final{/ Skip in this example. /}
    // Note 2std::stringgen(NodeRef&ref){this->code=””;if(ref->IsInstance()){this->visit(Downcast(ref));}elseif(ref->IsInstance()){relay::Modulemod=Downcast(ref);for(constauto&it:mod->functions){this->visit(Downcast(it.second));}}else{LOG(FATAL)<<”The input ref is expected to be a Relay function or module”;}returnthis->code;}
    private:/! \brief The function id that represents a C source function. /std::stringcode;}注意1:我们再次实现相应的访问者函数,以生成ExampleJSON代码并将其存储到类变量中code(在此示例中,我们跳过了访问者函数的实现,因为它们的概念与C代码生成基本相同)。完成图访问之后,我们应该在中有一个ExampleJSON图code注意2:我们定义了一个内部APIgen来获取子图并生成ExampleJSON代码。该API可以采用您喜欢的任意名称。下一步是实施自定义运行时,以利用的输出ExampleJsonCodeGen。### 实施自定义的运行时 在本节中,我们将逐步实现自定义的TVM运行时并将其注册到TVM运行时模块。自定义的运行时应位于src/runtime/contrib/<your-runtime-name>/。在我们的示例中,我们将运行时命名为“ exampleext_runtime”,并将其放在。随时检查此文件以获取完整的实现。再次,我们首先定义一个自定义的运行时类,如下所示。该类必须从TVM派生ModuleNode,以便与其他TVM运行时模块兼容。#include#include#include#include#include#include#include#include
    #include#include#include#include#include#include
    namespacetvm{namespaceruntime{classExampleJsonModule:publicModuleNode{public:explicitExampleJsonModule(std::stringgraph_json);
    PackedFuncGetFunction(conststd::string&name,constObjectPtr&sptr_to_self)final;
    constchartype_key()const{return”examplejson”;}
    voidSaveToBinary(dmlc::Stream
    stream)final;
    staticModuleLoadFromBinary(voidstrm);
    staticModuleCreate(conststd::string&path);
    std::stringGetSource(conststd::string&format=””);
    voidRun(intid,conststd::vector&inputs,intoutput);
    voidParseJson(conststd::string&json);
    private:/
    \brief The json string that represents a computational graph. */std::stringgraph_json;/ \brief The subgraph that being processed. /std::stringcurrsubgraph;/! \brief A simple graph from subgraph id to node entries. /std::map>graph;/ \brief A simple pool to contain the tensor for each node in the graph. /std::vectordata_entry;/ \brief A mapping from node id to op name. /std::vectoropid;};特别是,ModuleNode我们必须在其中实现一些衍生的功能ExampleJsonModule
    • 构造函数:此类的构造函数应接受一个子图(以您的表示形式),以所需的任何格式对其进行处理和存储。保存的子图可由以下两个功能使用。
    • GetFunction:这是此类中最重要的功能。当TVM运行时要使用您的编译器标签执行子图时,TVM运行时会从您的自定义运行时模块调用此函数。它提供函数名称以及运行时参数,并且GetFunction应返回打包的函数实现以供TVM运行时执行。
    • SaveToBinaryLoadFromBinarySaveToBinary将运行时模块序列化为二进制格式,以供以后部署。用户使用export_libraryAPI时,TVM将调用此函数。另一方面,由于我们现在使用自己的图表示形式,因此必须确保LoadFromBinary能够通过采用生成的序列化二进制文件来构造相同的运行时模块SaveToBinary
    • GetSource(可选):如果您想查看生成的ExampleJSON代码,则可以实现此函数以将其转储;否则,您可以跳过实施。 其他功能和类变量将与上述必备功能的实现一起引入。#### 实施构造函数 explicitExampleJsonModule(std::stringgraphjson){this->graph_json=graphjson;ParseJson(this->graph_json);}然后,我们实现ParseJson解析为ExampleJSON格式的子图,并在内存中构造一个图供以后使用。由于在此示例中我们不支持带有分支的子图,因此我们仅使用数组按顺序存储子图中的每个节点。voidParseJson(conststd::string&json){std::stringline;std::stringcurrsubgraph;std::stringstreamss(json);
      while(std::getline(ss,line,’\n’)){std::stringstreamss2(line);std::stringtoken;intid=0;
      ss2>>token;if(token.find(“subgraph
      “)!=std::string::npos){currsubgraph=token;continue;}
      ss2>>id;if(op_id
      .size()<=staticcast(id)){op_id.resize(id+1);dataentry.resize(id+1);}
      int64ttotal_elements=1;std::vectorshape;if(token==”input”){int64_tsize=0;while(ss2>>size){total_elements*=size;shape.push_back(size);}}else{op_id[id]=token;// Note 1boolshapedata=false;NodeEntryentry;while(ss2>>token){if(token==”shape:”){shape_data=true;}elseif(shape_data){total_elements*=std::stoll(token);shape.push_back(std::stoll(token));}elseif(token!=”inputs:”){entry.inputs.push_back(std::stoi(token));}}entry.id=id;entry.output=id;graph[currsubgraph].push_back(entry);// Note 2}DLContextctx;ctx.device_type=static_cast(1);ctx.device_id=0;data_entry[id]=NDArray::Empty(shape,DLDataType{kDLFloat,32,1},ctx);// Note 3}}注意1:我们使用类变量op_id_将子图节点ID映射到运算符名称(例如add),以便我们可以在运行时调用相应的运算符函数。注意2:我们使用类变量graph_将子图名称映射到节点数组。GetFunction将在运行时通过子图ID查询图节点。注3:我们使用类变量dataentry将子图节点ID映射到张量数据占位符。我们将在运行时将输入和输出放入相应的数据条目。#### 实现GetFunction 构造后,我们应该准备好上述类变量。然后GetFunction,我们实现为TVM运行时提供可执行的子图函数:PackedFuncGetFunction(conststd::string&name,constObjectPtr&sptrto_self)final{if(this->graph.find(name)!=this->graph.end()){this->curr_subgraph=name;returnPackedFunc(sptr_to_self,this{
      // Copy input tensors to corresponding data entries.for(autoi=0;idataentry[i].CopyFrom(arg);}else{NDArrayarg=args[i];this->dataentry[i].CopyFrom(arg);}}
      // Execute the subgraph.for(constauto&it:this->graph[this->curr_subgraph]){this->Run(it.id,it.inputs,it.output);}CHECKGT(graph.count(this->currsubgraph),0U);
      // Copy the output from a data entry back to TVM runtime argument.autooutidx=graph[this->currsubgraph].back().output;if(args[args.size()-1].typecode()==kArrayHandle){DLTensor*arg=args[args.size()-1];this->data_entry[outidx].CopyTo(arg);}else{NDArrayarg=args[args.size()-1];this->data_entry[outidx].CopyTo(arg);}*rv=data_entry.back();});}else{LOG(FATAL)<<”Unknown subgraph: “<&inputs,intoutput){// Make a list data entry indexs.std::vectorargs(inputs.begin(),inputs.end());args.pushback(output);
      // Initialize data holders.std::vectorvalues(args.size());std::vectortype_codes(args.size());
      // Initialize a TVM arg setter with TVMValue and its type code.TVMArgsSettersetter(values.data(),type_codes.data());
      // Set each argument to its corresponding data entry.if(op_id
      [id]==”add”||opid[id]==”sub”||opid[id]==”mul”){for(sizeti=0;i<args.size();i++){setter(i,data_entry[args[i]]);}}
      // Invoke the corresponding operator function.if(opid[id]==”add”){Add(values.data(),typecodes.data(),args.size());}elseif(op_id[id]==”sub”){Sub(values.data(),typecodes.data(),args.size());}elseif(op_id[id]==”mul”){Mul(values.data(),typecodes.data(),args.size());}else{LOG(FATAL)<<”Unknown op: “<<op_id[id]<<”\n”;}}Run功能主要有两个部分。第一部分分配的列表TVMValue,并映射相应的数据输入块。这将成为我们运算符功能的参数。第二部分将调用我们的运算符函数。虽然我们使用相同的C函数与前面的例子,可以更换AddSub以及Mul用自己的引擎。您只需要确保引擎将结果存储到最后一个参数,就可以将其传输回TVM运行时。通过实现上述功能,我们定制的代码生成和运行时现在可以执行子图。最后一步是注册一个API(examplejson_module_create)以创建此模块:TVMREGISTER_GLOBAL(“module.examplejson_module_create”).set_body_typed({auton=make_object(code);returnruntime::Module(n);});#### 实现SaveToBinary和LoadFromBinary 到目前为止,我们已经实现了自定义运行时的主要功能,以便可以将其用作其他TVM运行时。但是,当用户要将已构建的运行时保存到磁盘以进行部署时,TVM不知道如何保存它。这就是我们要实现SaveToBinary和的原因LoadFromBinary,它们告诉TVM如何保留和还原此自定义运行时。我们首先实现SaveToBinary允许用户将该模块保存在磁盘中的功能。voidSaveToBinary(dmlc::Stream*stream)final{stream->Write(this->graph_json);}我们可以发现此功能非常简单。回想一下,我们在构造函数中使用的唯一参数是子图表示,这意味着我们只需要一个子图表示即可构造/恢复此自定义运行时模块。结果,SaveToBinary只需将子图写入输出DMLC流。也就是说,当用户使用export_libraryAPI导出模块时,自定义模块将是子图的ExampleJSON流。相似,LoadFromBinary读取子图流并重新构建自定义的运行时模块:staticModuleLoadFromBinary(voidstrm){dmlc::Streamstream=static_cast(strm);std::stringgraph_json;stream->Read(&graph_json);auton=tvm::runtime::make_object(graph_json);returnModule(n);}我们还需要注册此功能以启用相应的Python API:TVM_REGISTER_GLOBAL(“module.loadbinary_examplejson”).set_body_typed(ExampleJsonModule::LoadFromBinary);上面的注册意味着当用户调用tvm.runtime.load(lib_path)API并且导出的库具有ExampleJSON流时,我们LoadFromBinary将被调用以创建相同的自定义运行时模块。另外,如果要直接从ExampleJSON文件支持模块创建,还可以实现一个简单的函数并注册Python API,如下所示:staticModuleCreate(conststd::string&path){std::ifstreamfilep;filep.open(path,std::ios::in);std::stringgraph_json;std::stringline;while(std::getline(filep,line)){graph_json+=line;graph_json+=”\n”;}filep.close();auton=tvm::runtime::make_object(graph_json);returnModule(n);}
      TVM_REGISTER_GLOBAL(“module.loadfile_examplejson”).set_body({*rv=ExampleJsonModule::Create(args[0]);});这意味着用户可以手动编写/修改ExampleJSON文件,并使用Python API来构建自定义模块。tvm.runtime.load(“mysubgraph.examplejson”,”examplejson”)## 摘要 总而言之,这是一份清单供您参考:
    • 衍生自甲代码生成类ExprVisitorCodegenCBase(仅对于C代码生成)与以下功能。
    • VisitExpr_(constCallNode*call)收集呼叫节点信息。
    • 收集子图信息所需的其他访问者功能。
    • JIT生成子图代码。
    • 注册代码生成器。
    • 创建的函数CSourceModule(用于C代码生成)。
    • ModuleNode以下功能派生的运行时模块类(用于图形表示)。
    • 构造函数。
    • GetFunction生成TVM运行时兼容的PackedFunc
    • Run执行子图。
    • 注册运行时创建API。
    • SaveToBinaryLoadFromBinary序列化/反序列化自定义的运行时模块。
    • 注册LoadFromBinaryAPI以支持tvm.runtime.load(your_module_lib_path)
    • (可选)Create以从表示形式的子图文件支持定制的运行时模块构造。
    • 一个用于对用户中继程序进行注释的注释器,以利用您的编译器和运行时(TBA)。