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可以总结为如下几步:
- 根据传入的fcompute,翻译成对应的表达式传入
- 生成ComputeOpNode,记录计算节点.
- 根据计算节点,返回计算节点输出对应的Tensor
Respect～
