其他分享
首页 > 其他分享> > 如何将自定义代码生成TVM

如何将自定义代码生成TVM

作者:互联网

如何将自定义代码生成TVM

如何将自定义代码生成TVM

本文参考链接:

https://tvm.apache.org/docs/dev/how_to/relay_bring_your_own_codegen.html

https://blog.csdn.net/weixin_42164269/article/details/104291635

简介

深度学习针对的硬件设备的数量不断增加,用户需要在各种设备上实现高性能所需的知识。硬件后端提供者要么提供像MKLDNN或cuDNN类的库,包含许多常用的深度学习运算符,要么提供诸如TensorRT这样的框架,用户以某种方式描述模型实现高性能。但是,用户尝试在新的库或设备上工作时,必须学习新的编程接口。结果,对统一编程接口的需求变得越来越重要。

1)让所有用户和硬件后端提供者站在同一页面上

2)提供一种可行的解决方案,允许专用硬件或库仅支持具有极高性能的广泛使用的运算符,但不支持的运算符回退到CPU / GPU等常规设备。

本文主要内容如下:

目录

简介

1. 生成C代码。

2. 生成任何其它图形表示。

实现一个C代码生成器

实现【CodegenC】

运算符代码生成

输入变量的代码生成

代码发送

实现【CSourceCodegen 】

实现【GenCFunc 】

实现【CreateCSourceModule 】

注册代码生成

实现一个代码生成表示

实现【ExampleJsonCodeGen 】

实现自定义runtime

实现构造函数

实现【GetFunction 】

实现运行

实现【SaveToBinary】和【LoadFromBinary 】

总结

在本开发人员指南中,演示了作为硬件后端提供者,如何轻松实现自定义代码生成,注册为Relay后端编译器,支持硬件设备/库。本文根据需要的不同图形表示形式,涵盖两种类型的代码生成器:

1. 要生成C代码。

如果硬件已经具有经过优化的C/C ++库,如对CPU拥有Intel CBLAS / MKL,GPU拥有NVIDIA CUBLAS,这就是所需要的。幸运的是,C源代码模块与TVM runtime模块完全兼容,生成的代码可以由具有适当编译标志的任何C / C ++编译器进行编译,唯一的任务就是实现一个为子图生成C代码的代码生成器和一个C源模块,集成到TVM runtime模块中。在下一节中,将演示如何为硬件实现C代码生成器。

2. 生成任何其它图形表示。

硬件可能需要其它形式的图形表示形式,如JSON。在这种情况下,不仅需要实现代码生成,还需要实现自定义的TVM runtime模块,使TVM runtime知道应如何执行图形表示。如果已经为硬件配备了完整的图形执行引擎,如用于GPU的TensorRT,可以考虑采用这种解决方案。

在完成代码生成和runtime后,可以让客户使用自定义标签注释模型使用。

实现一个C代码生成器

在这一部分中,演示如何实现使用预实现的运算符函数,生成C代码的代码生成器。简化起见,示例代码生成器不依赖于第三方库。相反,在C中手动实现了两个宏:

#define CSOURCE_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)         \

    extern "C" void p_ID_(float* a, float* b, float* out) { \

        for (int64_t i = 0; i < p_DIM1_; ++i) {             \

            out[i] = a[i] p_OP_ b[i];                       \

        }                                                   \

    }

 

#define CSOURCE_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \

    extern "C" void p_ID_(float* a, float* b, float* out) {   \

        for (int64_t i = 0; i < p_DIM1_; ++i) {               \

            for (int64_t j = 0; j < p_DIM2_; ++j) {           \

                int64_t k = i * p_DIM2_ + j;                  \

                out[k] = a[k] p_OP_ b[k];                     \

            }                                                 \

        }                                                     \

    }

使用这两个宏,可以为一维和二维张量,生成二进制运算符。如给定一个子图如下。假设所有输入都是二维张量,形状为(10,10)。

c_compiler_input0

       |

      add <-- c_compiler_input1

       |

    subtract <-- c_compiler_input2

       |

    multiply <-- c_compiler_input3

       |

      out

目标是生成以下可编译代码执行子图:

#include <tvm/runtime/c_runtime_api.h>

#include <tvm/runtime/packed_func.h>

#include <dlpack/dlpack.h>

#include <cstdint>

#include <cstring>

#include <iostream>

 

#define GCC_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)           \

  extern "C" void p_ID_(float* a, float* b, float* out) { \

    for (int64_t i = 0; i < p_DIM1_; ++i) {               \

      out[i] = a[i] p_OP_ b[i];                           \

    }                                                     \

  }

 

#define GCC_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \

  extern "C" void p_ID_(float* a, float* b, float* out) { \

    for (int64_t i = 0; i < p_DIM1_; ++i) {               \

      for (int64_t j = 0; j < p_DIM2_; ++j) {             \

        int64_t k = i * p_DIM2_ + j;                      \

        out[k] = a[k] p_OP_ b[k];                         \

      }                                                   \

    }                                                     \

  }

 

