CANN算子基础框架库opbase的算子开发与扩展机制深度解析

CANN算子基础框架库opbase的算子开发与扩展机制深度解析

前言

在深度学习框架中，算子（Operator）是构成计算图的基本单元。opbase是CANN框架中负责算子定义、注册、管理和调度的基础框架库。它提供了一套完整的算子开发基础设施，使开发者能够高效地创建、注册和部署自定义算子。本文将深入剖析opbase的架构设计和核心机制。

相关链接：

• CANN组织链接：https://atomgit.com/cann
• opbase仓库链接：https://atomgit.com/cann/opbase

---

一、opbase框架概述

1.1 设计理念

opbase框架遵循以下核心设计理念：

• 分层解耦：算子接口定义与内核实现分离
• 多后端支持：支持CPU、NPU等多种计算后端
• 类型安全：强类型系统确保数据类型一致性
• 可扩展性：插件化架构支持动态算子注册
• 性能优先：零拷贝数据传递、内核融合优化

1.2 架构层次

opbase/
├── 算子定义层 (Operator Definition)
│   ├── OpDef - 算子元数据定义
│   ├── Attr - 属性类型系统
│   └── DataType - 数据类型系统
├── 算子注册层 (Operator Registration)
│   ├── OpRegistry - 算子注册表
│   ├── OpKernel - 算子内核注册
│   └── Dialect - 方言系统
├── 形状推导层 (Shape Inference)
│   ├── ShapeFunc - 形状计算函数
│   ├── TensorShape - 张量形状表示
│   └── Dimension - 维度系统
├── 算子执行层 (Operator Execution)
│   ├── OpExecutor - 算子执行器
│   ├── ExecutionContext - 执行上下文
│   └── MemoryManager - 内存管理器
└── 工具支持层 (Utility)
    ├── OpBuilder - 算子构建器
    ├── OpPrinter - 算子打印/调试
    └── OpValidator - 算子验证

---

二、核心API详解

2.1 算子定义API

OpDef结构

/**
 * @brief 算子定义结构体
 * 描述算子的元数据，包括名称、输入输出、属性等
 */
class OpDef {
public:
    // 构造函数
    explicit OpDef(const std::string& name) : name_(name) {}

    // 设置算子名称
    OpDef& Name(const std::string& name) {
        name_ = name;
        return *this;
    }

    // 添加输入
    OpDef& Input(const std::string& name, DataType dtype,
                 TensorShape shape = TensorShape::Unknown()) {
        inputs_.push_back({name, dtype, shape});
        return *this;
    }

    // 添加输出
    OpDef& Output(const std::string& name, DataType dtype,
                  TensorShape shape = TensorShape::Unknown()) {
        outputs_.push_back({name, dtype, shape});
        return *this;
    }

    // 添加属性
    template<typename T>
    OpDef& Attr(const std::string& name, const T& value) {
        attrs_[name] = Attribute(value);
        return *this;
    }

    // 设置描述
    OpDef& Description(const std::string& desc) {
        description_ = desc;
        return *this;
    }

    // 设置文档
    OpDef& Doc(const std::string& doc) {
        doc_string_ = doc;
        return *this;
    }

    // 获取方法
    const std::string& GetName() const { return name_; }
    const std::vector<IOInfo>& GetInputs() const { return inputs_; }
    const std::vector<IOInfo>& GetOutputs() const { return outputs_; }
    const AttributeMap& GetAttributes() const { return attrs_; }

    // 验证算子定义合法性
    bool Validate() const;

    // 生成算子签名
    std::string GetSignature() const;

private:
    struct IOInfo {
        std::string name;
        DataType dtype;
        TensorShape shape;
    };

    std::string name_;
    std::string description_;
    std::string doc_string_;
    std::vector<IOInfo> inputs_;
    std::vector<IOInfo> outputs_;
    AttributeMap attrs_;
};

// 使用示例
OpDef conv2d_def("Conv2D");
conv2d_def
    .Input("input", DT_FLOAT)
    .Input("filter", DT_FLOAT)
    .Input("bias", DT_FLOAT)
    .Output("output", DT_FLOAT)
    .Attr<int>("strides", {1, 1})
    .Attr<std::string>("padding", "SAME")
    .Attr<bool>("use_bias", true)
    .Description("2D convolution operator");

属性系统

/**
 * @brief 属性类型枚举
 */
enum AttrType {
    ATTR_INT = 0,
    ATTR_FLOAT,
    ATTR_STRING,
    ATTR_BOOL,
    ATTR_INT_LIST,
    ATTR_FLOAT_LIST,
    ATTR_STRING_LIST,
    ATTR_BOOL_LIST,
    ATTR_TENSOR_SHAPE,
    ATTR_DATA_TYPE
};

/**
 * @brief 类型安全的属性包装类
 */
