0x0. 前言

在【從零開始學(xué)深度學(xué)習(xí)編譯器】一，深度學(xué)習(xí)編譯器及TVM 介紹我們已經(jīng)知道TVM可以將各種深度學(xué)習(xí)訓(xùn)練框架的模型（計算圖）轉(zhuǎn)化為內(nèi)部的Graph IR（Relay），然后通過TVM提供的指令生成模塊將Graph IR翻譯成特定硬件可執(zhí)行的指令或者代碼?？偟膩碚f的TVM的思想可以總結(jié)為表示和調(diào)度分離，所謂表示就是IR，調(diào)度就是scheduler。同時，在高性能計算方面TVM提供了多種調(diào)度源語（scheduler），包含了大多數(shù)常見的優(yōu)化手段如算子融合，讀寫緩存，分塊計算，并行計算等等，這些計算方法都可以通過scheduler進行實現(xiàn)。所以這一節(jié)，我們就一起來探索一下TVM中的scheduler。

0x01. 介紹

我們知道TVM的核心就是自動代碼生成技術(shù)，而scheduler則是自動代碼生成技術(shù)的核心概念。scheduler我們可以簡單理解為是一系列優(yōu)化選擇的集合，這些選擇不會影響整個計算的結(jié)果，但對計算的性能卻至關(guān)重要。一個常見的例子是矩陣乘法，給定輸入矩陣A和B，維度分別為和，然后獲得結(jié)果矩陣C，維度為，我在之前的道阻且長_再探矩陣乘法優(yōu)化 詳細列出了為了加速這個計算所采用的一系列優(yōu)化方法，注意這里是以Arm端為例。具體如下：

可以看到雖然這些實現(xiàn)都完成了矩陣乘法這個計算任務(wù)，也就是說輸入輸出都是完全相同的，但在計算過程中卻使用了一系列不同的優(yōu)化手段，這些優(yōu)化算法的集合就可以統(tǒng)稱為scheduler。

接下來我們明確一下scheduler在整個TVM軟件棧中的位置，最近一直在找這樣一張圖，然后OpenMMLab最新放出的介紹Ansor文章里的一張圖剛好能完美解釋這個問題，這里我就抄過來了。以深度學(xué)習(xí)中一個常見的MatMul+Add+Relu計算圖為例，看一下TVM做代碼生成的一個過程。首先TVM將接受的計算圖轉(zhuǎn)換為TVM中的領(lǐng)域特定語言Tensor Expression，即圖中的黃色部分。接下來用戶可以手動指定計算策略即scheduler，然后TVM會自動生成特定后端的代碼，注意圖中的tiling和binding分別代表拆分和綁定的意思，也是scheduler。我們現(xiàn)在明確了scheduler在TVM軟件棧中的位置，也應(yīng)該清楚TVM能否產(chǎn)生高性能的代碼關(guān)鍵就在于scheduler是否指定合理，即優(yōu)化算法在指定后端是否work and efiicient。

0x02. 從Tensor Expression開始看TVM是如何生成CUDA代碼的

我們以chentianqi大佬在TVM文檔中的介紹Tensor Expression例子初步感受一下上面那張圖中描述的TVM代碼生成過程，這里面也包含了scheduler。這一節(jié)之后我們再列舉一些其它的例子來感受scheduler的更多用法?，F(xiàn)在我們從Tensor Expression開始，看看TVM是如何生成代碼的，以及我們具體是如何指定scheduler的。首先導(dǎo)入一堆要用到的包。

import tvm
import tvm.testing
from tvm import te
import numpy as np

# 全局環(huán)境定義

tgt_host = "llvm"
# 如果啟用了GPU，則將其更改為相應(yīng)的GPU，例如：cuda、opencl、rocm
tgt = "cuda"

然后使用向量加法來演示TVM的工作流程。作為第一步，我們需要描述我們的計算。TVM采用Tensor Expression，每個中間結(jié)果表示為一個多維數(shù)組。用戶需要描述生成張量的計算規(guī)則。我們首先定義一個符號變量n來表示形狀。然后我們定義兩個占位符張量，A和B，具有給定的形狀。然后我們用一個計算函數(shù)來描述結(jié)果張量C。計算函數(shù)采用張量的形式，以及描述張量每個位置的計算規(guī)則的lambda函數(shù)。在這個階段沒有計算發(fā)生，因為我們只是聲明應(yīng)該如何進行計算。代碼如下：

