PTO ISA 指令详解：TSELS —— 基于 Mask 的 Tile 与标量逐元素选择

PTO ISA 指令详解TSELS —— 基于 Mask 的 Tile 与标量逐元素选择【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isaTSELS是 CANN PTOParallel Tile Operation虚拟指令集体系中一条典型的逐元素选择指令它借助一个Mask Tile在源 Tilesrc与一个**标量scalar**之间按元素做出选择将结果写入目标 Tiledst。该指令与双 Tile 选择的 TSEL 互为姊妹指令广泛用于掩码裁剪如x[mask] ? x : 0、条件置零、ReLU 类分段函数的标量分支等场景。阅读本文后你将掌握 TSELS 的数学语义、两级汇编语法、C 内建接口的完整签名、Atlas A2/A3 与 Ascend 950PR/950DT 两套后端在数据类型、布局、掩码编码与临时空间上的差异并能够直接照抄文中的 Auto/Manual 示例开始开发。指令语义数学定义TSELS对目标 Tile 有效区域内的每个元素(i, j)执行如下选择$$ \mathrm{dst}{i,j} \begin{cases} \mathrm{src}{i,j} \text{if } \mathrm{mask}_{i,j}\ \text{为真} \ \mathrm{scalar} \text{否则} \end{cases} $$即mask 位为真时拷贝源 Tile 对应元素否则写入给定标量。迭代域由dst.GetValidRow()/dst.GetValidCol()决定与 PTO ISA 的通用有效区域约定一致——数学语义仅对0 i valid_row且0 j valid_col的元素做出定义有效区域之外的值为未指定。指令示意图图中展示了 mask Tile 与 src Tile 逐元素交织最终输出 dst 的完整数据流。汇编语法两级抽象TSELS提供同步形式以及 AS Level 1SSA与 AS Level 2DPS三种写法同步形式PTO 汇编%dst tsels %mask, %src, %scalar : !pto.tile...AS Level 1SSA资源由编译器/运行时放置与调度%dst pto.tsels %mask, %src, %scalar : (!pto.tile..., !pto.tile..., dtype) - !pto.tile...AS Level 2DPS显式声明输入/输出缓冲区pto.tsels ins(%mask, %src, %scalar : !pto.tile_buf..., !pto.tile_buf..., dtype) outs(%dst : !pto.tile_buf...)在手动Manual模式下发射指令前需要通过pto.tassign显式绑定 tile 资源对应 C 层的TASSIGN# 手动模式先显式绑定资源再发射指令。 # pto.tassign %arg0, tile(0x1000) # pto.tassign %arg1, tile(0x2000) %dst pto.tsels %mask, %src, %scalar : (!pto.tile..., !pto.tile..., dtype) - !pto.tile...C 内建接口TSELS的 C 内建接口声明于 include/pto/common/pto_instr.hpptemplate typename TileDataDst, typename TileDataMask, typename TileDataSrc, typename TileDataTmp, typename... WaitEvents PTO_INST RecordEvent TSELS(TileDataDst dst, TileDataMask mask, TileDataSrc src, TileDataTmp tmp, typename TileDataSrc::DType scalar, WaitEvents ... events);返回类型为PTO_INST RecordEvent与 PTO 其余指令一致支持通过可变参数WaitEvents ... events挂接事件同步接口内部先执行detail::PtoWaitEvents(events...)再通过MAP_INSTR_IMPL(TSELS, dst, mask, src, tmp, scalar)分发到各后端实现。公共包含头为pto/pto-inst.hpp内部声明位于pto/common/pto_instr.hpp。各后端通过TSELS_IMPL提供实现Atlas A2/A3 见 include/pto/npu/a2a3/TSels.hppAscend 950PR/950DT 见 include/pto/npu/a5/TSels.hppCPU 仿真参考实现见 include/pto/cpu/TSelS.hpp。约束与实现检查Atlas A2/A3 训练系列产品 / Atlas A2/A3 推理系列产品sizeof(TileDataDst::DType)必须是 2 或 4 字节。支持的数据类型int16_t、uint16_t、int32_t、uint32_t、half、bfloat16_t、float。dst与src必须使用相同的元素类型且均为行主序。运行时检查src.GetValidRow()/GetValidCol()必须与dst.GetValidRow()/GetValidCol()一致对应源码中的PTO_ASSERT断言见 include/pto/npu/a2a3/TSels.hpp。Ascend 950PR / Ascend 950DTsizeof(TileDataDst::DType)可以为 1、2、4 或 8 字节。支持的数据类型int8_t、uint8_t、int16_t、uint16_t、int32_t、uint32_t、int64_t、uint64_t、half、bfloat16_t、float。dst、mask、src必须均为行主序。运行时检查src.GetValidRow()/GetValidCol()必须与dst.GetValidRow()/GetValidCol()一致。对应实现中A5 后端通过static_assert对数据类型与布局做编译期校验见 include/pto/npu/a5/TSels.hpp并按元素字节宽度8/4/2/1 字节分别分派到Int64SelectScalar、TSels_b32与TSels_b16_8三条内部路径。有效区域操作以dst.GetValidRow()/dst.GetValidCol()作为迭代域src的有效形状必须与之严格相等。掩码编码64 位数据Ascend 950PR/950DT使用 RowMajor 打包掩码每个元素对应一位列j使用第j / 8字节的第j % 8位行地址按掩码物理字节步长计算。可直接使用 TCMP 或 TCMPS 产生的掩码掩码容量必须覆盖目标有效区域数据 Tile 的有效列数仍按数据元素计数。掩码 Tile 被解释为目标定义布局中的打包谓词位packed predicate bits。临时空间tmp的使用差异tmp参数在两个后端上的行为截然不同这是迁移代码时最容易踩坑的点Atlas A2/A3 训练系列产品 / Atlas A2/A3 推理系列产品tmp被使用作为小型缓冲区存放set_cmpmask操作所需的标量值并保存比较掩码。在选择循环开始前标量被写入tmp[0]对应源码 include/pto/npu/a2a3/TSels.hpp 中的*scalarPtr scalar; set_cmpmask(scalarPtr);。tmp的元素类型必须与TileDataSrc::DType一致即与数据元素类型相同。tmp大小要求至少 1 个元素用于存放标量典型声明为TileTileType::Vec, float, 1, 16或类似形态。Ascend 950PR / Ascend 950DTtmp被接口接受但不被使用。A5 后端通过vdup将标量广播到向量寄存器、再用vsel完成选择不需要暂存 Tile 存储tmp仅为了与 Atlas A2/A3 的 API 保持兼容而保留在 C 内建接口签名中对应源码 include/pto/npu/a5/TSels.hpp 中的vdup(vregScalar, scalar, pReg, MODE_ZEROING)与逐块vsel循环。代码示例自动模式Auto自动模式下Tile 的放置与调度由编译器/运行时负责用户只需声明类型并调用指令#include pto/pto-inst.hpp using namespace pto; void example_auto() { using TileDst TileTileType::Vec, float, 16, 16; using TileSrc TileTileType::Vec, float, 16, 16; using TileTmp TileTileType::Vec, float, 16, 16; using TileMask TileTileType::Vec, uint8_t, 16, 32, BLayout::RowMajor, -1, -1; TileDst dst; TileSrc src; TileTmp tmp; TileMask mask(16, 2); float scalar 0.0f; TSELS(dst, mask, src, tmp, scalar); }注意掩码 Tile 的声明要点uint8_t元素类型配合Cols32、-1, -1步长即每 8 个数据元素对应 1 字节打包位mask(16, 2)的构造参数为有效形状16 行 × 2 字节恰好覆盖16 × 16个 float 元素对应的打包掩码容量。手动模式Manual手动模式下需要先用TASSIGN显式绑定各 Tile 的物理地址单位字节再发射指令#include pto/pto-inst.hpp using namespace pto; void example_manual() { using TileDst TileTileType::Vec, float, 16, 16; using TileSrc TileTileType::Vec, float, 16, 16; using TileTmp TileTileType::Vec, float, 16, 16; using TileMask TileTileType::Vec, uint8_t, 16, 32, BLayout::RowMajor, -1, -1; TileDst dst; TileSrc src; TileTmp tmp; TileMask mask(16, 2); float scalar 0.0f; TASSIGN(src, 0x1000); TASSIGN(tmp, 0x2000); TASSIGN(dst, 0x3000); TASSIGN(mask, 0x4000); TSELS(dst, mask, src, tmp, scalar); }后端实现原理从源码结构看两个后端的实现思路体现了不同的硬件抽象A2A3 路径以vsel(..., SELMODE::VSEL_TENSOR_SCALAR_MODE)为内核先将标量写入tmp[0]通过set_cmpmask建立比较掩码再按行循环执行向量选择include/pto/npu/a2a3/TSels.hpp。cmpmaskLen与数据类型宽度相关16 位类型取 4128 bit32 位类型取 264 bit。A5 路径完全基于向量寄存器vdup广播标量、plds加载打包掩码、vsel完成选择并利用pintlv_b16/punpack等位操作将 8 位打包掩码展开为谓词寄存器按elementsPerRepeatCCE_VL / sizeof(T)分块循环处理天然规避了tmp的存储开销。CPU 参考实现include/pto/cpu/TSelS.hpp 以GetTileElementOffset计算行主序偏移按(c % bitsPerElement)逐位解析掩码直接给出与硬件语义一致的标量参考便于在 CPU 仿真环境参见 docs/coding/cpu_sim_zh.md中验证逻辑正确性。与 TSEL 的对比TSELS与 TSEL 共享掩码编码与两套后端的约束框架核心区别在于选择对象TSEL在src0/src1两个 Tile 之间选择而TSELS在 Tile 与标量之间选择因此接口中多出一个scalar参数相应地A2A3 后端需要借用tmp暂存标量以构造set_cmpmask而TSEL的tmp用于保存从掩码 Tile 拷贝的逐行比较掩码其元素类型为uint32_t容量需求为每行cmpmaskLen个元素。在实现条件置零mask ? src : 0、标量饱和截断等算子时二者可按需互换使用。总结TSELS以极简的“Tile × 标量 打包掩码”语义覆盖了算子开发中最常见的条件改写场景。开发者在移植或新写代码时应重点核对目标平台的元素类型宽度与布局约束、掩码 Tile 的打包位数与容量、以及tmp在 A2A3 上“必用”、在 A5 上“仅为兼容保留”的行为差异。完整指令族与更多数据流操作可继续查阅 docs/isa 目录下的指令参考文档。【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考