生成TilingCase
在自动融合技术中,tiling解空间生成是实现高效计算调度的关键环节,其核心目标是为复杂计算任务提供多样化的分块(tiling)策略选项,以便后续优化器从中筛选出最优解。简单来说,tiling解空间生成的过程可以理解为对输入数据或计算任务进行“分块可能性”的系统枚举,每一个解空间被称为一个TilingCase。
切分方式的设计与算子实现特性紧密关联,为了实现对多样化算子的系统性切分策略枚举,首先基于算子的实现特性将其抽象为9种compute_type(如下图中的compute类和view类算子)。同一compute_type类别的算子具有相似的计算逻辑与数据访问模式,因此可以共享一套tiling切分策略框架。

为具象化这一策略框架,我们对算子的轴进行归一化分组,将所有轴统一划分为(xgroup、ygroup、rgroup)三个维度集合,具体定义如下:
- xgroup:专为Transpose、Concat、Split等视图类算子设计的分组。例如,对Concat会以concat_dim进行划分,将concat_dim前的轴划分到xgroup,将concat_dim及其以后的轴划分到ygroup。
- ygroup:对应Elementwise、Broadcast等类型算子的循环轴分组。
- rgroup:Reduce类操作通常对reduce轴有特殊的切分要求,因此会将所有reduce轴单独放入rgroup。
引入 xgroup、ygroup、rgroup 的核心原因是为了支持复杂场景下的“双切分”需求。例如,在包含reduce混合的计算图中,ygroup控制elementwise的循环切分,rgroup中的轴控制reduce方向的循环切分。
在完成单算子的轴分组后,需通过预设的合并规则(Merge)对计算图中所有算子的分组策略进行合并。合并结果将作为适用于全图所有算子的统一切分策略,为后续解空间生成提供基础。
通过上述分组与合并机制,可实现两大核心功能:
- 筛选出适用于计算图中所有节点的切分方式,形成有效的TilingCase。
- 通过判断不同AscGraph的tiling分组是否能成功合并,验证两张图的可融合性。
以下面case为例,介绍通过tiling分组合并生成TilingCase的原理:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | data = ascir.ops.Data('data', graph) data.y.dtype = ascir.dtypes.float32 # 声明load算子 load = ascir.ops.Load('load') load.attr.sched.axis = [z0, z1] # 调度轴 load.x = data.y load.y.axis = [z0, z1] # tensor的输出轴 load.y.size = [s0, s1] # tensor输出大小 load.y.strides = [s1, 1] # tensor的输出步长 # 声明abs算子 abs = ascir.ops.Abs('abs') abs.attr.sched.axis = [z0, z1] abs.x = load.y abs.y.axis = [z0, z1] abs.y.size = [s0, s1] abs.y.strides = [s1, 1] # 声明max算子 max = ascir.ops.Max('max') max.attr.sched.axis = [z0, z1] max.x = abs.y max.y.axis = [z0, z1] max.y.size = [s0, 1] # 对z1轴进行reduce操作 max.y.strides = [1, 0] |
abs是一个elewise算子,每个轴在计算上并无差异,因此只要内存连续,可以将多根轴合并成一根轴进行切分:例如图1所示,先在轴上做block切分,15被分到3个block上,block0为紫色的部分,在block0内再进行tiling分块,此时tiling块未占满block0分配的部分,因此还需要在block内加一个for循环。
reduce类的切分较为复杂,实现上需要双切分:行方向上是elewise轴,列方向上是reduce轴;首先在行方向上分block,block内写循环,然后在列方向上再加一层for循环。

通过TilingGroup的合并规则:
1 | ()(z0,z1)() Merge ()(z0)(z1) => ()(z0)(z1) |
将abs算子原本同属ygroup的(z0,z1)轴进行拆分,其中z1轴被调整至rgroup,这一调整的核心目的是使abs算子与后续的reduce类型算子保持统一的切分策略。