// Note 1

GCC_BINARY_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 2

extern "C" void gcc_0_(float* gcc_input0, float* gcc_input1,

                       float* gcc_input2, float* gcc_input3, float* out) {

  float* buf_0 = (float*)malloc(4 * 100);

  float* buf_1 = (float*)malloc(4 * 100);

  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 3

extern "C" int gcc_0_wrapper(DLTensor* arg0, DLTensor* arg1, DLTensor* arg2,

                             DLTensor* arg3, DLTensor* out) {

  gcc_0_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),

         static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),

         static_cast<float*>(out->data));

  return 0;

}

TVM_DLL_EXPORT_TYPED_FUNC(gcc_0, gcc_0_wrapper);

在这里,突出显示上面代码中标记的注释:

Note1,子图中三个节点的函数实现。

Note2,一个函数,通过分配中间缓冲区,调用相应函数执行子图。

Note3,TVM runtime兼容的包装函数。接受一个输入张量和一个输出张量的列表(最后一个参数),转换为正确的数据类型,调用Note2中描述的子图函数。此外,【TVM_DLL_EXPORT_TYPED_FUNC】是一个TVM宏,生成另一个函数【gcc_0】,【gcc_0】具有统一的函数参数,通过把所有的参数张量打包成【TVMArgs】。结果,TVM runtime可以直接调用gcc_0执行子图,无需付出额外的努力。使用上面生成的代码,TVM可以与图的其余部分一起编译,导出单个库以进行部署。

在本节的其余部分,将逐步实现一个codegen生成上述代码。自定义代码源必须位于src/relay/backend/contrib/<your-codegen-name>/。在示例中,将代码源命名为“ codegen_c”,将放在“此处<https://github.com/apache/incubator-tvm/blob/master/src/relay/backend/contrib/codegen_c/codegen.cc>下`_。可以随时检查文件,获取完整的实现。

在此文件中实现两个类,这是相互关系:

                     subgraph                                subgraph

TVM backend -----------------------------> CSourceCodegen -------------> CodegenC

       ^                                       |    ^                       |

       |                                       |    |                       |

       ----------------------------------------      ------------------------

          generated C source runtime module              generated C code

当TVM后端在Relay中找到一个函数(子图)时,使用已注册的编译器标记进行注释(【ccompiler】在此示例中),TVM后端将调用【CSourceCodegen】,转换该子图。【CSourceCodegen】的成员函数【CreateCSourceModule】将

1)为子图生成C代码,

2)将生成的C代码包装到C源runtime模块中,供TVM后端编译和部署。

特别地,C代码生成对于【CodegenC】类是透明的,因为提供了许多有用的实用程序,简化代码生成的实现。以下各节将以自底向上的顺序实现这两个类。

实现【CodegenC】

在src/relay/backend/contrib/codegen_c/codegen.cc中,先在【tvm.relay.contrib】名称空间下,创建一个代码生成类骨架:

#include <tvm/relay/expr_functor.h>

#include <tvm/relay/transform.h>

#include <tvm/relay/type.h>

#include <tvm/runtime/module.h>

#include <tvm/runtime/object.h>

 

#include <fstream>

#include <sstream>

 

#include "codegen_c.h"

 

namespace tvm {

namespace relay {

namespace contrib {

 

class CodegenC : public ExprVisitor, public CodegenCBase {

  public:

    explicit CodegenC(const std::string& id) { this->ext_func_id_ = id; }

 

    void VisitExpr_(const VarNode* node) { ; }

    void VisitExpr_(const CallNode* call) final { ; }

    std::string JIT() { ; }

 

  private:

    /*! \brief The function id that represents a C source function. */

    std::string ext_func_id_ = "";

    /*! \brief The index of a wrapped C function. */

    int func_idx = 0;

    /*! \brief The index of allocated buffers. */

    int buf_idx_ = 0;

    /*! \brief The arguments of a C compiler compatible function. */

    std::vector<std::string> ext_func_args_;

    /*! \brief The statements of a C compiler compatible function. */

    std::vector<std::string> ext_func_body;

    /*! \brief The declaration statements of a C compiler compatible function. */

    std::vector<std::string> func_decl_;

    /*! \brief The declaration statements of buffers. */

    std::vector<std::string> buf_decl_;

    /*! \brief The name and index pairs for output. */

