TVM开发三个示例分析

TVM开发三个示例分析

把自主生成的代码生成TVM

把自主生成的代码生成TVM

简介

1. 要生成C代码。

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

实现一个C代码生成器

实现【CodegenC】

运算符代码生成

输入变量的代码生成

代码发送

实现【CSourceCodegen 】

实现【GenCFunc 】

实现【CreateCSourceModule 】

注册代码生成

为表示实现一个代码生成

实现【ExampleJsonCodeGen 】

实现自定义运行时

实现构造函数

实现【GetFunction 】

实现运行

实现【SaveToBinary】和【LoadFromBinary 】

总结

简介

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

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

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

本文演示了作为硬件后端提供者，如何轻松实现自主生成的代码生成并注册为Relay后端编译器，以支持硬件设备/库。根据需要的不同图形表示形式涵盖两种类型的代码生成器：

要生成C代码。

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

要生成任何其它图形表示。

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

在完成代码生成和运行时之后，可以让客户使用自定义标签，注释模型使用。最终用户注释和启动特定代码生成。

实现一个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运行时兼容的包装函数。接受一个输入张量和一个输出张量的列表（最后一个参数），将转换为正确的数据类型，调用Note2中描述的子图函数。此外，【TVM_DLL_EXPORT_TYPED_FUNC】是一个TVM宏，生成另一个函数【gcc_0】，【gcc_0】具有统一的函数参数，通过把所有的参数张量打包成【TVMArgs】。结果，TVM运行时可以直接调用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源运行时模块中，以供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运行时可以调用和传递数据。幸运的是，继承的基类已经提供了实现【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运行时兼容函数【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】，只是实现遍历Relay函数（子图）并获得生成的C代码即可。内置函数【GetExtSymbol】在Relay 函数中，检索唯一的符号名称（如gcc_0），必须用作C函数名称，因为该符号将用于DSO运行时查找。

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 】

该函数为外部库创建一个运行时模块。在此示例中，创建了一个【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::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】类，已经实现了自主生成的图形表示，实现定制的运行时模块，使TVM运行时知道，如何执行该图形表示。

为了简化，定义了一个名为“ 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运行时模块，执行【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：此行获得指向用于创建定制运行时模块的函数的指针。采用了刚刚生成的ExampleJSON格式的子图代码，初始化了运行时模块。

在以下各节中，将介绍

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可以采用喜欢的任意名称。

下一步是实施自定义的运行时，输出ExampleJsonCodeGen。

实现自定义运行时

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

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

#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运行时要使用编译器标记执行子图时，TVM运行时会从自定义运行时模块调用此函数。提供函数名称以及运行时参数，【GetFunction】应返回打包的函数实现，供TVM运行时执行。

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

【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】），以便可以在运行时调用相应的运算符函数。

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

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

实现【GetFunction 】

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

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运行时参数复制到在构造函数中分配的相应数据条目。

第二部分使用【Run】函数（将在以后实现）执行子图并将结果保存到另一个数据条目中。

第三部分将结果从输出数据条目复制回相应的TVM运行时参数以进行输出。

实现运行

现在让实现【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运行时。

通过实现上述功能，自定义的代码生成和运行时，现在可以执行子图。最后一步是注册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 】

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

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

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

stream->Write(this->graph_json_);

}

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

同理，【LoadFromBinary】读取子图流并重新构建自定义的运行时模块：

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】调用创建相同的自定义运行时模块。

另外，如果想直接从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】派生的运行时模块类，具有下面的函数（用于图形表示）。

构造函数。

【GetFunction】生成TVM运行时兼容的【PackedFunc】。

【Run 】执行子图。

注册运行时创建API。

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

（可选）【Create】以从表示中的子图文件，支持定制的运行时模块构造。

一个用于对用户Relay程序进行注释的注释器，利用编译器和运行时（TBA）。

参考链接：

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

TVM代码库演练示例

代码库结构概述

向量添加示例

了解新的代码库可能是一个挑战。对于像TVM这样的代码库，尤其如此，其中不同的组件以非显而易见的方式交互。在本指南中，尝试通过一个简单的示例来说明构成编译的关键元素。对于每个重要步骤，都会显示在代码库中的哪个位置。目的是让新开发人员和感兴趣的用户更快地进入代码库。

代码库结构概述

在TVM库的根目录中，具有以下子目录，这些子目录一起构成了大部分代码库。

src -用于操作符编译和部署运行时的C ++代码。

src/relay -Relay实现，深度学习框架的新功能IR。

python-Python前端，封装【src】中C ++函数和对象实现。

topi -计算标准神经网络操作符的定义和后端调度。

