前言

任何一个深度学习框架要将计算图映射到昇腾 NPU 上执行,中间都必须经历一个从高层语义到低层执行的翻译过程。昇腾 昇腾NPU上的CANN 的 Graph Engine(GE)承担了这个翻译角色,而 MetaDef 正是 GE 赖以运转的元定义基础设施。如果把 GE 想象成一条编译流水线,那么 MetaDef 就是这条流水线上的"零件目录"——它告诉编译器每个算子长什么样、需要什么样的输入输出格式、在图优化阶段可以施加哪些变换。没有 MetaDef,GE 就无从知道一个 Add 算子的输入张量支持哪些 Format,也无法在编译期决定是否将相邻的 MatMul 和 Add 融合为单个算子。

MetaDef 这个名字本身揭示了它的职责边界:Meta 代表元信息定义,Def 代表 Definition。它并不直接执行计算,而是为计算执行提供结构化描述。在 CANN 全面开源之后,MetaDef 的内部实现变得可追溯可审计,这对于需要深度定制算子库或自行编写图优化 Pass 的工程师而言,是一个关键的研究入口。

本文将以深度实践的视角,从算子注册、图定义、格式推导、Pass 注册与执行四个维度,剖析 MetaDef 的架构设计及其在昇腾 NPU 计算流水线中的实际作用。所有代码段均来源于社区开源仓库的真实结构,经过精简以便聚焦核心逻辑。

MetaDef 在 CANN 编译层的位置

CANN 的计算编译层由 Graph Compiler(GE)、BiSheng 编译器和 ATC 编译器共同组成。其中 GE 是前端图编译的核心引擎,负责接收来自 PyTorch、TensorFlow 等训练框架的计算图,经过一系列优化 Pass 后,交给 BiSheng 或 ATC 进行后端代码生成。MetaDef 在这套流程中扮演"信息中枢"的角色——GE 的每一个优化决策都建立在 MetaDef 提供的元信息之上。

从依赖关系看,GE 直接依赖 MetaDef,而算子仓库(ops-math、ops-nn、ops-tensor 等)通过 opbase 间接依赖 MetaDef 的注册接口。Runtime 在执行阶段也会查询 MetaDef 中记录的算子原型信息,用于校验输入输出的合法性。这种"三方共同依赖"的设计使得 MetaDef 成为编译层和算子层之间的唯一契约点。

值得强调的是,MetaDef 并不是一个独立运行的进程或服务。它以静态库和头文件的形式被 GE、算子库和运行时链接,在编译期和加载期分别提供元信息查询能力。这种静态链接的设计避免了运行时的额外进程间通信开销,对于需要频繁查询算子属性的编译流水线来说至关重要。

算子原型注册:Operator 的结构化描述

MetaDef 最基础的能力是算子原型注册。每一个算子在进入 CANN 计算流水线之前,必须先在 MetaDef 中注册其原型信息——包括算子名称、输入输出张量的数量与类型、支持的数据格式(Format)以及属性参数列表。

注册过程通常发生在算子库初始化阶段。以 ops-math 仓库中的加法算子为例,其注册逻辑会通过 MetaDef 提供的宏接口完成声明。注册信息的核心数据结构是 OperatorProto,它描述了算子的"接口签名"。

// 算子原型注册示例(简化后的核心结构)
IMPLEMENT_OPERATOR(Add)
    .Input("x1", "Tensor", "第一个输入张量")
    .Input("x2", "Tensor", "第二个输入张量")
    .Output("y", "Tensor", "输出张量")
    .Attr("dtype", "DataType", "计算时的数据类型")
    .SetTypeInfer(AddInferShape)      // WHY: 形状推导函数让 GE 在编译期推断输出维度,避免运行时才发现维度不匹配
    .SetFormatInfer(AddInferFormat)    // WHY: 格式推导函数决定输出使用 ND 或 NC1HWC0,直接影响后续算子能否融合
    .SetSupportFormat({FORMAT_ND, FORMAT_NC1HWC0, FORMAT_NCHW});

这段代码展示了一个典型算子注册的骨架。Input 和 Output 宏声明了算子的张量端口,Attr 定义了可配置的属性参数。SetTypeInfer 和 SetFormatInfer 分别绑定了形状推导和格式推导函数——这两个函数是 GE 在图优化阶段最频繁调用的元信息接口。

