Operator Implementation

Operator Analysis and Kernel Function Definition

The operator analysis procedure is as follows:

1. Specify the function and compute logic of the operator.

   The gather operator obtains data of specified index rows from the input tensor. That is, the operator obtains m rows of data of a specified index from the two-dimensional vector input with the shape of M × N. The row index of the m rows is specified by the input index. The formula for computing the data in the i row of the operator output is as follows:

   output[i] = input[index[i]]

2. Specify the input and output of the operator.
   • The gather operator has two inputs, input and index, and the output is output.
   • In this example, the supported operator input data types are float, half, and int32_t, and the data type of index is uint32_t. The operator output data type is the same as the input data type.
   • Each thread processes one row of data. The length of each row of data (in_width) and the total number of rows to be processed (index_total_length) need to be passed to ensure that the tail thread does not perform invalid operations.
   • During operator implementation, you do not need to use a large number of temporary variables. To improve performance, you can increase the maximum number of threads of the kernel function based on the default maximum number of threads (1024).

3. Specify the function name and parameters.
   • Customize the kernel function name. In this example, the kernel function is named gather_custom.
   • By analyzing operator input and output, you can use template parameters to support different types of input and output data.

     Template Parameter	Type	Definition
     type_data	typename	Data type of the input/output.
     type_idx	typename	Data type of the index.

     The input parameters of the function are defined as follows:

     Parameter	Type	Definition
     input	type_data*	Memory address of the input data in the Global Memory
     index	type_idx*	Memory address of the index data in the Global Memory
     gather_output	type_data*	Memory address of the output data in the Global Memory
     in_width	uint32_t	Length of the second dimension of the input data (column width)
     index_total_length	uint32_t	Total length of the index data

4. Specify the solution for setting dynamic parameters (such as gridDim and blockDim) of the SIMT kernel function.
   • In this example, the even splitting solution is used. Based on the number of available cores and the maximum number of threads, compute and adjust gridDim (number of enabled thread blocks) and blockDim (number of threads enabled by a thread block). Ensure that the value of gridDim does not exceed 65535 and that of blockDim does not exceed the maximum number of threads (2048).
   • The dynamic UB space is not required in the implementation logic of this operator.

Based on the preceding analysis, the design specifications of the SIMT Gather operator are as follows:

• Operator type (OpType): Gather
• Operator input and output:

  Table 1 Input and output specifications of the Gather operator

  name	shape	data type	format
  input (input)	(M, N)	float/half/int32_t	ND
  index (input)	(m), m < M	uint32_t	ND
  output (output)	(m, N)	float/half/int32_t	ND

• Kernel function name: gather_custom

The kernel function is defined as follows:

1
2
3
4
5
6
7
8
9

constexpr uint32_t MAX_THREAD_COUNT = 2048;

template <typename type_data, typename type_idx>
__global__ __launch_bounds__(MAX_THREAD_COUNT) void gather_custom(
    type_data* input,
    type_idx* index,
    type_data* gather_output,
    uint32_t in_width,
    uint32_t index_total_length)

Thread Splitting Compute on the Host

This example uses a simple even splitting solution to describe how to compute dynamic splitting parameters.

1. Set the initial value of gridDim.
   If the value of gridDim is set to a value less than the actual number of AIV cores, idle cores will be wasted. Therefore, set the initial value of gridDim to the actual number of AIV cores of the current chip. To obtain the number of AIVs, perform the following steps:

   1 2 3 4

   uint32_t real_core_num = 0; const auto& platformInfoMgr = platform_ascendc::PlatformAscendCManager::GetInstance(); real_core_num = platformInfoMgr->GetCoreNumAiv(); block_num = real_core_num; // block_num indicates the initial value of gridDim.

2. Compute blockDim.

   Compute the number of threads enabled for a thread block (blockDim) based on the length of the input index (index_total_length) and the initial value of gridDim.

   // thread_num_per_block is the value of blockDim.
   thread_num_per_block = (index_total_length + block_num - 1) / block_num;

3. Adjust blockDim.

   If blockDim exceeds the maximum number of threads, adjust blockDim to the maximum number of threads.

   1 2 3

   if (thread_num_per_block > MAX_THREAD_COUNT) { thread_num_per_block = MAX_THREAD_COUNT; }

