如何创建一个接受 sycl 内核的可定制函子？答案

【问题标题】：How can i create a customizable functor accepting sycl kernel?如何创建一个接受 sycl 内核的可定制函子？
【发布时间】：2022-11-03 03:44:16
【问题描述】：

在 sycl 中，我们创建一个像这样的内核：

queue.submit(
[&d_cells, &d_count_occupied](sycl::handler& cgh)
{
     auto cells_accessor = d_cells.get_access<sycl::access_mode::read>(cgh);
     auto count_accessor =
         d_count_occupied.get_access<sycl::access_mode::write>(cgh);
     cgh.parallel_for(
         d_cells.range(),
         [cells_accessor,
          count_accessor](sycl::id<3> id, sycl::kernel_handler kh)
         {
             auto cell = cells_accessor.at(kh, id);
             if (cell.is_occupied())
             {
                 sycl::atomic_ref<
                     unsigned,
                     sycl::memory_order::relaxed,
                     sycl::memory_scope::device>
                     count{count_accessor[0]};
                 count++;
             }
         }
     );
}
)

该内核采用 2 个缓冲区，其中 1 个保存单元信息，另一个用于计算“占用”单元的数量。现在想象一下，我将d_cells 缓冲区包装到一个知道或不知道占用单元格的类中。我们可以想象提供一个智能函数，它需要用户提供的 lambda 来对单元格进行操作：

class Cell {
   bool is_occupied() const;
   int get_position() const;

   // implementation details.

};

class Grid {

// Apply some user function to all of the occupied cells.
template <typename TFunctor, typename... TArgs>
sycl::event apply_all_occupied(sycl::queue q, TFunctor&& function, TArgs... args);

private: 
sycl::buffer<Cell> d_cells;

};

预期的调用模式将是这样的：

sycl::buffer<unsigned> d_count_occupied{
    count_occupied.data(), count_occupied.size()};
auto function = [](auto grid_cell, sycl::kernel_handler, auto count_accessor)
{
    sycl::atomic_ref<
        unsigned,
        sycl::memory_order::relaxed,
        sycl::memory_scope::device>
        count{count_accessor[0]};
    count++;
};
grid.apply_all_occupied(queue, function, d_count_occupied).wait_and_throw();

这将非常酷，它大大简化和抽象了“网格”的实现，这很好。但是这里我们有一个问题。用户提供的函子的实现必须能够在设备上运行。因此，提供的缓冲区需要在传递给用户提供的函数之前转换为“访问器”。我们也许可以通过一些元编程来解决它，例如：