class Attribute {
public:
    // 构造函数
    Attribute() : type_(ATTR_INT) {}
    explicit Attribute(int value) : type_(ATTR_INT), int_val_(value) {}
    explicit Attribute(float value) : type_(ATTR_FLOAT), float_val_(value) {}
    explicit Attribute(const std::string& value)
        : type_(ATTR_STRING), string_val_(value) {}
    explicit Attribute(bool value) : type_(ATTR_BOOL), bool_val_(value) {}

    // 列表构造
    explicit Attribute(const std::vector<int>& value)
        : type_(ATTR_INT_LIST), int_list_val_(value) {}
    explicit Attribute(const std::vector<float>& value)
        : type_(ATTR_FLOAT_LIST), float_list_val_(value) {}

    // 获取值
    int GetInt() const { return int_val_; }
    float GetFloat() const { return float_val_; }
    const std::string& GetString() const { return string_val_; }
    bool GetBool() const { return bool_val_; }
    const std::vector<int>& GetIntList() const { return int_list_val_; }
    const std::vector<float>& GetFloatList() const { return float_list_val_; }

    // 类型检查
    AttrType GetType() const { return type_; }
    bool IsInt() const { return type_ == ATTR_INT; }
    bool IsFloat() const { return type_ == ATTR_FLOAT; }
    bool IsString() const { return type_ == ATTR_STRING; }
    bool IsBool() const { return type_ == ATTR_BOOL; }
    bool IsList() const {
        return type_ >= ATTR_INT_LIST && type_ <= ATTR_BOOL_LIST;
    }

    // 类型转换
    template<typename T>
    T As() const;

    template<>
    int As<int>() const { return GetInt(); }

    template<>
    float As<float>() const { return GetFloat(); }

    template<>
    std::string As<std::string>() const { return GetString(); }

    template<>
    bool As<bool>() const { return GetBool(); }

private:
    AttrType type_;
    union {
        int int_val_;
        float float_val_;
        bool bool_val_;
    };
    std::string string_val_;
    std::vector<int> int_list_val_;
    std::vector<float> float_list_val_;
    std::vector<std::string> string_list_val_;
    std::vector<bool> bool_list_val_;
};

// 属性映射表
using AttributeMap = std::unordered_map<std::string, Attribute>;

2.2 算子注册API

/**
 * @brief 算子注册表
 * 单例模式，管理所有已注册的算子
 */
class OpRegistry {
public:
    // 获取单例
    static OpRegistry& Global() {
        static OpRegistry instance;
        return instance;
    }

    // 注册算子定义
    Status RegisterOpDef(const OpDef& op_def) {
        std::unique_lock<std::shared_mutex> lock(mutex_);

        const std::string& name = op_def.GetName();
        if (op_defs_.find(name) != op_defs_.end()) {
            return Status::Error("Op " + name + " already registered");
        }

        if (!op_def.Validate()) {
            return Status::Error("Invalid op definition: " + name);
        }

        op_defs_[name] = op_def;
        return Status::OK();
    }

    // 查找算子定义
    StatusOr<const OpDef*> LookUpOpDef(const std::string& name) const {
        std::shared_lock<std::shared_mutex> lock(mutex_);

        auto it = op_defs_.find(name);
        if (it == op_defs_.end()) {
            return Status::Error("Op not found: " + name);
        }
        return &it->second;
    }

    // 获取所有已注册算子名称
    std::vector<std::string> GetRegisteredOps() const {
        std::shared_lock<std::shared_mutex> lock(mutex_);

        std::vector<std::string> names;
        names.reserve(op_defs_.size());
        for (const auto& pair : op_defs_) {
            names.push_back(pair.first);
        }
        return names;
    }

    // 检查算子是否已注册
    bool IsRegistered(const std::string& name) const {
        std::shared_lock<std::shared_mutex> lock(mutex_);
        return op_defs_.find(name) != op_defs_.end();
    }

private:
    OpRegistry() = default;
    ~OpRegistry() = default;

    mutable std::shared_mutex mutex_;
    std::unordered_map<std::string, OpDef> op_defs_;
};

/**
 * @brief 算子内核注册器
 * 管理算子的具体实现内核
 */
class OpKernelRegistry {
public:
    // 算子内核接口
    class OpKernel {
    public:
        virtual ~OpKernel() = default;

        // 计算输出形状
        virtual Status ComputeShape(OpKernelContext* context) = 0;

        // 执行计算
        virtual Status Compute(OpKernelContext* context) = 0;

        // 获取内核签名
        virtual std::string GetKernelSignature() const = 0;
    };

    // 算子内核构建器
    using KernelFactory = std::function<std::unique_ptr<OpKernel>()>;

    // 注册内核
    Status RegisterKernel(const std::string& op_name,
                         const KernelSignature& signature,
                         KernelFactory factory) {
        std::unique_lock<std::shared_mutex> lock(mutex_);

        std::string key = op_name + ":" + signature.ToString();
        if (kernels_.find(key) != kernels_.end()) {
            return Status::Error("Kernel already registered: " + key);
        }

        kernels_[key] = std::move(factory);
        return Status::OK();
    }