使用标准的深度学习术语，【src/relay】是管理计算图的组件，并且图中的节点是使用【src】其余部分中实现的基础结构来编译和执行的。python为用户可用来执行编译的C ++ API和驱动程序代码提供python绑定。操作符对应【src/relay/op】中注册的每一个节点。操作符的实现位于【topi】，并且使用C ++或Python进行编码。

当用户通过【relay.build(...)】调用图编译时，图中的每个节点都会发生以下操作序列：

通过查询操作符注册表来查找操作符实现

为操作符生成计算表达式和调度

将运算符编译为目标代码

TVM代码库有趣的方面之一是C ++和Python之间的互操作性不是单向的。通常，所有执行繁重工作的代码都是用C ++实现的，并且为用户界面提供了Python绑定。在TVM中也是如此，但是在TVM代码库中，C ++代码也可以调用Python模块中定义的函数。例如，卷积运算符是用Python实现的，其实现是从Relay中的C ++代码调用的。

向量加法示例

使用一个直接使用低级TVM API的简单示例。该示例是矢量加法，【https://docs.tvm.ai/tutorials/get_started.html#sphx-glr-tutorials-get-started-py】进行详细介绍。

n = 1024

A = tvm.placeholder((n,), name='A')

B = tvm.placeholder((n,), name='B')

C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C")

在这里，A，B，C的类型是【tvm.tensor.Tensor】，定义在【python/tvm/tensor.py】中。Python中的【Tensor】是由C ++中的【Tensor】包装的，在【include/tvm/tensor.h】和【src/lang/tensor.cc】中实现。TVM中的所有Python类型都可以视为具有相同名称的基础C ++类型的句柄。如果在下面看到Python 【Tensor】类型的定义，可以看到它是【Object】的子类。

@register_object

class Tensor(Object, _expr.ExprOp):

"""Tensor object, to construct, see function.Tensor"""

def __call__(self, *indices):

...

对象协议是将C ++类型公开给前端语言（包括Python）的基础。TVM实现Python包装的方法并不简单。【https://docs.tvm.ai/dev/runtime.html#tvm-node-and-compiler-stack】简要介绍了它，如果有兴趣，请参阅【python/tvm/_ffi/】详细信息。

使用【TVM_REGISTER_*】宏，以PackedFunc的形式将C ++函数公开给前端语言。【PackedFunc】是TVM在C ++和Python之间实现互操作性的另一种机制。特别的，这使得从C ++代码库调用Python函数非常容易。还可以检查【 FFI Navigator（https://github.com/tqchen/ffi-navigator）】，该导航器使可以在python和c ++ FFI调用之间进行导航。

【Tensor】对象具有【Operation】与其相关联，定义在【python/tvm/te/tensor.py】，【include/tvm/te/operation.h】和【src/tvm/te/operation】子目录。【Tensor】是【Operation】对象的输出。每个【Operation】对象都有相应的【input_tensors()】方法，该方法返回输入【Tensor】列表。这样就可以跟踪【Operation】之间的依赖关系。

传递与输出张量【C】相对应的运算以到【python/tvm/te/schedule.py】中的【tvm.create_schedule()】函数。

s = tvm.create_schedule(C.op)

此函数映射到【include/tvm/schedule.h】中的C ++函数。

inline Schedule create_schedule(Array<Operation> ops) {

return ScheduleNode::make(ops);

}

【Schedule】由【Stage】和输出【Operation】的集合组成。

【Stage】对应一个【Operation】。在上面的矢量加法示例中，有两个占位符操作和一个计算操作，因此调度【s】包含三个阶段。各【Stage】保持关于循环嵌套结构的信息，每个循环的类型（Parallel，Vectorized，Unrolled），并且下一个【Stage】循环嵌套执行其计算，如果有的话。

【Schedule】和【Stage】被定义在【tvm/python/te/schedule.py】，【include/tvm/te/schedule.h】和【src/te/schedule/schedule_ops.cc】。

为简单起见，在上述【create_schedule()】函数创建的默认调度中调用【tvm.build(...)】函数。

target = "cuda"

fadd = tvm.build(s, [A, B, C], target)

定义在【python/tvm/driver/build_module.py】中的【tvm.build()】，接受一个调度，输入和输出【Tensor】以及目标，然后返回一个【tvm.runtime.Module】对象。一个【tvm.runtime.Module】对象包含一个可以使用函数调用语法调用的已编译函数。

【tvm.build()】的过程可以分为两个步骤：

降低，将高级别的初始循环嵌套结构转换为最终的低级别IR

代码生成，其中从低级IR生成目标机器代码

降低是通过【tvm.lower()】函数完成的，定义在【python/tvm/build_module.py】中。首先，执行边界推断，并创建初始循环嵌套结构。

def lower(sch,

args,

name="default_function",

binds=None,

simple_mode=False):

...

bounds = schedule.InferBound(sch)

stmt = schedule.ScheduleOps(sch, bounds)

...

