Reg Vector Computation Programming

Overview

Reg vector computation APIs are developed for the Regbase architecture. You can use these APIs to directly operate registers involved in vector computation on the chip, achieving greater flexibility and better performance. Reg vector computation APIs have similar functions as basic APIs. However, unlike basic APIs, whose input and output data must be LocalTensors, the input or output data of Reg vector computation APIs is Reg vector computation registers. Computation APIs are used to obtain data from a given register, perform computation, and save the result in the given register. Movement APIs are used to implement data movement between the UB and registers. Compared with basic APIs, the Reg vector computation APIs hand over the data movement and Reg computation processes to users for independent control, thereby achieving greater development flexibility.

Regbase Programming Model

The Regbase-based programming model allows users to write and call vector functions. These functions are marked with __simd_vf__ and are sent to the vector unit in the hardware for execution. Within the SIMD VF function, computation operations are implemented using Reg vector computation APIs. Figure 1 shows the memory hierarchy and programming architecture.

In the memory architecture of the SIMD Vector, the VF Reg is closest to the vector unit. It is the private memory of the SIMD and contains multiple types of Reg vector computation registers, which are used to store multiple data elements for parallel processing. All VF Reg registers in a single core share a local memory resource UB. The SIMD architecture does not support loading data from the global memory to the Reg vector computation register. Instead, data is first transferred from the global memory to the unified buffer, and then loaded to the Reg vector computation register through explicit Load/Store instructions.

Figure 1 SIMD Reg vector computation memory hierarchy

In the SIMD Reg vector computation programming architecture, instructions are sent to the Reg vector computation execution unit. The execution unit reads data from the registers, performs computation, and writes the computation result back to the registers. The DMA unit copies data between the registers and the local memory.

Figure 2 SIMD Reg vector computation programming architecture

Regbase and Membase Programming Call Hierarchy

In the Membase architecture, the basic APIs call the framework APIs or directly calls the compiler's built-in APIs to implement functions, while the high-level APIs call the basic APIs to implement functions. In the Regbase architecture, Reg vector computation APIs are added. Users can directly call these APIs during operator implementation, and the high-level APIs and basic APIs can also call these APIs to implement functions. The Reg vector computation APIs directly call the compiler's built-in APIs to implement functions.

In the Regbase architecture, intermediate results can be temporarily stored in the register, eliminating the overhead of copying data out to the local memory. In the Membase architecture, all operations are performed based on the memory. This means that data needs to be loaded from the local memory for each computation, and the result needs to be copied back to the local memory after the computation is complete. All intermediate computation results need to be temporarily stored in the local memory.

In the Regbase architecture, the maximum data length that a register can accommodate is the vector length (VL). Due to the limited capacity of the register, only data of the VL length can be processed at a time. Therefore, data needs to be split. Each time, data of the VL length is copied from the local memory to the register for computation. After the computation is complete, the result is copied back to the local memory. In the Membase architecture, LocalTensors of the complete length can be directly processed without the need for data splitting, thereby simplifying the data processing process.

Reg Vector Computation Call Hierarchy

Kernel functions are identified by __global__ __aicore__, which are entry functions on the device and can be called by the host using the <<<...>>> syntax.
__aicore__ functions are executed on the device. Kernel functions can call __aicore__ functions.
The SIMD VF function is identified by __simd_vf__ and can be called by kernel functions. Only the __simd_callee__ function and constexpr aicore can be called in the SIMD VF function.
__simd_callee__ subfunctions can be called in the SIMD VF function. These subfunctions may need to return values or pass parameters by reference. Such subfunctions are identified by __simd_callee__. Only the __simd_callee__ function and constexpr aicore function can be called in the __simd_callee__ function.

The following figure shows the specific call relationship.

The following is the only valid function call chain.

In the Regbase programming model, SIMD VF functions can be defined and marked using __simd_vf__. This design has the following advantages:

The code of __aicore__ and __simd_vf__ is clearly isolated, so that the compiler can check whether the usage of its built-in APIs is valid.
Comprehensive checks and error reporting are performed for function calls. For example, calling the __aicore__ or SIMT function within the __simd_vf__ function is an incorrect usage.
When using the __simd_vf__ function for programming, you can control certain optimization options (such as the fusion of multiple SIMD VF functions) to take effect only for specific functions, or disable certain optimization options for specific functions.

In this example, the VF function AddVF is called within Compute of the __aicore__ function to perform vector addition.

      
           template <typename T> 
__aicore__ inline void Compute()      
{  
     // Allocate an output queue and read the input result.
     ...
     // Call the SIMD VF function.
     asc_vf_call<AddVF<T>>(dstAddr, src0Addr, src1Addr, count, oneRepeatSize, repeatTimes);
     // Write the result to the output queue and free the memory of the input queue.
     ...
}