    // 创建内核实例
    StatusOr<std::unique_ptr<OpKernel>> CreateKernel(
        const std::string& op_name,
        const KernelSignature& signature) {

        std::shared_lock<std::shared_mutex> lock(mutex_);

        std::string key = op_name + ":" + signature.ToString();
        auto it = kernels_.find(key);
        if (it == kernels_.end()) {
            return Status::Error("Kernel not found: " + key);
        }

        return it->second();
    }

private:
    mutable std::shared_mutex mutex_;
    std::unordered_map<std::string, KernelFactory> kernels_;
};

// 内核签名
class KernelSignature {
public:
    KernelSignature(const std::vector<DataType>& input_dtypes,
                   const std::vector<DataType>& output_dtypes,
                   const std::string& device)
        : input_dtypes_(input_dtypes),
          output_dtypes_(output_dtypes),
          device_(device) {}

    std::string ToString() const {
        std::string sig = device_ + "/";
        for (auto dt : input_dtypes_) {
            sig += DataTypeToString(dt) + ",";
        }
        sig += "->";
        for (auto dt : output_dtypes_) {
            sig += DataTypeToString(dt) + ",";
        }
        return sig;
    }

    bool operator==(const KernelSignature& other) const {
        return input_dtypes_ == other.input_dtypes_ &&
               output_dtypes_ == other.output_dtypes_ &&
               device_ == other.device_;
    }

private:
    std::vector<DataType> input_dtypes_;
    std::vector<DataType> output_dtypes_;
    std::string device_;
};

2.3 形状推导API

/**
 * @brief 形状推导上下文
 */
class ShapeInferenceContext {
public:
    ShapeInferenceContext(const OpDef& op_def,
                         const std::vector<TensorShape>& input_shapes)
        : op_def_(op_def), input_shapes_(input_shapes) {}

    // 获取输入形状
    const TensorShape& GetInputShape(int index) const {
        return input_shapes_.at(index);
    }

    // 获取输入数量
    int GetNumInputs() const {
        return input_shapes_.size();
    }

    // 获取属性值
    template<typename T>
    T GetAttr(const std::string& name) const {
        return op_def_.GetAttributes().at(name).As<T>();
    }

    // 设置输出形状
    void SetOutputShape(int index, const TensorShape& shape) {
        if (output_shapes_.size() <= index) {
            output_shapes_.resize(index + 1);
        }
        output_shapes_[index] = shape;
    }

    // 获取输出形状
    const std::vector<TensorShape>& GetOutputShapes() const {
        return output_shapes_;
    }

    // 输入形状完整检查
    bool AllInputsKnown() const {
        for (const auto& shape : input_shapes_) {
            if (shape.IsUnknown()) return false;
        }
        return true;
    }

private:
    const OpDef& op_def_;
    std::vector<TensorShape> input_shapes_;
    std::vector<TensorShape> output_shapes_;
};

/**
 * @brief 形状推导函数类型
 */
using ShapeFunc = std::function<Status(ShapeInferenceContext* context)>;

/**
 * @brief 注册形状推导函数
 */
class ShapeFuncRegistry {
public:
    static ShapeFuncRegistry& Global() {
        static ShapeFuncRegistry instance;
        return instance;
    }

    Status Register(const std::string& op_name, ShapeFunc func) {
        std::unique_lock<std::mutex> lock(mutex_);
        shape_funcs_[op_name] = std::move(func);
        return Status::OK();
    }

    StatusOr<ShapeFunc> Get(const std::string& op_name) const {
        std::shared_lock<std::mutex> lock(mutex_);
        auto it = shape_funcs_.find(op_name);
        if (it == shape_funcs_.end()) {
            return Status::Error("Shape function not found: " + op_name);
        }
        return it->second;
    }

private:
    mutable std::shared_mutex mutex_;
    std::unordered_map<std::string, ShapeFunc> shape_funcs_;
};

// 形状推导辅助函数
namespace shape_inference {

// 一对一形状传递
inline Status PassThrough(ShapeInferenceContext* context) {
    const TensorShape& input_shape = context->GetInputShape(0);
    context->SetOutputShape(0, input_shape);
    return Status::OK();
}

// 广播形状计算
inline Status Broadcast(ShapeInferenceContext* context) {
    const TensorShape& shape_a = context->GetInputShape(0);
    const TensorShape& shape_b = context->GetInputShape(1);

    TensorShape result_shape = BroadcastShapes(shape_a, shape_b);
    context->SetOutputShape(0, result_shape);
    return Status::OK();
}

// 卷积输出形状计算
inline Status Conv2DShape(ShapeInferenceContext* context) {
    const TensorShape& input_shape = context->GetInputShape(0);
    const TensorShape& filter_shape = context->GetInputShape(1);

    auto strides = context->GetAttr<std::vector<int>>("strides");
    auto padding_str = context->GetAttr<std::string>("padding");

    // NCHW格式: [N, C, H, W]
    int batch = input_shape.dim(0);
    int in_channels = input_shape.dim(1);
    int in_height = input_shape.dim(2);
    int in_width = input_shape.dim(3);

    int out_channels = filter_shape.dim(0);
    int filter_height = filter_shape.dim(2);
    int filter_width = filter_shape.dim(3);

    int padding_h, padding_w;
    if (padding_str == "SAME") {
        int out_h = (in_height + strides[0] - 1) / strides[0];
        int out_w = (in_width + strides[1] - 1) / strides[1];
        padding_h = std::max(0, (out_h - 1) * strides[0] +
                             filter_height - in_height);
        padding_w = std::max(0, (out_w - 1) * strides[1] +
                             filter_width - in_width);
    } else {  // VALID
        padding_h = 0;
        padding_w = 0;
    }

    int out_height = (in_height + 2 * padding_h - filter_height) / strides[0] + 1;
    int out_width = (in_width + 2 * padding_w - filter_width) / strides[1] + 1;

    TensorShape output_shape({batch, out_channels, out_height, out_width});
    context->SetOutputShape(0, output_shape);
    return Status::OK();
}

}  // namespace shape_inference

