[PHI] Two Stage Scatter/Gather Kernel for Fully Synced Results

[Enigmatisms]

PR Category

Operator Mechanism

PR Types

Improvements

Description

• [PHI] Major fix for gather/scatter related CUDA kernels #74922 修复了 scatter/gather 的主要逻辑问题：offset 计算存在问题。本 PR 在此基础上进一步对正确性以及 PyTorch 一致性进行调优。修复了：paddle.put_along_axis 系列的 API 在数据量稍大时，include_self=False模式下，部分计算结果就不正确 #72803 中提到的问题。

本 PR 的三个主要改动：

• scatter，reduce!=assign 模式 forward下转为两（mean是三）阶段计算：（1）如果include self=False，对应位置赋值对应操作的的 init value（比如max是对应的 numeric_limit，mean是0）。（2）reduce op 计算，无额外原子操作使用。（3）如果是mean，则进行 cast div（后续与PyTorch完全对齐，修改int类型舍入规则将在对应 kernel 内增加）。这样的修改使得 reduce!=assign 的情况进一步削减了内存开销：
  • include self 情况下，除mean外均不需要额外分配一个与self tensor等大的tensor，内存开销最大可降为原来的50%
  • mean 情况下仅分配单个（而非两个）与 self tensor 等大的tensor作为 atomic counter，内存开销最大可降为原来的66.7%
• 新增了一个 ScatterScalarKernel，用于在输入的 values 仅为单个 scalar 时，不再通过python端的 broadcast_to 产生对应的view。此 kernel 除了在 reduce!=assign，include_self=False 的情况有赋初值的作用外，还可以作为 scalar 的一个faster pass。faster pass相关的进一步优化工作拟转社区开发。
• 此前的 grad 计算中，大部分kernel都是一阶段的：

  ComputeOffset<true>(smem_shape_strides,
                      src_strides,
                      grad_strides,
                      &replace_index_value,
                      &replace_index,
                      tid,
                      ndim,
                      dim,
                      index);

  if (value_data[replace_index_value] == out_data[replace_index])
    phi::CudaAtomicAdd(aux_buffer + replace_index, 1);
  __syncthreads();
  if (out_data[replace_index] != x_data[replace_index]) {
    grad_data[replace_index] = 0;
  } else {
    grad_data[replace_index] = self_data[replace_index] /
                               static_cast<tensor_t>(aux_buffer[replace_index]);
  }

在 tensor 较大的情况下，blocks 较多时，对 aux_buffer 的操作无法仅用 __syncthreads() 来同步（intra-grid同步无法通过intra-block sync达成）。目前采用的方法与 #74854 类似，拆为两个 kernel。这么做有如下考虑：

• 本 PR 过程中尝试过实现 global memory 上的 semaphore，多个 block 等待 semaphore done，或考虑使用 cooperative_group 的 grid.sync()，但由于原始代码写得实在难以令人接受：没有显式控制 block 数量，完全依赖调度，导致两种方法都没办法简单实用（前者hang，后者报错）。正确的修改思路是在 kernel 内开 for 循环（spatial inner loop），外部只维护少量的（所有 SM 能放下的）blocks。这样就可以将两阶段 kernels 融合成单阶段 kernel。考虑到这个工作包含的：
  • 所有 kernel 内部改 for
  • 所有 block 分配机制需要添加对当前 dev 上的可容纳 block 数量预估
  • 重新添加 semaphore & cg 混合机制（部分低CUDA版本的机器无法用 cg）

具有一定的重复性，并且工作量也不小，拟后续转社区任务。

完备的前、反向测试脚本见这个 gist：Enigmatisms/scatter_gather_forward_backward.py

---

• CPU 改动：待推进
• 更完备的单测：待提交

Pcard-89620

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[Enigmatisms]

/re-run all-failed