题目:TENET: A Framework for Modeling Tensor Dataflow Based on Relation-centric Notation
时间:2021
会议:ISCA
研究机构:北大
TENET
总结
有一个迭代域 D S D_S DS,访问函数 A S , F A_{S,F} AS,F,以及PE Interconnection。我们定义一个仿射变换,用 m × n m\times n m×n的矩阵代表,这个矩阵就是我们将要优化的dataflow,即S和spacetime-stamp的relation。由此,可以得到一系列关系集合:data assignment relation(时空戳和张量的关系,来自dataflow)、spacetime-stamp map relation(将在性能计算中作为约束,由PE interconnection得到)。从而,根据这些集合计算出性能。关系集合的计算存在约束,也即优化问题的约束。利用整数规划的方法,计算出最优解。
I. Introduction
典型的空间架构由处理元件 (PE) 阵列和暂存器组成。空间架构的主要特征是多样化的硬件数据流选择。
hardware dataflow描述了:1)循环实例到PE array的分配和2)PE中循环实例的执行顺序。
现有一些工作的总结:硬件数据流对于实现高吞吐量和低延迟至关重要,因为它决定了 PE 利用率、数据访问模式和片上带宽要求。 不同的张量计算偏好不同的硬件数据流。 例如,谷歌的张量处理单元 (TPU) [22] 使用脉动数据流连接 PE,其中每个 PE 负责一个乘加运算。 而 Camircon [31] 通过多播通信网络连接 PE,其中每个 PE 执行一个点积。 其他空间架构,如 DySER [19] 和 Plasticine [42],以灵活的方式集成 PE 及其互连,因此可以支持更广泛的应用。
现在,a formal notation is still strongly desired to represent hardware dataflow。一个记号(notation)应该能够系统地覆盖整个数据流设计空间,并促进简单而准确的性能建模。现有的两种:compute-centric的度量无法度量数据传输、重用;data-centric的度量的MAESTRO模型有缺陷。
本文定义TENET,其核心部分是relation-centric notation。我们定义以下的关系:1)循环实例和执行计算的 PE 之间的关系,2)循环实例及其在 PE 中的执行顺序,3)PE 和相应分配的张量元素,以及 4)PE与互连网络相连。总体而言,TENET 能够估计各种硬件指标,包括数据重用、延迟、PE 通信带宽和片上内存带宽。
II. Background
A. Spatial Architectures
The PE also contains register files for data storage.
通常,空间架构具有三层存储器层次结构,即 PE 寄存器级、片上暂存器和片外存储器。 为简单起见,我们在对硬件行为进行建模时做出以下两个假设。 首先,ALU 能够执行一次乘加 (MAC) 操作。 其次,通过相邻PE之间互连的数据传输需要一个周期。
在设计数据流时,数据重用是实现高性能和低能耗的关键因素,可进一步分为时间重用和空间重用。时间重用发生在相同数据在不同周期重用时,而空间重用发生在相同数据在不同 PE 处重用时。
B. Notation Basics
In this paper, TENET supports tensor applications with perfectly-nest loops and single unconditional statement.
What is a perfectly nested loop?
A perfectly nested loop is a nested loop, where all the assignment instructions are inside the innermost loop.
Iteration domain:迭代域。例如,
D
S
=
{
S
[
i
,
j
]
:
0
≤
i
<
4
,
0
≤
j
<
3
}
D_S=\{ S[i, j]: 0 \leq i<4,0 \leq j<3 \}
DS={S[i,j]:0≤i<4,0≤j<3}。
Access function:Given a loop instance, the access function returns the tensor elements accessed by the statement S. We use a relation to represent the access function of tensor F.
A
S
,
F
=
{
S
[
n
⃗
]
→
F
[
f
⃗
]
}
A_{S,F} = \{ S[\vec{n}] \rightarrow F[\vec{f}] \}
AS,F={S[n]→F[f]}
例如,在图 1 的 1D-CONV 算子中,张量 Y 的访问函数是
{
S
[
i
,
j
]
→
Y
[
i
]
:
0
≤
i
<
4
,
0
≤
j
<
3
}
\{ S[i,j] \rightarrow Y[i]: 0 \leq i<4,0 \leq j<3 \}
{S[i,j]→Y[i]:0≤i<4,0≤j<3},表示循环实例S[i, j] 访问张量元素 Y[i]。
C. Limitations of Existing Dataflow Notations
III. TENET Overview
IV. 关系为中心的记号
TENET定义了四种关系,包括循环实例到PE阵列的映射(Section IV-A)、数据分配(Section IV-B)、PE互连(Section IV-C)、不同时空标记之间的映射(Section IV- D)。 使用以关系为中心的符号,我们可以准确地确定每个循环实例在空间架构中的执行位置和时间,访问张量元素的位置和时间,以及张量元素如何跨 PE 移动。
A. Dataflow Relation
定义1:Dataflow:Given a statement S with iteration domain DS and iteration vector
n
⃗
\vec{n}
n, the dataflow is defined as
Θ
S
,
D
=
{
S
[
n
⃗
]
→
(
P
E
[
p
⃗
]
∣
T
[
t
⃗
]
)
}
,
S
[
n
⃗
]
∈
D
S
\Theta_{S,D} = \{ S[\vec{n}] \rightarrow (PE[\vec{p}] \, | \, T[\vec{t}]) \}, \quad S[\vec {n}]∈D_S
ΘS,D={S[n]→(PE[p]∣T[t])},S[n]∈DS
它将循环实例 S[n] 分配给一个时空戳,它是一对空间戳PE[p]和时间戳T[t]。Θ的脚标D应该指的就是时空戳的集合,而不是S的迭代域DS。这在后文IV-D部分也能够体现。
上图,
Θ
S
,
D
=
{
S
[
i
,
j
,
k
]
→
(
P
E
[
i
,
j
]
∣
T
[
i
+
j
+
k
]
)
}
\Theta_{S,D} = \{ S[i,j,k] \rightarrow (PE[i,j] \, | \, T[i+j+k]) \}
ΘS,D={S[i,j,k]→(PE[i,j]∣T[i+j+k])}。一般而言,
Θ
\Theta
Θ相当于把loop iterator S通过仿射变换到时空戳。
Dataflow设计空间:n是循环的层数,
n
×
n
n\times n
n×n是仿射变换矩阵的大小,矩阵的每个元素为0或1,所以是2的指数,设计空间大小为
O
(
2
n
2
)
O(2^{n^2})
O(2n2)。我觉得这里应该用
m
×
n
m \times n
m×n的矩阵乘以
n
×
1
n \times 1
n×1的向量(S的索引向量),结果为
m
×
1
m\times 1
m×1的向量,m是时空戳的总维度。例如上面的m=3,PE有两个维度,T有一个维度。
B. Data Assignment Relation
定义2:Data assignment:给定Dataflow
Θ
S
,
D
\Theta_{S,D}
ΘS,D,定义data assignment为:
A的脚标D是时空戳的集合。以图3为例,张量Y的data assignment是:
对于这个例子,我们观察到每个 PE 总是在不同的时间戳计算相同的输出张量 (Y),这意味着张量 Y[i, j] 保持静止,并以不同的时间戳迭代重用,直到计算完成。
C. Interconnection Relation
定义3:PE Interconnection:给定一个 PE array,interconnection是一个集合,其中每个元素描述从一个 PE 到另一个 PE 的映射。
c1, … , ck是用来描述PE拓扑的条件。PE只能通过其他PE或暂存器(scratchpad)获取数据。本文,我们建模了3种常见的PE互连的拓扑:
D. Spacetime-stamp Map Relation
定义4:Spacetime-map:给定一个数据流,spacetime-map是一个关系集,它将一组时空戳 D 映射到另一组时空标记 D’。 D 和 D’ 都来自给定的数据流。
我们可以假设the time distance of D and D’ is within 1, and the PE specified by D and D’ are interconnected. 从而写出一些map,以图3为例:
结合data assignment,有:
上图说明Y[0,0]的复用情况。
上图说明A[0,0]、B[0,0]的复用情况。
实际应用中,我们可以:1)限制D和D’包含相同PE,则可以捕获数据的时间复用(temporal reuse);2)限制D和D’包含相连PE,则可以捕获数据的空间复用(spatial reuse)。
V. Performance Model
A. Data reuse and Volume Model
表2中的sum表示计算集合的元素个数。都是针对某个特定张量F而言的。
TotalVolume:每次用到的数据都被计算。例如,the TotalVolume of tensor A in Figure 3 can be calculated as
ReuseVolume:复用数据。
UniqueVolume:每个时间点的新数据,从暂存器取。
ReuseFactor 描述了从暂存器内存中获取数据后平均重复使用的次数。
SpatialReuseVolume: The SpatialReuseVolume sums up the volume of data with spatial reuse at different space-stamps, where D and D’ contain interconnected PEs only.
TemporalReuseVolume: only refers to the temporal reuse at the same PE, where D and D’ contain the same PE.
因此,有ReuseVolume = SpatialReuseVolume + TemporalReuseVolume。
B. Latency and Bandwidth Model
数据流的延迟包括每个 PE 中的通信延迟和计算延迟。 我们假设buffers、networkds和算术单元以流水线方式工作,并且使用double buffering等技术来隐藏延迟。 那么,整体延迟只是通信延迟和计算延迟的最大值。
PE阵列和scratchpad的通信延迟:
计算延迟:
其中
D
S
D_S
DS是迭代域,sum为迭代域总数,
U
t
i
l
P
E
Util_{PE}
UtilPE是PE的平均利用率。
IBW(interconnection bandwidth):
SBW(scratchpad bandwidth):
C. 模型实现
使用了C++的ISL库和Barvinok库。
VI. Evaluation
benchmark: GEMM, 2D-CONV, Matrix multiplication chain (MMc), Matricized tensor times Khatri-Rao product (MTTKRP), Jacobi-2D