SelectOp::ShapeProp 实现说明

1. 文档目的

本文详细说明 DVM 中 SelectOp::ShapeProp() 的实现、调用链、广播规则、动态 Shape 处理方式，以及它对 GraphKernel 符号式切分的影响。

本函数对应的代码修改包含：

1. 在 src/ops.cc 的算子元信息表中，为 Select 注册 SelectOp::ShapeProp；
2. 在 src/ops.cc 中实现 SelectOp::ShapeProp()；
3. 在 src/ops.h 中声明该函数。

该修改不改变 Select 的计算语义，也不改变 V_SEL_* 字节码的生成方式，主要用于补齐符号式切分阶段的输出 Shape 推导。

2. Select 算子的基本含义

DVM 中的 Select 可以表示为：

Select(condition, lhs, rhs)

含义是：

condition 为真 → 选择 lhs
condition 为假 → 选择 rhs

Select 有三个输入：

DVM 成员	含义
xhs_	condition，条件输入
lhs_	真分支数据
rhs_	假分支数据
shape_	Select 输出的逻辑 Shape
shape_ref_	指向 shape_ 的 Shape 引用

Select 的三个输入都可能发生广播，因此输出 Shape 需要同时考虑三个输入，而不是只看 lhs_ 或 rhs_。

例如：

condition: [B, 1, N]
lhs:       [1, H, N]
rhs:       [B, H, 1]

output:    [B, H, N]

3. 代码位置和调用关系

3.1 相关代码位置

代码	位置	作用
SelectOp::ShapeProp 声明	src/ops.h	声明 Shape 推导回调
ShapeProp 注册	src/ops.cc 的 GenObjectMeta()	将 Select 和回调函数绑定
ShapeProp 实现	src/ops.cc	计算 Select 输出 Shape
通用 ShapeProp 分发	src/ops.h 的 NDObject::ShapeProp()	根据算子类型调用回调
符号式切分入口	src/gkernel.cc 的 GraphSpliter::Run()	遍历图并重新推导 Shape

3.2 调用流程

NewKernel("split:priv1")
        │
        ▼
GraphKernel
        │
        ▼
GraphKernel::CodeGen()
        │
        ▼
GraphSpliter::Run()
        │
        ▼
NDObject::ShapeProp()
        │
        ├─ 根据 obj_id_ 查询 ObjectMeta
        └─ 找到 SelectOp::ShapeProp
                    │
                    ▼
          计算 Select 输出 shape_

GraphSpliter::Run() 中的关键逻辑是：

for (auto op : kernel->build_ops_) {
  op->ShapeProp(sym_dim_next);
}

NDObject::ShapeProp() 根据 ObjectMeta::shape_prop[] 查找回调：

void ShapeProp(int64_t &sym_dim_next) {
  if (auto func = meta_.shape_prop[obj_id_]) {
    func(this, sym_dim_next);
  }
}

原来 Select 在元信息表中对应的回调是 nullptr，因此 GraphKernel 重新推导 Shape 时，Select 不会执行任何操作。

4. 变更前后对比

4.1 变更前

// Select
{kGenSimd3, F_IP | F_NS | F_LR | F_RR,
 nullptr, nullptr, nullptr}

最后一个字段是 shape_prop，值为 nullptr。

因此：

GraphSpliter 调用 Select.ShapeProp()
        │
        └─ 没有回调，直接返回

Select 的 shape_ 不会在该阶段重新计算

4.2 变更后

// Select
{kGenSimd3, F_IP | F_NS | F_LR | F_RR,
 nullptr, nullptr, nullptr, SelectOp::ShapeProp}

调用关系变为：

GraphSpliter
    │
    └─ Select.ShapeProp()
          │
          ├─ 读取 condition、lhs、rhs 的 Shape
          ├─ 按广播规则计算每一维
          └─ 写回 Select.shape_

5. 完整实现

当前实现如下：

