https://github.com/apache/tvm/pull/8909/files
TOPI中的算子分为其算子的定义和算子的schedule两部分。算子的定义是唯一的,而对于不同的后端(x86、cuda等)schedule可以不同。可以不定义schedule,此时一切按照算子的定义执行且不加入任何优化,执行性能较差。
python/tvm/relay/op/strategy/目录下定义了一系列算子+调度的组合策略(strategy)。它们属于relay中的op。
比如,python/tvm/relay/op/strategy/x86.py是专为x86后端设计的策略,而python/tvm/relay/op/strategy/cuda.py则是为了cuda后端:
@softmax_strategy.register("cpu")
def softmax_strategy_cpu(attrs, inputs, out_type, target):
"""softmax x86 strategy"""
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_softmax(topi.nn.softmax),
wrap_topi_schedule(topi.x86.schedule_softmax),
name="softmax.x86",
)
return strategy
@softmax_strategy.register(["cuda", "gpu"])
def softmax_strategy_cuda(attrs, inputs, out_type, target):
"""softmax cuda strategy"""
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_softmax(topi.nn.softmax),
wrap_topi_schedule(topi.cuda.schedule_softmax),
name="softmax.cuda",
)
if target.kind.name == "cuda" and "cudnn" in target.libs:
strategy.add_implementation(
wrap_compute_softmax(topi.cuda.softmax_cudnn),
wrap_topi_schedule(topi.cuda.schedule_softmax_cudnn),
name="softmax.cudnn",
plevel=15,
)
return strategy
可以观察到,wrap_compute_softmax的参数都是topi.nn.softmax,而wrap_topi_schedule包裹的调度函数则不同。这直接体现了TVM计算和调度分离的核心思想。
调度不仅仅需要考虑当前算子本身如何高效执行,还需要考虑算子融合,即当前算子与其他算子融合之后的执行是否高效。这让调度的开发变得更加复杂。
例如:
python/tvm/topi/x86/dense.py
def dense_vnni_schedule(cfg, s, C, O, do_parallel=True):
"""Schedule dense compute using VNNI vpdpbusd instruction"""
# C: The output of GEMM
# O: The output of the fused op
def split_y(out):
default_y_split_factor = 32
a_y = out.op.axis[-2]
if cfg.is_fallback:
return s[out].split(a_y, factor=default_y_split_factor)
return cfg["tile_y"].apply(s, out, a_y)
(a_k,) = C.op.reduce_axis
a_yo, a_yi = split_y(C)
a_xo, a_xi = s[C].split(C.op.axis[-1], factor=16)
a_ko, a_ki = s[C].split(a_k, factor=4)
s[C].reorder(a_yo, a_xo, a_yi, a_ko, a_xi, a_ki)
pc = dot_16x1x16_uint8_int8_int32_cascadelake()
s[C].tensorize(a_xi, pc)
if C == O:
fused = s[O].fuse(a_yo, a_xo)
else:
a_yo, a_yi = split_y(O)
a_xo, a_xi = s[O].split(O.op.axis[-1], factor=16)
s[O].reorder(a_yo, a_xo, a_yi, a_xi)
s[O].vectorize(a_xi)
s[C].compute_at(s[O], a_yi)
fused = s[O].fuse(a_yo, a_xo)
if do_parallel:
s[O].parallel(fused)
return s, fused
把C≠O作为判断是否是fuse op的条件:C==O则当前算子(dense)是未被融合的算子,反之,算子之后还存在某个依赖于当前算子计算结果的算子,此时C≠O。
另外,vnni dense的调度策略可以看这篇。
s[C].tensorize(a_xi, pc)
调用张量化调度策略。而在C≠O时,需要将输出变量O的调度安排到C上:
a_yo, a_yi = split_y(O)
a_xo, a_xi = s[O].split(O.op.axis[-1], factor=16)
s[O].reorder(a_yo, a_xo, a_yi, a_xi)
s[O].vectorize(a_xi)
s[C].compute_at(s[O], a_yi)
fused = s[O].fuse(a_yo, a_xo)
简单来说,就是在最内层的乘加(vpdpbusd)后加上另一个被融合的算子(例如biais_add):
vpdpbusd + biais_add
test case中就测试了dense+biais_add的情况:
tests/python/relay/test_op_level1.py
@pytest.mark.skip("Requires cascadelake")
def test_dense_vnni():
data_shape = (32, 96)
weight_shape = (128, 96)
for data_dtype in ["uint8", "int8"]:
data = relay.var("data", shape=data_shape, dtype=data_dtype)
weight = relay.var("weight", shape=weight_shape, dtype="int8")
bias = relay.var("bias", shape=(weight_shape[0],), dtype="int32")
dense = relay.nn.dense(data, weight, out_dtype="int32")
out = relay.nn.bias_add(dense, bias)
mod = tvm.IRModule.from_expr(out)
target = "llvm -mcpu=cascadelake"
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target)
asm = lib.lib.get_source("asm")
assert "vpdpbusd" in asm
dev = tvm.device(target, 0)
runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
a = np.random.uniform(1, 10, size=data_shape).astype(data_dtype)
b = np.random.uniform(1, 10, size=weight_shape).astype("int8")
c = np.random.uniform(1, 10, size=(weight_shape[0],)).astype("int32")
runtime.set_input("data", a)
runtime.set_input("weight", b)
runtime.set_input("bias", c)
runtime.run()
out = runtime.get_output(0).numpy()
ref = np.dot(a.astype("int32"), b.transpose().astype("int32")) + c
np.testing.assert_equal(out, ref)
官网上有很好的例子:
https://tvm.apache.org/docs/topic/vta/tutorials/optimize/convolution_opt.html?highlight=schedule
这里我们着重来看看TVM中如何实现这些调度。
以softmax为例:
softmax用numpy实现:
def softmax_python(a_np, axis=1):
"""Softmax operator.
Parameters
----------
a_np : numpy.ndarray
N-D input data
Returns
-------
output_np : numpy.ndarray
N-D output with same shape
"""
max_elem = np.amax(a_np, axis=axis, keepdims=True)
e = np.exp(a_np - max_elem)
expsum = np.sum(e, axis=axis, keepdims=True)
out_np = e / expsum
return out_np
不使用调度的TVM算子:
@main = primfn(A_1: handle) -> ()
attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
buffers = {A: Buffer(A_2: Pointer(float32), float32, [12], [])}
buffer_map = {A_1: A}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [3, 4], [])} {
allocate(T_softmax_maxelem: Pointer(global float32), float32, [3]), storage_scope = global;
allocate(T_softmax_exp: Pointer(global float32), float32, [12]), storage_scope = global {
for (i0: int32, 0, 3) {
T_softmax_maxelem_1: Buffer(T_softmax_maxelem, float32, [3], [], align=8)[i0] = -3.40282e+38f32
for (k: int32, 0, 4) {
T_softmax_maxelem_1[i0] = max(T_softmax_maxelem_1[i0], A[((i0*4) + k)])
}
}
for (i0_1: int32, 0, 3) {
for (i1: int32, 0, 4) {
let cse_var_1: int32 = ((i0_1*4) + i1)
T_softmax_exp_1: Buffer(T_softmax_exp, float32, [12], [], align=32)[cse_var_1] = @tir.exp((A[cse_var_1] - T_softmax_maxelem_1[i0_1]), dtype=float32)
}
}
for (i0_2: int32, 0, 3) {
T_softmax_maxelem_2: Buffer(T_softmax_maxelem, float32, [3], [], align=8)[i0_2] = 0f32
for (k_1: int32, 0, 4) {
T_softmax_maxelem_2[i0_2] = (T_softmax_maxelem_2[i0_2] + T_softmax_exp_1[((i0_2*4) + k_1)])
}
}
for (i0_3: int32, 0, 3) {
for (i1_1: int32, 0, 4) {
let cse_var_2: int32 = ((i0_3*4) + i1_1)
T_softmax_exp_2: Buffer(T_softmax_exp, float32, [12], [], align=32)[cse_var_2] = (T_softmax_exp_1[cse_var_2] / T_softmax_maxelem_2[i0_3])
}
}
}
}
可以观察到:原函数实现生成了整整4个for语句。增加了程序的复杂度。解决方案是,合并4个for语句并对outer axis做并行化处理。
python/tvm/topi/x86/nn.py
# only parallelize outer dimensions up to axis
outer_axes = [s[softmax_op].op.axis[i] for i in range(0, axis)]
fused_outer_axes = s[softmax_op].fuse(*outer_axes)
s[softmax_op].parallel(fused_outer_axes)
其中
axis=1
softmax_op= compute(T_softmax_norm, body=[(T_softmax_exp[i0, i1]/T_softmax_expsum[i0])], axis=[iter_var(i0, range(min=0, ext=3)), iter_var(i1, range(min=0, ext=4))], reduce_axis=[], tag=softmax_output, attrs={"axis": 1})
s[softmax_op].op.axis= [iter_var(i0, range(min=0, ext=3)), iter_var(i1, range(min=0, ext=4))]
由softmax_op的body可知,softmax_op对应的是最后一行的计算:
T_softmax_exp_2: Buffer(T_softmax_exp, float32, [12], [], align=32)[cse_var_2] = (T_softmax_exp_1[cse_var_2] / T_softmax_maxelem_2[i0_3])
s[softmax_op]也仅仅作用于最后一个for。
outer_axes是所有迭代变量。i0是最外部的for中的迭代变量,长度为3。i1则是内部for,长度为4。其中,
for i in range(0, axis)
保证了outer_axes只能取到所有外部的轴而只保留了最里层的for。当前的axis=1,则outer_axes储存了唯一一个轴i0。当axis=2或更大时,外部的轴有若干个。不管外部的轴有几个,它们都会通通被fuse到一起然后并行掉。
合并用到的函数是fuse。
https://tvm.apache.org/docs/how_to/work_with_schedules/schedule_primitives.html?highlight=primitives
得到:
@main = primfn(A_1: handle) -> ()
attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
buffers = {A: Buffer(A_2: Pointer(float32), float32, [12], [])}
buffer_map = {A_1: A}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [3, 4], [])} {
allocate(T_softmax_maxelem: Pointer(global float32), float32, [3]), storage_scope = global;
allocate(T_softmax_exp: Pointer(global float32), float32, [12]), storage_scope = global {
for (i0: int32, 0, 3) {
T_softmax_maxelem_1: Buffer(T_softmax_maxelem, float32, [3], [], align=8)[i0] = -3.40282e+38f32
for (k: int32, 0, 4) {
T_softmax_maxelem_1[i0] = max(T_softmax_maxelem_1[i0], A[((i0*4) + k)])
}
}
for (i0_1: int32, 0, 3) {
for (i1: int32, 0, 4) {
let cse_var_1: int32 = ((i0_1*4) + i1)
T_softmax_exp_1: Buffer(T_softmax_exp, float32, [12], [], align=32)[cse_var_1] = @tir.exp((A[cse_var_1] - T_softmax_maxelem_1[i0_1]), dtype=float32)
}
}
for (i0_2: int32, 0, 3) {
T_softmax_maxelem_2: Buffer(T_softmax_maxelem, float32, [3], [], align=8)[i0_2] = 0f32
for (k_1: int32, 0, 4) {
T_softmax_maxelem_2[i0_2] = (T_softmax_maxelem_2[i0_2] + T_softmax_exp_1[((i0_2*4) + k_1)])
}
}
for (i0_3: int32, 0, 3) "parallel" {
for (i1_1: int32, 0, 4) {
let cse_var_2: int32 = ((i0_3*4) + i1_1)
T_softmax_exp_2: Buffer(T_softmax_exp, float32, [12], [], align=32)[cse_var_2] = (T_softmax_exp_1[cse_var_2] / T_softmax_maxelem_2[i0_3])
}
}
}
}
可以观察到由i0表示的外部for被并行掉了,而内部for没有任何变动。
但这显然不够。我们要求的是所有外部for都并行掉,而不是最后一个for。
因此,我们需要找出其他的变量,然后重复上面的代码即可。
取出变量却没有那么容易。变量以input_tensors的形式储存着,通常是嵌套的。
softmax_op的input_tensors有两个:T_softmax_exp和T_softmax_expsum。
softmax_op= compute(T_softmax_norm, body=[(T_softmax_exp[i0, i1]/T_softmax_expsum[i0])], axis=[iter_var(i0, range(min=0, ext=3)), iter_var(i1, range(min=0, ext=4))], reduce_axis=[], tag=softmax_output, attrs={"axis": 1})
softmax_op.input_tensors= [Tensor(shape=[3, 4], op.name=T_softmax_exp), Tensor(shape=[3], op.name=T_softmax_expsum)]
它们都在softmax_op的body中。由于重命名规则,在打印出的调度中T_softmax_exp和T_softmax_expsum被分别重命名为了T_softmax_exp_1和T_softmax_maxelem_2。
T_softmax_exp和T_softmax_expsum被称为softmax_op的两个producer(生产者)。生产者为当前算子提供输入。取出它们的方法就很清晰了:
exp = softmax_op.input_tensors[0]
expsum = softmax_op.input_tensors[1]
现在,我们有了T_softmax_exp_1和T_softmax_maxelem_2,还缺少T_softmax_maxelem_1。同理,T_softmax_maxelem_1是T_softmax_exp_1的第二个生产者(第一个生产者是A)。
max_elem = s[exp].op.input_tensors[1]
现在我们完成了变量的获得,现在需要将其全部合并和改为并行。
# move computations with the same outer dimensions under the same root
s[max_elem].compute_at(s[softmax_op], fused_outer_axes)
s[expsum].compute_at(s[softmax_op], fused_outer_axes)
if exp is not None:
s[exp].compute_at(s[softmax_op], fused_outer_axes)
这里用到了compute_at方法。compute_at将expsum变量的计算移动到softmax_op的fused_outer_axes轴上。
得到:
@main = primfn(A_1: handle) -> ()
attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
buffers = {A: Buffer(A_2: Pointer(float32), float32, [12], [])}
buffer_map = {A_1: A}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [3, 4], [])} {
allocate(T_softmax_norm: Pointer(global float32), float32, [12]), storage_scope = global;
for (i0: int32, 0, 3) "parallel" {
allocate(T_softmax_maxelem: Pointer(global float32), float32, [1]), storage_scope = global;
allocate(T_softmax_exp: Pointer(global float32), float32, [4]), storage_scope = global;
allocate(T_softmax_expsum: Pointer(global float32), float32, [1]), storage_scope = global {
T_softmax_maxelem_1: Buffer(T_softmax_maxelem, float32, [1], [], align=4)[0] = -3.40282e+38f32
for (k: int32, 0, 4) {
T_softmax_maxelem_1[0] = max(T_softmax_maxelem_1[0], A[((i0*4) + k)])
}
for (i1: int32, 0, 4) {
T_softmax_exp_1: Buffer(T_softmax_exp, float32, [4], [], align=16)[i1] = @tir.exp((A[((i0*4) + i1)] - T_softmax_maxelem_1[0]), dtype=float32)
}
T_softmax_expsum_1: Buffer(T_softmax_expsum, float32, [1], [], align=4)[0] = 0f32
for (k_1: int32, 0, 4) {
T_softmax_expsum_1[0] = (T_softmax_expsum_1[0] + T_softmax_exp_1[k_1])
}
for (i1_1: int32, 0, 4) {
T_softmax_norm_1: Buffer(T_softmax_norm, float32, [12], [], align=32)[((i0*4) + i1_1)] = (T_softmax_exp_1[i1_1] / T_softmax_expsum_1[0])
}
}
}
}
可以看到其他3个for也合并到了一起。
截止到本篇文章的TVM版本,可以观察到一个仍需优化的点: