内容隐藏

3.1. topi.add

3.2. topi.nn.relu

3.3. tvm::te::compute

topi 是 tvm 的一个张量算子库，提供了很多常见的算子操作，例如 conv2D，transpose，等等。

tvm 在编译计算图的时候，会先将计算图从 relay ir 翻译成 tir，再将 tir 翻译成目标设备代码（比如llvm，比如c，比如cuda），但是 relay ir 和 tir 之间还有一种中间语言叫 tensor expression，这是一种专门为张量计算而设计的语言。relay ir 的绝大多数部分 op 其实都是通过 tensor expression 这种语言来实现的

https://tvm.apache.org/docs/reference/langref/relay_expr.html

An operator is a primitive operation, such as add or conv2d, not defined in the Relay language. Operators are declared in the global operator registry in C++. Many common operators are backed by TVM’s Tensor Operator Inventory.

topi 和 te 的关系：

te( tensor expression ) 是一门函数式编程语言，面向张量计算的
topi 是基于 te 实现的张量算子库

类似于编程语言和标准库之间的关系

1. topi 算子列表

源码目录：src/topi

src/topi 只是把算子通过 TVM_REGISTER_GLOBAL 注册到 tvm 的全局函数中，并没有太多细节

函数的实现，全部是以 inline 函数的形式，实现在 include/tvm/topi 中，由于算子都是通过 tensor expression 来实现的，因此所有算子的实现最终都会调用 tvm::te::compute 来进行计算

最常见的有4类操作：broadcast，elemwise，nn，reduction

broadcast	elemwise	nn	reduction
topi.add topi.subtract topi.multiply topi.divide topi.floor_divide topi.mod topi.floor_mod topi.maximum topi.minimum topi.power topi.left_shift topi.logical_and topi.logical_or topi.logical_xor topi.bitwise_and topi.bitwise_or topi.bitwise_xor topi.right_shift topi.greater topi.less topi.equal topi.not_equal topi.greater_equal topi.less_equal topi.broadcast_to	topi.acos topi.acosh topi.asin topi.asinh topi.atanh topi.exp topi.fast_exp topi.erf topi.fast_erf topi.tan topi.cos topi.cosh topi.sin topi.sinh topi.tanh topi.fast_tanh topi.atan topi.sigmoid topi.sqrt topi.rsqrt topi.log topi.log2 topi.log10 topi.identity topi.negative topi.clip topi.cast topi.reinterpret topi.elemwise_sum topi.sign topi.full topi.full_like topi.logical_not topi.bitwise_not	topi.nn.relu topi.nn.leaky_relu topi.nn.prelu topi.nn.pad topi.nn.space_to_batch_nd topi.nn.batch_to_space_nd topi.nn.nll_loss topi.nn.dense topi.nn.bias_add topi.nn.dilate topi.nn.flatten topi.nn.scale_shift_nchw topi.nn.scale_shift_nhwc topi.nn.pool_grad topi.nn.global_pool topi.nn.adaptive_pool topi.nn.adaptive_pool3d topi.nn.pool1d topi.nn.pool2d topi.nn.pool3d topi.nn.softmax topi.nn.log_softmax topi.nn.lrn topi.nn.binarize_pack topi.nn.binary_dense	topi.sum topi.min topi.max topi.argmin topi.argmax topi.prod topi.all topi.any

2. topi 编程示例

由于 topi 实现的所有算子，最终都会注册到 tvm 的全局函数中，并已 relay.op 的形式存在

因此，我们在描述计算图的时候，可以直接调用相关的算子，相当于我们平时写 c 代码的时候，有很多编译器的 builtin 函数一样

比如 topi.add 对应的 relay.op 是 relay.add，其他类推

a1 = relay.var("a1", shape=(1,), dtype="float32")
c1 = relay.const(10, 'float32')
c2 = relay.add(c1, a1)

3. topi 算子实现

topi 算子的代码，主要是分在3个地方：

include/tvm/topi/：虽说是头文件，但是为了优化性能，其实 topi 的绝大多数算子都是以 inline 形式在头文件里直接实现的（当然底层绝大部分工作还是 tvm::te::compute）
src/topi/：将 topi 算子注册到 tvm 的全局作用域上（让 c++ relay 或者 python 可以很方便调用）
src/relay/op 和 include/tvm/relay/op 把 topi 算子封装成 relay op

