懒猫后花园

哦，这该死的代码

矩阵乘法的对称性：自用 Tile-based 硬件符号分析系统

从算法到具体编译执行，大概包括三个部分：

  • 算法表示： 这里限制到矩阵乘法，本篇相关包括爱因斯坦表示法中心以及对称性；
  • 执行表示： 具体是在目标平台上是怎么执行的，类似目标平台的程序或者编程模型，本篇相关包括 mapping 、 操作、依赖原语等；
  • cost model： 用于量化评估某种参数/程序的效果，简单行为可以用解析公式，复杂硬件一般需要模拟执行，simulator或者是真实硬件，本篇不涉及相关内容。

本篇 blog 的符号系统主要涉及前两个部分，主要关注如果要描述一个矩阵乘法到硬件上执行的“序”，应该定义什么原语，而不关注具体执行的数据延时等。

但即使只包括两层也相当复杂了，即使本身矩阵乘法计算相对规整，tilling / mapping / compute communication overlap / spatial arch 系统相当高维且 NP 复杂，相比 CPU 猜测执行什么的动态复杂度，GPGPU / NPU 的复杂度转移到了在将每一个晶体管都利用到极致降低 “dark silicon” 之上（这里 dark silicon 不是指 TDP 导致罢工，而是指部署 utilization 过低罢工）。

量化的建模在对称性和性能之间平衡， 既要通过对称性降低表达、调度、编译复杂度，又要保证去掉的非对称性不过于影响性能发挥。给出舒服的建模不仅仅和 GEMM 算法本身特性有关，也和具体平台相关。举个例子，简化 sptial mapping 一个方式是通过沿着某个维度均匀分割任务到不同的单元上，但这种做法仅当单元之间平权效果较好，比如集群交换机-服务器总线架构或是一些 mesh dataflow 架构也会有一个总线 global buffer sram ，而放到同构 chiplet 乃至 wafer scale chip，这种对称性假设则会带来很大的 overhead。本篇 blog 不涉及具体建模优劣分析，仅从一些常见对称性给出符号体系。

如何表达执行？从算法到执行

算法既数据和操作的组合，\(A=(D,O)\) 。包含了什么操作、什么数据，以及操作数据之间的依赖关系。理论上程序只要满足依赖关系，就可以以任意先后顺序执行操作，程序并没有描述规定执行的顺序。

而具体程序在硬件上的执行，则是（1）首先将算法切块粒度适应硬件粒度；（2）其次把算法映射到时间和空间之上。为了清晰表达，我们将时间和空间定义为两个符号，实际上时间和空间可能并不存在正交的划分方式（spatial-temporal mapping），为了研究方便，我们先局限在 spatial-only mapping，用正交符号表达 。

映射到空间的函数叫做 spatial mapping，而映射到时间的函数则描述了先后顺序，定义为 temporal mapping。可定义 \(E=(D_h,O_h, M_s, M_t)\)。 其中 \((D_h,O_h)\) 是根据硬件粒度、拓扑适配并映射的程序，这个映射过程定义为 \(A \rightarrow A_h \stackrel{\text{def}}{=} (D\rightarrow D_h, O \rightarrow O_h)\)。

所以表达一个算法的执行包括以下三部分：

  • \(A\rightarrow A_h\) ：tilling，数据和操作都可以通过爱因斯坦标记法的维度表征，对维度切块即同时对数据以及操作切块；
  • \(M_t\) : termporal mapping，定义同一个硬件单元时间上执行的先后顺序；
  • \(M_s\)：spatial mapping，定义数据和维度怎么映射到硬件。

如果机器完全依靠数据依赖执行则是图处理器，所有的执行都变成一张图混杂在一起，硬件上可能一个资源面临很多相同来源请求，仲裁路由等调度算法使得程序执行更难以分辨。为了编程清晰不掉 san 方便往往会引入 barrier 等的同步操作，引入全局或者指定局部区域的同步点互相等待； 亦或者依赖图出入节点太多，超出了指令长度，更改依赖关系便于表达[1]。即实际执行的依赖关系既可能来自原始的算法依赖关系，也可能来自人为插入的依赖关系。 为了方便分析，仅讨论来自原始算法的依赖关系。

怎么简化表达？矩阵乘法以及对称性

刚刚的分析介绍了表达执行需要什么原语，那么将 workload 限制到 tensor 处理，怎么简化表达呢？先从矩阵乘法以及 nested loop 开始引入，如下是一个两层系统（off-chip memory, on-chip SRAM & compute），只表达计算的矩阵乘法循环。为表达方便，一般省略了初始值赋 0 的表达。

