feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3-GAMMALab-dgl-ascend-AtomGit

已合并

feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3 #21

lhp_lhp创建于 7 天前从lhp_lhp/dgl-ascend： feat/scatter-add-ascend 合入到 AI4Science/dgl-ascend： master

feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3 #21

已合并

lhp_lhp创建于 7 天前

从lhp_lhp/dgl-ascend feat/scatter-add-ascend合入到AI4Science/dgl-ascend master

共 6 个文件变更+1949-15

Adocs/ascend/operators/scatter_add.md+204-0

@@ -0,0 +1,204 @@
	1	+# ScatterAdd
	2	+
	3	+> Ascend NPU 算子文档。
	4	+
	5	+## 功能描述
	6	+
	7	+ScatterAdd 是沿第一维的散射累加算子，数学语义为：
	8	+
	9	+```
	10	+对每个 i ∈ [0, N):
	11	+ out[idx[i], ] += feat[i, ]
	12	+```
	13	+
	14	+即把 `feat` 的每一行按 `idx` 指定的目标行号累加到 `out` 对应行上。当多个 `i` 映射到同一目标行时，所有命中的 `feat[i,*]` 全部累加到该 `out` 行（跨核写冲突由 GM 原子累加保证正确性）。
	15	+
	16	+该算子为 DGL Ascend 后端的原生 kernel，替换原先的 CPU 回退路径（D2H + `cpu::ScatterAdd` + H2D），消除 host-device 间数据搬运开销。累加策略采用 Strategy A：`SetAtomicAdd<float>()` 开启 GM 原子 read-modify-write，通过 `DataCopyPad` UB→GM 实现原子加，语义与 CUDA `atomicAdd` 的 `+=` 逐字一致。kernel 不清零 `out`，仅做 `+=`；`out` 由调用方预清零（与 CUDA/CPU 一致，DGL Python 调用层 `F.zeros` 已保证）。
	17	+
	18	+## 接口原型
	19	+
	20	+```cpp
	21	+// host 侧 NDArray 接口（src/array/kernel_decl.h 声明，与 CUDA/CPU 同构）
	22	+template <int XPU, typename IdType, typename DType>
	23	+void ScatterAdd(NDArray feat, NDArray idx, NDArray out);
	24	+
	25	+// Ascend 特化（本次新增）
	26	+template <>
	27	+void ScatterAdd<kDGLAscend, IdType, DType>(NDArray feat, NDArray idx, NDArray out);
	28	+
	29	+// device kernel 入口（scatter_add_kernel.cpp）
	30	+extern "C" __global__ __aicore__ void kernel_scatter_add(
	31	+ GM_ADDR feat, GM_ADDR idx, GM_ADDR output, GM_ADDR tiling_ptr);
	32	+
	33	+// host launcher（scatter_add.cc，ACLRT_LAUNCH_KERNEL 宏展开）
	34	+extern "C" uint32_t aclrtlaunch_kernel_scatter_add(
	35	+ uint32_t blockDim, aclrtStream stream,
	36	+ void* feat, void* idx, void* out, void* tiling);
	37	+```
	38	+
	39	+> 上层 Python 通路：`dgl.backend.scatter_add(x, idx, m)` → `_CAPI_DGLKernelScatterAdd`（`_sparse_ops.py`），host 侧 `F.zeros` 预清零 `out` 后调用本算子。
	40	+
	41	+## 参数说明
	42	+
	43	+\| 参数 \| 内存位置 \| 方向 \| 类型 \| 说明 \|
	44	+\|------\|----------\|------\|------\|------\|
	45	+\| feat \| Device (GM) \| 输入 \| DType* \| 特征张量，shape `[N, ]`，行优先；`N = feat->shape[0]`，`` 为第 1..k 维乘积记为 `dim`。按行散射，每行被唯一核读取。 \|
	46	+\| idx \| Device (GM) \| 输入 \| IdType* \| 目标行号张量，shape `[N]`，元素须满足 `0 <= idx[i] < M`；越界行为由上游保证（kernel 不显式校验，与 CUDA 一致）。int64 时 host 侧经 `AsNumBits(32)` 位宽转换后送 kernel。 \|
	47	+\| out \| Device (GM) \| 输出 \| DType* \| 输出张量，shape `[M, ]`，`M = out->shape[0]`。调用前必须由调用方清零*；kernel 仅做 `+=`，不清零。 \|
	48	+\| tiling \| Device (GM) \| 输入（host 下发） \| `ScatterAddTilingData*` \| 形状/切分参数 `{ N, M, featDim }`，由 host `aclrtMalloc` + `aclrtMemcpy` 下发，kernel `Init` 时读取。 \|
	49	+
	50	+### Tiling 结构体
	51	+
	52	+```cpp
	53	+struct ScatterAddTilingData {
	54	+ uint32_t N; // feat 行数（= idx 长度）
	55	+ uint32_t M; // out 行数
	56	+ uint32_t featDim; // 每行元素数 dim = ∏(out->shape[1..ndim])
	57	+};
	58	+```
	59	+
	60	+## 支持数据类型
	61	+
	62	+\| 数据类型（IdType × DType） \| 支持状态 \| 说明 \|
	63	+\|----------\|------\|------\|
	64	+\| int32 × float32 \| 支持 \| 原生 kernel 直走，首版落地。 \|
	65	+\| int64 × float32 \| 支持 \| host 侧 `idx.CopyTo(CPU) → AsNumBits(32) → CopyTo(NPU)` 位宽转换后走 int32 内核（与 `segment_reduce.cc` 同范式，`AsNumBits` 为位宽保持转换非截断）。 \|
	66	+\| int32/int64 × half (float16) \| 不支持 \| `LOG(FATAL)` 占位，提示「Current Ascend scatter_add kernel only supports float features」。 \|
	67	+\| int32/int64 × bfloat16 \| 不支持 \| 非 CUDA 构建下由 `ATEN_FLOAT_TYPE_SWITCH_16BITS` 宏在到达特化前 `LOG(FATAL)`，与 `segment_reduce.cc` 一致。 \|
	68	+\| int32/int64 × double (float64) \| 不支持 \| `LOG(FATAL)` 占位。 \|
	69	+\| 其他 dtype (int8/uint8/...) \| 不支持 \| dispatch 走 `ATEN_FLOAT_TYPE_SWITCH_16BITS`，不进入本算子。 \|
	70	+
	71	+## 约束说明
	72	+
	73	+### Shape 约束
	74	+
	75	+- `feat->shape[0] == idx->shape[0] == N`。
	76	+- `out->ndim == feat->ndim`，且 `out->shape[1..ndim]` 与 `feat->shape[1..ndim]` 逐维相等（即 `dim` 相同）。
	77	+- `M = out->shape[0]`，`idx[i]` 须满足 `0 <= idx[i] < M`（上游保证，kernel 不校验）。
	78	+- `dim = ∏(out->shape[1..ndim])`，允许 `dim % 8 != 0`（非对齐），`DataCopyPad` 自动处理行尾对齐填充。
	79	+- `N`、`M`、`dim` 受 GM/HBM 容量约束，无固定上限；多核切分 `blockDim = min(N, 40)`，`N=0` 时 `blockDim=1`、kernel 提前返回（空操作）。
	80	+
	81	+### dtype 约束
	82	+
	83	+- 仅支持 `IdType ∈ {int32, int64}` × `DType = float32`，其余组合 `LOG(FATAL)`（见上表）。
	84	+
	85	+### 清零契约
	86	+
	87	+- `out` 必须由调用方预清零。kernel 仅执行 `+=`（`SetAtomicAdd` 原子累加），不清零也不初始化 `out`。DGL Python 调用层 `_sparse_ops.py` 已在调用 `_CAPI_DGLKernelScatterAdd` 前 `out = F.zeros(out_shp, dtype, ctx)`。若调用方未清零，原子加会在旧值上累加，结果错误——此行为与 CUDA/CPU 实现一致。
	88	+
	89	+### 平台约束
	90	+
	91	+- 目标芯片：Ascend 910B3（dav-2201，`--npu-arch=ascend910b`）。
	92	+- 代码架构：SIMD（实现载体 MemBase，`kernel_operator.h` + `GlobalTensor` + `TPipe`/`TQue`），纯 Vector 任务（`KERNEL_TYPE_AIV_ONLY`）。
	93	+- `SetAtomicAdd<float>` 要求芯片支持 float 类型 GM 原子加，dav-2201 支持。
	94	+
	95	+### 确定性约束
	96	+
	97	+- 累加顺序非确定：`SetAtomicAdd` 的原子 RMW 顺序由硬件调度，同一输入多次执行累加顺序可能不同（尤其跨核冲突时）。此为预期行为，与 CUDA `atomicAdd` 同范式，不保证 bit-exact，精度由混合容差覆盖。
	98	+
	99	+### 调用前置条件
	100	+
	101	+- host 侧 `ScatterAddDispatch`（`src/array/kernel.cc`）经 `ATEN_XPU_SWITCH_CUDA_ASCEND` + `ATEN_ID_TYPE_SWITCH` + `ATEN_FLOAT_TYPE_SWITCH_16BITS` 三级分派进入特化。调用前 host 侧执行 `aclrtSynchronizeDevice()` 排空 PyTorch NPU op，确保 `feat/idx/out` 数据就绪。
	102	+
	103	+## 返回值 / 错误码
	104	+
	105	+\| 返回值 \| 含义 \|
	106	+\|--------\|------\|
	107	+\| （host 接口 `void`） \| 无返回值。kernel launch 失败时 host 侧 `LOG(FATAL)` 终止（释放 tiling 后报错码）。 \|
	108	+\| `aclrtlaunch_kernel_scatter_add` 返回 `ACL_SUCCESS` \| kernel launch 成功。 \|
	109	+\| `aclrtlaunch_kernel_scatter_add` 返回非 `ACL_SUCCESS` \| launch 失败，host `aclrtFree(tiling)` 后 `LOG(FATAL)` 输出错误码。 \|
	110	+\| `aclrtMalloc` / `aclrtMemcpy` / `aclrtSynchronizeStream` / `aclrtFree` 失败 \| `CHECK` 宏终止，输出 `Ascend Error, code: <e>`。 \|
	111	+\| 不支持 dtype \| `LOG(FATAL)`："Current Ascend scatter_add kernel only supports float features." \|
	112	+\| 未编译 Ascend 支持 \| `LOG(FATAL)`："Ascend support is not compiled. Please compile with -DUSE_ASCEND=ON"。 \|
	113	+
	114	+## 调用示例
	115	+
	116	+### C++ 直接调用（host 侧）
	117	+
	118	+```cpp
	119	+#include <dgl/array.h>
	120	+#include <dgl/runtime/device_api.h>
	121	+
	122	+using dgl::aten::ScatterAdd;
	123	+using dgl::kDGLAscend;
	124	+
	125	+// feat: NDArray shape [N, dim], dtype float32, ctx=NPU
	126	+// idx : NDArray shape [N], dtype int32, ctx=NPU
	127	+// out : NDArray shape [M, dim], dtype float32, ctx=NPU
	128	+//
	129	+// 调用方必须预清零 out（与 CUDA/CPU 一致，kernel 仅 +=）。
	130	+// dgl::aten::FillZeros(out); // 或上层 F.zeros
	131	+
	132	+ScatterAdd<kDGLAscend, int32_t, float>(feat, idx, out);
	133	+
	134	+// int64 idx 同样支持，host 侧自动转 int32：
	135	+// ScatterAdd<kDGLAscend, int64_t, float>(feat, idx_int64, out);
	136	+```
	137	+
	138	+### Python 上层调用（DGL 通路）
	139	+
	140	+```python
	141	+import torch
	142	+import dgl
	143	+from dgl import backend as F
	144	+
	145	+N, M, dim = 1024, 1024, 64
	146	+feat = torch.randn(N, dim, device="npu", dtype=torch.float32)
	147	+idx = torch.randint(0, M, (N,), device="npu", dtype=torch.int32)
	148	+
	149	+# DGL 内部 F.zeros 预清零 out，再调用 _CAPI_DGLKernelScatterAdd
	150	+out = F.zeros((M, dim), feat.dtype, feat.ctx)
	151	+dgl.backend.scatter_add(feat, idx, M, out=out)
	152	+# out[i] == sum(feat[j] for j where idx[j]==i), 累加顺序非确定
	153	+```
	154	+
	155	+## 实现要点
	156	+
	157	+> 本节为使用者提供实现概览，便于理解性能与确定性特征。
	158	+
	159	+- 累加策略（Strategy A）：`SetAtomicAdd<float>()` 开启 GM 原子累加 → 分批 `DataCopyPad` 搬 `feat` 行到 UB → `DataCopyPad` UB→GM（在原子模式下执行 atomic_add）→ `SetAtomicNone()` 复位。与 CUDA `atomicAdd` 语义逐字一致。
	160	+- 多核切分：按 `N`（feat 行数）切分到各核，`blockDim = min(N, 40)`，核 `c` 负责 feat 行 `[c×ceil(N/blockDim), min((c+1)×ceil(N/blockDim), N))`。各核写出地址可能重叠（同一目标行被多行累加），由 `SetAtomicAdd` 保证原子正确性。
	161	+- idx 数据流：idx 一次性预搬 UB（`DataCopyPad` + `TBuf`），`PipeBarrier<PIPE_ALL>` 显式同步 MTE2→Scalar 后，主循环内 `LocalTensor.GetValue` 标量读 idx 值（生产代码禁用 `GlobalTensor::GetValue` 红线 API，改用 UB 预搬 + LocalTensor 读）。
	162	+- 流水线同步：feat 采用 Double Buffer（`BUFFER_NUM=2`）。每批 `FreeTensor` 前插入 `PipeBarrier<PIPE_ALL>`（多队列汇聚点）排空在途 MTE3 写出，防止 Double Buffer 复用覆盖在途原子加源数据；feat `DeQue` 后用 `SetFlag/WaitFlag<HardEvent::MTE2_MTE3>` 事件同步跨 pipe 依赖（MTE2 搬入完成才能 MTE3 原子写出）。
	163	+- 原子态复位：`Process()` 末尾 `PipeBarrier<PIPE_ALL>` 排空本核 MTE3 后调用 `SetAtomicNone()`，复位设备原子加模式寄存器，杜绝跨 launch 泄漏。
	164	+
	165	+## 精度标准
	166	+
	167	+\| 项目 \| 内容 \|
	168	+\|------\|------\|
	169	+\| 输出 dtype \| float32 \|
	170	+\| 判定方式 \| 混合容差（非逐比特），逐元素 `\|actual - golden\| ≤ atol + rtol × \|golden\|` \|
	171	+\| rtol \| 2^-10 ≈ 9.77e-4 \|
	172	+\| atol \| 2^-16 ≈ 1.53e-5 \|
	173	+\| required_matched_ratio \| ≥ 0.99 \|
	174	+\| max_abs_error_limit \| 1e-2 \|
	175	+\| golden 来源 \| CPU `cpu::ScatterAdd`（`src/array/cpu/segment_reduce.h`），float32 `+=` 累加 \|
	176	+\| 精度差异来源 \| 累加顺序（硬件原子调度 vs CPU OpenMP 线程调度），ULP 级，远小于容差阈值 \|
	177	+\| 验收结论 \| 42/42 精度用例 PASS（含多核 + 单核，确定性验证通过）。不要求 bit-exact。 \|
	178	+
	179	+## 性能特征
	180	+
	181	+\| 项目 \| 内容 \|
	182	+\|------\|------\|
	183	+\| 对比基线 \| CPU 回退路径（D2H + `cpu::ScatterAdd` + H2D）端到端时延 \|
	184	+\| 加速比范围 \| 17x – 44x（代表性 shape，中位数口径，详见 `tests/ascend/perf_output/scatter_add/scatter_add_perf.md`） \|
	185	+\| 典型 case \| int32×float32，N=4096/dim=128/M=4096：原生 472 µs vs CPU 回退 20833 µs → 44.1x；N=1024/dim=16/M=1：511 µs vs 8878 µs → 17.4x \|
	186	+\| 瓶颈维度 \| host dispatch/launch（scalar 占比 ~90-99%），device kernel 计算占比 < 10%。AICore 利用率 64%-97%。 \|
	187	+\| 验收结论 \| 性能验收通过（17x-44x 加速比，全部用例不劣于 CPU 回退）。 \|
	188	+
	189	+## 支持芯片
	190	+
	191	+\| 项目 \| 内容 \|
	192	+\|------\|------\|
	193	+\| 芯片型号 \| Ascend 910B3（dav-2201，c220 微架构） \|
	194	+\| `--npu-arch` \| `ascend910b` \|
	195	+\| `NpuArch` \| 2201 \|
	196	+\| 仓库默认 SOC_VERSION \| Ascend910B4（`CMakeLists.txt`），运行时检测为 910B3 / dav-2201 \|
	197	+\| CANN 版本 \| 9.0.0（Ascend-cann-toolkit 9.0.0） \|
	198	+\| 代码架构 \| SIMD / MemBase（`kernel_operator.h`），不支持 SIMT / RegBase / Cube（dav-2201 无这些能力，且本算子非 Matmul 类） \|
	199	+
	200	+## 修订记录
	201	+
	202	+\| 版本 \| 日期 \| 修改内容 \|
	203	+\|------\|------\|----------\|
	204	+\| v1.0 \| 2026-08-26 \| 首版算子文档。接口、参数、约束、dtype、精度、性能、示例齐全。 \|

