XLA layout_assignment.cc 对 PTOAS VMI vreg layout 分配的借鉴意义

[Zhendong404]

调研报告

主题
评估 XLA layout_assignment.cc 对 PTOAS VMI vreg layout 分配的借鉴意义。
参考源码：XLA layout_assignment.cc
https://sources.debian.org/src/chromium/150.0.7871.100-1/third_party/tflite/src/third_party/xla/xla/service/layout_assignment.cc

一、结论

这套算法有明显借鉴价值，但不适合直接照搬。
最值得借的是它的“全局约束求解思路”，不值得直接搬的是它围绕 LogicalBuffer / alias / copy 的实现模型。

对我们来说，更合适的方向不是做一个 XLA 式 buffer-layout pass，而是在现有 VMI solver 上继续演进：

• 给 layout 请求引入优先级/强弱约束
• 支持多候选关系下的选择
• 用代价模型减少 ensure_layout / ensure_mask_layout
• 保留当前 SSA value 为中心的设计，不转成 buffer 中心模型

二、XLA 这类 LayoutAssignment 的核心思想

从问题建模上看，XLA 处理的是“buffer 的物理 layout 约束统一”：

• 某些 op 或 backend 对输入/输出 layout 有硬约束
• 某些值通过 alias / tuple / call / while / conditional 跨 computation 关联
• 当 producer layout 和 consumer layout 不一致时，可以插 copy 修复

它的关键价值不在某个局部规则，而在于：

• 它是全图传播，不是单 op 就地决策
• 它区分强约束和偏好约束
• 它允许多轮传播与回收敛
• 它把“无法统一”的地方留给显式物化去修补

这套思路和我们 VMI 当前的目标是相通的。

三、我们当前 VMI layout assignment 的现状

现有实现已经有一套相当完整的全局框架。

在 VMILayoutAssignment.cpp 里：

• 用 DataLayoutSeedPhase 把 seed 分成 GroupLoad / Reduce / Cast / Store 等阶段
• 为 value 维护 naturalLayout 和 preferredLayout
• 先收集约束，再按 phase 逐轮请求 propagator 传播

在 VMILayoutPropagation.h 和 VMILayoutPropagation.cpp 里：

• 用 op-specific Transfer 描述 layout 关系
• 用 request(value, layout) / request(operand, layout) 做传播
• 用 VMIValueLayoutAssignment 记录 primary layout 和 use conflicts
• 最后通过 ensure_layout / ensure_mask_layout 做物化修复

在 VMIMaskGranularityAssignment.cpp 里，mask granularity 也采用了同样模式：

• 主值保留一个 primary 决策
• use-site 冲突通过 ensure_mask_granularity 修补

所以从架构上说，我们已经走在和 XLA 同方向的路上，只是现在还偏“规则驱动 + phase 顺序驱动”。

四、和 XLA 相比，我们当前实现的优势

1. 更贴合 VMI 语义

VMI 的 layout 更像“寄存器视图/解释方式”，不是 buffer 的存储布局。
因此我们用 ensure_layout 做显式桥接，比 XLA 的 copy 模型更自然。

2. SSA value 为中心，模型更轻

我们当前直接围绕 SSA value、operand、result 建图，不需要引入额外 buffer 抽象。
对 VMI 这种 IR 层寄存器 layout 问题，这样更直接、实现成本也更低。

3. 已经覆盖了跨 region / call 的大部分结构性约束

VMILayoutAssignment.cpp 到 VMILayoutAssignment.cpp 已经把：

• if
• while
• for
• switch
• branch
• return
• call
  都纳入了等价传播或边界物化范围。

这一点已经很接近“全局求解器”的基本形态了。

五、我们当前实现和 XLA 思路相比的主要差距

1. 缺少显式优先级模型

当前 request(value, layout) 一旦给 value 设了 primary layout，后续不同请求基本不会替换它，见 VMILayoutPropagation.cpp。
这意味着结果很依赖：

• seed 进入顺序
• phase 排序
• 哪个 op 先传播到该 value