4. Adjust gridDim.

   Recompute gridDim and ensure gridDim × blockDim > index_total_length, that is, ensure that all enabled threads can process data of a specified number of rows.

   1	block_num = (index_total_length + thread_num_per_block - 1) / thread_num_per_block;

The complete splitting compute code is as follows:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

constexpr uint32_t MAX_THREAD_COUNT = 2048;
constexpr uint32_t MAX_BLOCK_COUNT = 65535;

bool block_split(uint32_t index_total_length, uint32_t &block_num, uint32_t &thread_num_per_block) {
    uint32_t real_core_num = 0;
    const auto& platformInfoMgr = platform_ascendc::PlatformAscendCManager::GetInstance();
    if (platformInfoMgr == nullptr) {
        std::cout << "[ERROR] Get platform info failed, please check device status."<< std::endl;
        return false;
    }
    real_core_num = platformInfoMgr->GetCoreNumAiv();
    block_num = real_core_num;
    thread_num_per_block = (index_total_length + block_num -1) / block_num;
    if (thread_num_per_block > MAX_THREAD_COUNT) {
        thread_num_per_block = MAX_THREAD_COUNT;
        block_num = (index_total_length + thread_num_per_block - 1) / thread_num_per_block;
        if (block_num > MAX_BLOCK_COUNT) {
        std::cout << "[ERROR] index_total_length: "<< index_total_length << " can not be bigger than "
            << MAX_THREAD_COUNT * MAX_BLOCK_COUNT<< "."<< std::endl;
        return false;
        }
    }
    return true;
}

Operator Implementation on the Kernel

1. Obtain the position offset of the current thread based on the even splitting algorithm.

   In this operator, only the first dimension of the thread dimensions such as gridDim and blockDim is used. Therefore, only the x dimension information needs to be considered during offset compute. As shown in the following code, threadIdx indicates the index of a thread in the thread block where the thread is located, blockDim indicates the number of threads set in a thread block, and blockIdx indicates the index of a thread block.

   1 2	// Compute the thread index. int32_t out_row = blockIdx.x * blockDim.x + threadIdx.x;

2. Based on the thread index, obtain the row index of the data to be processed by the current thread, and compute the corresponding input and output position offsets to collect data of the entire row.

   1 2 3 4 5 6 7 8

   uint32_t in_row = index[out_row]; int input_idx = in_row * in_width; int output_idx = out_row * in_width; for (int32_t col = 0; col < in_width; col++) { gather_output[output_idx] = input[input_idx]; input_idx += 1; output_idx += 1; }

 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 MAX_THREAD_COUNT = 2048;
constexpr uint32_t MAX_BLOCK_COUNT = 65535;

template <typename type_data, typename type_idx>
__global__ __launch_bounds__(MAX_THREAD_COUNT) void gather_custom(
    type_data* input,
    type_idx* index,
    type_data* gather_output,
    uint32_t in_width,
    uint32_t index_total_length)
{
    // Calculate global thread ID
    int32_t out_row = blockIdx.x * blockDim.x + threadIdx.x;
    // Maps to the row index of output tensor
    if (out_row >= index_total_length) {
        return;
    }
    // Single thread processes entire row (all columns) - enables coalesced memory access
    uint32_t in_row = index[out_row];
    int input_idx = in_row * in_width;
    int output_idx = out_row * in_width;
    for (int32_t col = 0; col < in_width; col++) {
        gather_output[output_idx] = input[input_idx];
        input_idx += 1;
        output_idx += 1;
    }
}

Runtime Verification

After the kernel function (operator kernel program) is developed, you can compile the kernel function calling program on the host to call the operator from the application on the host for runtime verification.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

std::vector<float> gather(std::vector<float>& input, const uint32_t* in_shape, std::vector<uint32_t>& index)
{
    ...
    // Compute the splitting parameters and set the dynamic UB memory.
   
    uint32_t block_num = 0;
    uint32_t thread_num_per_block = 0;
    block_split(index_total_length, block_num, thread_num_per_block))
    
    ...
    // Compute the splitting parameter and set the dynamic UB memory.
    uint32_t dyn_ubuf_size = 0;  // No need to alloc dynamic memory.
    
    // Use the memory launch symbol <<<...>>> to call the kernel function to complete specified operations.
    
    gather_custom<<<block_num, thread_num_per_block, dyn_ubuf_size, stream>>>(
              input_device, index_device, output_device, in_shape[1], index_total_length);
    ...
}