TVM中的Compute操作

定义

TVM从Halide继承了计算与调度分离的思想，并在其内部重用了部分Halide的调度原语，也引入了一些新的调度原语，用于优化GPU和专用加速器性能。

先举个例子吧：

import tvm
from tvm import te

n = 1024
dtype = "float32"
A = te.placeholder((n, n), dtype=dtype, name='A')
K = te.reduce_axis((0, n), name='k')
B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name='B')
S = te.create_schedule(B.op)

• Compute(计算)：在tvm中，te.compute() 接口的作用是通过指定的计算规则，在指定的矩阵形状上构造一个新的张量，它并不会触发实际的计算，而是采用基于lambda表达式的张量表达式进行张量计算，即仅声明了进行何种计算。

• Schedule(调度)：用于指定计算操作的顺序和方式，如通过Schedule对象，对计算操作进行调度，包括并行化、优化内存访问模式、设备特定优化等。Schedule提供了丰富的接口和方法，可以帮助优化和控制计算的执行方式，以提高性能和效率。TVM中通过te.create_schedule()创建调度，该接口返回的是tvm.te.schedule.Schedule类对象，其中包含了若干个阶段，每个阶段对应一个描述调度方式的Compute(操作)。

简单说，Compute 主要用于定义计算图中计算操作之间的关系，而 Schedule 则用于调度和优化这些计算操作的执行方式

作用

这种分离的设计是为了提高计算图的灵活性、可扩展性和性能

• 模块化设计：将compute和Schedule分开可以使得用户更容易理解和管理计算图中的操作，降低复杂度。compute主要关注计算单元的定义，而Schedule则负责对计算进行调度和优化，这种模块化设计简化了代码结构，提高了可维护性。

• 优化灵活性：通过将计算和调度分隔开来，用户可以根据不同的需求针对计算和调度分别进行优化。这种灵活性使得用户可以更轻松地测试不同的优化策略，找到最佳的性能表现。

• 性能提升：分开compute和Schedule也有助于实现更高效的编译和优化过程。通过将计算和调度分别处理，TVM可以更精确地对计算进行优化，从而达到更好的性能表现

怎么做的

本节主要来看下TVM中Compute是如何实现的

代码实现见：python/tvm/te/operation.py

def compute(shape, fcompute, name="compute", tag="", attrs=None, varargs_names=None):
    if _tag.TagScope.get_current() is not None:
        if tag != "":
            raise ValueError("nested tag is not allowed for now")
        tag = _tag.TagScope.get_current().tag
    shape = (shape,) if isinstance(shape, tvm.tir.PrimExpr) else shape
    # for python3
    shape = tuple([int(s) if isinstance(s, float) else s for s in shape])
    out_ndim = len(shape)

    argspec = inspect.getfullargspec(fcompute)
    if len(argspec.args) == 0 and argspec.varargs is None:
        arg_names = [f"i{i}" for i in range(out_ndim)]
    elif argspec.varargs is not None:
        # if there is a varargs, it takes the remaining dimensions of out_ndim
        num_remaining_args = out_ndim - len(argspec.args)
        if varargs_names is not None:
            if len(varargs_names) != num_remaining_args:
                raise RuntimeError(
                    f"Number of varargs ({num_remaining_args}) does not match number"
                    f"of varargs_names ({len(varargs_names)})"
                )
            arg_names = argspec.args + varargs_names
        else:
            arg_names = argspec.args + [f"i{i}" for i in range(out_ndim - len(argspec.args))]
    else:
        arg_names = argspec.args
        # if there are fewer args than out dimensions, the remaining dimensions
        # are implicitly broadcast
        out_ndim = len(arg_names)
    assert argspec.varkw is None, "Variable keyword arguments not supported in fcompute"
    assert argspec.defaults is None, "Default arguments not supported in fcompute"
    assert len(argspec.kwonlyargs) == 0, "Keyword arguments are not supported in fcompute"

    if out_ndim != len(arg_names):
        raise ValueError(
            "Number of args to fcompute does not match dimension, "
            f"args={len(arg_names)}, dimension={out_ndim}"
        )

    dim_var = [tvm.tir.IterVar((0, s), x, 0) for x, s in zip(arg_names, shape[:out_ndim])]
    body = fcompute(*[v.var for v in dim_var])

    if isinstance(body, _tensor.TensorIntrinCall):
        for i, s in enumerate(shape[out_ndim:]):
            var_name = "ax" + str(i)
            dim_var.append(tvm.tir.IterVar((0, s), var_name, 4))
        op_node = _ffi_api.TensorComputeOp(
            name,
            tag,
            dim_var,
            body.reduce_axis,
            out_ndim,
            body.intrin,
            body.tensors,
            body.regions,
            body.scalar_inputs,
        )
    else:
        if not isinstance(body, (list, tuple)):
            body = [body]
        body = convert(body)
        op_node = _ffi_api.ComputeOp(name, tag, attrs, dim_var, body)

    num = op_node.num_outputs
    outputs = tuple(op_node.output(i) for i in range(num))
    return outputs[0] if num == 1 else outputs

fcompute(*[v.var for v in dim_var])会用dim_var替换compute.op的lambda表达式中的数值。

TVM Compute操作有ComputeOp()和TensorComputeOp()

ComputeOp()：每次操作一个标量的计算操作

TensorComputeOp()：每次操作一个张量切片的计算操作

通过FFI接口调用C++端实现，可分别返回ComputeOp对象 和TensorComputeOp对象

ComputeOp()实现如下：

ComputeOp::ComputeOp(std::string name, std::string tag, Map<String, ObjectRef> attrs,
                     Array<IterVar> axis, Array<PrimExpr> body) {
  if (!attrs.defined()) {
    attrs = Map<String, ObjectRef>();
  }
  auto n = make_object<ComputeOpNode>();
  n->name = std::move(name);
  n->tag = std::move(tag);
  n->attrs = std::move(attrs);
  n->axis = std::move(axis);
  n->body = std::move(body);
  if (n->body[0]->IsInstance<tir::ReduceNode>()) {
    const tir::ReduceNode* reduce = n->body[0].as<tir::ReduceNode>();
    n->reduce_axis = reduce->axis;
  }
  VerifyComputeOp(n.get());
  data_ = std::move(n);
}
// 注册
TVM_REGISTER_GLOBAL("te.ComputeOp")
    .set_body_typed([](std::string name, std::string tag, Map<String, ObjectRef> attrs,
                       Array<IterVar> axis,
                       Array<PrimExpr> body) { return ComputeOp(name, tag, attrs, axis, body); });

data_为成员变量，定义如下：

ObjectPtr<Object> data_;

此处，对ComputeOp做继续分析，TensorComputeOp类似

在ComputeOp对象调用 num_outputs 时，实现如下：

int ComputeOpNode::num_outputs() const { return body.size(); }

另：类ComputeOp继承自类Operation

当在调用对象的 output()方法 时，调用的是父类Operation方法，实现如下(src/te/tensor.cc)：

Tensor Operation::output(size_t i) const {
  auto node = make_object<TensorNode>();
  node->op = *this;
  node->value_index = i;
  node->dtype = (*this)->output_dtype(i);
  node->shape = (*this)->output_shape(i);
  return Tensor(node);
}

通过上述代码，可将te.compute可以总结为如下几步:

1. 根据传入的fcompute,翻译成对应的表达式传入
2. 生成ComputeOpNode,记录计算节点.
3. 根据计算节点,返回计算节点输出对应的Tensor

Respect～