Reg Vector Computation Registers

The basic data types for Reg vector computation APIs are described as follows. For details about the APIs, see Reg Vector Computation.

RegTensor
Vector data register, which is the basic storage unit of Reg vector computation and is used for vector computation. The bit width of RegTensor is the vector length (VL), which can store VL/sizeof(T) data (T indicates the data type).
MaskReg
Mask register, which is used to select elements involved in vector computation. The bit width of MaskReg is VL/8.
UnalignRegForLoad & UnalignRegForStore
Unaligned register, which is used as a buffer to optimize the overhead of contiguous unaligned address access between the UB and RegTensor. Before reading an unaligned address, UnalignReg should be initialized through LoadUnAlignPre and then used by LoadUnAlign. When writing an unaligned address, use StoreUnAlign first and then StoreUnAlignPost for post-processing.

AddrReg

Address register, which is used to store address offsets. AddrReg is initialized through CreateAddrReg and then used to store address offsets in the loop. AddrReg is incremented based on the configured stride in each loop.

In this example, the AddVF function uses the Reg vector computation API add to add two groups of data, achieving efficient and flexible vector computation. By setting the mask register MaskReg, a mask is generated based on the actual valid data length (count) to control the number of data elements involved in the computation. The LoadAlign and StoreAlign APIs are used to copy data between the UB and the Reg vector computation register.

In this example, the scenario is contiguous aligned copy-in and copy-out. The types of used registers are RegTensor, MaskReg, and AddrReg.

        
         
           
           
             template<typename T>
__simd_vf__ inline void AddVF(__ubuf__ T* dstAddr, __ubuf__  T* src0Addr, __ubuf__ T* src1Addr, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes)
{
    AscendC::Reg::RegTensor<T> srcReg0;
    AscendC::Reg::RegTensor<T> srcReg1;
    AscendC::Reg::RegTensor<T> dstReg;
    AscendC::Reg::MaskReg mask;
    AscendC::Reg::AddrReg aReg;
    for (uint16_t i = 0; i < repeatTimes; ++i) {
        aReg = AscendC::Reg::CreateAddrReg<T>(i, oneRepeatSize);        
        mask = AscendC::Reg::UpdateMask<T>(count);
        AscendC::Reg::LoadAlign(srcReg0, src0Addr, aReg);
        AscendC::Reg::LoadAlign(srcReg1, src1Addr, aReg);
        AscendC::Reg::Add(dstReg, srcReg0, srcReg1, mask);
        AscendC::Reg::StoreAlign(dstAddr, dstReg, aReg, mask);
    }
}

            

          

        
       

In this example, the scenario is contiguous non-aligned copy-in and copy-out. The types of used registers are RegTensor, MaskReg, AddrReg, UnalignRegForLoad, and UnalignRegForStore.

        
             template <typename T>
__simd_vf__ inline void LoadUnAlignVF(__ubuf__ T* dstAddr, __ubuf__ T* srcAddr, uint32_t oneRepeatSize, uint16_t repeatTimes)
{
    AscendC::Reg::RegTensor<T> srcReg;
    AscendC::Reg::UnalignRegForLoad ureg0;
    AscendC::Reg::UnalignRegForStore ureg1;
    AscendC::Reg::AddrReg aReg;
    for (uint16_t i = 0; i < repeatTimes; ++i) {
        aReg = AscendC::Reg::CreateAddrReg<T>(i, oneRepeatSize);
        AscendC::Reg::LoadUnAlignPre(ureg0, srcAddr, aReg);
        AscendC::Reg::LoadUnAlign(srcReg, ureg0, srcAddr, aReg, 0);
        AscendC::Reg::StoreUnAlign(dstAddr, srcReg, ureg1, aReg);
    }
    AscendC::Reg::StoreUnAlignPost(dstAddr, ureg1, aReg);
}

Pipeline Synchronization Control

When writing the VF function of SIMD, you may need to write different values to the same address based on the loop, or the destination (dst) and source (src) are the same address. This involves synchronization instructions for different pipelines. The synchronization instruction between different pipelines in the SIMD VF function is represented by LocalMemBar. This synchronization instruction specifies the source pipeline (src) and destination pipeline (dst). As shown in the following figure, the destination pipeline waits for all instructions on the source pipeline to complete before execution. In the write-read scenario, if the register used by the write instruction is the same as that used by the read instruction, register order-preserving can be triggered. In this case, instructions are executed in the code sequence, and no synchronization instruction needs to be inserted. If the register used by the write instruction is different from that used by the read instruction, a synchronization instruction needs to be inserted to ensure that the two instructions are executed in sequence. The same applies to the write-write scenario.

Function prototype:

