Operator Implementation

This sample aims to demonstrate the SIMD and SIMT programming model through a simple operator implementation, rather than providing the best practice of the operator functionality.

The following figure shows the process for implementing vector operator kernel functions based on the SIMD and SIMT programming model.

Figure 1 SIMD and SIMT kernel function implementation process

The following uses the gather & adds operator as an example to detail the preceding steps. The gather & adds operator extracts 8,192 pieces of data with specified indexes from a 1D vector with a length of 100,000 and adds 1 to each piece of extracted data. For details about the complete code of the operator described in this sample, see the sample of implementing the gather&adds operator using hybrid programming of SIMD and SIMT.

Operator Analysis

The operator analysis procedure is as follows:

  1. Specify the input and output of the operator.

    • The gather & adds operator has two inputs: input and index. input indicates the original data, and index indicates the index of the data to be obtained in input. The output is output.
    • In this sample, the supported data type of input is float, the supported data type of index is uint32_t, and the data type of output is the same as that of input.
    • The shape supported by input is [100000]. The shape supported by index is [8192], and the data value of index is within the range of [0, 100000). The shape of output is the same as that of index.
    • The supported input format is ND.

  2. Specify the mathematical expression and compute logic of the operator.

    The i-th piece of data in output of the gather & adds operator is computed as follows:

    output[i] = input[index[i]] + 1

    Compute logic:

    • Use the SIMT programming model to obtain the data with specified indexes from input (Global Memory) and store the data to the Unified Buffer.
    • Use the SIMD programming model to add 1 to the data in the on-chip memory (Unified Buffer).
    • Move the compute result in the Unified Buffer to the external storage (Global Memory).
    Figure 2 Operator compute logic

    The operation of adding 1 in simd_adds can be quickly implemented in the simt_gather function. This sample aims to demonstrate the SIMD and SIMT programming model through a simple operator implementation, rather than providing the best practice of the operator functionality.

  3. Define the kernel function name and parameters.

    • In this sample, the kernel function is named gather_and_adds_kernel.
    • Based on the operator input and output analysis, the kernel function has five parameters: input, index, output, input_total_length, and index_total_length. input and index indicate the memory addresses of the inputs in the Global Memory, output indicates the memory address of the output in the Global Memory, input_total_length indicates the data length of input, and index_total_length indicates the data length of index and also the data length of output.

  4. Specify the core allocation policy, SIMT thread configuration, and the number of cyclic call times of the SIMD Reg vector compute API.

    In this sample, the shape of the operator input index is 8192. You can set the number of cores to 8 and the amount of data processed by each core to 1,024.

    For SIMT implementation, the number of threads can be set to 1,024, and each thread processes one piece of data. A single core needs to call the simt_gather function only once to complete the gather operation.

    For SIMD Reg vector compute implementation, the amount of data processed by a single core is 1,024. The length of data processed by the Reg vector compute API at a time (one_repeat_size) is GetVecLen/sizeof(float), and the number of cyclic call times of the API (repeat_times) is 1024/one_repeat_size.

  5. Define the SIMT VF function name and parameters.

    • In this sample, the SIMT VF function is named simt_gather.
    • According to the SIMT thread configuration policy, the SIMT VF function has six parameters: input, index, gather_output, input_total_length, index_total_length, and output_total_length. input and index indicate the memory addresses of the inputs in the Global Memory, gather_output indicates the memory address of the output in the Unified Buffer, input_total_length indicates the data length of input, index_total_length indicates the data length of index, and output_total_length indicates the data length of gather_output on a single core.

  6. Define the SIMD VF function name and parameters.

    • In this sample, the SIMD VF function is named simd_adds.
    • According to the preceding SIMD policy, the SIMD VF function has five parameters: output, input, count, one_repeat_size, and repeat_times. output indicates the memory address of the output in the Unified Buffer, input indicates the memory address of the input in the Unified Buffer, count indicates the total amount of data processed by a single core, one_repeat_size indicates the amount of data processed in a single call, and repeat_times indicates the number of cyclic call times of the Reg vector compute API.

Based on the preceding analysis, the design specifications of the Ascend C gather & adds operator are as follows.

Kernel Function Definition and Implementation

Define kernel functions based on the rules described in Kernel Function.

  1. Define the function prototype.

    In this sample, the function name is gather_and_adds_kernel (the kernel function name can be customized). Based on the preceding analysis, the function prototype is defined as follows:

    1
2
3
    		 
    __global__ __aicore__ void gather_and_adds_kernel(__gm__ float* input, __gm__ uint32_t* index, __gm__ float* output, uint32_t input_total_length, uint32_t index_total_length)
{
}

  2. Start the SIMT VF function simt_gather to obtain the data with specified indexes from input.

    1. Compute the amount of data to be processed by a single core. index_total_length indicates the total data amount. Divide it by the number of cores to determine the amount of data to be processed by a single core. 
      1
      		 
      uint32_t index_total_length_per_block = index_total_length / AscendC::GetBlockNum();
    2. Call Alloc to allocate the Unified Buffer memory space and use the tensor as the output of the simt_gather function.
    3. Call asc_vf_call to start the SIMT VF function simt_gather. The first parameter is of the dim3 structure, indicating the 3D hierarchy of threads. In this example, the parameter is initialized to dim3(1024). The 1D definition mode is used, and the total number of threads is 1,024.
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
    		 
    constexpr uint32_t THREAD_COUNT = 1024;