void SelectOp::ShapeProp(NDObject *op, int64_t &sym_dim_next) {
  auto *select = static_cast<SelectOp *>(op);
  const IntArrayRef *inputs[] = {
      select->xhs_->shape_ref_,
      select->lhs_->shape_ref_,
      select->rhs_->shape_ref_,
  };
  size_t max_size = 0;
  for (auto input : inputs) {
    max_size = std::max(max_size, input->size);
  }
  select->shape_.Resize(max_size);
  for (size_t i = 0; i < max_size; ++i) {
    int64_t result = 1;
    for (auto input : inputs) {
      size_t offset = max_size - input->size;
      int64_t dim = i < offset ? 1 : input->data[i - offset];
      if (dim == 1 || result == dim) {
        continue;
      }
      if (result == 1) {
        result = dim;
      } else {
        result = sym_dim_next--;
      }
    }
    select->shape_[i] = result;
  }
}

6. 逐段解读

6.1 将通用对象转换为 SelectOp

auto *select = static_cast<SelectOp *>(op);

所有 ShapeProp 回调都使用统一的参数类型：

void (*shape_prop)(NDObject *, int64_t &)

因此传入的参数是通用基类 NDObject*。由于该回调只会注册给 Select，这里可以转换为 SelectOp*，从而访问：

select->xhs_
select->lhs_
select->rhs_
select->shape_

6.2 收集三个输入的 Shape

const IntArrayRef *inputs[] = {
    select->xhs_->shape_ref_,
    select->lhs_->shape_ref_,
    select->rhs_->shape_ref_,
};

这里使用的是三个输入的逻辑 Shape，而不是内部执行用的 nd_。

数组顺序是：

inputs[0] → condition
inputs[1] → lhs
inputs[2] → rhs

三个输入都必须参与广播计算，因为 condition 自己也可能是标量、低 rank Tensor 或带有广播维度的 Tensor。

6.3 计算最大 rank

size_t max_size = 0;
for (auto input : inputs) {
  max_size = std::max(max_size, input->size);
}

input->size 表示 Shape 的维度数量，也就是 rank。

例如：

condition: [B, 1, N]  rank=3
lhs:       [H, N]     rank=2
rhs:       [B, H, N]  rank=3

因此：

max_size = 3

输出 rank 必须至少和输入中 rank 最大的输入一致。

6.4 为输出 Shape 分配空间

select->shape_.Resize(max_size);

这一步为输出 Shape 分配 max_size 个维度。

SelectOp 的构造函数中已经设置：

shape_ref_ = &shape_;

因此更新 select->shape_[i]，就等价于更新 Select 的 shape_ref_。

后续 Area、Stage 和算子切分逻辑读取 shape_ref_ 时，就能看到新的输出 Shape。

6.5 遍历每个输出维度

for (size_t i = 0; i < max_size; ++i) {
  int64_t result = 1;
  ...
  select->shape_[i] = result;
}

外层循环逐维计算输出 Shape。

result 初始值设为 1，是因为 1 是广播计算的单位维度：

1 和 X 广播后得到 X

6.6 对齐不同 rank 的输入

size_t offset = max_size - input->size;
int64_t dim = i < offset ? 1 : input->data[i - offset];

广播按照尾维对齐，rank 较小的输入在高维方向补 1。

例如：

max_size = 3
input shape = [H, N]

逻辑上对齐为：

[1, H, N]

此时：

offset = 3 - 2 = 1

i=0 → dim=1
i=1 → dim=H
i=2 → dim=N

这样就可以把 [H, N] 和 [B, H, N] 按标准广播规则进行比较。

6.7 处理可以广播的维度

if (dim == 1 || result == dim) {
  continue;
}

有两种情况可以直接跳过。

当前输入维度为 1

1 和 B → B
1 和 H → H

维度为 1 的输入可以广播到其他输入的维度，因此不会改变 result。

当前输入维度和结果相同

B 和 B → B
H 和 H → H

说明两个输入在这一维大小一致，也不需要修改结果。

6.8 当前结果为 1 时采用输入维度

if (result == 1) {
  result = dim;
}

例如：

result = 1
dim = B

那么当前输出维度变为：

result = B

再例如：

condition: 1
lhs:       H
rhs:       H

处理过程是：

初始 result = 1
读取 condition 的 1 → result 仍为 1
读取 lhs 的 H       → result = H
读取 rhs 的 H       → 与 result 相同
最终输出 H

6.9 无法静态确定时创建符号维度

else {
  result = sym_dim_next--;
}

如果当前已有一个非 1 维度，而新输入又提供了另一个编译期无法确认是否相同的非 1 维度，就使用新的符号维度表示结果。

