外部张量函数

虽然TVM支持透明代码生成，但有时将手动编写的代码合并到管道中也很有帮助。例如，我们想去为部分卷积和使用cuDNN和定义其他阶段。

TVM原生支持黑盒函数调用。TVM支持兼容DLPack的所有张量函数。这意味着我们可以使用POD类型（pointer，int，float）或指向DLTensor的指针作为参数调用任何函数。

from __future__ import absolute_import, print_function

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

使用外部张量函数

在下面的示例中，我们使用tvm.extern添加外部数组函数调用。在外部调用中，我们声明输出张量的形状。在第二个参数中，我们提供输入列表。

用户需要提供描述如何计算结果的函数。计算函数获取输入的符号占位符（placeholder）列表，输出的符号占位符列表，并返回执行语句。

在这种情况下，我们只需调用一个注册的TVM函数，该函数调用CBLAS库。 TVM不控制外部数组函数的内部并将其视为黑盒子。我们可以进一步混合可策略的TVM调用，为结果添加bias项。

n = 1024
l = 128
m = 235
bias = tvm.var('bias',dtype=tvm.float32)
#定义输入
A = tvm.placeholder((n,l),name='A')
B = tvm.placeholder((l,m),name='B')
#矩阵乘法（调用外部张量函数）
C = tvm.extern((n,m),[A,B],lambda ins,outs:tvm.call_packed("tvm.contrib.cblas.matmul",ins[0],ins[1],outs[0],False,False),nams="C")
#偏置加
D = tvm.compute(C.shape, lambda i, j: C[i,j]+bias,name="D")
s=tvm.create_schedule(D.op)

验证结果

ctx = 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),ctx)
b = tvm.nd.array(np.random.uniform(size=(l, m)).astype(B.dtype), ctx)
d = tvm.nd.array(np.zeros((n, m), dtype=D.dtype), ctx)
bb = 10.0
f(a, b, d, bb)
tvm.testing.assert_allclose(d.asnumpy(), np.dot(a.asnumpy(), b.asnumpy())+10, rtol=1e-5)

外部contrib封装器

下面例子等同于前面的例子

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

钩子python函数作为外部调用

在TVM能够调用任何包函数（PackedFunc）。我们能用外部函数去调用python。

以下示例，将python函数注册到TVM运行时系统，并使用它来完成计算的一个阶段。这使得TVM更加灵活。例如，我们可以插入前端回调函数来检查中间结果或将自定义代码与TVM混合。

@tvm.register_func("tvm.contrib.my_tvm_addone")
def my_tvm_addon(x,y):
    print("my_tvm_addone signatures: %s, %s" % (type(x), type(y)))
    tvm.nd.array(x.asnumpy() + 1).copyto(y)
    
A = tvm.placeholder((n,), name='A')
#调用上面注册的函数来计算
B = tvm.extern(A.shape, [A], lambda ins, outs: tvm.call_packed(
    "tvm.contrib.my_tvm_addone", ins[0], outs[0]), name="C")
s = tvm.create_schedule(B.op)
f = tvm.build(s, [A, B], "llvm")
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)
f(a, b)
tvm.testing.assert_allclose(b.asnumpy(), a.asnumpy() + 1, rtol=1e-5)

总结

• TVM通过tvm.extern调用外部张量函数
• 为了外部张量函数调用的好处使用contrib封装器
• 我们能使用前端钩子函数作为外部张量函数调用（别人写好的python函数直接注册到TVM用于调用）