__global__ __aicore__ void gather_and_adds_kernel(__gm__ float* input, __gm__ uint32_t* index, __gm__ float* output, uint32_t input_total_length, uint32_t index_total_length)
{
    // Set the kernel type to AIV_ONLY.
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    
    // Define the UB memory allocation object.
    AscendC::LocalMemAllocator<AscendC::Hardware::UB> ub_allocator;

    // Compute the amount of data to be processed by a single core.
    uint32_t index_total_length_per_block = index_total_length / AscendC::GetBlockNum();
    // Allocate UB memory for the output of simt_gather.
    AscendC::LocalTensor<float> gather_output = ub_allocator.Alloc<float>(index_total_length_per_block);
    // 1. Call the SIMT function to obtain 1,024 pieces of data with specified indexes.
    asc_vf_call<simt_gather>(dim3(THREAD_COUNT), input, index,
                             (__ubuf__ float *)gather_output.GetPhyAddr(),
                             input_total_length,
                             index_total_length,
                             index_total_length_per_block);

    // 2. Call the SIMD function to add 1 to the data.
    ...

    // 3. Move the data to the GM.
    ...
}

  3. Start the SIMD VF function simd_adds to add 1 to the value of data in the Unified Buffer.

    1. Call Alloc to allocate the Unified Buffer memory space and use the tensor as the output of the simt_adds function.
    2. Call GetVecLen and divide the return value by the length of a single piece of data to compute the amount of data processed by the Reg vector compute API at a time (one_repeat_size). Divide the amount of data to be processed by a single core (index_total_length_per_block) by the amount of data processed at a time (one_repeat_size) to compute the number of cyclic call times of the Reg vector compute API.
    3. Call asc_vf_call to start the SIMD VF function simd_adds.
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
    		 
    __global__ __aicore__ void gather_and_adds_kernel(__gm__ float *input, __gm__ uint32_t *index, __gm__ float *output, uint32_t input_total_length, uint32_t index_total_length)
{
    // 1. Call the SIMT function to obtain 1,024 pieces of data with specified indexes.
    ...

    // Allocate UB memory for the output of simd_adds.
    AscendC::LocalTensor<float> adds_output = ub_allocator.Alloc<float>(index_total_length_per_block);
    // Compute the amount of data processed by the Reg vector compute API per call.
    constexpr uint32_t one_repeat_size = AscendC::GetVecLen() / sizeof(float);
    // Compute the number of cyclic call times of the Reg vector compute API.
    uint16_t repeat_times = (index_total_length_per_block + one_repeat_size - 1) / one_repeat_size;
    // 2. Call the SIMD function to add 1 to the data.
    asc_vf_call<simd_adds>((__ubuf__ float *)adds_output.GetPhyAddr(),
        (__ubuf__ float *)gather_output.GetPhyAddr(), index_total_length_per_block, one_repeat_size, repeat_times);

    // Depend on the synchronization between the PIPE_V and PIPE_MTE3 pipelines.
    AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(0);
    AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(0);

    // 3. Move the data to the GM.
    ...
}

  4. Call DataCopy to move the result data to the Global Memory. 

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
    		 
    __global__ __aicore__ void gather_and_adds_kernel(__gm__ float* input, __gm__ uint32_t* index, __gm__ float* output, uint32_t input_total_length, uint32_t index_total_length)
{
    // 1. Call the SIMT function to obtain 1,024 pieces of data with specified indexes.
    ...

    // 2. Call the SIMD function to add 1 to the data.
    ...

    // 3. Move the data to the GM.
    // Define a GlobalTensor object for data movement.
    AscendC::GlobalTensor<float> output_global_tensor;
    // Initialize the GlobalTensor address based on the core offset.
    output_global_tensor.SetGlobalBuffer(output + index_total_length_per_block * AscendC::GetBlockIdx());
    // Call the data movement API to move the data to the GM.
    AscendC::DataCopy(output_global_tensor, adds_output, index_total_length_per_block);
}