例如：

lhs: [M]
rhs: [N]

如果 M 和 N 在编译期无法确认相同，输出维度会使用新的符号编号：

output: [S0]

在 DVM 内部，符号维度通常使用负数表示。GraphSpliter::Run() 会初始化一个符号维度计数器，例如：

int64_t sym_dim_next = -10000;

每生成一个新的符号维度就递减：

-10000 → -10001 → -10002

这样可以避免不同未知维度使用同一个编号。

该行为和 DVM 中 BinaryShapeProp() 的处理方式保持一致。

需要注意：当前函数本身不负责对所有不兼容的具体 Shape 抛出异常，而是沿用现有符号 Shape 推导策略，将编译期无法确定的结果表示为新的符号维度。输入 Shape 的合法性需要由上层图构造和运行时 Shape 约束保证。

6.10 写回输出 Shape

select->shape_[i] = result;

当前维度计算完成后，写回 shape_。

当外层循环结束后：

select->shape_ 中的所有维度都已经确定或符号化

后续 GraphSpliter 可以使用这个结果进行 Area 和 Stage 切分。

7. 算法伪代码

输入：condition_shape、lhs_shape、rhs_shape
输出：select_shape

max_rank = max(rank(condition), rank(lhs), rank(rhs))

for axis in [0, max_rank):
    result = 1

    for input_shape in [condition, lhs, rhs]:
        将 input_shape 按尾维对齐到 max_rank
        当前维度为 dim

        if dim == 1:
            继续

        if result == dim:
            继续

        if result == 1:
            result = dim
        else:
            result = 新符号维度

    select_shape[axis] = result

8. 广播示例

8.1 三个输入都为三维

输入：

condition: [B, 1, N]
lhs:       [1, H, N]
rhs:       [B, H, 1]

逐维计算：

维度	condition	lhs	rhs	输出维度
0	B	1	B	B
1	1	H	H	H
2	N	N	1	N

最终：

Select 输出：[B, H, N]

8.2 rank 不一致

输入：

condition: [H, N]
lhs:       [B, H, N]
rhs:       [1, H, N]

对齐后：

condition: [1, H, N]
lhs:       [B, H, N]
rhs:       [1, H, N]

输出：

[B, H, N]

8.3 标量输入

如果某个输入是标量，逻辑上可以看作 rank 为 0：

condition: scalar
lhs:       [B, H]
rhs:       [B, H]

标量输入在每个维度上都相当于 1，输出为：

[B, H]

8.4 动态符号维度

输入：

condition: [B, 1, K]
lhs:       [1, H, K]
rhs:       [B, H, 1]

输出：

[B, H, K]

如果某一维无法在编译期确认：

lhs: [M]
rhs: [N]

则输出对应维度使用新的符号维度：

output: [S0]

这表示该维度由运行时 Shape 决定，而不是表示生成了额外的运行时算子。

9. 与 Normalize 和 Emit 的区别

Select 相关的三个函数负责不同阶段的工作。

9.1 ShapeProp()

编译期 Shape 分析

职责：

• 读取 condition、lhs、rhs 的逻辑 Shape；
• 计算广播后的输出 Shape；
• 为动态、不确定维度分配符号维度；
• 为 GraphKernel 的 Area/Stage 切分提供 Shape 信息。

不会：

• 插入 Broadcast 算子；
• 生成字节码；
• 发起 kernel launch。

9.2 Normalize()

执行图规范化

职责：

• 恢复 Select 的原始输入；
• 判断三个输入是否需要广播；
• 插入隐式 Broadcast 辅助算子；
• 更新内部执行维度 nd_。

9.3 Emit()

DVM 字节码生成

职责：

• 使用三个输入的 xbuf；
• 设置输出 xbuf；
• 根据数据类型选择 V_SEL_* 指令；
• 将 Select 指令编码到 DVM instruction buffer。

流程关系：

ShapeProp()
    └─ 输出 Shape 分析

Normalize()
    └─ 插入广播辅助节点并准备执行图

Emit()
    └─ 生成 V_SEL_* 字节码

10. 对 GraphKernel 切分的影响

GraphKernel 会根据算子 Shape 建立 Area，并执行 Broadcast、Reduce、Reshape、Cube 等模式合并。