2.4 执行上下文API

/**
 * @brief 算子执行上下文
 * 提供算子执行时所需的所有资源和信息
 */
class OpKernelContext {
public:
    OpKernelContext(const std::vector<Tensor*>& inputs,
                   const std::vector<Attribute>& attrs,
                   MemoryManager* memory_manager)
        : inputs_(inputs),
          attrs_(attrs),
          memory_manager_(memory_manager) {}

    // 获取输入
    const Tensor* GetInput(int index) const {
        return inputs_.at(index);
    }

    int GetNumInputs() const { return inputs_.size(); }

    // 分配输出
    StatusOr<Tensor*> AllocateOutput(int index,
                                     const TensorShape& shape,
                                     DataType dtype) {
        if (outputs_.size() <= index) {
            outputs_.resize(index + 1);
        }

        if (outputs_[index] == nullptr) {
            size_t bytes = shape.num_elements() * DataTypeSize(dtype);
            void* buffer = memory_manager_->Allocate(bytes);
            outputs_[index] = new Tensor(buffer, shape, dtype);
        }

        return outputs_[index];
    }

    // 获取输出
    Tensor* GetOutput(int index) {
        return outputs_.at(index);
    }

    // 获取属性
    template<typename T>
    T GetAttr(const std::string& name) const {
        return attrs_.at(name).As<T>();
    }

    // 获取设备信息
    const DeviceContext& GetDeviceContext() const {
        return device_context_;
    }

    // 获取执行流
    const Stream& GetStream() const {
        return stream_;
    }

    // 设置状态信息
    void SetStatus(const Status& status) {
        status_ = status;
    }

    const Status& GetStatus() const {
        return status_;
    }

private:
    std::vector<Tensor*> inputs_;
    std::vector<Tensor*> outputs_;
    std::vector<Attribute> attrs_;
    MemoryManager* memory_manager_;
    DeviceContext device_context_;
    Stream stream_;
    Status status_;
};

/**
 * @brief 算子执行器
 * 负责算子的调度和执行
 */
class OpExecutor {
public:
    // 执行算子
    Status Execute(const std::string& op_name,
                   const std::vector<Tensor*>& inputs,
                   std::vector<Tensor*>& outputs,
                   const AttributeMap& attrs) {

        // 1. 查找算子定义
        ASSIGN_OR_RETURN(const OpDef* op_def,
                        OpRegistry::Global().LookUpOpDef(op_name));

        // 2. 检查输入输出数量
        if (inputs.size() != op_def->GetInputs().size()) {
            return Status::Error("Input count mismatch");
        }

        // 3. 推导输出形状
        std::vector<TensorShape> input_shapes;
        for (auto* input : inputs) {
            input_shapes.push_back(input->shape());
        }

        ShapeInferenceContext shape_ctx(*op_def, input_shapes);
        ASSIGN_OR_RETURN(ShapeFunc shape_func,
                        ShapeFuncRegistry::Global().Get(op_name));
        RETURN_IF_ERROR(shape_func(&shape_ctx));

        // 4. 选择并创建内核
        std::vector<DataType> input_dtypes, output_dtypes;
        for (auto* input : inputs) {
            input_dtypes.push_back(input->dtype());
        }
        for (const auto& out_shape : shape_ctx.GetOutputShapes()) {
            // 简化：使用输入类型
            output_dtypes.push_back(input->dtype());
        }

        KernelSignature signature(input_dtypes, output_dtypes, "NPU");
        ASSIGN_OR_RETURN(auto kernel,
                        OpKernelRegistry::Global().CreateKernel(
                            op_name, signature));

        // 5. 创建执行上下文
        std::vector<Attribute> attr_list;
        for (const auto& pair : attrs) {
            attr_list.push_back(pair.second);
        }
        OpKernelContext kernel_ctx(inputs, attr_list, memory_manager_);

        // 6. 执行内核
        RETURN_IF_ERROR(kernel->Compute(&kernel_ctx));

        // 7. 返回输出
        outputs = kernel_ctx.GetOutputs();

        return Status::OK();
    }

private:
    MemoryManager* memory_manager_;
};

