[Feature] [Performance] VPTO 未将受条件保护的 SIMT 循环不变量地址计算外提

[LLMZhangYC]

Summary

A5 VPTO 管线在 scf.for 内部嵌套了依赖循环变量的 scf.if 时，没有将其中可安全推测执行、且与循环变量无关的整数和地址计算外提。

Motivation / use case

module attributes {
  pto.target_arch = "a5",
  pto.kernel_kind = #pto.kernel_kind<vector>
} {
  func.func @simt_guarded_licm(
      %dst: !pto.ptr<i32, ub>,
      %w: i32,
      %valid_end: i32,
      %local_end: i32,
      %total_start: i32,
      %trip_count: index) attributes {pto.entry} {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c31_i32 = arith.constant 31 : i32
    %c32_i32 = arith.constant 32 : i32
    %c3_i64 = arith.constant 3 : i64
    %c31_i64 = arith.constant 31 : i64
    %c32_i64 = arith.constant 32 : i64
    %c128_i64 = arith.constant 128 : i64
    %c1024_i64 = arith.constant 1024 : i64
    %minus_one = arith.constant -1 : i32

    pto.section.simt<<<512, 1, 1>>> {
      %tid = pto.get_tid_x : i32
      %lane = arith.andi %tid, %c31_i32 : i32
      scf.for %i = %c0 to %trip_count step %c1 {
        %i32 = arith.index_cast %i : index to i32
        %iter_offset = arith.muli %i32, %c32_i32 : i32
        %logical_index = arith.addi %iter_offset, %lane : i32
        %in_bounds = arith.cmpi slt, %logical_index, %local_end : i32
        scf.if %in_bounds {
          %w64 = arith.extsi %w : i32 to i64
          %slot = arith.andi %w64, %c3_i64 : i64
          %slot_base = arith.muli %slot, %c1024_i64 : i64
          %valid64 = arith.extsi %valid_end : i32 to i64
          %local64 = arith.extsi %local_end : i32 to i64
          %range_base = arith.minsi %valid64, %local64 : i64
          %tid64 = arith.extsi %tid : i32 to i64
          %lane64 = arith.andi %tid64, %c31_i64 : i64
          %total_start64 = arith.extsi %total_start : i32 to i64
          %base0 = arith.addi %slot_base, %range_base : i64
          %base1 = arith.addi %base0, %lane64 : i64
          %base2 = arith.addi %base1, %c128_i64 : i64
          %base = arith.subi %base2, %total_start64 : i64

          %i64 = arith.extsi %i32 : i32 to i64
          %dynamic_offset = arith.muli %i64, %c32_i64 : i64
          %offset = arith.addi %base, %dynamic_offset : i64
          %index = arith.index_cast %offset : i64 to index
          pto.store %minus_one, %dst[%index] : !pto.ptr<i32, ub>, i32
        }
      }
    }
    return
  }
}

Proposed API / behavior

这些与循环变量无关、可安全推测执行的地址表达式应在每个 SIMT thread 中只计算一次，即外提到 scf.for 之前，或通过等价的地址强度削减消除重复计算。循环体内应只保留依赖循环变量的 i * 32 部分。

最终 SIMT IR 的结构预期类似：

%base = ... // 仅依赖 w、valid_end、local_end、tid、total_start
scf.for %i = ... {
  %in_bounds = ... // 依赖 %i
  scf.if %in_bounds {
    %dynamic_offset = ... // 依赖 %i
    %offset = arith.addi %base, %dynamic_offset
    pto.store ...
  }
}

对于原始算子，消除重复的 SIMT 地址计算后，预期可以降低目前约 47%–50% 的额外 vector pipe 活跃周期，并使 SIMT 代码体积更接近 AscendC。这里不要求与 AscendC 达到完全相同的性能；核心预期是循环不变量不再在每个受保护的循环迭代中重复执行。

Alternatives considered

最终 VPTO IR 仍将所有不变的基址计算保留在 scf.for 及其嵌套的 scf.if 内：

scf.for %arg6 = %c0 to %arg5 step %c1 {
  %2 = arith.index_cast %arg6 : index to i32
  %3 = arith.muli %2, %c32_i32 : i32
  %4 = arith.addi %3, %1 : i32
  %5 = arith.cmpi slt, %4, %arg0 : i32
  scf.if %5 {
    %6 = arith.extsi %arg1 : i32 to i64
    %7 = arith.andi %6, %c3_i64 : i64
    %8 = arith.muli %7, %c1024_i64 : i64
    %9 = arith.extsi %arg2 : i32 to i64
    %10 = arith.extsi %arg0 : i32 to i64
    %11 = arith.minsi %9, %10 : i64
    %12 = arith.extsi %0 : i32 to i64
    %13 = arith.andi %12, %c31_i64 : i64
    %14 = arith.extsi %arg3 : i32 to i64
    %15 = arith.addi %8, %11 : i64
    %16 = arith.addi %15, %13 : i64
    %17 = arith.addi %16, %c128_i64 : i64
    %18 = arith.subi %17, %14 : i64
    %19 = arith.extsi %2 : i32 to i64
    %20 = arith.muli %19, %c32_i64 : i64
    %21 = arith.addi %18, %20 : i64
    %22 = arith.index_cast %21 : i64 to index
    pto.store %c-1_i32, %arg4[%22] : !pto.ptr<i32, ub>, i32
  }
}

Additional context

No response

[LLMZhangYC]

PR 确实改善了最终 VPTO IR，但没有显著减少最终 SIMT 指令开销，PTOAS 需要继续处理：

• 对已知非负数的 floordivsi(x, 32) 做强度削弱，生成 shrui(x, 5)，避免额外的符号修正指令。
• 将符合范围条件的 SIMT scf.for 循环变量彻底收窄到 i32，避免循环中的 i64/index 转换。
• 将 TID、循环变量、长度等非负/范围信息保留到 LLVM IR。
• 在 ObjectEmission.cpp 中对比 AscendC 使用的 Bisheng 优化配置，验证 -O3、MI scheduler、post-RA scheduler 等选项。