opbase:CANN 所有算子的公共地基

翻开 CANN 任何一个算子仓库------ops-math、ops-nn、ops-blas、ops-fft、ops-rand------它们都依赖同一个基础库:opbase。

opbase 不提供面向用户的算子能力。它提供的是所有算子共享的基础设施:Tensor 的内存管理和形状描述、Kernel 的注册和加载机制、Buffer 的分配和生命周期管理。


opbase 在 CANN 中的位置

复制代码
上层算子库:ops-math、ops-nn、ops-blas、ops-fft、ops-rand ......(各自管一类算子)
  ↓ 全部依赖
opbase:
  ┌──────────────────────────────┐
  │ Tensor 数据结构              │
  │ Buffer 内存管理              │
  │ Kernel 注册框架              │
  │ 通用类型定义                 │
  └──────────────────────────────┘
  ↓
CANN 底层:Runtime、driver

每个新算子仓库启动时不需要从零写 Tensor 管理、Buffer 分配、Kernel 注册------这些在 opbase 中已就绪,新仓库直接引用。


为什么所有算子仓库都依赖 opbase

一个算子从写到 Runtime 执行,需要的公共能力包括:

Tensor 描述。 每个算子需要知道输入 Tensor 的形状、数据类型、内存地址。opbase 定义了统一的 TensorDesc 结构------所有算子仓库都使用这个描述。ops-math 和 ops-nn 的算子之间传递 Tensor 信息时不需要做格式转换。

Buffer 分配。 算子执行过程中的临时 Tensor 需要分配显存。opbase 提供了 opbase::AllocBuffer------封装了 Runtime 的 aclrtMalloc,但增加了算子级别的生命周期追踪。Buffer 在算子执行完后自动回收。

Kernel 注册。 写好的算子 Kernel 需要注册到 CANN 的算子表中才能被 GE 识别和调度。opbase 提供了 REGISTER_KERNEL 宏------开发者只需要在 Kernel 实现文件中加一行注册语句。

cpp 复制代码
#include "opbase/opbase.h"

// 自定义 Add Kernel
__aicore__ void AddKernel(...) { ... }

// 注册到算子表------一行搞定
REGISTER_KERNEL("CustomAdd", AddKernel);
// GE 在编译时就能在算子表中找到 "CustomAdd"

Tensor 基础能力如何复用

所有算子仓库中 Tensor 的创建和销毁走同一条路径:

cpp 复制代码
// opbase 的 Tensor 创建------所有算子仓库共用
opbase::Tensor tensor = opbase::AllocTensor(
    {batch, seq_len, hidden_dim},  // Shape
    opbase::FLOAT16,               // 数据类型
    opbase::ND                     // 格式
);

// tensor 的内存由 opbase 管理
// 算子执行完自动回收

每个算子仓库不需要重新实现 Tensor 分配逻辑。

Buffer 的生命周期管理也是复用的:

cpp 复制代码
// ops-math 的 ReduceSum 算子内部
void ReduceSumKernel(...) {
    // opbase 管理的临时 Buffer
    auto temp = opbase::AllocBuffer(partial_sum_size);
    // 计算...
    // temp 在函数结束时自动回收
}

算子开发中的基础设施

写一个新的 CANN 算子时,opbase 提供的工作流:

  1. 继承 OpBase。 新算子继承 opbase::Operator 基类,获得输入输出解析、参数校验、生命周期管理等默认实现
  2. 实现 Compute。 只需要实现 Compute() 虚函数------算子的核心计算逻辑
  3. 注册。 使用 REGISTER_KERNEL 宏将算子注册到算子表
cpp 复制代码
class MyNewOp : public opbase::Operator {
public:
    OpBase::Status Init(const OpDesc& desc) override {
        // opbase 自动解析输入输出描述
        return opbase::SUCCESS;
    }
    
    OpBase::Status Compute(const std::vector<opbase::Tensor>& inputs,
                           std::vector<opbase::Tensor>& outputs) override {
        // 只有这段代码需要开发者写
        auto input = inputs[0].data<float16>();
        auto output = outputs[0].data<float16>();
        // 计算逻辑...
        return opbase::SUCCESS;
    }
};

REGISTER_KERNEL("MyNewOp", MyNewOp);

opbase 的类型系统

opbase 定义了一套跨仓库统一的类型系统:

  • DataType:FLOAT32、FLOAT16、INT8、INT32、INT64、BOOL
  • Format:ND、NZ、NHWC、NCHW
  • Shape:动态 Shape 用 -1 表示,支持 ShapeRange

所有算子仓库都使用同一套类型。ops-math 的 ReduceSum 输出的 DataType 跟 ops-nn 的 Conv 输入的 DataType 是同一个枚举值------不需要类型转换。

opbase 的 Kernel 注册表

Kernel 注册表是 opbase 的另一个核心功能。每个算子库在加载时调用 REGISTER_KERNEL 把 Kernel 注册到全局表中。GE 在编译时查表找到对应算子的 Kernel。

注册表是进程级的------跨模型共享。模型 A 加载了 ops-nn 的 Conv Kernel,模型 B 加载同一个 Operator 时不需要重复注册。这个机制在多模型服务中节省了重复加载 Kernel 的时间。

参考仓库

opbase 算子基础组件库

Ascend C 算子编程指南

相关推荐
知识燃料2 分钟前
企业技术决策者落地 AI Agent 体系,AWS 中国峰会有哪些战略与工程内容可看?
大数据·人工智能
一个王同学8 分钟前
从零到一 | CV转多模态大模型 | week22 | 实战项目-DocuMind-VL:基于 OCR 与 Qwen-VL 的文档多模态问答系统(二)
人工智能·深度学习·机器学习·计算机视觉·ocr
Canace11 分钟前
笔记本都合上了,Claude 为什么还能在手机上执行电脑上装的技能?
前端·人工智能·ai编程
u01030552711 分钟前
昇腾+AtomGit打造AI甘特图应用生产线
人工智能
美狐美颜SDK开放平台12 分钟前
第三方美颜SDK接入教程:直播APP实现实时美颜、滤镜与美型效果的方法
人工智能·深度学习·音视频·sdk·美颜sdk·视频美颜sdk·直播app开发
乱世刀疤12 分钟前
AI Weekly 8.3-8.9
人工智能
AI导出鸭13 分钟前
怎么让豆包做表格?AI导出鸭苹果版将豆包输出的管道表格智能解析为二维结构,一键导出为Excel或Word标准表格。
人工智能·chatgpt·word·excel·ai导出鸭
TDengine (老段)18 分钟前
TDengine 第三方工具 — Telegraf、Kafka Connect、Flink、Spark
大数据·数据库·物联网·flink·kafka·时序数据库·tdengine
霸道流氓气质2 小时前
普通办公电脑基于 Ollama 的本地 AI 能力技术文档
人工智能·电脑
点云-激光雷达-Slam-三维牙齿2 小时前
速度起飞 笔记本电脑6G显卡llama运行Qwen3.6 35BA3B MTP 大模型
人工智能·python·电脑·llama