SIMT VF Function Definition and Implementation

  1. Define the function prototype.

    Define the SIMT VF function prototype based on the preceding parameter analysis of the SIMT VF function. Use the __simt_vf__ function type qualifier to identify the entry of the SIMT VF kernel function so that it can be called by asc_vf_call.

    In SIMT programming, __launch_bounds__(thread_num) is optional and is used to specify the maximum number of threads that can be started by the kernel function during compilation. If this parameter is not configured, the default value of thread_num is 1024. Ensure that the value of thread_num ≥ x * y * z (that is, the first parameter of asc_vf_call: dim3{x, y, z}). The value of thread_num ranges from 1 to 2,048. The maximum number of threads determines the number of registers that can be allocated to each thread. For details about the mapping, see Table 5. Registers are used to store local variables in threads. If the number of local variables exceeds the number of registers, issues such as stack overflow may occur.

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
    		 
    constexpr uint32_t THREAD_COUNT = 1024;

__simt_vf__ __launch_bounds__(THREAD_COUNT) inline void simt_gather(
    __gm__ float* input,
    __gm__ uint32_t* index,
    __ubuf__ float* gather_output,
    uint32_t input_total_length,
    uint32_t index_total_length,
    uint32_t output_total_length)
{
}

  2. Implement the function.

    The simt_gather function implementation obtains data with specified indexes from input (Global Memory). Based on the preceding data tiling policy, compute the index of the data to be processed by a thread first, and then store the data to the Unified Buffer through assignment operations.

    In this example, the number of cores is set to 8, the thread hierarchy is {1024, 1, 1}, and the total amount of data is 8192 (8 × 1024). Each thread needs to process only one piece of data. The compute logic of the index of data to be processed in index is as follows: Current core ID × Number of threads per core + ID of the current thread. The code is as follows: 
    1
    		 
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    blockIdx is used to obtain the current core ID. blockDim is used to obtain the 3D thread hierarchy {x, y, z}. In this example, the hierarchy is {1024, 1, 1}, where both the second and third dimensions are 1. A 1D hierarchy is used, so the number of threads can be expressed as blockDim.x. threadIdx is used to obtain the 3D thread index {x, y, z}. In this example, only the first dimension x is used. You can obtain the ID of the current thread through threadIdx.x.

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
    		 
    __simt_vf__ __launch_bounds__(THREAD_COUNT) inline void simt_gather(
    __gm__ float* input,
    __gm__ uint32_t* index,
    __ubuf__ float* gather_output,
    uint32_t input_total_length,
    uint32_t index_total_length,
    uint32_t output_total_length)
{
    // Check for exceptions to prevent out-of-bounds access.
    if (threadIdx.x >= output_total_length) {
        return;
    }

    // Compute the index of the data to be processed by the thread in index.
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    // Check for exceptions to prevent out-of-bounds access.
    if (idx >= index_total_length) {
        return;
    }

    // Read the index of the data to be obtained by the thread in input.
    uint32_t gather_idx = index[idx];
    // Check for exceptions to prevent out-of-bounds access.
    if (gather_idx >= input_total_length) {
        return;
    }

    // Store the data whose index is gather_idx in input to the UB.
    gather_output[threadIdx.x] = input[gather_idx];
}

SIMD VF Function Definition and Implementation

  1. Define the function prototype.

    Define the SIMD VF function prototype based on the preceding parameter analysis of the SIMD VF function. Use the __simd_vf__ function type qualifier to identify the SIMD VF entrypoint function so that it can be called by the asc_vf_call keyword. 
    1
2
3
4
    		 
    __simd_vf__ inline void simd_adds(__ubuf__ float *output, __ubuf__ float *input,
    uint32_t count, uint32_t one_repeat_size, uint16_t repeat_times)
{
}

  2. Call the Reg vector computation API for repeat_times times to perform the operations of adding 1.

    1. Call Contiguous Aligned Data Move-in to move data from the Unified Buffer to the basic Reg vector compute unit RegTensor.
    2. Call Adds to add 1 to the data.
    3. Call Contiguous Aligned Data Move-out to move the data from RegTensor to the Unified Buffer.
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
    		 
     constexpr float ADDS_ADDEND = 1.0f;

__simd_vf__ inline void simd_adds(__ubuf__ float *output, __ubuf__ float *input,
    uint32_t count, uint32_t one_repeat_size, uint16_t repeat_times)
{
    // Initialize the Reg vector compute unit as the source operand.
    AscendC::Reg::RegTensor<float> src_reg0;
    // Initialize the Reg vector compute unit as the destination operand.
    AscendC::Reg::RegTensor<float> dst_reg0;
    // Initialize the Reg vector compute mask register.
    AscendC::Reg::MaskReg mask;

    for (uint16_t i = 0; i < repeat_times; i++) {
        // Move the data from the UB to the basic Reg vector compute unit.
        mask = AscendC::Reg::UpdateMask<float>(count);
        AscendC::Reg::LoadAlign(src_reg0, input + i * one_repeat_size);
        // Call Adds to add 1 to the data.
        AscendC::Reg::Adds(dst_reg0, src_reg0, ADDS_ADDEND, mask_reg);
        // Move the data from the basic Reg vector compute unit to the UB.
        AscendC::Reg::StoreAlign(output + i * one_repeat_size, dst_reg0, mask_reg);
    }
}

Parent topic: SIMD and SIMT Operator Implementation