3.1. topi.add

add 函数，最主要的就是实现2个 tensor 的相加。在深度学习里面，所有的数据都是以 tensor 的方式来描述的

topi.add 的实现主要靠2个关键的宏，如下：

TOPI_DEFINE_BCAST_OP(add, { return a + b; });
TOPI_DEFINE_OP_OVERLOAD(operator+, add);

1）TOPI_DEFINE_BCAST_OP

TOPI_DEFINE_BCAST_OP 宏定义了4类操作：

PrimExpr + PrimExpr -> PrimExpr
Tensor + Tensor -> Tensor
Tensor + PrimExpr -> Tensor
PrimExpr + Tensor -> Tensor

如下：

#define TOPI_DEFINE_BCAST_OP(Name, ComputeRule)                                                   \
  inline tvm::PrimExpr Name(const tvm::PrimExpr& a, const tvm::PrimExpr& b) { ComputeRule; }      \
  inline tvm::te::Tensor Name(const tvm::te::Tensor& A, const tvm::te::Tensor& B,                 \
                              std::string name = "T_" #Name, std::string tag = kBroadcast) {      \
    auto l = [](tvm::PrimExpr a, tvm::PrimExpr b) { ComputeRule; };                               \
    return detail::WithBroadcast(l, A, B, name, tag);                                             \
  }                                                                                               \
  inline tvm::te::Tensor Name(const tvm::te::Tensor& A, const tvm::PrimExpr& B,                   \
                              std::string name = "T_" #Name, std::string tag = kElementWise) {    \
    auto l = [](tvm::PrimExpr a, tvm::PrimExpr b) { ComputeRule; };                               \
    return tvm::te::compute(                                                                      \
        A->shape, [&](const ::tvm::Array<::tvm::tir::Var>& i) { return l(A(i), B); }, name, tag); \
  }                                                                                               \
  inline tvm::te::Tensor Name(const tvm::PrimExpr& A, const tvm::te::Tensor& B,                   \
                              std::string name = "T_" #Name, std::string tag = kElementWise) {    \
    auto l = [&](tvm::PrimExpr a, tvm::PrimExpr b) { ComputeRule; };                              \
    return tvm::te::compute(                                                                      \
        B->shape, [&](const ::tvm::Array<::tvm::tir::Var>& i) { return l(A, B(i)); }, name, tag); \
  }

TOPI_DEFINE_BCAST_OP 的第一个参数是操作符名字，这里是 add，第二个参数是计算规则

2）TOPI_DEFINE_OP_OVERLOAD

这个宏主要用来实现 tensor 的操作符重载，比如 TOPI_DEFINE_OP_OVERLOAD(operator+, add) 展开得到

inline tvm::te::Tensor operator+(const tvm::te::Tensor& A, const tvm::te::Tensor& B) {
  return topi::add(A, B);
}
inline tvm::te::Tensor operator+(const tvm::PrimExpr& A, const tvm::te::Tensor& B) {
  return topi::add(A, B);
}

操作符重载是一个c++的常见用法，我们实现2个 tensor 相加的时候，不用直接写 topi.add(a, b)，而是可以直接写 a + b

3.2. topi.nn.relu

代码：include/tvm/topi/nn.h

nn.relu 是一个神经网络里常用的激活函数，f(x) = max(0, x)

/*!
 * \brief Creates an operation that performs a rectified linear unit
 *
 * \param t The input tensor
 * \param threshold The relu threshold (default 0)
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the relu operation
 */
template 
tvm::te::Tensor relu(const tvm::te::Tensor& t, T threshold = static_cast(0),
                            std::string name = "T_relu", std::string tag = kElementWise) {
    sleep(100);
    return tvm::te::compute(
      t->shape,
      [&](const tvm::Array& i) {
        auto threshold_const = tvm::tir::make_const(t->dtype, threshold);
        return tvm::max(t(i), threshold_const);
      },
      name, tag);
}

threshold 默认是 0，这里 tvm::max 取 max 值

3.3. tvm::te::compute

无论是 topi.add 还是 topi.nn.relu，最终都会调用 tvm::te::compute 返回一个新的 Tensor 代表计算的结果

但是这里并不执行真正的计算，或者这里的计算并不是我们所理解的那个常规意义上的计算

