​作者​：昇腾实战派
​知识地图​：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

背景概述

Gather 算子是深度学习推理中常用的数据重排操作，用于从输入张量中按指定索引收集数据。在 MoE（混合专家）路由、Embedding 查找等场景中，Gather 的性能直接影响端到端推理延迟。

本文从基础实现出发，逐步优化至支持多维输入和 batch_dims 的通用版本，并分享在 Atlas 800I A3 上的性能调优经验，帮助开发者掌握 Ascend C 算子开发与优化的核心方法。

实现 basic gather

方案设计

  • 算子功能：
    根据torch.gather的定义，为方便理解，我们先实现gather算子最基础的功能，从固定二维输入张量中采集指定的m行数据，先不考虑多维input和batch_dims批量处理功能，数据切分也先设置为固定值。

    在本案例中。gather算子实现了从shape=(100000 * 128)的二维向量中获取指定索引的12288行数据的功能。算子输出output第i行数据计算公式为：

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

  • 算子规格：

    算子类型(OpType) gather
    算子输入 name shape data type format
    input 100000 * 128 float ND
    index 12288 uint32_t ND
    算子输出 output 12288 * 128 float ND
    核函数名 gather_custom

  • 数据切分：

    • 核数：48核
    • 每核线程数：256线程
    • 单线程处理：1行 x 128列
    • 总处理能力：48×256=12288行（覆盖索引长度）

  • 算子实现：
    gather算子的实现流程为从输入input（Global Memory）中获取指定索引的数据。基于上述数据切分，首先计算线程应处理数据的索引，然后通过赋值操作将一行数据存储到Global Memory上。

  • 调用实现：
    使用内核调用符<<<>>>调用核函数。

tiling设计

  1. 特定场景的优化：

    • 因为输入数据是二维的，且 Gather 操作是针对行进行的（即 axis=0）。
    • 这种假设让代码可以跳过复杂的多维坐标转换逻辑，直接操作内存地址。

  2. Grid-Stride Loop 模式：

    • 线程分工：每个线程负责输出矩阵中的一整行。
    • 循环处理：核函数中使用了 if (out_row >= index_total_length) 保护，并结合线程 ID 直接映射到输出行。代码无需显式的 for(stride) 循环，通过配置 numBlocks 和 thread_num_per_block，隐式地覆盖了所有需要处理的行。

Host侧代码实现如下：

/**
 * @brief Host 端主函数 (负责内存管理与核函数启动)
 */