Asrc/array/ascend/scatter_add.cc+182-0

@@ -0,0 +1,182 @@
	1	+/*
	2	+ * @file src/array/ascend/scatter_add.cc
	3	+ * @brief Host launcher + template specializations for Scatter Add on Ascend.
	4	+ *
	5	+ * Mirrors segment_reduce.cc:
	6	+ * - aclrtlaunch_scatter_add: tiling malloc'd on GM + aclrtMemcpy H2D +
	7	+ * launch + aclrtSynchronizeStream + aclrtFree.
	8	+ * - int64 idx -> int32 conversion via CopyTo(CPU) + AsNumBits(32) + CopyTo(NPU),
	9	+ * identical to segment_reduce.cc:196-199.
	10	+ * - dtype coverage: int32/int64 x float32 supported; half(uint16_t)/double
	11	+ * LOG(FATAL) placeholders (matches segment_reduce.cc specializations).
	12	+ *
	13	+ * NOTE on BFloat16: under the non-CUDA build (USE_CUDA=OFF), the
	14	+ * ATEN_FLOAT_TYPE_SWITCH_16BITS macro LOG(FATAL)s Ascend+BFloat16 before it
	15	+ * reaches any specialization, so no BFloat16 specialization is required here,
	16	+ * consistent with segment_reduce.cc.
	17	+ */
	18	+#include <dgl/array.h>
	19	+#include "../kernel_decl.h"
	20	+
	21	+#include <algorithm>
	22	+#include <cstdint>
	23	+#include <string>
	24	+#include <vector>
	25	+
	26	+#ifdef DGL_USE_ASCEND
	27	+#include <acl/acl.h>
	28	+#include <acl/acl_rt.h>
	29	+#define ASCEND_CALL(func) \
	30	+ { \
	31	+ aclError e = (func); \
	32	+ CHECK(e == ACL_SUCCESS) << "Ascend Error, code: " << e; \
	33	+ }
	34	+
	35	+struct ScatterAddTilingData {
	36	+ uint32_t N; // feat rows
	37	+ uint32_t M; // out rows
	38	+ uint32_t featDim; // elements per row
	39	+};
	40	+
	41	+#ifndef ACLRT_LAUNCH_KERNEL
	42	+#define ACLRT_LAUNCH_KERNEL(kernel_func) aclrtlaunch_##kernel_func
	43	+#endif
	44	+
	45	+// The ACLRT_LAUNCH_KERNEL macro expands to aclrtlaunch_<kernel_func_name>.
	46	+// The device kernel function is named `kernel_scatter_add` (see
	47	+// scatter_add_kernel.cpp), so the host-side launcher symbol is
	48	+// `aclrtlaunch_kernel_scatter_add`.
	49	+extern "C" uint32_t aclrtlaunch_kernel_scatter_add(
	50	+ uint32_t blockDim, aclrtStream stream, void* feat, void* idx,
	51	+ void* out, void* tiling);
	52	+#endif // DGL_USE_ASCEND
	53	+
	54	+namespace dgl {
	55	+namespace aten {
	56	+
	57	+#ifdef DGL_USE_ASCEND
	58	+namespace {
	59	+
	60	+// blockDim = min(N, 40), consistent with segment_reduce.cc:75-76.
	61	+template <typename IdType, typename DType>
	62	+void LaunchScatterAddKernel(
	63	+ int64_t num_items, int64_t num_segments, int64_t feat_dim,
	64	+ const void* feat_ptr, const void* idx_ptr, void* out_ptr,
	65	+ aclrtStream stream) {
	66	+ uint32_t block_dim = static_cast<uint32_t>(
	67	+ num_items > 0 ? std::min<int64_t>(num_items, 40) : 1);
	68	+
	69	+ ScatterAddTilingData tiling = {
	70	+ static_cast<uint32_t>(num_items),
	71	+ static_cast<uint32_t>(num_segments),
	72	+ static_cast<uint32_t>(feat_dim)};
	73	+
	74	+ void* tiling_device = nullptr;
	75	+ ASCEND_CALL(aclrtMalloc(
	76	+ &tiling_device, sizeof(ScatterAddTilingData),
	77	+ ACL_MEM_MALLOC_HUGE_FIRST));
	78	+ ASCEND_CALL(aclrtMemcpy(
	79	+ tiling_device, sizeof(ScatterAddTilingData), &tiling,
	80	+ sizeof(ScatterAddTilingData), ACL_MEMCPY_HOST_TO_DEVICE));
	81	+
	82	+ aclError launch_err = aclrtlaunch_kernel_scatter_add(
	83	+ block_dim, stream, const_cast<void*>(feat_ptr),
	84	+ const_cast<void*>(idx_ptr), out_ptr, tiling_device);
	85	+ if (launch_err != ACL_SUCCESS) {
	86	+ ASCEND_CALL(aclrtFree(tiling_device));
	87	+ LOG(FATAL) << "ScatterAdd kernel launch failed with error code: "
	88	+ << launch_err;
	89	+ }
	90	+
	91	+ ASCEND_CALL(aclrtSynchronizeStream(stream));
	92	+ ASCEND_CALL(aclrtFree(tiling_device));
	93	+}
	94	+
	95	+// Shared implementation for int32 idx x float feat.
	96	+template <typename IdType, typename DType>
	97	+void ScatterAddAscendImpl(NDArray feat, NDArray idx, NDArray out) {
	98	+ DGLContext ctx = feat->ctx;
	99	+ ASCEND_CALL(
	100	+ aclrtSynchronizeDevice()); // ensure PyTorch NPU ops complete before
	101	+ // touching feat/idx/out
	102	+ ASCEND_CALL(aclrtSetDevice(ctx.device_id));
	103	+
	104	+ int64_t num_items = feat->shape[0];
	105	+ int64_t num_segments = out->shape[0];
	106	+ int64_t feat_dim = 1;
	107	+ for (int i = 1; i < out->ndim; ++i) feat_dim *= out->shape[i];
	108	+
	109	+ // Use the default ACL stream (nullptr). Correctness is ensured by the
	110	+ // aclrtSynchronizeDevice() above, not by stream alignment. Matches
	111	+ // segment_reduce.cc:139.
	112	+ aclrtStream stream = nullptr;
	113	+
	114	+ LaunchScatterAddKernel<IdType, DType>(
	115	+ num_items, num_segments, feat_dim, feat->data, idx->data, out->data,
	116	+ stream);
	117	+}
	118	+
	119	+} // namespace
	120	+#endif // DGL_USE_ASCEND
	121	+
	122	+// ---- Supported: int32/int64 x float32 ----
	123	+
	124	+template <>
	125	+void ScatterAdd<kDGLAscend, int32_t, float>(
	126	+ NDArray feat, NDArray idx, NDArray out) {
	127	+#ifdef DGL_USE_ASCEND
	128	+ ScatterAddAscendImpl<int32_t, float>(feat, idx, out);
	129	+#else
	130	+ LOG(FATAL) << "Ascend support is not compiled. Please compile with "
	131	+ "-DUSE_ASCEND=ON";
	132	+#endif
	133	+}
	134	+
	135	+template <>
	136	+void ScatterAdd<kDGLAscend, int64_t, float>(
	137	+ NDArray feat, NDArray idx, NDArray out) {
	138	+#ifdef DGL_USE_ASCEND
	139	+ // int64 idx -> int32 (AsNumBits is a width-preserving cast, not truncation).
	140	+ // Same pattern as segment_reduce.cc:196-199.
	141	+ NDArray idx_cpu = idx.CopyTo(DGLContext{kDGLCPU, 0});
	142	+ IdArray idx_i32_cpu = aten::AsNumBits(idx_cpu, 32);
	143	+ NDArray idx_i32 = idx_i32_cpu.CopyTo(idx->ctx);
	144	+ ScatterAddAscendImpl<int32_t, float>(feat, idx_i32, out);
	145	+#else
	146	+ LOG(FATAL) << "Ascend support is not compiled. Please compile with "
	147	+ "-DUSE_ASCEND=ON";
	148	+#endif
	149	+}
	150	+
	151	+// ---- Unsupported dtypes: half(uint16_t) / double placeholders ----
	152	+
	153	+template <>
	154	+void ScatterAdd<kDGLAscend, int32_t, uint16_t>(
	155	+ NDArray feat, NDArray idx, NDArray out) {
	156	+ LOG(FATAL) << "Current Ascend scatter_add kernel only supports float "
	157	+ "features.";
	158	+}
	159	+
	160	+template <>
	161	+void ScatterAdd<kDGLAscend, int64_t, uint16_t>(
	162	+ NDArray feat, NDArray idx, NDArray out) {
	163	+ LOG(FATAL) << "Current Ascend scatter_add kernel only supports float "
	164	+ "features.";
	165	+}
	166	+
	167	+template <>
	168	+void ScatterAdd<kDGLAscend, int32_t, double>(
	169	+ NDArray feat, NDArray idx, NDArray out) {
	170	+ LOG(FATAL) << "Current Ascend scatter_add kernel only supports float "
	171	+ "features.";
	172	+}
	173	+
	174	+template <>
	175	+void ScatterAdd<kDGLAscend, int64_t, double>(
	176	+ NDArray feat, NDArray idx, NDArray out) {
	177	+ LOG(FATAL) << "Current Ascend scatter_add kernel only supports float "
	178	+ "features.";
	179	+}
	180	+
	181	+} // namespace aten
	182	+} // namespace dgl

Asrc/array/ascend/scatter_add_kernel.cpp+187-0

@@ -0,0 +1,187 @@
	1	+/*
	2	+ * @file src/array/ascend/scatter_add_kernel.cpp
	3	+ * @brief Ascend C device kernel for Scatter Add on first dimension.
	4	+ *
	5	+ * Math: out[idx[i], ] += feat[i, ]
	6	+ *
	7	+ * - SIMD / MemBase paradigm, mirroring segment_reduce_sum_kernel.cpp.
	8	+ * - Accumulation: SetAtomicAdd<float>() + DataCopyPad UB->GM (atomic +=),
	9	+ * kernel does NOT zero out; caller pre-zeros out (matches CPU/CUDA contract).
	10	+ * - SetAtomicAdd / SetAtomicNone paired to manage the device-persistent atomic
	11	+ * mode register lifecycle (SetAtomicAdd is per-MTE3-pipe / per-core; the
	12	+ * reset after this core's MTE3 is drained does not race with other cores'
	13	+ * in-flight atomic writes).
	14	+ * - Multi-core split over N (feat rows); blockDim = min(N, 40) set on host.
	15	+ *
	16	+ * Key synchronization mechanisms:
	17	+ * - idx pre-loaded to UB via TBuf + DataCopyPad, followed by
	18	+ * PipeBarrier<PIPE_ALL> for explicit MTE2->S (Scalar) synchronization
	19	+ * (one-shot load outside the main loop, non-hot-path).
	20	+ * - In ProcessBatch, after feat DeQue, SetFlag/WaitFlag<HardEvent::MTE2_MTE3>
	21	+ * event synchronization (hot-path single cross-pipe dependency; Set enqueued
	22	+ * on MTE2 after DataCopyPad, Wait enqueued on MTE3 before DataCopyPad so MTE3
	23	+ * blocks until MTE2 completes — covers the dependency that EnQue/DeQue's
	24	+ * MTE2->V-only sync does not).
	25	+ * - Before FreeTensor in ProcessBatch, PipeBarrier<PIPE_ALL> for multi-queue
	26	+ * convergence (MTE3 atomic write + V pipe TQue internal events + MTE2 buffer
	27	+ * reuse). PIPE_ALL is the only reliable option here; narrower
	28	+ * PipeBarrier<PIPE_MTE3> or SetFlag/WaitFlag<MTE3_MTE2> exhibit intermittent
	29	+ * races (V pipe TQue internal events not drained causes AllocTensor
	30	+ * wait_flag<V_MTE2> to return early).
	31	+ * - At the end of Process(), PipeBarrier<PIPE_ALL> followed by SetAtomicNone()
	32	+ * to reset the device atomic-add mode (outside the main loop, non-hot-path).
	33	+ */
	34	+#include "kernel_operator.h"
	35	+
	36	+constexpr uint32_t BUFFER_NUM = 2;
	37	+constexpr uint32_t UB_TOTAL_SIZE = 192 * 1024;
	38	+constexpr uint32_t UB_RESERVED = 2048;
	39	+constexpr uint32_t UB_AVAILABLE = UB_TOTAL_SIZE - UB_RESERVED;
	40	+constexpr uint32_t BYTE_ALIGN = 32;
	41	+constexpr uint32_t BLOCK_COUNT_MAX = 4095;
	42	+
	43	+class KernelScatterAdd {
	44	+public:
	45	+ __aicore__ inline void Init(GM_ADDR feat, GM_ADDR idx, GM_ADDR output,
	46	+ GM_ADDR tiling_ptr) {
	47	+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
	48	+
	49	+ AscendC::GlobalTensor<uint32_t> tilingGm;
	50	+ tilingGm.SetGlobalBuffer((__gm__ uint32_t*)tiling_ptr, 3);
	51	+ this->N = tilingGm.GetValue(0);
	52	+ this->M = tilingGm.GetValue(1);
	53	+ this->featDim = tilingGm.GetValue(2);
	54	+
	55	+ this->rowBytes = this->featDim * sizeof(float);
	56	+ uint32_t rowAlignedBytes =
	57	+ (this->rowBytes + BYTE_ALIGN - 1) / BYTE_ALIGN * BYTE_ALIGN;
	58	+ this->rowAlignedElems = rowAlignedBytes / sizeof(float);
	59	+ this->rightPadding = this->rowAlignedElems - this->featDim;
	60	+
	61	+ uint32_t blockIdx = AscendC::GetBlockIdx();
	62	+ uint32_t blockNum = AscendC::GetBlockNum();
	63	+ this->rowsPerCore = (this->N + blockNum - 1) / blockNum;
	64	+ this->startRow = blockIdx * this->rowsPerCore;
	65	+ this->endRow = (this->startRow + this->rowsPerCore > this->N)
	66	+ ? this->N
	67	+ : (this->startRow + this->rowsPerCore);
	68	+
	69	+ idxGm.SetGlobalBuffer((__gm__ uint32_t*)idx, this->N);
	70	+ featGm.SetGlobalBuffer((__gm__ float)feat, this->N this->featDim);
	71	+ outputGm.SetGlobalBuffer((__gm__ float)output, this->M this->featDim);
	72	+
	73	+ uint32_t myRows = (this->endRow > this->startRow)
	74	+ ? (this->endRow - this->startRow)
	75	+ : 0;
	76	+ this->myRows = myRows;
	77	+ uint32_t idxRowsForBuf = myRows > 0 ? myRows : 1;
	78	+ uint32_t idxAlignedBytes =
	79	+ ((idxRowsForBuf * sizeof(uint32_t) + BYTE_ALIGN - 1) / BYTE_ALIGN)
	80	+ * BYTE_ALIGN;
	81	+
	82	+ uint32_t ubForFeat = (UB_AVAILABLE > idxAlignedBytes)
	83	+ ? (UB_AVAILABLE - idxAlignedBytes)
	84	+ : 0;
	85	+ this->batchItems = ubForFeat / (BUFFER_NUM * rowAlignedBytes);
	86	+ if (this->batchItems == 0) {
	87	+ this->batchItems = 1;
	88	+ }
	89	+ if (this->batchItems > BLOCK_COUNT_MAX) {
	90	+ this->batchItems = BLOCK_COUNT_MAX;
	91	+ }
	92	+
	93	+ pipe.InitBuffer(featQueue, BUFFER_NUM,
	94	+ this->batchItems * rowAlignedBytes);
	95	+ pipe.InitBuffer(idxBuf, idxAlignedBytes);
	96	+
	97	+ // Event ID for MTE2->MTE3 cross-pipe sync.
	98	+ evtMte2ToMte3_ = static_cast<int32_t>(
	99	+ pipe.AllocEventID<AscendC::HardEvent::MTE2_MTE3>());
	100	+ }
	101	+
	102	+ __aicore__ inline void Process() {
	103	+ if (this->myRows == 0) {
	104	+ return;
	105	+ }
	106	+ AscendC::LocalTensor<uint32_t> idxLocal = idxBuf.Get<uint32_t>();
	107	+ AscendC::DataCopyExtParams idxParams = {
	108	+ 1, this->myRows * (uint32_t)sizeof(uint32_t), 0, 0, 0};
	109	+ AscendC::DataCopyPadExtParams<uint32_t> idxPad = {false, 0, 0, 0};
	110	+ AscendC::DataCopyPad<uint32_t>(
	111	+ idxLocal, idxGm[this->startRow], idxParams, idxPad);
	112	+ // MTE2->S sync: ensure idx DataCopyPad completes before scalar GetValue.
	113	+ AscendC::PipeBarrier<PIPE_ALL>();
	114	+
	115	+ AscendC::SetAtomicAdd<float>();
	116	+ for (uint32_t row = this->startRow; row < this->endRow;
	117	+ row += this->batchItems) {
	118	+ uint32_t remaining = this->endRow - row;
	119	+ uint32_t itemCount = remaining > this->batchItems
	120	+ ? this->batchItems
	121	+ : remaining;
	122	+ ProcessBatch(row, itemCount, idxLocal);
	123	+ }
	124	+ // Drain MTE3 before resetting atomic mode.
	125	+ AscendC::PipeBarrier<PIPE_ALL>();
	126	+ AscendC::SetAtomicNone();
	127	+
	128	+ pipe.ReleaseEventID<AscendC::HardEvent::MTE2_MTE3>(evtMte2ToMte3_);
	129	+ }
	130	+
	131	+private:
	132	+ __aicore__ inline void ProcessBatch(
	133	+ uint32_t startRow, uint32_t itemCount,
	134	+ const AscendC::LocalTensor<uint32_t>& idxLocal) {
	135	+ AscendC::LocalTensor<float> featBatch = featQueue.AllocTensor<float>();
	136	+
	137	+ AscendC::DataCopyExtParams copyParams = {
	138	+ (uint16_t)itemCount, this->rowBytes, 0, 0, 0};
	139	+ AscendC::DataCopyPadExtParams<float> padParams = {
	140	+ true, 0, (uint8_t)this->rightPadding, 0.0f};
	141	+ AscendC::DataCopyPad<float>(
	142	+ featBatch, featGm[startRow * this->featDim], copyParams, padParams);
	143	+ featQueue.EnQue(featBatch);
	144	+
	145	+ AscendC::LocalTensor<float> featBuf = featQueue.DeQue<float>();
	146	+ // MTE2->MTE3 sync: MTE3 atomic writes wait for MTE2 source load.
	147	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE3>(evtMte2ToMte3_);
	148	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(evtMte2ToMte3_);
	149	+ for (uint32_t i = 0; i < itemCount; ++i) {
	150	+ uint32_t writeRow = idxLocal.GetValue(startRow - this->startRow + i);
	151	+ AscendC::DataCopyExtParams outParams = {1, this->rowBytes, 0, 0, 0};
	152	+ AscendC::DataCopyPad<float>(
	153	+ outputGm[writeRow * this->featDim],
	154	+ featBuf[i * this->rowAlignedElems], outParams);
	155	+ }
	156	+ // Multi-queue convergence: drain MTE3+V before FreeTensor.
	157	+ AscendC::PipeBarrier<PIPE_ALL>();
	158	+ featQueue.FreeTensor(featBuf);
	159	+ }
	160	+
	161	+private:
	162	+ AscendC::TPipe pipe;
	163	+ AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> featQueue;
	164	+ AscendC::TBuf<AscendC::TPosition::VECCALC> idxBuf;
	165	+ AscendC::GlobalTensor<uint32_t> idxGm;
	166	+ AscendC::GlobalTensor<float> featGm;
	167	+ AscendC::GlobalTensor<float> outputGm;
	168	+ int32_t evtMte2ToMte3_; // MTE2->MTE3 cross-pipe event ID.
	169	+ uint32_t N;
	170	+ uint32_t M;
	171	+ uint32_t featDim;
	172	+ uint32_t rowBytes;
	173	+ uint32_t rowAlignedElems;
	174	+ uint32_t rightPadding;
	175	+ uint32_t batchItems;
	176	+ uint32_t rowsPerCore;
	177	+ uint32_t startRow;
	178	+ uint32_t endRow;
	179	+ uint32_t myRows;
	180	+};
	181	+
	182	+extern "C" __global__ __aicore__ void kernel_scatter_add(
	183	+ GM_ADDR feat, GM_ADDR idx, GM_ADDR output, GM_ADDR tiling_ptr) {
	184	+ KernelScatterAdd op;
	185	+ op.Init(feat, idx, output, tiling_ptr);
	186	+ op.Process();
	187	+}

Msrc/array/kernel.cc+5-15

	@@ -401,22 +401,12 @@ void SegmentReduceDispatch(
401	});	401	});
402	}	402	}
403		403
404	-/** @brief Scatter Add (on first dimension) dispatch function. */	404	+/** @brief Scatter Add (on first dimension) dispatch function.
		405	+ *
		406	+ * Native NPU kernel for Ascend; replaces former CPU fallback path.
		407	+ */
405	void ScatterAddDispatch(NDArray feat, NDArray idx, NDArray out) {	408	void ScatterAddDispatch(NDArray feat, NDArray idx, NDArray out) {
406	- if (feat->ctx.device_type == kDGLAscend) {	409	+ ATEN_XPU_SWITCH_CUDA_ASCEND(feat->ctx.device_type, XPU, "ScatterAdd", {
407	- DGLContext cpu_ctx{kDGLCPU, 0};
408	- NDArray feat_cpu = feat.CopyTo(cpu_ctx);
409	- NDArray idx_cpu = idx.CopyTo(cpu_ctx);
410	- NDArray out_cpu = out.CopyTo(cpu_ctx);
411	- ATEN_ID_TYPE_SWITCH(idx_cpu->dtype, IdType, {
412	- ATEN_FLOAT_TYPE_SWITCH_16BITS(feat_cpu->dtype, Dtype, kDGLCPU, "Feature data", {
413	- ScatterAdd<kDGLCPU, IdType, Dtype>(feat_cpu, idx_cpu, out_cpu);
414	- });
415	- });
416	- out_cpu.CopyTo(out);
417	- return;
418	- }
419	- ATEN_XPU_SWITCH_CUDA(feat->ctx.device_type, XPU, "ScatterAdd", {
420	ATEN_ID_TYPE_SWITCH(idx->dtype, IdType, {	410	ATEN_ID_TYPE_SWITCH(idx->dtype, IdType, {
421	ATEN_FLOAT_TYPE_SWITCH_16BITS(feat->dtype, Dtype, XPU, "Feature data", {	411	ATEN_FLOAT_TYPE_SWITCH_16BITS(feat->dtype, Dtype, XPU, "Feature data", {
422	ScatterAdd<XPU, IdType, Dtype>(feat, idx, out);	412	ScatterAdd<XPU, IdType, Dtype>(feat, idx, out);

Atests/ascend/scatter_add_perf.py+653-0

@@ -0,0 +1,653 @@
	1	+"""
	2	+Performance collection framework for scatter_add on Ascend NPU.
	3	+
	4	+Produces:
	5	+ - multi-round wall-time collection (warmup-excluded; >=30 rounds; records
	6	+ median / min / max / p90 / mean / cv for volatility handling)
	7	+ - bottleneck dimension breakdown: compute (device kernel time, extracted via
	8	+ torch_npu.profiler `kernel_details.csv`) / H2D (data copy-in) / D2H (data
	9	+ copy-out) / scalar (host launch + dispatch + sync residual) -- as us and as
	10	+ shares of the op wall time
	11	+ - CPU fallback baseline (D2H + cpu::ScatterAdd + H2D) end-to-end, for direct
	12	+ native-vs-fallback comparison
	13	+ - JSON + JSONL + markdown report
	14	+
	15	+Coverage (int32/int64 x float32):
	16	+ - L1-8 : N=1024 dim=256 int32 (aligned large, compute-bound)
	17	+ - L1-13 : N=4096 dim=128 int32 (GNN typical scale)
	18	+ - L1-5 : N=1024 dim=1 int32 (narrow, scalar/launch-bound)
	19	+ - L1-6 : N=64 dim=17 int32 (non-aligned, small)
	20	+ - L1-11 : N=1024 dim=16 int32 (all hit same row, atomic-conflict bound)
	21	+ - L1-10 : N=64 dim=16 int64 (IdType=int64 representative)
	22	+ - PERF-I64-L : N=4096 dim=128 int64 (IdType=int64 large scale — covers
	23	+ int64 × float32 at representative GNN scale)
	24	+ - L2-3 : N=1024 dim=16 int32 (M=1 all hit unique row)
	25	+
	26	+Usage (standalone):
	27	+
	28	+ PYTHONPATH=python python tests/ascend/scatter_add_perf.py \
	29	+ --warmup 5 --rounds 30 --device 4
	30	+
	31	+The profiler path (compute/scalar split) needs torch_npu.profiler; if it is
	32	+unavailable the framework falls back to an isolated-timing estimate but flags
	33	+the breakdown as heuristic.
	34	+"""
	35	+from __future__ import annotations
	36	+
	37	+import argparse
	38	+import csv
	39	+import glob
	40	+import json
	41	+import os
	42	+import shutil
	43	+import statistics
	44	+import sys
	45	+import time
	46	+import warnings
	47	+from dataclasses import dataclass, field
	48	+from typing import Dict, List, Optional, Tuple
	49	+
	50	+# Allow sibling imports.
	51	+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
	52	+
	53	+import torch # noqa: E402
	54	+
	55	+import dgl.backend as F # noqa: E402
	56	+
	57	+from scatter_add_golden import CASES # noqa: E402
	58	+
	59	+
	60	+# Final perf data lands in the test-dir collection output.
	61	+TEST_DIR = os.path.dirname(os.path.abspath(__file__))
	62	+PERF_RESULTS_DIR = os.environ.get(
	63	+ "SCATTER_ADD_PERF_RESULTS_DIR",
	64	+ os.path.join(TEST_DIR, "perf_output", "scatter_add"),
	65	+)
	66	+PROF_SCRATCH_DIR = os.environ.get(
	67	+ "SCATTER_ADD_PROF_SCRATCH_DIR",
	68	+ os.path.join(TEST_DIR, "prof_scratch"),
	69	+)
	70	+
	71	+# A large-scale int64 case not present in the golden table (perf-only).
	72	+def _perf_i64_large():
	73	+ n = 4096
	74	+ feat = torch.randn(n, 128, dtype=torch.float32)
	75	+ idx = torch.randint(0, n, (n,), dtype=torch.int64)
	76	+ return feat, idx, n, "int64"
	77	+
	78	+
	79	+# Ordered perf-relevant case set covering int32/int64 × float32.
	80	+# (case_id, builder-or-None). None ⇒ resolve from CASES table.
	81	+PERF_CASES: List[Tuple[str, Optional[callable]]] = [
	82	+ ("L1-8", None),
	83	+ ("L1-13", None),
	84	+ ("L1-5", None),
	85	+ ("L1-6", None),
	86	+ ("L1-11", None),
	87	+ ("L1-10", None),
	88	+ ("PERF-I64-L", _perf_i64_large),
	89	+ ("L2-3", None),
	90	+]
	91	+
	92	+# Kernel name substrings matched in kernel_details.csv for device-time split.
	93	+KERNEL_SCATTER = "kernel_scatter_add"
	94	+KERNEL_ZEROS = "aclnnInplaceZero" # F.zeros memset kernel (part of op)
	95	+
	96	+
	97	+@dataclass
	98	+class PerfRecord:
	99	+ case_id: str
	100	+ level: str
	101	+ desc: str
	102	+ shape: Dict[str, int]
	103	+ idtype: str
	104	+ rounds: int
	105	+ # Native NPU op wall-time (warmup excluded), microseconds.
	106	+ op_us_median: float
	107	+ op_us_std: float
	108	+ op_us_min: float
	109	+ op_us_max: float
	110	+ op_us_p50: float
	111	+ op_us_p90: float
	112	+ op_us_mean: float
	113	+ op_us_cv: float
	114	+ # CPU-fallback (D2H+CPU ScatterAdd+H2D) wall-time, microseconds.
	115	+ cpu_fallback_us_median: float = 0.0
	116	+ cpu_fallback_us_min: float = 0.0
	117	+ cpu_fallback_us_max: float = 0.0
	118	+ speedup_vs_cpu: float = 0.0 # cpu_fallback / native op
	119	+ # Bottleneck breakdown (microseconds).
	120	+ compute_us: float = 0.0 # device kernel time (scatter_add + zeros)
	121	+ scatter_kernel_us: float = 0.0
	122	+ zeros_kernel_us: float = 0.0
	123	+ h2d_us: float = 0.0 # isolated input copy-in
	124	+ d2h_us: float = 0.0 # isolated output copy-out
	125	+ scalar_us: float = 0.0 # host launch+dispatch+sync residual
	126	+ # AICore utilization (from op_statistic.csv Ratio%).
	127	+ aicore_util_pct: float = 0.0 # scatter_add kernel AICore utilization %
	128	+ zeros_util_pct: float = 0.0
	129	+ # Shares of op wall time (%).
	130	+ compute_share: float = 0.0
	131	+ h2d_share: float = 0.0
	132	+ d2h_share: float = 0.0
	133	+ scalar_share: float = 0.0
	134	+ samples: List[float] = field(default_factory=list)
	135	+ cpu_samples: List[float] = field(default_factory=list)
	136	+ volatile: bool = False
	137	+ breakdown_method: str = ""
	138	+ notes: str = ""
	139	+
	140	+
	141	+# ----------------------------- device helpers ----------------------------- #
	142	+
	143	+def _npu_available() -> bool:
	144	+ return hasattr(torch, "npu") and torch.npu.is_available()
	145	+
	146	+
	147	+def _get_device(device_id: int) -> torch.device:
	148	+ dev = torch.device(f"npu:{device_id}")
	149	+ torch.npu.set_device(dev)
	150	+ return dev
	151	+
	152	+
	153	+def _sync(dev: torch.device) -> None:
	154	+ if dev.type == "npu":
	155	+ torch.npu.synchronize(dev)
	156	+
	157	+
	158	+def _measure(fn, dev, warmup, rounds) -> List[float]:
	159	+ """Run `fn()` warmup times then `rounds` timed rounds (us per round)."""
	160	+ for _ in range(warmup):
	161	+ fn()
	162	+ _sync(dev)
	163	+ samples = []
	164	+ for _ in range(rounds):
	165	+ _sync(dev)
	166	+ t0 = time.perf_counter()
	167	+ fn()
	168	+ _sync(dev)
	169	+ t1 = time.perf_counter()
	170	+ samples.append((t1 - t0) * 1_000_000.0)
	171	+ return samples
	172	+
	173	+
	174	+def _stats(samples: List[float]) -> Dict[str, float]:
	175	+ s = sorted(samples)
	176	+ n = len(s)
	177	+
	178	+ def _pct(p):
	179	+ if n == 0:
	180	+ return 0.0
	181	+ k = max(0, min(n - 1, int(round((p / 100.0) * (n - 1)))))
	182	+ return s[k]
	183	+
	184	+ median = statistics.median(samples) if samples else 0.0
	185	+ mean = statistics.mean(samples) if samples else 0.0
	186	+ std = statistics.stdev(samples) if len(samples) > 1 else 0.0
	187	+ return {
	188	+ "median": float(median),
	189	+ "mean": float(mean),
	190	+ "std": float(std),
	191	+ "min": float(min(samples)) if samples else 0.0,
	192	+ "max": float(max(samples)) if samples else 0.0,
	193	+ "p50": float(_pct(50)),
	194	+ "p90": float(_pct(90)),
	195	+ "cv": float(std / median) if median > 0 else 0.0,
	196	+ }
	197	+
	198	+
	199	+def _resolve_case(case_id: str) -> Tuple:
	200	+ """Return (feat_cpu, idx_cpu, m, idtype, level, desc)."""
	201	+ for cid, builder in PERF_CASES:
	202	+ if cid == case_id:
	203	+ if builder is not None:
	204	+ feat, idx, m, idtype = builder()
	205	+ return feat, idx, m, idtype, "PERF", "perf int64 large scale"
	206	+ break
	207	+ for c in CASES:
	208	+ if c.case_id == case_id:
	209	+ feat, idx, m, idtype = c.build()
	210	+ return feat, idx, m, idtype, c.level, c.desc
	211	+ raise KeyError(f"case {case_id} not found")
	212	+
	213	+
	214	+# ----------------------------- profiler split ----------------------------- #
	215	+
	216	+def _profiler_kernel_us(dev: torch.device, feat_npu, idx_npu, m: int,
	217	+ case_id: str, out_dir: str) -> Dict[str, float]:
	218	+ """
	219	+ Capture a traced run (Level1 + PipeUtilization) and extract device kernel
	220	+ durations + AICore utilization from op_statistic.csv. Returns:
	221	+ scatter_us, zeros_us, compute_us (scatter+zeros Avg),
	222	+ aicore_util_pct (scatter_add kernel AICore Ratio%),
	223	+ zeros_util_pct.
	224	+
	225	+ op_statistic.csv gives per-op Avg(us) and Ratio(%) over the recorded
	226	+ iterations — more robust than per-step kernel_details medians.
	227	+ """
	228	+ try:
	229	+ import torch_npu.profiler as npp
	230	+ from torch_npu.profiler import (AiCMetrics, ProfilerActivity,
	231	+ ProfilerAction, ProfilerLevel,
	232	+ profile, schedule)
	233	+ from torch_npu.profiler.experimental_config import \
	234	+ _ExperimentalConfig as EC
	235	+ except Exception as e: # noqa: BLE001
	236	+ return {"scatter_us": 0.0, "zeros_us": 0.0, "compute_us": 0.0,
	237	+ "error": f"profiler unavailable: {e}"}
	238	+
	239	+ case_prof_dir = os.path.join(out_dir, case_id)
	240	+ if os.path.exists(case_prof_dir):
	241	+ shutil.rmtree(case_prof_dir, ignore_errors=True)
	242	+ os.makedirs(case_prof_dir, exist_ok=True)
	243	+
	244	+ # Warmup so the traced iterations are steady.
	245	+ for _ in range(3):
	246	+ F.scatter_add(feat_npu, idx_npu, m)
	247	+ _sync(dev)
	248	+
	249	+ cfg = EC(profiler_level=ProfilerLevel.Level1,
	250	+ aic_metrics=AiCMetrics.PipeUtilization)
	251	+
	252	+ def _sch(step):
	253	+ if step < 3:
	254	+ return ProfilerAction.RECORD
	255	+ if step == 3:
	256	+ return ProfilerAction.RECORD_AND_SAVE
	257	+ return ProfilerAction.NONE
	258	+
	259	+ try:
	260	+ with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
	261	+ schedule=_sch, experimental_config=cfg,
	262	+ on_trace_ready=npp.tensorboard_trace_handler(case_prof_dir)) as prof:
	263	+ for i in range(5):
	264	+ F.scatter_add(feat_npu, idx_npu, m)
	265	+ _sync(dev)
	266	+ prof.step()
	267	+ except Exception as e: # noqa: BLE001
	268	+ return {"scatter_us": 0.0, "zeros_us": 0.0, "compute_us": 0.0,
	269	+ "error": f"profile failed: {e}"}
	270	+
	271	+ os_csv = glob.glob(os.path.join(case_prof_dir, "**", "op_statistic.csv"),
	272	+ recursive=True)
	273	+ if not os_csv:
	274	+ return {"scatter_us": 0.0, "zeros_us": 0.0, "compute_us": 0.0,
	275	+ "error": "op_statistic.csv not found"}
	276	+
	277	+ scatter_us = zeros_us = 0.0
	278	+ scatter_util = zeros_util = 0.0
	279	+ with open(os_csv[0]) as f:
	280	+ reader = csv.DictReader(f)
	281	+ for row in reader:
	282	+ op_type = row.get("OP Type", "")
	283	+ try:
	284	+ avg = float(row.get("Avg Time(us)", "0"))
	285	+ except (ValueError, TypeError):
	286	+ avg = 0.0
	287	+ try:
	288	+ ratio = float(row.get("Ratio(%)", "0"))
	289	+ except (ValueError, TypeError):
	290	+ ratio = 0.0
	291	+ if KERNEL_SCATTER in op_type:
	292	+ scatter_us = avg
	293	+ scatter_util = ratio
	294	+ elif "Zero" in op_type or "zero" in op_type:
	295	+ zeros_us = avg
	296	+ zeros_util = ratio
	297	+
	298	+ return {
	299	+ "scatter_us": scatter_us,
	300	+ "zeros_us": zeros_us,
	301	+ "compute_us": scatter_us + zeros_us,
	302	+ "aicore_util_pct": scatter_util,
	303	+ "zeros_util_pct": zeros_util,
	304	+ }
	305	+
	306	+
	307	+# ----------------------------- collection ----------------------------- #
	308	+
	309	+def collect_one(case_id: str, dev: torch.device, warmup: int,
	310	+ rounds: int, prof_dir: str) -> PerfRecord:
	311	+ feat_cpu, idx_cpu, m, idtype, level, desc = _resolve_case(case_id)
	312	+
	313	+ feat_npu = feat_cpu.to(dev)
	314	+ idx_npu = idx_cpu.to(dev)
	315	+ _sync(dev)
	316	+
	317	+ # ---- 1) Native NPU op wall-time (resident inputs) ----
	318	+ def _op_fn():
	319	+ return F.scatter_add(feat_npu, idx_npu, m)
	320	+
	321	+ op_samples = _measure(_op_fn, dev, warmup, rounds)
	322	+ op_stats = _stats(op_samples)
	323	+
	324	+ # ---- 2) CPU fallback: D2H(feat,idx) + cpu::ScatterAdd + H2D(out) ----
	325	+ # End-to-end realistic fallback as it would run if dispatch fell back.
	326	+ def _cpu_fallback_fn():
	327	+ f_c = feat_npu.cpu()
	328	+ i_c = idx_npu.cpu()
	329	+ out_c = F.scatter_add(f_c, i_c, m) # CPU dispatch → cpu::ScatterAdd
	330	+ _ = out_c.to(dev)
	331	+
	332	+ cpu_samples = _measure(_cpu_fallback_fn, dev, warmup=max(2, warmup // 2),
	333	+ rounds=max(10, rounds // 2))
	334	+ cpu_stats = _stats(cpu_samples)
	335	+
	336	+ # ---- 3) H2D (both inputs) ----
	337	+ def _h2d_fn():
	338	+ _ = feat_cpu.to(dev)
	339	+ _ = idx_cpu.to(dev)
	340	+
	341	+ h2d_samples = _measure(_h2d_fn, dev, warmup=max(1, warmup // 2),
	342	+ rounds=max(10, rounds // 2))
	343	+ h2d_stats = _stats(h2d_samples)
	344	+
	345	+ # ---- 4) D2H (op output) ----
	346	+ out_ref = F.scatter_add(feat_npu, idx_npu, m)
	347	+ _sync(dev)
	348	+
	349	+ def _d2h_fn():
	350	+ _ = out_ref.cpu()
	351	+
	352	+ d2h_samples = _measure(_d2h_fn, dev, warmup=max(1, warmup // 2),
	353	+ rounds=max(10, rounds // 2))
	354	+ d2h_stats = _stats(d2h_samples)
	355	+
	356	+ # ---- 5) compute/scalar split via profiler ----
	357	+ prof = _profiler_kernel_us(dev, feat_npu, idx_npu, m, case_id, prof_dir)
	358	+ aicore_util = prof.get("aicore_util_pct", 0.0)
	359	+ zeros_util = prof.get("zeros_util_pct", 0.0)
	360	+ if prof.get("compute_us", 0.0) > 0 and "error" not in prof:
	361	+ compute_us = prof["compute_us"]
	362	+ scatter_us = prof["scatter_us"]
	363	+ zeros_us = prof["zeros_us"]
	364	+ method = ("torch_npu.profiler op_statistic.csv (Avg Duration + "
	365	+ "AICore Ratio%, Level1+PipeUtilization)")
	366	+ else:
	367	+ # Fallback heuristic (flagged): compute ≈ op minus host launch-only.
	368	+ def _launch_only_fn():
	369	+ return F.scatter_add(feat_npu, idx_npu, m)
	370	+ launch_samples = _measure(_launch_only_fn, dev, warmup=max(1, warmup // 2),
	371	+ rounds=max(10, rounds // 2))
	372	+ launch_med = _stats(launch_samples)["median"]
	373	+ compute_us = max(0.0, op_stats["median"] - launch_med)
	374	+ scatter_us = compute_us
	375	+ zeros_us = 0.0
	376	+ method = ("heuristic isolated (op_sync - op_launch_no_sync); "
	377	+ f"profiler unavailable: {prof.get('error','')}")
	378	+
	379	+ op_med = op_stats["median"] or 1e-9
	380	+ scalar_us = max(0.0, op_med - compute_us)
	381	+
	382	+ speedup = (cpu_stats["median"] / op_med) if op_med > 0 else 0.0
	383	+ volatile = op_stats["cv"] > 0.10
	384	+
	385	+ shape_info = {
	386	+ "N": int(feat_cpu.shape[0]),
	387	+ "dim": int(feat_cpu.shape[1]) if feat_cpu.dim() > 1 else 1,
	388	+ "M": int(m),
	389	+ }
	390	+
	391	+ return PerfRecord(
	392	+ case_id=case_id,
	393	+ level=level,
	394	+ desc=desc,
	395	+ shape=shape_info,
	396	+ idtype=idtype,
	397	+ rounds=rounds,
	398	+ op_us_median=op_stats["median"],
	399	+ op_us_std=op_stats["std"],
	400	+ op_us_min=op_stats["min"],
	401	+ op_us_max=op_stats["max"],
	402	+ op_us_p50=op_stats["p50"],
	403	+ op_us_p90=op_stats["p90"],
	404	+ op_us_mean=op_stats["mean"],
	405	+ op_us_cv=op_stats["cv"],
	406	+ cpu_fallback_us_median=cpu_stats["median"],
	407	+ cpu_fallback_us_min=cpu_stats["min"],
	408	+ cpu_fallback_us_max=cpu_stats["max"],
	409	+ speedup_vs_cpu=speedup,
	410	+ compute_us=compute_us,
	411	+ scatter_kernel_us=scatter_us,
	412	+ zeros_kernel_us=zeros_us,
	413	+ h2d_us=h2d_stats["median"],
	414	+ d2h_us=d2h_stats["median"],
	415	+ scalar_us=scalar_us,
	416	+ aicore_util_pct=aicore_util,
	417	+ zeros_util_pct=zeros_util,
	418	+ compute_share=compute_us / op_med,
	419	+ h2d_share=h2d_stats["median"] / op_med,
	420	+ d2h_share=d2h_stats["median"] / op_med,
	421	+ scalar_share=scalar_us / op_med,
	422	+ samples=op_samples,
	423	+ cpu_samples=cpu_samples,
	424	+ volatile=volatile,
	425	+ breakdown_method=method,
	426	+ notes=(
	427	+ "compute=device kernel time (scatter_add+zeros) from torch_npu "
	428	+ "profiler op_statistic.csv (Avg); AICore利用率=scatter_add kernel "
	429	+ "AICore Ratio%; scalar=op_wall−compute (host dispatch+launch+sync); "
	430	+ "h2d=.to(npu) of inputs; d2h=.cpu() of output."
	431	+ ),
	432	+ )
	433	+
	434	+
	435	+# ----------------------------- reporting ----------------------------- #
	436	+
	437	+def _record_to_dict(r: PerfRecord) -> Dict:
	438	+ return {
	439	+ "case_id": r.case_id,
	440	+ "level": r.level,
	441	+ "desc": r.desc,
	442	+ "shape": r.shape,
	443	+ "idtype": r.idtype,
	444	+ "rounds": r.rounds,
	445	+ "op_us": {
	446	+ "median": r.op_us_median,
	447	+ "mean": r.op_us_mean,
	448	+ "std": r.op_us_std,
	449	+ "min": r.op_us_min,
	450	+ "max": r.op_us_max,
	451	+ "p50": r.op_us_p50,
	452	+ "p90": r.op_us_p90,
	453	+ "cv": r.op_us_cv,
	454	+ "samples": r.samples,
	455	+ },
	456	+ "cpu_fallback_us": {
	457	+ "median": r.cpu_fallback_us_median,
	458	+ "min": r.cpu_fallback_us_min,
	459	+ "max": r.cpu_fallback_us_max,
	460	+ "samples": r.cpu_samples,
	461	+ },
	462	+ "speedup_vs_cpu": r.speedup_vs_cpu,
	463	+ "bottleneck_breakdown_us": {
	464	+ "compute": r.compute_us,
	465	+ "scatter_kernel": r.scatter_kernel_us,
	466	+ "zeros_kernel": r.zeros_kernel_us,
	467	+ "h2d": r.h2d_us,
	468	+ "d2h": r.d2h_us,
	469	+ "scalar": r.scalar_us,
	470	+ },
	471	+ "aicore_utilization_pct": {
	472	+ "scatter_add_kernel": r.aicore_util_pct,
	473	+ "zeros_kernel": r.zeros_util_pct,
	474	+ },
	475	+ "bottleneck_share_of_op": {
	476	+ "compute": r.compute_share,
	477	+ "h2d": r.h2d_share,
	478	+ "d2h": r.d2h_share,
	479	+ "scalar": r.scalar_share,
	480	+ },
	481	+ "volatile": r.volatile,
	482	+ "breakdown_method": r.breakdown_method,
	483	+ "notes": r.notes,
	484	+ }
	485	+
	486	+
	487	+def _write_report(records: List[PerfRecord], out_dir: str,
	488	+ warmup: int, rounds: int, device: int) -> None:
	489	+ os.makedirs(out_dir, exist_ok=True)
	490	+ payload = [_record_to_dict(r) for r in records]
	491	+
	492	+ with open(os.path.join(out_dir, "scatter_add_perf.jsonl"), "w") as f:
	493	+ for rec in payload:
	494	+ f.write(json.dumps(rec, default=str) + "\n")
	495	+ with open(os.path.join(out_dir, "scatter_add_perf.json"), "w") as f:
	496	+ json.dump(payload, f, indent=2, default=str)
	497	+
	498	+ lines = [
	499	+ "# scatter_add performance report",
	500	+ "",
	501	+ f"- 设备: NPU {device} (910B3)",
	502	+ f"- 采集口径: 预热 {warmup} 轮剔除; 每用例 {rounds} 轮稳态采集; "
	503	+ f"取值=中位数 (不报单次最优); 每次 op 前后 torch.npu.synchronize "
	504	+ f"排空流水; 设备 NPU{device} 与其他进程共享 (非完全空闲), 波动用例 "
	505	+ f"已做 ≥30 轮并以中位数稳态为准.",
	506	+ f"- dtype 覆盖: int32 / int64 × float32",
	507	+ f"- 瓶颈分解口径: compute=device kernel time (scatter_add + zeros, "
	508	+ f"torch_npu.profiler op_statistic.csv Avg Duration); "
	509	+ f"AICore利用率=scatter_add kernel AICore Ratio% (Level1+PipeUtilization); "
	510	+ f"scalar=op_wall−compute (host dispatch+launch+sync); "
	511	+ f"h2d=输入 .to(npu); d2h=输出 .cpu().",
	512	+ f"- 波动判定: cv (std/median) > 0.10 视为波动, 已记录多轮 min/max/cv.",
	513	+ "",
	514	+ "## 1. 原生 NPU vs CPU 回退 (D2H+cpu::ScatterAdd+H2D)",
	515	+ "",
	516	+ "\| Case \| IdType \| N \| dim \| M \| 原生 op 中位(us) \| CPU回退中位(us) \| 加速比(CPU/原生) \|",
	517	+ "\|------\|--------\|---\|-----\|---\|------------------\|------------------\|------------------\|",
	518	+ ]
	519	+ for r in records:
	520	+ lines.append(
	521	+ f"\| {r.case_id} \| {r.idtype} \| {r.shape['N']} \| {r.shape['dim']} \| "
	522	+ f"{r.shape['M']} \| {r.op_us_median:.3f} \| "
	523	+ f"{r.cpu_fallback_us_median:.3f} \| {r.speedup_vs_cpu:.2f}x \|"
	524	+ )
	525	+
	526	+ lines += [
	527	+ "",
	528	+ "## 2. 瓶颈维度分解",
	529	+ "",
	530	+ "\| Case \| op 中位(us) \| compute(us) \| scatter_kernel(us) \| zeros_kernel(us) \| h2d(us) \| d2h(us) \| scalar(us) \| compute% \| h2d% \| d2h% \| scalar% \| AICore利用率% \| cv \| 轮数 \|",
	531	+ "\|------\|-----------\|-------------\|---------------------\|------------------\|---------\|---------\|------------\|----------\|------\|------\|---------\|-------------\|-----\|------\|",
	532	+ ]
	533	+ for r in records:
	534	+ lines.append(
	535	+ f"\| {r.case_id} \| {r.op_us_median:.3f} \| {r.compute_us:.3f} \| "
	536	+ f"{r.scatter_kernel_us:.3f} \| {r.zeros_kernel_us:.3f} \| "
	537	+ f"{r.h2d_us:.3f} \| {r.d2h_us:.3f} \| {r.scalar_us:.3f} \| "
	538	+ f"{r.compute_share100:.1f} \| {r.h2d_share100:.1f} \| "
	539	+ f"{r.d2h_share100:.1f} \| {r.scalar_share100:.1f} \| "
	540	+ f"{r.aicore_util_pct:.1f} \| {r.op_us_cv:.3f} \| {r.rounds} \|"
	541	+ )
	542	+
	543	+ lines += [
	544	+ "",
	545	+ "## 3. 方差处理",
	546	+ "",
	547	+ "\| Case \| op 中位(us) \| min(us) \| max(us) \| p90(us) \| std(us) \| cv \| 波动? \|",
	548	+ "\|------\|-----------\|---------\|---------\|---------\|---------\|-----\|-------\|",
	549	+ ]
	550	+ for r in records:
	551	+ lines.append(
	552	+ f"\| {r.case_id} \| {r.op_us_median:.3f} \| {r.op_us_min:.3f} \| "
	553	+ f"{r.op_us_max:.3f} \| {r.op_us_p90:.3f} \| {r.op_us_std:.3f} \| "
	554	+ f"{r.op_us_cv:.3f} \| {'是' if r.volatile else '否'} \|"
	555	+ )
	556	+
	557	+ volatile = [r for r in records if r.volatile]
	558	+ lines += [
	559	+ "",
	560	+ "## 4. 瓶颈分解方法说明",
	561	+ "",
	562	+ "compute (设备计算耗时) = torch_npu.profiler op_statistic.csv 中 "
	563	+ "`kernel_scatter_add` 与 `ZerosLike` 两条设备 kernel 的 Avg Time(us) "
	564	+ "之和 (Level1+PipeUtilization, 3 轮记录取平均). 这是 op 在 AICore 上 "
	565	+ "真正的计算耗时, 非启发式估计. AICore利用率 = kernel_scatter_add 的 "
	566	+ "AICore Ratio% (PipeUtilization).",
	567	+ "",
	568	+ "scalar (host 标量开销) = op_wall_median − compute, 即 host 侧 dispatch "
	569	+ "(ATEN_ID_TYPE_SWITCH/ATEN_FLOAT_TYPE_SWITCH) + kernel launch + "
	570	+ "torch.npu.synchronize 的残差开销.",
	571	+ "",
	572	+ "h2d = feat_cpu.to(npu) + idx_cpu.to(npu) 独立采集 (输入搬入).",
	573	+ "d2h = out.cpu() 独立采集 (输出搬出).",
	574	+ "",
	575	+ "> share = 各维度耗时 / op_wall_median. h2d/d2h 为独立采集的参考耗时, "
	576	+ "> 其 share 可能 >100% (因 op 测量时输入已驻留, 不含 h2d). "
	577	+ "> compute+scalar share 之和约为 100%.",
	578	+ ]
	579	+ if volatile:
	580	+ lines += [
	581	+ "",
	582	+ "## 5. 波动用例",
	583	+ "",
	584	+ "以下用例 cv>0.10, 已做 ≥30 轮稳态采集, 以中位数为准:",
	585	+ ]
	586	+ for r in volatile:
	587	+ lines.append(
	588	+ f"- {r.case_id}: 中位 {r.op_us_median:.3f}us, "
	589	+ f"min/max {r.op_us_min:.3f}/{r.op_us_max:.3f}, cv={r.op_us_cv:.3f}"
	590	+ )
	591	+
	592	+ with open(os.path.join(out_dir, "scatter_add_perf.md"), "w") as f:
	593	+ f.write("\n".join(lines) + "\n")
	594	+
	595	+
	596	+def main():
	597	+ parser = argparse.ArgumentParser(
	598	+ description="scatter_add NPU performance collection"
	599	+ )
	600	+ parser.add_argument("--cases", default="",
	601	+ help="逗号分隔 case_id; 留空则用默认 perf 集")
	602	+ parser.add_argument("--warmup", type=int, default=5)
	603	+ parser.add_argument("--rounds", type=int, default=30,
	604	+ help="每用例采集轮数 (>=30)")
	605	+ parser.add_argument("--device", type=int, default=4)
	606	+ args = parser.parse_args()
	607	+
	608	+ if not _npu_available():
	609	+ print("NPU not available; skipping performance collection.")
	610	+ return
	611	+
	612	+ if args.cases:
	613	+ case_ids = [s.strip() for s in args.cases.split(",") if s.strip()]
	614	+ else:
	615	+ case_ids = [cid for cid, _ in PERF_CASES]
	616	+
	617	+ dev = _get_device(args.device)
	618	+ os.makedirs(PROF_SCRATCH_DIR, exist_ok=True)
	619	+
	620	+ records: List[PerfRecord] = []
	621	+ for cid in case_ids:
	622	+ try:
	623	+ rec = collect_one(cid, dev, args.warmup, args.rounds, PROF_SCRATCH_DIR)
	624	+ records.append(rec)
	625	+ print(
	626	+ f"[{cid}] {rec.idtype} N={rec.shape['N']} dim={rec.shape['dim']} "
	627	+ f"M={rec.shape['M']} \| native={rec.op_us_median:.3f}us "
	628	+ f"cv={rec.op_us_cv:.3f} \| cpu_fb={rec.cpu_fallback_us_median:.3f}us "
	629	+ f"speedup={rec.speedup_vs_cpu:.2f}x \| "
	630	+ f"compute={rec.compute_us:.3f}us({rec.compute_share*100:.0f}%) "
	631	+ f"scatter_k={rec.scatter_kernel_us:.3f} "
	632	+ f"h2d={rec.h2d_us:.3f} d2h={rec.d2h_us:.3f} "
	633	+ f"scalar={rec.scalar_us:.3f}us({rec.scalar_share*100:.0f}%)"
	634	+ )
	635	+ except Exception as e: # noqa: BLE001
	636	+ import traceback
	637	+ traceback.print_exc()
	638	+ print(f"[{cid}] collection FAILED: {e}")
	639	+
	640	+ _write_report(records, PERF_RESULTS_DIR, args.warmup, args.rounds, args.device)
	641	+ print(
	642	+ f"\nWrote {len(records)} records to {PERF_RESULTS_DIR}/"
	643	+ f"scatter_add_perf.{{json,jsonl,md}}"
	644	+ )
	645	+
	646	+ volatile = [r for r in records if r.volatile]
	647	+ if volatile:
	648	+ print(f"\nVolatile cases (cv>0.10): {[r.case_id for r in volatile]}")
	649	+ print(" -> 多轮采集已记录中位数/min/max/cv.")
	650	+
	651	+
	652	+if __name__ == "__main__":
	653	+ main()

Atests/ascend/test_scatter_add_npu.py+718-0

@@ -0,0 +1,718 @@
	1	+"""
	2	+Test scatter_add (ScatterAdd) on Ascend NPU.
	3	+
	4	+Tests the Ascend ScatterAdd kernel: out[idx[i], ] += feat[i, ] on the first
	5	+dimension. Mirrors the pytest style of test_segment_reduce_npu.py /
	6	+test_index_select_npu.py.
	7	+
	8	+Coverage: 46 cases — L0(6) + L1(24) + L2(9) + 补-1(1) blackbox + WB(6) whitebox.
	9	+ - L0/L1/补-1/WB: precision cases, asserted via mixed tolerance gate.
	10	+ - L2-1..L2-5: boundary cases.
	11	+ - L2-6..L2-9: dtype LOG(FATAL) exception cases.
	12	+
	13	+Golden: `dgl.backend.scatter_add` CPU dispatch → cpu::ScatterAdd (semantic
	14	+authority). torch.Tensor.scatter_add_ is an auxiliary cross-check only.
	15	+"""
	16	+from __future__ import annotations
	17	+
	18	+import os
	19	+import sys
	20	+from dataclasses import dataclass
	21	+from typing import Callable, Dict, List, Tuple
	22	+
	23	+import numpy as np
	24	+import pytest
	25	+import torch
	26	+
	27	+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
	28	+
	29	+import dgl.backend as F # noqa: E402
	30	+
	31	+
	32	+# ---------------------------------------------------------------------------
	33	+# Precision constants (float32 mixed tolerance).
	34	+# ---------------------------------------------------------------------------
	35	+
	36	+RTOL = 2 ** (-10)
	37	+ATOL = 2 ** (-16)
	38	+REQUIRED_MATCHED_RATIO = 0.99
	39	+MAX_ABS_ERROR_LIMIT = 1e-2
	40	+_ULP32 = 2 ** (-23)
	41	+_MAX_ABS_ERROR_LIMIT_EFF = max(MAX_ABS_ERROR_LIMIT, 32 * _ULP32)
	42	+AUX_RTOL = 1e-4
	43	+AUX_ATOL = 1e-4
	44	+
	45	+RESULTS_DIR = os.environ.get(
	46	+ "SCATTER_ADD_PREC_RESULTS_DIR",
	47	+ os.path.join(os.path.dirname(__file__), "results"),
	48	+)
	49	+
	50	+
	51	+# ---------------------------------------------------------------------------
	52	+# Precision gate
	53	+# ---------------------------------------------------------------------------
	54	+
	55	+def check_precision(out_npu: torch.Tensor, out_golden: torch.Tensor) -> Dict:
	56	+ npu = out_npu.detach().cpu().to(torch.float32).numpy()
	57	+ golden = out_golden.detach().cpu().to(torch.float32).numpy()
	58	+
	59	+ if npu.shape != golden.shape:
	60	+ raise ValueError(
	61	+ f"Shape mismatch: npu {npu.shape} vs golden {golden.shape}"
	62	+ )
	63	+
	64	+ g64 = golden.astype(np.float64)
	65	+ a64 = npu.astype(np.float64)
	66	+ finite_mask = np.isfinite(g64)
	67	+ g64_safe = np.where(finite_mask, g64, np.float64(0.0))
	68	+ a64_safe = np.where(finite_mask, a64, np.float64(0.0))
	69	+ abs_err_finite = np.abs(a64_safe - g64_safe)
	70	+ elem_threshold = ATOL + RTOL * np.abs(g64_safe)
	71	+ finite_passed = (abs_err_finite <= elem_threshold) & finite_mask
	72	+
	73	+ both_nonfinite = (~finite_mask) & (~np.isfinite(a64))
	74	+ inf_match = (
	75	+ np.isinf(a64) & np.isinf(g64)
	76	+ & (np.signbit(a64) == np.signbit(g64))
	77	+ )
	78	+ nan_match = np.isnan(a64) & np.isnan(g64)
	79	+ nonfinite_passed = (inf_match \| nan_match) & both_nonfinite
	80	+
	81	+ elem_passed = finite_passed \| nonfinite_passed
	82	+
	83	+ finite_count = int(np.sum(finite_mask))
	84	+ max_abs_error = float(np.max(abs_err_finite)) if finite_count > 0 else 0.0
	85	+
	86	+ total = npu.size
	87	+ if total == 0:
	88	+ matched_ratio = 1.0
	89	+ max_abs_error = 0.0
	90	+ else:
	91	+ matched_ratio = float(np.sum(elem_passed)) / total
	92	+
	93	+ ratio_pass = matched_ratio >= REQUIRED_MATCHED_RATIO
	94	+ max_err_pass = max_abs_error <= _MAX_ABS_ERROR_LIMIT_EFF
	95	+ is_pass = bool(ratio_pass and max_err_pass)
	96	+
	97	+ aux_pass = bool(torch.allclose(
	98	+ out_npu.detach().cpu().to(torch.float32),
	99	+ out_golden.detach().cpu().to(torch.float32),
	100	+ rtol=AUX_RTOL, atol=AUX_ATOL, equal_nan=True,
	101	+ ))
	102	+
	103	+ return {
	104	+ "output_tensor": "out",
	105	+ "dtype": "float32",
	106	+ "metric": "mixed_tolerance",
	107	+ "rtol": RTOL,
	108	+ "atol": ATOL,
	109	+ "matched_ratio": matched_ratio,
	110	+ "required_matched_ratio": REQUIRED_MATCHED_RATIO,
	111	+ "max_abs_error": max_abs_error,
	112	+ "max_abs_error_limit": _MAX_ABS_ERROR_LIMIT_EFF,
	113	+ "verdict": "PASS" if is_pass else "FAIL",
	114	+ "ratio_pass": bool(ratio_pass),
	115	+ "max_err_pass": bool(max_err_pass),
	116	+ "aux_allclose": aux_pass,
	117	+ "aux_divergent_from_authoritative": bool(is_pass != aux_pass),
	118	+ "shape": tuple(npu.shape),
	119	+ "total_elements": int(total),
	120	+ }
	121	+
	122	+
	123	+def assert_precision(
	124	+ case_id: str,
	125	+ out_npu: torch.Tensor,
	126	+ out_golden: torch.Tensor,
	127	+ *,
	128	+ persist: bool = True,
	129	+) -> Dict:
	130	+ result = check_precision(out_npu, out_golden)
	131	+ result["case_id"] = case_id
	132	+ result["authoritative_source"] = "float32 mixed tolerance"
	133	+
	134	+ if persist:
	135	+ try:
	136	+ import json
	137	+ os.makedirs(RESULTS_DIR, exist_ok=True)
	138	+ path = os.path.join(RESULTS_DIR, f"{case_id}.json")
	139	+ with open(path, "w") as f:
	140	+ json.dump(result, f, indent=2, default=str)
	141	+ except Exception as e:
	142	+ print(f"[warn] failed to persist precision record for {case_id}: {e}")
	143	+
	144	+ assert result["verdict"] == "PASS", (
	145	+ f"[{case_id}] precision FAIL: "
	146	+ f"matched_ratio={result['matched_ratio']:.6f} "
	147	+ f"(req {REQUIRED_MATCHED_RATIO}), "
	148	+ f"max_abs_error={result['max_abs_error']:.6g} "
	149	+ f"(limit {_MAX_ABS_ERROR_LIMIT_EFF})."
	150	+ )
	151	+ return result
	152	+
	153	+
	154	+def summarize(records: List[Dict]) -> Dict:
	155	+ total = len(records)
	156	+ passed = sum(1 for r in records if r["verdict"] == "PASS")
	157	+ return {
	158	+ "total": total,
	159	+ "passed": passed,
	160	+ "failed": total - passed,
	161	+ "matched_ratio_min": (
	162	+ min(r["matched_ratio"] for r in records) if records else 0.0
	163	+ ),
	164	+ "max_abs_error_max": (
	165	+ max(r["max_abs_error"] for r in records) if records else 0.0
	166	+ ),
	167	+ }
	168	+
	169	+
	170	+# ---------------------------------------------------------------------------
	171	+# Golden — authoritative path via cpu::ScatterAdd (dgl.backend.scatter_add CPU)
	172	+# ---------------------------------------------------------------------------
	173	+
	174	+def golden_scatter_add_cpu(
	175	+ feat_cpu: torch.Tensor, idx_cpu: torch.Tensor, m: int
	176	+) -> torch.Tensor:
	177	+ out = F.scatter_add(feat_cpu, idx_cpu, m)
	178	+ return out.detach().clone()
	179	+
	180	+
	181	+def aux_torch_scatter_add(
	182	+ feat_cpu: torch.Tensor, idx_cpu: torch.Tensor, m: int
	183	+) -> torch.Tensor:
	184	+ """Auxiliary cross-check only (NOT a pass gate)."""
	185	+ out = torch.zeros((m,) + tuple(feat_cpu.shape[1:]), dtype=feat_cpu.dtype)
	186	+ idx_long = idx_cpu.to(torch.int64)
	187	+ view_shape = [idx_long.shape[0]] + [1] * (feat_cpu.ndim - 1)
	188	+ idx_expanded = idx_long.view(view_shape).expand_as(feat_cpu)
	189	+ out.scatter_add_(0, idx_expanded, feat_cpu)
	190	+ return out
	191	+
	192	+
	193	+# ---------------------------------------------------------------------------
	194	+# Case table
	195	+# ---------------------------------------------------------------------------
	196	+
	197	+@dataclass
	198	+class Case:
	199	+ case_id: str
	200	+ level: str
	201	+ desc: str
	202	+ build: Callable[[], tuple] # () -> (feat_cpu, idx_cpu, M, idtype_str)
	203	+
	204	+
	205	+FLT_MAX = float(np.finfo(np.float32).max)
	206	+FLT_MIN_NORMAL = float(np.finfo(np.float32).tiny)
	207	+FLT_SUBNORM_MIN = float(np.finfo(np.float32).smallest_subnormal)
	208	+
	209	+
	210	+def _randn(shape, *kw):
	211	+ return torch.randn(shape, *kw)
	212	+
	213	+
	214	+def _make_feat_with_row_value(n, dim, row_idx, value):
	215	+ f = torch.randn(n, dim, dtype=torch.float32)
	216	+ f[row_idx] = float(value)
	217	+ return f
	218	+
	219	+
	220	+# ---- L0 (6) ----
	221	+
	222	+def _l0_1():
	223	+ feat = _randn(4, 16, dtype=torch.float32)
	224	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	225	+ return feat, idx, 4, "int32"
	226	+
	227	+def _l0_2():
	228	+ feat = _randn(2, 4, dtype=torch.float32)
	229	+ idx = torch.tensor([0, 1], dtype=torch.int64)
	230	+ return feat, idx, 2, "int64"
	231	+
	232	+def _l0_3():
	233	+ feat = _randn(4, 8, dtype=torch.float32)
	234	+ idx = torch.tensor([0, 0, 1, 1], dtype=torch.int32)
	235	+ return feat, idx, 2, "int32"
	236	+
	237	+def _l0_4():
	238	+ feat = _randn(2, 8, dtype=torch.float32)
	239	+ idx = torch.tensor([1, 3], dtype=torch.int32)
	240	+ return feat, idx, 4, "int32"
	241	+
	242	+def _l0_5():
	243	+ feat = _randn(4, 1, dtype=torch.float32)
	244	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	245	+ return feat, idx, 4, "int32"
	246	+
	247	+def _l0_6():
	248	+ feat = _randn(4, 17, dtype=torch.float32)
	249	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	250	+ return feat, idx, 4, "int32"
	251	+
	252	+
	253	+# ---- L1 (24) ----
	254	+
	255	+def _l1_1():
	256	+ feat = _randn(4, 16, dtype=torch.float32)
	257	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	258	+ return feat, idx, 4, "int32"
	259	+
	260	+def _l1_2():
	261	+ feat = _randn(8, 16, dtype=torch.float32)
	262	+ idx = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int32)
	263	+ return feat, idx, 4, "int32"
	264	+
	265	+def _l1_3():
	266	+ feat = _randn(4, 16, dtype=torch.float32)
	267	+ idx = torch.tensor([1, 3, 5, 7], dtype=torch.int32)
	268	+ return feat, idx, 8, "int32"
	269	+
	270	+def _l1_4():
	271	+ feat = _randn(4, 16, dtype=torch.float32)
	272	+ idx = torch.tensor([3, 2, 1, 0], dtype=torch.int32)
	273	+ return feat, idx, 4, "int32"
	274	+
	275	+def _l1_5():
	276	+ n = 1024
	277	+ feat = _randn(n, 1, dtype=torch.float32)
	278	+ idx = torch.randint(0, n, (n,), dtype=torch.int32)
	279	+ return feat, idx, n, "int32"
	280	+
	281	+def _l1_6():
	282	+ n = 64
	283	+ feat = _randn(n, 17, dtype=torch.float32)
	284	+ idx = torch.randint(0, n, (n,), dtype=torch.int32)
	285	+ return feat, idx, n, "int32"
	286	+
	287	+def _l1_7():
	288	+ n = 64
	289	+ feat = _randn(n, 33, dtype=torch.float32)
	290	+ idx = torch.randint(0, n, (n,), dtype=torch.int32)
	291	+ return feat, idx, n, "int32"
	292	+
	293	+def _l1_8():
	294	+ n = 1024
	295	+ feat = _randn(n, 256, dtype=torch.float32)
	296	+ idx = torch.randint(0, n, (n,), dtype=torch.int32)
	297	+ return feat, idx, n, "int32"
	298	+
	299	+def _l1_9():
	300	+ n = 64
	301	+ feat = _randn(n, 16, dtype=torch.float32)
	302	+ idx = torch.randint(0, n, (n,), dtype=torch.int32)
	303	+ return feat, idx, n, "int32"
	304	+
	305	+def _l1_10():
	306	+ n = 64
	307	+ feat = _randn(n, 16, dtype=torch.float32)
	308	+ idx = torch.randint(0, n, (n,), dtype=torch.int64)
	309	+ return feat, idx, n, "int64"
	310	+
	311	+def _l1_11():
	312	+ n = 1024
	313	+ feat = _randn(n, 16, dtype=torch.float32)
	314	+ idx = torch.zeros(n, dtype=torch.int32)
	315	+ return feat, idx, 1, "int32"
	316	+
	317	+def _l1_12():
	318	+ feat = _randn(4, 16, dtype=torch.float32)
	319	+ idx = torch.tensor([0, 2, 4, 6], dtype=torch.int32)
	320	+ return feat, idx, 8, "int32"
	321	+
	322	+def _l1_13():
	323	+ n = 4096
	324	+ feat = _randn(n, 128, dtype=torch.float32)
	325	+ idx = torch.randint(0, n, (n,), dtype=torch.int32)
	326	+ return feat, idx, n, "int32"
	327	+
	328	+def _l1_14():
	329	+ feat = torch.tensor(
	330	+ [[0.0, -0.0, 0.0, -0.0],
	331	+ [-0.0, 0.0, -0.0, 0.0],
	332	+ [0.0, -0.0, 0.0, -0.0],
	333	+ [-0.0, 0.0, -0.0, 0.0]],
	334	+ dtype=torch.float32,
	335	+ )
	336	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	337	+ return feat, idx, 4, "int32"
	338	+
	339	+def _l1_15():
	340	+ feat = _randn(4, 4, dtype=torch.float32)
	341	+ feat[0] = float("inf")
	342	+ feat[1] = float("inf")
	343	+ idx = torch.tensor([0, 0, 1, 1], dtype=torch.int32)
	344	+ return feat, idx, 2, "int32"
	345	+
	346	+def _l1_16():
	347	+ feat = _randn(4, 4, dtype=torch.float32)
	348	+ feat[0] = float("-inf")
	349	+ feat[1] = float("-inf")
	350	+ idx = torch.tensor([0, 0, 1, 1], dtype=torch.int32)
	351	+ return feat, idx, 2, "int32"
	352	+
	353	+def _l1_17():
	354	+ feat = _randn(2, 4, dtype=torch.float32)
	355	+ feat[0] = float("inf")
	356	+ feat[1] = float("-inf")
	357	+ idx = torch.tensor([0, 0], dtype=torch.int32)
	358	+ return feat, idx, 1, "int32"
	359	+
	360	+def _l1_18():
	361	+ feat = _randn(4, 4, dtype=torch.float32)
	362	+ feat[0, 0] = float("nan")
	363	+ feat[1, 1] = float("nan")
	364	+ idx = torch.tensor([0, 0, 1, 1], dtype=torch.int32)
	365	+ return feat, idx, 2, "int32"
	366	+
	367	+def _l1_19():
	368	+ n = 1024
	369	+ feat = torch.zeros(n, 16, dtype=torch.float32)
	370	+ idx = torch.randint(0, n, (n,), dtype=torch.int32)
	371	+ return feat, idx, n, "int32"
	372	+
	373	+def _l1_20():
	374	+ base = torch.tensor(
	375	+ [1e-3, 1e-2, 1e-1, 1.0, 10.0, 1e3, -1e-3, -1e-2, -1e-1,
	376	+ -1.0, -10.0, -1e3, 5e-2, -5e-2, 0.0, 1.0],
	377	+ dtype=torch.float32,
	378	+ )
	379	+ feat = base.repeat(4, 1)
	380	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	381	+ return feat, idx, 4, "int32"
	382	+
	383	+def _l1_21():
	384	+ feat = torch.tensor(
	385	+ [[FLT_MAX, -FLT_MAX, 1.0, 2.0],
	386	+ [-FLT_MAX, FLT_MAX, 3.0, 4.0],
	387	+ [FLT_MAX, FLT_MAX, 5.0, 6.0],
	388	+ [-FLT_MAX, -FLT_MAX, 7.0, 8.0]],
	389	+ dtype=torch.float32,
	390	+ )
	391	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	392	+ return feat, idx, 4, "int32"
	393	+
	394	+def _l1_22():
	395	+ feat = torch.tensor(
	396	+ [[FLT_MIN_NORMAL, -FLT_MIN_NORMAL, 1.0, 2.0],
	397	+ [-FLT_MIN_NORMAL, FLT_MIN_NORMAL, 3.0, 4.0],
	398	+ [FLT_MIN_NORMAL, FLT_MIN_NORMAL, 5.0, 6.0],
	399	+ [-FLT_MIN_NORMAL, -FLT_MIN_NORMAL, 7.0, 8.0]],
	400	+ dtype=torch.float32,
	401	+ )
	402	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	403	+ return feat, idx, 4, "int32"
	404	+
	405	+def _l1_23():
	406	+ feat = torch.tensor(
	407	+ [[FLT_SUBNORM_MIN, -FLT_SUBNORM_MIN, 1.0, 2.0],
	408	+ [-FLT_SUBNORM_MIN, FLT_SUBNORM_MIN, 3.0, 4.0],
	409	+ [FLT_SUBNORM_MIN, FLT_SUBNORM_MIN, 5.0, 6.0],
	410	+ [-FLT_SUBNORM_MIN, -FLT_SUBNORM_MIN, 7.0, 8.0]],
	411	+ dtype=torch.float32,
	412	+ )
	413	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	414	+ return feat, idx, 4, "int32"
	415	+
	416	+def _l1_24():
	417	+ feat = _randn(8, 16, dtype=torch.float32)
	418	+ idx = torch.tensor([0, 7, 0, 7, 0, 7, 0, 7], dtype=torch.int32)
	419	+ return feat, idx, 8, "int32"
	420	+
	421	+
	422	+# ---- L2 boundary (5) ----
	423	+
	424	+def _l2_1():
	425	+ feat = torch.tensor([[0.5]], dtype=torch.float32)
	426	+ idx = torch.tensor([0], dtype=torch.int32)
	427	+ return feat, idx, 1, "int32"
	428	+
	429	+def _l2_2():
	430	+ feat = torch.zeros(0, 16, dtype=torch.float32)
	431	+ idx = torch.zeros(0, dtype=torch.int32)
	432	+ return feat, idx, 4, "int32"
	433	+
	434	+def _l2_3():
	435	+ n = 1024
	436	+ feat = _randn(n, 16, dtype=torch.float32)
	437	+ idx = torch.zeros(n, dtype=torch.int32)
	438	+ return feat, idx, 1, "int32"
	439	+
	440	+def _l2_4():
	441	+ feat = _randn(1, 16, dtype=torch.float32)
	442	+ idx = torch.tensor([torch.randint(0, 1024, (1,)).item()],
	443	+ dtype=torch.int32)
	444	+ return feat, idx, 1024, "int32"
	445	+
	446	+def _l2_5():
	447	+ feat = torch.tensor([[1.0]], dtype=torch.float32)
	448	+ idx = torch.tensor([0], dtype=torch.int32)
	449	+ return feat, idx, 1, "int32"
	450	+
	451	+
	452	+# ---- 补充 (1) ----
	453	+
	454	+def _supp_1():
	455	+ n = 1024
	456	+ feat = _randn(n, 16, dtype=torch.float32)
	457	+ idx = torch.full((n,), 1023, dtype=torch.int32)
	458	+ return feat, idx, 1024, "int32"
	459	+
	460	+
	461	+# ---- Whitebox (6) ----
	462	+
	463	+def _wb_1():
	464	+ n = 40
	465	+ feat = torch.randn(n, 16, dtype=torch.float32)
	466	+ idx = torch.arange(n, dtype=torch.int32)
	467	+ return feat, idx, n, "int32"
	468	+
	469	+def _wb_2():
	470	+ n = 41
	471	+ feat = torch.randn(n, 16, dtype=torch.float32)
	472	+ idx = torch.arange(n, dtype=torch.int32)
	473	+ return feat, idx, n, "int32"
	474	+
	475	+def _wb_3():
	476	+ n = 8
	477	+ feat_dim = 20001
	478	+ feat = torch.randn(n, feat_dim, dtype=torch.float32)
	479	+ idx = torch.arange(n, dtype=torch.int32)
	480	+ return feat, idx, n, "int32"
	481	+
	482	+def _wb_4():
	483	+ feat = torch.randn(4, 2, 8, dtype=torch.float32)
	484	+ idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
	485	+ return feat, idx, 4, "int32"
	486	+
	487	+def _wb_5():
	488	+ n = 80
	489	+ feat = torch.randn(n, 16, dtype=torch.float32)
	490	+ idx = torch.arange(n, dtype=torch.int32)
	491	+ return feat, idx, n, "int32"
	492	+
	493	+def _wb_6():
	494	+ n = 240
	495	+ feat_dim = 4097
	496	+ feat = torch.randn(n, feat_dim, dtype=torch.float32)
	497	+ idx = torch.arange(n, dtype=torch.int32)
	498	+ return feat, idx, n, "int32"
	499	+
	500	+
	501	+# ---- Case registries ----
	502	+
	503	+CASES: List[Case] = [
	504	+ Case("L0-1", "L0", "基础顺序 N=M=4 dim=16 int32", _l0_1),
	505	+ Case("L0-2", "L0", "IdType=int64 最小", _l0_2),
	506	+ Case("L0-3", "L0", "N>M 重复命中", _l0_3),
	507	+ Case("L0-4", "L0", "N<M 间隙", _l0_4),
	508	+ Case("L0-5", "L0", "dim=1", _l0_5),
	509	+ Case("L0-6", "L0", "dim 非对齐 17", _l0_6),
	510	+ Case("L1-1", "L1", "N4 正常组合", _l1_1),
	511	+ Case("L1-2", "L1", "N>M 重复命中", _l1_2),
	512	+ Case("L1-3", "L1", "N<M 间隙", _l1_3),
	513	+ Case("L1-4", "L1", "idx 逆序", _l1_4),
	514	+ Case("L1-5", "L1", "dim=1 大规模", _l1_5),
	515	+ Case("L1-6", "L1", "dim=17 非对齐", _l1_6),
	516	+ Case("L1-7", "L1", "dim=33 非对齐", _l1_7),
	517	+ Case("L1-8", "L1", "dim=256 对齐大规模", _l1_8),
	518	+ Case("L1-9", "L1", "IdType=int32 中规模", _l1_9),
	519	+ Case("L1-10", "L1", "IdType=int64 中规模", _l1_10),
	520	+ Case("L1-11", "L1", "全部命中同一行", _l1_11),
	521	+ Case("L1-12", "L1", "部分行未命中", _l1_12),
	522	+ Case("L1-13", "L1", "GNN 典型规模", _l1_13),
	523	+ Case("L1-14", "L1", "feat 含 ±0", _l1_14),
	524	+ Case("L1-15", "L1", "feat 含 +inf", _l1_15),
	525	+ Case("L1-16", "L1", "feat 含 -inf", _l1_16),
	526	+ Case("L1-17", "L1", "feat 含 ±inf 同行", _l1_17),
	527	+ Case("L1-18", "L1", "feat 含 nan", _l1_18),
	528	+ Case("L1-19", "L1", "feat 全 0", _l1_19),
	529	+ Case("L1-20", "L1", "feat 量级带覆盖", _l1_20),
	530	+ Case("L1-21", "L1", "feat 含 ±FLT_MAX", _l1_21),
	531	+ Case("L1-22", "L1", "feat 含最小正规", _l1_22),
	532	+ Case("L1-23", "L1", "feat 含次正规", _l1_23),
	533	+ Case("L1-24", "L1", "idx 含 0 与 M-1", _l1_24),
	534	+ Case("L2-1", "L2", "最小 shape 单元素", _l2_1),
	535	+ Case("L2-2", "L2", "空 idx/空 feat", _l2_2),
	536	+ Case("L2-3", "L2", "M=1 全命中唯一行", _l2_3),
	537	+ Case("L2-4", "L2", "N=1 单行散射", _l2_4),
	538	+ Case("L2-5", "L2", "dim=1 且 N=M=1", _l2_5),
	539	+ Case("补-1", "补充", "idx 全 M-1", _supp_1),
	540	+ Case("WB-1", "WB", "N=40 blockDim 边界", _wb_1),
	541	+ Case("WB-2", "WB", "N=41 尾核+空闲核", _wb_2),
	542	+ Case("WB-3", "WB", "大 featDim batchItems=1", _wb_3),
	543	+ Case("WB-4", "WB", "3D feat 多维", _wb_4),
	544	+ Case("WB-5", "WB", "N=80 整除 blockDim", _wb_5),
	545	+ Case("WB-6", "WB", "大 featDim 满批+尾批", _wb_6),
	546	+]
	547	+
	548	+EXCEPTION_CASES: List[tuple] = [
	549	+ ("L2-6", torch.float16, torch.int32, "DType=half LOG(FATAL)"),
	550	+ ("L2-7", torch.bfloat16, torch.int64, "DType=bfloat16 LOG(FATAL)"),
	551	+ ("L2-8", torch.float64, torch.int32, "DType=double LOG(FATAL)"),
	552	+ ("L2-9", torch.int8, torch.int32, "DType=非浮点 LOG(FATAL)"),
	553	+]
	554	+
	555	+
	556	+# ---------------------------------------------------------------------------
	557	+# NPU device helpers
	558	+# ---------------------------------------------------------------------------
	559	+
	560	+def _npu_available() -> bool:
	561	+ return hasattr(torch, "npu") and torch.npu.is_available()
	562	+
	563	+def _get_npu_device() -> torch.device:
	564	+ if not _npu_available():
	565	+ pytest.skip("NPU device not available; skipping NPU tests.")
	566	+ dev = torch.device("npu:0")
	567	+ torch.npu.set_device(dev)
	568	+ return dev
	569	+
	570	+def _sync(dev: torch.device) -> None:
	571	+ if dev.type == "npu" and hasattr(torch, "npu"):
	572	+ torch.npu.synchronize(dev)
	573	+
	574	+
	575	+# ---------------------------------------------------------------------------
	576	+# Test runner — precision cases
	577	+# ---------------------------------------------------------------------------
	578	+
	579	+_SUITE_RECORDS: list = []
	580	+
	581	+def _run_precision_case(case):
	582	+ dev = _get_npu_device()
	583	+ feat_cpu, idx_cpu, m, _idtype = case.build()
	584	+ assert feat_cpu.dtype == torch.float32
	585	+ assert idx_cpu.dtype in (torch.int32, torch.int64)
	586	+
	587	+ out_golden = golden_scatter_add_cpu(feat_cpu, idx_cpu, m)
	588	+
	589	+ feat_npu = feat_cpu.to(dev)
	590	+ idx_npu = idx_cpu.to(dev)
	591	+ out_npu = F.scatter_add(feat_npu, idx_npu, m)
	592	+ _sync(dev)
	593	+ out_npu_cpu = out_npu.detach().cpu()
	594	+
	595	+ rec = assert_precision(case.case_id, out_npu_cpu, out_golden)
	596	+ _SUITE_RECORDS.append(rec)
	597	+
	598	+ aux = aux_torch_scatter_add(feat_cpu, idx_cpu, m)
	599	+ aux_match = torch.equal(
	600	+ out_golden.to(torch.float32), aux.to(torch.float32)
	601	+ ) or torch.allclose(
	602	+ out_golden.to(torch.float32), aux.to(torch.float32),
	603	+ rtol=1e-4, atol=1e-4, equal_nan=True,
	604	+ )
	605	+ if not aux_match:
	606	+ print(
	607	+ f"[{case.case_id}] aux torch.scatter_add_ diverges from "
	608	+ f"cpu::ScatterAdd golden — CPU golden is authoritative."
	609	+ )
	610	+
	611	+ if case.case_id == "L2-2":
	612	+ assert out_npu_cpu.numel() > 0
	613	+ assert torch.all(out_npu_cpu == 0.0)
	614	+ if case.case_id in ("L1-12", "L0-4"):
	615	+ assert (out_npu_cpu == 0.0).any()
	616	+ if case.case_id == "WB-3":
	617	+ assert torch.allclose(
	618	+ out_npu_cpu.to(torch.float32), feat_cpu.to(torch.float32),
	619	+ rtol=1e-4, atol=1e-4,
	620	+ )
	621	+ if case.case_id == "WB-5":
	622	+ assert out_npu_cpu.shape[0] == 80
	623	+ assert not torch.all(out_npu_cpu == 0.0)
	624	+ if case.case_id == "WB-4":
	625	+ assert out_npu_cpu.ndim == 3
	626	+ assert tuple(out_npu_cpu.shape) == (4, 2, 8)
	627	+
	628	+
	629	+def _make_precision_test(case):
	630	+ def test_fn(self=None):
	631	+ _run_precision_case(case)
	632	+ test_fn.__name__ = f"test_{case.case_id.replace('-', '_')}"
	633	+ test_fn.__doc__ = f"[{case.level}] {case.case_id}: {case.desc}"
	634	+ return test_fn
	635	+
	636	+for _c in CASES:
	637	+ fn = _make_precision_test(_c)
	638	+ globals()[fn.__name__] = fn
	639	+
	640	+
	641	+# ---------------------------------------------------------------------------
	642	+# Test runner — exception cases
	643	+# ---------------------------------------------------------------------------
	644	+
	645	+def _run_exception_case(case_id, feat_dtype, idx_dtype, desc):
	646	+ dev = _get_npu_device()
	647	+ if feat_dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
	648	+ feat = torch.randn(4, 4).to(feat_dtype)
	649	+ else:
	650	+ feat = torch.tensor([[1, 2, 3, 4]] * 4, dtype=feat_dtype)
	651	+ idx = torch.tensor([0, 0, 1, 1], dtype=idx_dtype)
	652	+ m = 2
	653	+
	654	+ feat_npu = feat.to(dev)
	655	+ idx_npu = idx.to(dev)
	656	+ with pytest.raises(Exception):
	657	+ F.scatter_add(feat_npu, idx_npu, m)
	658	+ _sync(dev)
	659	+
	660	+def _make_exception_test(case_id, feat_dtype, idx_dtype, desc):
	661	+ def test_fn(self=None):
	662	+ _run_exception_case(case_id, feat_dtype, idx_dtype, desc)
	663	+ test_fn.__name__ = f"test_{case_id.replace('-', '_')}"
	664	+ test_fn.__doc__ = f"[L2] {case_id}: {desc}"
	665	+ return test_fn
	666	+
	667	+for _cid, _fd, _id, _desc in EXCEPTION_CASES:
	668	+ fn = _make_exception_test(_cid, _fd, _id, _desc)
	669	+ globals()[fn.__name__] = fn
	670	+
	671	+
	672	+# ---------------------------------------------------------------------------
	673	+# Suite-level summary
	674	+# ---------------------------------------------------------------------------
	675	+
	676	+@pytest.fixture(scope="session", autouse=True)
	677	+def _persist_suite_summary():
	678	+ yield
	679	+ try:
	680	+ import json
	681	+ os.makedirs(RESULTS_DIR, exist_ok=True)
	682	+ summary = summarize(_SUITE_RECORDS)
	683	+ with open(os.path.join(RESULTS_DIR, "suite_summary.json"), "w") as f:
	684	+ json.dump(summary, f, indent=2, default=str)
	685	+ print(
	686	+ f"\n[scatter_add suite] {summary['passed']}/{summary['total']} "
	687	+ f"PASS, {summary['failed']} FAIL"
	688	+ )
	689	+ except Exception as e:
	690	+ print(f"[warn] suite summary persistence failed: {e}")
	691	+
	692	+
	693	+# ---------------------------------------------------------------------------
	694	+# Entry point
	695	+# ---------------------------------------------------------------------------
	696	+
	697	+if __name__ == "__main__":
	698	+ if not _npu_available():
	699	+ print("NPU not available; exiting.")
	700	+ sys.exit(0)
	701	+ fn_names = sorted(globals().keys())
	702	+ failures = 0
	703	+ total = 0
	704	+ for name in fn_names:
	705	+ if not name.startswith("test_"):
	706	+ continue
	707	+ fn = globals()[name]
	708	+ if not callable(fn):
	709	+ continue
	710	+ total += 1
	711	+ try:
	712	+ fn()
	713	+ print(f" PASS [{name}]")
	714	+ except Exception as e:
	715	+ print(f" FAIL [{name}] {e}")
	716	+ failures += 1
	717	+ print(f"\nResults: {total - failures}/{total} passed, {failures} failed")
	718	+ sys.exit(1 if failures else 0)