template <typename TFunctor, typename... TArgs>
sycl::event apply_all_occupied(sycl::queue q, TFunctor&& function, TArgs... args) {

queue.submit(
[this, function, &args...](sycl::handler& cgh)
{
     auto cells_accessor = d_cells_.get_access<sycl::access_mode::write>(cgh);

     // Somehow get the access to all of the arguments here?
     std::tuple accessors = {args.get_access<sycl::access_mode::read>(cgh), ...};

     cgh.parallel_for(
         d_cells.range(),
         [cells_accessor,
          accessors, function](sycl::id<3> id, sycl::kernel_handler kh)
         {
             auto cell = cells_accessor.at(kh, id);
             function(kh, cell, accessors);
         }
     );
}

但这有严重的问题：

用户需要他们的 lambda 来接收一些带有访问器的模棱两可的元组类型。
无法为每个get_access 调用自定义访问模式。
有没有一种明智的方法来实现这种行为？

【问题讨论】：

标签： c++ c++17 intel sycl dpc++

【解决方案1】：

是的，有办法。您自定义访问模式的第二个要求意味着您希望将每个缓冲区的转换操作传递到 apply_all_occupied() 而不是缓冲区本身。 IE。你会收到一个参数包BufferAccessFuncsT &&... get_access_funcs，其中每个元素都是可调用的。例如：

int main()
{
  sycl::buffer d_count_occupied;
  sycl::other_buffer other_buf;
  sycl::queue q;
  Grid grid;

  auto function = [](auto grid_cell,
                     sycl::kernel_handler & kh,
                     sycl::buffer::accessor & count_accessor,
                     sycl::other_buffer::accessor & buf2) {
    std::cout << "Called" << std::endl;
    // Do stuff, e.g.:
    // sycl::atomic_ref<...> count{count_accessor[0]};
    // count++;
  };

  grid.apply_all_occupied(
      q,
      function,
      [&d_count_occupied](sycl::handler & cgh) { return d_count_occupied.get_access<sycl::access_mode::write>(cgh); },
      [&other_buf](sycl::handler & cgh) { return other_buf.get_access<sycl::access_mode::read>(cgh); });
}

在此示例中，我将两个 lambda 传递给 apply_all_occupied()，它们返回每个缓冲区的访问器。当然，它也适用于仅一个或零个或多个访问器。 main() 中的 function 期望在相同的顺序因为 lambda 被传递到 apply_all_occupied()。

关于你的第一个要求，用户定义的函数不应该直接接收元组，而是直接接收参数，你基本上想要一个“本地参数包变量”之类的东西

// Invalid, does not compile
auto &&... accessors = (std::forward<BufferAccessFuncsT>(get_access_funcs)(cgh))...;

然后您可以转发到您的实际功能。据我所知，这样的事情是不存在的。但是，您可以进行转换并将结果直接传递给另一个辅助函数。像这样：

struct Grid
{
  template <class FuncT, class... BufferAccessFuncsT>
  void apply_all_occupied(sycl::queue & q, FuncT && func, BufferAccessFuncsT &&... get_access_funcs)
  {
    q.submit([&](sycl::handler & cgh) {
      auto cells_accessor = 0; // Or whatever

      // Helper function that receives the transformed arguments in the parameter pack get_access_funcs.
      auto call_parallel_for_with_accessors = [&](auto &&... accessors) {
        cgh.parallel_for([&](sycl::kernel_handler & kh) {
          int grid_cell = cells_accessor; // Or whatever
          func(grid_cell, kh, accessors...);
        });
      };

      call_parallel_for_with_accessors((std::forward<BufferAccessFuncsT>(get_access_funcs)(cgh))...);
    });
  }
};

Grid::apply_all_occupied() 中的 call_parallel_for_with_accessors 是接收访问器的辅助函数。

请注意，我从您的原始代码中删除了一些非必要的东西，以获得一个最小的示例。

完整示例（直播于godbolt）：

#include <iostream>
#include <utility>

namespace sycl
{
struct kernel_handler
{
};

struct handler
{
  kernel_handler kh;

  template <class FuncT, class... ArgsT>
  void parallel_for(FuncT && func, ArgsT &&... args)
  {
    func(kh, std::forward<ArgsT>(args)...);
  }
};


enum class access_mode
{
    read,
    write
};

struct buffer
{
  struct accessor
  {
  };

  template <access_mode mode>
  accessor get_access(handler &)
  {
    return accessor{};
  }
};

// Just to have another buffer type.
struct other_buffer
{
  struct accessor
  {
  };

  template <access_mode mode>
  accessor get_access(handler &)
  {
    return accessor{};
  }
};


struct queue
{
  handler cgh;

  template <class FuncT>
  void submit(FuncT func)
  {
    func(cgh);
  }
};
} // namespace sycl


struct Grid
{
  template <class FuncT, class... BufferAccessFuncsT>
  void apply_all_occupied(sycl::queue & q, FuncT && func, BufferAccessFuncsT &&... get_access_funcs)
  {
    q.submit([&](sycl::handler & cgh) {
      auto cells_accessor = 0; // Or whatever

      // Helper function that receives the transformed arguments in the parameter pack get_access_funcs.
      auto call_parallel_for_with_accessors = [&](auto &&... accessors) {
        cgh.parallel_for([&](sycl::kernel_handler & kh) {
          int grid_cell = cells_accessor; // Or whatever
          func(grid_cell, kh, accessors...);
        });
      };

      call_parallel_for_with_accessors((std::forward<BufferAccessFuncsT>(get_access_funcs)(cgh))...);
    });
  }
};


int main()
{
  sycl::buffer d_count_occupied;
  sycl::other_buffer other_buf;
  sycl::queue q;
  Grid grid;

  auto function = [](auto grid_cell,
                     sycl::kernel_handler & kh,
                     sycl::buffer::accessor & count_accessor,
                     sycl::other_buffer::accessor & buf2) {
    std::cout << "Called" << std::endl;
    // Do stuff, e.g.:
    // sycl::atomic_ref<...> count{count_accessor[0]};
    // count++;
  };

  grid.apply_all_occupied(
      q,
      function,
      [&d_count_occupied](sycl::handler & cgh) { return d_count_occupied.get_access<sycl::access_mode::write>(cgh); },
      [&other_buf](sycl::handler & cgh) { return other_buf.get_access<sycl::access_mode::read>(cgh); });
}

编辑：如果不需要灵活地将完整的 lambda 传递给 apply_all_occupied()，但每个缓冲区只应指定 access_mode，则可以引入额外的辅助函数

template <sycl::access_mode mode, class BufferT>
auto AccessAs(BufferT & buffer)
{
    return [&] (sycl::handler & cgh) { 
        return buffer.template get_access<mode>(cgh); 
    };
};

并像这样调用apply_all_occupied()：

  grid.apply_all_occupied(
      q,
      function,
      AccessAs<sycl::access_mode::write>(d_count_occupied),
      AccessAs<sycl::access_mode::read>(other_buf));

godbolt 上的完整示例。

您可以通过定义进一步缩写

template <class BufferT>
auto AsWritableBuffer(BufferT & buffer)
{
    return AccessAs<sycl::access_mode::write>(buffer);
};

template <class BufferT>
auto AsReadableBuffer(BufferT & buffer)
{
    return AccessAs<sycl::access_mode::read>(buffer);
};

并像这样使用它

  grid.apply_all_occupied(
      q,
      function,
      AsWritableBuffer(d_count_occupied),
      AsReadableBuffer(other_buf));

godbolt 上的完整示例。

【讨论】：

这是一个很好的建议。不过，我希望不要为每个功能走上单独功能的路线。我会考虑的。
@FantasticMrFox 我使用示例帮助函数AccessAs 编辑了我的答案，它缩写了传递给apply_all_occupied() 的参数。但也许我没有正确理解你：如果你说你不想将单个函数传递给apply_all_occupied()，你的意思是你只是想直接传递对缓冲区的引用吗？但是，您如何想象调用者可以“自定义访问模式”的要求？为此，您一定需要以某种方式传递有关access_mode 的信息吗？你要找的界面是什么？