边界推断是推断所有循环边界和中间缓冲区大小的过程。如果以CUDA后端为目标并且使用共享内存，则会在此处自动确定所需的最小大小。绑定推理在【src/te/schedule/bound.cc】，【src/te/schedule/graph.cc】和【src/te/schedule/message_passing.cc】中实现。有关绑定推理如何工作的更多信息，请参见【http://docs.tvm.ai/dev/inferbound.html】。

【stmt】，是【ScheduleOps()】的输出，代表初始的循环嵌套结构。如果已将【reorder】原语和【split 】原语应用到调度中，则初始循环嵌套已经反映了这些更改。【ScheduleOps()】在【src/te/schedule/schedule_ops.cc】中定义。

接下来，将多个降低转换应用于【stmt】。这些过程在【src/tir/pass】子目录中实现。例如，如果已对时间表应用了【vectorize】或【unroll】原语，则将被应用到循环矢量化和下面的展开过程中。

...

stmt = ir_pass.VectorizeLoop(stmt)

...

stmt = ir_pass.UnrollLoop(

stmt,

cfg.auto_unroll_max_step,

cfg.auto_unroll_max_depth,

cfg.auto_unroll_max_extent,

cfg.unroll_explicit)

...

降低完成后，【build()】函数从降低的函数生成目标机器代码。如果以x86为目标，则此代码可以包含SSE或AVX指令，或以CUDA为目标的PTX指令。除了目标特定的机器代码之外，TVM还生成主机侧代码，该代码负责内存管理，内核启动等。

代码生成由【python/tvm/target/codegen.py】中定义的【build_module()】函数完成。在C ++侧，代码生成在【src/target/codegen】子目录中实现。【build_module()】Python函数将达到【src/target/codegen/codegen.cc】中的【Build()】函数：

runtime::Module Build(const Array<LoweredFunc>& funcs,

const std::string& target) {

std::string build_f_name = "codegen.build_" + target;

const PackedFunc* bf = runtime::Registry::Get(build_f_name);

runtime::Module m = (*bf)(funcs, target);

return m;

}

【Build()】函数在【PackedFunc】注册表中查找给定目标的代码生成器，并调用找到的函数。例如，【codegen.build_cuda】函数在【src/codegen/build_cuda_on.cc】中注册，如下所示：

TVM_REGISTER_GLOBAL("codegen.build_cuda")

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

*rv = BuildCUDA(args[0]);

});

上述使用【CodeGenCUDA 】类从降低IR生成的CUDA源码核【BuildCUDA()】定义在【src/codegen/codegen_cuda.cc】，和使用NVRTC内核编译。如果针对使用LLVM（包括x86，ARM，NVPTX和AMDGPU）的后端，则代码生成主要由【src/codegen/llvm/codegen_llvm.cc】中定义的类【CodeGenLLVM】完成。【CodeGenLLVM】将TVM IR转换为LLVM IR，运行大量LLVM优化遍历，并生成目标机器代码。

【src/codegen/codegen.cc】中的【Build()】函数返回定义在【include/tvm/runtime/module.h】和【src/runtime/module.cc】中定义的对象【runtime::Module】。【Module】对象是一个容器，装载特定于目标的【ModuleNode】对象。每个后端都实现【ModuleNode】子类，以添加目标特定的运行时API调用。例如，CUDA后端在【src/runtime/cuda/cuda_module.cc】中实现【CUDAModuleNode】类，该类管理CUDA驱动程序API。上面的【BuildCUDA()】函数用【runtime::Module】装饰【CUDAModuleNode】，并返回到Python端。LLVM后端【LLVMModuleNode】在【src/codegen/llvm/llvm_module.cc】中实现，它处理已编译代码的JIT执行。【ModuleNode】的其他子类可以在【src/runtime】的子目录下找到，与每个后端相对应。

返回的模块（可以认为是已编译函数和设备API的组合）可以在TVM的NDArray对象上调用。

ctx = tvm.context(target, 0)

a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)

b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)

c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)

fadd(a, b, c)

output = c.asnumpy()

在幕后，TVM会自动分配设备内存并管理内存传输。为此，每个后端都需要继承【DeviceAPI】类，定义在【include/tvm/runtime/device_api.h】中，并重写内存管理方法以使用特定于设备的API。例如，在【src/runtime/cuda/cuda_device_api.cc】中实现的【CUDADeviceAPI】CUDA后端，以使用【cudaMalloc】，【cudaMemcpy】等。

首次使用【fadd(a, b, c)】调用已编译的模块时，【ModuleNode】的【GetFunction()】方法被调用，来获得一个可用于内核调用的【PackedFunc 】方法。例如，在【src/runtime/cuda/cuda_device_api.cc】中，CUDA后端【CUDAModuleNode::GetFunction()】实现如下：