因此 Select Shape 不正确时，可能导致：

Select 输出 Shape 错误
        ↓
Area Shape 判断错误
        ↓
Broadcast/Reduce 合并判断错误
        ↓
Stage 边界错误
        ↓
子图结构或字节码不正确

增加 ShapeProp 后：

Select 三输入广播 Shape 正确
        ↓
GraphSpliter 获得正确 shape_ref_
        ↓
Area 和 border 判断更准确
        ↓
Stage 依赖关系更可靠

该修改主要提升符号式切分的正确性和稳定性。

11. 对性能和字节码的影响

11.1 直接影响

SelectOp::ShapeProp() 只在构图/编译阶段运行，因此本身不会：

• 增加运行时 kernel；
• 增加一次 DVM launch；
• 增加一条 V_SEL_* 指令；
• 增加运行时 Tensor 搬运。

11.2 间接影响

Shape 推导结果会影响 GraphKernel 的切分方式，因此可能间接改变：

• Area 的归属；
• Stage 的数量和顺序；
• Stage 之间的 Store/Load 边界；
• 最终生成的 DVM 子图结构；
• 字节码中的 Load/Store 数量。

因此，ShapeProp 本身没有运行时开销，但修正 Shape 后，最终图结构可能发生变化。

12. 当前实现的边界

12.1 只处理 Shape 推导

该函数不负责判断 Select 是否被 PostGrad 白名单放行，也不负责判断 emitter 是否支持 Select。

当前 Select 已经有：

SelectOp::Emit()
vSelect::Encode()
V_SEL_* 指令

本次只是补充符号式 Shape 推导。

12.2 不等于支持 unsqueeze

该修改与 unsqueeze 支持无关：

Select ShapeProp ≠ unsqueeze emitter 注册
Select ShapeProp ≠ PostGrad 融合白名单

如果图中存在未注册的 unsqueeze，仍然需要在对应的 PostGrad、DVM emitter 或图切分层单独处理。

12.3 当前调用时机

当前仓库中 NDObject::ShapeProp() 的实际调用入口是 GraphSpliter::Run()。

因此通常只有进入：

GraphKernel / GraphSpliter / 符号式 gsplit

时，SelectOp::ShapeProp() 才会被执行。

普通 VKernel、普通 SplitGraphS、SplitGraphD、Sequence 和 Eager 路径不会因为这次注册而自动调用该函数。

12.4 不是每次 launch 都调用

ShapeProp 属于编译/构图阶段，不属于运行阶段：

首次 CodeGen → ShapeProp
后续 launch   → 直接执行已经生成的代码

它不会在每次 DVM kernel launch 时重复执行。

13. 与普通 Binary ShapeProp 的关系

DVM 已经有通用的二元广播 Shape 推导：

BinaryOp::ShapeProp()
CompareOp::ShapeProp()
PowerOp::ShapeProp()

这些函数处理两个输入：

lhs、rhs

Select 需要处理三个输入：

condition、lhs、rhs

因此 SelectOp::ShapeProp() 复用了相同的思想，但将广播合并从两个输入扩展到了三个输入。

14. 代码修改清单

DVM 仓库

src/ops.cc
    1. 在 GenObjectMeta() 中注册 SelectOp::ShapeProp
    2. 新增 SelectOp::ShapeProp() 实现

src/ops.h
    3. 增加 SelectOp::ShapeProp() 声明

没有修改的内容

SelectOp::Normalize()
SelectOp::Emit()
SelectOp::Clone()
vSelect::Encode()
PostGrad 融合白名单
unsqueeze 支持
普通 split 的切分算法

15. 总结

SelectOp::ShapeProp() 做的事情可以概括为：

读取 Select 的三个输入 Shape
        ↓
按尾维对齐
        ↓
逐维执行三输入广播合并
        ↓
无法静态确定的维度使用符号维度
        ↓
写回 Select.shape_
        ↓
供 GraphKernel 符号式切分使用

它解决的是：

GraphKernel 在重新分析符号图时，Select 输出 Shape 没有被正确推导的问题。

它不是运行时优化，也不是新的 Select 执行实现，而是 GraphKernel 切分阶段所需的 Shape 元数据补全。