Format 推导Infer Format特性分析【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1. 特性背景1.1 问题的本质深度学习框架PyTorch、TensorFlow 等在构造计算图时用户关注的是计算语义——张量的维度、算子的数学含义以及数据依赖关系。但昇腾 AI 处理器Ascend NPU的硬件架构对数据的内存布局有特定要求例如Conv2D 的图片输入在硬件上亲和 NC1HWC0 格式将 C 轴按 16 对齐拆分MatMul 的权重亲和 FRACTAL_NZ 格式不同算子对格式有不同的支持能力和性能偏好用户以 NCHW 或 NHWC 等通用格式描述的模型在昇腾设备上执行时需要被转换为硬件亲和的内存布局。这涉及两个核心问题语义理解如何正确还原用户在整张计算图中表达的格式语义执行优化如何为算子选择合适的执行格式并尽量减少数据重排TransData的开销1.2 为什么需要两套格式字段GE 引入了Origin Format和Storage Format也称 Running Format两套表示体系其根本原因是语义正确性和执行效率需要独立建模视角字段职责来源Originorigin_format_表达用户构造计算图时的原始格式语义前端框架或用户显式指定Storageformat_描述实际执行时的内存布局编译过程中推导得到如果只用一个 Format 字段会导致以下问题语义丢失当 FE融合引擎将 NCHW 转为 NC1HWC0 后原始的 NCHW 语义无处保存后续优化 Pass 无法判断这个张量原本代表什么优化受限Data Dump、Profiling 等调试场景需要将 NC1HWC0 的数据转回 NCHW 供用户理解没有 Origin Format 就无法正确还原格式传播混乱整网格式推导需要锚点如 Conv2D 的 data_format 属性如果执行格式覆盖了语义格式推导的起点就丢失了具体而言以一个 NCHW 张量[8, 3, 224, 224]为例字段值含义OriginFormatNCHW用户定义的语义格式OriginShape[8, 3, 224, 224]用户理解的维度StorageFormatNC1HWC0实际内存布局StorageShape[8, 1, 224, 224, 16]C3 向上对齐到 16 后的实际存储形态仅从 StorageShape[8, 1, 224, 224, 16]无法唯一还原其 OriginFormat ——它可能来自 NCHWC3也可能来自 NHWC因此两个字段必须共存。2. 用户使用场景2.1 离线编译场景atc用户使用 atc 工具将 ONNX/PB 模型编译为 OM 文件时GE 需要自动完成整网的格式推导模型输入以 NCHW 或 NHWC 格式定义GE 推导出每个算子的 OriginFormatFE 根据算子能力和性能偏好选择 StorageFormat最终生成包含正确内存布局信息的 OM 模型2.2 在线训练/推理场景通过 TorchAir 或 TFA 集成时框架传入的计算图可能未显式标注所有算子的格式。GE 需要在图编译阶段推导图中所有张量的 OriginFormat在PrepareRunningFormatRefiner阶段根据用户设置的storage_format属性刷新 Data/NetOutput 节点确保 InferShape 在正确的格式上下文中执行2.3 Data Dump 与 Profiling用户在调试时需要查看算子的输入输出数据。数据在设备上以 StorageFormat 存储但用户理解的是 OriginFormat。GE 通过保存ATTR_NAME_DATA_DUMP_ORIGIN_FORMAT属性在 Dump 时将数据从 StorageFormat 转回 OriginFormat。3. 对外接口3.1 GeTensorDesc 上的格式接口GeTensorDesc定义于inc/graph_metadef/graph/ge_tensor.h是张量描述的核心类提供以下格式相关接口GeTensorDesc ├── GetFormat() / SetFormat() // StorageFormat 的读写 ├── GetOriginFormat() / SetOriginFormat() // OriginFormat 的读写 ├── GetShape() / SetShape() // StorageShape 的读写 ├── GetOriginShape() / SetOriginShape() // OriginShape 的读写在graph_metadef/graph/normal_graph/tensor.cc的TensorDescImpl类中可以看到这两个字段独立存储format_对应 StorageFormat默认FORMAT_NDorigin_format_对应 OriginFormat默认FORMAT_NDorigin_format_is_set_标记 OriginFormat 是否被显式设置3.2 StorageFormat 描述体在运行时gert 命名空间StorageFormat是一个同时携带 Origin 和 Storage 信息的描述体定义于inc/graph_metadef/external/graph/types.h其构造方式为StorageFormat(origin_format, storage_format, expand_dims_type)在graph_metadef/register/shape_inference.cc的GetTensorHolder函数中可以看到创建 Tensor 时会将GeTensorDesc的两个格式字段分别映射{input_desc.GetOriginFormat(), input_desc.GetFormat(), {}}即StorageFormat描述体的第一个参数是 OriginFormat第二个参数是 StorageFormat。3.3 算子 InferFormat 注册接口算子开发者可通过以下方式注册格式推导函数IMPL_OP(OpType).InferFormat(infer_format_func)其中infer_format_func签名为UINT32(InferFormatContext *context)。V2 接口通过InferFormatContext提供更结构化的输入输出访问。在graph_metadef/register/shape_inference.cc的UpdateOpDescOutFormat函数中V2 推导完成后会将结果写回 OpDescdesc-SetOriginFormat(format-GetOriginFormat()); desc-SetFormat(format-GetStorageFormat());3.4 用户指定 StorageFormat用户可通过在 TensorDesc 上设置属性来指定算子的 StorageFormatAttrUtils::SetInt(tensor_desc, ATTR_NAME_STORAGE_FORMAT, format_value); AttrUtils::SetListInt(tensor_desc, ATTR_NAME_STORAGE_SHAPE, shape_dims);这些属性在compiler/graph/preprocess/graph_prepare.cc的PrepareRunningFormatRefiner阶段被消费用于刷新 Data 节点和 NetOutput 节点的格式。3.5 Format 枚举定义GE 支持的格式类型定义于inc/framework/executor_c/types.h核心格式包括格式说明典型用途FORMAT_NDN 维张量默认格式不携带特殊语义FORMAT_NCHWN-C-H-W用户侧常见的卷积格式FORMAT_NHWCN-H-W-CTensorFlow 默认格式FORMAT_NC1HWC0N-C1-H-W-C0Conv2D 在昇腾上的亲和格式FORMAT_FRACTAL_ZC1HW-N1-N0-C0Conv2D Filter 的亲和格式FORMAT_FRACTAL_NZN1-N0-C1-C0MatMul 权重的亲和格式4. 具体实现4.1 整体流程格式推导在 GE 编译流程中的位置如下关键顺序先推导 OriginFormat再做 InferShape最后处理 StorageFormat。这是因为 InferShape 需要在正确的 OriginFormat 上下文中执行Shape 的维度含义依赖 Format而 StorageFormat 的选择发生在 FE 阶段。4.2 Origin Format 推导FormatRefinerOrigin Format 推导的核心实现在graph_metadef/graph/refiner/format_refiner.cc的FormatRefiner::InferOrigineFormat函数中。4.2.1 推导算法推导采用锚点扩散策略具体步骤锚点识别GetAnchorPoints遍历图中所有节点找到输入/输出中存在非 ND 格式的节点作为锚点。这些节点通常是 Conv2D、Pooling 等对格式敏感的算子它们通过属性如data_format已携带格式信息。锚点刷新RefreshOriginFormatOfAnchor对锚点节点如果origin_format仍为 ND 或 RESERVED则将其format值复制到origin_format。这确保锚点自身的 OriginFormat 被正确建立。双向扩散AnchorProcess向后推导BackInferProcess从锚点的输入端出发沿数据流反向传播格式。对每个上游节点如果其origin_format为 ND 且未锁定则将锚点的格式传递过去。向前推导ForwardInferProcess从锚点的输出端出发沿数据流正向传播格式。Data 节点兜底DataNodeFormatProcess对于推导过程中未被触及的 Data 节点通常是因为缺少格式锚点使用图的全局data_format参数统一设置格式。4.2.2 格式传播的中断条件推导过程在以下情况会中断节点格式已锁定ATTR_NAME_FORMAT_LOCKED为 true某些算子的格式不应被推导过程覆盖节点类型为维度变化算子PERMUTE、EXPANDDIMS、SQUEEZE当维度数小于 4 时维度语义不确定不应传播格式遇到标量dim_num 0标量无格式语义遇到 ND 格式ND 表示无特定格式语义传播到此处停止遇到 NetOutput作为图的边界不继续向前传播4.2.3 Ref 反射机制对于 If/Case 等控制流算子GE 通过RefRelations建立子图与主图之间 Data 节点的反射关系。当主图中某个 Data 的格式被推导后通过ReflectionProcess将格式同步到子图对应的 Data 节点及其父节点的输入。4.2.4 算子自定义推导除了默认的格式传播机制算子还可注册自定义的 InferFormat 函数。在NodeUtilsEx::InferOriginFormat中会调用OpDescUtilsEx::CallInferFormatFunc优先使用算子注册的推导函数否则使用DefaultInferFormat将第一个非 ND 格式传播到所有输入输出。4.3 Storage Format 的确定Storage Format 的确定分为两条路径4.3.1 用户显式指定路径在compiler/graph/preprocess/graph_prepare.cc的UpdateDataNetOutputByStorageFormat函数中对 Data 节点从ATTR_NAME_STORAGE_FORMAT和ATTR_NAME_STORAGE_SHAPE属性读取用户指定的 StorageFormat调用ModifyTensorDescStorageFormatAndShape刷新 TensorDesc对 NetOutput 节点同样读取属性并刷新确保输出格式正确对 ConstPlaceHolder 节点处理常量的存储格式ModifyTensorDescStorageFormatAndShape函数的核心操作根据 StorageFormat 和 OriginFormat 计算存储形态StorageShape包括维度的扩展和对齐调用SetFormat设置 StorageFormat注意不是SetOriginFormat调用SetShape设置 StorageShape计算并设置 Tensor 的内存大小4.3.2 FE 自动选择路径对于计算密集型算子如 Conv2D、MatMulFE融合引擎在算子编译阶段根据算子的能力选择最优的 StorageFormat。这一过程由compiler/engines/nn_engine/optimizer/format_selector/下的 FormatSelector 体系实现FormatDtypeOpBuiltinSelector内置算子的格式选择FormatDtypeOpKernelSelector基于算子内核信息的格式选择FormatDtypeOpCustomizeSelector自定义算子的格式选择这些 Selector 通过FormatDtypeManagerBase协调最终由FormatDtypeSetter将选择的格式设置到图节点上。4.4 运行时格式推导InferShape 阶段在运行时的 InferShape 阶段格式信息通过InferFormatContextV2 接口传递给算子的推导函数。在graph_metadef/register/shape_inference.cc的InferFormatOnCompile函数中构造InferFormatContext为每个输入/输出创建CompileTimeTensorDesc其中同时包含origin_format和storage_format调用算子注册的infer_format_func将推导结果通过UpdateOpDescOutFormat写回 OpDesc在GetTensorHolder函数中可以看到 Tensor 的创建方式gert::Tensor(storage_shape, {input_desc.GetOriginFormat(), input_desc.GetFormat(), {}}, input_desc.GetDataType())其中StorageShape描述体同时携带 Origin 和 Storage 的 Shape 信息。4.5 格式差异检测与 TransData 插入在runtime/v2/graph_builder/storage_format.cc中DiffStorageFormat函数检测一个张量的 OriginFormat 与 StorageFormat 是否不同或 Shape 是否不同return td-GetFormat() ! td-GetOriginFormat() || td-GetShape().GetDims() ! td-GetOriginShape().GetDims()AnyDiffStorageFormat检查一个节点的所有输入输出中是否存在任何格式差异。如果存在差异则在 Lowering 阶段插入 TransData 算子来完成实际的格式转换。4.6 InferShape 中的格式刷新在runtime/v2/kernel/common_kernel_impl/infer_shape_compatible.cc的兼容性 InferShape 中有一个关键的格式处理// RT1时算子的infershape只能拿到format字段但是却需要用origin format input_desc-SetFormat(input_desc_in_context-GetOriginFormat()); input_desc-SetOriginFormat(input_desc_in_context-GetOriginFormat());这表明在 RT1运行时第一版兼容模式下InferShape 只能获取到format字段但实际需要的是 OriginFormat。因此需要从 Context 中正确取出 OriginFormat 并同步设置到format和origin_format两个字段。在runtime/v2/kernel/common_kernel_impl/infer_shape.h的TransformOutputShape函数中当 OriginFormat 与 StorageFormat 不同时会调用ShapeTransferAccordingToFormat::TransferShape将 OriginShape 转换为 StorageShapeif (output_td-GetOriginFormat() output_td-GetStorageFormat()) { // 格式相同无需转换 return GRAPH_SUCCESS; } // 格式不同需要根据 StorageFormat 计算 StorageShape TransferShape(origin_format, storage_format, data_type, storage_shape)4.7 整网编译流程中的调用时机在compiler/graph/manager/graph_manager.cc中格式相关的阶段按以下顺序执行PrepareRunningFormatRefiner ← StorageFormat 刷新 → UpdateDataNetOutputByStorageFormat → VariablePrepareOpPass → UpdateInputOutputByOptions → UpdateVariableFormats在此之前的 GraphPrepare::GenerateInfershapeGraph 中InferOriginFormat ← OriginFormat 推导 → FormatRefiner::InferOrigineFormat4.8 格式优化 Pass在 OriginFormat 推导和 StorageFormat 确定之后compiler/graph/passes/format_optimize/下的多个 Pass 负责优化 TransData 的插入和消除Pass功能TransOpSymmetryEliminationPass消除对称的格式转换对如 NCHW→NC1HWC0→NCHWTransOpBreadthFusionPass将同一节点的多个输出侧 TransData 合并TransOpWithoutReshapeFusionPass融合不含 Reshape 的连续 TransDataTransposeTransDataPass将 Transpose 与 TransData 合并优化UnchangedTransposeRemovePass移除不改变数据的 TransposeCastRemovePass移除不必要的 Cast5. 关键设计决策5.1 为什么 FormatRefiner 使用锚点扩散而非全图遍历锚点扩散的优势在于避免无意义传播大量 ElementWise 算子如 Add、ReLU对格式不敏感其格式应与上游保持一致不需要单独推导降低复杂度只在格式语义发生变化的节点锚点附近做推导而非遍历全图支持控制流锚点扩散结合 RefRelations 可以自然处理子图间的格式同步5.2 为什么 OriginFormat 和 StorageFormat 在 GeTensorDesc 中都存储为format_和origin_format_这种设计的核心考虑是origin_format_是 FormatRefiner 阶段确定的表达用户语义一旦确定不应被修改format_初始与origin_format_相同后续在 FE 阶段被修改为 StorageFormat两个字段共存允许任何时刻对比它们来判断是否需要 TransData5.3 为什么 StorageFormat 描述体同时携带 Origin 信息StorageFormatgert 命名空间虽然名字叫Storage但实际上是一个复合描述体。原因是仅靠 StorageFormat 的值如 NC1HWC0无法唯一还原其 OriginFormat可能是 NCHW 或 NHWC必须同时保存两者。这使得运行时无需回溯到 OpDesc 就能获得完整的格式上下文。6. 数据流总结以下是一个典型 Conv2D 网络中格式信息的完整生命周期【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考