    std::vector<std::pair<std::string, int>> out_;

}

【CodegenC】类继承两个类:【ExprVisitor】提供遍历子图,收集所需的信息并生成子图的功能的能力,如【gcc_0_】; 【CodegenCBase】提供了生成包装函数的功能和用法,如gcc_0上面的示例。可以看出,只需要在此codegen类中,实现三个函数即可工作。

运算符代码生成

先实现【VisitExpr_(const CallNode* call)】。遍历子图时,此函数访问所有调用节点。每个调用节点都包含一个要卸载到硬件上的运算符。结果,需要按照拓扑顺序使用正确的运算符,生成相应的C代码。按以下步骤逐步实现此功能。

1. 生成函数声明

结果示例:【GCC_BINARY_OP_2D(gcc_0_0, *, 10, 10);】

如上所示,要生成函数声明,需要

1)函数名称(例如gcc_0_0)

2)运算符的类型(如*)

3)输入张量形状(如(10, 10))。

幸运的是,可以从【CallNode】位置轻松获取此信息:

 

std::ostringstream macro_stream;

std::ostringstream decl_stream;

std::ostringstream buf_stream;

 

// Generate a unique function name you like.

std::string func_name = ext_func_id_ + "_" + std::to_string(func_idx++);

 

// Make function declaration string.

macro_stream << "CSOURCE_BINARY_OP_" << call->args.size() << "D(" << func_name << ", ";

 

// Check the operator type.

if (IsOp(call, "add")) {

  macro_stream << "+";

} else if (IsOp(call, "subtract")) {

  macro_stream << "-";

} else if (IsOp(call, "multiply")) {

  macro_stream << "*";

} else {

  LOG(FATAL) << "Unrecognized op";

}

 

// Extract the input tensor shape.

auto in_shape = GetShape(call->args[0]->checked_type());

for (size_t i = 0; i < in_shape.size(); ++i) {

  macro_stream << ", " << in_shape[i];

}

macro_stream << ");";

func_decl_.push_back(macro_stream.str());

可以看出,将生成的代码放到类成员变量【func_decl_】。在完成遍历整个子图后,已经收集了所有必需的函数声明,唯一需要做的就是让由GCC进行编译。【VisitExpr_(const CallNode* call)】的实现,也遵循此概念。

2. 生成函数调用

结果示例:【gcc_0_0(buf_1, gcc_input3, out);】

生成函数声明后,需要生成具有正确输入和输出的函数调用。要知道在调用此函数时,应放置哪些输入或缓冲区,必须访问参数:

bool first = true;

decl_stream << func_name << "(";

for (size_t i = 0; i < call->args.size(); ++i) {

  VisitExpr(call->args[i]); // Note 1

  for (auto out : out_) {

    if (!first) {

      decl_stream << ", ";

    }

    first = false;

    decl_stream << out.first;

  }

}

// Note 2

同样,要突出显示以上代码中的注释:

Note1:【VisitExpr(call->args[i])】是递归调用,访问当前函数的参数。参数可以是另一个节点的输出或输入张量。在示例实现中,确保每个节点在离开访问器前,都更新一个类变量【out_】。这是一个例子:

  arg_node                 arg_node <- Visit arg (Note 1)       arg_node

     |                        |                                    |

 curr_node <- Process      curr_node                            curr_node <- Put "buf_0" as an input buffer

 

(a) out_ = {}            (b) out_ = {}                   (c) out_ = {("buf_0", 20)}

可以在上图中看到,在访问参数节点前,类变量【out_】为空,填充了【arg_node】输出缓冲区的名称和大小。结果,当完成访问参数节点时,可以通过查看【out_】,应该放置适当的输入缓冲区。将在本节末尾和下一节中找到更新【out_】的方式。