# ---- off-chip loop ----
for tt in Tt
	for ot in Ot
		for it in It
#		---- inner-chip loop ----
			for ts in Ts
				for os in Os
					for is in Is
						y[tt*Ts+ts][ot*Os+os] += w[ot][it][os][is] * x[tt*Ts+ts][it*Is+is]

从爱因斯坦标记法的角度矩阵乘法可用三个维度表达，一般 infra 常用符号 M、N、K，我个人一般还是喜欢用 T（Token）、I（Input Channel）、O（Output Channel） ，语义性更强一点。从个人经验来看，简化抽象主要依据几个对称性：

  • 维度对称： 三个维度 T、I、O 互相平权没有差异，此时可以将三个维度看作统一整体，用原语进行操作；
  • 数据对称： 矩阵乘法三个数据 X/W/Y 互相平权，享用相同层次的维度切分方式，如下图来自 ZigZag DSE 框架对 X/W/Y 平权或平权 tilling 的示意图 [2]；
  • 可分离域： 嵌套循环的迭代空间是矩形迭代域/可分离域，其迭代空间可以写成笛卡尔积的形式；反之则称作耦合域和不可分离域；
  • 均匀 tilling： 对于一个维度是按等比例切多个格点，还是按不均匀步长切割格点。考虑到数据在空间的映射，均匀 tilling 一般是在全局性较好的架构上比较合适，比如有统一 global buffer 总线等，对于分布式 mesh 拓扑天然决定了不同位置算力资源利用特性不同。 对于常见空间并行切割数据方法，比如 tensor parallism 或是 data parallism ，都是均匀切割策略。

这三个对称性也不总是满足，而且大都时候要么同时满足，要么同时不满足。

Extra/Images/IMG_20260804152102941.png

Extra/Images/IMG_20260812185732972.png

基础原语介绍

Tiling： 划分多层循环顺序的语句，也可以根据选择不显式写出。强行每个维度划分的层次相同，如果某个变量该层没有，则取1

Tiling(tile_size1=(TT,OT,IT),tile_size=(TS,OS,IS)) # 切两层
Tiling(tile_size1=(TD,OD,ID),tile_size2=(TT,OT,IT),tile_size=(TS,OS,IS)) # 切三层

或者用变量式方便表达语义

D = (TD,OD,ID) # Off-chip Temporal Mapping
T = (TT,OT,IT) # On-chip Temporal Mapping
S = (TS,OS,IS) # On-chip Spatial Mapping
Tiling(tile_size1=D,tile_size2=S,tile_size=S)

实际上，如果知道完整循环的维度 T,O,I ，那么切分N 层，只需要 N-1 自由度就可以表征。

数据对象： tile 原语，一个数据应该用 index 和 size 描述，index 是 hierarchy 的可以扩展，比如 index_1, index_2, ...，默认维度顺序按 (T,O,I) 排序。默认 id 越低，共用的人越多，越靠近高层 (off-chip Memory)，id 越高越靠近底层（Compute Unit）。结合前文所说的自由度，理论上只用 tile_id 也足够表达某个 tile。

Tile(tile_id1=(tt,ot,it), tile_size=(TS,OS,IS))
Tile(tile_id1=(tt,ot,it), tile_id2=(td,od,id), tile_size=(TS,OS,IS))

当非单 core 架构中，一个数据可能在系统中存在多个副本，可添加额外标记符指定。

Tile(tile_id1=(tt,ot,it),tile_size=(TS,OS,IS), location=Core(0,1))
Tile(tile_id1=(tt,ot,it),tile_size=(TS,OS,IS), location=Any)

以及 tile 存在不同粒度，因此 tile 之间存在包含关系，定义了划分方法归档了两个 tile 之间只存在包含或者无交集关系，不会存在存在交集都不是包含关系的情况。

操作原语： 分为计算原语和访存原语，访存原语不区分具体的方向和变量，根据 tile 表征语义自动拆开

# 原始 load 
load(src=DDR, tar=SRAM, data=X(size=(TS,IS),id=(tt,it)))
load(src=DDR, tar=SRAM, data=W(size=(OS,IS),id=(ot,it)))
store(src=SRAM, tar=DDR, data=Y(size=(TS,OS),id=(ot,ot)))

# 统一访存语句
load_store(src=DDR, tar=SRAM, data=Tile(tile_id1=(tt,ot,it),tile_size=(TS,OS,IS)))

