Loop原理及其描述#

在编译阶段会将loop编译为控制流，将loop body转换为计算流，分别对应kernel_aicpu / kernel_aicore。kernel_aicpu负责控制流的执行，创建kernel_aicore任务，包括使用到的内存分配以及执行参数准备，同时分析输入输出的依赖将kernel_aicore任务合并成一个更大的调度单元，然后提交给调度单元进行调度。

原型介绍#

def loop(start, end, step=1, name=None, idx_name=None,
         unroll_list=[1], submit_before_loop=False):

def loop_roll(start, end, step=1, name=None, idx_name=None,
              unroll_list=[1], submit_before_loop=False):

参数说明#

start, end, step：分别表示循环的起始值、结束值和步长，类型可以为SymbolicScalar或者int，和Python中的range语法基本保持一致。
name：循环的名称，默认值为loop_{id}，对实际使用运行效果无影响，仅用于调试和生成代码中注释信息。
idx_name：循环索引的名称，默认值为loop_idx_{id}，对于嵌套的Loop使用相同的idx会产生覆盖行为，目前会在前端进行检查报错。

unroll_list：主要用于循环展开，产生更大的loop body，降低调度开销。

对于unroll_list=2，loop大概会产生如下代码：

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)  # 需要用户一次处理step=1的步长
    new_start = stop - left

loop_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)，以提高效率。