Select（灵活标量位置）

产品支持情况

功能说明

给定两个源操作数src0和src1，根据selMask（用于选择的Mask掩码）的比特位值选取元素，得到目的操作数dst。选择的规则为：当selMask的比特位是1时，从src0中选取，比特位是0时从src1选取。

对于tensor高维切分计算接口，支持根据mask参数对上述选取结果，再次进行过滤，有效位填入最终的dst，无效位则保持dst原始值。例如：src0为[1,2,3,4,5,6,7,8]，src1为[9,10,11,12,13,14,15,16]，selMask为[0,0,0,0,1,1,1,1]，mask为[1,1,1,1,0,0,0,0]，dst原始值为[-1,-2,-3,-4,-5,-6,-7,-8]，则根据selMask的比特位选取后的结果dst_temp为：[9,10,11,12,5,6,7,8]，然后再根据mask进行过滤，dst的最终输出结果为[9,10,11,12,-5,-6,-7,-8]。

本选择功能支持三种模式：

• 模式0：根据selMask在两个tensor中选取元素。selMask中有效数据的个数存在限制，具体取决于源操作数的数据类型。在每一轮迭代中，根据selMask的有效位数据进行选择操作，每一轮迭代采用的selMask，均为相同数值，即selMask的有效数值。
• 模式1：根据selMask在1个tensor和1个scalar标量中选取元素，selMask无有效数据限制。多轮迭代时，每轮迭代连续使用selMask的不同部分。
• 模式2：根据selMask在两个tensor中选取元素，selMask无有效数据限制。多轮迭代时，每轮迭代连续使用selMask的不同部分。

针对模式1，提供灵活标量位置的接口。

函数原型

• tensor前n个数据计算

  1 2

  template <typename T0 = BinaryDefaultType, typename T1 = BinaryDefaultType, const BinaryConfig &config = DEFAULT_BINARY_CONFIG, typename T2, typename T3, typename T4> __aicore__ inline void Select(const T2& dst, const LocalTensor<T1>& selMask, const T3& src0, const T4& src1, SELMODE selMode, uint32_t count)

• tensor高维切分计算
  • mask逐bit模式

    1 2

    template <typename T0 = BinaryDefaultType, typename T1 = BinaryDefaultType, bool isSetMask = true, const BinaryConfig &config = DEFAULT_BINARY_CONFIG, typename T2, typename T3, typename T4> __aicore__ inline void Select(const T2& dst, const LocalTensor<T1>& selMask, const T3& src0, const T4& src1, SELMODE selMode, uint64_t mask[], uint8_t repeatTime, const BinaryRepeatParams& repeatParams)

  • mask连续模式

    1 2

    template <typename T0 = BinaryDefaultType, typename T1 = BinaryDefaultType, bool isSetMask = true, const BinaryConfig &config = DEFAULT_BINARY_CONFIG, typename T2, typename T3, typename T4> __aicore__ inline void Select(const T2& dst, const LocalTensor<T1>& selMask, const T3& src0, const T4& src1, SELMODE selMode, uint64_t mask, uint8_t repeatTime, const BinaryRepeatParams& repeatParams)

参数说明

1
2
3
4

struct BinaryConfig {
    int8_t scalarTensorIndex = 1; // 用于指定标量为LocalTensor单点元素时标量的位置，0表示左操作数，1表示右操作数
};
constexpr BinaryConfig DEFAULT_BINARY_CONFIG = {1};

1
2
3
4
5

enum class SELMODE : uint8_t {
    VSEL_CMPMASK_SPR = 0, 
    VSEL_TENSOR_SCALAR_MODE,
    VSEL_TENSOR_TENSOR_MODE,
};

返回值说明

无

约束说明

• 调用灵活标量位置接口且源操作数为LocalTensor单点元素的场景，不支持源操作数和目的操作数地址重叠。
• 操作数地址对齐要求请参见通用地址对齐约束。
• 针对Atlas 350 加速卡，int8_t/uint8_t/uint64_t/int64_t/complex32/complex64数据类型仅支持tensor前n个数据计算接口。
• 左操作数及右操作数中，必须有一个为矢量；当前不支持左右操作数同时为标量。
• 本接口传入LocalTensor单点数据作为标量时，idx参数需要传入编译期已知的常量，传入变量时需要声明为constexpr。
• 模式1场景使用灵活标量位置接口时需要填写模板参数config避免接口匹配到其他模式。

调用示例

• Select-tensor前n个数据计算样例（模式1）

  1 2 3 4 5 6 7

  // 灵活标量位置，src1Local[0]作为标量 static constexpr AscendC::BinaryConfig config = { 1 }; AscendC::Select<BinaryDefaultType, uint8_t, config>(dstLocal, maskLocal, src0Local, src1Local[0], AscendC::SELMODE::VSEL_TENSOR_SCALAR_MODE, dataSize); // 灵活标量位置，src0Local[0]作为标量 static constexpr AscendC::BinaryConfig config = { 0 }; AscendC::Select<BinaryDefaultType, uint8_t, config>(dstLocal, maskLocal, src0Local[0], src1Local, AscendC::SELMODE::VSEL_TENSOR_SCALAR_MODE, dataSize);