CANN/GE 零拷贝（Zero Copy）特性

GE 零拷贝Zero Copy特性【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1. 问题背景为什么需要零拷贝在 AI 推理场景中每次模型执行都需要将用户的输入数据拷贝到模型的内存空间再将模型输出拷贝回用户缓冲区。这种 Host-to-Device / Device-to-Host 的拷贝是推理延迟的重要来源。零拷贝的核心目标就是消除不必要的中间拷贝让用户直接在模型可见的内存区域读写数据。GE 的零拷贝覆盖两个方向输入零拷贝Input Zero Copy用户直接将数据写入模型输入内存避免 H2D 拷贝输出零拷贝Output Zero Copy模型计算结果直接写入用户预分配的输出缓冲区避免 D2H/D2D 拷贝2. 架构总览零拷贝特性横跨三层涉及编译期规划和运行时执行3. 用户侧接口与 Option3.1 模型加载选项ge.exec.reuseZeroCopyMemory— 定义于api/acl/acl_model/model/model.cpp通过aclmdlConfigHandle设置用于控制是否复用零拷贝内存。// model.cpp constexpr const char_t *OPTION_EXEC_REUSE_ZERO_COPY_MEMORY ge.exec.reuseZeroCopyMemory; // model.cpp acl::UpdateGraphOptions(OPTION_EXEC_REUSE_ZERO_COPY_MEMORY, std::to_string(handle-reuseZeroCopy));3.2 RT2 LoweringOption定义于inc/graph_metadef/external/exe_graph/lowering/lowering_opt.hstruct LoweringOption { bool trust_shape_on_out_tensor false; // 信任用户输出 shape bool always_zero_copy false; // 总是零拷贝强模式 bool always_external_allocator false; // 总是使用外部 allocator bool enable_single_stream false; };always_zero_copy的含义lowering_opt.h默认关闭。打开后外部调用者必须保证正确申请输出内存大小 ≥ shape 计算的 Tensor 大小、placement 正确打开后可提升 Host 调度性能但零拷贝失效时不做回退处理直接报错3.3 单算子执行场景在api/acl/acl_op_executor/single_op/op_executor.cppACL 单算子执行时强制启用零拷贝gert::LoweringOption oOption; oOption.always_zero_copy true; oOption.always_external_allocator true; auto streamExecutor gert::LoadStreamExecutorFromModelData(modelData, oOption, ret).release();4. 编译器侧实现4.1 内存分配阶段识别零拷贝区域核心问题用户输入/输出的内存地址在编译期不可知但模型内部算子的内存规划必须预留出这块区域。解决方案是将模型输入/输出标记为独立的MemoryBlock不与其他内存块合并,并且将模型输入/输出逻辑地址放在feature map逻辑地址的后面如果用户开启了零拷贝功能模型加载的时候GE就不需要分配这块内存。MemoryBlock.is_zero_copy_ 标记定义于compiler/graph/build/memory/block_mem_assigner.his_zero_copy_为 true 的内存块有以下约束block_mem_assigner.cc不可合并block_mem_assigner.cc零拷贝块之间、零拷贝块与普通块之间不能做内存复用合并因为多个用户输入地址可能不连续graph_mem_assigner.h独立分配偏移graph_mem_assigner.cc零拷贝块在AssignZeroCopyMemory中单独分配 offset且图输入块会按大小排序后放在一起put graph-input-blocks together, so input data can merge H2D copyAssignZeroCopyMemory 流程compiler/graph/build/memory/graph_mem_assigner.ccATTR_IS_ZERO_COPY_BLOCKgraph_mem_assigner.cc(void)ge::AttrUtils::GetBool(tensor_desc, ge::ATTR_IS_ZERO_COPY_BLOCK, is_zero_block);运行时根据此属性判断该 Data 节点的输出是否为零拷贝块。如果不是零拷贝块且启用了reuseZeroCopyMemory则对该地址调用DisableZeroCopydavinci_model.cc。ATTR_MODEL_ZERO_COPY_MEMORY_SIZE编译器在序列化模型时将零拷贝内存总量写入模型属性compiler/graph/build/model_builder.ccGE_CHK_BOOL_EXEC(ge::AttrUtils::SetInt(model, ATTR_MODEL_ZERO_COPY_MEMORY_SIZE, zero_copy_mem_size_), ...);5. 运行时 V1 实现5.1 ZeroCopyOffset偏移映射核心runtime/v1/graph/load/model_manager/zero_copy_offset.h是零拷贝在 V1 运行时的核心数据结构。输入初始化zero_copy_offset.cc从 OpDesc 读取ATTR_ZERO_COPY_BASIC_OFFSET和ATTR_ZERO_COPY_RELATIVE_OFFSET根据是否有 L2 Fusion 分两种路径无 Fusiondata_info {size, virtual_addr}relative_offset 0有 Fusion遍历 basic_offset 匹配计算out_offset virtual_addr relative_offset地址注册zero_copy_offset.ccSetInputOutsideAddrs/SetOutputOutsideAddrs将虚拟地址注册到outside_addrs_映射表运行时地址刷新zero_copy_offset.ccSetOutsideAddrsValue当用户传入新的输入地址时遍历outside_addrs_将所有引用该虚拟地址的任务参数地址task args记录下来。执行时只需修改这些地址即可。5.2 DavinciModel 的零拷贝管理runtime/v1/graph/load/model_manager/davinci_model.hstd::mapuint32_t, ZeroCopyOffset input_data_info_; std::mapuint32_t, ZeroCopyOffset output_data_info_;加载阶段davinci_model.cc为每个 Data 节点创建ZeroCopyOffset并初始化调用SetInputOutsideAddrs注册虚拟地址到real_virtual_addrs_如果启用了reuseZeroCopyMemory检查ATTR_IS_ZERO_COPY_BLOCK非零拷贝块调用DisableZeroCopy执行阶段 — 输入davinci_model.ccCopyInputData→CopyPlainData遍历input_data_info_将用户数据data_buf.data拷贝到data_info.GetBasicAddr()即模型内部地址。这是零拷贝模式下的直接拷贝——用户 buffer → 模型可见内存省去了中间缓冲。执行阶段 — 输出 类似地模型输出通过output_data_info_映射直接将结果写入用户提供的输出 buffer。SetZeroCopyAddrdavinci_model.cc 这是零拷贝运行时的关键操作——将算子任务参数中的地址记录到input_data_info_/output_data_info_的outside_addrs_中。当用户每次传入新地址时只需遍历这些记录的位置替换为新的用户地址即可。DisableZeroCopydavinci_model.cc 将某个地址加入copy_only_addrs_只拷贝不同步地址意味着该地址处的数据每次执行都会被拷贝不做地址替换。5.3 内存分类与任务参数刷新MemoryAppTyperuntime/v1/graph/load/model_manager/memory_app_type_classifier.henum class MemoryAppType : int32_t { kMemoryTypeFix, // 固定地址权重、常量 kMemoryTypeFeatureMap, // Feature Map 内存 kMemoryTypeModelIo, // 零拷贝的模型输入/输出 };TaskArgsRefreshTypeClassifiertask_args_refresh_type_classifier.hstatic constexpr uint64_t kRefreshByModelIo 1UL 0U; // 由 Model IO 触发刷新 static constexpr uint64_t kRefreshByFm 1UL 1U; // 由 Feature Map 触发刷新这两个分类器协同工作MemoryAppTypeClassifier根据逻辑地址判断内存类型TaskArgsRefreshTypeClassifier根据内存类型决定任务参数的刷新策略对于kMemoryTypeModelIo类型的地址任务参数在每次执行时需要被刷新为用户传入的新地址kRefreshByModelIo。5.4 HCCL 零拷贝支持HCCL通信库任务有特殊的零拷贝处理runtime/v1/graph/load/model_manager/task_info/hccl/hccl_task_info.ccHCCL 算子通过input_zero_copy_flag/output_zero_copy_flag标记是否支持零拷贝不支持时调用DisableZeroCopy确保通过拷贝而非地址直通来传递数据支持时通过SetZeroCopyAddr注册地址映射5. 运行时 V2 实现V2 运行时采用了全新的 Lowering 架构零拷贝通过图变换 Pass 实现更加优雅。5.1 ModelOutTensorZeroCopy Passruntime/v2/lowering/pass/model_out_tensor_zero_copy.cc这个 Pass 在 Lowering 阶段运行将模型输出从分配→计算→拷贝到用户模式改为直接在用户 buffer 上计算。图变换过程找到EnsureTensorAtOutMemory节点输出内存保障节点沿数据边找到上游的AllocMemory节点将AllocMemory的类型改为AllocModelOutTensor增加OutputData作为输入AllocModelOutTensorruntime/v2/kernel/outputs/model_outputs.cc 不重新分配内存而是直接引用输出 Tensor 的数据。输出 Tensor 的地址就是用户传入的 buffer 地址。5.2 always_zero_copy 模式当LoweringOption.always_zero_copy true时model_out_tensor_zero_copy.cc进一步将AllocModelOutTensor变换为SplitTensorForOutputDataSplitTensorForOutputDataruntime/v2/kernel/common_kernel_impl/build_tensor.cc不分配任何内存直接将用户传入的 Tensor 切片为输出 Tensor严格校验如果用户 buffer 为 null 或大小不足直接报错不回退不增加引用计数后面不需要 FreeMemory5.3 EnsureTensorAtOutMemoryruntime/v2/kernel/common_kernel_impl/memory_copy.cc这是非always_zero_copy模式下的安全保障。如果输出 Tensor 数据为空零拷贝失败它会尝试从 allocator 分配内存如果用户 buffer 已有数据直接引用零拷贝成功如果没有分配新内存并设置输出6. 使用场景分析场景 1单算子执行ACL Single Op// op_executor.cpp gert::LoweringOption oOption; oOption.always_zero_copy true; oOption.always_external_allocator true;单算子场景下用户直接提供输入输出 buffer强制零拷贝。因为调用者ACL 框架完全控制内存生命周期可以保证输出 buffer 正确。场景 2模型推理V1 Runtime用户通过aclmdlExecute执行模型。此时输入用户的DataBuffer.data被aclrtMemcpy到input_data_info_中的basic_addr模型内存区输出模型计算完后从output_data_info_映射的地址拷贝到用户的DataBuffer在 V1 中零拷贝更多是指地址直接映射而非完全不拷贝——运行时直接操作模型内存区的用户可见部分省去中间缓冲。启用reuseZeroCopyMemory后零拷贝内存可在多次执行间复用。场景 3V2 Runtime 模型加载通过LoadStreamExecutorFromModelData加载可传入LoweringOptionalways_zero_copy false默认输出零拷贝失败时回退到分配新内存always_zero_copy true强制零拷贝失败直接报错关键源文件索引层次文件职责APIapi/acl/acl_model/model/model.cppreuseZeroCopyMemoryoptionAPIapi/acl/acl_op_executor/single_op/op_executor.cpp单算子强制零拷贝公共定义inc/graph_metadef/external/exe_graph/lowering/lowering_opt.hLoweringOption结构体公共定义inc/graph_metadef/graph/debug/ge_attr_define.h零拷贝相关属性常量编译器compiler/graph/build/memory/graph_mem_assigner.ccAssignZeroCopyMemory编译器compiler/graph/build/memory/block_mem_assigner.his_zero_copy_标记编译器compiler/graph/build/model_builder.cc序列化零拷贝大小编译器compiler/graph/passes/memory_conflict/set_input_output_offset_pass.cc偏移映射建立编译器compiler/graph/passes/memory_conflict/memcpy_addr_async_pass.cc硬件兼容性保护编译器compiler/graph/build/task_generator.cc零拷贝偏移表运行时 V1runtime/v1/graph/load/model_manager/zero_copy_offset.h偏移映射核心运行时 V1runtime/v1/graph/load/model_manager/davinci_model.cc零拷贝生命周期管理运行时 V1runtime/v1/graph/load/model_manager/memory_app_type_classifier.h内存类型分类运行时 V1runtime/v1/graph/load/model_manager/task_args_refresh_type_classifier.h任务参数刷新策略运行时 V1runtime/v1/graph/load/model_manager/task_info/hccl/hccl_task_info.ccHCCL 零拷贝运行时 V2runtime/v2/lowering/pass/model_out_tensor_zero_copy.cc输出零拷贝图变换运行时 V2runtime/v2/kernel/common_kernel_impl/build_tensor.ccSplitTensorForOutputData运行时 V2runtime/v2/kernel/common_kernel_impl/memory_copy.ccEnsureTensorAtOutMemory运行时 V2runtime/v2/kernel/outputs/model_outputs.ccAllocModelOutTensor【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考