PyPTO 数据拷贝粒度控制：DataCopyMode 枚举在 vf.load_align / vf.store_align / vf.gather 中的用法详解
【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto
DataCopyMode 是 PyPTO（Parallel Tensor/Tile Operation 编程范式）中用于控制 Vector Function（VF）数据搬移粒度的枚举类型，它决定了vf.load_align、vf.store_align、vf.gather三条核心数据搬移指令是"逐元素拷贝"还是"按 DataBlock 块拷贝"。本文从该枚举的定义出发，结合仓库中的 Python 绑定实现（ir.cpp）、VF API 声明（_vf_api.py）与 ST 测试用例（test_vf_basic_ops.py），系统讲解三种模式的语义、适用指令、参数约束与实战示例，帮助读者在编写 VF kernel 时正确选择数据拷贝粒度。
产品支持情况
DataCopyMode 所依赖的 VF 数据搬移指令仅在新一代昇腾硬件上可用：
| 产品 | 支持情况 |
|---|---|
| Ascend 950PR / Ascend 950DT | 支持 |
| Atlas A3 训练系列产品 / Atlas A3 推理系列产品 | 不支持 |
| Atlas A2 训练系列产品 / Atlas A2 推理系列产品 | 不支持 |
这意味着所有data_copy_mode相关用法（含vf.gather的 DATA_BLOCK_LOAD 模式）都需要在 Ascend 950 系列硬件上运行，迁移到 A2/A3 平台时请先确认目标指令可用性。
功能说明：DataCopyMode 控制什么
DataCopyMode 定义了 vf.load_align、vf.store_align 和 vf.gather 的数据拷贝粒度模式，用于控制是逐元素拷贝还是按 DataBlock 块拷贝：
- 逐元素拷贝（NORM）：以单个数据元素为最小搬移单位，索引（index）以元素为单位，逻辑直观，适合常规连续数据搬移与按元素收集。
- 按 DataBlock 块拷贝：以 32 字节（32B）DataBlock 为最小搬移单位，一次搬移一整块数据，可显著减少搬移指令条数与地址计算开销，适合处理块状布局的数据。
从源码看，DataCopyMode由 C++ 侧定义并通过 pybind11 暴露给 Python：在 ir.cpp 中，py::enum_<ir::DataCopyMode>注册了与文档一致的三个成员，并在 Python 侧通过pypto.ir模块导出（_vf_api.py顶部from pypto.ir import DataCopyMode），最终用户以pl.DataCopyMode.XXX形式访问。
原型定义
class DataCopyMode(enum.Enum): NORM = ... # 普通逐元素拷贝（默认） DATA_BLOCK_LOAD = ... # DataBlock 加载（用于 vf.gather：按 32B DataBlock 粒度 gather） DATA_BLOCK_COPY = ... # 非连续 DataBlock 拷贝（基于 block_stride）三个成员的语义定位与典型使用场景：
| 枚举成员 | 语义 | 典型使用指令 | 底层指令 |
|---|---|---|---|
NORM | 普通逐元素拷贝（默认值） | vf.load_align/vf.store_align/vf.gather | vlds / vsts / vgather 常规模式 |
DATA_BLOCK_LOAD | 按 32B DataBlock 粒度加载/收集 | vf.gather（Tile→reg 形式） | 按 DataBlock 粒度 gather |
DATA_BLOCK_COPY | 基于block_stride的非连续 DataBlock 拷贝 | vf.load_align/vf.store_align | vsldb / vsstb |
需要特别说明DATA_BLOCK_LOAD与DATA_BLOCK_COPY的区别：前者用于vf.gather的按块收集，索引直接对应 DataBlock 位置；后者用于load_align/store_align的非连续块拷贝，块与块之间的地址间隔由block_stride参数指定。DATA_BLOCK_LOAD在load_align的上下文中还被接受为DATA_BLOCK_COPY的等价旧别名（详见下文 load_align 小节）。
调用示例
以下是关联文档给出的最小调用示例，展示在 VF kernel 内使用DATA_BLOCK_COPY模式完成寄存器到 UB Tile 的非连续块存储：
import pypto_pro.language as pl @pl.vector_function def vf_kernel(): vf.store_align(ub_tile, reg, data_copy_mode=pl.DataCopyMode.DATA_BLOCK_COPY)需要注意，vf.store_align是语句形式（无返回值，目标是 UB Tile 而非寄存器），实际使用时还需要提供谓词寄存器preg与可选的block_stride/repeat_stride参数，完整形态见下文。
结合源码的深度解析：三个枚举成员的实际用法
vf.gather 中的 NORM 与 DATA_BLOCK_LOAD：按元素 vs 按 32B 块收集
vf.gather 根据索引将源操作数收集到目的操作数中，是DATA_BLOCK_LOAD最典型的使用场景。在 _vf_api.py 的声明中，data_copy_mode参数注释明确指出：
NORM mode: per-element gather. DATA_BLOCK_LOAD mode: per-32B-DataBlock gather.
gather 会根据src参数类型自动分发为两种形式：
- Tile→reg 形式（src 为 Tile）：从 UB Tile 中按索引收集数据到 reg_tensor，必须提供
preg。此时可通过data_copy_mode选择收集粒度：NORM（默认）：按元素收集，index 单位为元素。例如对 FP32 数据收集 64 个元素，index 取值范围对应元素下标。DATA_BLOCK_LOAD：按 DataBlock（32B）收集，index 单位为字节且需 32B 对齐，一个索引值对应 1 个 DataBlock。
- reg→reg 形式（src 为 reg_tensor）：寄存器到寄存器按元素收集，无需 mask，且不支持
data_copy_mode参数（传入会被忽略或不符合调用约束）。
两种形式的另一差异：Tile→reg 形式下，8 位宽源数据（DT_INT8、DT_UINT8）会被零扩展到 16 位宽；而 reg→reg 形式保持源数据类型不变。这一约束在编写 gather 前后处理时需特别注意——NORM 模式下若 src 为 8 位宽、dst 为 16 位宽，目的寄存器低 8 位为有效数据、高 8 位自动补 0。
NORM 模式数据类型约束（Tile→reg 形式）：
| dst | src | index |
|---|---|---|
| DT_INT16 | DT_INT8 | DT_UINT16 |
| DT_INT16 | DT_INT16 | DT_UINT16 |
| DT_UINT16 | DT_UINT8 | DT_UINT16 |
| DT_UINT16 | DT_UINT16 | DT_UINT16 |
| DT_FP16 | DT_FP16 | DT_UINT16 |
| DT_BF16 | DT_BF16 | DT_UINT16 |
| DT_INT32 | DT_INT32 | DT_UINT32 |
| DT_UINT32 | DT_UINT32 | DT_UINT32 |
| DT_FP32 | DT_FP32 | DT_UINT32 |
| DT_INT64 | DT_INT64 | DT_UINT32 / DT_UINT64 |
| DT_UINT64 | DT_UINT64 | DT_UINT32 / DT_UINT64 |
| DT_INT16 / DT_UINT16 / DT_FP16 / DT_BF16 | 同类型 | DT_UINT32 |
DATA_BLOCK_LOAD 模式数据类型约束：src 支持 DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_FP16、DT_BF16、DT_INT32、DT_UINT32、DT_FP32、DT_INT64、DT_UINT64；index 仅支持 DT_UINT32。
reg→reg 形式：src/dst 支持 8 位宽（DT_INT8、DT_UINT8）、16 位宽（DT_INT16、DT_UINT16、DT_FP16、DT_BF16）、32 位宽（DT_INT32、DT_UINT32、DT_FP32）；index 支持 DT_UINT8、DT_UINT16、DT_UINT32。
一个完整的 NORM 模式 gather 示例（来源：gather.md）：
import os import pypto_pro.language as pl import torch import torch_npu @pl.vector_function def example_vf(src_tile, index_tile, dst_tile): preg = vf.create_mask(pattern=pl.MaskPattern.ALL, dtype=pl.DT_FP32) index_reg = vf.load_align(index_tile, 0) dst_reg = vf.gather(src_tile, index_reg, preg) # 默认 NORM，按元素收集 vf.store_align(dst_tile, dst_reg, preg)若要切换为按 32B DataBlock 收集，只需显式传入data_copy_mode=pl.DataCopyMode.DATA_BLOCK_LOAD，此时 index 寄存器中的值必须以字节为单位且 32B 对齐。
vf.load_align 中的 DATA_BLOCK_COPY：非连续 DataBlock 加载（vsldb）
vf.load_align 将 UB Tile 中连续数据加载到 VF 寄存器。在其声明 _vf_api.py 中：
data_copy_mode=pl.DataCopyMode.DATA_BLOCK_COPY选择vsldb 指令（非连续 DataBlock 加载）；DATA_BLOCK_LOAD被接受为等价的旧别名（legacy alias）；block_stride：DataBlock 步长，单位为字节，用于指定非连续块之间在源地址上的间隔；post_update：置 True 时使用后递增寻址，加载后自动推进目标地址。
调用形式：
dst = vf.load_align(src_tile, offset, data_copy_mode=pl.DataCopyMode.DATA_BLOCK_COPY, block_stride=32)该用法在 ST 测试中有直接验证：test_vf_basic_ops.py 以block_stride=32测试DATA_BLOCK_LOAD（vsldb）路径，第 3103 行则以DATA_BLOCK_COPY+block_stride=32验证非连续块加载（对应日志中的 "Kernel 73 (LoadAlign DATA_BLOCK_LOAD/vsldb) PASSED" 与 "Kernel 86 (LoadAlign DATA_BLOCK_COPY) PASSED"）。
vf.store_align 中的 DATA_BLOCK_COPY：非连续 DataBlock 存储（vsstb）
vf.store_align 将 VF 寄存器数据存储到 UB Tile，DATA_BLOCK_COPY对应vsstb 指令。其声明 _vf_api.py 提供了两个块拷贝专用参数：
block_stride：DataBlock 拷贝的块步长；repeat_stride：DataBlock 拷贝的重复步长（配合 repeat 次数控制多组块的地址跨度）。
语句形式（无赋值形式，目标是 UB Tile）：
vf.store_align(dst_tile, src_reg, preg, data_copy_mode=pl.DataCopyMode.DATA_BLOCK_COPY, block_stride=32)与 load 方向呼应：load_align的DATA_BLOCK_COPY用block_stride描述源端非连续间隔，store_align的DATA_BLOCK_COPY用block_stride/repeat_stride描述目的端非连续间隔。两个方向配合，即可实现 UB 内任意块步长的搬运。
使用注意事项
- 默认行为：
data_copy_mode是可选的 keyword 参数，缺省时取NORM（逐元素），保持与常规 VF 搬移一致，不显式传入不会破坏既有代码。 - 对齐要求：DATA_BLOCK 系列模式以 32B 为粒度，索引/地址必须 32B 对齐；gather 的
DATA_BLOCK_LOAD模式下 index 单位是字节而非元素。 - 作用范围限制：
data_copy_mode仅对vf.gather的 Tile→reg 形式有效，reg→reg 形式不接受该参数；load_align/store_align的 DataBlock 模式需要配合block_stride使用。 - 硬件约束：全部 DataCopyMode 特性仅 Ascend 950PR/950DT 支持，A2/A3 平台不支持。
- 数据类型匹配：gather 的 DATA_BLOCK_LOAD 模式要求源与目的数据类型一致；NORM 模式下 8 位宽源数据零扩展到 16 位，存储时若需按 16 位粒度写回可使用
pl.StoreDist.NORM_B16控制写回粒度。
总结
DataCopyMode 是 PyPTO VF 编程中控制数据搬移粒度的关键开关：NORM提供直观的逐元素语义，DATA_BLOCK_LOAD面向vf.gather的按 32B 块收集，DATA_BLOCK_COPY面向load_align/store_align的基于block_stride的非连续块搬移。三者共享同一个枚举定义（ir.cpp），并在 _vf_api.py 中统一以data_copy_mode关键字参数暴露。编写高性能 VF kernel 时，建议先根据数据布局判断是否具备 32B 对齐的块状结构：连续密集数据用NORM，块状稀疏布局用DATA_BLOCK_COPY/DATA_BLOCK_LOAD，并始终以 test_vf_basic_ops.py 中的用例为参照完成正确性验证。
【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
