compaction.cuh

#pragma once

#include <cub/cub.cuh>
#include <cuda/atomic>
#include <cuda/std/limits>

/*
 * @author Kevin Kristensen
 */

/// T should be an unsigned integer type
template <typename T>
constexpr __host__ __device__ __forceinline__ bool evaluate_top_digit(T value)
{
  return (value >> (cuda::std::numeric_limits<T>::digits - 1));
}

template <typename T>
constexpr __host__ __device__ __forceinline__ T set_top_digit(T value)
{
  return (value | (static_cast<T>(1) << (cuda::std::numeric_limits<T>::digits - 1)));
}

template <typename T>
constexpr __host__ __device__ __forceinline__ T remove_top_digit(T value)
{
  return value & ~(static_cast<T>(1) << (cuda::std::numeric_limits<T>::digits - 1));
  ;
}

/// Compare op for uint64_t indices with MSB potentially set
struct IndexCompareOpT
{
  __host__ __device__ __forceinline__ bool operator()(const uint64_t& l, const uint64_t& r) const
  {
    return remove_top_digit(l) < remove_top_digit(r);
  }
};

/// Initializes state for decoupled look-back
template <typename ScanTileStateT>
__global__ void scan_tile_state_init_kernel(ScanTileStateT scan_tile_state, int32_t num_tiles)
{
  scan_tile_state.InitializeStatus(num_tiles);
}

/// Determines the partitions of the allocation table along which to merge with the memory pool
template <int32_t TileSize, typename AllocationIndexIteratorT, typename MemoryPoolIndexIteratorT>
__global__ void partition_merge_path_kernel(AllocationIndexIteratorT allocation_indices,
                                            uint64_t num_allocation_indices,
                                            MemoryPoolIndexIteratorT memory_pool_indices,
                                            uint64_t num_memory_pool_indices,
                                            uint64_t* merge_partitions,
                                            uint64_t num_merge_partitions)
{
  uint64_t partition_idx = blockDim.x * blockIdx.x + threadIdx.x;
  if (partition_idx < num_merge_partitions)
  {
    uint64_t diagonal =
      (cub::min)(TileSize * partition_idx, num_allocation_indices + num_memory_pool_indices);
    merge_partitions[partition_idx] = cub::MergePath(allocation_indices,
                                                     memory_pool_indices,
                                                     num_allocation_indices,
                                                     num_memory_pool_indices,
                                                     diagonal,
                                                     IndexCompareOpT{});
  }
}

/*
 * The compaction kernel. This is a complicated kernel and should be decomposed into modular
 * components in the system.
 */
template <int32_t BlockThreads,
          int32_t ItemsPerThread,
          typename ChunkT,
          typename ScanTileStateT,
          typename CounterT>