形状推导函数的逻辑相对直观:根据输入张量的 shape 和属性参数,计算出输出张量的 shape。而格式推导函数的设计则要复杂得多,因为它需要考虑昇腾达芬奇架构中不同计算单元对数据排布的要求。Cube 单元偏好 NC1HWC0 五维格式,Vector 单元偏好 ND 四维格式,同一个算子在不同执行路径上可能需要不同的格式输出。MetaDef 通过 AllowExplicitFormat 机制允许算子开发者标注"本算子支持哪些格式转换",从而让 GE 的格式推导 Pass 有足够的信息做出决策。

在实际工程中,算子注册的完整性直接影响 GE 的编译成功率。如果某个自定义算子遗漏了格式推导函数的注册,GE 在遇到该算子时就无法自动选择合适的格式转换策略,最终导致编译失败或运行时格式不匹配错误。这是自定义算子开发中最常见的踩坑点之一。

图定义与中间表示:Graph、Tensor 和 Format 的协同

算子注册完成后,MetaDef 提供的另一组核心能力是图定义与中间表示。在 GE 的编译流水线中,计算图由 Graph 对象表示,图中的节点是 Operator,边是 Tensor。这三者构成了 MetaDef IR(中间表示)的基本三角关系。

Graph 对象本身并不存储具体的张量数据,它存储的是计算拓扑和类型信息。当 GE 从训练框架接收到一张计算图时,首先会将其转换为一个 MetaDef 格式的 Graph 对象。转换过程中,GE 会查询 MetaDef 中已注册的算子原型,校验每个节点的输入输出是否与原型定义一致。

Tensor 在 MetaDef IR 中是一个轻量级的描述对象,包含 shape、dtype、format 和 origin_format 四个关键字段。其中 shape 描述张量的维度信息,dtype 描述数据类型,format 描述当前的数据排布格式,origin_format 记录原始的框架侧格式。这种四字段设计的原因在于:训练框架通常使用 NCHW 或 NHWC 等通用格式,而昇腾 NPU 内部为了适配达芬奇架构的 Cube 和 Vector 单元,会使用 NC1HWC0 等硬件友好格式。origin_format 字段确保在需要将张量数据回传给框架侧时,能够正确还原格式。

// TensorDesc 的核心字段(简化描述)
struct TensorDesc {
    Shape shape;              // 张量形状,如 [1, 3, 224, 224]
    DataType dtype;           // 数据类型,如 DT_FLOAT16
    Format format;            // 当前数据排布,如 FORMAT_NC1HWC0
    Format origin_format;     // 原始框架侧格式,如 FORMAT_NCHW
};

// WHY: 保存 origin_format 是因为昇腾 NPU 内部会将 NCHW 自动拆分重排为 NC1HWC0
// 但当 GE 需要将部分结果反馈给 PyTorch 时,必须知道原始格式才能正确还原
// 如果丢失这个信息,就会出现推理结果维度正确但数值完全错乱的隐蔽 bug

Format 转换在 MetaDef 中不是自动发生的,而是由 GE 的 FormatTransfer 机制驱动的。FormatTransfer 注册了一系列格式转换规则,每条规则描述了"从格式 A 到格式 B 的转换条件"和"转换算法"。GE 在图优化阶段遍历 Graph 中的每条边,检查上下游算子的格式要求是否匹配。如果不匹配,就在两个算子之间插入一个 TransData 节点,执行格式转换。

但频繁插入 TransData 节点会带来性能问题——每次格式转换都需要显存读写,在数据量大的场景下会成为瓶颈。因此 MetaDef 的格式推导设计强调"格式传播"而非"格式转换":尽可能让算子直接支持上下游的格式,减少 TransData 节点的插入。这就是前面提到的 SetSupportFormat 和 SetFormatInfer 的工程价值所在。

图优化 Pass 注册与执行引擎

MetaDef 的第三大能力模块是图优化 Pass 的注册与执行框架。GE 的图编译过程本质上是多轮 Pass 依次作用于计算图的过程——每一轮 Pass 对图施加某种优化变换,最终生成适合后端执行的优化图。

MetaDef 将每个 Pass 抽象为一个独立对象,通过注册机制让 GE 的 PassManager 统一调度。一个 Pass 的注册信息包含名称、执行阶段、依赖关系和作用范围。执行阶段决定了该 Pass 在编译流水线中的位置,依赖关系决定了 Pass 之间的执行顺序,作用范围决定了该 Pass 是全局作用还是仅作用于特定子图。

Pass 的执行遵循"注册-匹配-应用"三段式流程。注册阶段,Pass 向 MetaDef 声明自己的触发条件。匹配阶段,PassManager 在当前图上搜索满足触发条件的子图模式。应用阶段,匹配成功后 Pass 对子图执行变换操作,并更新 Graph 对象。

