Tensor Comprehensions

TC框架也是一个能够降低高性能机器学习代码编写门槛的工具,它能够直接将高级语言生成GPU代码。(填补研究人员于数学运算领域的沟通鸿沟)


适用情况

1、你的 Pytorch 层又大又慢,你打算为此写一段 C++ 或者 CUDA 的优化代码,但是你又不擅长。

2、你写好了 CUDA 代码,但是你还要花费大量时间去优化。你希望可以在最短时间内搞定这些。

3、为了加速,你想在网络中融合多个层,例如 Conv-ReLU-BatchNorm 或者 Linear-ReLU-Linear-ReLU。但是这很难理解。

4、你的研究涉及 CuDNN 和 MKL 未能优化的不寻常的张量。例如,你要使用 13 x 24 的卷积核对 143 x 55 的输入图像进行卷积。你试着用 CUDNN 跑,并且发现它慢的超乎想象。

5、通过调整 Tensors 以适应特定的内存布局,你的代码会变慢。你希望编写在你的输入布局上高效运行的自定义代码很容易。


背景:大量的具有卷积和循环网络的DL模型,诸如TensorFlow、Pytorch这些框架也分别就可用性和表达性之间做了不同权衡。它们在计算运算符的DAG上操作,包装成高性能库,如CUDNN,自动进行内存分配、同步和分发。但如果计算不适合现有的库调用,则需要自定义运算符,通常需要高昂的工程成本。

论文TC:Framework-Agnostic High-Performance Machine Learning Abstractions提出贡献:(TC适合例如GPU这种有分区内存的硬件设备)

1、提出了TC(Tensor Comprehensions):一个用于ML tensor计算的高级语言。(语法是泛化爱因斯坦Einstein表示法)。支持形状和大小推理;

2、端到端编译流程:能够将TC“降低”到高效的GPU代码;

3、Polyhedral编译算法:与通用的并行编译器不同,我们主要通过kernel融合优化以减少启动和同步的开销,还支持多级并行提升到内存层次的更深层次;

4、Autotuning框架:利用JIT编译代码缓存的自动调整框架。它包括非标准大小的specialization,消除了控制和地址生成逻辑,并拥有从ML框架到代码生成器的所有优化旋钮 (knobs)。

5、Polyhedral JIT(实时)编译器:用于将深度学习DAG的数学描述转换为CUDA kernel,并提供了诸如算子融合(operator fusion)和的specialization;

6、集成到两个通用的ML框架PyTorch和Caffe2。并且在原则上,这个系统是足够通用集成到其他ML框架中的。

TC框架自动结合了仿射affine变换,包括层次平铺 (hierarchical tiling)、映射 (mapping)、移动 (shifting)、融合 (fusion)、分布 (distribution)、交换 (interchange),涉及参数化或完全实例化的问题,这些问题对Halide、Latte或XLA都是不可达的。

Tensor Comprehension 将 Halide 编译器作为调用库。FAIR 研究员构建了 Halide 的中间表征(IR)和分析工具,并与多面编译进行技术配对,因此,用户可以在无需理解运行原理的条件下使用相似的高级语法编写层。此外,FAIR 研究员也找到了简化语言的方法,不需要为缩减运算制定循环边界。


发展

Phase1:现有的active library或者BTO (built-to-order)都是孤立地调整库内核,因此会错过上下文依赖,并且创建一个新的库,包含各个独立kernel的所有组合是不可行的。

Phase2:因此,Halide这种domain-specific的语言被生成,在imaging方面取得了成功,因为它能在不混淆底层算法的情况下融合大型的流水线pipelines。然而,在GPU上使用使用Halide时,所有调度转换都必须手动指定,并且通过正确的组合实现高性能超出了大多数用户的能力。(Halide 是一种最近在高性能图像处理领域颇受欢迎的语言,它采用类似的高级函数语法来描述一个图像处理的 pipeline,随后在单独代码块中调度到硬件上,并且详细到如何平铺、矢量化、并行化和融合。对于具有专业知识的人而言,这是一种非常高效的语言;但对于机器学习从业者来说,这一难度并不小。Halide 的自动调度在研究上非常活跃,但对于 GPU 上运行的机器学习代码,目前还没有很好的解决方案。)