std::vector<float> gather(std::vector<float>& input, const uint32_t* in_shape, std::vector<uint32_t>& index)
{
    // 1. 初始化变量
    uint32_t input_total_length = input.size(); //获取容器中元素的个数
    size_t input_total_byte_size = input_total_length * sizeof(float);

    uint32_t index_total_length = index.size();
    size_t index_total_byte_size = index_total_length * sizeof(uint32_t);

    // 2. 计算输出形状
    // 输出矩阵的行数 = 索引数组的长度
    // 输出矩阵的列数 = 输入矩阵的列数
    uint32_t output_total_length = index.size() * in_shape[1];
    size_t output_total_byte_size = output_total_length * sizeof(float);

    // 3. Ascend C 运行时环境初始化
    int32_t device_id = 0;
    aclrtStream stream = nullptr;

    // 指针声明
    uint8_t* input_host = reinterpret_cast<uint8_t *>(input.data());
    uint8_t* index_host = reinterpret_cast<uint8_t *>(index.data());
    uint8_t* output_host = nullptr; // 用于存放结果的 Host 内存
    float* input_device = nullptr;   // Device 上的 Input
    uint32_t* index_device = nullptr; // Device 上的 Index
    float* output_device = nullptr; // Device 上的 Output

    // 4. 初始化 Device 并创建流 (Stream)
    aclInit(nullptr);
    aclrtSetDevice(device_id);
    aclrtCreateStream(&stream);

    // 5. 内存分配与数据传输 (Malloc & Memcpy)
    // 分配 Host 端内存 (用于接收结果)
    aclrtMallocHost((void **)(&output_host), output_total_byte_size);

    // 分配 Device 端内存 (HUGE_FIRST: 优先分配大页内存，减少缺页中断)
    aclrtMalloc((void **)&input_device, input_total_byte_size, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void **)&index_device, index_total_byte_size, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void **)&output_device, output_total_byte_size, ACL_MEM_MALLOC_HUGE_FIRST);

    // 数据拷贝 (Host to Device)
    aclrtMemcpy(input_device, input_total_byte_size, input_host, input_total_byte_size, ACL_MEMCPY_HOST_TO_DEVICE);
    aclrtMemcpy(index_device, index_total_byte_size, index_host, index_total_byte_size, ACL_MEMCPY_HOST_TO_DEVICE);

    // 6. 计算 Kernel 启动参数 (Launch Configuration)
    // 这是一个固定的配置，通常基于 Ascend 芯片的物理核心数(Cores)或计算单元(CU)进行调优
    uint32_t numBlocks = 48;         // 启动 48 个 Block
    uint32_t thread_num_per_block = 256; // 每个 Block 256 个线程
    uint32_t dyn_ubuf_size = 0;      // 无需动态共享内存

    // 7. 核函数调用 (Kernel Launch)
    // <<<Grid, Block, SharedMem, Stream>>>
    gather_custom<<<numBlocks, thread_num_per_block, dyn_ubuf_size, stream>>>(
              input_device, 
              index_device, 
              output_device, 
              in_shape[1],           // 传入列数 (in_width)
              index_total_length     // 传入索引行数
    );

    // 8. 同步与结果回传
    aclrtSynchronizeStream(stream); // 等待 GPU 计算完成

    // Device to Host 拷贝
    aclrtMemcpy(output_host, output_total_byte_size, output_device, output_total_byte_size, ACL_MEMCPY_DEVICE_TO_HOST);

    // 构造返回的 Vector
    std::vector<float> output((float *)output_host, (float *)(output_host + output_total_byte_size));

    // 9. 资源释放 (Cleanup)
    // 遵循 RAII 原则，最后释放所有申请的资源
    aclrtFree(input_device);
    aclrtFree(index_device);
    aclrtFree(output_device);
    aclrtFreeHost(output_host);

    // 运行时去初始化
    aclrtDestroyStream(stream);
    aclrtResetDevice(device_id);
    aclFinalize();

    return output;
}

Kernel 实现

/**
 * @brief Gather 核函数 (Kernel Function)
 * 
 * @tparam type_data 数据类型 (如 float)
 * @tparam type_idx 索引类型 (如 uint32_t)
 * @param input 输入张量指针
 * @param index 索引数组指针 (存储的是行号)
 * @param gather_output 输出张量指针
 * @param in_width 输入矩阵的宽度 (列数)
 * @param index_total_length 索引数组的总长度 (即需要抓取多少行)
 */
template <typename type_data, typename type_idx>
__global__ void gather_custom(
    type_data* input,
    type_idx* index,
    type_data* gather_output,
    uint32_t in_width,
    uint32_t index_total_length)
{
    // 1. 计算全局线程 ID
    // blockIdx.x: 当前 Block 的索引
    // blockDim.x: 每个 Block 中的线程数 = 256
    // threadIdx.x: 线程在 Block 内的索引
    int32_t out_row = blockIdx.x * blockDim.x + threadIdx.x;

    // 2. 边界检查 (Guard Clause)
    // 如果当前线程计算出的行号超过了索引总数，直接返回
    // 这种写法允许我们启动比实际数据量更多的线程，简化了计算逻辑
    if (out_row >= index_total_length) {
        return;
    }

    // 3. 核心逻辑：单行全量拷贝
    // 目的：保证内存访问的连续性 (Coalesced Memory Access)
    // 获取当前行在输入中的源行号
    uint32_t in_row = index[out_row]; 

    // 关键优化点：这里不是简单的 gather_output[out_row] = input[in_row]
    // 而是循环拷贝该行的所有列数据
    for (int32_t col = 0; col < in_width; col++) {
        // 计算输入数据的线性地址: 源行首地址 + 列偏移
        int input_idx = in_row * in_width + col;
        // 计算输出数据的线性地址: 目标行首地址 + 列偏移
        int output_idx = out_row * in_width + col;

        // 执行拷贝
        gather_output[output_idx] = input[input_idx];
    }
    // 优势：这种连续的 for 循环访问内存，可以让 Ascend 芯片的 DMA 引擎进行高效搬运。
}

性能优化点1：动态block配置

