Star67
10
代码
代码
Issues1
Pull Requests5
Actions
项目讨论
Wiki
分析
Star67
10
已合并
新增 Reformer LSH 桶排序算子 #30
海阔天空创建于 14 天前
新增 Reformer LSH 桶排序算子 #30
已合并
海阔天空创建于 14 天前
从haikuo/mat-chem-sim-predfeature/reformer-lsh-bucket-sort合入到cann/mat-chem-sim-predmaster
海阔天空
海阔天空成员
14 天前

变更概述

本 MR 新增 ReformerLshBucketSort，用于精确替换 Reformer LSH attention
中的两次通用排序。算子输入确定性的“桶编号 + token 位置”编码，同时输出：

  • 稳定排序后的 key；
  • 排序后位置到原位置的正向 permutation；
  • 从排序结果恢复 token 原顺序的 inverse permutation。

该算子替换 reformer-pytorch 1.4.4 中由 TSLib Reformer encoder 调用的
sort_key_val 与 inverse-sort 热点。它不替换 attention score、softmax、
value aggregation 或预测头，不改变 Reformer 的模型结构和预测语义。

Native、TorchAir 与自定义算子三路审计

当前 TSLib 源码在 hash 路径中创建 CPU torch.Generator，会阻断官方模型的
fullgraph 捕获。TorchAir 审计恢复了上游 torch.randn(..., device=npu) 的
NPU 随机旋转表达，不改变 hashing、sorting、inverse sorting、attention、
模型权重或输入。

测试形状 Native eager，非 TorchAir 兼容 TorchAir ACLNN 自定义算子 判定
B4/L96 4.316 ms 3.008 ms 非目标提交形状 TorchAir 最快
B16/L96 10.291 ms 11.087 ms 非目标提交形状 Native 最快
B32/L96 18.088 ms 32.048 ms 非目标提交形状 Native 最快
B4/L336 9.039 ms 9.015 ms 非目标提交形状 基本持平
B16/L336 28.604 ms 50.493 ms 非目标提交形状 Native 最快
B32/L336 55.528 ms 99.733 ms 17.577 ms 相对最快 Native 降低 68.35%

六档兼容 TorchAir 图均能正确执行，观察到的最大图模式/reference 误差为
3.40e-5。TorchAir 在最小 B4/L96 档占优，但没有吸收本算子所针对的
B32/L336 排序热点。

算子阶段与完整测试集证据

  • B32/L336 可替换阶段：37.2931 -> 0.6023 ms，延迟降低 98.38%，
    加速 61.92 倍，三个输出均逐元素一致。
  • 同轮三路完整模型：Native 55.5278 ms、TorchAir 99.7328 ms、
    自定义算子 17.5769 ms，相对最快 Native 降低 68.35%。
  • 独立交付 gate：55.1197 -> 17.5769 ms，延迟降低 68.11%。
  • 完整 ETTh1 测试集共 2,857 个窗口：Native 1558.82 ms -> 自定义算子
    682.56 ms，延迟降低 56.21%，预测指标保持不变。

完整测试集 headline 使用 Native 作为基线，因为同形状 TorchAir 图在替换前
比 Native 慢 79.62%。同轮三路 68.35% 与独立 gate 68.11% 来自两次独立
测试口径，不进行混算。

对外接口与验证

  • aclnnReformerLshBucketSortGetWorkspaceSize；
  • aclnnReformerLshBucketSort。

交付件包含 Host API、Ascend C kernel、稳定顺序与 inverse 索引契约测试、
独立构建、ACL smoke、输出所有权测试、文档、形状矩阵证据、checkpoint 模型
E2E 以及完整测试集性能证据。

likedislike
Pull Request已成功合入, 合并人@CANN-robot
（感谢 海阔天空 的贡献）
海阔天空海阔天空成员
14 天前 创建了 pull request，commit e86b42c5
atomgit-bot
atomgit-bot
14 天前 评论：

变更摘要