Phase3:XLA和Latte:它们结合了计算图和算子,允许跨运算符进行优化以及利用数据大小进行优化。然而,GPU上的性能还不行。这些框架的转换语言无法表示复杂调度 (complex schedule)和映射转换 (mapping transformations),而这些转换对于使用分区内存架构的GPU来说往往至关重要。

综上,计算图引擎的有效编程语言必须同时解决以下两个难题:

1、确保abstraction不仅能提高程序员的工作效率,而且能使编译器及其支持执行环境消除与目标平台无关的顾虑,通过靠近机器的中间表示来改进代码,并自动探索优化空间。换言之,该系统必须能够提供"abstraction without regret",同时传递编译时可用的丰富的语义信息;

2、选择适当的中间表示优化算法,处理深度并行性内存层次,以及硬件功能,如向量指令和专用内存。


TC语法

1、在表达式中使用索引变量隐式定义,从索引内容推断索引变量的范围;

2、表达式右边出现但左边没有的索引假设为reduction维度; 

3、迭代空间中各点的评估顺序不影响输出。

实例

1、矩阵向量积

矩阵向量积

在TC中,定义 +=! 来替代 += ,这样就可以避免初始化操作,即

+=!

函数mv,引入了两个索引变量i和k。变量未在任何地方定义,隐式地成为索引变量。通过索引使用的方式推断它们的范围;并且,因为k只出现在表达式右侧部分,所以存入C的存储将对k进行reduce通过reduction运算符“+”。Reductions可以发生在多个变量之间,但是它们都共享同一类型的关联和交换运算符(例如, +=),以确保计算顺序不影响计算值。

矩阵向量积伪代码

TC支持in-place就地更新,但要求在LHS左手边上分配任何元素之前读取完整的RHS右边表达式。如果LHS张量也在RHS中出现:编译器负责检查元素依赖项的就地更新的因果关系。只允许在点态定义和张量收缩上进行就地更新(即a(i,j) = a(j,i)是不可以的)。

2、SGEMM矩阵乘法、Relu、conv2D、max_pool等。注意where 语句用来确定变量的变化范围,从而程序不会出越界错。(如果没有where语句,那么循环将默认从0开始。)下标表达式可以是迭代器的任何仿射函数,或下标表达式的下标及其组合。后者捕获与数据相关的访问,如收集操作:

gather

TC还可以很简洁的实现 strided convolution(跳格卷积,即步长不为1),sh代表竖直步长,sw代表水平步长,H、W分别代表图像数据(I)的长、宽。

跳格卷积

范围推导

——在表达式中使用索引变量隐式定义,从索引内容推断索引变量的范围。以矩阵向量积为例,output的C的尺寸是通过i,k可取的最大值推理而来的,即M和K,所以C[M*K]。

全局数据布局转换

TC使全局数据布局转换变得非常容易。]ML研究者对于这类原语的主要用法是算法平铺层次分解。前者与数据平铺具有强连接,现在在频域中实现高性能卷积以控制内存占用【83】或适合快速本地内存【2】中无处不在。TC支持广义张量转置(即,应用n-D置换矩阵,其中n>2),并且可以通过简单地重塑张量和调整TC索引表达式来实现数据平铺。范围推断和检查保证了这种重塑将始终在整个TC一致。结构数组到结构数组的转换以及短向量上的类似操作是可用的副产品:它们是维交换和数据平铺的特殊情况。此时,数据布局和TC转换由领域专家在源代码级别进行, TC推理过程保证了表达式具有适当的范围。我们将在将来引入TC中的自动数据布局转换时重新评估这一假设。