Tensor compute(Array shape, FCompute fcompute, std::string name, std::string tag,
               Map<String, ObjectRef> attrs) {
  // compute dimension.
  size_t ndim = shape.size();
  std::vector axis;
  std::vector args;
  for (size_t i = 0; i < ndim; ++i) {
    std::ostringstream os;
    os << "ax" << i; axis.emplace_back(IterVar(Range(0, shape[i]), Var(os.str(), shape[i].dtype()), kDataPar)); args.push_back(axis.back()->var);
  }

  return ComputeOp(name, tag, attrs, axis, {fcompute(args)}).output(0);
}

这里的 Fcompute 就是我们前面的：

auto l = [&](tvm::PrimExpr a, tvm::PrimExpr b) { return a + b; };

这个也是：

[&](const tvm::Array& i) {
    auto threshold_const = tvm::tir::make_const(t->dtype, threshold);
    return tvm::max(t(i), threshold_const);
},

fcompute(args) 会返回一个 PrimExpr，实际上可以理解为是一个 tir 算子。

注意这里的 tvm::max 的 c++ namespace，像这种基础运算 tvm 里有很多同名的实现，比如 tvm::max 和 tvm::topi::max，前者实现在 tir 里，返回一个 PrimExpr，后者实现在 te 里，返回一个 Tensor

后面深入再了解一下

4. topi 的编译(lower)过程

relay op 到 topi 的编译过程在 tvm 里叫 lower

最前面我们讲编译的基本流程的时候，有一个 codegen 的过程，topi 就是在里面处理的

简单来说，就是在编译 relay 计算图的时候，发现如果 relay op 是一个内置的 topi 的话，直接转而变成 topi 的具体实现

为了看下这个具体的过程，我在 topi.add 的实现里增加了一个 sleep(100)，用 gdb 得到完整的调用栈

TOPI_DEFINE_BCAST_OP(add, { sleep(100); return a + b; });

测试代码：

a1 = relay.var("a1", shape=(1,), dtype="float32")
a2 = relay.var("a2", shape=(1,), dtype="float32")
add_op = relay.add(a1, a2)  // -> topi.add 算子
f1 = relay.Function([a1, a2], add_op)
mod = tvm.IRModule({'main': f1})
graph, lib, params = relay.build(mod, 'llvm')

调用栈如下（省略非关键的）：

35: tvm::relay::backend::RelayBuildModule::BuildRelay(tvm::IRModule, tvm::runtime::String const&)

33: tvm::relay::backend::GraphExecutorCodegen::Codegen(tvm::IRModule, tvm::relay::Function, tvm::runtime::String)

26: tvm::relay::tec::LowerTE(tvm::IRModule const&, tvm::runtime::String const&, std::function, tvm::CompilationConfig)

13: tvm::relay::transform::DeviceAwareExprMutator::VisitExpr_(tvm::relay::CallNode const*)

12: tvm::relay::tec::LowerTensorExprMutator::DeviceAwareVisitExpr_(tvm::relay::CallNode const*)

tec::LowerTensorExprMutator::DeviceAwareVisitExpr_

这里对 CallNode 的处理分很多 case，但其实最关键的其实就这地方：

Expr DeviceAwareVisitExpr_(const CallNode* call_node) override {
    ...
    BaseFunc primitive_func = ResolveToPrimitive(call_node->op);
    if (primitive_func->HasNonzeroAttr(attr::kExtern)) {
      // ...
    } else {
      // Cases 1 and 2: lower the primitive function for the desired target, possibly using external
      // codegen.
      CCacheKey key(Downcast(primitive_func), target,
                    GetVirtualDevice(GetRef(call_node)));
      CachedFunc cfunc = compiler_->Lower(key);   // 这里
      return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
                             call_node->span, target, cfunc->funcs->functions);
    }
}

compiler_->Lower(key); 就会完成 relay op 到 topi 的转换，compiler_ 是一个 TECompiler 实例，用来做 topi 算子的编译

注意这里的 key 会包含 target，因为同一个算子在不同的平台上，有不同的实现，所以需要考虑target信息

继续跟踪 TECompiler::Lower 的实现

8: tvm::relay::tec::TECompilerImpl::Lower(tvm::relay::tec::CCacheKey const&)

7: tvm::relay::tec::TECompilerImpl::LowerInternal(tvm::relay::tec::CCacheKey const&, tvm::GlobalVarSupply)