# 原始计算语句
Y(size=(TS,OS),id=(ot,ot)) += X(size=(TS,IS),id=(tt,it)) *  W(size=(OS,IS),id=(ot,it))

# 统一计算语句
compute(data=Tile(tile_id1=(tt,ot,it),tile_size=(TS,OS,IS)))

Temporal Mapping：NestLoop 循环原语： 将多层循环用统一原语表征，减少重新循环语句，其中可以将符号看作有序，根据 t/i/o 区分。也可以看作无序只是通用表征。根据具体使用来定。

# 原始嵌套循环
for tt in Range(TT):
for ot in Range(OT):
for it in Range(IT):

# 等价嵌套循环
for (tt,ot,it) in Range(TT,OT,IT): # 这里表示了循环从外层到内层按照 tt,ot,it 的顺序，并且每个顺序是从 it=0, it=1, ... 这样的递增顺序

如果破坏了三个维度对称性的时空映射：

# Output Stationary Systolic Array
for (td,od,id) in Range(TD,OD,ID)
for (tt,ot) in Range(TT,OT)
for tau in Range(StreamingEquation) # 这里根据 it / row / col 的值，给出 tau 从 0 到最大值的范围，Streaming Equation 在后续破坏对称性介绍

Intra-tile 依赖原语：

同一个数据 tile，不同硬件单元之间存在依赖，比如一个操作是另一个操作的 producer 。一个数据块假设存在互斥的4种状态，未诞生、正在写入、就绪、正在读取。其中值得注意的是 reading 状态，由于 reading 的时候允许有多个单元同时读取（multicast/broadcast），所以 reading 状态中会维护一个计数器，每次受到对同一个数据的 start_reading 计数器 +1，每次受到 finish_reading 计数器 -1，当计数器归 0 返回 ready 状态。

之前 torchbit[3] 和 RTL 电路的开发经历告诉我，想要在统一 nestloop 里面表征各个错位的流水线复杂度极高，借鉴类似 FIFO 的依赖语义是更舒适的表征，虽然代码语言还是不如流水线图直观

Extra/Images/IMG_20260730165355986.png

# 状态转移语句，这些语句自动带等待语义，等待前置状态达成
start_writing(tile=xx)
finish_writing(tile=xx)
start_reading(tile=xx)
finish_reading(tile=xx)
destory(tile=xx)

Inter-tile 依赖原语： 不同的 tile 之间也会产生竞争关系，比如对 SRAM 容量存储的竞争、对链路带宽的竞争。这里同样需要引入 buffer 原语或者带宽原语等。比如假设 pingpong buffer ，每次需要占用 lock 一个 buffer 才能获取带宽。

buf = Buffer(slot=2)
buf.acquire()
buf.release()
buf.full
buf.empty

又或者 NoC 请求，有两种策略，一种是当请求产生后或者完整链路带宽使用，另一种则是一个链路允许同时有多个请求，做分组交换时分复用资源。无论那种只是具体任务完成时长不同，在依赖层面仍然是对资源的依赖构建。不过 NoC 的访问吸收在 load_store 原语。

Spatial Arch 原语扩展

因为 spatial 相关原语破坏了对称性，这里专门分小节介绍。

Spatial Mapping: 刚才的符号直接划分了 mapping 属于特殊情况，tile 的维度和 mapping 划分方式完全相同，tile 层次和 mapping 有对应关系（Spatial-only Mapping）

D = (TD,OD,ID)
T = (TT,OT,IT)
S = (TS,OS,IS)
Tiling(tile_size1=D,tile_size2=S,tile_size=S)
TemporalMapping(tile=D)
TemporalMapping(tile=T)
SpatialMapping(tile=S)

如果数据在时空间流动（Spatial-Temporal Mapping），那么便无法划分出这种形式，需要时空耦合的时空流方程，此时会破坏三个维度作为整体出现的对称形式。

D = (TD,OD,ID)
T = (TT,OT,IT)
S = (TS,OS,IS)
Tiling(tile_size1=D,tile_size2=S,tile_size=S)

# Systolic Array(Output Stationary)
Spatial = (ROW, COL)
SpatialMapping("S.ts->Spatial.row","S.os->Spatial.col") # IS=1
StreamingEquation("tau=Spatial.row+Spatial.col+T.it")
TemporalMapping(tile=(T.tt,T.ot))
TemporalMapping(tile=D)

