torch_npu.npu.NPUPluggableAllocator

定义文件

torch_npu/npu/memory.py

函数原型

功能说明

从so文件加载的NPU内存分配器。

参数说明

• path_to_so_file：(str) so文件路径。
• alloc_fn_name：(str)内存申请函数名（与c/c++文件中函数名一致）。
• free_fn_name：(str)内存释放函数名（与c/c++文件中函数名一致）。

约束说明

alloc_fn_name内存申请函数名必须与c/c++文件中函数名一致。

free_fn_name内存释放函数名必须与c/c++文件中函数名一致。

支持的型号

• [object Object] Atlas 训练系列产品[object Object]
• [object Object] Atlas A2 训练系列产品[object Object]
• [object Object] Atlas A3 训练系列产品[object Object]
• [object Object] Atlas 推理系列产品[object Object]

调用示例

完整调用示例可参考undefined。

Python代码示例：

风险提示：自定义的so建议严格参照安全代码示例，内存申请，释放函数必须正确实现。错误的so文件可能出现内存申请失败，内存泄漏等问题。

建议处理方式：用户可在内存申请，释放等内存相关操作函数中，增加日志记录内存行为，方便后续定位问题。

c++安全代码示例：

日志记录示例：可使用Ascend Extension for PyTorch自带的debug级别日志打印。日志需注意存放至安全路径。