Element-Wise算子模板：一套代码实现加、减、乘、除

在训练营的进阶课程中，我们将把这种思维应用到更复杂的场景：如何设计通用的Reduce模板、卷积模板、甚至自定义的融合算子模板。回想这次Element-Wise模板的实践，我完成的不仅仅是一套技术方案，更是一次编程思维的升级。它在Ascend C算子开发中的成功应用，证明了好的软件设计原则是跨领域、跨平台的。，像搭积木一样快速生成所有Element-Wise算子，体验从"工匠"到"架构师"的思维跃迁

2501_94317725

987人浏览 · 2025-11-19 18:38:46

2501_94317725 · 2025-11-19 18:38:46 发布

在CANN训练营里，当我按照标准化流程成功实现ReLU算子后，导师随即布置了新任务：“接下来，请实现加法、减法、乘法、除法四个算子。”

我心想，这还不简单？Ctrl+C、Ctrl+V，改改计算逻辑不就行了？但当我真的开始复制粘贴时，突然意识到：这四个算子的代码结构、内存管理、数据搬运逻辑几乎一模一样，唯一不同的就是计算部分的那一个操作符。这种重复劳动不仅低效，更可怕的是，一旦发现一个公共bug，我需要在四个地方修改四次！

在 [2025年昇腾CANN训练营第二季] 的“码力全开特辑”中，导师引入了C++模板元编程来解决这个问题。今天，我就带你领略如何用一套模板代码，像搭积木一样快速生成所有Element-Wise算子，体验从"工匠"到"架构师"的思维跃迁。

>> 高效的开发技巧是训练营的核心价值：点击加入，学习更智能的编程方式

第一章：从"复制粘贴"的困境到模板的曙光

让我先展示一下传统方式的"痛苦"。下面是加法算子的核心计算部分：

// 加法算子
for (uint32_t i = 0; i < copyLength; ++i) {
    localZ[i] = localX[i] + localY[i];
}

然后是减法算子：

// 减法算子  
for (uint32_t i = 0; i < copyLength; ++i) {
    localZ[i] = localX[i] - localY[i];
}

还有乘法、除法… 看到问题了吗？95%的代码都是重复的，只有那个操作符不同。

这种重复带来了三大痛点：

维护噩梦：修复一个数据搬运的bug，需要修改所有算子的代码。
测试负担：每个算子都需要独立的测试用例。
扩展困难：想要新增一个"平方"算子，又得从头复制一遍。

模板技术的核心思想：将不变的部分（架构）与变化的部分（计算逻辑）分离，让编译器在编译期间为我们生成具体的代码。

第二章：设计Element-Wise通用模板——打造可复用的"乐高积木"

第一步：抽象算子行为

我们首先定义所有Element-Wise算子的共同特征：

输入：两个Tensor（A和B），shape相同
输出：一个Tensor（C），shape与输入相同
计算：对每个位置的元素进行二元运算：C[i] = A[i] op B[i]

第二步：设计模板接口

我们需要一个方式来表示那个"变化的部分"——具体的运算。在C++中，有几种实现方式：

方案1：函数指针（传统但有效）

typedef half (*BinaryOp)(half, half);

方案2：函数对象（现代C++推荐）

template<typename T>
struct BinaryOp {
    virtual T operator()(T a, T b) const = 0;
};

方案3：C++11 Lambda + std::function（灵活但可能有性能开销）

在Ascend C的高性能场景下，我们选择方案2：函数对象，因为它既类型安全，又可以被编译器内联优化。

第三章：实现通用模板核函数——构建"万能工厂"

现在，让我们实现这个通用的模板核函数。

核心模板定义：

// element_wise_template.h
#include "kernel_operator.h"

using namespace AscendC;

// 定义Tiling结构
struct ElementWiseTiling {
    uint32_t totalLength;
    uint32_t tileNum;
};