template <MemType src, MemType dst> 
__simd_callee__ inline void LocalMemBar()

How to Use the Reg Vector Computation API

The register-based programming model means that in each loop, data of a VL length is loaded from a LocalTensor to a register through a data movement instruction, and then moved out to the LocalTensor after complex mathematical computation. All computation logic is completed in the register, thereby reducing data movement between LocalTensors and greatly improving overall performance. The specific process is as follows.

Take the AddVF function as an example. First, define three vector data registers srcReg0, srcReg1, and dstReg, and a mask register mask. Each time, use the data movement function to move data of a VL length from src0 and src1 to the data registers srcReg0 and srcReg1. The address offsets are src0Addr + i × oneRepeatSize and src1Addr + i × oneRepeatSize. Then, call the Add function to store the result in dstReg (dstReg = srcReg0 + srcReg1). mask indicates the number of elements involved in the Add computation. Finally, call the data movement function to move the result from dstReg to dst.

The prototype definition of the Add function is as follows:

      
           template <typename T = DefaultType, MaskMergeMode mode = MaskMergeMode::ZEROING, typename U>
__simd_callee__ inline void Add(U& dstReg, U& srcReg0, U& srcReg1, MaskReg& mask)

The template parameter T indicates the operand data type, and MaskMergeMode indicates that the elements that are not filtered by mask are set to 0 or retain the original value in dst. The UpdateMask function is used to update the mask elements involved in the computation. Each loop consumes an element of the VL length. The LoadAlign and StoreAlign functions are used to move data in and out. LoadAlign(srcReg0, src0Addr + i * oneRepeatSize) indicates that data is moved from the LocalTensor to the srcReg0 register, with the start address being src0Addr + i * oneRepeatSize. StoreAlign(dstAddr + i * oneRepeatSize, dstReg, mask) indicates that dstReg is moved out to the LocalTensor, with the destination address being dstAddr + i * oneRepatSize. The mask parameter indicates the number of elements involved in the move-out operation.

Reg Vector Computation Programming Example

Take the Add function as an example. The macro function AddVF is marked with __simd_vf__, and such functions are also called SIMD VF functions. AddVF contains six parameters. dstAddr indicates the output data, src0Addr and src1Addr indicate the input data, __ubuf__ indicates the local memory (Unified Buffer) used for vector computation (which is the physical location where the LocalTensor is actually stored), count indicates the number of elements in the input data that participate in computation, repeatTimes indicates the number of loops, and oneRepeatSize indicates the amount of data involved in each loop. The Add function first calculates the amount of data (oneRepeatSize) that can be moved in to the register in each loop and the number of loops (repeatTimes). Then, it uses GetPhyAddr to obtain the UB addresses of the input and output data, and calls the AddVF macro function through asc_vf_call<AddVF<T>> for computation.

       
        
          
          
            // SIMD function
template <typename T> 
__simd_vf__ inline void AddVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ T* src1Addr, uint32_t count,  uint32_t oneRepeatSize, uint16_t repeatTimes) 
{     
    AscendC::Reg::RegTensor<T> srcReg0;
    AscendC::Reg::RegTensor<T> srcReg0;
    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(srcReg0, src0Addr + i * oneRepeatSize);  
        AscendC::Reg::LoadAlign(srcReg1, src1Addr + i * oneRepeatSize);
        AscendC::Reg::Add(dstReg, srcReg0, srcReg1, mask);
        AscendC::Reg::StoreAlign(dstAddr + i * oneRepeatSize, dstReg , mask);
    }
}

template <typename T> 
__aicore__ inline void Compute()      
{  
     AscendC::LocalTensor<T> dst = outQueueZ.AllocTensor<T>();     
     AscendC::LocalTensor<T> src0 = inQueueX.DeQue<T>();
     AscendC::LocalTensor<T> src1 = inQueueY.DeQue<T>();
     constexpr uint32_t oneRepeatSize = AscendC::GetVecLen()/sizeof(T);
     uint32_t count = 512;
     // Round up and compute repeat.
     uint16_t repeatTimes = AscendC::CeilDivision(count, oneRepeatSize);
     __ubuf__ T* dstAddr = (__ubuf__ T*)dst.GetPhyAddr();
     __ubuf__ T* src0Addr = (__ubuf__ T*)src0.GetPhyAddr();
     __ubuf__ T* src1Addr = (__ubuf__ T*)src1.GetPhyAddr();
     asc_vf_call<AddVF<T>>(dstAddr, src0Addr, src1Addr, count, oneRepeatSize, repeatTimes);
     outQueueZ.EnQue(dst);
     inQueueX.FreeTensor(src0);
     inQueueY.FreeTensor(src1);
}

           

         

       
      

Parent topic: SIMD Programming