注意2:可能会注意到,在此步骤中没有关闭函数调用字符串。当前的函数调用字符串,如下所示:【gcc_0_0(buf_1, gcc_input3】。这是因为没有将最后一个参数(即输出)放入此调用。函数调用的输出可以是分配的临时缓冲区,也可以是子图输出张量。简化起见,在此示例中,为每个调用节点分配一个输出缓冲区(下一步),将结果从最后一个缓冲区复制到输出张量。

3.生成输出缓冲区

结果示例: 【float* buf_0 = (float*)malloc(4 * 100);】

如上一步所述,除了子图输入和输出张量外,可能还需要缓冲区来保留中间结果。为了生成缓冲区,提取形状信息以确定缓冲区的类型和大小:

// This example only supports single output.

auto type_node = call->checked_type().as<TensorTypeNode>();

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::string out = "buf_" + std::to_string(buf_idx_++);

 

// Extract the shape to be the buffer size.

auto out_shape = GetShape(call->checked_type());

int out_size = 1;

for (size_t i = 0; i < out_shape.size(); ++i) {

  out_size *= out_shape[i];

}

 

// Make the buffer allocation and push to the buffer declarations.

buf_stream << "float* " << out << " = (float*)std::malloc(4 * " << out_size << ");";

buf_decl_.push_back(buf_stream.str());

分配输出缓冲区后,现在可以关闭函数调用字符串,将生成的函数调用,放到类变量【ext_func_body】。

decl_stream << ", " << out << ");";

ext_func_body.push_back(decl_stream.str());

4. 更新输出缓冲区

为了接受当前调用节点的输出,作为输入的下一个节点,知道应使用的缓冲区,需要在离开此访问函数前,更新类变量【out_】。

out_.clear();

out_.push_back({out, out_size});

恭喜!已经完成了本文中最困难的功能。在接下来的两节中,只需要组成此函数中的一些次要缺失部分。

输入变量的代码生成

回想一下,通过访问调用节点的参数,收集输入缓冲区的信息(上一节的第二步),处理了参数是另一个调用节点的情况(第四步)。在本节中,以【VarNode】示例,演示如何处理其它节点。

【VarNode】表示模型中的输入张量。拥有的唯一的,但重要的信息是名称提示(如data,weight等)。在访问【VarNode】时,只需更新类变量【out_】传递名称提示,以便后代调用节点可以生成正确的函数调用。

void VisitExpr_(const VarNode* node) {

  ext_func_args_.push_back(node->name_hint());

  out_.clear();

  out_.push_back({node->name_hint(), 0});

}

在此示例中,假设要卸载的子图仅具有调用节点和变量节点。如果子图包含其它类型的节点,如TupleNode,需要访问并绕过输出缓冲区信息。

代码发送

该【codegen】类的最后一部分是一个【JIT】函数,该函数为子图发送C函数,将刚生成的C代码用作函数体。除了前面几节中生成的子图函数外,需要一个包装器函数,该函数具有统一的参数,TVM runtime可以调用和传递数据。幸运的是,继承的基类已经提供了实现【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. TVM runtime兼容函数【gcc_0】具有TVM统一函数参数,可解压缩TVM打包的张量,调用【gcc_0__wrapper_】。

 

因此,【JIT】实现过程中唯一需要做的,就是将生成的所有子图函数代码传递给【JitImpl】:

std::string JIT() {

  // Write function macros

  for (auto decl : func_decl_) {

    code_stream_ << decl << "\n";

  }

  return JitImpl(ext_func_id_, ext_func_args_, buf_decl_, ext_func_body, out_);

}

传递的所有的变量(【ext_func_id】等)都是类变量,在遍历子图时会被填充。

实现【CSourceCodegen 】

同样,​​创建一个类框架,实现所需的功能。注意,继承【CSourceModuleCodegenBase】

class CSourceCodegen : public CSourceModuleCodegenBase {

 public:

  // Pass a subgraph function, and generate the C code.

  void GenCFunc(const Function& func) { ; }

 

  // Use GenCFunc to generate the C code and wrap it as a C source module.

  runtime::Module CreateCSourceModule(const NodeRef& ref) override { ; }

 

 private:

  std::ostringstream code_stream_;

};

实现【GenCFunc 】

【GenCFunc】只需使用【CodegenC】,只是实现遍历Rel​​ay函数(子图),获得生成的C代码即可。内置函数【GetExtSymbol】在Relay 函数中,检索唯一的符号名称(例如gcc_0),用作C函数名称,因为该符号将用于DSOruntime查找。

 

void GenCFunc(const Function& func) {

  CHECK(func.defined()) << "Input error: expect a Relay function.";

 

  // Record the external symbol for runtime lookup.

  auto sid = GetExtSymbol(func);

 

  CodeGenC builder(sid);

  builder.VisitExpr(func->body);

  code_stream_ << builder.JIT();

}

实现【CreateCSourceModule 】

该函数为外部库创建一个runtime模块。在此示例中,创建了一个【CSourceModule】,可以直接编译,与TVM生成的DSOModule链接在一起。实现【CodegenC】后,实现此功能相对简单:

runtime::Module CreateCSourceModule(const NodeRef& ref) override {

  // Create headers

  code_stream_ << "#include <cstdint>\n";

  code_stream_ << "#include <iostream>\n";

  code_stream_ << "#include <cstdlib>\n";

  code_stream_ << "#include <stdio.h>\n";

  code_stream_ << "#include <cstring>\n";

  code_stream_ << "#include <tvm/runtime/c_runtime_api.h>\n";

  code_stream_ << "#include <dlpack/dlpack.h>\n";

 

  // Append some common macro for operator definition.

  const char* operator_macro = R"op_macro(

  #define CSOURCE_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)       \

    extern "C" void p_ID_(float* a, float* b, float* out) { \

      for (int64_t i = 0; i < p_DIM1_; ++i) {               \

        out[i] = a[i] p_OP_ b[i];                           \

      }                                                     \

    }

 

  #define CSOURCE_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \

    extern "C" void p_ID_(float* a, float* b, float* out) {     \

      for (int64_t i = 0; i < p_DIM1_; ++i) {                   \

        for (int64_t j = 0; j < p_DIM2_; ++j) {                 \

          int64_t k = i * p_DIM2_ + j;                          \

          out[k] = a[k] p_OP_ b[k];                             \

        }                                                       \

      }                                                         \

    }

  )op_macro";

 

  code_stream_ << operator_macro << "\n\n";

 

  // Generate C code for the subgraph.

  if (ref->IsInstance<FunctionNode>()) {

    GenCFunc(Downcast<Function>(ref));

  } else if (ref->IsInstance<relay::ModuleNode>()) {

    relay::Module mod = Downcast<relay::Module>(ref);

    for (const auto& it : mod->functions) {

      GenCFunc(Downcast<Function>(it.second));

    }

  } else {

    LOG(FATAL) << "The input ref is expected to be a Relay function or module"

               << "\n";

  }

 

  // Create a CSourceModule

  const auto* pf = runtime::Registry::Get("module.csource_module_create");

  CHECK(pf != nullptr) << "Cannot find csource module to create the external runtime module";

  return (*pf)(code_stream_.str(), "cc");

}