n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
print(type(C))

打印出的信息為：<class 'tvm.te.tensor.Tensor'>

接著，雖然上面的幾行描述了計算規(guī)則，但是我們可以用很多方法來計算C，因為C可以在軸上用數(shù)據(jù)并行的方式來計算。TVM要求用戶提供一個稱為schedule的計算描述，即等效于下面的代碼：

for (int i = 0; i < n; ++i) {
  C[i] = A[i] + B[i];
}

我們調(diào)用te.create_schedule來創(chuàng)建scheduler，然后使用split構(gòu)造來拆分C的軸，這將把原來的一個迭代軸拆分成兩個迭代軸的乘積，代碼如下：

s = te.create_schedule(C.op)
bx, tx = s[C].split(C.op.axis[0], factor=64)

這等效于下面的代碼：

for (int bx = 0; bx < ceil(n / 64); ++bx) {
  for (int tx = 0; tx < 64; ++tx) {
    int i = bx * 64 + tx;
    if (i < n) {
      C[i] = A[i] + B[i];
    }
  }
}

最后，我們將迭代軸bx和tx綁定到GPU計算grid中的線程。這些是特定于GPU的構(gòu)造，允許我們生成在GPU上運行的代碼。

if tgt == "cuda" or tgt == "rocm" or tgt.startswith("opencl"):
    s[C].bind(bx, te.thread_axis("blockIdx.x"))
    s[C].bind(tx, te.thread_axis("threadIdx.x"))

上面我們已經(jīng)完成了指定scheduler，接下來我們就可以將上面的所有代碼編譯成一個TVM的函數(shù)了。默認情況下，TVM會將其編譯成一個類型擦除函數(shù)，可以直接從Python端調(diào)用。下面我們使用tvm,build來創(chuàng)建一個編譯函數(shù)，編譯函數(shù)接收scheduler，函數(shù)簽名（包含輸入輸出）以及我們需要編譯到的目標語言。編譯fadd的結(jié)果是一個GPU設(shè)備函數(shù)（如果涉及GPU）以及一個調(diào)用GPU函數(shù)的host端包裝器。fadd是生成的host端包裝函數(shù)，它在內(nèi)部包含對生成的設(shè)備函數(shù)的引用。代碼如下：

fadd = tvm.build(s, [A, B, C], tgt, target_host=tgt_host, name="myadd")

編譯后的TVM函數(shù)生成了一個簡潔的C API，可以被任何語言調(diào)用。TVM在python中提供了一個最小的array API來幫助快速測試和原型開發(fā)。array API基于DLPack（https://github.com/dmlc/dlpack）標準。要運行這個函數(shù)，首先需要創(chuàng)建一個GPU context，然后使用tvm.nd.array將數(shù)據(jù)拷貝到GPU，再使用我們編譯好的函數(shù)fadd來執(zhí)行計算，最后再用asnumpy()將GPU端的array拷貝回CPU使用numpy進行計算，最后比較兩者計算結(jié)果的差距。這部分的代碼如下：

ctx = tvm.context(tgt, 0)

n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
fadd(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())

到了這里整個計算過程就已經(jīng)完成了，但是我們相信大家一定對TVM生成的代碼長什么樣子非常感興趣，TVM也提供了對應(yīng)的接口來讓用戶查看生成的代碼。tvm.build的結(jié)果是一個TVM Module。fadd是包含host包裝器的模塊，同時它也包含了用于CUDA（GPU）設(shè)備的功能模塊。我們將使用下面的代碼打印生成的代碼：

if tgt == "cuda" or tgt == "rocm" or tgt.startswith("opencl"):
    dev_module = fadd.imported_modules[0]
    print("-----GPU code-----")
    print(dev_module.get_source())
else:
    print(fadd.get_source())

輸出為：