__global__ void compact_kernel(const uint64_t* allocation_indices,
                               uint64_t num_allocations,
                               ChunkT* memory_pool,
                               uint64_t num_memory_pool_chunks,
                               const uint64_t* merge_partitions,
                               ScanTileStateT scan_tile_state,
                               CounterT* counter)
{
  /// Static values
  static constexpr int32_t items_per_tile        = BlockThreads * ItemsPerThread;
  static constexpr auto default_allocation_index = set_top_digit<uint64_t>(0);

  /// Using declarations
  using block_load_t =
    cub::BlockLoad<ChunkT, BlockThreads, ItemsPerThread, cub::BLOCK_LOAD_STRIPED>;
  using block_load_storage_t     = typename block_load_t::TempStorage;
  using block_scan_t             = cub::BlockScan<int32_t, BlockThreads>;
  using block_scan_storage_t     = typename block_scan_t::TempStorage;
  using block_exchange_t         = cub::BlockExchange<uint64_t, BlockThreads, ItemsPerThread>;
  using block_exchange_storage_t = typename block_exchange_t::TempStorage;
  using block_store_t =
    cub::BlockStore<ChunkT, BlockThreads, ItemsPerThread, cub::BLOCK_STORE_STRIPED>;
  using block_store_storage_t = typename block_store_t::TempStorage;
  using scan_tile_state_t     = ScanTileStateT; // Should be cub::ScanTileState<uint64_t>;
  using scan_op_t             = cuda::std::plus<uint64_t>;
  using tile_prefix_op_t      = cub::TilePrefixCallbackOp<uint64_t, scan_op_t, scan_tile_state_t>;
  using prefix_op_storage_t   = typename tile_prefix_op_t::TempStorage;
  using compare_op_t          = IndexCompareOpT;

  /// Shared memory
  __shared__ union
  {
    uint64_t shared_previous_allocation_index;
    uint64_t shared_indices[items_per_tile + 1];
    block_scan_storage_t scan;
    block_exchange_storage_t exchange;
    prefix_op_storage_t prefix_op;
    block_store_storage_t store;
    block_load_storage_t load;
  } temp_storage;
  __shared__ uint64_t write_offset;

  /// Thread memory
  uint64_t thread_indices[ItemsPerThread];
  int32_t thread_valid_flags[ItemsPerThread];
  int32_t thread_ranks[ItemsPerThread];
  ChunkT thread_chunks[ItemsPerThread];
  compare_op_t compare_op{};
  scan_op_t scan_op{};
  int32_t num_valid = 0;
  bool is_partial_tile =
    ((blockIdx.x + 1) * items_per_tile) > (num_allocations + num_memory_pool_chunks);
  bool is_last_tile =
    ((blockIdx.x + 1) * items_per_tile) >= (num_allocations + num_memory_pool_chunks);

  /// Determine the bounding box for the range of allocation and memory pool indices to be processed
  /// by this thread block
  uint64_t allocation_indices_beg = merge_partitions[blockIdx.x + 0];
  uint64_t allocation_indices_end = merge_partitions[blockIdx.x + 1];

  uint64_t diag0 = items_per_tile * blockIdx.x;
  uint64_t diag1 = (cub::min)(num_allocations + num_memory_pool_chunks, diag0 + items_per_tile);

  uint64_t memory_pool_indices_beg = diag0 - allocation_indices_beg;
  uint64_t memory_pool_indices_end = diag1 - allocation_indices_end;

  auto num_allocation_indices_tile =
    static_cast<int32_t>(allocation_indices_end - allocation_indices_beg);
  auto num_memory_pool_indices_tile =
    static_cast<int32_t>(memory_pool_indices_end - memory_pool_indices_beg);

  // To carry forward the valid flag from a previous allocation index to a current memory pool
  // index, load the previous allocation index into shared memory
  if (threadIdx.x == 0)
  {
    if (allocation_indices_beg > 0)
    {
      temp_storage.shared_previous_allocation_index =
        allocation_indices[allocation_indices_beg - 1];
    }
    else
    {
      temp_storage.shared_previous_allocation_index = default_allocation_index;
    }
  }
  __syncthreads();
  uint64_t previous_allocation_index = temp_storage.shared_previous_allocation_index;

  // Optimized code path for when the entire tile's memory pool belongs to a single allocation
  // if (num_allocation_indices_tile == 0 ||
  //     (num_allocation_indices_tile == 1 && !compare_op(memory_pool[memory_pool_indices_beg],
  //                                                      allocation_indices[allocation_indices_beg])))
  if (num_allocation_indices_tile == 0) // Skip the merge
  {
    // uint64_t allocation_index = num_allocation_indices_tile == 0
    //                               ? previous_allocation_index
    //                               : allocation_indices[allocation_indices_beg];
    // int32_t num_valid = evaluate_top_digit(allocation_index) ? num_memory_pool_indices_tile : 0;
    num_valid = evaluate_top_digit(previous_allocation_index) ? num_memory_pool_indices_tile : 0;

    // We must load the data before doing the lookback!
    if (num_valid != 0)
    {
      block_load_t{temp_storage.load}.Load(memory_pool + memory_pool_indices_beg,
                                           thread_chunks,
                                           num_valid);
      __syncthreads();
    }
  }
  else // Do the merge
  {
    // Load the allocation and memory pool indices into registers and then from registers to shared
    // memory
    if (!is_partial_tile)
    {
#pragma unroll ItemsPerThread
      for (int32_t i = 0; i < ItemsPerThread; ++i)
      {
        int32_t shared_idx = threadIdx.x + BlockThreads * i;
        thread_indices[i]  = shared_idx < num_allocation_indices_tile
                               ? allocation_indices[allocation_indices_beg + shared_idx]
                               : static_cast<uint64_t>(memory_pool_indices_beg + shared_idx -
                                                      num_allocation_indices_tile);
      }
#pragma unroll ItemsPerThread
      for (int32_t i = 0; i < ItemsPerThread; ++i)
      {
        int32_t shared_idx                      = threadIdx.x + BlockThreads * i;
        temp_storage.shared_indices[shared_idx] = thread_indices[i];
      }
    }
    else
    { // partial tile
#pragma unroll ItemsPerThread
      for (int32_t i = 0; i < ItemsPerThread; ++i)
      {
        int32_t shared_idx = threadIdx.x + BlockThreads * i;
        if (shared_idx < (num_allocation_indices_tile + num_memory_pool_indices_tile))
        {
          thread_indices[i] = shared_idx < num_allocation_indices_tile
                                ? allocation_indices[allocation_indices_beg + shared_idx]
                                : static_cast<uint64_t>(memory_pool_indices_beg + shared_idx -
                                                        num_allocation_indices_tile);
        }
      }
#pragma unroll ItemsPerThread
      for (int32_t i = 0; i < ItemsPerThread; ++i)
      {
        int32_t shared_idx = threadIdx.x + BlockThreads * i;
        if (shared_idx < (num_allocation_indices_tile + num_memory_pool_indices_tile))
        {
          temp_storage.shared_indices[shared_idx] = thread_indices[i];
        }
      }
      // Initialize the flags (Maybe unnecessary?)
#pragma unroll ItemsPerThread
      for (int32_t i = 0; i < ItemsPerThread; ++i)
      {
        thread_valid_flags[i] = 0;
      }
    }
    __syncthreads();

    // Do binary search in shared memory to merge the set of keys for each thread
    auto diag0_thread =
      static_cast<int32_t>((cub::min)(num_allocation_indices_tile + num_memory_pool_indices_tile,
                                      ItemsPerThread * threadIdx.x));

    int32_t allocation_indices_beg_thread =
      cub::MergePath(&temp_storage.shared_indices[0],
                     &temp_storage.shared_indices[num_allocation_indices_tile],
                     num_allocation_indices_tile,
                     num_memory_pool_indices_tile,
                     diag0_thread,
                     compare_op);
    int32_t allocation_indices_end_thread =
      num_allocation_indices_tile; // I'm stealing this assignment from CUB, which I don't fully
                                   // understand
    int32_t memory_pool_indices_beg_thread = diag0_thread - allocation_indices_beg_thread;
    int32_t memory_pool_indices_end_thread = num_memory_pool_indices_tile;

    /// Do the 'merge'
    uint64_t allocation_index = temp_storage.shared_indices[allocation_indices_beg_thread];
    uint64_t memory_pool_index =
      temp_storage.shared_indices[num_allocation_indices_tile + memory_pool_indices_beg_thread];
    if (allocation_indices_beg_thread > 0)
    {
      previous_allocation_index = temp_storage.shared_indices[allocation_indices_beg_thread - 1];
    }
    bool is_valid = evaluate_top_digit(previous_allocation_index);

#pragma unroll ItemsPerThread
    for (int32_t i = 0; i < ItemsPerThread; ++i)
    {
      bool is_memory_pool_index =
        (memory_pool_indices_beg_thread < memory_pool_indices_end_thread) &&
        ((allocation_indices_beg_thread >= allocation_indices_end_thread) ||
         compare_op(memory_pool_index, allocation_index));

      thread_valid_flags[i] = is_memory_pool_index && is_valid;
      thread_indices[i]     = is_memory_pool_index ? memory_pool_index : allocation_index;

      if (is_memory_pool_index)
      {
        ++memory_pool_indices_beg_thread;
        memory_pool_index =
          temp_storage.shared_indices[num_allocation_indices_tile + memory_pool_indices_beg_thread];
      }
      else
      { // allocation index
        is_valid = evaluate_top_digit(allocation_index);
        ++allocation_indices_beg_thread;
        allocation_index = temp_storage.shared_indices[allocation_indices_beg_thread];
      }
    }
    __syncthreads();

    // Now we will scatter the thread_indices based on their validity before loading data for
    // compaction
    block_scan_t(temp_storage.scan).ExclusiveSum(thread_valid_flags, thread_ranks, num_valid);
    __syncthreads();

    block_exchange_t(temp_storage.exchange)
      .ScatterToStripedFlagged(thread_indices, thread_indices, thread_ranks, thread_valid_flags);
    __syncthreads();

    // Load the data from the memory pool (coalesced)
#pragma unroll ItemsPerThread
    for (int32_t i = 0; i < ItemsPerThread; ++i)
    {
      int32_t shared_idx = threadIdx.x + BlockThreads * i;
      if (shared_idx < num_valid)
      // if (thread_valid_flags[i])
      {
        thread_chunks[i] = memory_pool[thread_indices[i]];
      }
    }
  }

  // Now do decoupled look-back
  if (blockIdx.x == 0)
  {
    if (threadIdx.x == 0)
    {
      scan_tile_state.SetInclusive(blockIdx.x, num_valid);
      write_offset = 0;

      // If it's the last tile, update the counter
      if (is_last_tile)
      {
        *counter = num_valid;
      }
    }
  }
  else
  {
    // The first warp does the look-back
    if (threadIdx.x < CUB_PTX_WARP_THREADS)
    {
      // Initialize the prefix op
      tile_prefix_op_t prefix_op(scan_tile_state, temp_storage.prefix_op, scan_op);

      // De the decoupled look-back
      uint64_t exclusive_num_valid = prefix_op(static_cast<uint64_t>(num_valid));

      // Update the tile-wide offset
      if (threadIdx.x == 0)
      {
        write_offset = exclusive_num_valid;

        // If it's the last tile, update the counter
        if (is_last_tile)
        {
          *counter = prefix_op.GetInclusivePrefix();
        }
      }
    }
  }

  // Store data, if necessary
  if (num_valid != 0)
  {
    __syncthreads(); // For write offset

    // Now that look-back is done, we can write back to the memory pool (if we need to)
    if (write_offset != memory_pool_indices_beg)
    {
      block_store_t(temp_storage.store).Store(memory_pool + write_offset, thread_chunks, num_valid);
      // #pragma unroll ItemsPerThread
      //       for (int32_t i = 0; i < ItemsPerThread; ++i)
      //       {
      //         if (thread_valid_flags[i])
      //         {
      //           memory_pool[write_offset + thread_ranks[i]] = thread_chunks[i];
      //         }
      //       }
    }
  }
}