---

三、应用实践

3.1 自定义算子开发

以下是一个完整的自定义算子开发示例，实现一个Leaky ReLU激活函数：

#include "opbase/op_registry.h"
#include "opbase/op_kernel.h"

// 1. 定义算子
OpDef& DefineLeakyReluOp() {
    static OpDef def("LeakyRelu");
    def.Input("x", DT_FLOAT, TensorShape::Unknown())
       .Output("y", DT_FLOAT, TensorShape::Unknown())
       .Attr<float>("alpha", 0.01f)
       .Description("Leaky ReLU activation: max(alpha * x, x)");
    return def;
}

// 2. 实现算子内核
class LeakyReluKernel : public OpKernelRegistry::OpKernel {
public:
    LeakyReluKernel(float alpha) : alpha_(alpha) {}

    Status ComputeShape(OpKernelContext* context) override {
        // LeakyReLU保持输入形状
        const Tensor* input = context->GetInput(0);
        TensorShape output_shape = input->shape();
        context->SetOutputShape(0, output_shape);
        return Status::OK();
    }

    Status Compute(OpKernelContext* context) override {
        const Tensor* input = context->GetInput(0);
        Tensor* output = context->GetOutput(0);

        // 获取输入数据
        const float* input_data = input->data<float>();
        float* output_data = output->data<float>();
        int64_t num_elements = input->shape().num_elements();

        // NPU执行路径
        if (input->on_device()) {
            return ComputeOnDevice(input, output, context);
        }

        // CPU回退路径
        for (int64_t i = 0; i < num_elements; ++i) {
            output_data[i] = std::max(alpha_ * input_data[i], input_data[i]);
        }

        return Status::OK();
    }

    std::string GetKernelSignature() const override {
        return "LeakyRelu_NPU_FLOAT32_FLOAT32";
    }

private:
    Status ComputeOnDevice(const Tensor* input, Tensor* output,
                          OpKernelContext* context) {
        // 调用NPU内核
        void* input_ptr = input->device_ptr();
        void* output_ptr = output->device_ptr();
        int64_t count = input->shape().num_elements();

        // 使用aclnn API
        aclnnStatus ret = aclnnInplaceLeakyRelu(
            (aclTensor*)input_ptr,
            alpha_,
            (aclTensor*)output_ptr,
            context->GetStream().GetHandle(),
            context->GetWorkspace());

        if (ret != ACLNN_SUCCESS) {
            return Status::Error("LeakyReLU kernel failed");
        }

        return Status::OK();
    }

    float alpha_;
};

// 3. 实现形状推导函数
Status LeakyReluShapeFunc(ShapeInferenceContext* context) {
    const TensorShape& input_shape = context->GetInputShape(0);
    context->SetOutputShape(0, input_shape);
    return Status::OK();
}

// 4. 注册算子
void RegisterLeakyReluOp() {
    // 注册算子定义
    OpRegistry::Global().RegisterOpDef(DefineLeakyReluOp());

    // 注册形状推导函数
    ShapeFuncRegistry::Global().Register("LeakyRelu", LeakyReluShapeFunc);

    // 注册算子内核
    OpKernelRegistry::Global().RegisterKernel(
        "LeakyRelu",
        KernelSignature({DT_FLOAT}, {DT_FLOAT}, "NPU"),
        []() -> std::unique_ptr<OpKernel> {
            return std::make_unique<LeakyReluKernel>(0.01f);
        }
    );
}

// 5. 使用注册宏自动注册
REGISTER_OP(LeakyRelu)
    .Input("x: float")
    .Output("y: float")
    .Attr("alpha: float = 0.01")
    .SetShapeFn(LeakyReluShapeFunc);

REGISTER_OP_KERNEL(LeakyRelu, NPU, FLOAT32)
    .KernelConstructor([](OpKernelConstruction* ctx) {
        float alpha;
        ctx->GetAttr("alpha", &alpha);
        return new LeakyReluKernel(alpha);
    });

3.2 复杂算子开发示例

实现一个Group Normalization算子：

// Group Normalization定义
class GroupNormOpDef {
public:
    static OpDef GetOpDef() {
        return OpDef("GroupNorm")
            .Input("input", DT_FLOAT, TensorShape::Unknown())
            .Input("gamma", DT_FLOAT, TensorShape::Unknown())
            .Input("beta", DT_FLOAT, TensorShape::Unknown())
            .Output("output", DT_FLOAT, TensorShape::Unknown())
            .Output("mean", DT_FLOAT, TensorShape::Unknown())
            .Output("variance", DT_FLOAT, TensorShape::Unknown())
            .Attr<int>("num_groups", 32)
            .Attr<float>("epsilon", 1e-5f)
            .Description("Group Normalization operator");
    }
};