-----GPU code-----

#ifdef _WIN32
  using uint = unsigned int;
  using uchar = unsigned char;
  using ushort = unsigned short;
  using int64_t = long long;
  using uint64_t = unsigned long long;
#else
  #define uint unsigned int
  #define uchar unsigned char
  #define ushort unsigned short
  #define int64_t long
  #define uint64_t ulong
#endif
extern "C" __global__ void myadd_kernel0(float* __restrict__ C, float* __restrict__ A, float* __restrict__ B, int n, int stride, int stride1, int stride2) {
  if (((int)blockIdx.x) < (n >> 6)) {
    C[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride2))] = (A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))] + B[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1))]);
  } else {
    if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
      C[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride2))] = (A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))] + B[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1))]);
    }
  }
}

好了，講到這里，我們就知道如何在TVM中定義scheduler并自動生成計算代碼了。

0x03. scheduler更詳細的例子

split

關(guān)于scheduler更詳細的例子可以看大神的這篇文章：https://zhuanlan.zhihu.com/p/94846767。我們這里簡單列舉幾個來理解一下，例如在循環(huán)優(yōu)化中，我們以split為例，代碼如下：

import tvm
from tvm import te

n = 1024
A = te.placeholder((n,), name='A')
k = te.reduce_axis((0, n), name='k')

B = te.compute((1,), lambda i: te.sum(A[k], axis=k), name='B')

s = te.create_schedule(B.op)

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

ko, ki = s[B].split(B.op.reduce_axis[0], factor=32)

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

生成的函數(shù)為：

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


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

可以看到split把iter以factor為間隔分成outer與inner兩層迭代，增加循環(huán)層數(shù)，用于將循環(huán)操作分割為更小的子任務(wù)。從Cuda的文檔中我們可以知道，gridDim和blockDim都可以最多是三維，因此可以通過split可以產(chǎn)生新的維度用于綁定到grid和block上。這個操作在生成CUDA代碼中是很常用的。

實驗代碼可以在https://github.com/BBuf/tvm_learn/blob/main/scheduler 這里找到，我使用的tvm版本為0.8.0-dev。

reorder

第二個想講一下的scheduler是reorder，我們貼出實驗代碼和經(jīng)TVM生成的代碼：

import tvm
from tvm import te

n = 1024
A = te.placeholder((n, n), name='A')
B = te.placeholder((n,n), name='B')
C = te.compute((n, n), lambda i, j: A[i, j] + B[i, j], name='C')

s = te.create_schedule(C.op)

xo, xi = s[C].split(s[C].op.axis[0], factor=32)
yo, yi = s[C].split(s[C].op.axis[1], factor=32)

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

s[C].reorder(xo, yo, yi, xi)

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

生成的函數(shù)為：

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], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i.outer: int32, 0, 32) {
    for (i.inner: int32, 0, 32) {
      for (j.outer: int32, 0, 32) {
        for (j.inner: int32, 0, 32) {
          C_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.inner)] = ((float32*)A_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.inner)] + (float32*)B_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.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], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i.outer: int32, 0, 32) {
    for (j.outer: int32, 0, 32) {
      for (j.inner: int32, 0, 32) {
        for (i.inner: int32, 0, 32) {
          C_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.inner)] = ((float32*)A_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.inner)] + (float32*)B_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.inner)])
        }
      }
    }
  }
}

可以看到reorder 方法重置了循環(huán)iter的內(nèi)外順序，根據(jù)局部性原理，這樣可以最大化利用cache中的現(xiàn)有數(shù)據(jù)，減少數(shù)據(jù)頻繁載入載出的情況，進而提高程序的性能。這也是我們之前探索矩陣乘法時，為什么要將K維放在最外層，而不是將M放在最外層的原因。

tile