/*
 * The new compaction kernel (under construction)
 */
// template <int32_t BlockThreads,
//           int32_t ItemsPerThread,
//           typename ChunkT,
//           typename ScanTileStateT,
//           typename CounterT>
// __global__ void compact_kernel_new(const uint64_t* old_allocation_indices,
//                                    const uint64_t* new_allocation_indices,
//                                    uint64_t num_new_allocation_indices,
//                                    ChunkT* memory_pool,
//                                    uint64_t num_memory_pool_chunks,
//                                    const uint64_t* merge_partitions,
//                                    ScanTileStateT scan_tile_state,
//                                    CounterT* counter)
// {
//   /// Static values
//   static constexpr int32_t items_per_tile = BlockThreads * ItemsPerThread;

//   /// Using declarations
//   using block_load_t =
//     cub::BlockLoad<ChunkT, BlockThreads, ItemsPerThread, cub::BLOCK_LOAD_STRIPED>;
//   using block_load_storage_t = typename block_load_t::TempStorage;
//   using block_store_t =
//     cub::BlockStore<ChunkT, BlockThreads, ItemsPerThread, cub::BLOCK_STORE_STRIPED>;
//   using block_store_storage_t = typename block_store_t::TempStorage;
//   using scan_tile_state_t     = ScanTileStateT; // Should be cub::ScanTileState<bool>;
//   using scan_op_t             = cuda::std::plus<uint64_t>;
//   using tile_prefix_op_t      = cub::TilePrefixCallbackOp<uint64_t, scan_op_t, scan_tile_state_t>;
//   using prefix_op_storage_t   = typename tile_prefix_op_t::TempStorage;
//   using compare_op_t          = IndexCompareOpT;

