Axpy

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	x
Atlas A2 training product/Atlas A2 inference product	x
Atlas 200I/500 A2 inference product	x
Atlas inference product AI Core	x
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

Performs the multiply-add operation on the input data dstReg, srcReg, and scalarValue by element based on mask, and writes the result to dstReg. The formula is as follows:

$\text{[math]}$

Prototype

template <typename T = DefaultType, typename U, MaskMergeMode mode = MaskMergeMode::ZEROING, typename S>
__simd_callee__ inline void Axpy(S& dstReg, S& srcReg, const U scalarValue, MaskReg& mask)

Parameters

**Table 1** Template parameters
Parameter	Description
T	Operand data type. For the Atlas 350 Accelerator Card, the supported data types are half, float, uint64_t, and int64_t.
U	Data type of the scalar source operand. For the Atlas 350 Accelerator Card, the supported data types are half, float, uint64_t, and int64_t.
mode	Set it to MERGING or ZEROING. Currently, only ZEROING is supported. ZEROING: The elements that are not filtered by mask are set to zero in dst. Only this mode is supported.
S	RegTensor type of the destination operand, for example, RegTensor<half>. It is automatically inferred by the compiler and does not need to be specified.

**Table 2** Parameters
Parameter	Input/Output	Description
dstReg	Output	Destination operand and source operand. The type is RegTensor.
srcReg	Input	Source operands. The type is RegTensor. Both source operands must have the same data type as the destination operand.
scalarValue	Input	Source operands. The type is scalar. Both source operands must have the same data type as the destination operand.
mask	Input	Valid indication of the source operand element operation. For details, see MaskReg.

Returns

None

Constraints

dstReg and srcReg can be the same RegTensor.

Example

template<typename T, typename U>
static __simd_vf__ inline void AxpyVF(__ubuf__ T* dstAddr, __ubuf__ T* srcAddr, U scalarValue, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes)
{
    AscendC::Reg::RegTensor<T> srcReg;
    AscendC::Reg::RegTensor<T> dstReg;
    AscendC::Reg::MaskReg mask;    
    for (uint16_t i = 0; i < repeatTimes; i++) {
        mask = AscendC::Reg::UpdateMask<T>(count);
        AscendC::Reg::LoadAlign(srcReg, srcAddr + i * oneRepeatSize);
        AscendC::Reg::LoadAlign(dstReg, dstAddr + i * oneRepeatSize);
        AscendC::Reg::Axpy(dstReg, srcReg, scalarValue, mask);
        AscendC::Reg::StoreAlign(dstAddr + i * oneRepeatSize, dstReg, mask);
    }
}

Parent topic: Compound Computation