将时间索引和 StreamingEquation 联立，可以得到每个时间步芯片内合理的数据 Tile 范围以及他们的空间坐标，对上一个时间索引联立求位置变化，即可得到数据流动起点终点。但这个表征只是规定了起点和终点，没有规定具体沿着什么路径流动。

举例

不考虑流水线错位的计算-访存表达

D = (TD,OD,ID) # Off-chip Temporal Mapping
T = (TT,OT,IT) # On-chip Temporal Mapping
S = (TS,OS,IS) # On-chip Spatial Mapping
Tilling(tile_size1=D,tile_size2=T,tile_size=S)

for (td,od,id) in Range(TD,OD,ID):
	load_store(src=DDR, tar=SRAM, tile_id1=(td,od,id),tile_size=(TT*TS,OT*OS,IT*IS))

for (tt,ot,it) in Range(TT,OT,IT):
	load_store(src=SRAM, tar=Compute,tile_id1=(td,od,id),tile_id2=(tt,ot,it), tile_size=(TS,OS,IS))
	compute(tile_id1=(td,od,id),tile_id2=(tt,ot,it), tile_size=(TS,OS,IS)) # 直接将 Spatial 看作一整块计算，不做 scalar 粒度的表征

Systolic Array

（Output Stationary）Systolic Array 输入数据需要产生偏斜电路，读出数据则需要读出 mux 将 \(O(n^2)\) 的单元映射到 \(O(n)\) 的 buffer 接口。天然不满足维度对称性，无法使用 load_store() 原语，下方伪代码感性感受一下计算过程的表达。

Extra/Images/IMG_20260730165356159.png

D = (TD,OD,ID) 
T = (TT,OT,IT) 
S = (TS,OS,IS)
Tilling(tile_size1=D,tile_size2=T,tile_size=S)

# Systolic Array(Output Stationary)
Spatial = (ROW, COL)
SpatialMapping("S.ts->Spatial.row","S.os->Spatial.col") # IS=1
StreamingEquation("tau=Spatial.row+Spatial.col+T.it")

TemporalMapping(tile=(T.tt,T.ot))
TemporalMapping(tile=D)

def Memory_SRAM_thread:
	for (td,od,id) in Range(TD,OD,ID):
		load_store(src=DDR, tar=SRAM, tile_id1=(td,od,id),tile_size=(TT*TS,OT*OS,IT*IS))

def SRAM_Skew_thread:
	for (td,od,id) in Range(TD,OD,ID):
		for (tt,ot) in Range(TT,OT):
			for tau in Range(ResolveRange(StreamingEquation)):
				if tau in Range(IT):
					load_store(src=SRAM, tar=Skew, tile_id1=(td,od,id), tile_id2=(tt,ot,tau), tile_size=(TS,OS,IS))

def Skew_Compute_thread:
	for (td,od,id) in Range(TD,OD,ID):
		for (tt,ot) in Range(TT,OT):
			last_active = Empty()
			for tau in Range(ResolveRange(StreamingEquation)):
				current_active = ResolveActive(StreamingEquation, tau)
				load_data = current_active - last_active
				load_store(src=Skew, tar=Compute,data=load_data)
				compute(data=current_active)

Pingpong Buffer

如果考虑访存和计算构成流水线，那么要同时建模流水线的前向和反压过程。Pingpong Buffer 可以看作 slot=2 的 Tile FIFO：Buffer 负责不同 Tile 竞争 SRAM 槽位的 Inter-tile 容量依赖，Tile 状态负责同一物理副本在 producer 和 consumer 之间的 Intra-tile 数据依赖。同一物理副本由 tile_id + location 确定；producer 先用 buf.acquire() 等待并占用空槽，finish_writing(tile) 将该副本置为 ready，consumer 的 start_reading(tile) 自动等待 ready，全部读取结束并销毁副本后再用 buf.release() 归还槽位。这样 producer 写下一个 Tile 时，consumer 可以同时计算当前 Tile。

# 有 Pingpong Buffer
D = (TD,OD,ID) # Off-chip Temporal Mapping
T = (TT,OT,IT) # On-chip Temporal Mapping
S = (TS,OS,IS) # On-chip Spatial Mapping
Tilling(tile_size1=D,tile_size2=T,tile_size=S)

core = Core(0,1)
sram = SRAM(core)
buf = Buffer(slot=2) # 两个槽位，统计 Inter-tile 容量依赖

