SYCL中使用atomic_ref实现Jacobi迭代单kernel同步方案咨询

问题背景 #

假设我们有两个长度为n的向量V和W，在SYCL中启动一个kernel对V的每个元素执行3次for循环迭代，循环规则如下：

• 首先基于当前迭代中V的4个随机取值计算对应索引的W[idx]，即W[idx] = sum (V[a] + V[b] + V[c]+ V[d])，其中a、b、c、d非连续，是每个idx对应的固定取值。
• 基于W[idx]更新V[idx]。但V[idx]的更新操作必须要等步骤1中所有需要用到V[idx]计算W的操作全部完成后才能执行。

当kernel内执行3次循环迭代时，可能出现如下场景：线程1处于迭代1，需要使用迭代1的V[2]计算迭代1的W[idx=18]；而线程2已经运行到迭代2，完成了W[2]的计算并要更新迭代2的V[2]。如果线程2运行速度快于线程1，会提前更新V[2]的值，此时如何在SYCL中保证线程1读取到的是迭代1的V[2]值？考虑到线程2只会在V[2]被线程1使用后才会写入，使用atomic_ref能否解决该同步问题？另外迭代1的V[2]还会被其他线程用于计算迭代1的其他W值，如何保证只有迭代1中所有需要用到V[2]的操作全部完成后，才会更新迭代2的V[2]值？

现有实现中每次迭代需要启动2个kernel来设置两次计算之间的同步点，且第二次计算结束后也需要等待。需要找到一种方案，仅启动1个kernel，在kernel内部完成所有迭代并实现所需的同步。

现有实现代码 #

void jacobi_relaxation(cl::sycl::queue& q, ProblemVar& obj, int current_level) {
  for (int iterations = 1; iterations <= mu1; iterations++) {
    // TODO   =>        v(k+1) = [(1 - omega) x I + omega x D^-1 x(-L-U)] x v(k) + omega x
    // D^-1
    // x
    // f
    //
    // step 1 =>        v* = (-L-U) x v
    // step 2 =>        v* = D^-1 x (v* + f)
    // step 3 =>        v = (1-omega) x v + omega x v*

    q.submit([&](cl::sycl::handler& h) {
      // Accessor for current_level matrix CSR values
      auto row = obj.A_sp_dict[current_level].row.get_access<cl::sycl::access::mode::read>(h);
      auto col = obj.A_sp_dict[current_level].col.get_access<cl::sycl::access::mode::read>(h);
      auto val = obj.A_sp_dict[current_level].values.get_access<cl::sycl::access::mode::read>(h);
      auto diag_indices
          = obj.A_sp_dict[current_level].diag_index.get_access<cl::sycl::access::mode::read>(h);

      auto vec = obj.vecs_dict[current_level].get_access<cl::sycl::access::mode::read>(h);
      auto f = obj.b_dict[current_level].get_access<cl::sycl::access::mode::read>(h);
      cl::sycl::accessor<double, 1, cl::sycl::access::mode::write> vec_star{
          obj.temp_dict[current_level], h, cl::sycl::noinit};

      // Require 2 kernels as we perform Jacobi Relaxations
      h.parallel_for(
          cl::sycl::range<1>{obj.num_dofs_per_level[current_level]}, [=](cl::sycl::id<1> idx) {
            // double diag_multiplier = 0.0;
            vec_star[idx[0]] = 0.0;
            for (std::int32_t i = row[idx[0]]; i < row[idx[0] + 1]; i++) {

              vec_star[idx[0]] += -1.0 * val[i] * vec[col[i]];

            }
            
            vec_star[idx[0]] = (1.0 / val[diag_indices[idx[0]]]) * (vec_star[idx[0]] + f[idx[0]])
                               + vec[idx[0]]; // step 2
          });
    });
    q.wait();
    q.submit([&](cl::sycl::handler& h) {
      // Accessor for current_level vector
      auto vec = obj.vecs_dict[current_level].get_access<cl::sycl::access::mode::read_write>(h);
      auto vec_star
          = obj.temp_dict[current_level].get_access<cl::sycl::access::mode::read_write>(h);

      h.parallel_for(cl::sycl::range<1>{obj.num_dofs_per_level[current_level]},
                     [=](cl::sycl::id<1> idx) {
                       vec[idx[0]] = (1.0 - omega) * vec[idx[0]] + omega * vec_star[idx[0]]; // step
                                                                                             // 3
                       vec_star[idx[0]] = 0.0;
                     });
    });
    q.wait();
  }
}

解决方案 #

核心逻辑说明 #

该场景属于雅可比迭代的典型全局数据依赖问题：迭代k的所有读操作必须全部完成后才能执行迭代k到k+1的写更新，单独使用atomic_ref无法解决该问题：

• atomic_ref仅能保证单个内存位置的读写原子性，无法批量等待所有工作项完成迭代k的读计算阶段，不满足「所有用到V[2]的迭代1操作全部完成再更新」的要求
• 需要结合SYCL 2020的nd_range执行模型 + 全局工作组屏障 + 双缓冲存储机制，才能在单kernel内实现所需同步

具体实现方案 #

1. 双缓冲机制：保留vec和vec_star两个存储作为双缓冲版本，迭代k读版本A写版本B，迭代k+1读版本B写版本A，完全避免读写同一块内存出现的数据竞争
2. nd_range执行模型：将原来的parallel_for的range执行模式改为nd_range模式，显式指定工作组大小，方便调用全局同步接口
3. 全局屏障同步：每次迭代完成vec_star的计算后，调用cl::sycl::group_barrier指定内存作用域为cl::sycl::memory_scope::device，确保所有工作项都读完旧版本vec的值，再统一执行vec的更新操作，更新完成后再执行一次全局屏障，确认所有写入完成再进入下一轮迭代

兼容方案 #

如果你的设备不支持SYCL 2020的全局工作组屏障，可以自行实现全局栅栏：每个工作组完成计算后对全局原子计数加1，当计数值等于总工作组数量时代表所有计算完成，所有工作项再进入更新阶段即可。注意不要用单独的atomic_ref管控每个V元素的更新，会带来极高的性能开销，执行效率远低于全局屏障。

内容的提问来源于stack exchange，提问作者NikhilT