//   /// Shared memory
//   __shared__ union
//   {
//     uint64_t shared_previous_allocation_index;
//     uint64_t shared_indices[items_per_tile + 1];
//     prefix_op_storage_t prefix_op;
//     block_store_storage_t store;
//     block_load_storage_t load;
//   } temp_storage;
//   __shared__ uint64_t write_offset;

//   /// Thread memory
//   uint64_t thread_new_indices[ItemsPerThread];
//   uint64_t thread_old_indices[ItemsPerThread];
//   ChunkT thread_chunks[ItemsPerThread];
//   compare_op_t compare_op{};
//   scan_op_t scan_op{};
//   int32_t num_valid = 0;
//   bool is_last_tile =
//     ((blockIdx.x + 1) * items_per_tile) >= (num_new_allocation_indices + num_memory_pool_chunks);

//   /// Determine the bounding box for the range of allocation and memory pool indices to be processed
//   /// by this thread block
//   uint64_t allocation_indices_beg = merge_partitions[blockIdx.x + 0];
//   uint64_t allocation_indices_end = merge_partitions[blockIdx.x + 1];

//   uint64_t diag0 = items_per_tile * blockIdx.x;
//   uint64_t diag1 =
//     (cub::min)(num_new_allocation_indices + num_memory_pool_chunks, diag0 + items_per_tile);

