通过 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)

通过 DLPack 构建跨框架深度学习编译器相关推荐

DLPack构建跨框架深度学习编译器-实现张量之间的相互转换【pytorch】
文章目录 1 概述 2 pytorch使用torch.utils.dlpack实现DLPack与张量之间的相互转换 2.1 将DLPack解码为张量. 2.2 将张量转换为DLPack 3 在 PyT ...
DLPack构建跨框架的深度学习编译器
DLPack构建跨框架的深度学习编译器 Tensorflow,PyTorch和ApacheMxNet等深度学习框架提供了一个功能强大的工具包,可用于快速进行原型设计和部署深度学习模型.易用性通常是以碎 ...
AI System 人工智能系统 TVM深度学习编译器 DSL IR优化计算图编译优化内存内核调度优化 DAG 图优化 DFS TaiChi 函数注册机 Registry
DSL 领域专用语言 TVM深度学习编译器 AI System 人工智能系统参考项目 TaiChi 三维动画渲染物理仿真引擎DSL TVM 深度学习DSL 密集计算DSL LLVM 模块化编译器编 ...
(三)大话深度学习编译器中的自动调优·Empirical Search
前面的第一篇"(一)大话深度学习编译器中的自动调优·前言"与第二篇"(二)大话深度学习编译器中的自动调优·DSL与IR"分别介绍了背景与一些相关概念,这第三篇我 ...
自动生成硬件优化内核：陈天奇等人发布深度学习编译器TVM
有了 TVM,业界与学界开发者们可以快速.轻松地在各个系统(包括手机.嵌入式设备与低功耗芯片)上部署深度学习应用程序,同时无须担心资源与速度的限制.「TVM 作为神经网络和硬件后端之间的共同层,消除了 ...
【从零开始学深度学习编译器】十三，如何在MLIR里面写Pass？
[GiantPandaCV导语]这篇文章是学习了比较久然后按照自己的理解步骤重新总结了下来,主要是MLIR Toy Tutorials第3,4篇文章的内容.这里主要讲解了如何在MLIR中自定义Pass ...
【从零开始学深度学习编译器】十二，MLIR Toy Tutorials学习笔记一
本笔记由学习MLIR Tutorials总结而成,欢迎批评指正. Chapter1: Toy语言和AST MLIR提供了一种Toy语言来说明MLIR的定义和执行的流程.Toy语言是一种基于张量的语言, ...
深度学习编译器Data Flow和Control Flow
深度学习编译器Data Flow和Control Flow 本文介绍了一下深度学习框架的Data Flow和Control Flow,基于TensorFlow解释了TensorFlow是如何在静态图中 ...
深度学习编译器综述The Deep Learning Compiler
深度学习编译器综述The Deep Learning Compiler The Deep Learning Compiler: A Comprehensive Survey 参考文献: https:/ ...

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

通过 DLPack 构建跨框架深度学习编译器相关推荐

最新文章

热门文章