注册代码生成

最后一步是将代码生成器注册到TVM后端。先实现一个简单的函数,调用代码生成器,生成一个runtime模块。

runtime::Module CCompiler(const NodeRef& ref) {

  CSourceCodegen csource;

  return csource.CreateCSourceModule(ref);

}

最后,将此功能注册到TVM后端:

TVM_REGISTER_GLOBAL("relay.ext.ccompiler").set_body_typed(CCompiler);

其中【ccompiler】是一个自定义标签,用于TVM知道这是在用【ccompiler】注释子图时,应使用生成和卸载子图的代码生成器。

最后,一个好的做法是设置CMake配置标志,仅为客户提供编译器。先创建一个cmake文件【cmake/modules/contrib/CODEGENC.cmake】:

if(USE_CODEGENC)

  file(GLOB CSOURCE_RELAY_CONTRIB_SRC src/relay/backend/contrib/codegen_c/codegen.cc)

  list(APPEND COMPILER_SRCS ${CSOURCE_RELAY_CONTRIB_SRC})

endif(USE_CODEGENC)

这样,用户可以在配置TVM时,使用【config.cmake】以下命令,配置是否包括编译器:

set(USE_CODEGENC ON)

为表示实现一个代码生成

尽管已经演示了如何实现C代码生成,但是硬件可能需要其它的图形表示形式,如JSON。在这种情况下,可以修改【CodegenC】类,已经实现了自定义图形表示,实现定制的runtime模块,使TVM runtime知道,如何执行该图形表示。

为了简化,在本文中定义了一个名为“ ExampleJSON”的图表示。ExampleJSON不是真正的JSON,仅仅是没有控制流的图的简单表示。例如,假设有一个名为【subgraph_0】的子图:

 input0

   |

  add <-- input1

   |

subtract <-- input2

   |

multiply <-- input3

   |

  out

然后,该子图的【ExampleJON】如下所示:

subgraph_0

  input 0 10 10

  input 1 10 10

  input 2 10 10

  input 3 10 10

  add 4 inputs: 0 1 shape: 10 10

  sub 5 inputs: 4 2 shape: 10 10

  add 6 inputs: 5 3 shape: 10 10

【input】关键字声明输入张量的ID和形状; 其它语句则以语法描述计算:

【<op> <output ID> inputs: [input ID] shape: [shape]】

在本节中,目标是实现以下定制的TVM runtime模块,执行【ExampleJSON】图。

runtime::Module ExampleJsonCompiler(const NodeRef& ref) {

    ExampleJsonCodeGen codegen(ref);

    std::string code = codegen.gen(); // Note 1

    const auto* pf = runtime::Registry::Get("module.examplejson_module_create"); // Note 2

    CHECK(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);

Note1:将实现自定义代码生成,通过子图生成ExampleJSON代码字符串。

Note2:此行获得指向用于创建定制runtime模块的函数的指针。可以看到采用了刚刚生成的ExampleJSON格式的子图代码,初始化了runtime模块。

在以下各节中,将介绍

1)如何实现【ExampleJsonCodeGen】

2)如何实现和注册【examplejson_module_create】。

实现【ExampleJsonCodeGen 】

类似于C代码生成器,从【ExprVisitor】派生了【ExampleJsonCodeGen】,利用访问者模式,进行子图遍历的方法。另一方面,不需要继承【CodegenCBase】,因为不需要TVM C ++装饰器。codegen类的实现如下:

#include <tvm/relay/expr_functor.h>

#include <tvm/relay/transform.h>

#include <tvm/relay/type.h>

#include <tvm/runtime/module.h>

#include <tvm/runtime/object.h>

 

#include <fstream>

#include <sstream>

 

namespace tvm {

namespace relay {

namespace contrib {

 

class ExampleJsonCodeGen : public ExprVisitor {

  public:

    explicit ExampleJsonCodeGen();

 

    // Note 1

    void VisitExpr_(const VarNode* node) { /* Skip in this example. */ }

    void VisitExpr_(const CallNode* call) final { /* Skip in this example. */ }

 

    // Note 2

