[2025秋季][T1-2-1]ttaohe

[ttaohe]

PR 背景 / 动机

rearrange 在 NVIDIA 上遇到典型 row-major ↔ col-major / full-transpose 的 stride pattern 时，通用 kernel 存在严重的非合并访存 + 索引开销，导致在一些 case（尤其 6D/大 2D）明显落后 PyTorch。
本 PR 引入 pattern detection + 专用 transpose fast-path kernels + 保底 fallback，对常见 transpose 类 case 提升显著，并保持不命中时回退到原通用实现。

主要改动

1) NVIDIA rearrange 新增 transpose fast-path（带 fallback）

• 在 Descriptor::calculate() 中识别并优先尝试 fast-path：
  • 2D：row-major ↔ col-major layout transform
  • 4D~6D：full-transpose（stride order 反转）pattern
• fast-path 失败/不匹配则 回退 到既有 prepareRearrangeParams + static/dynamic kernel 路径，保证正确性。

涉及文件：

• src/infiniop/ops/rearrange/nvidia/rearrange_nvidia.cu

2) 新增/增强 transpose kernels（2D / 5D / 6D）

新增专用 kernel 文件并持续扩展：

• src/infiniop/ops/rearrange/nvidia/rearrange_transpose_kernel.cuh

包含：

• 2D shared-memory tiled transpose（32×32 tile，避免 bank conflict）
• 2D small-matrix 16×16 tile：优化 100×100 这类小矩阵的 launch/sync 占比
• 6D VEC+carry kernel（F16/F32）：每线程处理 4 个连续元素，减少 div/mod 索引开销（F16 用 uint16_t bitwise copy）
• 5D VEC+carry kernel（F16/F32）：覆盖 (3,4,7,53,9) 这类中等规模 full-transpose

3) 调整 full-transpose pattern 的启发式阈值

• isFullTransposePattern() 维持严格的 stride-order 反转判断
• 阈值按维度区分：让 5D case 可以命中，但总体仍然保守，避免小 case 误判带来回退开销/波动。

涉及文件：

• src/infiniop/ops/rearrange/nvidia/rearrange_nvidia.cu

4) 测试/Benchmark 集成与修复

• 增加 test/infinicore/ops/rearrange.py，支持 run.py --ops rearrange --bench 跑功能与性能对比
• 调整为 预分配 out + 只 copy_ 的 bench 方式，更贴近算子核心开销
• 修复测试框架里 rearrange_tensor() 对非 contiguous tensor 使用 .view(-1) 报错，改为 .reshape(-1)：
  • test/infinicore/framework/utils.py

Benchmark（示例）

在 srun --gres=gpu:nvidia:1 python test/infinicore/run.py --ops rearrange --nvidia --bench 下：

• (2000,2000) F16/F32：显著快于 PyTorch（2D tiled fast-path 命中）
• (3,4,50,50,5,7) F32：由 6D VEC+carry fast-path 拉起（多数情况下可超过 PyTorch）
• (3,4,7,53,9) F16/F32：5D fast-path 基本追平或略优（视环境波动）

如何构建 & 复现

# 编译安装 C++/CUDA
xmake build infiniop && xmake install infiniop

# 运行单算子 bench（device timing）
srun --gres=gpu:nvidia:1 --cpus-per-task=16 --mem=256G \
  python test/infinicore/run.py --ops rearrange --nvidia --bench

风险与回退策略

• fast-path 仅在严格 pattern 匹配时启用；否则自动 fallback
• kernel launch 失败会返回错误并回退到原通用实现（避免影响正确性）
• 小矩阵采用小 tile，降低 overhead；大矩阵保持 32×32 tile

HONER_CODE.md

HONOR_CODE.md

REFERENCE.md

REFERENCE.md