PackedFunc CUDAModuleNode::GetFunction(

const std::string& name,

const std::shared_ptr<ModuleNode>& sptr_to_self) {

auto it = fmap_.find(name);

const FunctionInfo& info = it->second;

CUDAWrappedFunc f;

f.Init(this, sptr_to_self, name, info.arg_types.size(), info.thread_axis_tags);

return PackFuncVoidAddr(f, info.arg_types);

}

【PackedFunc】的超载【operator()】将被调用，这反过来又调用实现在【src/runtime/cuda/cuda_module.cc】中的【CUDAWrappedFunc】的【operator()】函数，在这里终于看到了【cuLaunchKernel】驱动调用：

class CUDAWrappedFunc {

public:

void Init(...)

...

void operator()(TVMArgs args,

TVMRetValue* rv,

void** void_args) const {

int device_id;

CUDA_CALL(cudaGetDevice(&device_id));

if (fcache_[device_id] == nullptr) {

fcache_[device_id] = m_->GetFunc(device_id, func_name_);

}

CUstream strm = static_cast<CUstream>(CUDAThreadEntry::ThreadLocal()->stream);

ThreadWorkLoad wl = thread_axis_cfg_.Extract(args);

CUresult result = cuLaunchKernel(

fcache_[device_id],

wl.grid_dim(0),

wl.grid_dim(1),

wl.grid_dim(2),

wl.block_dim(0),

wl.block_dim(1),

wl.block_dim(2),

0, strm, void_args, 0);

}

};

总结了TVM如何编译和执行函数。尽管没有详细介绍TOPI或Relay，但是最后，所有神经网络操作符都经过与上述相同的编译过程。鼓励深入研究其余代码库的细节。

参考链接：

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

TVM Operator Inventory (TOPI)简介

TOPI简介

这是 TVM Operator Inventory (TOPI) 的介绍。TOPI 提供了比 TVM 具有更高抽象的 numpy 风格的，通用操作和调度。TOPI 如何在 TVM 中，编写样板代码。

from __future__ import absolute_import, print_function

import tvm

import tvm.testing

from tvm import te

from tvm import topi

import numpy as np

基本示例

重新审视行总和操作（相当于B=numpy.sum(A,axis=1)），要计算二维 TVM 张量 A 行总和，应该指定符号操作及调度。

n = te.var("n")

m = te.var("m")

A = te.placeholder((n, m), name="A")

k = te.reduce_axis((0, m), "k")

B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name="B")

s = te.create_schedule(B.op)

以人类可读的格式，检查 IR 代码，可以这样做。

print(tvm.lower(s, [A], simple_mode=True))

输出：

primfn(A_1: handle) -> ()

attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}

buffers = {A: Buffer(A_2: Pointer(float32), float32, [n: int32, m: int32], [stride: int32, stride_1: int32], type="auto")}

