webgpu_radix_sort.hpp

Go to the documentation of this file.

#pragma once
 
#include "utils_macros.hpp"
#include <string>
#include <webgpu/webgpu.hpp>
 
namespace netdem {
namespace webgpu {
 
class RadixSort {
public:
  RadixSort(wgpu::Device &device);
 
  void Init(wgpu::Device &device);
 
  void Solve(wgpu::Buffer data_keys, wgpu::Buffer data_vals, size_t array_size);
 
private:
  wgpu::Device &device;
 
  VecXT<wgpu::BindGroupLayout> bg_layouts;
 
  wgpu::ComputePipeline cmp_pl_radix_count;
  wgpu::ComputePipeline cmp_pl_prefix_sum_sub;
  wgpu::ComputePipeline cmp_pl_add_block_sums;
  wgpu::ComputePipeline cmp_pl_radix_reorder;
 
  void Init();
  void InitBindGroupLayout();
 
  wgpu::BindGroup GetBindGroup0(wgpu::Buffer data_keys, wgpu::Buffer data_vals,
                                wgpu::Buffer data_keys_bak,
                                wgpu::Buffer data_vals_bak, wgpu::Buffer bit_id,
                                size_t array_size);
 
  wgpu::BindGroup GetBindGroup1(
      wgpu::Buffer buffer_data_src_0, wgpu::Buffer buffer_presum_src_0,
      wgpu::Buffer buffer_dst_0, wgpu::Buffer buffer_data_src_1,
      wgpu::Buffer buffer_presum_src_1, wgpu::Buffer buffer_dst_1,
      size_t offset_data_src, size_t offset_presum_src, size_t offset_dst,
      size_t array_size_src, size_t array_size_dst);
 
  void InitPipeline();
 
  void RadixCount(wgpu::Buffer buffer_keys, wgpu::Buffer buffer_vals,
                  wgpu::Buffer buffer_keys_bak, wgpu::Buffer buffer_vals_bak,
                  wgpu::Buffer buffer_bit_id, wgpu::Buffer buffer_data_src_0,
                  wgpu::Buffer buffer_presum_src_0, wgpu::Buffer buffer_dst_0,
                  wgpu::Buffer buffer_data_src_1,
                  wgpu::Buffer buffer_presum_src_1, wgpu::Buffer buffer_dst_1,
                  size_t array_size_src, size_t array_size_dst,
                  wgpu::ComputePassEncoder cmp_pass);
 
  void PrefixSumSub(wgpu::Buffer buffer_data_src_0,
                    wgpu::Buffer buffer_presum_src_0, wgpu::Buffer buffer_dst_0,
                    wgpu::Buffer buffer_data_src_1,
                    wgpu::Buffer buffer_presum_src_1, wgpu::Buffer buffer_dst_1,
                    size_t offset_data_src, size_t offset_presum_src,
                    size_t offset_dst, size_t array_size_src,
                    size_t array_size_dst, wgpu::ComputePassEncoder cmp_pass);
 
  void AddBlockSums(wgpu::Buffer buffer_data_src_0,
                    wgpu::Buffer buffer_presum_src_0, wgpu::Buffer buffer_dst_0,
                    wgpu::Buffer buffer_data_src_1,
                    wgpu::Buffer buffer_presum_src_1, wgpu::Buffer buffer_dst_1,
                    size_t offset_data_src, size_t offset_presum_src,
                    size_t offset_dst, size_t array_size_src,
                    size_t array_size_dst, wgpu::ComputePassEncoder cmp_pass);
 
  void RadixReorder(wgpu::Buffer buffer_data_src_0,
                    wgpu::Buffer buffer_presum_src_0, wgpu::Buffer buffer_dst_0,
                    wgpu::Buffer buffer_data_src_1,
                    wgpu::Buffer buffer_presum_src_1, wgpu::Buffer buffer_dst_1,
                    size_t array_size_src, size_t array_size_dst,
                    wgpu::ComputePassEncoder cmp_pass);
 
  int GetMaxBitID(wgpu::Buffer data_vals, size_t array_size);
 
  const size_t workgroup_size = 256;
  const size_t ele_size = sizeof(u_int32_t);
  const size_t align_size = 256 / ele_size;
};
 
} // namespace webgpu
} // namespace netdem