    std::string gen(NodeRef& ref) {

        this->code = "";

        if (ref->IsInstance<FunctionNode>()) {

            this->visit(Downcast<Function>(ref));

        } else if (ref->IsInstance<relay::ModuleNode>()) {

            relay::Module mod = Downcast<relay::Module>(ref);

            for (const auto& it : mod->functions) {

                this->visit(Downcast<Function>(it.second));

            }

        } else {

            LOG(FATAL) << "The input ref is expected to be a Relay function or module";

        }

        return this->code;

    }

 

  private:

      /*! \brief The function id that represents a C source function. */

     std::string code;

}

Note1:再次实现相应的访问者函数,生成ExampleJSON代码,将存储到类变量【code】中(在本示例中,跳过了访问器函数的实现,因为概念与C代码基本相同)。完成图访问后,应该在【code】中有一个ExampleJSON图。

Note2:定义了一个内部API gen,获取子图生成ExampleJSON代码。该API可以采用喜欢的任意名称。

下一步是实施自定义的runtime,输出ExampleJsonCodeGen。

实现自定义runtime

在本节中,将逐步实现自定义的TVM runtime,将注册到TVM runtime模块。自定义的runtime应位于src/runtime/contrib/<your-runtime-name>/。在示例中,将runtime命名为“ example_ext_runtime”,将放在“ here <src / runtime / contrib / example_ext_runtime / example_ext_runtime.cc>” _下。随时检查此文件,获取完整的实现。

再次,先定义一个自定义的runtime类,如下所示。该类必须从TVM派生【ModuleNode】,以便与其它TVM runtime模块兼容。

#include <dmlc/logging.h>

#include <tvm/runtime/c_runtime_api.h>

#include <tvm/runtime/memory.h>

#include <tvm/runtime/module.h>

#include <tvm/runtime/ndarray.h>

#include <tvm/runtime/object.h>

#include <tvm/runtime/packed_func.h>

#include <tvm/runtime/registry.h>

 

#include <fstream>

#include <cmath>

#include <map>

#include <sstream>

#include <string>

#include <vector>

 

namespace tvm {

namespace runtime {

class ExampleJsonModule : public ModuleNode {

 public:

  explicit ExampleJsonModule(std::string graph_json);

 

  PackedFunc GetFunction(const std::string& name,

                         const ObjectPtr<Object>& sptr_to_self) final;

 

  const char* type_key() const { return "examplejson"; }

 

  void SaveToBinary(dmlc::Stream* stream) final;

 

  static Module LoadFromBinary(void* strm);

 

  static Module Create(const std::string& path);

 

  std::string GetSource(const std::string& format = "");

 

  void Run(int id, const std::vector<int>& inputs, int output);

 

  void ParseJson(const std::string& json);

 

 private:

  /* \brief The json string that represents a computational graph. */

  std::string graph_json_;

  /* \brief The subgraph that being processed. */

  std::string curr_subgraph_;

  /*! \brief A simple graph from subgraph id to node entries. */

  std::map<std::string, std::vector<NodeEntry> > graph_;

  /* \brief A simple pool to contain the tensor for each node in the graph. */

  std::vector<NDArray> data_entry_;

  /* \brief A mapping from node id to op name. */