buffer_map = {A_1: A} {

allocate(B: Pointer(global float32), float32, [n]), storage_scope = global;

for (i: int32, 0, n) {

B[i] = 0f32

for (k: int32, 0, m) {

B[i] = ((float32*)B[i] + (float32*)A_2[((i*stride) + (k*stride_1))])

}

对于这样一个常见的操作，必须定义 reduce 轴，以及使用 te.compute进行显式计算。对于更复杂的操作，需要提供多少细节。可以用简单topi.sum的，如numpy.sum，替换这两行。

C = topi.sum(A, axis=1)

ts = te.create_schedule(C.op)

print(tvm.lower(ts, [A], simple_mode=True))

输出：

primfn(A_1: handle) -> ()

attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}

buffers = {A: Buffer(A_2: Pointer(float32), float32, [n: int32, m: int32], [stride: int32, stride_1: int32], type="auto")}

buffer_map = {A_1: A} {

allocate(A_red: Pointer(global float32), float32, [n]), storage_scope = global;

for (ax0: int32, 0, n) {

A_red[ax0] = 0f32

for (k1: int32, 0, m) {

A_red[ax0] = ((float32*)A_red[ax0] + (float32*)A_2[((ax0*stride) + (k1*stride_1))])

}

Numpy 风格的算子重载

可以使用topi.broadcast_add具有正确（可广播特定）shape的张量，添加两个张量。TOPI 为此类常见操作，提供了算子重载。例如，

x, y = 100, 10

a = te.placeholder((x, y, y), name="a")

b = te.placeholder((y, y), name="b")

c = a + b # same as topi.broadcast_add

d = a * b # same as topi.broadcast_mul

使用相同的语法重载，TOPI 处理，将原语（int，float）广播到 tensor d-3.14。

通用调度和融合操作

TOPI 如何免于在较低级别的 API 中，编写显式计算。像以前一样进行调度，TOPI根据给定的上下文，提供更高级别的调度方法。例如，对于 CUDA，可以using only topi.generic.schedule_reduce，调度topi.sum结尾的一系列操作。

e = topi.elemwise_sum([c, d])

f = e / 2.0

g = topi.sum(f)

with tvm.target.cuda():

sg = topi.cuda.schedule_reduce(g)

print(tvm.lower(sg, [a, b], simple_mode=True))

输出：

primfn(a_1: handle, b_1: handle) -> ()

attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}

buffers = {b: Buffer(b_2: Pointer(float32), float32, [10, 10], []),

a: Buffer(a_2: Pointer(float32), float32, [100, 10, 10], [])}

buffer_map = {a_1: a, b_1: b} {

allocate(T_divide_red: Pointer(global float32), float32, [1]), storage_scope = global;

attr [IterVar(threadIdx.x: int32, [0:1024], "ThreadIndex", "threadIdx.x")] "thread_extent" = 1024;

allocate(T_divide_red.rf: Pointer(local float32), float32, [1]), storage_scope = local;

allocate(reduce_temp0: Pointer(local float32), float32, [1]), storage_scope = local {

T_divide_red.rf[0] = 0f32

for (k0.k1.fused.k2.fused.outer: int32, 0, 10) {

if @tir.likely((((((k0.k1.fused.k2.fused.outer*1024) + threadIdx.x) < 10000) && (((k0.k1.fused.k2.fused.outer*1024) + threadIdx.x) < 10000)) && (((k0.k1.fused.k2.fused.outer*1024) + threadIdx.x) < 10000)), dtype=bool) {

T_divide_red.rf[0] = ((float32*)T_divide_red.rf[0] + ((((float32*)a_2[((k0.k1.fused.k2.fused.outer*1024) + threadIdx.x)] + (float32*)b_2[floormod(((k0.k1.fused.k2.fused.outer*1024) + threadIdx.x), 100)]) + ((float32*)a_2[((k0.k1.fused.k2.fused.outer*1024) + threadIdx.x)]*(float32*)b_2[floormod(((k0.k1.fused.k2.fused.outer*1024) + threadIdx.x), 100)]))*0.5f32))

}

attr [meta[tir.CommReducer][0]] "reduce_scope" = @tir.reinterpret(0u64, dtype=handle);

@tir.tvm_thread_allreduce(1u32, (float32*)T_divide_red.rf[0], True, reduce_temp0, threadIdx.x, dtype=handle)

if (threadIdx.x == 0) {

T_divide_red[0] = (float32*)reduce_temp0[0]

}

计算的预定阶段已经累积，可以通过以下方式检查。

print(sg.stages)

输出：

[stage(a, placeholder(a, 0xd9c0fa00)), stage(b, placeholder(b, 0xe225cf70)), stage(T_add, compute(T_add, body=[(a[ax0, ax1, ax2] + b[ax1, ax2])], axis=[iter_var(ax0, range(min=0, ext=100)), iter_var(ax1, range(min=0, ext=10)), iter_var(ax2, range(min=0, ext=10))], reduce_axis=[], tag=broadcast, attrs={})), stage(T_multiply, compute(T_multiply, body=[(a[ax0, ax1, ax2]*b[ax1, ax2])], axis=[iter_var(ax0, range(min=0, ext=100)), iter_var(ax1, range(min=0, ext=10)), iter_var(ax2, range(min=0, ext=10))], reduce_axis=[], tag=broadcast, attrs={})), stage(T_elemwise_sum, compute(T_elemwise_sum, body=[(T_add[ax0, ax1, ax2] + T_multiply[ax0, ax1, ax2])], axis=[iter_var(ax0, range(min=0, ext=100)), iter_var(ax1, range(min=0, ext=10)), iter_var(ax2, range(min=0, ext=10))], reduce_axis=[], tag=elemwise, attrs={})), stage(T_divide, compute(T_divide, body=[(T_elemwise_sum[ax0, ax1, ax2]/2f)], axis=[iter_var(ax0, range(min=0, ext=100)), iter_var(ax1, range(min=0, ext=10)), iter_var(ax2, range(min=0, ext=10))], reduce_axis=[], tag=elemwise, attrs={})), stage(T_divide_red.rf, compute(T_divide_red.rf, body=[reduce(combiner=comm_reducer(result=[(x + y)], lhs=[x], rhs=[y], identity_element=[0f]), source=[T_divide[floordiv(floordiv((k0.k1.fused.k2.fused.inner + (k0.k1.fused.k2.fused.outer*1024)), 10), 10), floormod(floordiv((k0.k1.fused.k2.fused.inner + (k0.k1.fused.k2.fused.outer*1024)), 10), 10), floormod((k0.k1.fused.k2.fused.inner + (k0.k1.fused.k2.fused.outer*1024)), 10)]], init=[], axis=[iter_var(k0.k1.fused.k2.fused.outer, range(min=0, ext=10))], where=tir.likely((((floordiv(floordiv((k0.k1.fused.k2.fused.inner + (k0.k1.fused.k2.fused.outer*1024)), 10), 10) < 100) && (floordiv((k0.k1.fused.k2.fused.inner + (k0.k1.fused.k2.fused.outer*1024)), 10) < 1000)) && ((k0.k1.fused.k2.fused.inner + (k0.k1.fused.k2.fused.outer*1024)) < 10000))), value_index=0)], axis=[iter_var(k0.k1.fused.k2.fused.inner, range(min=0, ext=1024))], reduce_axis=[iter_var(k0.k1.fused.k2.fused.outer, range(min=0, ext=10))], tag=, attrs={})), stage(T_divide_red, compute(T_divide_red.repl, body=[reduce(combiner=comm_reducer(result=[(x + y)], lhs=[x], rhs=[y], identity_element=[0f]), source=[T_divide_red.rf[k0.k1.fused.k2.fused.inner.v]], init=[], axis=[iter_var(k0.k1.fused.k2.fused.inner.v, range(min=0, ext=1024))], where=(bool)1, value_index=0)], axis=[], reduce_axis=[iter_var(k0.k1.fused.k2.fused.inner.v, range(min=0, ext=1024))], tag=, attrs={}))]

可以通过与numpy结果进行比较，测试正确性，如下所示。

func = tvm.build(sg, [a, b, g], "cuda")

dev = tvm.cuda(0)

a_np = np.random.uniform(size=(x, y, y)).astype(a.dtype)

b_np = np.random.uniform(size=(y, y)).astype(b.dtype)

g_np = np.sum(np.add(a_np + b_np, a_np * b_np) / 2.0)

a_nd = tvm.nd.array(a_np, dev)

b_nd = tvm.nd.array(b_np, dev)

g_nd = tvm.nd.array(np.zeros(g_np.shape, dtype=g_np.dtype), dev)

func(a_nd, b_nd, g_nd)

tvm.testing.assert_allclose(g_nd.numpy(), g_np, rtol=1e-5)

TOPI 提供常用的神经网络操作，如 _softmax_ 优化调度

tarray = te.placeholder((512, 512), name="tarray")

softmax_topi = topi.nn.softmax(tarray)

with tvm.target.Target("cuda"):

sst = topi.cuda.schedule_softmax(softmax_topi)

print(tvm.lower(sst, [tarray], simple_mode=True))

输出：

primfn(tarray_1: handle) -> ()

attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}

buffers = {tarray: Buffer(tarray_2: Pointer(float32), float32, [512, 512], [])}

buffer_map = {tarray_1: tarray} {

allocate(T_softmax_norm: Pointer(global float32x4), float32x4, [65536]), storage_scope = global;

attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = 512;

allocate(normal_reduce_temp0: Pointer(local float32), float32, [1]), storage_scope = local;

allocate(reduce_temp0: Pointer(local float32), float32, [1]), storage_scope = local;

allocate(T_softmax_exp: Pointer(warp float32), float32, [512]), storage_scope = warp;

allocate(normal_reduce_temp0_1: Pointer(local float32), float32, [1]), storage_scope = local;

allocate(reduce_temp0_1: Pointer(local float32), float32, [1]), storage_scope = local {

attr [IterVar(threadIdx.x: int32, [0:32], "ThreadIndex", "threadIdx.x")] "thread_extent" = 32 {

normal_reduce_temp0[0] = -3.40282e+38f32

for (k.inner: int32, 0, 16) {

normal_reduce_temp0[0] = max((float32*)normal_reduce_temp0[0], (float32*)tarray_2[(((blockIdx.x*512) + (threadIdx.x*16)) + k.inner)])

}

attr [meta[tir.CommReducer][0]] "reduce_scope" = @tir.reinterpret(0u64, dtype=handle);

@tir.tvm_thread_allreduce(1u32, (float32*)normal_reduce_temp0[0], True, reduce_temp0, threadIdx.x, dtype=handle)

for (i1.inner.outer: int32, 0, 4) {

T_softmax_exp[ramp(((threadIdx.x*16) + (i1.inner.outer*4)), 1, 4)] = @tir.exp(((float32x4*)tarray_2[ramp((((blockIdx.x*512) + (threadIdx.x*16)) + (i1.inner.outer*4)), 1, 4)] - broadcast((float32*)reduce_temp0[0], 4)), dtype=float32x4)

}

attr [IterVar(threadIdx.x, [0:32], "ThreadIndex", "threadIdx.x")] "thread_extent" = 32 {

normal_reduce_temp0_1[0] = 0f32

for (k.inner_1: int32, 0, 16) {

normal_reduce_temp0_1[0] = ((float32*)normal_reduce_temp0_1[0] + (float32*)T_softmax_exp[((threadIdx.x*16) + k.inner_1)])

}

attr [meta[tir.CommReducer][1]] "reduce_scope" = @tir.reinterpret(0u64, dtype=handle);

@tir.tvm_thread_allreduce(1u32, (float32*)normal_reduce_temp0_1[0], True, reduce_temp0_1, threadIdx.x, dtype=handle)

for (i1.inner.outer_1: int32, 0, 4) {

T_softmax_norm[ramp((((blockIdx.x*512) + (threadIdx.x*16)) + (i1.inner.outer_1*4)), 1, 4)] = ((float32x4*)T_softmax_exp[ramp(((threadIdx.x*16) + (i1.inner.outer_1*4)), 1, 4)] / broadcast((float32*)reduce_temp0_1[0], 4))

}

融合卷积

可以融合topi.nn.conv2d和topi.nn.relu在一起。

TOPI 函数都是通用函数。对不同的后端，有不同的实现优化性能。对于每个后端，有必要在计算声明和调度的目标范围内调用。TVM 将选择正确的函数，调用目标信息。

data = te.placeholder((1, 3, 224, 224))

kernel = te.placeholder((10, 3, 5, 5))

with tvm.target.Target("cuda"):

conv = topi.cuda.conv2d_nchw(data, kernel, 1, 2, 1)

out = topi.nn.relu(conv)

sconv = topi.cuda.schedule_conv2d_nchw([out])

print(tvm.lower(sconv, [data, kernel], simple_mode=True))

Out:

primfn(placeholder_2: handle, placeholder_3: handle) -> ()

attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}

buffers = {placeholder_1: Buffer(placeholder_4: Pointer(float32), float32, [10, 3, 5, 5], []),

placeholder: Buffer(placeholder_5: Pointer(float32), float32, [1, 3, 224, 224], [])}

buffer_map = {placeholder_2: placeholder, placeholder_3: placeholder_1} {

allocate(compute: Pointer(global float32), float32, [501760]), storage_scope = global;

attr [IterVar(blockIdx.z: int32, (nullptr), "ThreadIndex", "blockIdx.z")] "thread_extent" = 5;

allocate(compute_1: Pointer(local float32), float32, [14]), storage_scope = local;

allocate(pad_temp.shared: Pointer(shared float32), float32, [112]), storage_scope = shared;

allocate(placeholder.shared: Pointer(shared float32), float32, [2]), storage_scope = shared;

attr [IterVar(blockIdx.y: int32, (nullptr), "ThreadIndex", "blockIdx.y")] "thread_extent" = 224;

attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = 2;

attr [IterVar(threadIdx.z: int32, (nullptr), "ThreadIndex", "threadIdx.z")] "thread_extent" = 1;

attr [IterVar(threadIdx.y: int32, (nullptr), "ThreadIndex", "threadIdx.y")] "thread_extent" = 1;

attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 16 {

compute_1[0] = 0f32

compute_1[2] = 0f32

compute_1[4] = 0f32

compute_1[6] = 0f32

compute_1[8] = 0f32

compute_1[10] = 0f32

compute_1[12] = 0f32

compute_1[1] = 0f32

compute_1[3] = 0f32

compute_1[5] = 0f32

compute_1[7] = 0f32

compute_1[9] = 0f32

compute_1[11] = 0f32

compute_1[13] = 0f32

for (rc.outer: int32, 0, 3) {

for (ry.outer: int32, 0, 5) {

attr [IterVar(threadIdx.z_1: int32, (nullptr), "ThreadIndex", "threadIdx.z")] "thread_extent" = 1;

attr [IterVar(threadIdx.y_1: int32, (nullptr), "ThreadIndex", "threadIdx.y")] "thread_extent" = 1;

attr [IterVar(threadIdx.x_1: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 16 {

pad_temp.shared[(threadIdx.x_1*7)] = @tir.if_then_else((((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)) && (2 <= ((blockIdx.x*112) + (threadIdx.x_1*7)))), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 450)], 0f32, dtype=float32)

pad_temp.shared[((threadIdx.x_1*7) + 1)] = @tir.if_then_else((((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)) && (1 <= ((blockIdx.x*112) + (threadIdx.x_1*7)))), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 449)], 0f32, dtype=float32)

pad_temp.shared[((threadIdx.x_1*7) + 2)] = @tir.if_then_else(((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 448)], 0f32, dtype=float32)

pad_temp.shared[((threadIdx.x_1*7) + 3)] = @tir.if_then_else(((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 447)], 0f32, dtype=float32)

pad_temp.shared[((threadIdx.x_1*7) + 4)] = @tir.if_then_else(((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 446)], 0f32, dtype=float32)

pad_temp.shared[((threadIdx.x_1*7) + 5)] = @tir.if_then_else(((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 445)], 0f32, dtype=float32)

pad_temp.shared[((threadIdx.x_1*7) + 6)] = @tir.if_then_else(((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 444)], 0f32, dtype=float32)

}

attr [IterVar(threadIdx.z_2: int32, (nullptr), "ThreadIndex", "threadIdx.z")] "thread_extent" = 1;

attr [IterVar(threadIdx.y_2: int32, (nullptr), "ThreadIndex", "threadIdx.y")] "thread_extent" = 1;

attr [IterVar(threadIdx.x_2: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 16;

if @tir.likely((threadIdx.x_2 < 2), dtype=bool) {

placeholder.shared[threadIdx.x_2] = (float32*)placeholder_4[((((blockIdx.z*150) + (threadIdx.x_2*75)) + (rc.outer*25)) + (ry.outer*5))]

}

compute_1[0] = ((float32*)compute_1[0] + ((float32*)pad_temp.shared[threadIdx.x]*(float32*)placeholder.shared[0]))

compute_1[2] = ((float32*)compute_1[2] + ((float32*)pad_temp.shared[(threadIdx.x + 16)]*(float32*)placeholder.shared[0]))

compute_1[4] = ((float32*)compute_1[4] + ((float32*)pad_temp.shared[(threadIdx.x + 32)]*(float32*)placeholder.shared[0]))

compute_1[6] = ((float32*)compute_1[6] + ((float32*)pad_temp.shared[(threadIdx.x + 48)]*(float32*)placeholder.shared[0]))

compute_1[8] = ((float32*)compute_1[8] + ((float32*)pad_temp.shared[(threadIdx.x + 64)]*(float32*)placeholder.shared[0]))

compute_1[10] = ((float32*)compute_1[10] + ((float32*)pad_temp.shared[(threadIdx.x + 80)]*(float32*)placeholder.shared[0]))

compute_1[12] = ((float32*)compute_1[12] + ((float32*)pad_temp.shared[(threadIdx.x + 96)]*(float32*)placeholder.shared[0]))

compute_1[1] = ((float32*)compute_1[1] + ((float32*)pad_temp.shared[threadIdx.x]*(float32*)placeholder.shared[1]))

compute_1[3] = ((float32*)compute_1[3] + ((float32*)pad_temp.shared[(threadIdx.x + 16)]*(float32*)placeholder.shared[1]))

compute_1[5] = ((float32*)compute_1[5] + ((float32*)pad_temp.shared[(threadIdx.x + 32)]*(float32*)placeholder.shared[1]))

compute_1[7] = ((float32*)compute_1[7] + ((float32*)pad_temp.shared[(threadIdx.x + 48)]*(float32*)placeholder.shared[1]))

compute_1[9] = ((float32*)compute_1[9] + ((float32*)pad_temp.shared[(threadIdx.x + 64)]*(float32*)placeholder.shared[1]))

compute_1[11] = ((float32*)compute_1[11] + ((float32*)pad_temp.shared[(threadIdx.x + 80)]*(float32*)placeholder.shared[1]))

compute_1[13] = ((float32*)compute_1[13] + ((float32*)pad_temp.shared[(threadIdx.x + 96)]*(float32*)placeholder.shared[1]))

attr [IterVar(threadIdx.z_1, (nullptr), "ThreadIndex", "threadIdx.z")] "thread_extent" = 1;

attr [IterVar(threadIdx.y_1, (nullptr), "ThreadIndex", "threadIdx.y")] "thread_extent" = 1;

attr [IterVar(threadIdx.x_1, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 16 {

pad_temp.shared[(threadIdx.x_1*7)] = @tir.if_then_else((((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)) && (1 <= ((blockIdx.x*112) + (threadIdx.x_1*7)))), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 449)], 0f32, dtype=float32)

pad_temp.shared[((threadIdx.x_1*7) + 1)] = @tir.if_then_else(((2 <= (blockIdx.y + ry.outer)) && ((blockIdx.y + ry.outer) < 226)), (float32*)placeholder_5[((((((rc.outer*50176) + (blockIdx.y*224)) + (ry.outer*224)) + (blockIdx.x*112)) + (threadIdx.x_1*7)) - 448)], 0f32, dtype=float32)

}

attr [IterVar(threadIdx.z_2, (nullptr), "ThreadIndex", "threadIdx.z")] "thread_extent" = 1;

attr [IterVar(threadIdx.y_2, (nullptr), "ThreadIndex", "threadIdx.y")] "thread_extent" = 1;

attr [IterVar(threadIdx.x_2, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 16;

if @tir.likely((threadIdx.x_2 < 2), dtype=bool) {

placeholder.shared[threadIdx.x_2] = (float32*)placeholder_4[(((((blockIdx.z*150) + (threadIdx.x_2*75)) + (rc.outer*25)) + (ry.outer*5)) + 1)]