LowerAutoVectorSplit Pass(向量自动拆分下降)¶
在 ExpandMixedKernel 之前,将带 AUTO pl.split 的混合 InCore 函数转换为
显式 split_aiv 形态:在 cube→vector 边界插入 tile.aiv_shard,在
vector→cube 边界插入 tile.aic_gather,仅对向量子区域沿拆分轴折半,注入
tile.get_subblock_idx(),并在函数上打 split + split_aiv 标记。
这是唯一的自动拆分下降路径:它始终运行,紧邻 ExpandMixedKernel 之前。运行后
每个拆分函数到达 SplitVectorKernel 时都已带
split_aiv 标记,因此该 pass 只打属性(其 split_aiv 分支)——其旧的逐算子折半驱动
已被删除,折半机制现仅存于 split_axis_utils,由本 pass 共享。
本 pass 就地下降一等公民区域节点 SplitAivScopeStmt 并保留其包装;
ExpandMixedKernel 在 codegen 前消费该结构。
为什么需要本 pass¶
用 pl.split 编写的混合 InCore 函数在同一函数体中描述 cube 与 vector 工作,拆分
意图仅由函数级 split 模式表达。实现该拆分有两种方式:
SplitVectorKernel中的后期逐算子折半 —— 在ExpandMixedKernel已经把函数 体分为带跨核tpush/tpop的 AIC + AIV 之后,再逐算子折半 AIV 函数体。这重复了tile.aiv_shard/tile.aic_gather已经编码的边界语义。- 早期显式下降(本 pass) —— 在
ExpandMixedKernel之前,把 AUTOpl.split函数体改写为手写显式核所用的同一split_aiv形态。随后ExpandMixedKernel中 单一的算子驱动边界分支会统一地把tile.aiv_shard/tile.aic_gather折叠为带 拆分标记的tpush/tpop——自动核与手写核走完全相同的下游路径。
方式 2 是当前路径。它与旧的逐算子折半逐字节一致(分阶段收敛期间已验证),因为两者调用
同一套 split_axis::ProcessStmts 机制,仅入口与边界处理不同。
物理切分轴长度为 1 的 load 是可复制的广播读取,保持原形状,与 AUTO 一致。它不会被标记为半宽值,不能为无关消费者提供半宽依据。
API¶
| C++ | Python | 层级 |
|---|---|---|
pass::LowerAutoVectorSplit() |
passes.lower_auto_vector_split() |
Program 级 |
Pass 属性¶
| 属性 | 值 |
|---|---|
| Required | SSAForm、IncoreTileOps、SplitIncoreOrch、TileOps2D、TileMemoryInferred、NormalizedStmtStructure、AivSplitValid |
| Produced | SSAForm、IncoreTileOps、SplitIncoreOrch、TileOps2D、TileMemoryInferred、NormalizedStmtStructure、AivSplitLoweredValid |
| Invalidated | AivSplitValid |
本 pass 以 AivSplitLoweredValid 接替源阶段的 AivSplitValid。共享验证器继续检查
区域结构,同时允许受支持的 flat lowered 形式。
来源:include/pypto/ir/transforms/pass_properties.h
(kLowerAutoVectorSplitProperties)。
作用范围¶
仅当全部满足时改写函数:
func_type_ == FunctionType::InCore,且- 带函数级拆分模式(
UpDown/LeftRight,mode != None),且 - 尚未为
split_aiv(手写显式核保持不动——它们已带显式 shard/gather 形态),且 - 确为混合(cube↔vector):其汇总亲和性为
MIXED,与ExpandMixedKernel判定is_mixed所用的ClassifyCallAffinity/CombineAffinity完全一致。
其余一律原样透传。最后一条很关键:纯向量 pl.split 函数(例如把一个逐元素算子
拆到两个 AIV lane,既无 cube 也无 C↔V 边界)没有可收敛的边界,故保持不动——
ExpandMixedKernel 会照旧把它转成普通 AIV 函数并剥掉其 split 属性,保留其原先
(未拆分)的行为。若在此处对其下降,剥离后它将只带 split_aiv 而无 split 模式,
SplitVectorKernel 会因此报错。
函数体中仍存在 ScopeStmt 的混合函数会被拒绝,理由与下文的区域路径相同:整函数折半同样
不跨越 scope 边界。汇总亲和性会穿过 scope 计算(scope 的亲和性即其函数体的亲和性),因此本
pass 能区分「被 scope 包裹的混合函数」与「纯向量函数」,只拒绝前者——纯向量函数体仍按原样透传,
不会报错。该情况通常不可达:本路径所依据的函数级 split 属性由 OutlineIncoreScopes 写入,而它
在同一步就消费掉了 scope;之所以校验而非假定,是因为另一种失败方式是静默的。在加入该校验之前,
RollupAffinity 根本没有 ScopeStmt 分支,于是任何被 scope 包裹的函数都汇总为 SHARED、被判定
为非混合,从而完全未拆分地透传且没有任何诊断。
显式 SplitAivScopeStmt 区域路径¶
除上述 AUTO 整函数路径外,函数体仍携带一个或多个 SplitAivScopeStmt 区域的
InCore 函数走一条独立的区域路径(LowerExplicitRegionFunction),它在 AUTO
路径之前判定。每个区域携带各自的 split_ 模式,因此可处理单一函数级模式无法表达
的多模式情形。区域局部的 tile_vars / var_replacements 映射保证折半后的变量不会泄漏
到同级区域或区域外的算子。任何区域之外的语句以全宽发出,且永不折半。所有区域下降后,
区域包装保留,函数标记为 split_aiv。ExpandMixedKernel 根据每个区域自身的模式
检查转置风险,并消费区域。
共享函数体检查与 AUTO 区域¶
AnalyzeSplitBody 同时检查变换后的 AUTO body 和显式 manual body。AUTO 在替换与克隆
之前提供形状变换已建立的 tile 事实;manual 从 shard、lane 地址、别名、tuple projection
和循环结果重建数据流。广播读取与只读 singleton 计算保持中立,不会替无关的全宽消费者
证明已分片。rank-1 load 与全宽 reshape/reinterpret view 可作为后续 lane-local slice
的中间值,但不会因此成为半宽事实。
控制流合并会逐 tuple 元素取两个分支的 shard 事实交集。循环回边必须保留从初值继承的 shard 事实。中性初值可以在循环体中变为按 lane 的值,但循环携带值与出口仍保持中性,避免仅根据 yield 将零次迭代的出口误判为按 lane 切分。
检查诊断分别记录全宽算子名(full_width_vec_ops)与循环携带变量名(carry_mismatches)。携带值不一致时,说明入口的 shard 事实在回边丢失,并提示保留按 lane 的 yield;算子未本地化时,报告算子名,并提示使用按 lane 的操作数或本地化读取地址。显式边界区域的检查失败使用面向用户的 CHECK_SPAN 文案;隐式或 AUTO 折半后的失败属于编译器后置条件违例,使用 INTERNAL_CHECK_SPAN,不附带作者修改建议。两类问题同时存在时,优先报告携带值不一致。
AUTO 完成下降后,只在同一结构验证器认可时包装单个直线向量阶段,保留计算顺序和 lane 变量身份,仅在区域外没有使用者时把 lane 绑定移入区域。
向量阶段之后、下一个计算阶段或 return 之前的 SHARED 调用归入该向量阶段。交错的 cube/vector 阶段、控制流、带 lane_stride 的重平衡边界、迁移后的边界轴
继续使用 flat lowered 形式。fallback 保留折半、偏移本地化与边界检查。
本 pass 消费源阶段的 AivSplitValid,产生 AivSplitLoweredValid;后者同时支持保留/合成
区域与 flat fallback,由 ExpandMixedKernel 要求并失效。
区域外契约(手动模式)¶
“以全宽发出”描述的是本 pass 对区域外语句所做的处理,而不是作者在那里可以写什么。
持有至少一个区域的函数即进入手动模式:区域对向量计算的放置具有决定权,而
AivSplitValid 验证器早在本 pass 之前就会强制这一分工:
| 算子 / 值 | 区域内 | 所有区域之外 |
|---|---|---|
| 向量计算 | AIV | 拒绝 —— 检查 (e) |
tile.load / tile.store |
AIV | 允许(编译器物化) |
| cube 计算 | 拒绝 —— 检查 (a) | AIC |
aiv_shard / aic_gather |
边界本身 | 拒绝 —— 检查 (c) |
no-duplicate 算子(pld.system.notify) |
钉在 AIV lane 上 | 复制到两条 lane 上(不会诊断) |
因此本 pass 在区域外真正会遇到的语句是:cube 计算、ConvertTensorToTileOps 从承载其
计算的区域中提升出来的 tile.load / tile.store 对,以及与核无关的标量 / 控制流语句。
全宽向量计算不在其列:若要让某个阶段保持全宽,请把它包进
for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE):,其含义恰是“两条 AIV lane 都运行
完整函数体”。因此多模式的目标是只有区域、每个向量阶段一个区域。没有任何区域的函数
不受手动模式影响。校验器的检查 (f)/(g) 还要求跨越区域边界的 tile 必须用边界算子写明,
因此隐式的 cube↔vector 跨越也不会到达本 pass。
写在所有区域之外的通信算子不会被拒绝,只是同样被复制到 cube lane 上(参见
verify_aiv_split.cpp 中的 “NOT CHECKED, DELIBERATELY” 注释)。下文的放置标记也并不能让
区域意味着“恰好一次”:AIV 函数带有 dual_aiv_dispatch,其函数体会在两条 AIV 子 lane 上
都运行,把只应发生一次的副作用在两条子 lane 之间分片是作者的职责,None 区域 V→C 跨越的
lane 规则同理(见作用域与放置)。
ExpandMixedKernel 消费区域放置¶
区域本身保留放置信息。ExpandMixedKernel 每个函数只擦除一次包装,并在消费后的 body
上用 pass 内临时映射记录需要 AIV 放置的调用,不再序列化放置或验证属性。
| 固有亲和性 | 区域的作用 |
|---|---|
SHARED、未声明 lane、set_no_duplicate()(notify) |
仅 AIV |
可安全复制的 SHARED(wait) |
保留两侧 |
VECTOR |
已属于 AIV |
显式声明 lane(tile.create、core_type) |
尊重声明 |
MIXED 边界 |
保留传输两端 |
CUBE 计算 |
区域内拒绝 |
该规则同时适用于 AUTO 合成区域、manual 区域与纯 AIV 函数。它不保证副作用在两条 AIV 子 lane 中只执行一次,作者仍需通过 lane 索引分配工作。
函数级 AUTO split(optimizations=[pl.split(mode)],包括 SplitMode.NONE)与显式
pl.split_aiv 区域是互斥的;若需在携带区域的作用域上指定自定义跨核槽位数,请使用
optimizations=[pl.cross_core_slot(slot_num=N)]:它只决定 pipe 大小,不标注任何拆分。
该检查在更早的 OutlineIncoreScopes 中执行,那里作用域自身
的 split_ 与其区域都仍可见;否则本区域路径会按区域下降并静默丢弃函数级 split。(提取后
二者会无法区分地合并:单个区域会合法地派生出一个函数级代表 split 模式。)
按区域的 split_ 模式处理三种区域体形态:
- 数据并行 · 全宽体(
UpDown/LeftRight,无显式边界算子):区域体持有全宽向量计算。 区域路径注入按区域的subblock_idx,将向量算子路由到共享的split_axis::ProcessStmts折半机制(按区域局部进行),并校验按区域的转置风险。这是自动收敛形态产生的范式。 - 数据并行 · 显式边界体(
UpDown/LeftRight,已存在tile.aiv_shard/tile.aic_gather):用户已手动切分 cube tile 并在每 lane 的半块上编写向量计算,故区域体 已是半宽形态。区域路径检测到这一点(RegionBodyHasExplicitBoundary)后原样透传区域 体——不再折半,也不注入重复的subblock_idx。若在此处再次折半将导致双重切分(下游的 Acc→Vec move 会被误判为新的 cube→vector 边界并被改写成第二个aiv_shard),从而产生一个 无任何分配的孤立 Acc memref 并使 PTO codegen 崩溃。ExpandMixedKernel会像处理手写 split_aiv 核一样,把显式边界折叠为tpush/tpop。 - 任务并行体(
None):没有拆分轴——两个 AIV lane 都运行完整区域体,由作者通过 区域的aiv_idlane 索引(例如按aiv_id跨步的循环)分派各自不相交的工作。区域路径 原样透传区域体(不折半、不本地化偏移、不注入subblock_idx;作者的aiv_id = get_subblock_idx()绑定已携带 lane 信息)。此处的tile.aiv_shard/tile.aic_gather是被接受的,并与其余语句一同透传:没有拆分轴时它只跨越 AIC/AIV 边界 而不切分,其split=0类型推导原样保留形状,因此没有可折半或可拼合的东西。本模式下会跳过ValidateMixedExplicitRegion——它拒绝的是「半宽边界算子与全宽向量算子混写」,而这里一切 都是全宽。该函数仍会被标记split_aiv,因此下游ExpandMixedKernel/SplitVectorKernel会把它派发到两个 AIV lane(经由dual_aiv_dispatch),而非 lane-0-only 的非拆分 replay——故从这类区域向外的 V→C 跨越上两个 lane 都会 push,写入同一个 共享槽位且没有任何仲裁,因此除非作者保证该值 lane-uniform,cube 收到的是二者之一且不确定 是哪一个。当区域的 tile 无法折半(单位维)或归约必须保持全宽时使用本模式。
显式边界区域内允许出现哪些算子¶
由于该区域体是原样拼接的,其中每个 vector 算子都必须已经是 per-lane 的——保持全宽的
算子会在两个 AIV lane 上运行出完全相同的结果。ValidateMixedExplicitRegion 负责校验这一
点,并以可操作的错误信息列出违规算子。满足以下任一条件的 tile 生成算子会被接受:
| 接受条件 | 原因 |
|---|---|
(传递地)消费了 tile.aiv_shard 的结果 |
依构造即处于 half-width 数据流中。 |
纯生成算子——tile.full / tile.ci / tile.random(以及 tile.create,它归类为 SHARED,本就不会被报告) |
其结果仅是自身属性的函数:不读取任何 tile、不读取内存,因此无论作者写的是什么 extent,per-lane 复制都是正确的。 |
携带地址的算子——tile.load / tile.slice / tile.extract / tile.gather_row——且其读地址引用了区域的 aiv_id |
作者已显式做了 per-lane 定位,例如 data[base + aiv_id * HALF : ...]。仅读偏移参数计入(tile.load 第 1 个、tile.slice 第 2 个、tile.extract 第 1–2 个、tile.gather_row 第 3 个即 src_offset)——出现在 shape、valid_shape 或目的槽位中的 lane 引用并不会移动窗口,因此不予接受。 |
该扫描按区域播种,并按程序顺序做一次前向遍历,因此它只能识别在当前扫描区域内定义的边界结果。从别处到达该区域的 tile.aiv_shard 结果——在兄弟区域中产生,或经由回边上的循环 iter_arg 传入——对它是不可见的,于是消费者会被报成全宽,尽管该值确实是按 lane 的。这两种写法都会被 AivSplitValid 验证器提前 12 个 pass 拒绝(检查 (i) 与 (j),见 99-verifier.md),这正是让上述误报不会出现在作者面前的原因;因此本扫描实际只会遇到同区域内的数据流,也正是它被设计来处理的范围。
tile.gather_row 是其中的 DMA 情形:它是 DPS,因此带有两个偏移,而只有 src_offset
决定两个 lane 是否在做不同的工作——src_offset 由 lane 派生意味着每个 lane 各自拉取属于
自己的散列 GM 行(接受);若只有 dst_offset 由 lane 派生,则两个 lane 会把相同的行取到
同一个全宽累加器的不同槽位(仍会被报告)。
其余归类为 VECTOR 的算子都会被报告。有两点需要注意:
- 生成算子仅对其自身被接受,它不会加入 half-width 数据流。
z = pl.full([FULL, N]); y = pl.add(z, z)仍会在y处被拒绝——全宽生成算子不得为其 消费者背书。 - lane 引用仅在携带地址的算子上被信任。在其他算子上,lane 派生的标量只是一个普通操作数,
并不能说明结果的宽度——因此
pl.set_validshape(full_width_tile, 1, aiv_id * HALF)无法把一个全宽 tile 洗白进区域。
该校验证明的是意图而非范围:偏移按 lane 跨步、但 extent 仍为全宽的 load 会被接受,
此时两个 lane 会读到重叠的窗口——这与本 pass 从不检查 tile.store 的 lane 相关偏移是同一种
信任。
由于区域经由通用的 BeginScope/EndScope 构建且不被提取,它可嵌套在 pl.range /
pl.pipeline 循环或 if 之内;区域路径会递归进入复合语句以下降每个区域,同时保留外围控制流。
Per-lane 散列 gather¶
pl.gather_row 是唯一能按任意运行时偏移读取 GM 的算子,因此它正是把 paged / top-k 行
集合切分到两个 AIV lane 上的手段:每个 lane 在 UB 中组装 tile 的一半,再由 pl.aic_gather
把重组后的 tile 交给 cube。
with pl.at(level=pl.Level.CORE_GROUP, name_hint="sparse_kv",
allow_early_resolve=True): # see the ring note below
for aiv in pl.split_aiv(2, mode=pl.SplitMode.UP_DOWN):
ub = pl.full([64, 512], dtype=pl.BF16, value=0.0) # per-lane HALF extent
for k in pl.range(64):
src = pl.cast(pl.read(idx, [aiv * 64 + k]), pl.INDEX)
ub = pl.gather_row(ub, pool, [k, 0], [src, 0], [1, 512]) # lane-derived src_offset
kv = pl.aic_gather(ub) # V2C -> [128, 512] in Mat
out[0:16, 0:128] = pl.matmul(q, kv, b_trans=True, out_dtype=pl.FP32)
有两条编写规则使其成立:
- 累加器按半 extent 编写。
pl.full是生成算子,因此无论给它什么 extent 都会被接受,而它 本身不会加入 half-width 数据流。gather 是凭其由 lane 派生的src_offset被接受的;该校验 证明的是意图而非范围,所以此处若写成全 extent 的累加器,gather 回来就会变成2 x FULL并在下游产生形状不匹配。 - 留意跨核 ring 的开销。 V2C ring 会在消费侧核的内存中预留
slot_size x slot_num字节 (V2C 为 L1,C2V 为 UB),其中slot_size是消费方弹出的完整 tile——此处为128 x 512 x 2 = 131072——而slot_num默认取 2,相当于在 512 KB 的 L1 中占用 256 KB;每次调用只 push 一次的 kernel 无需超过该深度。pl.cross_core_slot(slot_num=N)可双向调整该值;若调得过大,AllocateMemoryAddr会报告溢出并指出被预留的字节数。
pl.aiv_shard 在这里不能替代半 extent 的 pl.full:它是 C→V 传输,要求操作数位于
Acc(cube 产出),因此无法对 vector lane 自己产生的值做切分。
区域必须不被 scope 包裹¶
区域下降会递归进入 ForStmt / WhileStmt / IfStmt / SeqStmts,但刻意不进入
ScopeStmt:scope 携带提取(outlining)与名字可见性语义,区域局部的折半不应跨越它。因此本
pass 运行时每个区域都必须已不被 scope 包裹——通常由
OutlineIncoreScopes(pass 9)保证,它会把外围的 InCore
scope 提取为独立函数。
pass 8 只对 Opaque / Orchestration 函数提取 scope,因此声明为
pl.FunctionType.InCore 的函数内部的 scope 会原封不动到达本 pass。该 scope 必须是
用户手写的:解析器不会在那里自行添加——InCore 函数中的顶层区域被裸露地发出,这样打印出来
的 *_incore_0 才能重新解析回同样的 IR;而在这类函数中书写区域本身,也会被更早的
AivSplitValid 检查 (h) 拒绝:
@pl.function(type=pl.FunctionType.InCore) # pass 8 会跳过该函数
def f(self, a: pl.Tensor[[128, 128], pl.FP32],
c: pl.Out[pl.Tensor[[128, 128], pl.FP32]]) -> pl.Tensor[[128, 128], pl.FP32]:
with pl.at(level=pl.Level.CORE_GROUP): # 不会被提取
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
base = aiv_id * 64
c = pl.store(pl.exp(pl.load(a, [base, 0], [64, 128])), [base, 0], c)
return c
下降完成后,LowerExplicitRegionFunction 会重新扫描函数体,对隐藏在非 split scope 后的区域抛出
ValueError,并把源位置指向 pl.split_aiv 那一行。修复方式:删掉这层多余的 scope,或改用
普通的 @pl.function / @pl.jit(Opaque)让 pass 8 提取它。
该重新扫描还会拒绝其他任何存活下来的 ScopeStmt,以覆盖对称情形:scope 嵌套在区域体
内部。此时 split 区域包装本身允许保留,故上一条检查会通过——但内层遍历(LowerStmts、
CheckNoCubeTileHalved、AnalyzeSplitBody)会跨过该 scope 而不进入,其中的向量算子会以
全宽被拼接出去,导致两条 AIV lane 都计算整块 tile。该情形从 DSL 不可达(pass 8 会把区域内的
with pl.at(...) 提取为独立函数,检查 (h) 又会拒绝在非提取器产生的 InCore 函数中书写区域),
因此它守护的是绕过 pass 8 的 IR——手工构造的,或反序列化的 .pto。
拆分轴分派¶
SplitMode(int) |
拆分轴 | 折半的向量子区域 |
|---|---|---|
None(0) |
—(无拆分轴) | 不折半——任务并行;tile 保持全宽,由 aiv_id 分派两个 lane |
UpDown(1) |
维 0(高度) | 行 |
LeftRight(2) |
维 1(宽度) | 列 |
SplitDimension(mode) 对 UpDown 返回 0,对 LeftRight 返回 1
(split_axis_utils);对 None 不调用(区域路径先对 None 分支——无轴可推导)。
split 模式与 pto-isa split code¶
SplitMode 描述的是作者选的轴。真正下发到设备的 split 属性含义更窄——它是 pto-isa
的 TileSplitAxis,还要说明两个 lane 的运行时 extent 之间的关系,因为消费侧正是
靠它在 FIFO 槽位中定位 lane 1 的数据段(popVecTileFromGMFiFo):
| Code | pto-isa | Lane 1 的数据段起点 | 要求 |
|---|---|---|---|
| 0 | TILE_NO_SPLIT |
— | 单一读者 |
| 1 / 2 | TILE_UP_DOWN / TILE_LEFT_RIGHT |
e1 * pitch |
e0 == e1 |
| 3 / 4 | TILE_UP_DOWN_ODD / TILE_LEFT_RIGHT_ODD |
(e1 + 1) * pitch |
e0 == e1 + 1 |
这两套词汇分别落在不同的算子上:
tile.aiv_shard/tile.aic_gather携带 MODE(0/1/2)。本 pass 只 盖int(mode)——作者选的轴,仅此而已。tile.tpush_*/tile.tpop_*/system.tfree_*携带 CODE(0..4),由 ExpandMixedKernel 在把边界折叠成传输对时,根据该 mode 与全宽 tile 的 extent 推导(split_axis::ShardSplitCode)。PTO codegen 原样打印为{split = N}。
奇数切分轴有两种来源,折半逻辑对二者完全一致:ComputeHalfDimSize 让两个 lane 都
拿到 ceil 折半作为物理 box,参差不齐的部分由逐 lane 的 valid extent 承载。
| 边界 tile | Lane box | Lane extent | Code |
|---|---|---|---|
[17, 128],全有效 |
[9, 128] |
9 / 8 | 3 |
[16, 128],valid [15, 128] |
[8, 128] |
8 / 7 | 3 |
[16, 128],全有效 |
[8, 128] |
8 / 8 | 1 |
跨边界的部分有效(partially-valid)操作数¶
valid_shape 不足物理 box 的跨界值,正是 kernel 的 ragged 尾部。Cube→Vector FIFO
把传输的列 extent 钉死为物理值,而行 extent 是自由的(推导与 ISA 依据见
PTO codegen)。切分轴上的收窄会让 extent 变成逐
lane 的——lane L 持有 clamp(V - L*half, 0, half)——因此哪种模式 ragged,就决定了
由哪个字段来承载:
下表是 shard(Cube→Vector)的契约;aic_gather 遵循本节末尾的几何规则。
| 收窄的轴 | 模式 | extent 性质 | 载体 | 状态 |
|---|---|---|---|---|
| 行(切分轴) | UpDown |
逐 lane | TPOP valid_row 操作数 |
两个 lane 可摆放时支持(见下) |
| 列(非切分轴) | UpDown |
两 lane 相同、静态 | 整 box 传输 + 静态 pto.treshape |
支持 |
| 行(非切分轴) | LeftRight |
两 lane 相同、静态 | TPOP valid_row 操作数 |
支持 |
| 列(切分轴) | LeftRight |
逐 lane | 无 | 拒绝 |
| 列为运行期值 | 任意 | 两 lane 相同、动态 | 无(treshape 不带操作数) |
拒绝 |
| 行为运行期值 | UpDown |
逐 lane、动态 | 边界算子保留整 box + 第一个消费者的 valid_shape |
支持(见下方说明) |
| 行逐 lane 且列收窄 | UpDown |
两者 | 无(treshape 会同时重写两个轴) |
拒绝 |
ReshapeSplitAxis 只能对切分轴做 ceil 折半,因为 lane 索引不属于 op 的类型函数。
LocalizeExplicitBoundaryValid 在本 pass 修正该猜测——区域自身的
aiv_id = tile.get_subblock_idx() 在作用域内——并把逐 lane extent 传给那些原样透传
valid_shape 的消费者;若消费者会改变逻辑矩形(reduction、slice),则连同 span 一并
报错。AUTO 分支通过 LocalizeShardValidForLane 施加同样的 extent 修正,但不含下面的
store 保护——它的消费者由折半遍历重建,而非本遍历。
- 两个 lane 的 extent 必须可摆放。 切分轴上的 extent 并不是自由字段:pto-isa 根据
被弹出 tile 自身的 valid extent 推导 lane 1 的数据段起点——
TILE_UP_DOWN下是e1,_ODD模式下是e1 + 1。因此可摆放的形态只有三种:e0 == e1(偶数 code)、e0 == e1 + 1(奇数 code),以及e1 == 0(lane 1 不弹出任何数据,其数据段永远不会 被解引用,偶数 code 依然精确)。box 分区对 ragged 边界并不保证这一点——16 行 box 上V = 13会得到 8 与 5——这正是下文均分分区要解决的问题;均分不适用时,ShardSplitCode会报错并给出可行的取值。 - 运行期的切分轴 valid extent 弹出完整 box。 split code 是编译期属性,而两个 lane
需要哪一个取决于它们的运行期 extent:16 行的轴上 valid 12 会让两 lane 变成 8 与 4,
valid 16 则是 8 与 8,没有哪个 code 对两者都正确。因此边界算子干脆不携带逐 lane
extent——
LocalizeExplicitBoundaryValid给它完整 box (split_axis::WithFullSplitAxisValid),并把 lane 的 extent 放到第一个消费者上。 这与偶数 code 恰好配套:生产者搬运的是完整物理 box,lane 1 的数据段就落在 box 的一半处, 偶数 code 正好指向那里,与运行期 extent 无关。已在 a2a3 上对 16 行边界的 1..16 全部 extent 验证。pto-isa 的 pop 确实按被弹出 tile 自身的 extent 放置 lane 1,与其源码读法一致——此前得出相反结论的实测, 其探针两个操作数都是常量,乘积的每一行每一列都相同,落位错误因而无法分辨。 - 手写
tile.tpop_from_aic上的同类行 extent 仍会落位错误。SplitVectorKernel的 折半会把用户声明的valid_shape局部化到 pop 自身,而 pto-isa 正是从这里读取数据段偏移。 仅加宽 pop 并不能修复该路径:其消费者继承作者的声明,随后从完整来源写入部分目标,实测 结果更差。tests/st/runtime/cross_core/test_cross_core_split_parity.py中标记 xfail 的参数记录了受影响的取值范围:UP_DOWN下half < V < box。 - 被收窄的列 extent 在所有路径上一律拒绝。 它根本没有承载者——槽位按生产者的物理列
间距写入,而 pop 依据 tile 自身的
validCol重建读取几何——并且在LeftRight下它就是 切分轴,必须逐 lane 取值。该契约由CheckSplitBoundaryCarriesValid(src/ir/op/tile_ops/cross_core.cpp)统一持有:它既在边界算子的类型推导中运行,也由ShardSplitCode调用,因此手写的tile.tpush_to_aiv/tile.tpop_from_aic同样受其约束。 - 空 lane 的 store 被保护。 ragged extent 覆盖不到的 lane,其 extent 为
0,而 零行TSTORE超出 pto-isa 契约(TSTORE_IMPL断言GetValidRow() > 0)。store 被加上运行时extent > 0判断;tpop与tfree保持无条件——两个 lane 都占用 槽位且都必须释放。 - gather 的限制源自几何而非 DMA。 V2C 的 pop 落到 NZ Mat tile
(
TLoadGm2L1Nd2nz),不读取任何 valid extent。真正限制aic_gather的是摆放: lanel位于偏移l*half,拼合后的数据是[0, v0) ∪ [half, half + v1)——只有两段 相邻时才是矩形。该规则在本 pass 而非类型推导中执行:推导发生在逐 lane extent 存在 之前,只能依据自己的 ceil 猜测来判断。由 localized shard 供给的 gather 会被精确定型 ——两段必然相邻,拼合 extent 即 shard 前的V——只有两个 lane 共享的部分 extent 才会 被拒绝。同样的几何原因决定了 gather 没有奇数形态:GatherSplitCode始终返回偶数 code,并拒绝两个 lane extent 不等的边界(请改为补齐该轴,再在 cube 侧收窄结果)。
把 ragged 边界在两个 lane 之间均分¶
切分会把切分轴分给两个 lane:lane L 拥有 [L*S, L*S + S)。可选的分区有两种,只有
边界 tile 是 ragged 时二者才不同:
| 分区 | S |
[16, …] box、valid 为 13 时的 lane extent |
可摆放? |
|---|---|---|---|
| box(默认) | ceil(box / 2) = 8 |
8 与 5 | 否 |
| valid(均分) | ceil(V / 2) = 7 |
7 与 6 | 是——TILE_UP_DOWN_ODD |
box 分区是通用的:无论 tile 的 valid extent 如何都成立,所以它是默认选择。它做不到的
是让两个 lane 相差不超过 1,而这恰恰是传输的硬要求(见上文可摆放规则)。改为均分
valid 区域后,e0 - e1 ∈ {0, 1} 由构造保证——并且真实工作量也被平均分配,而不是把
余数全丢给 lane 1。
split_axis::ResolveLaneStride 在改写函数体之前扫描一遍,只有在既必要又安全时才返回
均分步长:
- 至少有一个 Cube→Vector 边界,且所有边界在切分轴上的静态
(box, valid)一致,并且ceil(V / 2) < ceil(box / 2)(否则 box 分区本身已经是均分的); - 没有 Vector→Cube 边界(gather 按各自的 extent 位置拼接两个 lane,只有 extent 相等时 两段才相邻);
- 函数体内其他切分轴上的 tile 全都派生自该边界——独立来源的 tile(
tile.load、 生成类算子)跨越整个 box,而均分分区覆盖不到。
其余情况一律保持 box 分区。选定的步长会贯穿整个折半流程——偏移(AdjustOffsets)、
逐 lane 的 valid extent(LocalizeValidDimForSplit)、shard 自身的类型
(LocalizeShardValidForLane)——而物理 box 仍是 ceil(box / 2),因此缓冲区和槽位大小
不变。步长以 lane_stride=S 盖在边界算子上,供
ExpandMixedKernel 用同一分区推导传输 code。在均分的函数体里
若出现迁移切分轴的 tile.reshape,会被拒绝:迁移后的轴带着自己的 half,无法表达按另一
条轴的 valid 做的均分。
# 16 行 box 里有 13 行有效,UP_DOWN。lane 0 取第 0-6 行,lane 1 取第 7-12 行。
popped: pl.Tile[[8, 128], pl.FP32, pl.Mem.Vec,
pl.TileView(valid_shape=[pl.min(pl.max(13, aiv_id * 7) - aiv_id * 7, 7), 128])
] = pl.tile.aiv_shard(qk, split=1, lane_stride=7)
out_store = pl.tile.store(popped, [0 + aiv_id * 7, 0], out_0)
显式 pl.split_aiv 区域永远不做均分:那里的逐 lane 偏移是作者自己写的
(out[aiv_id * HALF : ...]),编译器不能用与区域索引方式不同的分区去切数据。
算法¶
LowerFunction 改写一个混合 InCore 函数:
1. split_dim = SplitDimension(mode);边界算子携带 int(mode)。
2. InjectSubblockIdx(func, is_aiv=true) 在函数体顶部插入
subblock_idx = tile.get_subblock_idx()
(若 'subblock_idx' 已占用则取新名)。
2a. ResolveLaneStride(body, split_dim) 选择分区:对于只有一个 ragged 跨界的
函数体取均分的 ceil(V / 2),否则为 null(box 分区)。它会贯穿下面每一处
偏移与逐 lane extent。
3. LowerStmts 遍历扁平函数体:
边界 tile.move(ClassifyMoveDirection):
CUBE_TO_VECTOR —— 将 move 替换为
tile.aiv_shard(full_cube_tile, split=int(mode)[, lane_stride=S]) -> 半
推导出的半类型已经带有消费侧 lane 内存(Vec):切分推导器让 memory_space
保持为空,由 OpRegistry::Create 用 tile.aiv_shard 的 set_output_memory
声明填充,与显式形式同源。将结果 var 连同其半尺寸种入 tile_vars,并记录
旧->新 var 重绑。cube 源(matmul / Acc 结果)保持全尺寸。
VECTOR_TO_CUBE —— 插入
tile.aic_gather(half_vector_tile, split=int(mode)) -> 全
将源解析到其折半后的 var 使 gather 把 半 -> 全 翻倍,随后保留对折叠后全尺寸
tile 的原 cube 放置 move(命名为 "<dest>_mat",以便 ExpandMixedKernel 的
V->C 边界据此命名其合成的 tpop)。
亲和性门控(ClassifyCallAffinity):
VECTOR 亲和叶子 —— 将单条语句送入
split_axis::ProcessStmts({stmt}, ..., is_aiv=true):与已删除的
SplitVectorKernel 驱动所用的同一机制。沿 split_dim 折半 tile.load /
tile.store / tile.slice / tile.reshape / 计算结果,按 subblock 本地化偏移,
在 tile_vars 中跟踪折半 var。
CUBE 亲和叶子 —— 全尺寸透传,绝不折半。
ForStmt / IfStmt —— 递归进入函数体处理向量内容。
4. CheckNoCubeTileHalved 重新遍历改写后的函数体,断言没有 CUBE 亲和算子消费或产生
tile_vars 中的 tile(亲和性门控绝不能把折半 tile 漏入 cube 操作数)——失败时
INTERNAL_CHECK。
5. transform_utils::Substitute 应用 var_replacements;DeepClone 脱离共享子树。
6. WithSplitAivAttrs 打 split + split_aiv(丢弃任何先前的 split / split_aiv /
dual_aiv_dispatch 条目)。
逐算子向量折半(沿拆分轴折半形状、按 subblock_idx * half 本地化偏移、tile.slice
静态形状参数与结果类型同步折半、rank-1 load 的 reshape 按 lane 切片、拒绝在拆分轴上
归约、保留单元素拆分维、循环 iter_arg/return_var 跟踪)全部由
split_axis::ProcessStmts / ProcessStmt 产生;同样的事实由
tests/ut/ir/transforms/test_lower_auto_vector_split.py 验证。
当拆分维不是单元素维时,自动折半会拒绝根生成器 tile.ci 和 tile.random。它们的
生成值依赖位置,因此仅修改结果类型并不充分;正确改写还需要按 lane 调整 shape 与生成器
状态,而本 pass 不会合成这些调整。请将算子移到会自动折半的拆分区域之外。单元素拆分维与
已有显式边界的半宽区域仍保持不变。
折半结果的可靠性由两个条件共同保证。它们按下面的顺序提问,因为第一个条件凡是能开口 的地方都已经把第二个包住了:
| # | 条件 | 由谁决定 | 拦住的情形 |
|---|---|---|---|
| 1 | 折半后的节点仍能与其自身实参通过类型检查 | 算子自己的 f_deduce_type,无需任何元数据 |
全宽的 tile.add 操作数、落在拆分轴上的 rank-1 bias、row_argmax 形状匹配的 tmp、tile.transpose_view 的轴置换 |
| 2 | 条件 1 看不见的操作数必须 lane 局部、沿拆分轴广播、或被显式声明 | 注册表声明 | tile.sel 的全宽 mask——其结果推导根本不读它 |
条件 1 —— 折半后的节点仍能通过类型检查¶
改写完实参之后,pass 会用尾部 Substitute 将要装入的实参(每个被跟踪的操作数换成其折半
替身)重新推导一次结果,并把推导出的物理 shape 与折半后的结果类型比对。valid shape 不
参与比较:HalveTileShape 会按 lane 做本地化,而类型推导无法复现这一点。若重新推导对折半
后的实参直接报错,同样计为不一致——算子正在拒绝它将要收到的那个节点。
这一问是问算子本身,因此不需要任何逐算子元数据,也不会随着算子新增或契约变化而过时。 最常见的情形正是由它拦下的:
def split_auto(vec: pl.Tile[[256, 128], pl.FP32, pl.Mem.Vec], ...):
vec_h = pl.tile.add(vec, vec) # 被拒绝:`vec` 是全宽的
vec 是全宽的 InCore 参数,区域内无人切分它。仅折半结果会产生
tile.add([256, 128], [256, 128]) -> [128, 128]:该节点声明的 shape 与对其自身实参做类型
推导的结果相矛盾,无法通过 print→parse,并会误导后续所有依据操作数重新推导类型的消费者。
请先得到 per-lane 的一半(在区域内使用 tile.load / tile.slice,或显式写
pl.tile.aiv_shard),或把两条 lane 共享的值放在区域之外。
同一个检查也能正确处理形状启发式两个方向都会判错的情形。操作数与结果从右对齐
(与 BroadcastShapes 一致):
tile.add([256, 128], bias) |
UP_DOWN(维 0) |
LEFT_RIGHT(维 1) |
|---|---|---|
bias: [128] |
接受——在拆分轴上是广播 | 拒绝——BroadcastShapes([256, 64], [128]) 不成立 |
bias: [1, 128] |
接受——拆分轴上是单元素维 | 拒绝——其维 1 是全宽 |
bias: [256, 128] |
拒绝 | 拒绝 |
它同样覆盖轴置换类算子,而后者是任何右对齐规则都描述不了的:
def split_auto(data: pl.Tile[[1, 256], pl.FP32, pl.Mem.Vec], ...):
viewed = pl.tile.transpose_view(data) # 被拒绝:[128, 1] 与推导出的 [256, 1] 矛盾
tile.transpose_view 交换末两维,于是折半后的 [128, 1] 结果与其操作数自身推导出的
[256, 1] 相矛盾。与 tile.transpose 不同,它是纯视图,不在 FindTransposeSplitHazard
的覆盖范围内(后者只匹配 tile.transpose)。
描述 shape 的算子(tile.full、tile.create、tile.slice、tile.reshape、
tile.reinterpret_view)在该检查之前就改写了自己的操作数,且全宽源在其语义中有明确含义。
tile.reshape 与 tile.reinterpret_view 的这一点来自一次更早的改写:对未被跟踪的源,
其无偏移视图会先以全宽发出,再跟一个 per-subblock 的 tile.slice,使每条 lane 读到自己
的那一半,而不是两条都读第一半。该 slice 只能由静态的半宽 extent 物化,因此全宽源上
动态的拆分 extent 会被直接拒绝,而不是落到普通的结果折半上。
条件 2 —— 类型推导看不见的操作数¶
推导器若从不读某个操作数的 extent,那么它"接受"这件事本身什么也不能说明。pass 会逐调用 就地判定这一点,同样无需元数据:把该操作数在拆分轴上加倍再推导一次。若推导报错或结果 发生变化,说明推导器盯着这个操作数,条件 1 已经替它做过判断;若答案不变,说明推导器对 它视而不见,此时只有声明才能回答全宽是否符合契约。
tile.sel 仅从 BroadcastShapes(lhs, rhs) 推导结果,对 mask 只检查它是不是 TileType;
于是"全宽 mask + 折半 lhs/rhs"能干干净净地重新推导通过——可两条 lane 都会去读 mask 的第
0..127 行。tile.sels 是同一类陷阱:它只校验 mask 覆盖 src 的有效行,而超宽的 mask
恰好满足覆盖。两者都只声明了自己的 tmp 位置,因此它们的 mask 会被拒绝。无人分类过的操作
数一律按 lane 数据对待——这是安全的默认值。
把这道关卡限定在"推导器看不见"的操作数上,正是它不越界的关键。否则每个带有合法全宽操作数
的算子都得写一条注册表声明,只为压掉一个本就不该运行的检查——tile.rsqrt 就是这么给一个
"推导器要求与输入逐维相等"的 scratch 加上声明的。
tests/ut/ir/operators/test_lane_invariant_arg_coverage.py 会度量哪些操作数处于盲区,
对推导器已经能决定、因而永远走不到的声明直接判失败,并钉住盲区集合,使新增算子或放松的
推导器必须被分类,而不是悄悄漏过。
"处于盲区"通常是症状,而推导器往往才是解药。 一个操作数落进盲区,是因为推导器从不读
它的形状——这常常只是契约校验不足,而非该维度真的自由。此时正确的做法是把算子本就写明的
关系补成检查,而不是把它声明为 lane 无关:补上检查后,任何调用方传入形状不符的操作数都会
被拒绝,而拆分 pass 也顺带拿到了一个"被钉住"的答案。tile.scatter_update 的 index 与
src 之所以在盲区,是因为推导只是把输入类型原样抄回;把 tensor 路径的下降早已断言的关系
(src 行数 == b * s、宽度相等)补进推导后,两个操作数都被钉住。若改为写声明,则既压掉
了拆分检查,又把形状缺陷原封不动留在原地。
是否需要声明的判据是位置对应关系:输出的第 i 个元素是否读取操作数的第 i 个元素?
有两类操作数的答案是否。
硬件 scratch。 全宽工作区符合契约,两条 lane 各自声明整块即可:
scratch = pl.tile.create([256, 128], dtype=pl.FP32, target_memory=pl.Mem.Vec)
sums = pl.tile.row_sum(v, scratch) # 接受:[128, 128] x [256, 128] -> [128, 1]
tile.row_sum 把 tmp_tile 定义为至少与输入一样大的 scratch,DeduceTileRowReductionType
从不读它,并通过 set_lane_invariant_arg(1) 声明。tile.create 注册为
CoreAffinity::SHARED,亲和性门控因此刻意让它以全宽通过、也从不跟踪它——
仅凭 extent 匹配会把整个 rms-norm 归约形状都拒掉。
它的同族 tile.row_argmax 需要与源完全同形的 tmp,DeduceTileRowReductionType 用
require_exact_tmp_shape 强制了这一点。因此它不带任何声明:条件 1 自己就会拒绝全宽
tmp,而一条声明反而会宣称一个算子并未给出的契约。
由绝对索引寻址的操作数。
indices = pl.tile.load(idx, [0, 0], [256, 128], target_memory=pl.Mem.Vec)
picked = pl.tile.gather(src, indices, tmp) # 接受:src 保持 [256, 128]
tile.gather 计算 out[i] = src[indices[i]],结果 shape 取自 indices。索引值是相对
src 的绝对偏移,因此折半 indices 本身就让每条 lane 拿到输出的各自一半,而两条
lane 读同一张完整表——这正是正确的降级结果,而不是漏掉的切分。tile.gatherb 是同一形态的
字节偏移版本。tile.scatter 与 tile.scatter_update 是其写侧对偶,用绝对索引寻址目的地。
声明只是"允许"操作数保持全宽,并不能"保证"它全宽。 没有任何机制阻止该操作数自己的 生产者被折半,而各类声明对此的含义不同——这正是注册表需要记录是哪一类的原因:
| 类别 | 全宽 | 生产者被折半 | 例子 |
|---|---|---|---|
Scratch |
允许 | 无害——输入折半时 scratch 一起折半仍然成立 | row_sum 的 tmp_tile |
IndexAddressedSource |
允许 | 拒绝——索引仍是绝对值,lane 1 读错半区 | gather 的 src |
AbsoluteIndexedDestination |
允许 | 拒绝——同一缺陷的写侧 | scatter 的 dst |
tile.scatter 执行 dst.flat[indexes[i, j]] = src[i, j],索引是相对整块目的地的扁平
偏移,因此按列写入所编码的 i * dst_cols + c 一旦行 stride 折半就不再成立。
src = pl.tile.load(tbl, [0, 0], [256, 128], target_memory=pl.Mem.Vec)
picked = pl.tile.gather(src, indices, tmp) # 被拒绝:表被折半了
这是常规路径而非角落:张量层的 pl.gather 会降级为"tile.load 表 + tile.gather",
所以在拆分区域内做 gather 的作者默认就会撞上。上面两个条件都看不见它——操作数被跟踪,
条件 2 会跳过;而 tile.gather 仅从 indices 推导结果,条件 1 也满足。这正是绝对索引这
两类即便推导器确实读了该操作数也仍要保留声明的原因:这个性质在两个方向上都无法用类型
表达。请把表放在拆分区域之外产生再传入;自动保持这类生产者全宽的重写尚未实现。
绝对索引检查跑在折半路径之前而非与之并列,因为即使什么都没被折半,尾部的 Substitute
依然会改写实参:
- 结果为单元素维时会提前返回——
indices为[1, N]的 gather 结果也是[1, N], 结果原封不动,而其下的表仍被折半; - 结果为 tuple 时走自己的路径(见下):通用块只跟随一个
result_split_dim,而 tuple 的拆分轴是逐元素的。
tuple 结果——每个元素各有一条拆分轴¶
tile.gather_compare 返回 Tuple[dst[rows, out_cols], cdst[1, rows]]。在按行拆分时,
这两个元素并不沿同一条轴移动:dst 在 dim 0 折半,而 cdst——按行计数连续排布的单行
——在 dim 1 折半。没有单一的结果轴可循,因此通用路径无法表达。
该映射是推导出来的,而非声明的:用折半后节点将要携带的实参重新做一次类型推导,读出每个
元素移动的是哪条轴。这样新增 tuple 返回算子无需任何注册,映射也不会像逐算子元数据那样过期
(见 gh#2612 的讨论)。重新推导只提供轴;折半本身仍由 HalveTileShape 完成,因此奇数
extent 的逐 lane 局部化与单 TileType 路径完全一致。
每个元素都必须恰好在一条轴上折半。元素原样返回才是危险情形而非无害情形——算子用每个
lane 只拥有一半的操作数产生了全宽输出,于是两个 lane 都没有完整答案,而形状看起来仍然正确。
这也正是 tile.gather_compare 的 LEFT_RIGHT 拆分被正确拒绝的原因:它的两个输出都由源的
行数定尺寸,折半列数不会让任何一个移动。
x = tup[i] 投影由 split_axis::RetypeTupleProjection 依据折半后的 tuple 重新定型,并记录
每个元素自己的轴,使后续 tile.store 偏移正确的维度。两条下降路径都要调用它:AUTO 路径
的亲和性门只把叶子 call 送进 ProcessStmts,投影若落到其"原样透传"兜底分支,就会在已折半
的 tuple 之上保留全宽的声明类型。
投影也可以完全不绑定:pl.tile.store(pair[0], [0, 0], out) 直接内联传入
TupleGetItemExpr,没有任何环节会把它提取成变量。因此任何在 tile 操作数上只匹配 Var 的
代码都会漏掉它,而该操作数之后仍会被替换,于是在逐 lane 数据之上留下全宽的声明类型。
split_axis::OperandSplitInfo 是"这个操作数是否被拆分、沿哪条轴"的唯一答案,绑定的 Var
与内联投影一视同仁;BuildHalvedCallArgs 则在折半后的 tuple 之上重建投影,使类型一致性探测
也看到逐 lane 的操作数。所有消费者都走它们——通用路径的被跟踪输入扫描、LocalizeStoreOffset,
以及 GetFirstTileArgMemory(它现在读操作数的类型,向量算子不会再被误判成 SHARED 并复制
到两条 lane)。
所有"这个操作数是否被拆分"的提问都走 OperandSplitInfo。这个问题的消费者远不止折半本身,
而每一处对内联投影答"否"的后果各不相同:
| 提问方 | 答错的代价 |
|---|---|
| 通用路径的被跟踪输入扫描 | 结果在逐 lane 操作数之上保留全宽类型 |
LocalizeStoreOffset |
两条 lane 写到同一批行 |
GetFirstTileArgMemory |
向量算子被判为 SHARED,复制到两条 lane |
| 绝对索引闸门 | tile.gather 的表已减半而索引仍绝对——无任何诊断,因为 gather 的结果尺寸取自 indices |
tile.reshape / tile.reinterpret_view |
在一个即将被折半的操作数上生成全宽视图 + 逐 lane 切片 |
tile.slice 偏移 |
在已是 lane 局部的偏移上再加 + subblock_idx * half,lane 1 读越界 |
| V→C 边界 | 合法的 tile.move(pair[0], target_memory=Mat) 被当作全宽拒绝 |
RepairIterArgs(循环初值) |
初值已折半,而携带值、出口及循环之后全部停留在全宽 |
YieldedTileInfo(回边 / 合并) |
两个方向都错:全宽携带值被喂 pl.yield_(pair[1]) 却放行,合法折半的反被拒绝 |
FindFullWidthOperand(条件 2) |
已分区的投影被报成全宽操作数——误拒 |
pass 里仅剩的 AsVarLike 查找只有两类:OperandSplitInfo / ReplacedOperand 自身的实现,
以及刻意按变量身份识别整个 tuple 的两处(YieldedHalvedTupleType、RetypeTupleProjection)。
除此之外任何询问操作数拆分状态的代码都应走该 helper——这才是阻止这类问题按调用点逐个复发的办法。
边界处的 tile.aic_gather 也改用 ReplacedOperand 构造——它在折半后的 tuple 之上重建投影,
使 gather 无论哪种写法都是 HALF → FULL 的加倍。
tuple 还会穿过分支合并与循环携带,两者都不从 tile_vars 读取拆分信息——tuple 变量从不
在其中。两者改为直接采用折半后的类型,其元素本就带着逐元素拆分,没有单一的轴可记录:
| 形态 | 修复者 | 一致性规则 |
|---|---|---|
if 合并 |
RepairIfReturnVars |
两个分支必须产出相同的折半 tuple |
| 循环携带 + 出口 | RepairIterArgs / RepairReturnVars |
回边 Yield 必须与携带值一致 |
DSL 对两者都无法写注解,但 ConvertToSSA 会为"在分支里被重新赋值的 tuple"合成正是这种 phi,
而 pl.range(..., init_values=(tup,)) 可以直接携带一个 tuple。
这里有两种不同的失败都会以拒绝告终,诊断信息把它们分开。一是算子直接拒绝折半后的
实参:可能是某条约束在折半后不再成立(tile.tquant_mx 要求 M % 16 == 0,而 per-lane 的
M 可能破坏它),也可能是 workspace 按完整源尺寸分配——LowerCompositeOps 分解之后,
tile.tquant_mx_raw 要求其 [1, groups] scratch 与由 src 推出的数量精确相等,而那个
单元素 dim 0 使通用路径不会折半它。重新划分这类 workspace 尚未实现:它的尺寸算在推导
函数内部,且其分配语句早已按全宽发射。两种情况下诊断都直接引用算子自己的报错,而不去猜测。
二是算子接受了折半实参但某个元素没有移动,即上面那种"内容错、形状对"的情形。
条件 2(盲操作数兜底)不在这条路径上运行——它是围绕单一结果轴写的。"每个元素都必须移动"这条
规则覆盖了常见情形:主逐 lane 操作数若保持全宽,元素就不再移动,从而被拒绝。它无法覆盖
任何元素形状都不依赖的次要逐 lane 操作数;当前已注册的 tuple 返回算子都没有这类操作数
(其余 tile 操作数都是已声明的工作区),而新增算子会出现在
test_lane_invariant_arg_coverage.py 的盲操作数清单里,迫使这个问题在合入前被回答。
只在部分 arity 下才是 scratch 的位置无法声明:tile.mrgsort_format2 的 tmp_or_src2
在 3/4 路归并里是第三个已排序输入、在 2 路里才是工作区,而 arity 由位置实参个数决定。
但实际上并不需要声明——推导函数总是用持有 tmp 的那个位置(永远是最后一个)来给结果定尺寸,
因此类型一致性在任何 arity 下都能判定该位置,其余位置则是必须被分片的真实排序输入。
循环携带值、分支归并与被丢弃的轴¶
有四处是在折半周围改写状态而非折半本身,它们必须遵循同样的轴映射与跟踪规则, 否则上面的条件会误判:
-
作为返回表达式出现的 store。
LocalizeStoreOffset把被跟踪 tile 的tile.store移到本 lane 所属的那半目的地,而语句形态决定了它是否会被触及。return pl.tile.store(v, [0, 0], out)是再普通不过的 DSL——没有任何环节把它归一化成 赋值——但它既不是AssignStmt也不是EvalStmt,AUTO 路径的亲和性门也不会转发它 (它自身不携带叶子 call)。尾部的Substitute仍然会换入折半后的 tile,于是两个 lane 用不同的数据写了相同的行。LocalizeReturnStores让两条路径都覆盖这种返回形态; "先把 store 绑定到一个名字"从来就不该是正确性的前提。 -
循环携带值。 一个携带值有三条边,三者必须一致。
iter_arg继承其 init 的跟踪信息, 因此 init 被折半时携带值也变为 lane 局部;循环出口的return_var再继承之,使后续tile.store拿到 per-lane 偏移。显式路径与 AUTO 亲和性门控路径共用RepairIterArgs/RepairReturnVars——AUTO 走的是自己的遍历,够不到显式路径的ForStmt分支,缺了这一步 就会把合法的 tile 累加器的携带值当成全宽操作数报错。 - 循环回边。 第三条边是循环体 yield 回携带值的那个值,由
ValidateCarryBackedge负责 校验:yield 的值当且仅当携带值是 lane 局部时才可以是 lane 局部。被折半的携带值收到一个 全宽值,就会声明一个与第二轮起流入它的值相矛盾的 per-lane 类型——和未分片操作数是同一个 缺陷,但所有操作数检查都看不见它,因为那些检查读的是Call的实参,而这里是Yield。 该校验是对称的,因此把 lane 局部的值 yield 进全宽携带值同样会被拒绝。它只校验、不修复: yield 的值若已被跟踪,尾部的Substitute本就会换上折半替身;若未被跟踪,则根本不存在 可替换的折半版本。
未绑定的回边——pl.yield_(pl.tile.add(acc, acc))——会经由另一条路径走到同一个拒绝点,
并给出专门的信息。传给 pl.yield_ 的表达式不会被任何环节提升出来,于是这个调用就内联留在
Yield 里;而本 pass 折半的是语句,永远够不到它。上面那条"两端不一致"的措辞会让作者去
审查一条本来没问题的携带边,真正可执行的指引是先把值绑定到名字上。
- 分支归并值。 IfStmt 的归并变量(return_vars_,一个 DefField)是否 lane 局部,
取决于两个分支 yield 的值是否 lane 局部。保持声明的全宽会同时与两个 Yield 矛盾,
而且因为它从不被登记进 tile_vars,后续 tile.store 拿不到 lane 偏移——两条 AIV
lane 都从输出第 0 行开始写,是覆盖写而不是各写一半。RepairIfReturnVars 依据降级后
的分支重新定型,两条路径都要调用它,原因和循环携带值需要两个调用点相同。两个分支必须
一致:一个 yield 被折半的值、另一个 yield 全宽值时不存在唯一的归并类型,此时直接拒绝,
而不是任选一边。这里两个分支必然都存在:只要定义了 return_vars_,SSA 就要求有 else,
而源码层没有 else 的 phi 会由 ConvertToSSA 合成 else Yield。
- 带 drop_dims 的 tile.slice。 这个算子有两个轴空间,而本 pass 在两个方向上
都要穿越它们。shape / offset / valid_shape 按丢弃前的 rank 索引;结果则是这些
轴去掉 drop_dims 之后、再通过在前面补单元素轴钳回 2D。因此
slice([1, 256, 128], drop_dims=[0]) -> [256, 128] 的结果维 0 对应源的维 1,反过来源的
维 0 也可能落到结果的维 1 上。
| 方向 | 用途 |
|---|---|
| 结果轴 → 丢弃前轴 | 改写按丢弃前 rank 索引的 shape / offset / valid_shape |
| 丢弃前轴 → 结果轴 | 把被跟踪源的拆分轴换算成索引结果类型的那个轴 |
第二个方向正是被跟踪的源所需要的:拆分轴记在操作数上,而折半路径索引的是结果。 假定两者一致,会让源的轴 0 读到结果里那个合成出来的单元素轴,走进单元素提前返回,于是 slice 原样放行、而它的源已经在下面被折半了——8 行的源上开出 16 行窗口。两个映射都建立在 同一份存活轴列表之上,因此在存活轴上互为逆映射。2D 钳位补出的那些合成单元素轴没有 丢弃前的对应轴:结果 → 丢弃前方向对它们原样返回,而丢弃前 → 结果方向根本不会产生它们, 所以在这些轴上两者不构成往返。
亲和性门控¶
仅折半向量工作,cube 工作保持全尺寸。亲和性由
core_affinity::ClassifyCallAffinity(按内存空间)决定:产生或消费 Vec tile 的算子
为 VECTOR;matmul 操作数与 Acc/Mat cube 结果为 CUBE。C→V 边界 tile.aiv_shard
是接缝:全尺寸 cube tile 是其输入,半尺寸向量 tile 是其输出。CheckNoCubeTileHalved
是兜底——若 cube 操作数被缩小则触发。
示例 —— cube→vector 边界,向量区域折半(UpDown)¶
混合核:cube tile(Mat)跨入 Vec,向量 add 在其上运行,结果被存储。
之前(InferTileMemorySpace 之后的混合 InCore):
@pl.function(type=pl.FunctionType.InCore, attrs={"split": pl.SplitMode.UP_DOWN})
def split_auto(qk: pl.Tile[[128, 128], pl.FP32, pl.Mem.Mat],
out_0: pl.Out[pl.Tensor[[128, 128], pl.FP32]]):
popped: pl.Tile[[128, 128], pl.FP32, pl.Mem.Vec] = pl.tile.move(qk, target_memory=pl.Mem.Vec)
y: pl.Tile[[128, 128], pl.FP32, pl.Mem.Vec] = pl.add(popped, popped)
return pl.store(y, [0, 0], out_0)
之后:
@pl.function(type=pl.FunctionType.InCore,
attrs={"split": pl.SplitMode.UP_DOWN, "split_aiv": True})
def split_auto(qk, out_0):
subblock_idx: pl.Scalar[pl.INDEX] = pl.tile.get_subblock_idx()
popped: pl.Tile[[64, 128], pl.FP32, pl.Mem.Vec] = pl.tile.aiv_shard(qk, split=1) # C->V, 半
y: pl.Tile[[64, 128], pl.FP32, pl.Mem.Vec] = pl.add(popped, popped)
return pl.store(y, [0 + subblock_idx * 64, 0], out_0)
cube 操作数 qk 保持 [128, 128];向量子区域折半为 [64, 128],store 偏移按
subblock 本地化。
示例 —— vector→cube 边界保持全尺寸(UpDown)¶
V→C tile.move 变为 tile.aic_gather;对折叠后 tile 的 cube 放置 move 保持全尺寸
[128, 128] Mat——cube 侧绝不会看到折半 tile:
# `v` 是 affinity gate 产出的每 lane 折半 tile,例如 [64, 128]。
gathered_mat: pl.Tile[[128, 128], pl.FP32, pl.Mem.Mat] = pl.tile.aic_gather(v, split=1)
gathered: pl.Tile[[128, 128], pl.FP32, pl.Mem.Mat] = pl.tile.move(gathered_mat,
target_memory=pl.Mem.Mat)
操作数必须是每 lane 的折半 tile。 tile.aic_gather 声明为 HALF → FULL,因此
gather 把 [64, 128] 加倍为 [128, 128],恰好等于 cube 放置 move 所保留的全尺寸
结果类型。这种一致性是前置条件而非保证:向量值可能以未折半的形式到达边界——
例如直接使用的 Vec 参数,或 split 维为 1 而被 affinity gate 特意保留的 tile。对
这类操作数做加倍会得到 [256, 128] 的 gather,而其后的 move 仍是 [128, 128],
这与 tile.move 的保形契约矛盾,且产生的 IR 无法通过 print→parse 往返。既然没有
折半就不存在正确的 gather,本 pass 会以可操作的 ValueError 拒绝它:
LowerAutoVectorSplit: the V->C boundary tile.move here carries a full-width
vector operand 'vec'. tile.aic_gather reassembles the two AIV lanes' per-lane
halves into the full tile the cube expects, so its operand must be a value the
split halving produced (a tile.load / tile.slice / elementwise result inside the
vector sub-region). An un-halved value has no half to gather — either derive the
per-lane half first (load or slice the value inside the split function) and move
that to the cube side, or, if the split axis is a singleton that cannot be
halved, keep the value on the vector side.
gather 沿操作数自身的 split 轴,而非函数的 split 轴。 tile.reshape 可能迁移
split 轴——rms_norm 的 [N, 1] ↔ [1, N] 列 reshape 会把它从 dim 0 移到 dim 1——
而 TileInfo::split_dim 记录了它最终所在的位置。gather 使用该维度对应的
split 编码发射(dim 0 → 1,dim 1 → 2),因此在 UpDown 函数中,lane 本地
的 [1, 8] 操作数经由 split=2 gather 为 [1, 16],与 move 一致;若改为对函数
轴加倍则会得到 [2, 8]。若追踪到的 split 维不在 {0, 1} 内,则无法表示为 2D
gather 的 split 属性,会被拒绝。
gather 结果是 Mat 而非 Vec:边界算子声明的类型指的是消费侧 lane 的空间,
而 AIC 会把 V→C 传输 pop 进 L1。(Vec 指的是生产侧 lane,与镜像算子
tile.aiv_shard 相矛盾——后者为其 cube 产出的操作数声明向量侧的 Vec。)随后的
cube 放置 move 才把 tile 放到最终的操作数空间——matmul 操作数为 Mat → Left;
此处的 Mat → Mat 是空操作,仅因本 pass 保留了作者原有的 move 而存在。
实现¶
头文件:include/pypto/ir/transforms/passes.h
实现:src/ir/transforms/lower_auto_vector_split_pass.cpp
LowerFunction/LowerStmts—— 边界改写 + 亲和性门控折半。MakeReshapeOpCall—— 构造tile.aiv_shard/tile.aic_gather调用。CheckNoCubeTileHalved—— cube 操作数完整性兜底。WithSplitAivAttrs—— 打split+split_aiv。
共享机制:src/ir/transforms/utils/split_axis_utils.cpp
(ProcessStmts、InjectSubblockIdx、SplitDimension、IsReduceOnSplitAxis)
—— 逐算子向量折半,与 SplitVectorKernel 的独立拆分分支
(ProcessStandaloneSplitFunction)以及 AivSplitValid 校验器
(SplitDimension / IsReduceOnSplitAxis)共享。
Python 绑定:python/bindings/modules/passes.cpp
测试:tests/ut/ir/transforms/test_lower_auto_vector_split.py 以及
tests/st/codegen/torch/test_torch_codegen_cross_core.py 中端到端 pl.split
golden 场景(test_lower_auto_vector_split_golden)。
相关¶
ResolveBackendOpLayouts—— 紧邻其前运行。ExpandMixedKernel—— 紧邻其后运行;把tile.aiv_shard/tile.aic_gather折叠为带拆分标记的tpush/tpop。SplitVectorKernel—— 下游;仅为本 pass 产生的split_aiv函数打属性,外加无拆分 dual-AIV 路径。