TSELS¶
指令示意图¶
简介¶
使用mask tile在源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}
\]
汇编语法¶
同步形式:
%dst = tsels %mask, %src, %scalar : !pto.tile<...>
AS Level 1(SSA)¶
%dst = pto.tsels %mask, %src, %scalar : (!pto.tile<...>, !pto.tile<...>, dtype) -> !pto.tile<...>
AS Level 2(DPS)¶
pto.tsels ins(%mask, %src, %scalar : !pto.tile_buf<...>, !pto.tile_buf<...>, dtype) outs(%dst : !pto.tile_buf<...>)
C++内建接口¶
声明于 include/pto/common/pto_instr.hpp:
公共包含头为
<pto/pto-inst.hpp>,内部声明位于pto/common/pto_instr.hpp。
template <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);
约束¶
- 实现检查 (Atlas A2/A3 训练系列产品/Atlas A2/A3 推理系列产品):
sizeof(TileDataDst::DType)必须是2或4字节。- 支持的数据类型为 2 或 4 字节类型:
int16_t、uint16_t、int32_t、uint32_t、half、bfloat16_t、float。 dst和src必须使用相同的元素类型。dst和src必须是行主序。- 运行时:
src.GetValidRow()/GetValidCol()必须与dst.GetValidRow()/GetValidCol()一致。
- 实现检查 (Ascend 950PR/Ascend 950DT):
sizeof(TileDataDst::DType)可以是1、2或4字节。- 支持的数据类型为
int8_t、uint8_t、int16_t、uint16_t、int32_t、uint32_t、half和float。 dst和src必须使用相同的元素类型。dst、mask和src必须是行主序。- 运行时:
src.GetValidRow()/GetValidCol()必须与dst.GetValidRow()/GetValidCol()一致。
- 有效区域:
- 该操作使用
dst.GetValidRow()/dst.GetValidCol()作为迭代域。
- 该操作使用
- 掩码编码:
- 掩码Tile被解释为目标定义布局中的打包谓词位。
临时空间¶
Atlas A2/A3 训练系列产品/Atlas A2/A3 推理系列产品¶
tmp 被使用作为小型缓冲区,用于存储 set_cmpmask 操作所需的标量值并保存比较掩码。在选择循环之前,标量被写入 tmp[0]。
tmp的元素类型必须与TileDataSrc::DType一致。tmp大小要求:至少1个元素(用于存放标量)。典型声明:Tile<TileType::Vec, float, 1, 16>或类似。
Ascend 950PR/Ascend 950DT¶
tmp 被接口接受但Ascend 950PR/Ascend 950DT实现不使用。Ascend 950PR/Ascend 950DT后端使用 vdup 将标量广播到向量寄存器,并使用 vsel 进行选择,不需要暂存Tile存储。tmp 仅为了与Atlas A2/A3 训练系列产品/Atlas A2/A3 推理系列产品的API兼容性而保留在C++内建接口签名中。
示例¶
自动(Auto)¶
#include <pto/pto-inst.hpp>
using namespace pto;
void example_auto() {
using TileDst = Tile<TileType::Vec, float, 16, 16>;
using TileSrc = Tile<TileType::Vec, float, 16, 16>;
using TileTmp = Tile<TileType::Vec, float, 16, 16>;
using TileMask = Tile<TileType::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);
}
手动(Manual)¶
#include <pto/pto-inst.hpp>
using namespace pto;
void example_manual() {
using TileDst = Tile<TileType::Vec, float, 16, 16>;
using TileSrc = Tile<TileType::Vec, float, 16, 16>;
using TileTmp = Tile<TileType::Vec, float, 16, 16>;
using TileMask = Tile<TileType::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);
}
汇编示例(ASM)¶
自动模式¶
# 自动模式:由编译器/运行时负责资源放置与调度。
%dst = pto.tsels %mask, %src, %scalar : (!pto.tile<...>, !pto.tile<...>, dtype) -> !pto.tile<...>
手动模式¶
# 手动模式:先显式绑定资源,再发射指令。
# 可选(当该指令包含 tile 操作数时):
# pto.tassign %arg0, @tile(0x1000)
# pto.tassign %arg1, @tile(0x2000)
%dst = pto.tsels %mask, %src, %scalar : (!pto.tile<...>, !pto.tile<...>, dtype) -> !pto.tile<...>
PTO汇编形式¶
%dst = tsels %mask, %src, %scalar : !pto.tile<...>
# AS Level 2 (DPS)
pto.tsels ins(%mask, %src, %scalar : !pto.tile_buf<...>, !pto.tile_buf<...>, dtype) outs(%dst : !pto.tile_buf<...>)