// Group Normalization内核
class GroupNormKernel : public OpKernelRegistry::OpKernel {
public:
    GroupNormKernel(int num_groups, float epsilon)
        : num_groups_(num_groups), epsilon_(epsilon) {}

    Status ComputeShape(OpKernelContext* context) override {
        const Tensor* input = context->GetInput(0);
        const TensorShape& input_shape = input->shape();

        // 输入形状: [N, C, H, W]
        // 输出形状与输入相同
        // mean和variance形状: [N, num_groups]

        TensorShape output_shape = input_shape;
        context->SetOutputShape(0, output_shape);

        TensorShape stat_shape({input_shape.dim(0), num_groups_});
        context->SetOutputShape(1, stat_shape);
        context->SetOutputShape(2, stat_shape);

        return Status::OK();
    }

    Status Compute(OpKernelContext* context) override {
        const Tensor* input = context->GetInput(0);
        const Tensor* gamma = context->GetInput(1);
        const Tensor* beta = context->GetInput(2);

        Tensor* output;
        Tensor* mean;
        Tensor* variance;

        // 分配输出
        context->AllocateOutput(0, input->shape(), DT_FLOAT);
        context->AllocateOutput(1, {input->shape().dim(0), num_groups_}, DT_FLOAT);
        context->AllocateOutput(2, {input->shape().dim(0), num_groups_}, DT_FLOAT);

        output = context->GetOutput(0);
        mean = context->GetOutput(1);
        variance = context->GetOutput(2);

        // 获取维度
        const int64_t N = input->shape().dim(0);
        const int64_t C = input->shape().dim(1);
        const int64_t H = input->shape().dim(2);
        const int64_t W = input->shape().dim(3);

        const int64_t group_size = C / num_groups_;
        const int64_t HW = H * W;
        const int64_t elements_per_group = group_size * HW;

        // 执行Group Normalization
        const float* input_data = input->data<float>();
        float* output_data = output->data<float>();
        float* mean_data = mean->data<float>();
        float* var_data = variance->data<float>();
        const float* gamma_data = gamma->data<float>();
        const float* beta_data = beta->data<float>();

        // 并行处理每个batch
        #pragma omp parallel for collapse(2)
        for (int64_t n = 0; n < N; ++n) {
            for (int64_t g = 0; g < num_groups_; ++g) {
                // 计算当前组的起始和结束通道
                int64_t c_start = g * group_size;
                int64_t c_end = c_start + group_size;

                // 计算均值
                float sum = 0.0f;
                for (int64_t c = c_start; c < c_end; ++c) {
                    for (int64_t hw = 0; hw < HW; ++hw) {
                        int64_t idx = n * C * HW + c * HW + hw;
                        sum += input_data[idx];
                    }
                }
                float group_mean = sum / elements_per_group;
                mean_data[n * num_groups_ + g] = group_mean;

                // 计算方差
                float var_sum = 0.0f;
                for (int64_t c = c_start; c < c_end; ++c) {
                    for (int64_t hw = 0; hw < HW; ++hw) {
                        int64_t idx = n * C * HW + c * HW + hw;
                        float diff = input_data[idx] - group_mean;
                        var_sum += diff * diff;
                    }
                }
                float group_var = var_sum / elements_per_group;
                var_data[n * num_groups_ + g] = group_var;

                // 归一化
                float std_dev = std::sqrt(group_var + epsilon_);
                for (int64_t c = c_start; c < c_end; ++c) {
                    for (int64_t hw = 0; hw < HW; ++hw) {
                        int64_t idx = n * C * HW + c * HW + hw;
                        int64_t c_idx = c;  // gamma/beta索引
                        float normalized = (input_data[idx] - group_mean) / std_dev;
                        output_data[idx] = gamma_data[c_idx] * normalized + beta_data[c_idx];
                    }
                }
            }
        }

        return Status::OK();
    }

    std::string GetKernelSignature() const override {
        return "GroupNorm_NPU_FLOAT32_FLOAT32";
    }

private:
    int num_groups_;
    float epsilon_;
};

// 形状推导
Status GroupNormShapeFunc(ShapeInferenceContext* context) {
    const TensorShape& input_shape = context->GetInputShape(0);

    // 输出形状与输入相同
    context->SetOutputShape(0, input_shape);

    // 统计量形状: [N, num_groups]
    int num_groups = context->GetAttr<int>("num_groups");
    TensorShape stat_shape({input_shape.dim(0), num_groups});
    context->SetOutputShape(1, stat_shape);
    context->SetOutputShape(2, stat_shape);

    return Status::OK();
}

// 注册
REGISTER_OP(GroupNorm)
    .Input("input: float")
    .Input("gamma: float")
    .Input("beta: float")
    .Output("output: float")
    .Output("mean: float")
    .Output("variance: float")
    .Attr("num_groups: int >= 1")
    .Attr("epsilon: float = 1e-5")
    .SetShapeFn(GroupNormShapeFunc);

3.3 算子融合开发

实现一个融合算子：Conv2D + BatchNorm + ReLU：

