Compare

函数功能

逐元素比较两个tensor大小，如果比较后的结果为真，则输出结果的对应比特位为1，否则为0。

支持多种比较模式：

• LT：小于（lower-than）
• GT：大于（greater-than）

• GE：大于或等于（greater-equal）
• EQ：等于（equal）
• NE：不等于（not-equal）
• LE：小于或等于（lower-equal）

函数原型

• 整个tensor参与计算

  1 2 3 4 5 6

  dstLocal = src0Local < src1Local; dstLocal = src0Local > src1Local; dstLocal = src0Local <= src1Local; dstLocal = src0Local >= src1Local; dstLocal = src0Local == src1Local; dstLocal = src0Local != src1Local;

  Atlas 200/500 A2推理产品暂不支持整个tensor参与计算的运算符重载。

• tensor前n个数据计算

  1 2	template <typename T, typename U> __aicore__ inline void Compare(const LocalTensor<U>& dstLocal, const LocalTensor<T>& src0Local, const LocalTensor<T>& src1Local, CMPMODE cmpMode, uint32_t calCount)

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

    1 2	template <typename T, typename U, bool isSetMask = true> __aicore__ inline void Compare(const LocalTensor<U>& dstLocal, const LocalTensor<T>& src0Local, const LocalTensor<T>& src1Local, CMPMODE cmpMode, const uint64_t mask[], uint8_t repeatTimes, const BinaryRepeatParams& repeatParams)

  • mask连续模式

    1 2	template <typename T, typename U, bool isSetMask = true> __aicore__ inline void Compare(const LocalTensor<U>& dstLocal, const LocalTensor<T>& src0Local, const LocalTensor<T>& src1Local, CMPMODE cmpMode, const uint64_t mask, uint8_t repeatTimes, const BinaryRepeatParams& repeatParams)

参数说明

返回值

无

支持的型号

Atlas 训练系列产品

Atlas推理系列产品AI Core

Atlas A2训练系列产品/Atlas 800I A2推理产品

Atlas 200/500 A2推理产品

注意事项

• 操作数地址偏移对齐要求请参见通用约束。

• dstLocal支持多种输入数据类型，内部会将其转换为uint8_t类型处理。
• dstLocal按照小端顺序排序成二进制结果，对应src中相应位置的数据比较结果。
• 使用整个tensor参与计算的运算符重载功能，src0Local和src1Local需满足256Byte对齐；使用tensor前n个数据参与计算的接口，设置calCount时，需要保证calCount个元素所占空间256字节对齐。

调用示例

本样例中，源操作数src0和src1各存储了256个float类型的数据。样例实现的功能为，逐元素对src0和src1中的数据进行比较，如果src0中的元素小于src1中的元素，dst结果中对应的比特位置1；反之，则置0。dst结果使用uint8_t类型数据存储。

本样例中只展示Compute流程中的部分代码。如果您需要运行样例代码，请将该代码段拷贝并替换样例模板中Compute函数的部分代码即可。

• 整个tensor参与计算

  1	dstLocal = src0Local < src1Local;

• tensor前n个数据计算

  1	AscendC::Compare(dstLocal, src0Local, src1Local, AscendC::CMPMODE::LT, srcDataSize);

样例模板

#include "kernel_operator.h"
class KernelCmp {
public:
    __aicore__ inline KernelCmp() {}
    __aicore__ inline void Init(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
    {
        src0Global.SetGlobalBuffer((__gm__ float*)src0Gm);
        src1Global.SetGlobalBuffer((__gm__ float*)src1Gm);
        dstGlobal.SetGlobalBuffer((__gm__ uint8_t*)dstGm);
        pipe.InitBuffer(inQueueSrc0, 1, srcDataSize * sizeof(float));
        pipe.InitBuffer(inQueueSrc1, 1, srcDataSize * sizeof(float));
        pipe.InitBuffer(outQueueDst, 1, dstDataSize * sizeof(uint8_t));
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();
    }
private:
    __aicore__ inline void CopyIn()
    {
        AscendC::LocalTensor<float> src0Local = inQueueSrc0.AllocTensor<float>();
        AscendC::LocalTensor<float> src1Local = inQueueSrc1.AllocTensor<float>();
        AscendC::DataCopy(src0Local, src0Global, srcDataSize);
        AscendC::DataCopy(src1Local, src1Global, srcDataSize);
        inQueueSrc0.EnQue(src0Local);
        inQueueSrc1.EnQue(src1Local);
    }
    __aicore__ inline void Compute()
    {
        AscendC::LocalTensor<float> src0Local = inQueueSrc0.DeQue<float>();
        AscendC::LocalTensor<float> src1Local = inQueueSrc1.DeQue<float>();
        AscendC::LocalTensor<uint8_t> dstLocal = outQueueDst.AllocTensor<uint8_t>();
 
        // 可根据实际使用接口Compare进行替换
        AscendC::Compare(dstLocal, src0Local, src1Local, AscendC::CMPMODE::LT, srcDataSize);
 
        outQueueDst.EnQue<uint8_t>(dstLocal);
        inQueueSrc0.FreeTensor(src0Local);
        inQueueSrc1.FreeTensor(src1Local);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<uint8_t> dstLocal = outQueueDst.DeQue<uint8_t>();
        AscendC::DataCopy(dstGlobal, dstLocal, dstDataSize);
        outQueueDst.FreeTensor(dstLocal);
    }
private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrc0, inQueueSrc1;
    AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueueDst;
    AscendC::GlobalTensor<float> src0Global, src1Global;
    AscendC::GlobalTensor<uint8_t> dstGlobal;
    uint32_t srcDataSize = 256;
    uint32_t dstDataSize = srcDataSize / AscendC::AscendCUtils::GetBitSize(sizeof(uint8_t));
};
extern "C" __global__ __aicore__ void main_cpu_cmp_sel_demo(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
{
    KernelCmp op;
    op.Init(src0Gm, src1Gm, dstGm);
    op.Process();
}