高层工作流程说明

首先将TC表达式集成到ML框架中。我们选择了“in process”方式的实现,以简化与基于它们的计算图引擎和ML应用程序的交互,这是全自动调度(fully-automated scheduling)映射流(mappling flow)的独特功能。我们提供了一个轻API,可将特定的张量对象模型转换为我们自己的模型。算子定义被重写以生成TC,而不是框架的后端实现,并为用户提供编写自己的TC的能力。在这种情况下,单个TC可能对应于ML框架中一个运算符DAG(有向无环图?)。目前,此匹配是手动完成的。(以后将会由我们的编译流程自动DAG分区,匹配和重写。)然后,将TC进行JIT编译,如图所示。

JIT编译流程将TC降低到Halide-IR,然后降低到Polyhedral-IR,然后进行优化、代码生成和执行。

从具有特定tensor大小和strides跨度的TC开始,我们将其降低为参数Halide表达式。直接将Halide-IR降低为Poly表示形式。

使用 Halide 和多面编译技术,Tensor Comprehensions 能通过委托内存管理和同步功能以自动合成 CUDA 内核。这一编译能够针对特定尺寸对一般操作符进行融合,对快速本地内存、快速缩减和 Just-in-Time 专业化都能优化。因不尝试涉足内存管理,因此这一流程能够轻松集成到机器学习的任意框架,以及任何允许调用 C++ 函数的语言中。与经典编译器技术和库方法相反的是,多面编译能让张量理解为每个新网络按需调度单个张量元素的计算。在 CUDA 层面,它结合了仿射循环转换、融合/分裂和自动并行处理,同时确保数据在存储器层次结构中正确移动。


Polyhedral JIT编译器

编译器弥合了高层张量操作的逻辑layout(尺寸排序)与Poly代码生成器期望的数据格式之间的不匹配。 “下降”步骤可确保大小和跨度的组合对应于非混叠数组和子数组语法; 它还确保没有边界访问,分析访问关系和推断的张量范围;它发出张量声明和重新排序表达式以匹配目标语言的数据模型,即行主数组。

请注意,张量规范可能具有输入和输出参数别名的特征,以便就地计算和隐式转换高维张量。 我们认为,这样的规范应该导致每个别名场景的多版本化和专业化。 而且,TC基于范围推断所采用的语义不同于XLA,PyTorch和MXNet以某种形式或其他形式采用的NumPy风格的“广播”语义。6TC不需要这种隐式语法糖。 例如,与所谓的外积矩阵乘积[p,q,r] matmul [1,s,r,t]→[p,s,q,t]对应的TC很简单:

如果需要,可以进一步转换布局并导出QPTS版本(以输出尺寸的顺序命名),而不是PSQT。

其他降低步骤包括卷积表达式的正向替换(存储/计算折衷),零填充,镜像和裁剪。

多面体框架提供了基于循环的程序的最先进的内部编译器表示形式,这些循环程序“大多是规则的” [27]。它以最基本的形式描述了由循环和分支包围的算术表达式,这些分支和分支的控制流可以静态确定。因此,据说多面体框架可在程序的静态控制部分(SCoP)上运行。更具体地说,循环边界和分支条件是外部循环迭代器的仿射表达式和象征性对待并称为参数的静态未知常数值。使用数组元素上的算术表达式描述计算,并且对下标的限制与对循环界限的约束相同。存在扩展以通过过度逼近[10]或用户定义的注释[3]处理不规则性。尽管该框架看似简单,但它捕获了大量的计算密集型代码,它在域和数组大小上是参数化的,并且比诸如Halide的特定于域的表示更具表达力。

Poly框架特别适合于深层神经网络,它与具有长依赖性链和不均匀或全部模式的大而深嵌套的循环相关联(由完全连接的层以及张量收缩和换位引起)。这些功能将优化问题推到了与Halide用于图像处理的启发式空间不同的地方,并且比单独的线性代数扩大了空间。

