TVM Schedule

http://tvm.apache.org/docs/api/python/te.html

https://github.com/StrongSpoon/tvm.schedule

TE (tensor expression) 位于 Relay IR 和 TIR 之间, TE 包含 compute + schedule, 其中 compute 用 TIR 来描述运算的核心操作，schedule 用来描述如何调度.

Q: 什么是调度 & 为什么需要调度

Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines

Learning to Optimize Tensor Programs

Relay IR 先经过 tvm.relay.transform 进行优化，然后转换为 TE, 然后通过 schedule 产生最佳性能的代码 (例如 vectorize scheduler 会使用 SIMD 指令加速)

import tvm
from tvm import te

M = 1024
N = 1024
A = te.placeholder((M, N), name="A")
B = te.placeholder((M, N), name="B")
C = te.compute((M, N), lambda x, y: A[x, y] + B[x, y], name="C")

s = te.create_schedule(C.op)
print(tvm.lower(s, [A, B, C]))

f = tvm.build(s, [A, B, C], "c")
print(f.get_source())

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (x: int32, 0, 1024) {
    for (y: int32, 0, 1024) {
      C_2[((x*1024) + y)] = ((float32*)A_2[((x*1024) + y)] + (float32*)B_2[((x*1024) + y)])
    }
  }
}


// tvm target: c -keys=cpu -link-params=0
#define TVM_EXPORTS
#include "tvm/runtime/c_runtime_api.h"
#include "tvm/runtime/c_backend_api.h"
#include <math.h>
#ifdef __cplusplus
extern "C"
#endif
TVM_DLL int32_t default_function(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode, void* resource_handle) {
  void* arg0 = (((TVMValue*)args)[0].v_handle);
  int32_t arg0_code = ((int32_t*)arg_type_ids)[(0)];
  void* arg1 = (((TVMValue*)args)[1].v_handle);
  int32_t arg1_code = ((int32_t*)arg_type_ids)[(1)];
  void* arg2 = (((TVMValue*)args)[2].v_handle);
  int32_t arg2_code = ((int32_t*)arg_type_ids)[(2)];
  void* A = (((DLTensor*)arg0)[0].data);
  void* arg0_shape = (((DLTensor*)arg0)[0].shape);
  void* arg0_strides = (((DLTensor*)arg0)[0].strides);
  int32_t dev_id = (((DLTensor*)arg0)[0].device.device_id);
  void* B = (((DLTensor*)arg1)[0].data);
  void* arg1_shape = (((DLTensor*)arg1)[0].shape);
  void* arg1_strides = (((DLTensor*)arg1)[0].strides);
  void* C = (((DLTensor*)arg2)[0].data);
  void* arg2_shape = (((DLTensor*)arg2)[0].shape);
  void* arg2_strides = (((DLTensor*)arg2)[0].strides);
  if (!(arg0_strides == NULL)) {
  }
  if (!(arg1_strides == NULL)) {
  }
  if (!(arg2_strides == NULL)) {
  }
  for (int32_t x = 0; x < 1024; ++x) {
    for (int32_t y = 0; y < 1024; ++y) {
      ((float*)C)[(((x * 1024) + y))] = (((float*)A)[(((x * 1024) + y))] + ((float*)B)[(((x * 1024) + y))]);
    }
  }
  return 0;
}

使用 split schedule 后：

s[C].split(C.op.axis[0], factor=20)
print(tvm.lower(s, [A, B, C]))

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (x.outer: int32, 0, 52) {
    for (x.inner: int32, 0, 20) {
      if (((x.outer*20) + x.inner) < 1024) {
        for (y: int32, 0, 1024) {
          C_2[(((x.outer*20480) + (x.inner*1024)) + y)] = ((float32*)A_2[(((x.outer*20480) + (x.inner*1024)) + y)] + (float32*)B_2[(((x.outer*20480) + (x.inner*1024)) + y)])
        }
      }
    }
  }
}

使用 vectorize schedule 后:

s = te.create_schedule(C.op)
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], 32, 32)

print(tvm.lower(s, [A, B, C], simple_mode=True))
print("---------cutting line---------")

s[C].vectorize(yi)

print(tvm.lower(s, [A, B, C], simple_mode=True))

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (x.outer: int32, 0, 32) {
    for (y.outer: int32, 0, 32) {
      for (x.inner: int32, 0, 32) {
        for (y.inner: int32, 0, 32) {
          C_2[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] = ((float32*)A_2[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] + (float32*)B_2[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)])
        }
      }
    }
  }
}


---------cutting line---------
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (x.outer: int32, 0, 32) {
    for (y.outer: int32, 0, 32) {
      for (x.inner: int32, 0, 32) {
        C_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)] = ((float32x32*)A_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)] + (float32x32*)B_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)])
      }
    }
  }
}

https://tvm.apache.org/docs/tutorials/language/schedule_primitives.html

schedule 主要分为以下几类:

1. memory tiling
2. loop transformation (splitting, reording, unrolling)
3. vectorization/tensorization
4. parallelization

https://en.wikipedia.org/wiki/Injective_function

injective function 即单射函数, 也称为 one-to-one function, 在 TVM 中, 有大量的符合 injective 定义的 op 都使用同一个 schedule: schedule_injective, 因为它们执行时除了具体的运算不同, schedule 的要求实际上是一样的.

TVM 提供了几个函数来指定 schedule_injective 做为 schedule:

1. register_injective_schedule
2. register_broadcast_schedule

例如:

