将 TVM 集成到您的项目中

TVM 运行时设计为轻量且可移植的。有几种方法可以将 TVM 集成到您的项目中。

本文介绍了将 TVM 作为 JIT 编译器集成到您的系统中以生成函数的几种可能方法。

DLPack 支持

TVM 生成的函数遵循 PackedFunc 的约定。它是一种可以接受位置参数的函数，这些参数包括诸如浮点数、整数、字符串等标准类型。PackedFunc 采用 DLPack 约定中的 DLTensor 指针。因此，您唯一需要解决的问题是创建相应的 DLTensor 对象。

集成用户定义的 C++ 数组

在 C++ 中唯一需要做的，就是将您的数组转换为 DLTensor，并将其地址作为 DLTensor* 传递给生成的函数。

集成用户定义的 Python 数组

假设您有 Python 对象 MyArray。您需要做以下三件事：

• 在您的数组中添加 _tvm_tcode 字段，该字段返回 tvm.TypeCode.ARRAY_HANDLE。

• 在您的对象中支持 _tvm_handle 属性，该属性返回 DLTensor 的地址，以 Python 整数的形式表示。

• 通过 tvm.register_extension 注册这个类。

# Example code
import tvm

class MyArray(object):
    _tvm_tcode = tvm.TypeCode.ARRAY_HANDLE

    @property
    def _tvm_handle(self):
        dltensor_addr = self.get_dltensor_addr()
        return dltensor_addr

# You can put registration step in a separate file mypkg.tvm.py
# and only optionally import that if you only want optional dependency.
tvm.register_extension(MyArray)