6: tvm::relay::tec::PrimFuncFor(tvm::relay::Function const&, tvm::Target const&, tvm::GlobalVarSupply)

5: tvm::relay::tec::ScheduleBuilder::Create(tvm::relay::Function const&, tvm::GlobalVarSupply)

4: tvm::relay::tec::LowerToTECompute::Lower(tvm::relay::Function const&)

LowerInternal 会调用 PrimFuncFor，而 PrimFuncFor 只是一个简单的 ScheduleBuilder 封装

CachedFunc ScheduleBuilder::Create(const Function& relay_func, GlobalVarSupply global_var_supply) {
    LowerToTECompute lower_te_compute(target_);
    Array tensor_outs = lower_te_compute.Lower(relay_func);
 // 这里
    ...
}

继续跟踪 LowerToTECompute，我们只需要关注这个类的 VisitExpr_(const CallNode* call_node) 方法

Array LowerToTECompute::VisitExpr_(const CallNode* call_node) final {
    static auto flower_call = tvm::runtime::Registry::Get("relay.backend.lower_call");
    ...
    LoweredOutput lowered_out = (*flower_call)(GetRef(call_node), inputs, target_);
    ...
}

在这个 VisitExpr_ 里就会去调用真正的 topi.add 的实现

relay.backend.lower_call 不是一个 c++ 实现的函数，是 python 实现的。（这是 tvm 比较牛逼的一个地方，很多函数 c++ & python 是可以互相调用的）

lower_call 会根据 call_node 的算子名称、target_，找到对应的具体的 topi.add 实现，细节可以看 lower_call 里的 select_implementation

前面我们知道 topi 的所有算子库，都是会通过 TVM_REGISTER_GLOBAL 注册到 tvm 的全局作用域上，select_implementation 就是从这里找到对应的 topi.add 实现的

5. topi 扩展自定义算子

参考算子实现，其实我们大概也就知道怎么扩展一个新的算子，并注册到 relay op 中了

以 topi.add2 为例，其函数功能和 topi.add 一致

5.1. 算子函数

由于 topi.add2 函数和 topi.add 是一样的，这个函数本身我们就不实现了

直接注册即可：

TOPI_REGISTER_BCAST_OP(“topi.add2”, topi::add);

这里其实如果我们不打算从 python 代码里直接调用 topi.add2 的话，这个注册其实可有可无

5.2. 注册 relay 算子

src/relay/op/tensor/binary.cc

增加如下代码：

RELAY_REGISTER_BINARY_OP("add2")
    .describe("Elementwise add with broadcasting")
    .set_support_level(1)
    .set_attr("FTVMCompute", RELAY_BINARY_COMPUTE(topi::add));

5.3. python 绑定

为了让 python 代码可以调用到这个算子，我们需要在 python 前端中简单声明一下这个函数，并通过 tvm 的底层机制调用 c++ 的实现

python/tvm/relay/op/tensor.py

def add2(lhs, rhs):
    ...
    return _make.add2(lhs, rhs)

5.4. 注册计算和调度函数

tvm 的计算和调度是分离的，我们还需要为 add2 注册调度策略

python/tvm/relay/op/_tensor.py

register_shape_func("add2", False, broadcast_shape_func)
register_broadcast_schedule("add2")

一	二	三	四	五	六	日
			1	2	3	4
5	6	7	8	9	10	11
12	13	14	15	16	17	18
19	20	21	22	23	24	25
26	27	28	29	30

成功，源于对美学的执著追求

深入浅出 tvm – (15) TVM Operator Inventory (TOPI)

1. topi 算子列表

2. topi 编程示例

3. topi 算子实现

3.1. topi.add

3.2. topi.nn.relu

3.3. tvm::te::compute

4. topi 的编译(lower)过程

5. topi 扩展自定义算子

5.1. 算子函数

5.2. 注册 relay 算子

5.3. python 绑定

5.4. 注册计算和调度函数

发表回复取消回复

成功，源于对美学的执著追求

1. topi 算子列表

2. topi 编程示例

3. topi 算子实现

3.1. topi.add

3.2. topi.nn.relu

3.3. tvm::te::compute

4. topi 的编译(lower)过程

5. topi 扩展自定义算子

5.1. 算子函数

5.2. 注册 relay 算子

5.3. python 绑定

5.4. 注册计算和调度函数

发表回复 取消回复

发表回复取消回复