以常量折叠 Pass 为例,该 Pass 在图的初始构建阶段执行,用于将编译期可确定的常量表达式提前计算。当一个 Add 节点的两个输入都是常量节点时,常量折叠 Pass 就会将其替换为一个单常量节点,消除运行时的无效计算。

// Pass 注册与匹配的核心逻辑(简化)
class ConstantFoldPass : public PassBase {
public:
    Status Run(NodePtr &node) override {
        // 检查当前节点是否为可折叠的运算
        if (!IsFoldable(node)) {
            return SUCCESS;  // WHY: 逐节点遍历图中所有节点,不可折叠的直接跳过
                              // 而非一次性收集所有可折叠节点再处理
                              // 这样可以尽早释放不需要继续处理的节点引用
        }

        // 提取常量输入并执行计算
        TensorPtr result = ComputeConstFold(node);

        // 用常量节点替换原来的运算节点
        ReplaceNodeWithConst(node, result);
        // WHY: 直接在原 Graph 上执行节点替换而非构建新图
        // 因为常量折叠是高频 Pass,每张图可能触发数十次
        // 构建-拷贝-替换的开销远大于原地修改

        return SUCCESS;
    }
};

REGISTER_PASS(ConstantFoldPass)
    .SetStage(PASS_STAGE_BEFORE_OPTIMIZE)
    .SetDepends({});  // WHY: 空依赖意味着该 Pass 可以在流水线最早阶段执行
                       // 常量折叠的结果可能触发后续更多优化机会
                       // 放在越前面,后续 Pass 的优化空间越大

Pass 之间的执行顺序并不是完全由开发者手动指定的。MetaDef 的 PassManager 实现了一套基于依赖拓扑排序的调度算法。每个 Pass 声明自己的前置依赖 Pass,PassManager 构建依赖图后进行拓扑排序,自动确定执行顺序。这种设计的优势在于:新增一个 Pass 时,只需声明它依赖哪些已有 Pass,无需手动修改全局执行序列,降低了多 Pass 协作开发的复杂度。

在昇腾 NPU 的实际推理场景中,Pass 的作用效果非常显著。一个未经优化的计算图可能包含数百个 TransData 节点、冗余的形状计算节点和不必要的内存拷贝操作。经过格式对齐、算子融合、内存优化等多轮 Pass 之后,最终图的结构会大幅精简,可执行的算子数量通常能减少一个数量级。这种编译期优化将运行时的计算和内存压力大幅前移,是昇腾 NPU 能够保持高推理吞吐量的关键因素之一。

格式推导的工程实践

格式推导是 MetaDef 框架中最具工程复杂度的模块。它需要同时满足三个约束:上游算子的输出格式与下游算子的输入格式兼容、格式转换的代价最小化、以及最终的格式选择对后续算子融合不会造成阻碍。

GE 在格式推导阶段会为每个算子维护一个"候选格式集合"。这个集合来源于算子注册时声明的 SupportFormat。推导算法从图的输出节点开始,逆向遍历到输入节点,在每条边上尝试寻找上下游格式的交集。如果交集为空,则需要引入 FormatTransfer 进行格式转换;如果交集包含多个格式,则需要通过代价模型选择最优格式。

代价模型的输入参数包括格式转换的计算开销(读取和重排数据的操作数)、格式转换的显存占用(中间缓冲区大小)、以及该格式对后续算子融合的影响。综合这三个维度打分后,代价模型选择得分最低的格式组合作为最终方案。

// 格式推导的核心判断逻辑(简化)
Status InferFormatForEdge(EdgePtr edge) {
    TensorDesc src_desc = edge->GetSrcTensorDesc();
    TensorDesc dst_desc = edge->GetDstTensorDesc();

    // 尝试寻找上下游都支持的格式交集
    std::vector<Format> intersection =
        FindFormatIntersection(src_desc.supported_formats,
                               dst_desc.supported_formats);

    if (intersection.empty()) {
        // WHY: 当格式完全不兼容时必须插入 TransData 节点
        // 但记录这条边的转换代价供全局优化参考
        InsertTransDataNode(edge, src_desc.format, dst_desc.format);
        return SUCCESS;
    }

    // 交集存在时,选择对后续融合最有利的格式
    Format optimal = SelectOptimalFormat(intersection, edge);
    // WHY: 不直接选第一个交集格式,而是通过代价模型评估
    // 因为 NC1HWC0 虽然对 Cube 算子友好,但可能导致相邻的
    // Vector 算子需要额外的格式回退开销
    // 全局最优不一定等于局部格式兼容

    src_desc.format = optimal;
    dst_desc.format = optimal;
    edge->UpdateTensorDesc(src_desc, dst_desc);
    return SUCCESS;
}

