TVM Relay softmax算子调度分析

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计算和调度分离的核心思想。

Fuse Ops 算子融合

调度不仅仅需要考虑当前算子本身如何高效执行，还需要考虑算子融合，即当前算子与其他算子融合之后的执行是否高效。这让调度的开发变得更加复杂。

例如：

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版本，可以观察到一个仍需优化的点：

• T_softmax_exp_1计算完成后可以立即加到T_softmax_expsum_1上，这样做可以省去一个for。
• 内层for没有split，当第二个轴很长时，cache miss会增加。（比如第二个轴不是4而是1024或2048这种值）
• 常量化。T_softmax_maxelem_1可以常量化，而不是一个长度为1的数组。