| feat: [ACLGRAPH]support pin memory
Co-authored-by: luochao60<luochao60@huawei.com>
# message auto-generated for no-merge-commit merge:
!34382 merge pta_add_pin_memory_pool_20260408_master into master
feat: [ACLGRAPH]support pin memory
Created-by: luochao60
Commit-by: luochao60
Merged-by: ascend-robot
Description: <!--
PR描述模板更新日期：20260203
-->
# 【合入来源】
> (如有)请关联需求文档/issue链接
- [x] 需求
- [ ] 问题单
- [ ] issue/工单
- [ ] 重构优化
- [ ] 资料更新
# 【修改方案】
> 请描述修改内容的具体实现，涉及哪些组件之间进行交互，可以用1、2、3、...进行罗列
> 如果是需求或者重构类的PR，需要补充详细设计文档（说明上下游组件关系、时序图、类图、DFX能力等内容）
本次改动让 ACLGraph 捕获期间也能正常使用 pin_memory（pinned host memory），将 host allocator 接入 NPUGraph 的 mempool 私有池机制，避免 capture 期间申请到的 pinned 内存被错误回收。
1. torch_npu/csrc/core/npu/CachingHostAllocator.cpp：
- NPUCachingHostAllocatorImpl 实现基类 CachingHostAllocator 的 get_current_stream() 与 stream_is_capturing() 虚接口：返回当前 NPU stream，并通过 AclmdlRICaptureGetInfo 判断该 stream 是否处于 capture 状态，使基类能正确识别 NPU 的 capture 语义。
- NPUCachingHostAllocator 重写 begin_allocate_to_pool / end_allocate_to_pool / release_pool 接口，转发到内部 impl，与设备侧 NPUCachingAllocator 的私有池接口一一对应。
2. torch_npu/csrc/core/npu/NPUGraph.cpp：
- capture_begin 增加约束：当 pin_memory_expandable_segments=True 时禁止开启 ACLGraph capture。原因是 NPUExpandableHostAllocatorImpl 重写了 allocate/free/empty_cache/record_event，不与基类 CachingHostAllocator 的私有池机制对接，会导致 capture 期间分配的 pinned block 在回放时被错误回收，引发数据损坏；通过 TORCH_CHECK 显式报错，提示用户调整 PYTORCH_NPU_ALLOC_CONF。
- 在调用 NPUCachingAllocator::beginAllocateToPool 之外，同步调用 host allocator 的 begin_allocate_to_pool，将 host pinned 内存也注册到同一个 mempool_id_ 私有池中；filter 通过 NPUStream::stream(false) 取出底层 aclrtStream 后复用同一份判断逻辑，避免触发 PTA 任务队列的 flush。
- capture_end / reset 配套调用 host allocator 的 end_allocate_to_pool / release_pool，与设备侧池的生命周期保持一致。
3. torch_npu/npu/graphs.py：graph 上下文管理器进入时，在 torch.npu.empty_cache() 之后追加一次 torch_npu.npu.host_empty_cache()，避免上一次 capture 残留的 pinned cache 影响本次 capture 的内存可见性。
4. test/npu/test_aclgraph_pin_memory.py：新增 ACLGraph + pin_memory 集成测试用例，覆盖 capture/replay 期间 pin_memory 申请/释放的正确性，以及 pin_memory_expandable_segments=True 与 capture 互斥时的报错路径。
# 【资料变更】
> 请确认是否涉及资料变更。如涉及，需要在PR中体现，并简要说明修改内容。如不涉及，需填写"不涉及"
不涉及
# 【接口变更】
> 请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及，需要详细说明接口以及对应的变更内容，同时需要在资料中体现。如不涉及，需填写"不涉及"
不涉及对外公开接口的变更。新增的 get_current_stream / stream_is_capturing / begin_allocate_to_pool / end_allocate_to_pool / release_pool 都是 PTA 内部 host allocator 对 PyTorch 基类 CachingHostAllocator 的虚函数实现/重写，不会改变面向用户的 Python API 行为。
# 【功能验证】
> 说明测试场景，测试方法。如果本次测试方式与常规单元测试不同，请详细说明您的测试步骤
> 新增/变更内容是否已新增/适配UT测试用例看护，并补充测试自验证截图
新增 test/npu/test_aclgraph_pin_memory.py 用例，覆盖以下场景：
- ACLGraph capture 期间申请 pin_memory tensor，capture 结束后 replay，结果与 eager 一致。
- 多次 capture/replay，pinned 内存释放后能被同一私有池复用，无内存泄漏。
- pin_memory_expandable_segments=True 时调用 capture_begin 应触发 TORCH_CHECK 报错，提示用户调整 PYTORCH_NPU_ALLOC_CONF。
- 通过 python test/npu/test_aclgraph_pin_memory.py 在本地 NPU 环境运行通过。
# 【CheckList】
> PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]
- [x] 代码注释完备，正确记录错误日志
- [x] 代码实现进行了返回值、空指针等校验
- [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
- [x] PR持续集成流水线（CI）执行通过，代码检查无异常
See merge request: Ascend/pytorch!34382 | 27 天前 |