  std::vector<std::string> op_id_;

};

特别的,必须在【ExampleJsonModule】中,实现一些【ModuleNode】派生的函数:

构造函数:此类的构造函数应接受一个子图(以表示形式),以所需的任何方式,进行处理和存储。保存的子图可由以下两个函数使用。

【GetFunction】:这是此类中最重要的函数。当TVM runtime要使用编译器标记执行子图时,TVM runtime会从自定义runtime模块调用此函数。提供函数名称以及runtime参数,【GetFunction】应返回打包的函数实现,供TVM runtime执行。

【SaveToBinary】和【LoadFromBinary】:【SaveToBinary】将runtime模块序列化为二进制格式,供以后部署。用户使用【export_libraryAPI 】时,TVM将调用此函数。另一方面,由于现在使用自定义图表示形式,因此必须确保【LoadFromBinary】能够通过采用【SaveToBinary】生成的序列化二进制文件,构造相同的runtime模块。

【GetSource】(可选):如果想查看生成的【ExampleJSON】代码,可以实现此函数转储;否则,可以跳过实施。

其它功能和类变量将与上述必备功能的实现一起引入。

实现构造函数

explicit ExampleJsonModule(std::string graph_json) {

  this->graph_json_ = graph_json;

  ParseJson(this->graph_json_);

}

然后,实现【ParseJson】解析ExampleJSON格式的子图,在内存中构造一个图,供以后使用。由于在此示例中不支持带有分支的子图,因此仅使用数组,按顺序存储子图中的每个节点。

void ParseJson(const std::string& json) {

  std::string line;

  std::string curr_subgraph;

  std::stringstream ss(json);

 

  while (std::getline(ss, line, '\n')) {

    std::stringstream ss2(line);

    std::string token;

    int id = 0;

 

    ss2 >> token;

    if (token.find("subgraph_") != std::string::npos) {

      curr_subgraph = token;

      continue;

    }

 

    ss2 >> id;

    if (op_id_.size() <= static_cast<size_t>(id)) {

      op_id_.resize(id + 1);

      data_entry_.resize(id + 1);

    }

 

    int64_t total_elements = 1;

    std::vector<int64_t> shape;

    if (token == "input") {

      int64_t size = 0;

      while (ss2 >> size) {

        total_elements *= size;

        shape.push_back(size);

      }

    } else {

      op_id_[id] = token; // Note 1

      bool shape_data = false;

      NodeEntry entry;

      while (ss2 >> token) {

        if (token == "shape:") {

          shape_data = true;

        } else if (shape_data) {

          total_elements *= std::stoll(token);

          shape.push_back(std::stoll(token));

        } else if (token != "inputs:") {

          entry.inputs.push_back(std::stoi(token));

        }

      }

      entry.id = id;

      entry.output = id;

      graph_[curr_subgraph].push_back(entry); // Note 2

    }

    DLContext ctx;

    ctx.device_type = static_cast<DLDeviceType>(1);

    ctx.device_id = 0;

    data_entry_[id] = NDArray::Empty(shape, DLDataType{kDLFloat, 32, 1}, ctx); // Note 3

  }

}

Note1:使用类变量【op_id_】将子图节点ID,映射到运算符名称(如【add】),可以在runtime调用相应的运算符函数。

Note2:使用类变量【graph_】,将子图名称映射到节点数组。【GetFunction】将在runtime通过子图ID查询图节点。

Note3:使用类变量【data_entry_】,将子图节点ID映射到张量数据占位符。将在runtime将输入和输出放入相应的数据条目。

实现【GetFunction 】

构造后,应该准备好上述类变量。然后,实现【GetFunction】为TVM runtime,提供可执行的子图函数:

PackedFunc GetFunction(const std::string& name,

                       const ObjectPtr<Object>& sptr_to_self) final {

  if (this->graph_.find(name) != this->graph_.end()) {

    this->curr_subgraph_ = name;

    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {

 

      // Copy input tensors to corresponding data entries.

      for (auto i = 0; i < args.size(); ++i) {

        CHECK(args[i].type_code() == kNDArrayContainer || args[i].type_code() == kArrayHandle)

            << "Expect NDArray or DLTensor as inputs\n";

        if (args[i].type_code() == kArrayHandle) {

          DLTensor* arg = args[i];

          this->data_entry_[i].CopyFrom(arg);

        } else {

          NDArray arg = args[i];

          this->data_entry_[i].CopyFrom(arg);

        }

      }

 

      // Execute the subgraph.

      for (const auto& it : this->graph_[this->curr_subgraph_]) {

        this->Run(it.id, it.inputs, it.output);

      }

      CHECK_GT(graph_.count(this->curr_subgraph_), 0U);

 

      // Copy the output from a data entry back to TVM runtime argument.

      auto out_idx = graph_[this->curr_subgraph_].back().output;

      if (args[args.size() - 1].type_code() == kArrayHandle) {

        DLTensor* arg = args[args.size() - 1];

        this->data_entry_[out_idx].CopyTo(arg);

      } else {

        NDArray arg = args[args.size() - 1];

        this->data_entry_[out_idx].CopyTo(arg);

      }

      *rv = data_entry_.back();

    });

  } else {

    LOG(FATAL) << "Unknown subgraph: " << name << "\n";

    return PackedFunc();

  }

}

可以看出,【GetFunction】由三个主要部分组成。第一部分将数据从TVM runtime参数复制到在构造函数中分配的相应数据条目。第二部分使用【Run】函数(将在以后实现)执行子图,将结果保存到另一个数据条目中。第三部分将结果从输出数据条目,复制回相应的TVM runtime参数进行输出。

实现运行

现在让实现【Run】函数。此函数接受:

1)一个子图ID;

2)输入数据条目索引的列表

3)输出数据条目索引。

void Run(int id, const std::vector<int>& inputs, int output) {

  // Make a list data entry indexs.

  std::vector<int> args(inputs.begin(), inputs.end());

  args.push_back(output);

 

  // Initialize data holders.

  std::vector<TVMValue> values(args.size());

  std::vector<int> type_codes(args.size());

 

  // Initialize a TVM arg setter with TVMValue and its type code.

  TVMArgsSetter setter(values.data(), type_codes.data());

 

  // Set each argument to its corresponding data entry.

  if (op_id_[id] == "add" || op_id_[id] == "sub" || op_id_[id] == "mul") {

    for (size_t i = 0; i < args.size(); i++) {

      setter(i, data_entry_[args[i]]);

    }

  }

 

  // Invoke the corresponding operator function.

  if (op_id_[id] == "add") {

    Add(values.data(), type_codes.data(), args.size());

  } else if (op_id_[id] == "sub") {

    Sub(values.data(), type_codes.data(), args.size());

  } else if (op_id_[id] == "mul") {

    Mul(values.data(), type_codes.data(), args.size());

  } else {

    LOG(FATAL) << "Unknown op: " << op_id_[id] << "\n";

  }

}

【Run】函数主要有两个部分。第一部分分配一个【TVMValue】列表,映射相应的数据条目块。这将成为运算符函数的参数。第二部分将调用运算符函数。虽然使用与前面的例子相同的C函数,可以用自定义引擎更换Add,Sub及Mul。只需要确保引擎将结果存储到最后一个参数,就可以将传输回TVM runtime。

通过实现上述功能,自定义的代码生成和runtime现在可以执行子图。最后一步是注册API(【examplejson_module_create】)以创建此模块:

TVM_REGISTER_GLOBAL("module.examplejson_module_create")

.set_body_typed([](std::string code){

    auto n = make_object<ExampleJsonModule>(code);

    return runtime::Module(n);

});

实现【SaveToBinary】和【LoadFromBinary 】

到目前为止,已经实现了自定义runtime的主要功能,可以用作其它TVM runtime。但是,当用户要将已构建的runtime保存到磁盘以进行部署时,TVM不知道如何保存。这就是要实现【SaveToBinary】和【LoadFromBinary】的原因,告诉TVM如何保留和恢复自定义的runtime。

先实现【SaveToBinary】,允许用户将该模块保存在磁盘中的功能。

void SaveToBinary(dmlc::Stream* stream) final {

    stream->Write(this->graph_json_);

}

可以发现此函数非常简单。回想一下,在构造函数中使用的唯一参数是一个子图表示,只需要一个子图表示,即可构造/恢复此定制的runtime模块。结果,【SaveToBinary】只需将子图写入输出DMLC流。当用户使用【export_library】API导出模块时,自定义模块将是子图的ExampleJSON流。

同样,【LoadFromBinary】读取子图流,重新构建自定义的runtime模块:

static Module LoadFromBinary(void* strm) {

  dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);