接下來我們再看一下tile這種scheduler，tile可以將stage（理解為一個OP，一個OP對應(yīng)了一個stage）的兩個維度按照各自的factor進行拆分，并以固定順序返回兩個outer和兩個inner的iter，從而增加循環(huán)層數(shù)，形成更小的計算任務(wù)。事實上，tile是可以由split和reorder來實現(xiàn)的，tile是矩陣乘法和卷積計算的重要schedule。在這篇文章的第二節(jié)介紹部分，我們貼出了在Arm端手寫各種優(yōu)化算法去優(yōu)化矩陣乘法，里面就多次用到了分塊的計算策略，也就是這里的tile scheduler，可以更好的利用緩存和寄存器，獲得更高的性能。

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], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, 1024) {
    for (j: int32, 0, 1024) {
      C_2[((i*1024) + j)] = 0f32
      for (K: int32, 0, 1024) {
        C_2[((i*1024) + j)] = ((float32*)C_2[((i*1024) + j)] + ((float32*)A_2[((i*1024) + K)]*(float32*)B_2[((K*1024) + j)]))
      }
    }
  }
}


---------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], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i.outer: int32, 0, 32) {
    for (j.outer: int32, 0, 32) {
      for (i.inner: int32, 0, 32) {
        for (j.inner: int32, 0, 32) {
          C_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.inner)] = 0f32
          for (K: int32, 0, 1024) {
            C_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.inner)] = ((float32*)C_2[((((i.outer*32768) + (i.inner*1024)) + (j.outer*32)) + j.inner)] + ((float32*)A_2[(((i.outer*32768) + (i.inner*1024)) + K)]*(float32*)B_2[(((K*1024) + (j.outer*32)) + j.inner)]))
          }
        }
      }
    }
  }
}

vectorize

我們最后再介紹一種scheduler，即向量化。這個也就是公眾號的【AI PC端算法優(yōu)化】介紹的一系列優(yōu)化方法，例如在Intel CPU上使用SSE或者AVX等指令集向量化普通的程序獲得更好的性能?，F(xiàn)在，我們看一下TVM里面是如何使用的吧。代碼如下：

import tvm
import numpy
import timeit
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)
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))

生成的函數(shù)為：

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], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_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], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_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)])
      }
    }
  }
}

我們可以看到vectorize將iter方向上的循環(huán)迭代替換成ramp，從而通過SIMD指令實現(xiàn)數(shù)據(jù)的批量計算，并且只有在數(shù)據(jù)size為常數(shù)、且分割的iter為2的冪（即滿足SIMD的計算數(shù)量）時才會發(fā)生替換，否則vectorize沒有效果，這是SIMD計算設(shè)備（如Intel CPU、Arm CPU）的常用schedule。

還有很多重要的scheduler介于篇幅原因就不一一列舉了，大家可以仔細讀這篇文章：https://zhuanlan.zhihu.com/p/94846767。如果要運行最新版本的TVM scheduler實驗，可以在https://github.com/BBuf/tvm_learn 這里找到代碼。

0x04. 小結(jié)

這篇文章主要結(jié)合了TVM中的一些實例來介紹了scheduler，其實寫到這里我們很自然的又會想出一些問題，例如對于一個深度學(xué)習(xí)模型，我們對于整個計算圖要如何應(yīng)用上面介紹的這些scheduler技巧才可以生成高效的特定后端的代碼，這個時候手動指定計算圖的scheduler就不現(xiàn)實了。這就和Auto-TVM和Auto-Scheduler(或者叫Ansor)有關(guān)了，不得不提的是Ansor是發(fā)表在OSDI會議上，目前比Auto-TVM擁有更好的表現(xiàn)，https://zhuanlan.zhihu.com/p/360041136 這篇近期發(fā)表的文章很好的介紹了Ansor的工作機制，推薦讀者閱讀。后面在理清相關(guān)概念之后，也會嘗試從源碼角度走進TVM，希望將前端和調(diào)度的具體過程嘗試理一下。

0x05. 參考資料

• https://tvm.apache.org/docs
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programming-model
• https://zhuanlan.zhihu.com/p/94846767

歡迎關(guān)注GiantPandaCV, 在這里你將看到獨家的深度學(xué)習(xí)分享，堅持原創(chuàng)，每天分享我們學(xué)習(xí)到的新鮮知識。( ? ?ω?? )?

有對文章相關(guān)的問題，或者想要加入交流群，歡迎添加BBuf微信：