通过 DLPack 构建跨框架深度学习编译器

通过 DLPack 构建跨框架深度学习编译器

深度学习框架，如Tensorflow, PyTorch, and ApacheMxNet，快速原型化和部署深度学习模型提供了强大的工具箱。不幸的是，易用性往往以碎片化为代价：孤立地使用每个框架是很容易的。纵向集成使开发简化为常用案例，但冒险走出困境可能比较棘手。

一个支持不力的方案是在内存中将算子从一个框架直接传递到另一个框架，而没有任何数据重复或复制。支持此类使用案例，将使用户能够将管道串联在一起，在一个框架（或更快）中，某些算子比在另一个框架中得到更好的支持。框架之间的共享数据表示也将弥补这一差距，并允许编译器堆栈在为算子生成代码时针对单一格式。

DLPack是拉伸数据结构的中间内存表示标准。以 DLPack 为共同表示形式，可以利用 TVM 编写的脚本，这些框架传统上只能依赖于供应商提供的库。TVM DLPack功能可以在 DLPack 算子上运行，提供DLPack，将 PyTorch 和 MxNet 等框架中的算子数据结构与零数据拷贝进行衔接。

DLPack 提供了一个简单、便携的内存数据结构：

/*!

 * \brief Plain C Tensor object, does not manage memory.

 */

typedef struct {

/*!

   * \brief The opaque data pointer points to the allocated data.

   *  This will be CUDA device pointer or cl_mem handle in OpenCL.

   *  This pointer is always aligns to 256 bytes as in CUDA.

   */

void* data;

/*! \brief The device context of the tensor */

DLContext ctx;

/*! \brief Number of dimensions */

int ndim;

/*! \brief The data type of the pointer*/

DLDataType dtype;

/*! \brief The shape of the tensor */

int64_t* shape;

/*!

   * \brief strides of the tensor,

   *  can be NULL, indicating tensor is compact.

   */

int64_t* strides;

/*! \brief The offset in bytes to the beginning pointer to data */

uint64_t byte_offset;

} DLTensor;

例如，在 TVM 中声明和编译矩阵乘法算子，并构建一个使用 DLPack 表示的wrapper ，使该算子能够支持 PyTorch 算子。还与 MxNet 重复此演示。此扩展允许机器学习开发人员，在不牺牲性能的情况下，将研究代码快速移植到相对不受支持的硬件平台。

DLPack 如何提供框架和 TVM 之间共享的中间wrapper 的插图：

Figure 1

First, we compute a reference output in PyTorch:

import torch

x = torch.rand(56,56)

y = torch.rand(56,56)

z = x.mm(y)

然后，使用默认调度表定义并构建 TVM 矩阵乘法算子：

n = tvm.convert(56)

X = tvm.placeholder((n,n), name='X')

Y = tvm.placeholder((n,n), name='Y')

k = tvm.reduce_axis((0, n), name='k')

Z = tvm.compute((n,n), lambda i,j : tvm.sum(X[i,k]*Y[k,j], axis=k))

s = tvm.create_schedule(Z.op)

fmm = tvm.build(s, [X, Y, Z], target_host='llvm', name='fmm')

简洁性，不涵盖 TVM 的大集合原型调度，可以优化矩阵乘法。

然后，将 TVM 功能转换为支持 PyTorch 算子的功能：

from tvm.contrib.dlpack import to_pytorch_func

# fmm is the previously built TVM function (Python function)

# fmm is the wrapped TVM function (Python function)

fmm_pytorch = to_pytorch_func(fmm)

z2 = torch.empty(56,56)

fmm_pytorch(x, y, z2)

np.testing.assert_allclose(z.numpy(), z2.numpy())

验证结果是否匹配。

可以重复相同的示例，但使用 MxNet 代替：

import mxnet

from tvm.contrib.mxnet import to_mxnet_func

ctx = mxnet.cpu(0)

x = mxnet.nd.uniform(shape=(56,56), ctx=ctx)

y = mxnet.nd.uniform(shape=(56,56), ctx=ctx)

z = mxnet.nd.empty(shape=(56,56), ctx=ctx)

f = tvm.build(s, [X, Y, Z], target_host='llvm', name='f')

f_mxnet = to_mxnet_func(f)

f_mxnet(x, y, z)

np.testing.assert_allclose(z.asnumpy(), x.asnumpy().dot(y.asnumpy()))

在PyTorch Example示例的hood下

由于TVM提供将dlpack张量转换为tvm s的功能，反之亦然，因此所有需要的是通过wrapper 功能来增加一些语法。使用带有dlpack支持的张量框架的通用转换器，可用于实现简易转换器。

NDArrayconvert_functo_pytorch_func

def convert_func(tvm_func, tensor_type, to_dlpack_func):

assert callable(tvm_func)

def _wrapper(*args):

args = tuple(ndarray.from_dlpack(to_dlpack_func(arg))\

if isinstance(arg, tensor_type) else arg for arg in args)

return tvm_func(*args)

return _wrapper

def to_pytorch_func(tvm_func):

import torch

import torch.utils.dlpack

return convert_func(tvm_func, torch.Tensor, torch.utils.dlpack.to_dlpack)