register_broadcast_schedule("log")
register_broadcast_schedule("log2")
register_broadcast_schedule("log10")
register_broadcast_schedule("tan")
register_broadcast_schedule("cos")
register_broadcast_schedule("cosh")
register_broadcast_schedule("sin")
register_broadcast_schedule("sinh")
# ...
register_broadcast_schedule("add")
register_broadcast_schedule("subtract")
register_broadcast_schedule("multiply")
register_broadcast_schedule("divide")
register_broadcast_schedule("floor_divide")
register_broadcast_schedule("power")
register_broadcast_schedule("copy")
register_broadcast_schedule("logical_not")
register_broadcast_schedule("logical_and")
register_broadcast_schedule("logical_or")
register_broadcast_schedule("logical_xor")
register_broadcast_schedule("bitwise_not")
register_broadcast_schedule("bitwise_and")
register_broadcast_schedule("bitwise_or")
register_broadcast_schedule("bitwise_xor")
register_broadcast_schedule("negative")
register_broadcast_schedule("mod")
register_broadcast_schedule("floor_mod")
# ...
register_injective_schedule("one_hot")
# ...

register_injective_schedule 会使用一个名为 schedule_injective 的 generic function, 不同的 target 会通过 schedule_injective.register 注册不同的实现

@schedule_injective.register("cpu")
def schedule_injective_cpu(attrs, outs, target):
    """schedule injective ops for x86"""
    with target:
        return topi.x86.schedule_injective(outs)


@schedule_injective.register(["arm_cpu", "micro_dev"])
def schedule_injective_arm_cpu(_, outs, target):
    """schedule injective ops for arm cpu"""
    with target:
        return topi.arm_cpu.schedule_injective(outs)

injective schedule 主要是利用 split+vectorize, 因为 vectorize (simd) 本来就是用来加速 injective 操作的

Reduce Schedule 与 Injective Schedule 类似, 但它针对的是 sum, max 等进行 reduction 的操作, reduce schedule 无法使用 vectorize, tvm 的实现主要是使用 parallel

register_reduce_schedule("argmax")
register_reduce_schedule("argmin")
register_reduce_schedule("sum")
register_reduce_schedule("all")
register_reduce_schedule("any")
register_reduce_schedule("max")
register_reduce_schedule("min")
register_reduce_schedule("prod")
register_reduce_schedule("mean")
# ...


@schedule_reduce.register("arm_cpu")
def schedule_reduce_cpu(attrs, outs, target):
    """schedule reduction ops for arm_cpu"""
    with target:
        return topi.x86.schedule_reduce(outs)


@schedule_reduce.register("cpu")
def schedule_reduce_cpu(attrs, outs, target):
    """schedule reduction ops for x86"""
    with target:
        return topi.x86.schedule_reduce(outs)

task = tvm.auto_scheduler.SearchTask(
    func=matmul_add, args=(N, L, M, "float32"), target=target
)
log_file = "/tmp/matmul.json"

print("-------------------- tunning --------------------")
tune_option = auto_scheduler.TuningOptions(
    num_measure_trials=10,
    measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
    verbose=1,
)
task.tune(tune_option)

-------------------- tunning --------------------
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Generate Sketches		#s: 3
Sample Initial Population	#s: 2015	fail_ct: 1	Time elapsed: 0.67
GA Iter: 0	Max score: 0.9998	Min score: 0.9348	#Pop: 128	#M+: 0	#M-: 0
GA Iter: 4	Max score: 1.0000	Min score: 0.9887	#Pop: 128	#M+: 1381	#M-: 72
EvolutionarySearch		#s: 128	Time elapsed: 2.90
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 10 programs to measure:
..........*T*T*T*T*T*T*T*T*T*T
Time elapsed for measurement: 101.94 s
----------------------------------------------------------------------
------------------------------  [ Done ]
----------------------------------------------------------------------
No valid state found in this search round. Check if it has traversed all of the search space.

print("-------------------- orig_scheuler--------------------")
args = matmul_add(N, L, M, "float32")
sch = te.create_schedule(args[-1].op)
eval(sch, args)

print("-------------------- auto_scheuler--------------------")
sch, args = task.apply_best(log_file)
eval(sch, args)

-------------------- orig_scheuler--------------------
Execution time of this operator: 287.769 ms
-------------------- auto_scheuler--------------------
Execution time of this operator: 6.305 ms

auto_scheduler 要解决的问题是:

定义一个 schedule 参数的搜索空间, 通过一个 measure 函数在这个空间中 sample, 然后基于这些 sample 用 cost_model 预测最佳的 schedule 参数

auto_scheduler 自带实现的:

1. builder

   • LocalBuilder

     多进程, 最终调用 build_module.build 及 export_library

2. runner

   • LocalRunner

     单进程, 在本地 CPU 上执行, 所以非本机 target 编译的 module 无法用 LocalRunner

   • RPCRunner

     多进程, 使用 RPC 在多个远端的 target 设备上执行, 所以这种方式可以并行执行, 而且可以执行为其它 target 编译的 module, 但现在的实现是依赖于一个基于 TCP 的中心 tracker 来连接 PC 和多个 target 设备, 和 MicroTVM 自己那一套基于串口 RPC 机制没有任何关系

虽然 MicroTVM 有一套自己的 Flasher, Transport 以及 RPC 机制, 现在看和 auto_scheduler 并不能一起工作, 为了让 auto_scheduler 支持 MicroTVM, 需要自己实现 builder 和 runner,