def Memory_SRAM_Thread:
    for td, od, id in Range(TD, OD, ID):
	    tile = Tile(
	        tile_id1=(td,od,id),
	        tile_size=(TT*TS,OT*OS,IT*IS),
	        location=core,
	    )
	    buf.acquire() # 自动等待空槽
	    start_writing(tile=tile)
	    load_store(src=DDR(0), tar=sram, data=tile)
	    finish_writing(tile=tile) # 发布 ready，唤醒等待该 Tile 的 consumer

def SRAM_Compute_Thread: 
    for td, od, id in Range(TD, OD, ID):
	    tile = Tile(
	        tile_id1=(td,od,id),
	        tile_size=(TT*TS,OT*OS,IT*IS),
	        location=core,
	    )
	    start_reading(tile=tile) # 自动等待相同 tile_id + location 的副本 ready
	    for tt, ot, it in Range(TT, OT, IT):
	        child_tile = Tile(
	            tile_id1=(td,od,id),
	            tile_id2=(tt,ot,it),
	            tile_size=(TS,OS,IS),
	            location=core,
	        )
	        load_store(src=sram, tar=Compute(core), data=child_tile)
	        compute(data=child_tile)
	    finish_reading(tile=tile)
	    destory(tile=tile)
	    buf.release() # consumer 完成后归还槽位，解除 producer 反压

Mesh Dataflow Chiplet

此案例施工中

Extra/Images/IMG_20260730165356326.png

# 分布式 DDR 颗粒 + Mesh Chiplet + Double Buffer
# 硬件描述
Chiplet = (COL,ROW) # Chiplet 的坐标空间，比如 (0,0) 是左上角第一个 Chiplet
DDR_Cluster = (COL_D, ROW_D) # DDR 颗粒的空间坐标，比如 (0,0) 是左上角第一个 DDR 
Buffer(slot=2, array=Chiplet) # Inter-tile 统计
# Adjacent 操作符，可以表达邻接的节点，比如 Adjacent(col_d,row_d) 得到临界的 Chiplet 节点

# Tilling
DS = (TDS, ODS, IDS) # Off-chip DDR Spatial Mapping
DT = (TDT,ODT,IDT) # Off-chip Temporal Mapping
C = (TC,OC,IC) # inter-chip chiplet Spatial Mapping
T = (TT,OT,IT) # On-chip Temporal Mapping
S = (TS,OS,IS) # On-chip Spatial Mapping
Tilling(tile_size1=DS,tile_size2=TD,tile_size3=C, tile_size4=T,tile_size=S)

# Mapping
SpatialMapping("col_d->DS.tds, row_d->DS.ods")
TemporalMapping(tile=TD)
SpatialMapping("col->C.tc, row->C.oc")
TemporalMapping(tile=T)
SpatialMapping(tile=S)

def Memory_SRAM_Thread:
	for tds, ods, ids in Range(DS):
	    for tdt, odt, idt in Range(DT):
		    adjacent_chiplet = Adjacent(tds,ods,ids)
		    if not Buffer[adjacent_chiplet].full():
			    Buffer[adjacent_chiplet].increase()
			    start_writing(tile_id1=(tds,ods,ids),tile_id2=(tdt,odt,idt))
		        load_store(src=DDR[tds,ods,ids], tar=SRAM[adjacent_chiplet], tile_id1=(tds,ods,ids),tile_id2=(tdt,odt,idt),tile_size=(TC*TT*TS,OC*OT*OS,IC*IT*IS))
				finish_writing(tile_id1=(tds,ods,ids),tile_id2=(tdt,odt,idt))
				
def Chiplet_Swap_Thread:
	for tc, oc, 

def SRAM_Compute_Thread: 
    for td, od, id in Range(TD, OD, ID):
	    start_reading(tile_id1=(td,od,id))
	    for tt, ot, it in Tile_Space(TT, OT, IT):
	        load_store(src=SRAM, tar=Compute,tile_id1=(td,od,id),tile_id2=(tt,ot,it), tile_size=(TS,OS,IS))
	        compute(tile_id1=(td,od,id),tile_id2=(tt,ot,it), tile_size=(TS,OS,IS))
		finish_reading(tile_id1=(td,od,id))
		destory(tile_id1=(td,od,id))
		Buffer.decrease()

  1. https://arxiv.org/abs/2503.20481 ↩︎

  2. https://arxiv.org/abs/2007.11360 ↩︎

  3. https://github.com/Devil-SX/torchbit ↩︎

posted @ 2026-07-28 15:03  DevilXXL  阅读(36)  评论(0)    收藏  举报