本 PR 新增 ReformerLshBucketSort 自定义算子，用于在 Ascend 910B3 上精确替换 Reformer LSH attention 中的两次通用排序（sort_key_val 与 inverse-sort）。算子输入为"桶编号 + token 位置"编码的 int64 键值矩阵，通过基于桶计数排序的 Ascend C kernel 同时输出稳定排序后的 key、正向 permutation（sticker）和逆向 permutation（inverse）。该算子不改变 Reformer 模型结构与预测语义，在 B32/L336 形状下相对框架原生排序实现最高加速 62.47 倍，完整 ETTh1 测试集端到端延迟降低 56.21% 且预测指标完全一致。

主要改动

  • 新增 Ascend C 内核实现：reformer_lsh_bucket_sort_kernel.cpp 实现了基于桶计数的稳定排序 kernel，按行并行执行"计数 → 前缀和 → 散射"三阶段，将通用排序替换为确定性桶排序，并同时产出 sorted_keys、sticker 和 inverse 三个输出张量。

  • 新增 Host 侧算子定义与 Tiling：reformer_lsh_bucket_sort_def.cpp 注册了 ReformerLshBucketSort 算子类，声明一个 int64 ND 输入和三个同形状输出，以及 sequence_length 和 total_buckets 两个必选整型属性；reformer_lsh_bucket_sort_tiling.h 定义了包含 rows、total_length、sequence_length、total_buckets 的 Tiling 数据结构，Host 端在 reformer_lsh_bucket_sort_host.cpp 中对输入形状和属性值域进行校验（total_buckets ≤ 4096，维度大小适配 uint32/int32）。

  • 新增 Python 适配器与守卫逻辑：integration/adapter.py 提供 supports_custom() 和 dispatch() 函数，严格校验输入 tensor 的 device、dtype（仅 torch.int64）、contiguous 属性及形状/参数值域上限，不满足条件时自动回退到框架原生 fallback，确保非目标场景不受影响。

  • 新增 CPU 参考实现与契约测试：reference/reference.py 提供了无依赖的纯 Python 参考实现，明确稳定排序契约（同桶内保持原始相对顺序）和 inverse 索引可逆性；tests/test_reference.py 覆盖标准用例、边界值、随机稳定性验证和最小值校验。

  • 新增独立构建 Gate 与性能证据：run_packaged_gate.sh 整合 msopgen 生成、源码替换、CMake 构建和 ACLNN smoke 测试的完整流水线；docs/evidence/ 目录下的 JSON 证据文件记录了 6 档形状的组件级速度提升（最高 62.47×）、B32/L336 模型级 68.11% 延迟降低及 ETTh1 完整测试集 2857 窗口的端到端结果。

likedislike
atomgit-bot
atomgit-bot
14 天前 评论：

代码审查

All major aspects have been checked. Now let me provide the closing summary.

审查结束 — 汇总

各文件审查结论

文件 结果
CMakeLists.txt 无问题
README.md 无问题
docs/algorithm.md 无问题
docs/api_reference.md 无问题
docs/benchmark.md 无问题
docs/evidence/checkpoint_etth1.json 无问题
docs/evidence/component_matrix.json 无问题
docs/evidence/model_ablation.json 无问题
docs/evidence/multistream.json 无问题
docs/evidence/node202_gate_summary.json 无问题
docs/test_report.md 无问题
docs/upstream_dedup.md 无问题
examples/test_aclnn_reformer_lsh_bucket_sort.cpp 2 个 P2 + 1 个 P3
integration/adapter.py 无问题
msopgen/reformer_lsh_bucket_sort_msopgen.json 无问题
op_host/reformer_lsh_bucket_sort_def.cpp 无问题
op_host/reformer_lsh_bucket_sort_host.cpp 1 个 P3
op_host/reformer_lsh_bucket_sort_tiling.h 无问题
op_kernel/reformer_lsh_bucket_sort_kernel.cpp 无问题
reference/reference.py 无问题
tests/CMakeLists.txt 无问题
tests/benchmark_reformer_lsh_bucket_sort_aclnn.cpp 无问题
tests/run_packaged_gate.sh 无问题
tests/test_adapter.py 无问题
tests/test_reference.py 无问题
tests/ut/CMakeLists.txt 无问题
tests/ut/op_kernel/CMakeLists.txt 无问题
tests/ut/op_kernel/test_reformer_lsh_bucket_sort.cpp 无问题