格式推导的一个典型踩坑场景是"格式死锁":当算子 A 只输出 NC1HWC0,算子 B 只接受 ND,而算子 C 只接受 NC1HWC0,且 B 位于 A 和 C 之间时,无论 B 选择哪种中间格式,都会导致至少一次格式转换。在这种情况下,MetaDef 的代价模型需要权衡"在 A-B 之间转换"还是"在 B-C 之间转换"哪个代价更低,而不是简单地报错终止编译。

对于需要自行开发图优化 Pass 的工程师来说,理解格式推导的机制尤为重要。自定义 Pass 在修改图结构时,必须确保修改后的格式约束仍然可满足。否则,即使图拓扑正确,也会在后续的格式推导阶段失败。

MetaDef 与算子库的协同工作机制

MetaDef 和各算子仓库之间的协同关系可以概括为"注册-查询-执行"三个阶段。注册阶段发生在算子库加载时,每个算子通过 opbase 提供的注册接口向 MetaDef 注册原型信息。查询阶段发生在 GE 编译期,GE 通过 MetaDef 的查询接口获取算子属性、格式支持和推导规则。执行阶段发生在运行时,Runtime 根据 MetaDef 中记录的算子原型校验实际输入的合法性。

opbase 作为所有算子仓库的基础依赖,封装了 MetaDef 的注册接口。算子开发者通常不需要直接调用 MetaDef 的底层 API,而是通过 opbase 提供的宏和辅助函数完成注册。这种分层设计的目的是降低算子开发者的心智负担——他们只需要关注算子的计算逻辑和接口声明,而不需要理解 MetaDef 内部的注册数据结构。

在 CANN 全面开源的背景下,MetaDef 的内部实现已经可以在社区仓库中直接查阅。这意味着工程师不仅可以了解算子的注册方式,还可以深入理解 GE 的 Pass 调度机制和格式推导算法的实现细节,为深度定制编译流水线提供了完整的参考路径。

效率对比:基于 MetaDef 的图优化效果

为了直观展示 MetaDef 驱动的图优化对昇腾 NPU 推理效率的影响,下面对比了启用与未启用 MetaDef 图优化 Pass 时的典型推理表现。测试场景为包含 MatMul、Add、ReLU、LayerNorm 等多个算子的 Transformer 编码器模块,在 Ascend 910 上执行推理任务。数据为概括性描述,反映典型场景下的量级差异。

指标 使用前(禁用图优化) 使用后(启用图优化)
可执行算子数量 数百个(含大量 TransData 和冗余节点) 数十个(经融合和消除后精简)
格式转换节点数量 占总节点数约三成以上 通常可完全消除或仅保留极少数必要转换
编译耗时 较短(跳过大部分优化) 适中(多轮 Pass 处理带来额外编译开销)
推理延迟 较高(运行时承担大量格式转换和冗余计算) 显著降低(编译期优化消除了运行时瓶颈)
显存占用 较高(中间结果全部保留) 有效降低(算子融合减少了中间缓冲区需求)
算子融合覆盖度 无融合 多组算子融合(如 MatMul+Add+ReLU)

从表格可以看出,图优化的核心收益并非来自单个算子的计算加速,而是来自图结构层面的精简。减少 TransData 节点消除了昂贵的显存重排操作,算子融合减少了中间结果的读写次数,这两者共同将推理延迟和显存占用压缩到接近理论下限。

MetaDef 与 GE 的协作机制

MetaDef 作为 GE 和算子库之间的中间层,其协作机制值得深入理解。当用户通过 MindSpore 或 PyTorch 构建一个计算图时,整个流程是这样的:

  • 前端框架将用户代码转换为计算图 IR,这个 IR 中的每个节点都是一个算子调用。
  • MetaDef 接收这个 IR,并为每个算子节点查找注册表中的原型信息。原型信息包括算子的输入输出描述、属性定义、支持的 Format 和 DataType 等。
  • GE 从 MetaDef 获取原型信息后,开始执行图优化 Pass。每个 Pass 都在 MetaDef 注册过,GE 按照注册的优先级依次调用。
  • 优化完成后,GE 根据最终的计算图调用算子库中的实现。

这个流程中 MetaDef 扮演的是"信息中介"的角色。它不执行任何计算,也不做优化决策,只提供算子元信息供 GE 查询。这种设计的好处是解耦——算子库可以独立更新,只要通过 MetaDef 注册的信息不变,GE 的优化逻辑就不需要修改。


仓库地址:https://atomgit.com/cann/metadef

Logo

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

更多推荐