//   uint64_t memory_pool_indices_beg = diag0 - allocation_indices_beg;
//   uint64_t memory_pool_indices_end = diag1 - allocation_indices_end;

//   auto num_allocation_indices_tile =
//     static_cast<int32_t>(allocation_indices_end - allocation_indices_beg);
//   auto num_memory_pool_indices_tile =
//     static_cast<int32_t>(memory_pool_indices_end - memory_pool_indices_beg);

//   // To carry forward the valid flag from a previous allocation index to a current memory pool
//   // index, load the previous allocation index into shared memory
//   if (threadIdx.x == 0)
//   {
//     if (allocation_indices_beg > 0)
//     {
//       temp_storage.shared_previous_allocation_index =
//         allocation_indices[allocation_indices_beg - 1];
//     }
//     else // You cannot go backwards from zero
//     {
//       temp_storage.shared_previous_allocation_index = 0;
//     }
//   }
//   __syncthreads();
//   uint64_t previous_allocation_index = temp_storage.shared_previous_allocation_index;

//   // Optimized code path for when the entire tile's memory pool belongs to a single allocation
//   // if (num_allocation_indices_tile == 0 ||
//   //     (num_allocation_indices_tile == 1 && !compare_op(memory_pool[memory_pool_indices_beg],
//   //                                                      allocation_indices[allocation_indices_beg])))
//   if (num_allocation_indices_tile == 0) // Skip the merge
//   {
//     // uint64_t allocation_index = num_allocation_indices_tile == 0
//     //                               ? previous_allocation_index
//     //                               : allocation_indices[allocation_indices_beg];
//     // int32_t num_valid = evaluate_top_digit(allocation_index) ? num_memory_pool_indices_tile : 0;

//     // We must load the data before doing the lookback!
//     block_load_t{temp_storage.load}.Load(memory_pool + memory_pool_indices_beg,
//                                          thread_chunks,
//                                          num_memory_pool_indices_tile);
//     __syncthreads();
//   }
//   else // Do the merge
//   {
//     // Load the allocation and memory pool indices into registers and then from registers to shared
//     // memory
//     if (!is_last_tile)
//     {
// #pragma unroll ItemsPerThread
//       for (int32_t i = 0; i < ItemsPerThread; ++i)
//       {
//         int32_t shared_idx    = threadIdx.x + BlockThreads * i;
//         thread_new_indices[i] = shared_idx < num_allocation_indices_tile
//                                   ? new_allocation_indices[allocation_indices_beg + shared_idx]
//                                   : static_cast<uint64_t>(memory_pool_indices_beg + shared_idx -
//                                                           num_allocation_indices_tile);
//       }
// #pragma unroll ItemsPerThread
//       for (int32_t i = 0; i < ItemsPerThread; ++i)
//       {
//         int32_t shared_idx                      = threadIdx.x + BlockThreads * i;
//         temp_storage.shared_indices[shared_idx] = thread_new_indices[i];
//       }
//     }
//     else // The last tile
//     {
// #pragma unroll ItemsPerThread
//       for (int32_t i = 0; i < ItemsPerThread; ++i)
//       {
//         int32_t shared_idx = threadIdx.x + BlockThreads * i;
//         if (shared_idx < (num_allocation_indices_tile + num_memory_pool_indices_tile))
//         {
//           thread_new_indices[i] = shared_idx < num_allocation_indices_tile
//                                     ? allocation_indices[allocation_indices_beg + shared_idx]
//                                     : static_cast<uint64_t>(memory_pool_indices_beg + shared_idx -
//                                                             num_allocation_indices_tile);
//         }
//       }
// #pragma unroll ItemsPerThread
//       for (int32_t i = 0; i < ItemsPerThread; ++i)
//       {
//         int32_t shared_idx = threadIdx.x + BlockThreads * i;
//         if (shared_idx < (num_allocation_indices_tile + num_memory_pool_indices_tile))
//         {
//           temp_storage.shared_indices[shared_idx] = thread_new_indices[i];
//         }
//       }
//       // Initialize the flags (Maybe unnecessary?)
// #pragma unroll ItemsPerThread
//       for (int32_t i = 0; i < ItemsPerThread; ++i)
//       {
//         thread_valid_flags[i] = 0;
//       }
//     }
//     __syncthreads();