按优先级统计

  • P0: 0
  • P1: 0
  • P2: 2（设备内存泄漏、未检查空 tensor 指针，均位于示例代码）
  • P3: 2（workspace 错误路径泄漏、InferShape 缺少空指针检查）

整体风险判断

低风险。 核心算子逻辑（kernel、tiling、host 校验）经审查未发现正确性问题：32 位乘法链首已正确提升到 64 位、边界条件在 host 侧有完整的参数校验（拒绝 0 值、超界 shape、非法桶数）、计数排序的稳定性实现正确、GM 偏移计算均在 uint64 域内。发现的问题全部集中在示例代码 test_aclnn_reformer_lsh_bucket_sort.cpp 的错误处理路径（资源泄漏与空指针），以及 InferShape 函数的一处防御性编码不一致，均不影响算子的生产正确性与性能。建议在合入前修复示例代码中的两个 P2 资源管理问题，避免该示例被其他算子开发者复制时扩散不良模式。

类型 数量
🔴 阻塞 0
🟡 建议 2

💬 仅评论

likedislike
CANN-robotCANN-robot成员
14 天前 添加了label：cann-cla/yes
CANN-robot
CANN-robot成员
14 天前 评论：

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here

PR Approval Progress

⚠️ This PR does not yet meet the following requirements:lgtm (requires ≥ 2 person(s) per module)、approve (requires ≥ 1 person(s) per module)

Module Approval Details

module lgtm status approve status
repo-cann/mat-chem-sim-pred ❌ (0/2)(You can also ask: 李姝漫, 赵俊, 高菲, 黄剑兴, 张强豪) ❌ (0/1)(You can also ask: 张玉橙, 高梓博, 刘达林, 张强豪, 李姝漫)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

haikuo, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
CANN-robot
CANN-robot成员
14 天前 评论：

CLA Signature Pass

haikuo, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
CANN-robotCANN-robot成员
14 天前 将gaozibo,Splendid2025,TokenDance,zhangqianghao,magiclf0716,zhaojun735,imSixSeven,LiizzMM,huawei_liudalin,Huang-jianxing,gcw_Vy5QKdan,gaofei06设为评审人
CANN-robotCANN-robot成员
14 天前 将gaozibo,Splendid2025,TokenDance,zhangqianghao,magiclf0716,zhaojun735,imSixSeven,LiizzMM,huawei_liudalin,Huang-jianxing,gcw_Vy5QKdan,gaofei06设为审查人
海阔天空
海阔天空成员
14 天前 评论：

/compile

likedislike
CANN-robotCANN-robot成员
14 天前 添加了label：ci-pipeline-running
CANN-robot
CANN-robot成员
14 天前 评论：

流水线任务触发成功
任务链接 [bc3e9e5bcfe34a5095ce6905c541ab08][流水线指导]

任务名称状态日志下载链接
SCA ✅ SUCCESS >>>>>
antipoison ✅ SUCCESS >>>>>
Check_Pr ✅ SUCCESS >>>>>
codecheck_style ✅ SUCCESS >>>>>
pre-commit ✅ SUCCESS
pre_comment ✅ SUCCESS >>>>>

[2026-08-06 00:22:35]    CI执行结束

likedislike
CANN-robotCANN-robot成员
14 天前 删除了label：ci-pipeline-running
CANN-robotCANN-robot成员
14 天前 添加了label：ci-pipeline-passed
atomgit-bot
atomgit-bot14 天前进行代码检视2
prediction/ProcessControl/TimeSeriesForecast/reformer_lsh_bucket_sort/examples/test_aclnn_reformer_lsh_bucket_sort.cpp
已过期
@@ -0,0 +101,4 @@
101+ aclTensor* keys = MakeTensor({1, 8}, ACL_INT64, dKeys);
102+ aclTensor* sorted = MakeTensor({1, 8}, ACL_INT64, dSorted);
103+ aclTensor* sticker = MakeTensor({1, 8}, ACL_INT64, dSticker);
104+ aclTensor* inverse = MakeTensor({1, 8}, ACL_INT64, dInverse);
atomgit-bot
atomgit-bot14 天前评论：

🟡 Medium Priority

