CANN Runtime 自定义内存分配器Custom Allocator实战AllocatorDesc 创建、注册与回调执行【免费下载链接】runtime本项目提供CANN运行时组件和维测功能组件。项目地址: https://gitcode.com/cann/runtime本文聚焦 CANN RuntimeAscendCL的自定义内存分配器能力围绕仓库 allocator 示例 下的 0_custom_allocator 样例系统讲解分配器描述符AllocatorDesc的创建、四类回调函数绑定、按 Stream 注册/查询/注销的完整流程。读完本文你将掌握aclrtAllocatorCreateDesc、aclrtAllocatorSetAllocFuncToDesc、aclrtAllocatorRegister、aclrtAllocatorGetByStream等关键接口的组合用法并理解其底层实现原理可直接照搬示例搭建自己的内存分配策略。一、能力概述为什么需要自定义分配器在默认情况下CANN Runtime 内部使用内置的内存管理策略完成 Host/Device 内存的申请与释放。但在特定业务场景下开发者往往需要接管内存分配行为——例如实现统一的内存池化、跟踪分配/释放次数用于内存审计、针对特定数据块附加元信息块地址、块大小、或对接自研内存管理框架。本目录聚焦的就是这套自定义分配器描述符、注册和回调函数绑定能力目录说明其核心思路是定义一个分配器描述符AllocatorDesc其中绑定用户自定义的分配器对象和函数指针把描述符注册到指定的 Stream运行时通过 Stream 查询出已注册的分配器与回调由用户自己触发申请、取地址与释放。该能力与仓库中 3_memory_advanced 章节强调的更细粒度的内存控制、自定义分配与流序内存池能力一脉相承是高级内存管理主题下的重要一环。产品支持情况本样例即自定义分配器能力支持以下产品产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√二、核心概念AllocatorDesc 与四类回调函数2.1 类型定义与回调签名自定义分配器的全部对外接口声明位于头文件 include/external/acl/acl_rt_allocator.h。该头文件定义了四种回调函数指针类型typedef void* (*aclrtAllocatorAllocFunc)(aclrtAllocator allocator, size_t size); typedef void (*aclrtAllocatorFreeFunc)(aclrtAllocator allocator, aclrtAllocatorBlock block); typedef void* (*aclrtAllocatorAllocAdviseFunc)(aclrtAllocator allocator, size_t size, aclrtAllocatorAddr addr); typedef void* (*aclrtAllocatorGetAddrFromBlockFunc)(aclrtAllocatorBlock block);回调类型函数签名职责aclrtAllocatorAllocFuncvoid* (*)(aclrtAllocator, size_t)申请内存返回分配块句柄aclrtAllocatorFreeFuncvoid (*)(aclrtAllocator, aclrtAllocatorBlock)释放内存块aclrtAllocatorAllocAdviseFuncvoid* (*)(aclrtAllocator, size_t, aclrtAllocatorAddr)按建议地址申请内存可空样例中复用 Alloc 逻辑aclrtAllocatorGetAddrFromBlockFuncvoid* (*)(aclrtAllocatorBlock)从分配块句柄中解析出实际内存地址其中aclrtAllocator是对用户自定义分配器对象的不透明句柄aclrtAllocatorBlock是分配块句柄aclrtAllocatorAddr是地址提示参数。注意函数返回的是分配块Block而非裸内存地址实际地址需要再通过GetAddrFromBlock回调解析——这样设计的目的是让分配器可以在块上附加元数据。2.2 描述符管理 APIAllocatorDesc 描述符的生命周期与内容由以下接口管理接口作用aclrtAllocatorCreateDesc()创建分配器描述符返回aclrtAllocatorDesc失败返回 nullaclrtAllocatorDestroyDesc(desc)释放分配器描述符返回ACL_SUCCESS表示成功aclrtAllocatorSetObjToDesc(desc, allocator)把自定义分配器对象绑定到描述符aclrtAllocatorSetAllocFuncToDesc(desc, func)绑定 alloc 回调函数指针aclrtAllocatorSetFreeFuncToDesc(desc, func)绑定 free 回调函数指针aclrtAllocatorSetAllocAdviseFuncToDesc(desc, func)绑定 allocAdvise 回调函数指针aclrtAllocatorSetGetAddrFromBlockFuncToDesc(desc, func)绑定 getAddrFromBlock 回调函数指针2.3 注册 / 查询 / 注销 API接口作用aclrtAllocatorRegister(stream, desc)把分配器描述符注册到指定 StreamaclrtAllocatorGetByStream(stream, descOut, allocatorOut, allocFunc, freeFunc, allocAdviseFunc, getAddrFunc)按 Stream 查询当前生效的分配器配置与回调函数aclrtAllocatorUnregister(stream)注销该 Stream 上注册的分配器三、工作流程总览示例 main.cpp 展示了完整闭环总体分六步初始化环境aclInit初始化 ACLaclrtSetDevice指定 Device创建 Context 与 Stream创建描述符aclrtAllocatorCreateDesc创建 AllocatorDesc绑定对象与回调把自定义状态对象CustomAllocatorState和四个回调函数写入描述符注册aclrtAllocatorRegister(stream, desc)将描述符绑定到 Stream查询并执行aclrtAllocatorGetByStream读回分配器与回调随后依次执行申请256B、按建议申请128B、取地址、释放清理aclrtAllocatorUnregister注销、aclrtAllocatorDestroyDesc销毁描述符最后销毁 Stream/Context、aclrtResetDeviceForce复位 Device、aclFinalize去初始化。四、示例代码逐段解析4.1 自定义状态与块结构struct CustomBlock { void* addr; // 实际内存地址 size_t size; // 申请大小 }; struct CustomAllocatorState { uint32_t allocCount; // 累计申请次数 uint32_t freeCount; // 累计释放次数 };示例用CustomBlock在分配块上携带地址 大小元数据用CustomAllocatorState记录分配/释放次数用于验证回调确实被执行。4.2 四类回调的实现void* CustomAlloc(aclrtAllocator allocator, size_t size) { auto* state static_castCustomAllocatorState*(allocator); auto* block static_castCustomBlock*(std::malloc(sizeof(CustomBlock))); if (block nullptr) return nullptr; block-addr std::malloc(size); if (block-addr nullptr) { std::free(block); return nullptr; } block-size size; if (state ! nullptr) state-allocCount; INFO_LOG(CustomAlloc size%zu, size); return block; } void CustomFree(aclrtAllocator allocator, aclrtAllocatorBlock block) { auto* state static_castCustomAllocatorState*(allocator); auto* typedBlock static_castCustomBlock*(block); if (typedBlock ! nullptr) { std::free(typedBlock-addr); std::free(typedBlock); } if (state ! nullptr) state-freeCount; INFO_LOG(CustomFree invoked); } void* CustomAllocAdvise(aclrtAllocator allocator, size_t size, aclrtAllocatorAddr addr) { (void)addr; INFO_LOG(CustomAllocAdvise size%zu, size); return CustomAlloc(allocator, size); // 忽略建议地址复用申请逻辑 } void* CustomGetAddrFromBlock(aclrtAllocatorBlock block) { auto* typedBlock static_castCustomBlock*(block); return typedBlock nullptr ? nullptr : typedBlock-addr; }实现要点CustomAlloc分配两块内存一块放CustomBlock元数据一块放实际size字节的用户内存返回的分配块句柄即CustomBlock*CustomFree先释放块内实际地址再释放块元数据避免内存泄漏CustomAllocAdvise接收建议地址参数addr本例直接忽略并复用CustomAlloc实际工程中可根据addr提示做地址对齐或固定地址映射CustomGetAddrFromBlock是后续用户访问内存的唯一入口通过它拿到block-addr。4.3 描述符创建、绑定与注册aclrtAllocatorDesc desc aclrtAllocatorCreateDesc(); CHECK_NOT_NULL(desc, aclrtAllocatorCreateDesc); CHECK_ERROR(aclrtAllocatorSetObjToDesc(desc, reinterpret_castaclrtAllocator(state))); CHECK_ERROR(aclrtAllocatorSetAllocFuncToDesc(desc, CustomAlloc)); CHECK_ERROR(aclrtAllocatorSetFreeFuncToDesc(desc, CustomFree)); CHECK_ERROR(aclrtAllocatorSetAllocAdviseFuncToDesc(desc, CustomAllocAdvise)); CHECK_ERROR(aclrtAllocatorSetGetAddrFromBlockFuncToDesc(desc, CustomGetAddrFromBlock)); CHECK_ERROR(aclrtAllocatorRegister(stream, desc));这里把CustomAllocatorState state {0, 0}的地址作为aclrtAllocator对象句柄写入描述符因此后续所有回调收到的allocator参数都指向该状态结构——这是回调能拿到自定义上下文的关键设计。随后把四个函数指针逐一写入描述符并注册到 Stream。4.4 按 Stream 查询并执行回调aclrtAllocatorDesc descOut nullptr; aclrtAllocator allocatorOut nullptr; aclrtAllocatorAllocFunc allocFunc nullptr; aclrtAllocatorFreeFunc freeFunc nullptr; aclrtAllocatorAllocAdviseFunc allocAdviseFunc nullptr; aclrtAllocatorGetAddrFromBlockFunc getAddrFunc nullptr; CHECK_ERROR(aclrtAllocatorGetByStream( stream, descOut, allocatorOut, allocFunc, freeFunc, allocAdviseFunc, getAddrFunc));GetByStream一次性读回描述符句柄、分配器对象和全部回调函数指针之后即可直接调用aclrtAllocatorBlock block allocFunc(allocatorOut, 256); // 申请 256B void* addr getAddrFunc(block); // 解析实际地址 std::fill_n(static_castunsigned char*(addr), 256, 0xAB); // 写入数据 if (allocAdviseFunc ! nullptr) { // 按建议地址申请 aclrtAllocatorBlock advisedBlock allocAdviseFunc(allocatorOut, 128, nullptr); if (advisedBlock ! nullptr) freeFunc(allocatorOut, advisedBlock); } freeFunc(allocatorOut, block); // 释放 INFO_LOG(Allocator desc queried: descOut%p allocCount%u freeCount%u, descOut, state.allocCount, state.freeCount);注意示例对allocAdviseFunc做了空指针判断——它属于可选回调注册时可以省略。最终日志中的allocCount2256B 一次 128B 一次、freeCount2直接证明四类回调被真实调用且state对象跨回调共享。4.5 清理与去初始化CHECK_ERROR(aclrtAllocatorUnregister(stream)); // 注销该 Stream 的分配器 CHECK_ERROR(aclrtAllocatorDestroyDesc(desc)); // 销毁描述符 CHECK_ERROR(aclrtDestroyStream(stream)); CHECK_ERROR(aclrtDestroyContext(context)); CHECK_ERROR(aclrtResetDeviceForce(deviceId)); CHECK_ERROR(aclFinalize());清理解析先注销注册关系再销毁描述符对象随后按Stream → Context → Device → ACL的逆序完成资源回收。五、编译与运行5.1 环境准备环境安装详情与运行说明见 example 目录 README。运行步骤如下# ${install_root} 替换为 CANN 安装根目录默认安装在 /usr/local/Ascend 目录 source ${install_root}/cann/set_env.sh export ASCEND_INSTALL_PATH${install_root}/cann # 编译运行 bash run.sh5.2 构建脚本解析run.sh 内部先source $ASCEND_INSTALL_PATH/bin/setenv.bash引入 CANN 环境再在本地build目录执行 CMake 配置与编译cmake .. -DASCEND_CANN_PACKAGE_PATH${_ASCEND_INSTALL_PATH} make -j$(nproc) ./build/main对应的 CMakeLists.txt 关键配置头文件路径${ASCEND_CANN_PACKAGE_PATH}/include与示例公共目录../../..即 example 顶层提供utils.h中的CHECK_ERROR、INFO_LOG宏链接库${ASCEND_CANN_PACKAGE_PATH}/lib64/libascendcl.so样例在 include 中同时引入acl/acl.h与acl/acl_rt_allocator.h自定义分配器 API 即由libascendcl.so导出编译选项-O2 -stdc17 -D_GLIBCXX_USE_CXX11_ABI0 -Wall -Werror要求代码零警告。5.3 示例输出编译运行成功后终端输出如下[INFO] CustomAlloc size256 [INFO] CustomAllocAdvise size128 [INFO] CustomAlloc size128 [INFO] CustomFree invoked [INFO] CustomFree invoked [INFO] Allocator desc queried: descOut..., allocCount2, freeCount2日志顺序与代码执行路径一一对应先申请 256B再走 allocAdvise 申请 128B内部复用 CustomAlloc最后两次释放结尾allocCount2、freeCount2验证了状态统计的准确性。六、底层实现原理源码级自定义分配器的 ACL 侧实现位于 src/acl/aclrt_impl/allocator.cpp。从源码结构看其内部维护了一个以 Stream 为键的注册表g_AllocatorDesMapstd::mapaclrtStream, std::pairaclrtAllocatorDesc, AllocatorDesc并以g_AllocatorDescMutex互斥锁保护并发访问。关键实现逻辑描述符绑定如aclrtAllocatorSetObjToDescImpl、aclrtAllocatorSetAllocFuncToDescImpl对传入的allocatorDesc做空指针校验后直接写入AllocatorDesc结构体的obj、allocFunc、freeFunc、allocAdviseFunc、getAddrFromBlockFunc字段注册aclrtAllocatorRegisterImpl强制要求stream非空并校验描述符中obj、allocFunc、freeFunc、getAddrFromBlockFunc均已绑定注意allocAdviseFunc不在必填校验之列与示例代码的空指针判断相印证校验通过后拷贝一份AllocatorDesc连同原描述符句柄一起写入g_AllocatorDesMap[stream]查询aclrtAllocatorGetByStreamImpl在注册表中按 Stream 查找若未注册则返回ACL_ERROR_INVALID_PARAM并上报错误信息 The stream is not registered with any allocator命中后按调用方传入的非空出参逐项回填描述符、分配器对象与各回调函数指针注销aclrtAllocatorUnregisterImpl直接g_AllocatorDesMap.erase(stream)删除注册项。同时描述符销毁实现aclrtAllocatorDestroyDescImpl通过delete释放AllocatorDesc对象并在 ACL 统计模块中记账ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC说明该能力与 ACL 的维测/统计体系ACL_PROFILING_REG 埋点是打通的。从上述实现可见AllocatorDesc 本质是对象 回调函数指针的聚合结构而 Stream 是注册与查询的索引键这是理解整套机制的核心。七、注意事项与最佳实践Stream 非空aclrtAllocatorRegister、aclrtAllocatorGetByStream、aclrtAllocatorUnregister均要求 Stream 已创建且非空未注册 Stream 直接GetByStream会返回ACL_ERROR_INVALID_PARAM必填回调注册前必须保证obj、allocFunc、freeFunc、getAddrFromBlockFunc四项已绑定否则注册失败allocAdviseFunc可选但调用前应判空块句柄语义alloc 类回调返回的是分配块Block而非裸地址业务侧必须通过getAddrFromBlock解析后再读写内存生命周期匹配描述符创建后须与DestroyDesc成对使用注册须与Unregister成对使用避免泄漏或悬垂对象所有权写入描述符的分配器对象如示例中的state由用户管理生命周期其存活时间必须覆盖所有回调执行期线程安全底层注册表使用互斥锁保护多线程下并发注册/查询是安全的但同一 Stream 重复注册会覆盖旧配置按需先注销。八、相关资源样例源码example/3_memory_advanced/allocator/0_custom_allocator/main.cpp样例说明0_custom_allocator/README.md含英文版 README_en.md对外 API 头文件include/external/acl/acl_rt_allocator.h底层实现src/acl/aclrt_impl/allocator.cpp本章目录说明example/3_memory_advanced/README.md可继续阅读同章节的 host_register、managed_memory、memory_pool 等主题构建完整的高级内存管理体系。当前已知 issue暂无。【免费下载链接】runtime本项目提供CANN运行时组件和维测功能组件。项目地址: https://gitcode.com/cann/runtime创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考