// 通用的Element-Wise核函数模板
template<typename BinaryOp>
__global__ __aicore__ void element_wise_template(
    ElementWiseTiling* tiling,
    half* x,
    half* y, 
    half* z,
    BinaryOp op)  // 关键：传入运算函数对象
{
    // 1. 任务划分（与ReLU中完全相同）
    uint32_t blockIdx = GET_BLOCK_IDX();
    uint32_t blockDim = GET_BLOCK_NUM();
    uint32_t totalLength = tiling->totalLength;
    uint32_t tileNum = tiling->tileNum;
    
    uint32_t dataPerBlock = totalLength / tileNum;
    uint32_t remainder = totalLength % tileNum;
    uint32_t currentLength = dataPerBlock + (blockIdx < remainder ? 1 : 0);
    uint32_t currentOffset = blockIdx * dataPerBlock + (blockIdx < remainder ? blockIdx : remainder);
    
    if (currentLength == 0) return;

    // 2. 内存指针定义
    __gm__ half* globalX = x + currentOffset;
    __gm__ half* globalY = y + currentOffset; 
    __gm__ half* globalZ = z + currentOffset;

    // 3. 本地缓冲区
    constexpr uint32_t BUFFER_SIZE = 256;
    half localX[BUFFER_SIZE];
    half localY[BUFFER_SIZE];
    half localZ[BUFFER_SIZE];

    // 4. 分块处理
    uint32_t processed = 0;
    while (processed < currentLength) {
        uint32_t copyLength = (currentLength - processed) > BUFFER_SIZE ? 
                             BUFFER_SIZE : (currentLength - processed);
        
        // 数据搬运
        DataCopy<LocalTensor, GM_ADDR>(localX, globalX + processed, copyLength / 16, 0, 0);
        DataCopy<LocalTensor, GM_ADDR>(localY, globalY + processed, copyLength / 16, 0, 0);

        // 5. 【核心变化点】使用模板参数进行通用计算
        for (uint32_t i = 0; i < copyLength; ++i) {
            localZ[i] = op(localX[i], localY[i]);  // 统一的调用接口！
        }

        // 结果回写
        DataCopy<GM_ADDR, LocalTensor>(globalZ + processed, localZ, copyLength / 16, 0, 0);
        
        processed += copyLength;
    }
}

这个模板的精妙之处在于：op(localX[i], localY[i]) 这一行代码替代了之前所有具体的运算逻辑。编译器会根据我们传入的不同的 BinaryOp 对象，在这里生成不同的机器指令。

第四章：实现具体运算——创建不同的"模具"

现在，我们为每种运算创建具体的函数对象：

// binary_ops.h
#ifndef BINARY_OPS_H
#define BINARY_OPS_H

#include <cstdint>

// 加法运算
struct AddOp {
    __host__ __device__ half operator()(half a, half b) const {
        return __hadd(a, b);  // 使用硬件优化的半精度加法
    }
};

// 减法运算  
struct SubOp {
    __host__ __device__ half operator()(half a, half b) const {
        return __hsub(a, b);
    }
};

// 乘法运算
struct MulOp {
    __host__ __device__ half operator()(half a, half b) const {
        return __hmul(a, b);
    }
};

// 除法运算
struct DivOp {
    __host__ __device__ half operator()(half a, half b) const {
        // 注意除零保护
        if (__heq(b, __float2half(0.0f))) {
            return __float2half(0.0f);  // 或者根据需求返回其他值
        }
        return __hdiv(a, b);
    }
};

// 你还可以轻松扩展其他运算！
struct MaximumOp {
    __host__ __device__ half operator()(half a, half b) const {
        return __hge(a, b) ? a : b;
    }
};

#endif

第五章：主机侧调用——像搭积木一样组合

现在，主机侧的调用变得异常简洁和统一：

// main.cpp
#include "element_wise_template.h"
#include "binary_ops.h"

// 统一的测试函数
template<typename BinaryOp>
void test_element_wise(const char* op_name, BinaryOp op) {
    std::cout << "测试 " << op_name << " 算子..." << std::endl;
    
    // 准备数据、内存分配等（与之前相同）
    // ...
    
    // 启动核函数 - 传入具体的运算对象
    element_wise_template<<<tiling.tileNum, nullptr>>>(deviceTiling, 
                                                      deviceX, deviceY, deviceZ, op);
    
    // 等待、验证、释放资源...
    // ...
}

int main() {
    // 一次性测试所有算子！
    test_element_wise("加法", AddOp{});
    test_element_wise("减法", SubOp{});
    test_element_wise("乘法", MulOp{});
    test_element_wise("除法", DivOp{});
    test_element_wise("最大值", MaximumOp{});  // 轻松扩展！
    
    return 0;
}

第六章：模板技术的威力与进阶技巧

1. 性能零开销
你可能担心模板会带来性能损失。实际上，由于函数对象的 operator() 通常被编译器内联，生成的汇编代码与直接写死运算逻辑完全一样，没有任何函数调用开销。

2. 编译期多态
这种方式实现了编译期多态，比运行时的虚函数调用高效得多，非常适合高性能计算场景。

3. 条件编译支持
你还可以通过模板特化，为特定数据类型或运算提供优化版本：

// 为整数类型特化的加法，避免浮点运算开销
template<>
struct AddOp<int32_t> {
    __host__ __device__ int32_t operator()(int32_t a, int32_t b) const {
        return a + b;
    }
};