【TVM 教程】外部張量函式

超神经HyperAI發表於2024-11-05
函式

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

作者：Tianqi Chen

雖然 TVM 支援透明程式碼生成，但有時也需將手寫的程式碼合併到流水線，例如對一些卷積核使用 cuDNN，並定義其餘階段。

原生 TVM 就支援黑盒函式呼叫。具體來說，TVM 支援所有與 DLPack 相容的張量函式。這意味著可以使用 POD 型別（指標、整數、浮點數），或者將指向 DLTensor 的指標作為引數，呼叫任何函式。

from __future__ import absolute_import, print_function

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

if not tvm.get_global_func("tvm.contrib.cblas.matmul", allow_missing=True):
    raise Exception("Not compiled with cblas support; can't build this tutorial")

使用外部張量函式

以下示例用 te.extern 來新增一個外部陣列函式呼叫。外部呼叫宣告瞭輸出張量的 shape，第二個引數給出了輸入列表。

使用者需要提供一個描述如何對結果進行計算的函式。計算函式獲取輸入和輸出的符號佔位符列表，並返回執行語句。

這種情況只需呼叫一個註冊的 TVM 函式，它會呼叫 CBLAS。TVM 不控制外部陣列函式的內部，將其視為黑盒。可以進一步混合可排程的 TVM 函式，為結果新增偏差項。

n = 1024
l = 128
m = 235
bias = te.var("bias", dtype="float32")
A = te.placeholder((n, l), name="A")
B = te.placeholder((l, m), name="B")
C = te.extern(
    (n, m),
    [A, B],
    lambda ins, outs: tvm.tir.call_packed(
        "tvm.contrib.cblas.matmul", ins[0], ins[1], outs[0], False, False
    ),
    name="C",
)
D = te.compute(C.shape, lambda i, j: C[i, j] + bias, name="D")
s = te.create_schedule(D.op)

驗證結果

驗證結果是否符合預期。

dev = tvm.cpu(0)
f = tvm.build(s, [A, B, D, bias], "llvm")
a = tvm.nd.array(np.random.uniform(size=(n, l)).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=(l, m)).astype(B.dtype), dev)
d = tvm.nd.array(np.zeros((n, m), dtype=D.dtype), dev)
bb = 10.0
f(a, b, d, bb)
tvm.testing.assert_allclose(d.numpy(), np.dot(a.numpy(), b.numpy()) + 10, rtol=1e-5)

外部 Contrib Wrappers

TVM 為外部呼叫提供了外部contrib Wrappers，以下程式碼與前面的示例等效。

from tvm.contrib import cblas

C = cblas.matmul(A, B)
D = te.compute(C.shape, lambda i, j: C[i, j] + bias, name="D")
s = te.create_schedule(D.op)

將 Python 函式 Hook 為 Extern

由於可以呼叫 TVM 中的任何 PackedFunc，所以可以用外部函式回撥到 Python 中。

以下示例將一個 Python 函式註冊到 TVM runtime 系統，並用它來完成一個階段的計算，這使得 TVM 更加靈活。例如，可透過插入前端回撥來檢查中間結果，或將自定義程式碼與 TVM 混合。

@tvm.register_func("tvm.contrib.my_tvm_addone")
def my_tvm_addone(x, y):
    print("my_tvm_addone signatures: %s, %s" % (type(x), type(y)))
    tvm.nd.array(x.numpy() + 1).copyto(y)

A = te.placeholder((n,), name="A")
B = te.extern(
    A.shape,
    [A],
    lambda ins, outs: tvm.tir.call_packed("tvm.contrib.my_tvm_addone", ins[0], outs[0]),
    name="C",
)
s = te.create_schedule(B.op)
f = tvm.build(s, [A, B], "llvm")
a = tvm.nd.array(np.random.uniform(size=(n,)).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=(n,)).astype(B.dtype), dev)
f(a, b)
tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1, rtol=1e-5)

輸出結果：

my_tvm_addone signatures: <class 'tvm.runtime.ndarray.NDArray'>, <class 'tvm.runtime.ndarray.NDArray'>

總結

  • TVM 透過 te.extern 呼叫外部張量函式。
  • 對外部張量呼叫使用 contrib wrappers。
  • 將前端函式 hook 為外部張量的回撥。

下載 Python 原始碼：extern_op.py

下載 Jupyter Notebook：extern_op.ipynb

相關文章