// 融合算子定义
OpDef DefineConvBNReLUFusedOp() {
    return OpDef("ConvBNReLU")
        .Input("input", DT_FLOAT)
        .Input("filter", DT_FLOAT)
        .Input("bias", DT_FLOAT)
        .Input("bn_scale", DT_FLOAT)
        .Input("bn_offset", DT_FLOAT)
        .Input("bn_mean", DT_FLOAT)
        .Input("bn_var", DT_FLOAT)
        .Output("output", DT_FLOAT)
        .Attr<int>("strides", {1, 1})
        .Attr<std::vector<int>>("padding", {0, 0})
        .Attr<float>("epsilon", 1e-5f)
        .Description("Fused Conv2D + BatchNorm + ReLU");
}

// 融合算子内核
class ConvBNReLUFusedKernel : public OpKernelRegistry::OpKernel {
public:
    ConvBNReLUFusedKernel(const std::vector<int>& strides,
                         const std::vector<int>& padding,
                         float epsilon)
        : strides_(strides), padding_(padding), epsilon_(epsilon) {}

    Status ComputeShape(OpKernelContext* context) override {
        // 使用卷积的形状推导逻辑
        return shape_inference::Conv2DShape(context);
    }

    Status Compute(OpKernelContext* context) override {
        const Tensor* input = context->GetInput(0);
        const Tensor* filter = context->GetInput(1);
        const Tensor* bias = context->GetInput(2);
        const Tensor* bn_scale = context->GetInput(3);
        const Tensor* bn_offset = context->GetInput(4);
        const Tensor* bn_mean = context->GetInput(5);
        const Tensor* bn_var = context->GetInput(6);

        // 分配输出
        Tensor* output;
        RETURN_IF_ERROR(context->AllocateOutput(0, ComputeOutputShape(input, filter),
                                               DT_FLOAT));
        output = context->GetOutput(0);

        // 1. 预计算融合的权重和偏置
        // 融合公式: output = ReLU(Conv(x) * BN)
        // 其中BN可以融合到卷积权重和偏置中
        // new_weight = weight * gamma / sqrt(var + epsilon)
        // new_bias = (bias - mean) * gamma / sqrt(var + epsilon) + beta

        int64_t out_channels = filter->shape().dim(0);
        std::vector<float> fused_weight(out_channels);
        std::vector<float> fused_bias(out_channels);

        const float* weight_data = filter->data<float>();
        const float* bias_data = bias->data<float>();
        const float* scale_data = bn_scale->data<float>();
        const float* offset_data = bn_offset->data<float>();
        const float* mean_data = bn_mean->data<float>();
        const float* var_data = bn_var->data<float>();

        for (int64_t oc = 0; oc < out_channels; ++oc) {
            float scale = scale_data[oc] / std::sqrt(var_data[oc] + epsilon_);
            fused_weight[oc] = scale;
            fused_bias[oc] = (bias_data[oc] - mean_data[oc]) * scale + offset_data[oc];
        }

        // 2. 执行融合卷积（一次性完成卷积、BN和ReLU）
        return ComputeFusedConv(input, filter, fused_weight, fused_bias,
                               output, context);
    }

    std::string GetKernelSignature() const override {
        return "ConvBNReLU_NPU_FLOAT32_FLOAT32";
    }

private:
    Status ComputeFusedConv(const Tensor* input, const Tensor* filter,
                          const std::vector<float>& fused_scale,
                          const std::vector<float>& fused_bias,
                          Tensor* output, OpKernelContext* context) {
        // 调用NPU融合算子API
        aclnnStatus ret = aclnnConvBnRelu(
            (aclTensor*)input->device_ptr(),
            (aclTensor*)filter->device_ptr(),
            fused_scale.data(),
            fused_bias.data(),
            (aclTensor*)output->device_ptr(),
            strides_.data(),
            padding_.data(),
            context->GetStream().GetHandle(),
            context->GetWorkspace());

        return ret == ACLNN_SUCCESS ? Status::OK() :
               Status::Error("Fused conv execution failed");
    }

    TensorShape ComputeOutputShape(const Tensor* input, const Tensor* filter) {
        // 简化版本
        return input->shape();
    }

    std::vector<int> strides_;
    std::vector<int> padding_;
    float epsilon_;
};

// 注册融合算子
REGISTER_OP(ConvBNReLU)
    .Input("input: T")
    .Input("filter: T")
    .Input("bias: T")
    .Input("bn_scale: T")
    .Input("bn_offset: T")
    .Input("bn_mean: T")
    .Input("bn_var: T")
    .Output("output: T")
    .Attr("T: {float, float16} = DT_FLOAT")
    .Attr("strides: list(int) = [1, 1]")
    .Attr("padding: list(int) = [0, 0]")
    .Attr("epsilon: float = 1e-5")
    .SetShapeFn(shape_inference::Conv2DShape);

---

四、高级特性

4.1 多后端支持