//     // Do binary search in shared memory to merge the set of keys for each thread
//     auto diag0_thread =
//       static_cast<int32_t>((cub::min)(num_allocation_indices_tile + num_memory_pool_indices_tile,
//                                       ItemsPerThread * threadIdx.x));

//     int32_t allocation_indices_beg_thread =
//       cub::MergePath(&temp_storage.shared_indices[0],
//                      &temp_storage.shared_indices[num_allocation_indices_tile],
//                      num_allocation_indices_tile,
//                      num_memory_pool_indices_tile,
//                      diag0_thread,
//                      compare_op);
//     int32_t allocation_indices_end_thread =
//       num_allocation_indices_tile; // I'm stealing this assignment from CUB, which I don't fully
//                                    // understand
//     int32_t memory_pool_indices_beg_thread = diag0_thread - allocation_indices_beg_thread;
//     int32_t memory_pool_indices_end_thread = num_memory_pool_indices_tile;

//     /// Do the 'merge'
//     uint64_t allocation_index = temp_storage.shared_indices[allocation_indices_beg_thread];
//     uint64_t memory_pool_index =
//       temp_storage.shared_indices[num_allocation_indices_tile + memory_pool_indices_beg_thread];
//     if (allocation_indices_beg_thread > 0)
//     {
//       previous_allocation_index = temp_storage.shared_indices[allocation_indices_beg_thread - 1];
//     }
//     uint64_t old_index = old_allocation_indices[previous_allocation_index];

// #pragma unroll ItemsPerThread
//     for (int32_t i = 0; i < ItemsPerThread; ++i)
//     {
//       bool is_memory_pool_index =
//         (memory_pool_indices_beg_thread < memory_pool_indices_end_thread) &&
//         ((allocation_indices_beg_thread >= allocation_indices_end_thread) ||
//          compare_op(memory_pool_index, allocation_index));

//       thread_valid_flags[i] = is_memory_pool_index && is_valid;
//       thread_new_indices[i] = is_memory_pool_index ? memory_pool_index : allocation_index;

//       if (is_memory_pool_index)
//       {
//         ++memory_pool_indices_beg_thread;
//         memory_pool_index =
//           temp_storage.shared_indices[num_allocation_indices_tile + memory_pool_indices_beg_thread];
//       }
//       else
//       { // allocation index
//         is_valid = evaluate_top_digit(allocation_index);
//         ++allocation_indices_beg_thread;
//         allocation_index = temp_storage.shared_indices[allocation_indices_beg_thread];
//       }
//     }
//     __syncthreads();

//     // Load the data from the memory pool (coalesced)
// #pragma unroll ItemsPerThread
//     for (int32_t i = 0; i < ItemsPerThread; ++i)
//     {
//       int32_t shared_idx = threadIdx.x + BlockThreads * i;
//       if (shared_idx < num_valid)
//       // if (thread_valid_flags[i])
//       {
//         thread_chunks[i] = memory_pool[thread_old_indices[i]];
//       }
//     }
//   }

//   // Now do decoupled look-back
//   if (blockIdx.x == 0)
//   {
//     if (threadIdx.x == 0)
//     {
//       scan_tile_state.SetInclusive(blockIdx.x, true);
//     }
//   }
//   else
//   {
//     // The first warp does the look-back
//     if (threadIdx.x < CUB_PTX_WARP_THREADS)
//     {
//       // Initialize the prefix op
//       tile_prefix_op_t prefix_op(scan_tile_state, temp_storage.prefix_op, scan_op);

//       // De the decoupled look-back
//       prefix_op(true);
//     }
//   }
//   __syncthreads();

//   // Store data (TODO)
// }