在basic gather中，我们将block固定为48核，每个核线程数固定为256个，但是，Ascend 芯片的型号可能更多样化，因此我们可以使用动态切分线程方法。

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 plateform 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 then "
            << MAX_THREAD_COUNT * MAX_BLOCK_COUNT<< "."<< std::endl;
        return false;
        }
    }
    return true;
}

对应Host侧修改为：

    // Calc splite params
    uint32_t numBlocks = 0;
    uint32_t thread_num_per_block = 0;
    if (!block_split(index_total_length, numBlocks, thread_num_per_block)) {
        return {};
    }

性能优化2：gather->gather_v2

我们现在已经实现了简化场景的gather算子，gather_v2和gather的区别是增加了“axis”、“batch_dims”两个参数，以及支持多维input。

  • 算子功能：
    gather_v2算子实现了从多维输入张量input中按照指定维度axis收集数据的功能，indices张量指定了要收集的索引位置。支持batch_dims批量处理模式，让不同的batch使用不同的索引集合。

  • 算子规格：

    算子类型(OpType) gather_v2
    算子输入 name data type format description
    input float ND 多维输入张量
    indices uint32_t / int32_t ND 索引张量，指定收集位置
    axis int32_t - 标量，用于指定收集维度
    batch_dims int32_t - 标量，用于指定批处理维度
    算子输出 output float ND 收集后的输出张量
    核函数名 gather_custom_v2

  • 约束说明：

    • indices：indices的前batch_dims维必须与input的前batch_dims维相同
    • axis：收集维度axis不能小于batch_dims，且不能超过input的维度数，即batch_dims <= axis < input.rank
    • batch_dims：批量维度数不能超过input和indices中较小的维度数
    • output.shape：input.shape[:axis] + indices.shape[batch_dims:] + input.shape[axis+1:]
    • 样例实现中，axis和batch_dims也支持传入负值，并会在计算前转换为对应的非负维度索引

Output Shape的拼接逻辑

规格说明中写道：output.shape = input.shape[:axis] + indices.shape[batch_dims:] + input.shape[axis+1:]，这可能不是很好理解，举个栗子。
这其实是一个“三段式拼接”的过程。我们可以把它想象成把 input 的 axis 维度挖掉，然后把 indices 的有效部分（去掉 batch 维度）塞进去。

假设：

  • Input shape: [B, M, H, W] (Rank = 4)
  • Indices shape: [B, N] (Rank = 2)
  • batch_dims: 1 (即维度 B 是共享的)
  • axis: 2 (即在 H 维度上进行收集)
拼接阶段 对应的形状切片 实际取值 说明
前段 input.shape[:axis] [B, M] 收集维度之前的部分（包含 Batch）
中段 indices.shape[batch_dims:] [N] Indices 去掉 Batch 后的部分
后段 input.shape[axis+1:] [W] 收集维度之后的部分
最终输出 拼接结果 [B, M, N, W] 这就是 Output 的完整形状

因此，我们需要新建一个struct记录各种关于shape的参数，方便后面传入kernel，代码如下。

struct GatherV2Params {
    uint32_t outer_size = 1;       // Product of dimensions before axis (excluding batch)
    uint32_t gather_dim_size = 1;  // Size of axis dimension
    uint32_t indices_size = 1;     // Product of indices dimensions (excluding batch)
    uint32_t slice_size = 1;       // Product of dimensions after axis
    uint32_t output_total_size = 0;// Total output elements
};

下面展示如何计算GatherV2Params中的变量。

GatherV2Params compute_gather_input(
    const std::vector<uint32_t>& input_shape,
    const std::vector<uint32_t>& indices_shape,
    int32_t axis,
    int32_t batch_dims)
{
    GatherV2Params params;
    uint32_t input_rank = input_shape.size();
    uint32_t indices_rank = indices_shape.size();

    uint32_t batch_size = 1;
    for (uint32_t i = 0; i < batch_dims; ++i) {
        batch_size *= input_shape[i];
    }

    for (uint32_t i = batch_dims; i < axis; ++i) {
        params.outer_size *= input_shape[i];
    }

    params.gather_dim_size = input_shape[axis];

    for (uint32_t i = batch_dims; i < indices_rank; ++i) {
        params.indices_size *= indices_shape[i];
    }

    for (uint32_t i = axis + 1; i < input_rank; ++i) {
        params.slice_size *= input_shape[i];
    }

    params.output_total_size = batch_size * params.outer_size * params.indices_size * params.slice_size;

    return params;
}

