/* Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once
#include <cuda_runtime.h>
#include <math.h>
#include <cmath>
#include <sstream>
#include <string>
#include "cub/cub.cuh"
#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
namespace phi {

static const float HALF_FLT_MAX = 65504.F;
static const float HALF_FLT_MIN = -65504.F;
static inline size_t AlignTo16(const size_t& input) {
  static constexpr int ALIGNMENT = 16;
  return ALIGNMENT * ((input + ALIGNMENT - 1) / ALIGNMENT);
}

class CubKeyValueSorter {
 public:
  inline CubKeyValueSorter();

  inline CubKeyValueSorter(cudaStream_t stream = 0);  // NOLINT

  inline explicit CubKeyValueSorter(const int num_experts);

  inline void update_num_experts(const int num_experts);

  inline size_t getWorkspaceSize(const size_t num_key_value_pairs,
                                 bool descending = false);

  template <typename KeyT>
  inline void run(void* workspace,
                  const size_t workspace_size,
                  const KeyT* keys_in,
                  KeyT* keys_out,
                  const int* values_in,
                  int* values_out,
                  const size_t num_key_value_pairs,
                  bool descending,
                  cudaStream_t stream);

 private:
  size_t num_key_value_pairs_;
  int num_experts_;
  int num_bits_;
  cudaStream_t stream_;
};

// ===== CUB Sorting things =====
CubKeyValueSorter::CubKeyValueSorter()
    : num_experts_(0), num_bits_(sizeof(int) * 8) {}

CubKeyValueSorter::CubKeyValueSorter(cudaStream_t stream)
    : num_experts_(0), num_bits_(sizeof(int) * 8), stream_(stream) {}

CubKeyValueSorter::CubKeyValueSorter(const int num_experts)
    : num_experts_(num_experts),
      num_bits_(static_cast<int>(log2(num_experts)) + 1) {}

void CubKeyValueSorter::update_num_experts(const int num_experts) {
  num_experts_ = num_experts;
  num_bits_ = static_cast<int>(log2(num_experts)) +
              3;  // 额外增加 3 位用于标记 topk的位置
}

size_t CubKeyValueSorter::getWorkspaceSize(const size_t num_key_value_pairs,
                                           bool descending) {
  num_key_value_pairs_ = num_key_value_pairs;
  size_t required_storage = 0;
  int* null_int = nullptr;
  if (descending) {
    cub::DeviceRadixSort::SortPairsDescending(NULL,
                                              required_storage,
                                              null_int,
                                              null_int,
                                              null_int,
                                              null_int,
                                              num_key_value_pairs,
                                              0,
                                              32,
                                              stream_);
  } else {
    cub::DeviceRadixSort::SortPairs(NULL,
                                    required_storage,
                                    null_int,
                                    null_int,
                                    null_int,
                                    null_int,
                                    num_key_value_pairs,
                                    0,
                                    num_bits_,
                                    stream_);
  }
  return required_storage;
}

template <typename KeyT>
inline void CubKeyValueSorter::run(void* workspace,
                                   const size_t workspace_size,
                                   const KeyT* keys_in,
                                   KeyT* keys_out,
                                   const int* values_in,
                                   int* values_out,
                                   const size_t num_key_value_pairs,
                                   bool descending,
                                   cudaStream_t stream) {
  size_t expected_ws_size = getWorkspaceSize(num_key_value_pairs);
  size_t actual_ws_size = workspace_size;

  if (expected_ws_size > workspace_size) {
    std::stringstream err_ss;
    err_ss << "[Error][CubKeyValueSorter::run]\n";
    err_ss
        << "Error. The allocated workspace is too small to run this problem.\n";
    err_ss << "Expected workspace size of at least " << expected_ws_size
           << " but got problem size " << workspace_size << "\n";
    throw std::runtime_error(err_ss.str());
  }
  if (descending) {
    cub::DeviceRadixSort::SortPairsDescending(workspace,
                                              actual_ws_size,
                                              keys_in,
                                              keys_out,
                                              values_in,
                                              values_out,
                                              num_key_value_pairs,
                                              0,
                                              32,
                                              stream);
  } else {
    cub::DeviceRadixSort::SortPairs(workspace,
                                    actual_ws_size,
                                    keys_in,
                                    keys_out,
                                    values_in,
                                    values_out,
                                    num_key_value_pairs,
                                    0,
                                    num_bits_,
                                    stream);
  }
}

template <>
inline void CubKeyValueSorter::run(void* workspace,
                                   const size_t workspace_size,
                                   const __nv_bfloat16* keys_in,
                                   __nv_bfloat16* keys_out,
                                   const int* values_in,
                                   int* values_out,
                                   const size_t num_key_value_pairs,
                                   bool descending,
                                   cudaStream_t stream) {}

// CubKeyValueSorter sorter_(stream);

// --------      initialize_expert_choice_route_kernel      -------- //
template <typename T>
__global__ void initialize_expert_choice_route_kernel(
    int* expert_for_source_row,
    int* source_row,
    int* expanded_source_row_to_expanded_dest_row,
    int64_t* total_rows_before_expert,
    T* attr_mask,
    const int cols,
    const int k,
    const int batch_size) {
  int64_t start = static_cast<int64_t>(cols) * static_cast<int64_t>(blockIdx.x);

  for (int i = threadIdx.x; i < cols; i += blockDim.x) {
    expert_for_source_row[start + i] = blockIdx.x;
    source_row[start + i] = start + i;
    expanded_source_row_to_expanded_dest_row[start + i] = -1;
    attr_mask[start + i] = (T)1.0f;
  }
  if (threadIdx.x == 0) {
    total_rows_before_expert[blockIdx.x] =
        static_cast<int64_t>(batch_size) * k * (blockIdx.x + 1);
  }
}

// --------      softmax_kernel      -------- //
template <int ITEMS_PER_THREAD, typename T>
__global__ void softmax_kernel_v4(
    T* qk_buf_,
    const T* qk_buf_src,  // shape [batch_size, seq_len]
    const T* attr_mask,   // shape [batch_size, seq_len]
    const int batch_size,
    const int seq_len) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
  float data[ITEMS_PER_THREAD];
  int64_t qk_offset;
  __shared__ float s_mean, s_max;
  float local_max = -1e20f;
  for (int i = 0; blockDim.x * i + threadIdx.x < seq_len; i++) {
    qk_offset = static_cast<int64_t>(blockIdx.y + blockIdx.z) * seq_len +
                static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
    int64_t mask_offset = static_cast<int64_t>(blockIdx.y) * seq_len +
                          static_cast<int64_t>(blockDim.x) * i + threadIdx.x;

    float qk = static_cast<float>(qk_buf_src[qk_offset]);
    float mask_val = static_cast<float>(__ldg(&attr_mask[mask_offset]));

    mask_val = (1.0f - mask_val) * -10000.0f;

    data[i] = qk + mask_val;
    local_max = fmax(local_max, data[i]);
  }

  float max_val = blockDim.x <= 32
                      ? funcs::WarpReduceMax<float>(local_max, 0xFFFFFFFF)
                      : funcs::BlockReduceMax<float>(local_max, 0xFFFFFFFF);
  if (threadIdx.x == 0) {
    s_max = max_val;
  }
  __syncthreads();

  float local_sum = 0;
  for (int i = 0; blockDim.x * i + threadIdx.x < seq_len; i++) {
    data[i] = __expf(data[i] - s_max);
    local_sum += data[i];
  }
  float sum_val = blockDim.x <= 32
                      ? funcs::WarpReduceSum<float>(local_sum, 0xFFFFFFFF)
                      : funcs::BlockReduceSum<float>(local_sum, 0xFFFFFFFF);
  if (threadIdx.x == 0) {
    s_mean = sum_val + 1e-6f;
    s_mean = __fdividef(1.0f, s_mean);
  }
  __syncthreads();

  for (int i = 0; blockDim.x * i + threadIdx.x < seq_len; i++) {
    qk_offset = static_cast<int64_t>(blockIdx.y + blockIdx.z) * seq_len +
                static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
    qk_buf_[qk_offset] = (T)(data[i] * s_mean);
  }
#endif
}

template <typename T, int ITEMS_PER_THREAD>
__global__ void softmax_kernel_v4_half2(T* qk_buf_,
                                        const T* attr_mask,
                                        const int batch_size,
                                        const int seq_len) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
  using T2 = half2;
  T2* qk_buf_half2 = reinterpret_cast<T2*>(qk_buf_);
  const T2* attr_mask_half2 = (const T2*)attr_mask;

  T2 data[ITEMS_PER_THREAD];
  int64_t qk_offset;
  __shared__ float s_mean, s_max;
  float local_max = -1e20f;
  for (int i = 0;
       blockDim.x * i + threadIdx.x < (seq_len / 2) && i < ITEMS_PER_THREAD;
       i++) {
    qk_offset = static_cast<int64_t>(blockIdx.y + blockIdx.z) * (seq_len / 2) +
                static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
    int64_t mask_offset = static_cast<int64_t>(blockIdx.y) * (seq_len / 2) +
                          static_cast<int64_t>(blockDim.x) * i + threadIdx.x;

    T2 qk = qk_buf_half2[qk_offset];
    T2 mask_val = __ldg(&attr_mask_half2[mask_offset]);
    mask_val = __hmul2(__hsub2(__float2half2_rn(1.0f), mask_val),
                       __float2half2_rn(-10000.0f));

    data[i] = __hadd2(qk, mask_val);

    local_max = fmax(
        local_max,
        fmax(static_cast<float>(data[i].x), static_cast<float>(data[i].y)));
  }

  float max_val = blockDim.x <= 32
                      ? funcs::WarpReduceMax<float>(local_max, 0xFFFFFFFF)
                      : funcs::BlockReduceMax<float>(local_max, 0xFFFFFFFF);
  if (threadIdx.x == 0) {
    s_max = max_val;
  }
  __syncthreads();

  float local_sum = 0;
  for (int i = 0;
       blockDim.x * i + threadIdx.x < (seq_len / 2) && i < ITEMS_PER_THREAD;
       i++) {
    data[i] = h2exp(__hsub2(data[i], __float2half2_rn(s_max)));
    local_sum += static_cast<float>(data[i].x + data[i].y);
  }

  float sum_val = blockDim.x <= 32
                      ? funcs::WarpReduceSum<float>(local_sum, 0xFFFFFFFF)
                      : funcs::BlockReduceSum<float>(local_sum, 0xFFFFFFFF);

  if (threadIdx.x == 0) {
    s_mean = sum_val + 1e-6f;
    s_mean = __fdividef(1.0f, s_mean);
  }
  __syncthreads();

  for (int i = 0;
       blockDim.x * i + threadIdx.x < (seq_len / 2) && i < ITEMS_PER_THREAD;
       i++) {
    qk_offset = static_cast<int64_t>(blockIdx.y + blockIdx.z) * (seq_len / 2) +
                static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
    qk_buf_half2[qk_offset] = __hmul2(data[i], __float2half2_rn(s_mean));
  }
#endif
}

template <typename T, int ITEMS_PER_THREAD, int NUM>
__global__ void softmax_kernel_v5_half2(T* qk_buf_,
                                        const T* attr_mask,
                                        const int batch_size,
                                        const int seq_len) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
  using T2 = half2;
  T2* qk_buf_half2 = reinterpret_cast<T2*>(qk_buf_);
  const T2* attr_mask_half2 = (const T2*)attr_mask;

  T2 data[NUM][ITEMS_PER_THREAD];

  int64_t qk_offset[NUM];

  __shared__ float s_sum[NUM], s_max[NUM];
  float local_max[NUM];
#pragma unroll
  for (int j = 0; j < NUM; j++) {
    local_max[j] = -1e20f;
  }

  const int MAX_NUM = min((1 + gridDim.x - 1) / gridDim.x, NUM);
  for (int i = 0;
       blockDim.x * i + threadIdx.x < (seq_len / 2) && i < ITEMS_PER_THREAD;
       i++) {
    int64_t mask_offset[NUM];
#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      qk_offset[j] =
          static_cast<int64_t>((blockIdx.y + blockIdx.z) + j * gridDim.x) *
              (seq_len / 2) +
          static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
      mask_offset[j] =
          static_cast<int64_t>(blockIdx.y + j * gridDim.x) * (seq_len / 2) +
          static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
    }

    T2 mask_val[NUM];
#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      mask_val[j] = __ldg(&attr_mask_half2[mask_offset[j]]);
    }

    T2 qk[NUM];
#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      qk[j] = qk_buf_half2[qk_offset[j]];
    }
#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      mask_val[j] = __hmul2(__hsub2(__float2half2_rn(1.0f), mask_val[j]),
                            __float2half2_rn(-10000.0f));
    }
#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      data[j][i] = __hadd2(qk[j], mask_val[j]);
      local_max[j] = fmax(local_max[j],
                          fmax(static_cast<float>(data[j][i].x),
                               static_cast<float>(data[j][i].y)));
    }
  }
  if (blockDim.x <= 32) {
    funcs::WarpReduceMaxV2<float, NUM>(local_max);
  } else {
    funcs::BlockReduceMaxV2<float, NUM>(local_max);
  }

  if (threadIdx.x == 0) {
#pragma unroll
    for (int j = 0; j < NUM; j++) {
      s_max[j] = local_max[j];
    }
  }
  __syncthreads();
  float local_sum[NUM];
#pragma unroll
  for (int j = 0; j < NUM; j++) {
    local_sum[j] = {0.f};
  }

  for (int i = 0;
       blockDim.x * i + threadIdx.x < (seq_len / 2) && i < ITEMS_PER_THREAD;
       i++) {
#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      data[j][i] = h2exp(__hsub2(data[j][i], __float2half2_rn(s_max[j])));
    }

#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      local_sum[j] += static_cast<float>(data[j][i].x + data[j][i].y);
    }
  }

  if (blockDim.x <= 32) {
    funcs::WarpReduceSumV2<float, NUM>(local_sum);

  } else {
    funcs::BlockReduceSumV2<float, NUM>(local_sum);
  }

  if (threadIdx.x == 0) {
#pragma unroll
    for (int j = 0; j < NUM; j++) {
      s_sum[j] = __fdividef(1.0f, local_sum[j] + 1e-6f);
    }
  }
  __syncthreads();

  for (int i = 0;
       blockDim.x * i + threadIdx.x < (seq_len / 2) && i < ITEMS_PER_THREAD;
       i++) {
#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      qk_offset[j] =
          static_cast<int64_t>((blockIdx.y + blockIdx.z) + j * gridDim.x) *
              (seq_len / 2) +
          static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
    }

#pragma unroll
    for (int j = 0; j < MAX_NUM; j++) {
      qk_buf_half2[qk_offset[j]] =
          __hmul2(data[j][i], __float2half2_rn(s_sum[j]));
    }
  }
#endif
}

// --------      transpose_kernel      -------- //
template <typename T>
__global__ void transposeAxis01(
    T* out, T* in, const int dim0, const int dim1, const int dim2) {
  int64_t index =
      static_cast<int64_t>(threadIdx.x) +
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x);
  if (index < static_cast<int64_t>(dim0) * dim1 * dim2) {
    const int input_dim2_index = index % dim2;
    index = (index - input_dim2_index) / dim2;
    const int input_dim1_index = index % dim1;
    index = (index - input_dim1_index) / dim1;
    const int input_dim0_index = index % dim0;

    out[static_cast<int64_t>(input_dim1_index) * dim0 * dim2 +
        static_cast<int64_t>(input_dim0_index) * dim2 + input_dim2_index] =
        in[static_cast<int64_t>(input_dim0_index) * dim1 * dim2 +
           static_cast<int64_t>(input_dim1_index) * dim2 + input_dim2_index];
  }
}

// --------      padding_kernel      -------- //
template <typename T>
__global__ void paddingKernel(T* output1,
                              int* output2,
                              const T* input1,
                              const int* input2,
                              const int* input_lengths,
                              const int num_tokens,
                              const int batch_size,
                              const int max_seq_len,
                              const int num_experts) {
  const bool IS_FP32 = std::is_same<T, float>::value;
  const T MIN_T_VAL = (!IS_FP32) ? (T)HALF_FLT_MIN : (T)FLT_MIN;
  int64_t offset1 = static_cast<int64_t>(blockIdx.x) * num_tokens;
  int64_t offset2 = static_cast<int64_t>(blockIdx.x) * batch_size * max_seq_len;
  for (int i = 0; i < batch_size; i++) {
    const T* in1_ptr = input1 + offset1;
    const int* in2_ptr = input2 + offset1;
    int input_length = input_lengths[i];
    offset1 += input_length;

    T* out1_ptr = output1 + offset2;
    int* out2_ptr = output2 + offset2;
    offset2 += max_seq_len;

    for (int j = threadIdx.x; j < max_seq_len; j += max_seq_len) {
      if (j < input_length) {
        out1_ptr[j] = in1_ptr[j];
        out2_ptr[j] = in2_ptr[j];
      } else {
        out1_ptr[j] = MIN_T_VAL;
        out2_ptr[j] = 0;
      }
    }
  }
}

// --------      general_topk_pair_sort_kernel      -------- //
template <typename T, int BLOCK_THREADS, int ITEMS_PER_THREAD>
__global__ void general_topk_pair_sort(T* out_keys,
                                       int* out_values,
                                       T* in_keys,
                                       int* in_values) {
  typedef cub::BlockRadixSort<T, BLOCK_THREADS, ITEMS_PER_THREAD, int>
      BlockRadixSort;
  typedef cub::
      BlockLoad<T, BLOCK_THREADS, ITEMS_PER_THREAD, cub::BLOCK_LOAD_TRANSPOSE>
          BlockLoadKey;
  typedef cub::
      BlockLoad<int, BLOCK_THREADS, ITEMS_PER_THREAD, cub::BLOCK_LOAD_TRANSPOSE>
          BlockLoadValue;
  typedef cub::
      BlockStore<T, BLOCK_THREADS, ITEMS_PER_THREAD, cub::BLOCK_STORE_TRANSPOSE>
          BlockStoreKey;
  typedef cub::BlockStore<int,
                          BLOCK_THREADS,
                          ITEMS_PER_THREAD,
                          cub::BLOCK_STORE_TRANSPOSE>
      BlockStoreValue;

  __shared__ union {
    typename BlockRadixSort::TempStorage sort;
    typename BlockLoadKey::TempStorage loadkey;
    typename BlockLoadValue::TempStorage loadvalue;
    typename BlockStoreKey::TempStorage storekey;
    typename BlockStoreValue::TempStorage storevalue;
  } temp_storage;

  int64_t block_offset =
      static_cast<int64_t>(blockIdx.x) * BLOCK_THREADS * ITEMS_PER_THREAD;

  T thread_keys[ITEMS_PER_THREAD];
  int thread_values[ITEMS_PER_THREAD];
  BlockLoadKey(temp_storage.loadkey).Load(in_keys + block_offset, thread_keys);
  BlockLoadValue(temp_storage.loadvalue)
      .Load(in_values + block_offset, thread_values);
  __syncthreads();

  BlockRadixSort(temp_storage.sort).SortDescending(thread_keys, thread_values);
  __syncthreads();

  BlockStoreKey(temp_storage.storekey)
      .Store(out_keys + block_offset, thread_keys);
  BlockStoreValue(temp_storage.storevalue)
      .Store(out_values + block_offset, thread_values);
}

// --------      finalize_moe_routing_kernel      -------- //
template <typename T>
__global__ void finalize_moe_routing_kernel(
    const T* expanded_permuted_rows,
    T* reduced_unpermuted_output,
    const T* skip,
    const T* bias,
    const T* scales,
    const int* expanded_source_row_to_expanded_dest_row,
    const int* expert_for_source_row,
    const int cols,
    const int k,
    bool ec_route) {
  const int original_row = blockIdx.x;
  const int num_rows = gridDim.x;
  T* reduced_row_ptr =
      reduced_unpermuted_output + static_cast<int64_t>(original_row) * cols;
  const T* skip_row_ptr = skip + static_cast<int64_t>(original_row) * cols;

  for (int tid = threadIdx.x; tid < cols; tid += blockDim.x) {
    T thread_output = skip_row_ptr[tid];
    for (int k_idx = 0; k_idx < k; ++k_idx) {
      const int64_t expanded_original_row =
          original_row + static_cast<int64_t>(k_idx) * num_rows;
      const int expanded_permuted_row =
          expanded_source_row_to_expanded_dest_row[expanded_original_row];

      if (ec_route && expanded_permuted_row == -1) continue;
      const int64_t k_offset =
          ec_route ? expanded_original_row
                   : static_cast<int64_t>(original_row) * k + k_idx;
      const T row_scale = scales[k_offset];
      const T* expanded_permuted_rows_row_ptr =
          expanded_permuted_rows +
          static_cast<int64_t>(expanded_permuted_row) * cols;

      const int expert_idx = ec_route ? k_idx : expert_for_source_row[k_offset];
      const T* bias_ptr = bias + static_cast<int64_t>(expert_idx) * cols;

      thread_output =
          thread_output +
          row_scale * (expanded_permuted_rows_row_ptr[tid] + bias_ptr[tid]);
    }
    reduced_row_ptr[tid] = thread_output;
  }
}

// --------      initialize_moe_routing_kernel      -------- //
template <typename T, int VecSize>
__global__ void initialize_moe_routing_kernel(
    const T* unpermuted_input,
    T* permuted_output,
    const int* expanded_dest_row_to_expanded_source_row,
    int* expanded_source_row_to_expanded_dest_row,
    const int num_rows,
    const int active_rows,
    const int cols,
    const int k,
    const int max_seq_len,
    bool ec_route) {
  // using LoadT = AlignedVector<T, VecSize>;
  // LoadT src_vec;

  // Reverse permutation map.
  // I do this so that later, we can use the source -> dest map to do the k-way
  // reduction and unpermuting. I need the reverse map for that reduction to
  // allow each threadblock to do 1 k-way reduce without atomics later in MoE. 1
  // thread block will be responsible for all k summations.
  const int expanded_dest_row = blockIdx.x;
  const int expanded_source_row =
      ec_route
          ? expanded_dest_row_to_expanded_source_row
                [static_cast<int64_t>(expanded_dest_row / k) * max_seq_len +
                 expanded_dest_row % k]
          : expanded_dest_row_to_expanded_source_row[expanded_dest_row];
  if (threadIdx.x == 0) {
    expanded_source_row_to_expanded_dest_row[expanded_source_row] =
        expanded_dest_row;
  }

  if (blockIdx.x < active_rows) {
    // Duplicate and permute rows
    const int source_row = expanded_source_row % num_rows;

    const T* source_row_ptr =
        unpermuted_input + static_cast<int64_t>(source_row) * cols;
    T* dest_row_ptr =
        permuted_output + static_cast<int64_t>(expanded_dest_row) * cols;

    for (int tid = threadIdx.x * VecSize; tid < cols;
         tid += blockDim.x * VecSize) {
      dest_row_ptr[tid] = source_row_ptr[tid];
      // Load<T, VecSize>(&source_row_ptr[tid], &src_vec);
      // Store<T, VecSize>(src_vec, &dest_row_ptr[tid]);
    }
  }
}

}  // namespace phi