调度树

Affine仿射图可以组成调度树,以促进从高级语言(TC)下游优化器的属性通信。调度树介绍特定的节点类型。一个band节点通过在迭代域上定义的一个或多个分段仿射函数来定义partial执行顺序。即可置换调度band,一维调度功能的元组,可以在保留程序语义的同时自由互换。band中的仿射功能通常称为band member或schedule dimension。过滤器节点(filter nodes)的集合划分了迭代空间,将其子树绑定到迭代域的子集。根据执行顺序是否必须为正确性而受约束,可以将它们安排在set或sequence节点中。上下文context节点提供有关变量和参数的其他信息,例如张量范围或GPU块的大小;它们还可能在子树中引入局部作用域和常量参数,这在将induction变量映射到CUDA中的块和线程标识符时很有用。最后,扩展extension节点引入了不属于原始迭代域的辅助计算,这对于例如引入将数据复制到shared memory。

TC的规范调度树由外部Sequence节点定义,然后是每个TC语句的Filter节点。 在每个filter的分支内,Band节点定义一个身份调度,该身份调度的一维调度功能与语句的循环迭代器一样多。 根据TC语义,隐式循环形成一个可置换band。 图3.a显示了sgemm TC的规范调度树,参数声明(N,M,K)→{...}。

可以识别2-D嵌套初始化语句,然后识别3-D嵌套更新语句。调度可以是输入大小的参数,也可以具有有关张量大小的其他上下文信息。在Band节点未定义内射injective调度的情况下,语句实例将按照其域坐标的字典顺序进行调度。Poly模型中的程序转换涉及定义不同的调度,该调度对应于遍历迭代域的不同(部分或全部)顺序。例如,观察到C张量在两个嵌套之间被重用,可以构造图3.b中的调度以利用访问局部性并提高性能。该树的特征是带有两个循环的i和j循环的外带节点,这对应于循环融合。序列节点确保在启用正确初始化的T的各个实例之前执行S的实例。第二band仅适用于T,并且对应于最里面的(减少)回路k。此外,树还引入了一个Context节点来陈述有关参数值的假设。

越界访问

Poly模型允许对张量访问进行相关编码。 将这些迭代域与表示为集合的迭代域进行组合,可以计算所有访问的张量元素的集合(即语句的足迹),并检查其是否适合(指定的或推断的)张量大小。 特别地,访问关系使得能够检测越界访问。 从张量形状推断出的属于足迹F = D.A而不属于张量元素T集合的元素对应于出界访问。 因此,(D.A)\ T =∅是不存在越界访问的条件。


Polyhedral Transformation and Mapping

以并行架构为目标时,程序转换涉及调度以及映射策略的改变,必须遵守依赖关系。除了保证转换的有效性之外,依赖关系还可以用于公开并行性(独立实例可以并行执行)或改善数据访问局部性(依赖实例在近距离执行)。存在几种有效的调度算法,着眼于并行性局部性矢量化的组合。我们的转换机制基于四个元素:

1、 isl提供核心的Poly调度:它可以自动优化(外部)循环并行性和局部性;我们调整了仿射调度启发法,将完整的TC程序折叠到单个GPU内核中;

2、调度被进一步平铺以促进在GPU的深度并行性和内存层次上的映射和时间重用;

3、借鉴PPCG实现的映射到GPU算法:并带有附加扩展以支持ML内核中出现的更复杂和不完美的嵌套结构。

4、内存提升通过显式数据往返于共享shared和私有private内存空间的传输。

Scheduling

