昇腾社区首页
中文
注册
开发者
下载

并行优化

并行优化的核心逻辑是通过“外层并行”+“内层向量化”的循环拆分策略以匹配昇腾AI处理器硬件的多级并行粒度,同时平衡计算效率与数据局部性。
  • 循环拆分

    循环拆分的核心作用是通过引入新的循环层级,明确划分出适合并行的外层循环和适合向量化的内层循环。

    针对每个TilingCase,会将xgroup,ygroup,rgroup中存在的轴切分成ub_out和ub_in:

    例如:{z0, z1,z2},切分在z1上,就把z1切分成z1T和z1t,z1T就是ub_out,z1t就是ub_in。

  • 向量化

    向量化是利用硬件SIMD(single-instruction, multiple-data stream processing,单指令流多数据流)单元提升数据并行计算效率的关键技术,通过将单元素操作转化为向量操作,显著减少指令执行次数并提高硬件利用率。

    例如,对于如下循环:

    1
    2
    3
    for (int i = 0; i < 256; i++) {
        c[i] = a[i] + b[i];
    }
    

    非向量化执行需要256条加法指令,向量化后只需要一条加法指令。

    在每个组内选择一根轴作为ub切分轴后,会将ub_in及其内轴作为向量化轴。此时,由于分组是按照xyr来生成的,按照这个顺序生成的向量化轴,轴序与内存排布不一致会引入非连续搬运,因此需要根据输出的轴序进行重排列。

    例如,输出tensor轴序是(a, b, c, d),轴分组是(a, c),(b, d)(),按照这个顺序生成的向量化轴分组是(a_in, c, b_in, d),需要调整成(a_in, b_in, c, d)。

    Schedule阶段全图统一设置了相同的向量化轴,对于部分API来说,由于指令等限制,并不能将所有向量化轴都进行向量化处理。此时,需要CodeGen阶段对无法向量化处理的轴进行外抛for循环处理。

    例如向量化轴为[z1,z2,z3],相当于3层循环,如果指令支持3层循环,则可以写成vector[z1,z2,z3];但如果指令只能支持两层循环,则需要CodeGen给出如下的代码:

    1
    2
    3
    for (i in z1) {
    	vector[z2,z3]
    }
    
  • 循环合并、循环绑核

    循环合并与循环绑核二者常常配合优化,循环切分阶段在每个分组内都产生了ub_out的轴。

    1.将所有非reduce轴合并为一根轴,所有reduce轴合并为一根轴以减少循环嵌套层数。

    2.对合并后的循环进行拆分,得到外层和内层。

    3.将拆分后的外层循环绑定到多个block上以实现并行,内层通过循环进行消化。

    例如:{z0, z1T, z1t,z2},会先把z0,z1T合轴成z0z1T,再在z0z1T上进行多核切分,由于z0z1T可能会超过blockdim,因此在分完核后还要多一层循环,z0z1T会被进一步拆成z0z1TB和z0z1Tb,值的具体大小由Auto Tiling在Tiling阶段计算得出。