CANN/pypto:Loop原理详解

📅 2026/7/10 20:22:12
CANN/pypto:Loop原理详解
Loop原理及其描述【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto在编译阶段会将loop编译为控制流将loop body转换为计算流分别对应kernel_aicpu/kernel_aicore。kernel_aicpu负责控制流的执行创建kernel_aicore任务包括使用到的内存分配以及执行参数准备同时分析输入输出的依赖将kernel_aicore任务合并成一个更大的调度单元然后提交给调度单元进行调度。原型介绍def loop(start, end, step1, nameNone, idx_nameNone, unroll_list[1], submit_before_loopFalse): def loop_roll(start, end, step1, nameNone, idx_nameNone, unroll_list[1], submit_before_loopFalse):参数说明start,end,step分别表示循环的起始值、结束值和步长类型可以为SymbolicScalar或者int和Python中的range语法基本保持一致。name循环的名称默认值为loop_{id}对实际使用运行效果无影响仅用于调试和生成代码中注释信息。idx_name循环索引的名称默认值为loop_idx_{id}对于嵌套的Loop使用相同的idx会产生覆盖行为目前会在前端进行检查报错。unroll_list主要用于循环展开产生更大的loop body降低调度开销。对于unroll_list2loop大概会产生如下代码new_start start for k in unroll_list: left (stop - start) % k for idx in loop(new_start, stop - left, k): for i in range(k): body(idx) # 需要用户一次处理step1的步长 new_start stop - leftloop_unroll大概会产生如下代码new_start start for k in unroll_list: left (stop - start) % k for idx in loop(new_start, stop - left, k): body(idx, k) # 需要用户一次处理k的步长 new_start stop - left原则上如果可以一次处理多个i使用loop_unroll会更高效如果一次只能处理1个i则需要使用loop。submit_before_loop表示是否在循环开始前提交任务默认值为False。如果设置为True则循环前的任务会先提交到调度队列中等待后续任务完成后再开始执行。过多的设置submit_before_loop会增加调度开销建议仅在必要时设置为True。unroll_list对pypto.cond的影响通常一个循环中有一个pypto.cond会产生两个分支。当unroll次数为4次时会产生2⁴ 16个路径分支通常每个分支都需要单独编译因此会大量增加编译时间和编译出的代码量。为了支持关键算子FA的编译优化提供了两个特殊的函数pypto.is_loop_begin()和pypto.is_loop_end()用于优化条件分支。考虑到对外层loop进行loop_unroll不能提升loop body的大小当前仅支持最内侧循环进行unroll。隐式loop示例框架会在编译阶段前端隐式地在function开始的位置插入一个loop(1)。循环直到下一个循环开始前结束pypto.frontend.jit def foo(a, b, c): c[:] a b # 等价于 pypto.frontend.jit def foo(a, b, c): for i in pypto.loop(1): c[:] a b pypto.frontend.jit def foo(a, b, c): t a 1 for i in pypto.loop(1): c[:] t b # 等价于 pypto.frontend.jit def foo(a, b, c): for i in pypto.loop(1): t a 1 for i in pypto.loop(1): c[:] t b框架当前不会自动进行loop(1)合并因此在实际使用中建议用户手动合并loop(1)以提高效率。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考