isl调度程序的核心部分迭代解决整数线性规划问题,以计算形成调度带的分段仿射函数。它还确保这些band是可置换的,并且可以进一步平铺。在内部,它构建一个数据依赖图(DDG),其中节点对应于语句,边表示它们之间的依赖。每个边缘都带有一组“类型化”的依赖关系注释。 isl调度程序设计用于更好的可伸缩性,因为在最坏的情况下整数线性编程具有指数复杂性。特别是,它介绍了仿射聚类技术,该技术基于分别为DDG的各个强连接组件计算调度范围,然后迭代地对这些组件进行聚类并相对于彼此进行调度。聚类不仅减少了调度程序必须解决的线性问题的大小,而且还为isl的循环融合启发式算法奠定了基础。

Poly扩展了isl调度程序,以便为调用方提供对调度过程的更细粒度的控制。对于仿射变换,用户可以提供其他任意约束以插入线性程序。对于聚类,在调度程序证明其有效之后,用户可以为成对依赖图组件组合提供决策函数。这些配置点用作调度策略的基础。利用这些策略,仿射变换可以被限制为:(1)非负调度系数and/or(2)非偏斜变换(即,循环交换和移位)。聚类决策允许控制常规的最小和最大融合目标,此外,最大融合保留了至少三个嵌套的并行循环(映射到CUDA块和线程)。可以通过自动调整过程配置和选择调度策略。在所有情况下,我们都强制生成一个GPU内核。

Imperfectly Nested Loop Tiling不完全嵌套循环平铺

进行调度后,循环平铺将作为一个单独的步骤实施,并作为调度树转换执行。本质上,它将可转换的调度band转换为两个band的链chain,其中外部频段包含tile loops,而内部频段包含point loops,且行程数固定。这可以看作是常规的strip-mine和sink transformation。例如,图3.c显示了融合和平铺的sgemm的进度树。

除了常规的循环嵌套平铺,调度树表示还允许对不完全嵌套的循环进行平铺。该技术基于以下观察:如果一个循环不携带依赖项(即并行),则它可能会沉入任何其他循环下面。在有效的调度中,所有依赖关系都由某个循环来承载(或满足),沿着这些循环具有正距离。仅当依赖关系在某个循环中被另一个循环承载之前为负值时,才会违反依赖关系。并行循环从定义上讲不会携带依赖关系,因此不会影响依赖关系的满足或违反。因此,不完全嵌套的切片是通过首先隔离所有band,然后将并行点循环沉入树中来实现的。在此过程中,将point band复制到序列(或集合)节点下的每个子树中,并更新其调度,以使其仅包含相关的域点。

sgemm的调度树故意有两个不完全嵌套的带。相关性分析表明,循环i和j是并行的。因此,我们可以将它们平铺,并将点环沉入约简k环的带以下,从而生成图3.d中的调度树。检查排列性后,可以将带点环的最内层嵌套带合并为一个band。

Mapping to Blocks and Threads

调度树也可以用于表示到加速器的映射,特别是具有多个块和线程的GPU。通过将某些调度程序band members和相应的循环与线程或块索引相关联来执行此操作。然后,如果可能,Poly代码生成器将省略循环,并相应地重写索引表达式。我们的映射算法是从最初为PPCG设计的算法衍生而来的,其中网格grid和块的大小独立于图块大小而指定。由于块和线程的语义,只能映射属于可置换调度范围的并行循环。如果将点循环映射到线程,则切片大小和块大小之间的比率将控制每个线程执行的迭代次数。请注意,如果块的点循环最终被映射到线程,则小于块大小的切片大小会导致某些线程不执行任何计算。

我们要求调度树至少具有一个最外层的band,其外部尺寸平行。与PPCG相反,PPCG将可能的多个外部带中的每个都映射到块和线程(平铺后),我们的映射算法将单个带映射到块,以便生成ML框架所期望的单个内核,而可以映射多个带线程。 (单个)最外面的band的平行尺寸映射到GPU块。在每个调度树分支中,通常由点循环组成的最里面的可替换带被映射到具有以下限制的GPU线程。跨分支的映射维数必须相等。这可能需要在某些分支中将某些线程尺寸映射为零。

