PTA自定义内存池

### Add a description # torch_npu 自定义内存池 (NPUPluggableAllocator) 机制分析 ## 1. 概述 torch_npu 提供了 `NPUPluggableAllocator` 机制，允许用户使用自定义的内存分配器替换默认的 `NPUCachingAllocator`。本文档详细分析该机制的实现原理、调用路径以及存在的问题。 ## 2. 核心架构 ### 2.1 分配器层次结构 ``` NPUAllocator (抽象接口) ├── NpuCachingAllocator (默认实现) │ └── DeviceCachingAllocator │ ├── BlockPool │ └── Block └── NPUPluggableAllocator (自定义分配器包装) └── 用户提供的函数指针 ├── alloc_fn_ ├── free_fn_ ├── record_stream_fn_ (可选) └── erase_stream_fn_ (可选) ``` ### 2.2 全局分配器切换机制 ```cpp // NPUCachingAllocator.h namespace c10_npu::NPUCachingAllocator { // 全局分配器指针 - 所有内存操作都通过它分发 C10_NPU_API extern std::atomic<NPUAllocator*> allocator; inline NPUAllocator* get() { return allocator.load(); } } // NPUCachingAllocator.cpp - 默认初始化 struct BackendStaticInitializer { BackendStaticInitializer() { allocator.store(&caching_allocator); // 默认使用 NpuCachingAllocator } }; // NPUPluggableAllocator.cpp - 切换到自定义分配器 void changeCurrentAllocator( const std::shared_ptr<NPUAllocator>& allocator) { TORCH_CHECK(!c10_npu::NPUCachingAllocator::allocator.load()->initialized(), "Can't swap an already initialized allocator"); c10_npu::NPUCachingAllocator::allocator.store(allocator.get()); current_custom_allocator = allocator; } ``` ## 3. NPUPluggableAllocator 实现细节 ### 3.1 核心数据结构 ```cpp // NPUPluggableAllocator.h struct _AllocationMetadata { size_t size; int device_idx; aclrtStream stream; // 注意：只记录分配时的stream }; struct NPUPluggableAllocator : public NPUAllocator { protected: std::function<void*(size_t, int, aclrtStream)> alloc_fn_; std::function<void(void*, size_t, int, aclrtStream)> free_fn_; std::function<void(void*, c10_npu::NPUStream)> record_stream_fn_; // 可选 std::function<void(void*, c10_npu::NPUStream)> erase_stream_fn_; // 可选 std::unordered_map<void*, _AllocationMetadata> allocation_metadata_; }; ``` ### 3.2 关键方法实现 #### recordStream 实现 ```…