changed line 101-104 → MakeTensor 内部调用 aclCreateTensor，分配失败时返回 nullptr。四路 MakeTensor 返回值（keys、sorted、sticker、inverse）均未做空检查，直接传入 aclnnReformerLshBucketSortGetWorkspaceSize。若任一 tensor 为 null，ACLNN API 将收到空指针，行为未定义，大概率触发段错误或 AIC error。

建议：在每个 MakeTensor 调用后增加空指针检查，若返回 nullptr 则释放已分配资源（dKeys/dSorted/dSticker/dInverse 及已创建的 tensor）后返回错误码。

likedislike
System
系统消息系统
6 天前 评论：

changed this line on 9036b1f5 view diff detail

atomgit-bot
atomgit-bot14 天前进行代码检视1
prediction/ProcessControl/TimeSeriesForecast/reformer_lsh_bucket_sort/examples/test_aclnn_reformer_lsh_bucket_sort.cpp
@@ -0,0 +97,4 @@
97+ !AllocateOutput<int64_t>(&dSticker, 8) || !AllocateOutput<int64_t>(&dInverse, 8))
98+ {
99+ return 1;
100+ }
atomgit-bot
atomgit-bot14 天前评论：

🟡 Medium Priority

changed line 96-100 → 当 CopyToDevice(&dKeys, hKeys) 成功（dKeys 已分配），但后续 AllocateOutput<int64_t>(&dSorted, 8) 或更后面的分配失败时，短路求值使整体条件为 true，直接 return 1，dKeys（以及可能已成功分配的 dSorted、dSticker）泄漏。此外，CopyToDevice 内部也存在同类泄漏：若 aclrtMalloc 成功但 aclrtMemcpy 失败，已分配内存同样泄漏。

建议：在 return 前释放已成功分配的资源。将分配步骤拆分为独立的 if 判断，每个失败分支释放已分配的资源；或使用 RAII 包装 aclrtMalloc/aclrtFree。

likedislike
H
Huang-jianxing成员
7 天前 评论：

/lgtm

likedislike
H
Huang-jianxing成员
7 天前 评论：

/approve

likedislike
CANN-robotCANN-robot成员
7 天前 添加了label：approved
LiizzMM
LiizzMM成员
6 天前 评论：

/lgtm

likedislike
CANN-robotCANN-robot成员
6 天前 添加了label：lgtm
CANN-robotCANN-robot成员
6 天前 删除了label：ci-pipeline-passed
CANN-robot
CANN-robot成员
6 天前 评论：

The following label is not ready.

ci-pipeline-passed: The ci-pipeline-passed label is expired as added over 72 hours. Please use /compile to recompile.

likedislike
LiizzMMLiizzMM成员
6 天前 将 cann-robot 设为合并人
LiizzMMLiizzMM成员
6 天前 通过了评审
LiizzMMLiizzMM成员
6 天前 通过审查
LiizzMMLiizzMM成员
6 天前 解决了最后一个问题
LiizzMM
LiizzMM成员
6 天前 评论：

compile

likedislike
海阔天空海阔天空成员
6 天前 update merge request[project id: 9128418, iid: 30, commit_id: 3286c0368cbf223dcab33328d796590518cd8de0] virtual merging success
此处折叠了11条事件消息 查看更多
CANN-robotCANN-robot成员
6 天前 添加了label：lgtmapprovedcann-cla/yes
CANN-robot
CANN-robot成员
6 天前 评论：

Notice

New code changes of the pull request are detected and remove these labels lgtm, approved. 😳

likedislike
海阔天空
海阔天空成员
6 天前 评论：

compile

likedislike
海阔天空海阔天空成员
6 天前 update merge request[project id: 9128418, iid: 30, commit_id: 5f8db7deaf8bd079f0e66e72bdfe443eeff8143e] virtual merging success
CANN-robotCANN-robot成员
6 天前 删除了label：ci-pipeline-failed
CANN-robotCANN-robot成员
6 天前 添加了label：ci-pipeline-running
CANN-robotCANN-robot成员
6 天前 删除了label：ci-pipeline-running
CANN-robotCANN-robot成员
6 天前 添加了label：ci-pipeline-passed
CANN-robotCANN-robot成员
6 天前 合入了pull request