尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

pyasc 内存池分配接口解析:TBufPool.init_buffer 为 TQue / TBuf 分配片上内存的完整实践

pyasc 内存池分配接口解析:TBufPool.init_buffer 为 TQue / TBuf 分配片上内存的完整实践 pyasc 内存池分配接口解析TBufPool.init_buffer 为 TQue / TBuf 分配片上内存的完整实践【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyascTBufPool.init_buffer是 CANN pyascAscend C 的 Python 编程接口中用于为队列TQue与临时缓冲区TBuf分配片上内存Unified Buffer / L1 Buffer的核心接口。本文围绕该接口的两种重载形式、参数语义、约束条件与底层实现展开并结合仓库源码与单元测试帮助开发者在多 stage 流水计算、片上内存复用的场景下正确使用TBufPool管理内存。读完本文你将掌握init_buffer与init_buf_pool/reset的配合方式并能够编写可运行的 pyasc 内存池初始化代码。一、为什么需要 TBufPool从 TPipe 到内存池在 pyasc 编程模型中TPipe 负责统一管理 Device 端内存资源一个 Kernel 函数必须且只能初始化一个TPipe对象。常规场景下开发者通过TPipe.init_buffer直接为TQue流水队列和TBuf临时变量缓冲区分配内存。但当一个算子存在多个 stage 计算、且每段计算都需要独立的输入/输出缓冲时Unified Buffer / L1 Buffer 的物理内存往往不够用。此时TBufPool便派上用场——正如 python/asc/language/fwk/tpipe.py 中类注释所述TPipe 可以管理全局内存资源而 TBufPool 可以手动管理或复用 Unified Buffer / L1 Buffer 物理内存主要用于多个 stage 计算中 Unified Buffer / L1 Buffer 物理内存不足的场景。TBufPool的本质是一块从TPipe整体资源中划分出来的子资源池它允许开发者把整块物理内存先划给一个TBufPool再在其中细分出多个队列/临时缓冲的分配在多个 stage 之间复用同一片物理内存通过share_buf共享从而在内存受限时完成流水编排。二、接口签名两种重载两种分配目标TBufPool.init_buffer对应文档见 asc.language.fwk.TBufPool.init_buffer.md其 Python 侧声明为两个重载# 重载一为队列 TQue 分配 num 块、每块 len 字节的内存 TBufPool.init_buffer(que: TQue, num: int 0, len: int 0) → None # 重载二为临时缓冲 TBuf 分配 len 字节的内存 TBufPool.init_buffer(buf: TBuf, len: int 0) → None它调用的是 Ascend C 的TBufPool::InitBuffer接口对应的 C 函数原型为template class T __aicore__ inline bool InitBuffer(T que, uint8_t num, uint32_t len) template TPosition pos __aicore__ inline bool InitBuffer(TBufpos buf, uint32_t len)两个重载解决两类不同对象的内存初始化重载分配对象语义init_buffer(que, num, len)流水队列TQue为队列分配num块内存每块大小为len字节init_buffer(buf, len)临时缓冲TBuf为临时缓冲区分配一块大小为len字节的内存三、参数说明接口涉及的参数含TBufPool构造时传入的pos语义如下posBuffer 的逻辑位置可选值为VECIN、VECOUT、VECCALC、A1、B1、C1。这些枚举定义在 python/asc/language/core/enums.py 的TPosition中其中VECIN/VECOUT/VECCALC对应向量单元Unified Buffer上的逻辑位置A1/B1/C1对应 Cube 单元L1 Buffer 相关上的逻辑位置。que需要分配内存的TQue对象。num分配内存块的个数块数。lenTQue 重载每个内存块的大小单位为 Bytes非 32 Bytes 对齐时会自动向上补齐至 32 Bytes 对齐。buf需要分配内存的TBuf对象。lenTBuf 重载为TBuf分配的内存大小单位为 Bytes同样非 32 Bytes 对齐会自动向上补齐至 32 Bytes 对齐。需要特别注意的是TBufPool自身的逻辑位置pos与队列que/ 缓冲buf的逻辑位置必须匹配——从源码看TBufPool构造时通过builder.get_tbuf_pool_type(pos, buf_id_size)生成带位置信息的类型tpipe.py其物理内存与后续分配的TQue/TBuf必须一致。四、约束说明可分配 Buffer 数量上限声明TBufPool时可通过构造参数buf_id_size指定可分配 Buffer 的最大数量默认上限为 4最大为 16见 asc.language.fwk.TBufPool.init.md。非共享模式的资源分配在本TBufPool上再次申请子TBufPool时子池的buf_id_size不能超过原池剩余可用的 Buffer 数量共享模式share_buf的资源分配子池的buf_id_size不能超过原池设置的 Buffer 数量。物理内存一致性TQue或TBuf的物理内存需要和TBufPool一致即它们必须来自同一个资源池的分配。长度对齐规则len参数非 32 Bytes 对齐时自动向上取整对齐因此开发者无需手工对齐但建议按 32B 倍数规划便于地址计算与get_with_offset等偏移操作。五、源码级实现从 Python 重载到底层 IR 算子TBufPool.init_buffer的 Python 实现位于 python/asc/language/fwk/tpipe.py。它并非一个普通函数而是借助OverloadDispatcher根据实参类型分发到不同重载require_jit set_tpipe_docstring(pipe_nameTBufPool, api_nameinit_buffer) def init_buffer(self, *args, **kwargs) - None: dispatcher OverloadDispatcher(__name__) dispatcher.register(queTQue, numRuntimeInt, lenRuntimeInt) def _(que: TQue, num: RuntimeInt 0, len: RuntimeInt 0): global_builder.get_ir_builder().create_asc_TBufPoolInitQueueOp( self.to_ir(), que.to_ir(), _mat(num, KnownTypes.int_).to_ir(), _mat(len, KnownTypes.int_).to_ir()) dispatcher.register(bufTBuf, lenRuntimeInt) def _(buf: TBuf, len: RuntimeInt 0): global_builder.get_ir_builder().create_asc_TBufPoolInitBufferOp( self.to_ir(), buf.to_ir(), _mat(len, KnownTypes.int_).to_ir()) dispatcher(*args, **kwargs)从实现可以看出两点关键信息编译期约束init_buffer被require_jit装饰即只能在asc.jit编译的 kernel 函数体内调用无法在 Host 侧 Python 环境中直接执行底层 IR 映射TQue 重载生成asc_TBufPoolInitQueueOp算子对应 C 的InitBuffer(T que, uint8_t num, uint32_t len)TBuf 重载生成asc_TBufPoolInitBufferOp算子对应 C 的InitBuffer(TBufpos buf, uint32_t len)num以int对标uint8_t、len以int对标uint32_t类型下发。这与文档中给出的 C 原型一一对应说明 pyasc 接口与 Ascend C 是严格对齐的。六、完整调用示例多 stage 内存复用流水原文档给出了一个非常典型的多 stage 内存复用场景先用TPipe.init_buf_pool从整体资源中划出大池tbuf_pool0再通过TBufPool.init_buf_pool细分出tbuf_pool1普通子池与tbuf_pool2共享tbuf_pool1物理内存的子池随后两个 stage 交替使用tbuf_pool1/tbuf_pool2为各自的TQue分配内存并配合reset()完成切换asc.jit def init(src0_gm: asc.GlobalAddress, src1_gm: asc.GlobalAddress, dst_gm: asc.GlobalAddress): src0_global.set_global_buffer(src0_gm) src1_global.set_global_buffer(src1_gm) dst_global.set_global_buffer(dst_gm) # 从 TPipe 整体资源中划分出 131072 字节的大资源池 pipe.init_buf_pool(tbuf_pool0, 131072) # 在大池中为 src0 队列分配 1 块 65536 字节内存 tbuf_pool0.init_buffer(quesrc_que0, num1, len65536) # Total src0 # 细分出子池 tbuf_pool165536 字节 tbuf_pool0.init_buf_pool(tbuf_pool1, 65536) # 细分出子池 tbuf_pool265536 字节并与 tbuf_pool1 共享物理内存 tbuf_pool0.init_buf_pool(tbuf_pool2, 65536, tbuf_pool1) asc.jit def Process(): # stage 1使用 tbuf_pool1 为输入/输出队列分配内存 tbuf_pool1.init_buffer(quesrc_que1, num1, len32768) tbuf_pool1.init_buffer(quedst_que0, num1, len32768) copy_in() compute() copy_out() tbuf_pool1.reset() # stage 2切换至 tbuf_pool2复用 tbuf_pool1 的物理内存 tbuf_pool2.init_buffer(src_que2, num1, len32768) tbuf_pool2.init_buffer(dst_que1, num1, len32768) copy_in1() compute1() copy_out1() tbuf_pool2.reset() tbuf_pool0.reset() pipe.reset()示例中展示的完整生命周期是建池pipe.init_buf_pool(tbuf_pool0, len)从TPipe整体资源中划分整块子资源池细分tbuf_pool0.init_buf_pool(tbuf_pool1, len)继续划分更小的池init_buf_pool(tbuf_pool2, len, tbuf_pool1)则创建与tbuf_pool1共享起始地址及长度的池相关约束见 asc.language.fwk.TBufPool.init_buf_pool.md新划分的资源池与被复用资源池物理内存一致、输入长度需小于等于被复用池长度分配各 stage 内用init_buffer(que..., num..., len...)为队列分配内存切换stage 切换前调用tbuf_pool.reset()结束当前资源池正在处理的相关事件。如 asc.language.fwk.TBufPool.reset.md 所述调用后资源池及其分配的 Buffer 仍然存在只是 Buffer 内容可能被改写切回该池后可重新使用这些 Buffer无需再次分配收尾pipe.reset()恢复TPipe的初始化状态。七、TBuf 场景init_buffer 与 TBuf.get 的配合当使用临时缓冲时init_buffer(buf, len)分配的内存后续通过TBuf.get取为可参与计算的LocalTensor。参考 asc.language.fwk.TBuf.get.md 的用法# 为 TBuf 初始化分配内存分配内存长度为 1024 字节 pipe asc.Tpipe() calc_buf asc.TBuf(asc.TPosition.VECCALC) byte_len 1024 pipe.init_buffer(calc_buf, byte_len) # 从 calc_buf 获取 TensorTensor 为 pipe 分配的所有内存大小为 1024 字节 temp_tensor1 calc_buf.get(asc.int32) # 从 calc_buf 获取 TensorTensor 为 128 个 int32_t 类型元素的内存大小为 512 字节 temp_tensor1 calc_buf.get(asc.int32, 128)在TBufPool场景下把示例中的pipe.init_buffer(calc_buf, byte_len)替换为tbuf_pool.init_buffer(bufcalc_buf, lenbyte_len)即可让临时缓冲从资源池中分配。注意约束len的数值是 Tensor 中元素的个数len * sizeof(T)不能超过TBuf初始化时的内存长度。八、测试验证仓库中如何验证该接口仓库在 python/test/unit/language/fwk/test_tbuf_pool.py 中提供了覆盖init_buffer两种重载的单元测试def test_init_buffer(mock_launcher_run): asc.jit def kernel_init_buffer() - None: que asc.TQue(asc.TPosition.VECIN, 1) tmp_buf asc.TBuf(asc.TPosition.VECCALC) buf_pool0 asc.TBufPool(posasc.TPosition.VECIN, buf_id_size4) buf_pool0.init_buffer(queque, num1, len256) buf_pool0.init_buffer(buftmp_buf, len256) kernel_init_buffer[1]() assert mock_launcher_run.call_count 1同一测试文件还覆盖了TBufPool构造test_init、init_buf_pool的普通/共享两种形态test_init_buf_pool以及resettest_reset完整验证了资源池从建池、细分、分配到复用的全流程。这些测试同时印证了接口的调用方式asc.TBufPool(pos..., buf_id_size...)构造资源池、asc.TQue(asc.TPosition.VECIN, depth)构造队列、asc.TBuf(asc.TPosition.VECCALC)构造临时缓冲与文档示例完全一致。九、实践要点小结分清TPipe.init_buffer与TBufPool.init_buffer前者由TPipe统一分配后者在手动划分的资源池内分配TPipe.init_buf_pool与TBufPool.init_buf_pool分别用于从整体划池与池内再细分。内存池适用于多 stage 内存不足场景通过share_buf共享物理内存的子池可以让不同 stage 的队列/缓冲复用同一片 UB/L1 内存从而在有限片上内存下完成多段流水。切换资源池必须 resettbuf_pool.reset()结束当前池相关事件后再切换到下一个池切回旧池时无需重新init_buffer但旧 Buffer 内容可能已被改写需重新填充数据。注意 32B 对齐与长度约束len自动向上补齐 32BTBuf.get时len * sizeof(T)不得超过初始化长度子池共享模式下长度不得超过被复用池长度。buf_id_size 上限默认 4、最大 16细分池时不能超过原池剩余或设置的可分配 Buffer 数。十、延伸阅读TBufPool 类总览与相关接口init_buf_pool、init_buffer、reset三个接口的完整说明TBufPool 构造与 buf_id_size 约束TBufPool.init_buf_pool子池划分与共享内存约束TBufPool.reset资源池切换语义实现源码TBufPool/TPipe/TQue/TBuf类的完整 Python 实现单元测试覆盖建池、细分、分配与 reset 的验证用例【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表