映射本身是通过插入特殊名称,通过上下文节点与代码生成器通信以及将它们与filter节点中的band维度相关联来执行的。对于矩阵乘法示例,我们的映射策略在图3.e中生成了调度树。我们在调度树中引入了一个上下文节点,以指示参数的有效大小以及网格和块大小(分别表示为bx,by和tx,ty,它们代表由blockIdx.xy和threadIdx可能采用的值。 xy)。当已知有效张量大小时,将立即执行此插入操作。还要注意引用bx,by,tx和ty参数的Filter节点。这些节点表示到GPU的映射。

Memory Promotion

我们有兴趣将部分张量提升为shared或private GPU内存。 虽然提升决策是通过试探法做出的,并且在稍后的阶段生成了相应的命令性代码,但是调度树提供了一个方便的接口,用于附加与内存相关的信息。 内存提升基于数组平铺array tile(一种用于软件控制的本地内存的数据切片形式)的概念。 它是数组中大小恒定的潜在跨步块,覆盖了给定(调度)图块中访问的所有元素。 我们重新访问并扩展了PPCG对内存提升的支持。

间接访问数组的提升

我们的内存提升方法还可以处理间接访问的数组。不失一般性,考虑访问O [l + Idx [i] [j]] [k]。我们将O称为外部数组,将Idx称为索引数组。如果是嵌套的间接寻址,则从最内层到最外层将迭代处理外/索引对。虽然外部数组的第一个索引表达式所取的值是静态未知的,但我们仍可以将它们本地缓存为shared_O [l] [i] [j] [k] = O [l + Idx [i] [j]] [k]。由于某些值可以重复,因此仅当外部数组和索引数组均被读取时才可以间接升级,因此对它们的写入可能会导致无法简单合并的不同值。通常,我们要求索引数组具有一个数组切片,即仅访问其固定大小的块。在为外部数组计算数组切片时,我们忽略下标的间接部分(仿射部分照常处理)。然后,我们在提升的外部数组中引入与索引数组中一样多的其他索引表达式。这些新维度上的数组范围完全对应于索引数组的数组切片大小。因此,升级后的数组的元素包含将使用给定索引数组访问的全局数组元素的副本。间接下标仅在从全局存储器复制时使用,而所有其他访问均在代码生成中重写。如果存在多个共享子表达式并在相应维度上具有相同图块大小的间接索引表达式,则对于所有相同的子表达式,在提升的数组中引入单个索引表达式就足够了。

提升启发法

如果存在固定大小的数组磁贴,单个元素被多次访问并且至少其中一个访问不具有内存合并功能,则直接访问的数组将提升为共享内存。 从访问关系和应用于该域的调度可以看到后者:最后一个访问维应与映射到x线程的调度维对齐。 对于间接数组,由于存在这些情况所带来的额外的长内存依赖关系,因此删除了合并要求。 共享内存的总量是固定的,如果所需的共享内存数量会超出可用数量,我们将采用简单的贪婪启发式方法并拒绝升级。

Mapping to Registers

当前没有方法实现比PPCG先前支持的更复杂的寄存器提升策略。基于以下观察,这是一个临时的,务实的选择:

1、无经验证据表明,根据高级规范自动生成低级代码会比汇编产生显着的性能提升;另外,现代的组装的generators现在已经公开可用并且可以重新定位; 

2、将寄存器优化泛化和重新定向到具有多个向量长度,对齐约束和模式的不同体系结构并不是一件容易的事;

3、存在许多不同的策略,包括在寄存器级别进行流水线处理、寄存器轮换、多缓冲、具有集体样式语义的新ISA,其中有些甚至不可用没有明显的变化;

4、通常,此类策略是首先以内部和概念验证的方式实施的,并且可以使用PeachPy之类的工具轻松地进行模板化和重用。 

实施的这种临时状态会对性能结果产生轻微的影响。我们还计划依靠外部库调用,例如,以加快减少速度。


性能

你可能感兴趣的:(Tensor Comprehensions)