根据我们的例子，可以计算出：

参数 计算逻辑 结果
gather_dim_size input_shape[0] 100,000
slice_size input_shape[1] (后面没维度了) 128
indices_size indices_shape 的总元素 12,288
outer_size batch_dims 到 axis 之间 1
output_total_size 1 * 1 * 12288 * 128 1,572,864

kernel实现

template <typename type_data, typename type_idx, bool is_batch_dims_zero, bool is_axis_zero>
__global__ __launch_bounds__(MAX_THREAD_COUNT) void gather_custom_v2(
    type_data* input,
    type_idx* indices,
    type_data* gather_output,
    uint32_t outer_size,
    uint32_t gather_dim_size,
    uint32_t indices_size,
    uint32_t slice_size,
    uint32_t output_total_size)
{
    uint32_t begin = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t stride = gridDim.x * blockDim.x;

    for (uint32_t i = begin; i < output_total_size; i += stride) {
        // Split the flattened output index into: [batch, outer, indices, slice].
        uint32_t slices_count = i / slice_size;
        uint32_t slice_i = i - slices_count * slice_size;

        uint32_t entries_count = slices_count / indices_size;
        uint32_t indices_i = slices_count - entries_count * indices_size;

        uint32_t batch_i = 0;
        uint32_t outer_i = 0;

        if constexpr (!is_batch_dims_zero) {
            batch_i = entries_count / outer_size;
        }
        if constexpr (!is_axis_zero) {
            outer_i = entries_count - batch_i * outer_size;
        }

        const type_idx gather_idx = indices[batch_i * indices_size + indices_i];

        if (gather_idx >= 0 && gather_idx < gather_dim_size) {
            const uint32_t input_i = (
                (batch_i * outer_size + outer_i) * gather_dim_size + gather_idx
            ) * slice_size + slice_i;
            gather_output[i] = input[input_i];
        } else {
            gather_output[i] = static_cast<type_data>(0);
        }
    }
}

假设input shape=(100000,128), indice shape=(12288), output shape=(12288,128)
num_block=56，block_dim=2048，总线程数 = 56 * 2048 = 114,688 个线程。

简化场景：1个线程负责搬运 1整行数据。input 的单行长度是 128，这意味着每个线程不仅要计算一次索引偏移，还要在内部开启一个循环，连续从全局内存搬运 128 个浮点数。 indice_shape = 12,288。只需要启动 12,288 个线程就能完成任务。这意味着我们配置的 114,688 个线程中，有 90% 的线程是空闲的，导致负载不均衡。
for (int32_t col = 0; col < in_width; col++) { //in_width=128
gather_output[output_idx] = input[input_idx];
input_idx += 1;
output_idx += 1;
}

泛化场景：gather_v2 不再让 1 个线程搬运 1 整行（128维），而是把这一行拆开，让多个线程协同搬运。它把输出张量看作一个长长的 1 维数组，每隔固定距离取一个元素。这在逻辑上等同于把 128 维的行拆开，让不同时间片的线程去处理。
uint32_t begin = blockIdx.x * blockDim.x + threadIdx.x;
uint32_t stride = gridDim.x * blockDim.x;
for (uint32_t i = begin; i < output_total_size; i += stride) {
// … 核心逻辑
}
其中stride=562048=144,688
output_total_size=12288128=1,572,864
循环内总共进行了10次，远小于原版本的128次循环。

# 性能优化 # 昇腾
Logo
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

【NVIDIA】Ascend Docker Runtime v26.0.1 之五 install/ — 超深度逐行分析

Ascend Docker Runtime v26.0.1 之五 install/ — 超深度逐行分析

avatar 鲲鹏昇腾开发者社区

【NPU】Ascend Docker Runtime v26.0.1 之六 mindxcheckutils/ — 超深度逐行分析

Ascend Docker Runtime v26.0.1 之六 mindxcheckutils/ — 超深度逐行分析

avatar 鲲鹏昇腾开发者社区

[AI][昇腾950] AscendC 硬件抽象架

David 系列DaVinci AICore。

avatar 鲲鹏昇腾开发者社区