#!/usr/bin/env python
# coding: utf-8

# # 外部张量函数(External Tensor Functions)
# 
# 虽然TVM支持透明的代码生成，但有时将手工编写的代码合并到流水线中也很有帮助。例如，我们可能想用cuDNN来做一些卷积内核，并定义其余的阶段。
# 
# TVM原生支持这些黑盒函数调用。确切地说，TVM支持所有与DLPack兼容的张量函数。这意味着我们可以调用任何具有POD类型(pointer、int、float)或指向DLTensor的指针作为参数的函数。

# In[1]:


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 buid this tutorial")

# # 使用外部张量函数
# 
# 在下面的例子中，我们使用te.extern添加外部数组函数调用。在外部调用中,我们声明了输出张量的形状.在第二个参数中，我们提供了输入列表。
# 
# 用户需要提供一个函数来描述如何计算结果。计算函数接受输入的符号占位符列表，输出的符号占位符列表，并返回正在执行的语句。
# 
# 在本例中，我们简单地调用一个注册的TVM函数，该函数调用CBLAS调用。TVM不控制外部数组函数的内部，并将其视为黑盒。我们可以进一步混合可调度的TVM调用，从而在结果中添加偏置项。

# In[2]:


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)

# # 验证结果
# 
# 我们可以验证结果是否符合我们的预期。

# In[3]:


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封装(Extern Contrib Wrappers)
# 
# TVM还为有用的extern调用提供了extern contrib包装器，下面的行等价于前面的示例。

# In[4]:


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)

# # Hook python函数作为外部调用(Hook Python Function as Extern)
# 
# 因为我们可以调用TVM中的任何封装函数。我们可以使用extern函数回调python。
# 
# 下面的示例将一个python函数注册到TVM runtime系统中，并使用它来完成计算的一个阶段。这使得TVM更加灵活。例如，我们可以插入前端回调来检查中间结果，或者将定制代码与TVM混合使用。

# In[5]:


@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.asnumpy() + 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), 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)

# Out[5]:
my_tvm_addone signatures: <class 'tvm.runtime.ndarray.NDArray'>, <class 'tvm.runtime.ndarray.NDArray'>

# In[ ]: