已合并
Add GatherV2 and ScatterElements DSL operators with generated tests #155
lds创建于 14 天前
Add GatherV2 and ScatterElements DSL operators with generated tests #155
已合并
共 8 个文件变更+4372-0
| @@ -0,0 +1,70 @@ | |||
| 1 | +# gather_v2 | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +- Ascend 950：支持。 | ||
| 6 | + | ||
| 7 | +## 功能说明 | ||
| 8 | + | ||
| 9 | +`gather_v2` 沿指定 `axis` 按 `indices` 取出 `x` 的元素，支持批次维、负轴和可选的负索引。实现位于 [`gather_v2.py`](gather_v2.py)。将 `axis` 和 `batch_dims` 归一化后，输出形状为 `x.shape[:axis] + indices.shape[batch_dims:] + x.shape[axis + 1:]`，输出 dtype 和设备与 `x` 相同。 | ||
| 10 | + | ||
| 11 | +设索引轴长度为 `A`。默认按 `[0, A)` 判断索引合法性；开启 `negative_index_support` 后，负索引先加一次 `A`。归一化后仍越界的位置输出零，索引轴为空时也输出零。数据通过同位宽整数视图搬运，保留 NaN payload、负零和 complex 位模式；输入不被原地修改，仅支持前向计算。 | ||
| 12 | + | ||
| 13 | +## 函数原型 | ||
| 14 | + | ||
| 15 | +```python | ||
| 16 | +gather_v2( | ||
| 17 | + x: torch.Tensor, | ||
| 18 | + indices: torch.Tensor, | ||
| 19 | + axis: int, | ||
| 20 | + *, | ||
| 21 | + batch_dims: int = 0, | ||
| 22 | + negative_index_support: bool = False, | ||
| 23 | + core_num: int | None = None, | ||
| 24 | + thread_num: int | None = None, | ||
| 25 | +) -> torch.Tensor | ||
| 26 | +``` | ||
| 27 | + | ||
| 28 | +`x` 至少一维，支持 float16、bfloat16、float32、float64、int8、uint8、int16、uint16、int32、int64、bool 和 complex64。`indices` 支持 int32/int64，可为标量；两输入必须连续且位于同一 NPU，`x.requires_grad` 必须为 `False`。 | ||
| 29 | + | ||
| 30 | +`axis` 范围为 `[-x.ndim, x.ndim)`。`batch_dims` 允许负值，按 `indices.ndim` 归一化；归一化后须满足 `0 <= batch_dims <= min(axis, indices.ndim)`，两输入前 `batch_dims` 维的形状必须相同。`core_num` 为可选 Vector Core 数上限。`thread_num` 为可选的每核 SIMT 线程数，支持 32、64、128、256、512、1024、2048，其中 2048 按 1024 启动；默认由入口选择。 | ||
| 31 | + | ||
| 32 | +## 测试 | ||
| 33 | + | ||
| 34 | +配置 CANN 环境并进入已安装 PyTorch、torch_npu、cannbotdsl 和 pytest 的 Python 环境，从仓库根目录执行： | ||
| 35 | + | ||
| 36 | +```bash | ||
| 37 | +python -m pytest -q test/gather_v2/test_gather_v2.py | ||
| 38 | +``` | ||
| 39 | + | ||
| 40 | +默认执行 15 条 NPU 用例和 5 条主机校验。NPU 用例包含原有 10 条指定规格的功能与性能用例，以及批次负索引、空 gather 轴、64 位特殊位模式、DMA 缓存复用与淘汰、短字节行尾块和极值索引 5 条典型回归。NPU 不可用时跳过设备测试。仅运行指定规格用例或主机侧测试可执行： | ||
| 41 | + | ||
| 42 | +```bash | ||
| 43 | +python -m pytest -q test/gather_v2/test_gather_v2.py -k test_functional_precision_case | ||
| 44 | +python -m pytest -q test/gather_v2/test_gather_v2.py -m "not npu" | ||
| 45 | +``` | ||
| 46 | + | ||
| 47 | +输入和 CPU 参考结果通过固定随机种子在内存中生成，不保存输入或参考数组。指定规格用例在功能校验通过后，预热 1 次、采集 2 次设备耗时，结果默认写入 pytest 临时目录；可用 `CANNBOT_CASE_RESULTS` 指定本次运行的输出目录。直接运行测试脚本同样执行全部测试，并额外汇总 `performance.csv`。 | ||
| 48 | + | ||
| 49 | +## 性能 | ||
| 50 | + | ||
| 51 | +以下为 2026-10-14 在 Ascend 950DT、Python 3.12 和同一 CANN 环境下配对采集的结果。Ascend C 与 DSL 使用相同输入；时间为设备侧 Task Duration 的最小值，不包含编译、主机调用和输入搬运耗时。Ascend C kernel 模式采集 31 个样本，ACLNN 模式采集 30 个样本；DSL 每例采集 30 个样本。所有 DSL 用例通过独立 CPU 参考校验。 | ||
| 52 | + | ||
| 53 | +`dsl/ascendc = DSL 时间 / Ascend C 时间`，小于 1 表示 DSL 更快。 | ||
| 54 | + | ||
| 55 | +| 用例 | Ascend C (μs) | DSL (μs) | dsl/ascendc | | ||
| 56 | +|:---|--:|--:|--:| | ||
| 57 | +| `gather_v2_0307` | 1.834 | 1.466 | 0.7993 | | ||
| 58 | +| `gather_v2_0016` | 17.978 | 13.964 | 0.7767 | | ||
| 59 | +| `gather_v2_0052` | 1.544 | 1.422 | 0.9210 | | ||
| 60 | +| `gather_v2_0343` | 68.647 | 67.843 | 0.9883 | | ||
| 61 | +| `gather_v2_0003` | 1.416 | 1.285 | 0.9075 | | ||
| 62 | +| `gather_v2_test_case_0` | 2.660 | 2.462 | 0.9256 | | ||
| 63 | +| `gather_v2_test_case_1` | 1.974 | 1.006 | 0.5096 | | ||
| 64 | +| `gather_v2_test_case_2` | 1.709 | 1.024 | 0.5992 | | ||
| 65 | +| `gather_v2_test_case_3` | 1.992 | 1.074 | 0.5392 | | ||
| 66 | +| `gather_v2_test_case_4` | 1.968 | 1.278 | 0.6494 | | ||
| 67 | + | ||
| 68 | +前 5 条 Ascend C 基线使用 `ops-nn` 内核，后 5 条使用 CANN ACLNN 接口；ACLNN 数值仍为设备算子耗时。越界索引按本算子的规则输出零。 | ||
| 69 | + | ||
| 70 | +此表使用配对输入及上述采样方式。测试脚本重新生成输入并使用 2 次采样，其单独测量结果不能与表中的 Ascend C 时间直接组合计算比值。 | ||
| @@ -0,0 +1,67 @@ | |||
| 1 | +# scatter_elements | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +- Ascend 950：支持。 | ||
| 6 | + | ||
| 7 | +## 功能说明 | ||
| 8 | + | ||
| 9 | +`scatter_elements` 根据 `indices` 沿指定 `axis` 把 `updates` 写入 `data` 的副本。实现位于 [`scatter_elements.py`](scatter_elements.py)。对于 `indices` 的位置 `p`，目标位置由 `p` 的非索引轴坐标与 `indices[p]` 的索引轴坐标组成。`reduction="none"` 替换目标值，`"add"` 和 `"mul"` 分别在原值上逐次累加、累乘。输出 shape、dtype 和设备与 `data` 相同，输入不被原地修改，仅支持前向计算。 | ||
| 10 | + | ||
| 11 | +索引须满足 `0 <= indices[p] < data.shape[axis]`；负索引不做归一化。设备内核会跳过越界值以避免非法访存，但越界值不属于有效输入。替换操作通过整数视图保留原始位模式，包括 NaN payload、负零和 complex 值。 | ||
| 12 | + | ||
| 13 | +## 函数原型 | ||
| 14 | + | ||
| 15 | +```python | ||
| 16 | +scatter_elements( | ||
| 17 | + data: torch.Tensor, | ||
| 18 | + indices: torch.Tensor, | ||
| 19 | + updates: torch.Tensor, | ||
| 20 | + axis: int, | ||
| 21 | + *, | ||
| 22 | + reduction: str = "none", | ||
| 23 | + core_num: int | None = None, | ||
| 24 | + deterministic: bool = False, | ||
| 25 | +) -> torch.Tensor | ||
| 26 | +``` | ||
| 27 | + | ||
| 28 | +三个输入必须连续、位于同一 NPU 且 rank 相同，rank 范围为 1～8；`data` 与 `updates` 的 dtype 相同，二者均不允许 `requires_grad=True`。`indices` 支持 int32/int64。每维须满足 `indices.shape[d] <= updates.shape[d]`，非索引轴还须满足 `indices.shape[d] <= data.shape[d]`。`axis` 范围为 `[-data.ndim, data.ndim)`，`core_num` 为可选 Vector Core 数上限。 | ||
| 29 | + | ||
| 30 | +`none` 支持 float16、bfloat16、float32、float64、int8、uint8、int16、int32、int64、bool 和 complex64；`add`、`mul` 支持其中除 float64、complex64 外的类型。空索引返回 `data` 的副本。 | ||
| 31 | + | ||
| 32 | +`deterministic=True` 时，重复目标按 `indices` 的展平顺序处理，替换操作最后一次写入生效。默认非确定性替换不保证重复目标的写入顺序，结果应来自该目标的某次合法更新。float32 原子加法的求和顺序也可能与 CPU 顺序参考不同。 | ||
| 33 | + | ||
| 34 | +## 测试 | ||
| 35 | + | ||
| 36 | +配置 CANN 环境并进入已安装 PyTorch、torch_npu、cannbotdsl 和 pytest 的 Python 环境，从仓库根目录执行： | ||
| 37 | + | ||
| 38 | +```bash | ||
| 39 | +python -m pytest -q test/scatter_elements/test_scatter_elements.py | ||
| 40 | +``` | ||
| 41 | + | ||
| 42 | +默认执行 10 条 NPU 用例和 13 条主机校验。NPU 用例包含原有 5 条指定规格的功能与性能用例，以及原始位模式保留、int64 高位加法与乘法、float16 有序乘法、axis=0 的 float32 确定性加法 5 条典型回归。NPU 不可用时跳过设备测试。仅运行指定规格用例或主机侧测试可执行： | ||
| 43 | + | ||
| 44 | +```bash | ||
| 45 | +python -m pytest -q test/scatter_elements/test_scatter_elements.py -k test_functional_precision_case | ||
| 46 | +python -m pytest -q test/scatter_elements/test_scatter_elements.py -m "not npu" | ||
| 47 | +``` | ||
| 48 | + | ||
| 49 | +输入和 CPU 参考结果通过固定随机种子在内存中生成，不保存输入或参考数组。指定规格用例在功能校验通过后，预热 1 次、采集 2 次设备耗时，结果默认写入 pytest 临时目录；可用 `CANNBOT_CASE_RESULTS` 指定本次运行的输出目录。直接运行测试脚本同样执行全部测试，并额外汇总 `performance.csv`。 | ||
| 50 | + | ||
| 51 | +## 性能 | ||
| 52 | + | ||
| 53 | +以下为 2026-10-14 在 Ascend 950DT、Python 3.12 和同一 CANN 环境下配对采集的结果。Ascend C 与 DSL 使用相同输入；时间为设备侧 Task Duration 的最小值，不包含编译、主机调用和输入搬运耗时。Ascend C kernel 模式采集 31 个样本，ACLNN 模式采集 30 个样本；DSL 每例采集 30 个样本。所有 DSL 用例通过独立 CPU 参考校验。 | ||
| 54 | + | ||
| 55 | +`dsl/ascendc = DSL 时间 / Ascend C 时间`，小于 1 表示 DSL 更快。 | ||
| 56 | + | ||
| 57 | +| 用例 | Ascend C (μs) | DSL (μs) | dsl/ascendc | | ||
| 58 | +|:---|--:|--:|--:| | ||
| 59 | +| `scatter_elements_determ_none_008` | 3.980 | 2.780 | 0.6985 | | ||
| 60 | +| `scatter_elements_none_027` | 5.662 | 4.474 | 0.7902 | | ||
| 61 | +| `scatter_elements_add_014` | 4.580 | 3.544 | 0.7738 | | ||
| 62 | +| `aclnnScatter_fp16_int64_nd_infnan_random_000016` | 5.276 | 4.052 | 0.7680 | | ||
| 63 | +| `ScatterElements_int64_int32_ND_fuzz_000075` | 4.343 | 3.362 | 0.7741 | | ||
| 64 | + | ||
| 65 | +5 条 Ascend C 基线均使用 `ops-nn` 内核，包括名称含 `aclnnScatter` 的用例。所有索引均已调整到合法范围。 | ||
| 66 | + | ||
| 67 | +此表使用配对输入及上述采样方式。测试脚本重新生成输入并使用 2 次采样，其单独测量结果不能与表中的 Ascend C 时间直接组合计算比值。 | ||
| @@ -43,3 +43,5 @@ operators: | |||
| 43 | - mega_recurrent_kda_replayssm | 43 | - mega_recurrent_kda_replayssm |
| 44 | - commit_recurrent_kda_replayssm | 44 | - commit_recurrent_kda_replayssm |
| 45 | - quant_mla_prolog | 45 | - quant_mla_prolog |
| 46 | + - gather_v2 | ||
| 47 | + - scatter_elements | ||
| @@ -0,0 +1,568 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 2 | +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 3 | +# CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 4 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 5 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 6 | +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 7 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 8 | + | ||
| 9 | +"""Host validation and NPU bit-exact tests for the pure DSL GatherV2.""" | ||
| 10 | + | ||
| 11 | +import csv | ||
| 12 | +import datetime | ||
| 13 | +import json | ||
| 14 | +import logging | ||
| 15 | +import os | ||
| 16 | +import sys | ||
| 17 | +import tempfile | ||
| 18 | +from itertools import product | ||
| 19 | +from pathlib import Path | ||
| 20 | + | ||
| 21 | +import pytest | ||
| 22 | +import torch | ||
| 23 | + | ||
| 24 | +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "samples" / "gather_v2")) | ||
| 25 | +from gather_v2 import GatherV2Tiling, gather_v2 | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +LOGGER = logging.getLogger(__name__) | ||
| 29 | +DEVICE_TASK_TYPES = { | ||
| 30 | + "AI_VECTOR_CORE", | ||
| 31 | + "AI_CORE", | ||
| 32 | + "KERNEL_AIVEC", | ||
| 33 | + "KERNEL_AICORE", | ||
| 34 | + "MIX_AIC", | ||
| 35 | + "MIX_AIV", | ||
| 36 | +} | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +def _word_dtype(element_bytes): | ||
| 40 | + return {1: torch.int8, 2: torch.int16, 4: torch.int32, 8: torch.int64}[ | ||
| 41 | + element_bytes | ||
| 42 | + ] | ||
| 43 | + | ||
| 44 | + | ||
| 45 | +def _reference(x, indices, axis, batch_dims=0, negative=False): | ||
| 46 | + axis %= x.ndim | ||
| 47 | + if batch_dims < 0: | ||
| 48 | + batch_dims += indices.ndim | ||
| 49 | + word_dtype = _word_dtype(x.element_size()) | ||
| 50 | + raw_x = x.view(word_dtype) | ||
| 51 | + inner_axis_start = axis + 1 | ||
| 52 | + shape = ( | ||
| 53 | + tuple(x.shape[:axis]) | ||
| 54 | + + tuple(indices.shape[batch_dims:]) | ||
| 55 | + + tuple(x.shape[inner_axis_start:]) | ||
| 56 | + ) | ||
| 57 | + output = torch.zeros(shape, dtype=word_dtype) | ||
| 58 | + if not output.numel() or not x.shape[axis]: | ||
| 59 | + return output.view(x.dtype) | ||
| 60 | + for batch in product(*(range(n) for n in x.shape[:batch_dims])): | ||
| 61 | + selected = indices[batch].long().flatten() | ||
| 62 | + if negative: | ||
| 63 | + selected = torch.where(selected < 0, selected + x.shape[axis], selected) | ||
| 64 | + valid = (selected >= 0) & (selected < x.shape[axis]) | ||
| 65 | + safe = selected.clamp(0, x.shape[axis] - 1) | ||
| 66 | + chunk = torch.index_select(raw_x[batch], axis - batch_dims, safe) | ||
| 67 | + mask_shape = [1] * chunk.ndim | ||
| 68 | + mask_shape[axis - batch_dims] = selected.numel() | ||
| 69 | + chunk = torch.where(valid.reshape(mask_shape), chunk, torch.zeros_like(chunk)) | ||
| 70 | + output[batch] = chunk.reshape(output[batch].shape) | ||
| 71 | + return output.view(x.dtype) | ||
| 72 | + | ||
| 73 | + | ||
| 74 | +@pytest.fixture(scope="module") | ||
| 75 | +def npu_ready(): | ||
| 76 | + pytest.importorskip("torch_npu") | ||
| 77 | + if not torch.npu.is_available(): | ||
| 78 | + pytest.skip("Ascend NPU is unavailable") | ||
| 79 | + return torch.device("npu:0") | ||
| 80 | + | ||
| 81 | + | ||
| 82 | +def _assert_npu_exact(x, indices, axis, device, **kwargs): | ||
| 83 | + expected = _reference( | ||
| 84 | + x, | ||
| 85 | + indices, | ||
| 86 | + axis, | ||
| 87 | + kwargs.get("batch_dims", 0), | ||
| 88 | + kwargs.get("negative_index_support", False), | ||
| 89 | + ) | ||
| 90 | + actual = gather_v2(x.to(device), indices.to(device), axis, **kwargs).cpu() | ||
| 91 | + assert actual.shape == expected.shape | ||
| 92 | + dtype = _word_dtype(x.element_size()) | ||
| 93 | + torch.testing.assert_close(actual.view(dtype), expected.view(dtype), rtol=0, atol=0) | ||
| 94 | + | ||
| 95 | + | ||
| 96 | +def test_host_batch_shape_and_thread_cap(): | ||
| 97 | + t = GatherV2Tiling( | ||
| 98 | + (2, 3, 4, 5), | ||
| 99 | + (2, 7, 9), | ||
| 100 | + -2, | ||
| @@ -0,0 +1,455 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 2 | +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 3 | +# CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 4 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 5 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 6 | +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 7 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 8 | + | ||
| 9 | +import csv | ||
| 10 | +import datetime | ||
| 11 | +import json | ||
| 12 | +import logging | ||
| 13 | +import os | ||
| 14 | +import sys | ||
| 15 | +import tempfile | ||
| 16 | +from itertools import product | ||
| 17 | +from pathlib import Path | ||
| 18 | + | ||
| 19 | +import pytest | ||
| 20 | +import torch | ||
| 21 | + | ||
| 22 | +sys.path.insert( | ||
| 23 | + 0, str(Path(__file__).resolve().parents[2] / "samples" / "scatter_elements") | ||
| 24 | +) | ||
| 25 | +from scatter_elements import ScatterElementsTiling, scatter_elements | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +LOGGER = logging.getLogger(__name__) | ||
| 29 | +DEVICE_TASK_TYPES = { | ||
| 30 | + "AI_VECTOR_CORE", | ||
| 31 | + "AI_CORE", | ||
| 32 | + "KERNEL_AIVEC", | ||
| 33 | + "KERNEL_AICORE", | ||
| 34 | + "MIX_AIC", | ||
| 35 | + "MIX_AIV", | ||
| 36 | +} | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +def test_tiling_and_strides(): | ||
| 40 | + tiling = ScatterElementsTiling((2, 3, 5), (2, 2, 4), (2, 3, 5), -2, "add") | ||
| 41 | + assert tiling.axis == 1 | ||
| 42 | + assert tiling.data_numel == 30 | ||
| 43 | + assert tiling.indices_numel == 16 | ||
| 44 | + assert tiling.axis_size == 3 | ||
| 45 | + assert tiling.data_strides == (15, 5, 1) | ||
| 46 | + assert tiling.index_strides == (8, 4, 1) | ||
| 47 | + assert tiling.update_strides == (15, 5, 1) | ||
| 48 | + assert len(tiling.stride_parameters) == 36 | ||
| 49 | + assert tiling.stride_parameters[24:27] == (2, 2, 4) | ||
| 50 | + assert tiling.stride_parameters[32:36] == (3, 1, 3, 1) | ||
| 51 | + | ||
| 52 | + | ||
| 53 | +@pytest.mark.parametrize( | ||
| 54 | + "data_shape,index_shape,update_shape,axis,reduction,match", | ||
| 55 | + [ | ||
| 56 | + ((), (), (), 0, "none", "rank"), | ||
| 57 | + ((2, 3), (2,), (2, 3), 0, "none", "same rank"), | ||
| 58 | + ((2, 3), (3, 3), (3, 3), 1, "none", "data.shape"), | ||
| 59 | + ((2, 3), (2, 3), (2, 2), 1, "none", "updates.shape"), | ||
| 60 | + ((2, 3), (2, 3), (2, 3), 2, "none", "axis"), | ||
| 61 | + ((2, 3), (2, 3), (2, 3), 1, "mean", "reduction"), | ||
| 62 | + ((2, 0), (2, 1), (2, 1), 1, "none", "cannot be empty"), | ||
| 63 | + ], | ||
| 64 | +) | ||
| 65 | +def test_invalid_tiling(data_shape, index_shape, update_shape, axis, reduction, match): | ||
| 66 | + with pytest.raises(ValueError, match=match): | ||
| 67 | + ScatterElementsTiling(data_shape, index_shape, update_shape, axis, reduction) | ||
| 68 | + | ||
| 69 | + | ||
| 70 | +@pytest.mark.parametrize("axis", [True, 1.5, "1"]) | ||
| 71 | +def test_invalid_axis_type(axis): | ||
| 72 | + with pytest.raises(TypeError, match="axis"): | ||
| 73 | + ScatterElementsTiling((2, 3), (2, 3), (2, 3), axis) | ||
| 74 | + | ||
| 75 | + | ||
| 76 | +def test_empty_indices_is_valid(): | ||
| 77 | + tiling = ScatterElementsTiling((2, 3), (0, 2), (1, 3), 1) | ||
| 78 | + assert tiling.indices_numel == 0 | ||
| 79 | + assert tiling.data_numel == 6 | ||
| 80 | + | ||
| 81 | + | ||
| 82 | +def test_overflow(): | ||
| 83 | + with pytest.raises(OverflowError, match="64-bit"): | ||
| 84 | + ScatterElementsTiling((2**62, 4), (0, 1), (0, 1), 1) | ||
| 85 | + | ||
| 86 | + | ||
| 87 | +@pytest.fixture(scope="module") | ||
| 88 | +def torch_npu_ready(): | ||
| 89 | + pytest.importorskip("torch_npu") | ||
| 90 | + if not torch.npu.is_available(): | ||
| 91 | + pytest.skip("Ascend NPU is unavailable") | ||
| 92 | + | ||
| 93 | + | ||
| 94 | +def reference(data, indices, updates, axis, reduction): | ||
| 95 | + output = data.clone() | ||
| 96 | + if not indices.numel(): | ||
| 97 | + return output | ||
| 98 | + for coordinate in product(*(range(size) for size in indices.shape)): | ||
| 99 | + destination = list(coordinate) | ||
| 100 | + destination[axis] = int(indices[coordinate]) | ||
| 101 | + destination = tuple(destination) | ||
| 102 | + if reduction == "none": | ||
| 103 | + output[destination] = updates[coordinate] | ||
| 104 | + elif reduction == "add": | ||
| 105 | + output[destination] = output[destination] + updates[coordinate] | ||
| 106 | + else: | ||
| 107 | + output[destination] = output[destination] * updates[coordinate] | ||
| 108 | + return output | ||
| 109 | + | ||
| 110 | + | ||
| 111 | +@pytest.mark.npu | ||
| 112 | +def test_none_preserves_bits(torch_npu_ready): | ||
| 113 | + bits = torch.tensor( | ||
| 114 | + [0x7FC01234, -2147483648, 0x7F800000, -8388608], dtype=torch.int32 | ||
| 115 | + ) | ||
| 116 | + data = bits.view(torch.float32) | ||
| 117 | + indices = torch.tensor([3, 1], dtype=torch.int64) | ||
| 118 | + update_bits = torch.tensor([0x7FA12345, -2147483648], dtype=torch.int32) | ||
| 119 | + updates = update_bits.view(torch.float32) | ||
| 120 | + actual = scatter_elements( | ||
| 121 | + data.npu(), indices.npu(), updates.npu(), 0, deterministic=True | ||
| 122 | + ) | ||
| 123 | + expected = reference(data, indices, updates, 0, "none") | ||
| 124 | + assert torch.equal(actual.cpu().view(torch.int32), expected.view(torch.int32)) | ||
| 125 | + | ||
| 126 | + | ||
| 127 | +@pytest.mark.npu | ||
| 128 | +@pytest.mark.parametrize("reduction", ["add", "mul"]) | ||
| 129 | +def test_int64_reduction_high_bits(torch_npu_ready, reduction): | ||
| 130 | + data = torch.tensor( | ||
| 131 | + [(1 << 63) - 1, -(1 << 63), -1, (1 << 32) - 1, 1 << 40], | ||
| 132 | + dtype=torch.int64, | ||
| 133 | + ) | ||
| 134 | + indices = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4], dtype=torch.int64) | ||
| 135 | + updates = torch.tensor( | ||
| 136 | + [-1, 1, (1 << 32) - 1, 2, -(1 << 63), -1, 3, 1 << 40, -1, 2], | ||
| 137 | + dtype=torch.int64, | ||
| 138 | + ) | ||
| 139 | + expected = reference(data, indices, updates, 0, reduction) | ||
| 140 | + actual = scatter_elements( | ||
| 141 | + data.npu(), indices.npu(), updates.npu(), 0, reduction=reduction | ||
| 142 | + ) | ||
| 143 | + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) | ||
| 144 | + | ||
| 145 | + | ||
| 146 | +@pytest.mark.npu | ||
| 147 | +@pytest.mark.parametrize("dtype_name", ["float16"]) | ||
| 148 | +@pytest.mark.parametrize("index_name", ["int32"]) | ||
| 149 | +@pytest.mark.parametrize("reduction", ["mul"]) | ||
| 150 | +@pytest.mark.parametrize("deterministic", [True]) | ||
| 151 | +def test_fractional_negative_ordered_reduction( | ||
| 152 | + torch_npu_ready, dtype_name, index_name, reduction, deterministic | ||
| 153 | +): | ||
| 154 | + """Fractional/negative payloads must not pass through integer loop state.""" | ||
| 155 | + dtype = getattr(torch, dtype_name) | ||
| 156 | + data = (torch.arange(30).remainder(9).float() / 8 - 0.75).reshape(2, 3, 5).to(dtype) | ||
| 157 | + updates = ( | ||
| 158 | + (torch.arange(30).remainder(7).float() / 8 - 0.5).reshape(2, 3, 5).to(dtype) | ||
| 159 | + ) | ||
| 160 | + indices = torch.tensor( | ||
| 161 | + [[[1, 1, 0, 2], [1, 0, 0, 2]], [[0, 1, 2, 1], [0, 1, 2, 1]]], | ||
| 162 | + dtype=getattr(torch, index_name), | ||
| 163 | + ) | ||
| 164 | + expected = reference(data, indices, updates, 1, reduction) | ||
| 165 | + actual = scatter_elements( | ||
| 166 | + data.npu(), | ||
| 167 | + indices.npu(), | ||
| 168 | + updates.npu(), | ||
| 169 | + 1, | ||
| 170 | + reduction=reduction, | ||
| 171 | + deterministic=deterministic, | ||
| 172 | + ) | ||
| 173 | + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) | ||
| 174 | + | ||
| 175 | + | ||
| 176 | +@pytest.mark.npu | ||
| 177 | +@pytest.mark.parametrize("index_name", ["int64"]) | ||
| 178 | +@pytest.mark.parametrize("deterministic", [True]) | ||
| 179 | +def test_axis0_float32_fractional_add(torch_npu_ready, index_name, deterministic): | ||
| 180 | + """Cover deterministic and atomic addition with repeated fractional updates.""" | ||
| 181 | + shape = (10, 17, 29) | ||
| 182 | + generator = torch.Generator().manual_seed(2701) | ||
| 183 | + data = torch.randint(-32, 33, shape, generator=generator).float() / 16 | ||
| 184 | + updates = torch.randint(-16, 17, shape, generator=generator).float() / 16 | ||
| 185 | + indices = torch.randint( | ||
| 186 | + 0, 10, shape, generator=generator, dtype=getattr(torch, index_name) | ||
| 187 | + ) | ||
| 188 | + expected = reference(data, indices, updates, 0, "add") | ||
| 189 | + actual = scatter_elements( | ||
| 190 | + data.npu(), | ||
| 191 | + indices.npu(), | ||
| 192 | + updates.npu(), | ||
| 193 | + 0, | ||
| 194 | + reduction="add", | ||
| 195 | + deterministic=deterministic, | ||
| 196 | + ) | ||
| 197 | + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) | ||
| 198 | + | ||
| 199 | + | ||
| 200 | +# data shape, index shape, update shape, data dtype, index dtype, axis, | ||
| 201 | +# reduction and reproducible seed. Indices are generated within axis bounds. | ||
| 202 | +_CASES = [ | ||
| 203 | + pytest.param( | ||
| 204 | + (10, 38, 26), | ||
| 205 | + (10, 18, 26), | ||
| 206 | + (10, 28, 96), | ||
| 207 | + torch.int32, | ||
| 208 | + torch.int32, | ||
| 209 | + 1, | ||
| 210 | + "none", | ||
| 211 | + 2001, | ||
| 212 | + id="scatter_elements_determ_none_008", | ||
| 213 | + ), | ||
| 214 | + pytest.param( | ||
| 215 | + (1, 3, 8, 16, 5, 3, 4), | ||
| 216 | + (1, 3, 8, 16, 5, 3, 4), | ||
| 217 | + (1, 3, 8, 16, 5, 3, 4), | ||
| 218 | + torch.bfloat16, | ||
| 219 | + torch.int32, | ||
| 220 | + 2, | ||
| 221 | + "none", | ||
| 222 | + 2002, | ||
| 223 | + id="scatter_elements_none_027", | ||
| 224 | + ), | ||
| 225 | + pytest.param( | ||
| 226 | + (10, 2, 18, 26, 5), | ||
| 227 | + (10, 2, 18, 26, 5), | ||
| 228 | + (10, 2, 18, 26, 5), | ||
| 229 | + torch.float32, | ||
| 230 | + torch.int64, | ||
| 231 | + 0, | ||
| 232 | + "add", | ||
| 233 | + 2003, | ||
| 234 | + id="scatter_elements_add_014", | ||
| 235 | + ), | ||
| 236 | + pytest.param( | ||
| 237 | + (5, 9, 15, 4, 6, 6, 11), | ||
| 238 | + (1, 1, 1, 1, 1, 1, 2), | ||
| 239 | + (7, 3, 12, 2, 4, 12, 9), | ||
| 240 | + torch.float16, | ||
| 241 | + torch.int64, | ||
| 242 | + -2, | ||
| 243 | + "none", | ||
| 244 | + 2004, | ||
| 245 | + id="aclnnScatter_fp16_int64_nd_infnan_random_000016", | ||
| 246 | + ), | ||
| 247 | + pytest.param( | ||
| 248 | + (672, 128), | ||
| 249 | + (672, 97), | ||
| 250 | + (672, 97), | ||
| 251 | + torch.int32, | ||
| 252 | + torch.int64, | ||
| 253 | + 1, | ||
| 254 | + "none", | ||
| 255 | + 2005, | ||
| 256 | + id="ScatterElements_int64_int32_ND_fuzz_000075", | ||
| 257 | + ), | ||
| 258 | +] | ||
| 259 | + | ||
| 260 | + | ||
| 261 | +def _check_precision_result(data, indices, updates, axis, reduction, actual): | ||
| 262 | + axis %= data.ndim | ||
| 263 | + assert bool(((indices >= 0) & (indices < data.shape[axis])).all()) | ||
| 264 | + assert actual.shape == data.shape and actual.dtype == data.dtype | ||
| 265 | + if reduction == "none": | ||
| 266 | + word = {1: torch.int8, 2: torch.int16, 4: torch.int32, 8: torch.int64}[ | ||
| 267 | + data.element_size() | ||
| 268 | + ] | ||
| 269 | + raw_data, raw_updates, raw_actual = [ | ||
| 270 | + value.view(word) for value in (data, updates, actual) | ||
| 271 | + ] | ||
| 272 | + choices = {} | ||
| 273 | + for coordinate in product(*(range(size) for size in indices.shape)): | ||
| 274 | + destination = list(coordinate) | ||
| 275 | + destination[axis] = int(indices[coordinate]) | ||
| 276 | + choices.setdefault(tuple(destination), set()).add( | ||
| 277 | + int(raw_updates[coordinate]) | ||
| 278 | + ) | ||
| 279 | + mask = torch.ones(data.shape, dtype=torch.bool) | ||
| 280 | + for destination, legal_values in choices.items(): | ||
| 281 | + mask[destination] = False | ||
| 282 | + assert int(raw_actual[destination]) in legal_values | ||
| 283 | + assert torch.equal(raw_actual[mask], raw_data[mask]), "Untouched bits changed" | ||
| 284 | + else: | ||
| 285 | + expected = reference(data, indices, updates, axis, reduction) | ||
| 286 | + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) | ||
| 287 | + | ||
| 288 | + | ||
| 289 | +@pytest.mark.npu | ||
| 290 | +@pytest.mark.parametrize( | ||
| 291 | + "shape,index_shape,update_shape,dtype,index_dtype,axis,reduction,seed", _CASES | ||
| 292 | +) | ||
| 293 | +def test_functional_precision_case( | ||
| 294 | + torch_npu_ready, | ||
| 295 | + tmp_path, | ||
| 296 | + request, | ||
| 297 | + shape, | ||
| 298 | + index_shape, | ||
| 299 | + update_shape, | ||
| 300 | + dtype, | ||
| 301 | + index_dtype, | ||
| 302 | + axis, | ||
| 303 | + reduction, | ||
| 304 | + seed, | ||
| 305 | +): | ||
| 306 | + generator = torch.Generator().manual_seed(seed) | ||
| 307 | + data = _make_precision_values(shape, dtype, generator) | ||
| 308 | + updates = _make_precision_values(update_shape, dtype, generator) | ||
| 309 | + if dtype.is_floating_point and reduction == "none": | ||
| 310 | + for value in (data, updates): | ||
| 311 | + value.reshape(-1)[:4] = torch.tensor( | ||
| 312 | + [float("nan"), float("inf"), -float("inf"), -0.0], dtype=dtype | ||
| 313 | + ) | ||
| 314 | + indices = torch.randint( | ||
| 315 | + 0, shape[axis], index_shape, generator=generator, dtype=index_dtype | ||
| 316 | + ) | ||
| 317 | + device = [value.npu() for value in (data, indices, updates)] | ||
| 318 | + | ||
| 319 | + def execute(): | ||
| 320 | + return scatter_elements(*device, axis, reduction=reduction, deterministic=False) | ||
| 321 | + | ||
| 322 | + actual = execute().cpu() | ||
| 323 | + _check_precision_result(data, indices, updates, axis, reduction, actual) | ||
| 324 | + _profile_precision_case(request.node.callspec.id, execute, tmp_path) | ||
| 325 | + | ||
| 326 | + | ||
| 327 | +def _make_precision_values(shape, dtype, generator): | ||
| 328 | + """Generate values in memory; preserve GatherV2's uint16 raw-bit inputs.""" | ||
| 329 | + if dtype.is_floating_point: | ||
| 330 | + return ( | ||
| 331 | + torch.randint(-7, 8, shape, generator=generator, dtype=torch.int32).float() | ||
| 332 | + / 8 | ||
| 333 | + ).to(dtype) | ||
| 334 | + return torch.randint(-7, 8, shape, generator=generator, dtype=dtype) | ||
| 335 | + | ||
| 336 | + | ||
| 337 | +def _device_durations(destination): | ||
| 338 | + """Read device durations from the available profiler CSV format.""" | ||
| 339 | + candidates = list(destination.rglob("kernel_details.csv")) | ||
| 340 | + if candidates: | ||
| 341 | + evidence = max(candidates, key=lambda path: path.stat().st_mtime) | ||
| 342 | + with evidence.open(encoding="utf-8", newline="") as stream: | ||
| 343 | + records = list(csv.DictReader(stream)) | ||
| 344 | + assert records, "Profiler returned no kernel samples" | ||
| 345 | + key = None | ||
| 346 | + for field in records[0]: | ||
| 347 | + if field.strip().lower() in ( | ||
| 348 | + "duration(us)", | ||
| 349 | + "duration(μs)", | ||
| 350 | + "duration(µs)", | ||
| 351 | + "duration", | ||
| 352 | + ): | ||
| 353 | + key = field | ||
| 354 | + break | ||
| 355 | + assert key is not None, "Profiler returned no duration column" | ||
| 356 | + durations = [float(row[key]) for row in records] | ||
| 357 | + else: | ||
| 358 | + candidates = list(destination.rglob("op_summary*.csv")) | ||
| 359 | + assert candidates, "No NPU device profiling CSV generated" | ||
| 360 | + evidence = max(candidates, key=lambda path: path.stat().st_mtime) | ||
| 361 | + with evidence.open(encoding="utf-8", newline="") as stream: | ||
| 362 | + durations = [ | ||
| 363 | + float(row["Task Duration(us)"]) | ||
| 364 | + for row in csv.DictReader(stream) | ||
| 365 | + if row.get("Task Type") in DEVICE_TASK_TYPES | ||
| 366 | + ] | ||
| 367 | + assert durations and all((value > 0 for value in durations)) | ||
| 368 | + return (durations, evidence) | ||
| 369 | + | ||
| 370 | + | ||
| 371 | +def _profile_precision_case(name, execute, tmp_path): | ||
| 372 | + """Warm up once and collect two device executions after functional validation.""" | ||
| 373 | + import torch_npu | ||
| 374 | + | ||
| 375 | + output_root = Path(os.environ.get("CANNBOT_CASE_RESULTS", str(tmp_path))) | ||
| 376 | + destination = output_root / name | ||
| 377 | + destination.mkdir(parents=True, exist_ok=True) | ||
| 378 | + torch.npu.synchronize() | ||
| 379 | + with torch_npu.profiler.profile( | ||
| 380 | + activities=[ | ||
| 381 | + torch_npu.profiler.ProfilerActivity.CPU, | ||
| 382 | + torch_npu.profiler.ProfilerActivity.NPU, | ||
| 383 | + ], | ||
| 384 | + schedule=torch_npu.profiler.schedule(wait=0, warmup=1, active=2, repeat=1), | ||
| 385 | + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(str(destination)), | ||
| 386 | + record_shapes=False, | ||
| 387 | + profile_memory=False, | ||
| 388 | + with_stack=False, | ||
| 389 | + ) as profiler: | ||
| 390 | + for _ in range(3): | ||
| 391 | + output = execute() | ||
| 392 | + torch.npu.synchronize() | ||
| 393 | + profiler.step() | ||
| 394 | + durations, evidence = _device_durations(destination) | ||
| 395 | + result = { | ||
| 396 | + "testcase_name": name, | ||
| 397 | + "dsl_us": sum(durations) / 2, | ||
| 398 | + "functional_validation": "PASS", | ||
| 399 | + "effective_samples": 2, | ||
| 400 | + "device_duration_samples_us": durations, | ||
| 401 | + "profiling_csv": str(evidence), | ||
| 402 | + "output_shape": list(output.shape), | ||
| 403 | + "output_dtype": str(output.dtype), | ||
| 404 | + } | ||
| 405 | + (destination / "result.json").write_text( | ||
| 406 | + json.dumps(result, indent=2), encoding="utf-8" | ||
| 407 | + ) | ||
| 408 | + LOGGER.info("%s: PASS, DSL=%.3f us", name, result["dsl_us"]) | ||
| 409 | + | ||
| 410 | + | ||
| 411 | +def _run_tests(): | ||
| 412 | + """Run one operator script and summarize its measured cases in a temporary directory.""" | ||
| 413 | + logging.basicConfig(level=logging.INFO, format="%(message)s") | ||
| 414 | + script = Path(__file__).resolve() | ||
| 415 | + timestamp = datetime.datetime.now(tz=datetime.timezone.utc).strftime( | ||
| 416 | + "%Y%m%d_%H%M%S_%f" | ||
| 417 | + ) | ||
| 418 | + output = ( | ||
| 419 | + Path(tempfile.gettempdir()) | ||
| 420 | + / "cannbot_operator_tests" | ||
| 421 | + / (script.parent.name + "_" + timestamp) | ||
| 422 | + ) | ||
| 423 | + os.environ.setdefault("CANNBOT_CASE_RESULTS", str(output)) | ||
| 424 | + result_dir = Path(os.environ["CANNBOT_CASE_RESULTS"]) | ||
| 425 | + result_dir.mkdir(parents=True, exist_ok=True) | ||
| 426 | + status = pytest.main([str(script), "-q", "-s", *sys.argv[1:]]) | ||
| 427 | + results = [] | ||
| 428 | + for case in _CASES: | ||
| 429 | + path = result_dir / case.id / "result.json" | ||
| 430 | + if path.exists(): | ||
| 431 | + results.append(json.loads(path.read_text(encoding="utf-8"))) | ||
| 432 | + if results: | ||
| 433 | + fields = [ | ||
| 434 | + "testcase_name", | ||
| 435 | + "dsl_us", | ||
| 436 | + "functional_validation", | ||
| 437 | + "effective_samples", | ||
| 438 | + ] | ||
| 439 | + with (result_dir / "performance.csv").open( | ||
| 440 | + "w", encoding="utf-8-sig", newline="" | ||
| 441 | + ) as stream: | ||
| 442 | + writer = csv.DictWriter(stream, fieldnames=fields, extrasaction="ignore") | ||
| 443 | + writer.writeheader() | ||
| 444 | + writer.writerows(results) | ||
| 445 | + LOGGER.info( | ||
| 446 | + "Measured %s/%s cases. Results: %s", | ||
| 447 | + len(results), | ||
| 448 | + len(_CASES), | ||
| 449 | + result_dir / "performance.csv", | ||
| 450 | + ) | ||
| 451 | + return status | ||
| 452 | + | ||
| 453 | + | ||
| 454 | +if __name__ == "__main__": | ||
| 455 | + raise SystemExit(_run_tests()) | ||
| @@ -281,6 +281,22 @@ flash_mla_with_kvcache: | |||
| 281 | options: | 281 | options: |
| 282 | - flash_mla_with_kvcache | 282 | - flash_mla_with_kvcache |
| 283 | 283 | ||
| 284 | +gather_v2: | ||
| 285 | + module: True | ||
| 286 | + src: | ||
| 287 | + - samples/gather_v2 | ||
| 288 | + - test/gather_v2 | ||
| 289 | + options: | ||
| 290 | + - gather_v2 | ||
| 291 | + | ||
| 292 | +scatter_elements: | ||
| 293 | + module: True | ||
| 294 | + src: | ||
| 295 | + - samples/scatter_elements | ||
| 296 | + - test/scatter_elements | ||
| 297 | + options: | ||
| 298 | + - scatter_elements | ||
| 299 | + | ||
| 284 | shared: | 300 | shared: |
| 285 | module: True | 301 | module: True |
| 286 | src: | 302 | src: |