  std::string graph_json;

  stream->Read(&graph_json);

  auto n = tvm::runtime::make_object<ExampleJsonModule>(graph_json);

  return Module(n);

}

需要注册此函数,启​​用相应的Python API:

TVM_REGISTER_GLOBAL("module.loadbinary_examplejson")

.set_body_typed(ExampleJsonModule::LoadFromBinary);

上面的注册当用户调用【tvm.runtime.load(lib_path)】API,导出的库具有ExampleJSON流时,【LoadFromBinary】将被调用,创建相同的自定义runtime模块。

另外,如果想直接从ExampleJSON文件支持模块创建,可以实现一个简单的函数,注册Python API,如下所示:

static Module Create(const std::string& path) {

    std::ifstream filep;

    filep.open(path, std::ios::in);

    std::string graph_json;

    std::string line;

    while (std::getline(filep, line)) {

        graph_json += line;

        graph_json += "\n";

    }

    filep.close();

    auto n = tvm::runtime::make_object<ExampleJsonModule>(graph_json);

    return Module(n);

}

 

TVM_REGISTER_GLOBAL("module.loadfile_examplejson")

.set_body([](TVMArgs args, TVMRetValue* rv) {

    *rv = ExampleJsonModule::Create(args[0]);

});

用户可以手动编写/修改ExampleJSON文件,使用Python API 【tvm.runtime.load("mysubgraph.examplejson", "examplejson")】构造自定义模块。

 

总结

总之,这是一份清单供参考:

派生自【ExprVisitor】和【CodegenCBase】的代码生成类和(仅对于C代码生成),具有以下函数。

【VisitExpr_(const CallNode* call)】 收集调用节点信息。

收集子图信息所需的其它访问器函数。

【JIT 】生成子图代码。

注册代码生成器。

创建【CSourceModule】的函数(用于C代码生成)。

从【ModuleNode】派生的runtime模块类,具有下面的函数(用于图形表示)。

构造函数。

【GetFunction】生成TVM runtime兼容的【PackedFunc】。

【Run 】执行子图。

注册runtime创建API。

【SaveToBinary】和【LoadFromBinary】序列化/反序列化自定义的runtime模块。

注册【LoadFromBinary】API以支持【tvm.runtime.load(your_module_lib_path)】。

(可选)【Create】从表示中的子图文件,支持定制的runtime模块构造。

一个用于对用户Relay程序进行注释的注释器,利用编译器和runtime(TBA)。

 

 

参考链接:

https://tvm.apache.org/docs/dev/how_to/relay_bring_your_own_codegen.html

https://blog.csdn.net/weixin_42164269/article/details/104291635

标签:std,代码生成,gcc,自定义,TVM,runtime,id,out
来源: https://www.cnblogs.com/wujianming-110117/p/15550639.html