// 多后端内核注册
template<typename Device>
class GenericAddKernel : public OpKernelRegistry::OpKernel {
public:
    Status Compute(OpKernelContext* context) override {
        const Tensor* a = context->GetInput(0);
        const Tensor* b = context->GetInput(1);
        Tensor* c;

        RETURN_IF_ERROR(context->AllocateOutput(0, a->shape(), a->dtype()));
        c = context->GetOutput(0);

        return Device::ComputeAdd(a, b, c);
    }
};

// CPU特化
template<>
class GenericAddKernel<CPUDevice> : public OpKernelRegistry::OpKernel {
public:
    Status Compute(OpKernelContext* context) override {
        // CPU实现
        auto* a = context->GetInput(0)->data<float>();
        auto* b = context->GetInput(1)->data<float>();
        auto* c = context->GetOutput(0)->data<float>();
        int64_t n = context->GetInput(0)->shape().num_elements();

        #pragma omp parallel for
        for (int64_t i = 0; i < n; ++i) {
            c[i] = a[i] + b[i];
        }
        return Status::OK();
    }
};

// NPU特化
template<>
class GenericAddKernel<NPUDevice> : public OpKernelRegistry::OpKernel {
public:
    Status Compute(OpKernelContext* context) override {
        // NPU实现
        return aclnnAdd(context->GetInput(0), context->GetInput(1),
                       context->GetOutput(0), context->GetStream());
    }
};

// 注册不同后端
REGISTER_OP_KERNEL(Add, CPU, FLOAT32).Kernel(GenericAddKernel<CPUDevice>);
REGISTER_OP_KERNEL(Add, NPU, FLOAT32).Kernel(GenericAddKernel<NPUDevice>);

4.2 自动微分支持

// 为算子添加梯度计算
class LeakyReluGradientOp : public OpKernelRegistry::OpKernel {
public:
    Status Compute(OpKernelContext* context) override {
        const Tensor* grad_output = context->GetInput(0);
        const Tensor* input = context->GetInput(1);
        Tensor* grad_input;

        RETURN_IF_ERROR(context->AllocateOutput(0, input->shape(),
                                               input->dtype()));
        grad_input = context->GetOutput(0);

        float alpha = context->GetAttr<float>("alpha");

        const float* grad_data = grad_output->data<float>();
        const float* input_data = input->data<float>();
        float* out_data = grad_input->data<float>();
        int64_t n = input->shape().num_elements();

        for (int64_t i = 0; i < n; ++i) {
            // dy/dx = alpha if x < 0 else 1
            out_data[i] = input_data[i] < 0 ?
                grad_data[i] * alpha : grad_data[i];
        }

        return Status::OK();
    }
};

// 注册梯度算子
REGISTER_OP(LeakyReluGrad)
    .Input("grad_output: T")
    .Input("input: T")
    .Output("grad_input: T")
    .Attr("T: {float, float16}")
    .Attr("alpha: float = 0.01");

REGISTER_OP_GRADIENT(LeakyRelu, LeakyReluGradientOp);

4.3 动态形状支持

// 处理动态形状的算子
class DynamicShapeExampleOp : public OpKernelRegistry::OpKernel {
public:
    Status ComputeShape(OpKernelContext* context) override {
        const TensorShape& input_shape = context->GetInputShape(0);

        // 即使输入形状包含未知维度，仍可推导输出形状模式
        TensorShape output_shape;
        if (input_shape.IsUnknown()) {
            output_shape = TensorShape::Unknown();
        } else {
            // 第一个维度可以动态变化
            output_shape = TensorShape({TensorShape::kUnknownDim,
                                       input_shape.dim(1),
                                       input_shape.dim(2) * 2});
        }

        context->SetOutputShape(0, output_shape);
        return Status::OK();
    }

    Status Compute(OpKernelContext* context) override {
        // 在运行时确定实际形状
        const Tensor* input = context->GetInput(0);
        const TensorShape& actual_shape = input->shape();

        // 根据实际形状分配输出
        TensorShape output_shape({actual_shape.dim(0),
                                  actual_shape.dim(1),
                                  actual_shape.dim(2) * 2});

        Tensor* output;
        RETURN_IF_ERROR(context->AllocateOutput(0, output_shape,
                                               input->dtype()));
        output = context->GetOutput(0);

        // 执行计算...
        return Status::OK();
    }
};

// 使用形状推导规则
REGISTER_OP(DynamicShapeExample)
    .Input("input: T")
    .Output("output: T")
    .Attr("T: {float}")
    .SetShapeFn([](ShapeInferenceContext* ctx) {
        // 定义形状推导规则
        ctx->SetOutputShape(0, ctx->GetInputShape(0));
        return Status::OK();
    });

---

五、总结

opbase框架提供了完整的算子开发基础设施，通过算子定义、注册、形状推导和执行等层次的抽象，使得开发者可以专注于算子逻辑本身。掌握opbase的使用，对于深度学习框架开发和自定义算子实现都至关重要。

相关链接：

• CANN组织链接：https://atomgit.com/cann
• opbase仓库链接：https://atomgit.com/cann/opbase