XLA 这类算法的一个关键点就是：约束不是“先到先得”，而是“有强弱之分”。

2. 多候选关系时，当前传播器基本跳过选择

在 VMILayoutPropagation.cpp：

• 如果一个 transfer query 返回多个 relation
• 当前 propagate() 只在 relations->size() == 1 时继续传播

这会直接放弃很多“本可通过全局选择解决”的场景。
而 XLA 的价值之一，正是它能在多种合法 layout 之间做全局一致化决策。

3. 缺少代价模型

当前冲突处理是可行性优先，不是优化优先。
例如：

• 哪个 primary layout 会导致更少 ensure_layout
• 哪个方案更贴近 load/store/native lowering
• 哪个方案能减少 layout 往返

这些还没有显式进入求解过程。

4. 默认 contiguous 过强

在 VMILayoutAssignment.cpp 之后，未定值最终会落到 contiguous。
这对保守正确性很好，但对某些 group/deinterleaved/lane_stride 友好的链路，可能会过早收缩到通用布局。

六、哪些点最值得借鉴

1. 把 request 升级为“带优先级的请求”

建议把当前：

• naturalLayout
• preferredLayout
• phase 顺序

进一步统一成明确的优先级体系，例如：

• Mandatory
• StrongPreference
• WeakPreference
• Default

这样 GroupLoad 的硬支持布局、Reduce 的自然布局、Store 的偏好布局，就能统一比较，而不是靠 phase 顺序隐式决定。

2. 支持多 relation 下的择优

当前 transfer 返回多 relation 时基本不传播。
建议增加一个轻量选择策略，例如优先：

• 不新增物化
• 与已分配邻居一致
• 更接近 lowering 支持表
• 更接近 producer natural layout

这一步会显著提高 solver 质量。

3. 引入局部代价模型

即使不做完整全局最优化，也可以做启发式打分：

• ensure_layout 次数
• ensure_mask_layout 次数
• ensure_mask_granularity 次数
• 破坏 group/deinterleaved 链的次数
• 退化到 contiguous 的次数

这会让分配结果从“可用”走向“更优”。

4. 保持“主布局 + use-site 修复”的总策略

这一点不建议改。
它和 VMIMaskGranularityAssignment.cpp 的模式一致，也非常符合 VMI IR 的语义。

七、不建议照搬的部分

1. 不建议引入 XLA 式 buffer/alias 主模型

VMI vreg layout 的核心不是 buffer physical layout。
如果强行引入 LogicalBuffer 类抽象，复杂度会上升很多，但未必带来对等收益。

2. 不建议把 copy 作为核心修复手段

VMI 里的 ensure_layout 已经是更准确的 IR 语义。
继续沿这个方向推进就好。

3. 不建议为了“全局最优”过度复杂化 pass

VMI 这里更适合：

• 规则表
• 启发式
• 多轮传播
• 局部代价优化

而不是上整数规划或重型 buffer analysis。

八、建议的演进路线

阶段 1：低风险增强

• 给 request 增加 priority
• 给 VMIValueLayoutAssignment 增加来源和强度信息
• 在冲突时允许高优先级请求覆盖低优先级 primary

阶段 2：提升解质量

• 处理 relations.size() > 1 的情形
• 增加一个简单 score function
• 优先减少 ensure_layout 和降级到 contiguous 的次数

阶段 3：面向 lowering 联动

• 把 VMIToVPTO / VMILayoutFold 的收益纳入代价模型
• 让 solver 更偏好能被后续 pass 折叠或直接 lower 的布局

九、最终判断

如果问题是：

“XLA 这套 layout_assignment 算法是否对 VMI vreg layout 分配有借鉴意义？”

答案是：

• 有，而且借鉴价值不小
• 但应该借它的求解思想，不该照搬它的实现模型

一句话概括就是：
我们现在的方向是对的，下一步最值得补的是“优先级 + 多候选选择 + 代价模型”。

如果你愿意，我下一步可以把这份报告继续落成一个更工程化的设计稿，直接给出 VMILayoutPropagation.h 和 VMILayoutPropagation.cpp 的具体改造方案。