Apache TVM 是一個端到端的深度學習編譯框架，適用於 CPU、GPU 和各種機器學習加速芯片。更多 TVM 中文文檔可訪問 → https://tvm.hyper.ai/

作者：Yongfeng Gu

本文介紹使用 TEDD（Tensor Expression Debug Display）對張量表達式進行可視化。

張量表達式使用原語進行調度，單個原語容易理解，但組合在一起時，就會變得複雜。在張量表達式中引入了調度原語的操作模型。

• 不同調度原語之間的交互，
• 調度原語對最終代碼生成的影響。

操作模型基於數據流圖、調度樹和 IterVar 關係圖。調度原語在這些計算圖上進行操作。

TEDD 從給定的 schedule 中呈現這三個計算圖，本教程演示瞭如何用 TEDD，以及如何解釋渲染的計算圖。

import tvm
from tvm import te
from tvm import topi
from tvm.contrib import tedd

使用 Bias 和 ReLU 定義和調度卷積

用 Bias 和 ReLU 為卷積構建一個張量表達式示例，首先連接 conv2d、add 和 relu TOPIs，然後創建一個 TOPI 通用 schedule。

batch = 1
in_channel = 256
in_size = 32
num_filter = 256
kernel = 3
stride = 1
padding = "SAME"
dilation = 1

A = te.placeholder((in_size, in_size, in_channel, batch), name="A")
W = te.placeholder((kernel, kernel, in_channel, num_filter), name="W")
B = te.placeholder((1, num_filter, 1), name="bias")

with tvm.target.Target("llvm"):
    t_conv = topi.nn.conv2d_hwcn(A, W, stride, padding, dilation)
    t_bias = topi.add(t_conv, B)
    t_relu = topi.nn.relu(t_bias)
    s = topi.generic.schedule_conv2d_hwcn([t_relu])

使用 TEDD 渲染計算圖

通過渲染計算圖來查看計算及其調度方式。若在 Jupyter Notebook 中運行本教程，則可以用以下注釋行來渲染 SVG 圖形，讓它直接在 Notebook 中顯示。

tedd.viz_dataflow_graph(s, dot_file_path="/tmp/dfg.dot")
# tedd.viz_dataflow_graph(s, show_svg = True)

第一個是數據流圖。每個節點代表一個階段，中間是名稱和內存範圍，兩邊是輸入/輸出信息。圖中的邊顯示節點的依賴關係。

tedd.viz_schedule_tree(s, dot_file_path="/tmp/scheduletree.dot")
# tedd.viz_schedule_tree(s, show_svg = True)

上面渲染了調度樹圖。注意範圍不可用的警告，它表明要調用 normalize() 來推斷範圍信息。跳過檢查第一個調度樹，推薦通過比較 normalize() 之前和之後的計算圖來了解其影響。

s = s.normalize()
tedd.viz_schedule_tree(s, dot_file_path="/tmp/scheduletree2.dot")
# tedd.viz_schedule_tree(s, show_svg = True)

仔細看第二個調度樹，ROOT 下的每一個 block 代表一個階段。階段名稱顯示在頂行，計算顯示在底行。中間行是 IterVars，外部越高，內部越低。

IterVar 行包含其索引、名稱、類型和其他可選信息。以 W.shared 階段為例，第一行是名稱「W.shared」和內存範圍「Shared」。它的計算是 W(ax0, ax1, ax2, ax3)。最外層循環 IterVar 是 ax0.ax1.fused.ax2.fused.ax3.fused.outer，以 kDataPar 的 0 為索引，綁定到 threadIdx.y，範圍（min=0，ext=8）。

還可以用索引框的顏色來判斷 IterVar 類型，如圖所示。

如果一個階段在任何其他階段都沒有計算，則它有直接到根節點的邊；否則，它有一條邊指向它所附加的 IterVar，例如 W.shared 在中間計算階段附加到 rx.outer。

tedd.viz_itervar_relationship_graph(s, dot_file_path="/tmp/itervar.dot")
# tedd.viz_itervar_relationship_graph(s, show_svg = True)

最後一個是 IterVar 關係圖。每個子圖代表一個階段，包含 IterVar 節點和轉換節點。

例如，W.shared 有三個拆分節點和三個融合節點。其餘的是與調度樹中的 IterVar 行格式相同的 IterVar 節點。 Root IterVars 是那些不受任何變換節點驅動的，例如 ax0；葉節點 IterVars 不驅動任何轉換節點，並且具有非負索引，例如索引為 0 的 ax0.ax1.fused.ax2.fused.ax3.fused.outer。

總結

本教程演示 TEDD 的用法。用一個 TOPI 構建的示例來顯示底層的 schedule，可在任何調度原語之前和之後用它來檢查其效果。

下載 Python 源代碼：tedd.py

下載 Jupyter Notebook：tedd.ipynb