// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

#include <cstdint>

// --- Hash Helper Functions ---

__device__ inline uint32_t murmur_32_scramble(uint32_t k) {
  k *= 0xCC9E2D51;
  k = (k << 15) | (k >> 17);
  k *= 0x1B873593;
  return k;
}

__device__ inline uint32_t _hash_murmur_impl(uint32_t h, uint32_t k) {
  h ^= murmur_32_scramble(k);
  h = (h << 13) | (h >> 19);
  h = h * 5 + 0xE6546B64;
  return h;
}

__device__ inline uint32_t _hash_murmur_finalize(uint32_t h, int length_bytes) {
  h ^= length_bytes;
  h ^= h >> 16;
  h *= 0x85EBCA6B;
  h ^= h >> 13;
  h *= 0xC2B2AE35;
  h ^= h >> 16;
  return h;
}

__device__ inline uint32_t _hash_fnv1a_impl(uint32_t hash_val, uint32_t key) {
  hash_val ^= key;
  hash_val *= 16777619;  // FNV prime
  return hash_val;
}

__device__ inline uint32_t _hash_city_impl(uint32_t hash_val, uint32_t key) {
  hash_val += key * 0x9E3779B9;
  hash_val ^= hash_val >> 16;
  hash_val *= 0x85EBCA6B;
  hash_val ^= hash_val >> 13;
  hash_val *= 0xC2B2AE35;
  hash_val ^= hash_val >> 16;
  return hash_val;
}

// --- Array Hash Functions ---

// key: pointer to the start of the integer array key
// key_dim: number of integers in the key
// capacity: size of the hash table (MUST be power of 2)
__device__ inline int hash_fnv1a_array(const int* key, int key_dim, int capacity) {
  uint32_t hash_val = 2166136261u;  // FNV offset basis
  for (int i = 0; i < key_dim; ++i) {
    hash_val = _hash_fnv1a_impl(hash_val, (uint32_t)key[i]);
  }
  return (int)(hash_val & (uint32_t)(capacity - 1));
}

__device__ inline int hash_city_array(const int* key, int key_dim, int capacity) {
  uint32_t hash_val = 0;
  for (int i = 0; i < key_dim; ++i) {
    hash_val = _hash_city_impl(hash_val, (uint32_t)key[i]);
  }
  return (int)(hash_val & (uint32_t)(capacity - 1));
}

__device__ inline int hash_murmur_array(const int* key, int key_dim, int capacity) {
  uint32_t h = 0x9747B28Cu;  // Seed
  for (int i = 0; i < key_dim; ++i) {
    h = _hash_murmur_impl(h, (uint32_t)key[i]);
  }
  // Finalize (assuming key_dim * 4 bytes)
  h = _hash_murmur_finalize(h, key_dim * 4);
  return (int)(h & (uint32_t)(capacity - 1));
}

// --- Hash Function Structs/Functors ---
// NOTE: capacity MUST be a power of 2 for correct behavior.

struct FNV1AHash {
  __device__ inline static int hash(const int* key, int key_dim, int capacity) {
    uint32_t hash_val = 2166136261u;  // FNV offset basis
    for (int i = 0; i < key_dim; ++i) {
      hash_val = _hash_fnv1a_impl(hash_val, (uint32_t)key[i]);
    }
    return (int)(hash_val & (uint32_t)(capacity - 1));
  }
};

struct CityHash {
  __device__ inline static int hash(const int* key, int key_dim, int capacity) {
    uint32_t hash_val = 0;
    for (int i = 0; i < key_dim; ++i) {
      hash_val = _hash_city_impl(hash_val, (uint32_t)key[i]);
    }
    return (int)(hash_val & (uint32_t)(capacity - 1));
  }
};

struct MurmurHash {
  __device__ inline static int hash(const int* key, int key_dim, int capacity) {
    uint32_t h = 0x9747B28Cu;  // Seed
    for (int i = 0; i < key_dim; ++i) {
      h = _hash_murmur_impl(h, (uint32_t)key[i]);
    }
    h = _hash_murmur_finalize(h, key_dim * 4);
    return (int)(h & (uint32_t)(capacity - 1));
  }
};

// --- Vector Comparison ---
// a, b: pointers to the start of the vectors
// dim: dimension of the vectors
template <typename T>
__device__ inline bool vec_equal(const T* a, const int* b, int dim) {
  for (int i = 0; i < dim; ++i) {
    if (a[i] != b[i]) {
      return false;
    }
  }
  return true;
}

// Vectorized comparison for 4D keys using 128-bit int4 loads
__device__ inline bool vec_equal_4d(const int* a, const int* b) {
  int4 va = *reinterpret_cast<const int4*>(a);
  int4 vb = *reinterpret_cast<const int4*>(b);
  return (va.x == vb.x) && (va.y == vb.y) && (va.z == vb.z) && (va.w == vb.w);
}

enum ExpandStatus : int {
  kExpandStatusSuccess = 0,
  kExpandStatusVectorOverflow = 1,
  kExpandStatusTableFull = 2,
};

__device__ inline void set_expand_status(int* status_ptr, ExpandStatus new_status) {
  if (status_ptr == nullptr) {
    return;
  }
  atomicCAS(status_ptr, kExpandStatusSuccess, static_cast<int>(new_status));
}

// --- Helper for finding or claiming a slot ---
template <typename HashFuncT>
__device__ inline int claim_slot_or_find(int* table_kvs,
                                         const int* vector_keys,
                                         const int* key,
                                         int key_dim,
                                         int capacity,
                                         bool* found_existing,
                                         int value_to_insert = -1) {
  int slot = HashFuncT::hash(key, key_dim, capacity);
  int initial_slot = slot;
  int attempts = 0;
  *found_existing = false;

  while (attempts < capacity) {
    if (value_to_insert != -1) {
      // Use 64-bit atomicCAS to claim and set value atomically.
      // This prevents the race condition where a slot is claimed but value is not yet written,
      // which causes other threads in the same warp to deadlock while waiting.
      unsigned long long* slot_addr_64 = (unsigned long long*)&table_kvs[slot * 2];
      unsigned long long old_val_64 = *slot_addr_64;

      // Extract the first int (slot marker) from the 64-bit value.
      // Assuming Little Endian (standard for NVIDIA GPUs).
      int slot_marker = (int)(old_val_64 & 0xFFFFFFFF);

      if (slot_marker == -1) {
        // Attempt to claim
        unsigned long long new_val_64 =
            ((unsigned long long)value_to_insert << 32) | (unsigned int)slot;
        unsigned long long prev_64 = atomicCAS(slot_addr_64, old_val_64, new_val_64);

        if (prev_64 == old_val_64) {
          return slot;  // Successfully claimed and written
        }
        // CAS failed, reload value
        old_val_64 = prev_64;
        slot_marker = (int)(old_val_64 & 0xFFFFFFFF);
      }

      // If we are here, the slot is occupied (or we failed CAS and it is now occupied)
      if (slot_marker != -1) {
        int vector_index = (int)(old_val_64 >> 32);
        if (vector_index < 0) {
          // This should not happen if all threads use 64-bit CAS, but if mixed with 32-bit CAS
          // (expand kernel), we might still hit this. Retry without advancing.
          continue;
        }

        const volatile int* existing_key = &vector_keys[vector_index * key_dim];
        if (vec_equal(existing_key, key, key_dim)) {
          *found_existing = true;
          return -1;  // Key found
        }

        // Collision
        slot = (slot + 1) & (capacity - 1);
        if (slot == initial_slot) {
          break;
        }
        attempts++;
        continue;
      }

      // If slot_marker is -1 here (CAS failed but value is still -1?), retry loop.
      continue;

    } else {
      // Original 32-bit logic for when value is not known yet (expand kernel)

      int* slot_address = &table_kvs[slot * 2];
      int prev = atomicCAS(slot_address, -1, slot);

      if (prev == -1) {
        return slot;  // Successfully claimed
      }

      // Slot occupied (or reserved by another thread just now)
      // Use volatile to ensure we read the latest value from memory
      volatile int* slot_value_ptr = &table_kvs[slot * 2 + 1];
      int vector_index = *slot_value_ptr;

      if (vector_index < 0) {
        // Another thread is writing to this slot. Retry without advancing.
        continue;
      }

      const volatile int* existing_key = &vector_keys[vector_index * key_dim];
      if (vec_equal(existing_key, key, key_dim)) {
        *found_existing = true;
        return -1;  // Key found
      }

      // Collision with different key
      slot = (slot + 1) & (capacity - 1);
      if (slot == initial_slot) {
        break;
      }
      attempts++;
    }
  }
  return -1;  // Table full or not found (and couldn't claim)
}

template <typename HashFuncT>
__device__ inline void insert_candidate_if_absent(int* table_kvs,
                                                  int* vector_keys,
                                                  const int* candidate_key,
                                                  int key_dim,
                                                  int table_capacity,
                                                  int vector_capacity,
                                                  int* num_entries_ptr,
                                                  int* status_ptr) {
  bool found = false;
  int slot = claim_slot_or_find<HashFuncT>(
      table_kvs, vector_keys, candidate_key, key_dim, table_capacity, &found);

  if (found) {
    return;  // Already present
  }

  if (slot == -1) {
    // Table full
    set_expand_status(status_ptr, kExpandStatusTableFull);
    return;
  }

  // Slot reserved, now safely allocate index
  int old_count, new_count;
  do {
    old_count = *num_entries_ptr;
    if (old_count >= vector_capacity) {
      // Vector full, roll back slot reservation
      atomicExch(&table_kvs[slot * 2], -1);
      table_kvs[slot * 2 + 1] = -1;
      set_expand_status(status_ptr, kExpandStatusVectorOverflow);
      return;
    }
    new_count = old_count + 1;
  } while (atomicCAS(num_entries_ptr, old_count, new_count) != old_count);

  int new_index = old_count;

  int* dst = &vector_keys[new_index * key_dim];
  for (int d = 0; d < key_dim; ++d) {
    dst[d] = candidate_key[d];
  }

  __threadfence();
  table_kvs[slot * 2 + 1] = new_index;
}

// --- Device Function for Hash Table Search ---
template <typename HashFuncT>
__device__ inline int search_hash_table(
    const int* __restrict__ table_kvs,    // Pointer to key-value store
    const int* __restrict__ vector_keys,  // Pointer to original vector keys
    const int* __restrict__ query_key,    // Pointer to the key being searched
    int key_dim,                          // Dimension of keys
    int table_capacity) {                 // Capacity of the hash table (must be power of 2)

  int slot = HashFuncT::hash(query_key, key_dim, table_capacity);
  const int capacity_mask = table_capacity - 1;
  int initial_slot = slot;
  int attempts = 0;

  while (attempts < table_capacity) {
    // Single 64-bit load for both slot marker and vector index.
    // table_kvs layout: [marker, vector_index] pairs, naturally 8-byte aligned.
    long long pair = *reinterpret_cast<const long long*>(&table_kvs[slot * 2]);
    int slot_marker = (int)(pair & 0xFFFFFFFF);
    int vector_index = (int)(pair >> 32);

    if (slot_marker == -1) {
      return -1;
    }

    if (vector_index >= 0) {
      // Use vectorized 128-bit comparison for the common 4D key case
      bool keys_match;
      if (key_dim == 4) {
        keys_match = vec_equal_4d(&vector_keys[vector_index * 4], query_key);
      } else {
        keys_match = vec_equal(&vector_keys[vector_index * key_dim], query_key, key_dim);
      }
      if (keys_match) {
        return vector_index;
      }
    }

    slot = (slot + 1) & capacity_mask;
    if (slot == initial_slot) {
      return -1;
    }
    attempts++;
  }
  return -1;
}

// --- Kernels ---

// Kernel to initialize the hash table slots to -1
extern "C" __global__ void prepare_key_value_pairs_kernel(int* table_kvs, int capacity) {
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  if (tid < capacity) {
    table_kvs[2 * tid + 0] = -1;  // Key hash slot indicator
    table_kvs[2 * tid + 1] = -1;  // Value (index into vector_keys)
  }
}

// --- Templated Insert Kernel ---
template <typename HashFuncT>
__device__ void insert_kernel_templated(
    int* table_kvs, const int* vector_keys, int num_keys, int key_dim, int table_capacity) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= num_keys) {
    return;
  }

  const int* key_to_insert = &vector_keys[idx * key_dim];
  bool found = false;
  int slot = claim_slot_or_find<HashFuncT>(
      table_kvs, vector_keys, key_to_insert, key_dim, table_capacity, &found, idx);

  if (found) {
    return;  // Already present (deduplication)
  }

  // If slot != -1, value was already written by claim_slot_or_find

  // If slot == -1 and !found, table is full (fail silently as per original logic, or could add
  // error handling)
}

// --- Templated Search Kernel ---
template <typename HashFuncT>
__device__ void search_kernel_templated(const int* __restrict__ table_kvs,
                                        const int* __restrict__ vector_keys,
                                        const int* __restrict__ search_keys,
                                        int* __restrict__ results,
                                        int num_search_keys,
                                        int key_dim,
                                        int table_capacity) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= num_search_keys) {
    return;
  }

  const int* query_key = &search_keys[idx * key_dim];
  results[idx] =
      search_hash_table<HashFuncT>(table_kvs, vector_keys, query_key, key_dim, table_capacity);
}

template <typename HashFuncT>
__device__ void expand_insert_kernel_templated(int* table_kvs,
                                               int* vector_keys,
                                               const int* __restrict__ base_coords,
                                               const int* __restrict__ offsets,
                                               int num_base_coords,
                                               int num_offsets,
                                               int key_dim,
                                               int table_capacity,
                                               int vector_capacity,
                                               int* num_entries_ptr,
                                               int* status_ptr) {
  extern __shared__ int shared_candidates[];

  // Shared memory probe cache: direct-mapped, 256 entries.
  // Placed after the shared_candidates region.
  // Layout: shared_candidates[0 .. blockDim.x * key_dim - 1] = candidate keys
  //         smem_cache[0..255] = cached 64-bit slot pairs
  //         smem_cache_tags[0..255] = tag (slot index) or -1 if invalid
  const int CACHE_SIZE = 256;
  const int CACHE_MASK = CACHE_SIZE - 1;
  long long* smem_cache = reinterpret_cast<long long*>(&shared_candidates[blockDim.x * key_dim]);
  int* smem_cache_tags = reinterpret_cast<int*>(&smem_cache[CACHE_SIZE]);

  // Initialize cache as invalid
  for (int i = threadIdx.x; i < CACHE_SIZE; i += blockDim.x) {
    smem_cache_tags[i] = -1;
  }
  __syncthreads();

  long long global_idx = blockIdx.x * static_cast<long long>(blockDim.x) + threadIdx.x;
  long long total = static_cast<long long>(num_base_coords) * static_cast<long long>(num_offsets);
  if (global_idx >= total) {
    return;
  }

  int offset_idx = static_cast<int>(global_idx / num_base_coords);
  int base_idx = static_cast<int>(global_idx % num_base_coords);

  int* candidate = &shared_candidates[threadIdx.x * key_dim];
  const int* base_ptr = &base_coords[base_idx * key_dim];
  const int* offset_ptr = &offsets[offset_idx * key_dim];
  for (int d = 0; d < key_dim; ++d) {
    candidate[d] = base_ptr[d] + offset_ptr[d];
  }

  // --- Search phase with shared memory cache before insert ---
  // Check if the candidate key already exists using the cache to reduce
  // redundant global memory reads within the same block.
  int slot = HashFuncT::hash(candidate, key_dim, table_capacity);
  const int capacity_mask = table_capacity - 1;
  int initial_slot = slot;
  int attempts = 0;
  bool found_existing = false;

  while (attempts < table_capacity) {
    long long pair;
    int cache_idx = slot & CACHE_MASK;

    // Check shared memory cache (read-only for search phase)
    if (smem_cache_tags[cache_idx] == slot) {
      pair = smem_cache[cache_idx];
    } else {
      pair = *reinterpret_cast<const long long*>(&table_kvs[slot * 2]);
      // Populate cache
      smem_cache[cache_idx] = pair;
      smem_cache_tags[cache_idx] = slot;
    }

    int slot_marker = (int)(pair & 0xFFFFFFFF);
    int vector_index = (int)(pair >> 32);

    if (slot_marker == -1) {
      // Empty slot — key is not present, proceed to insert
      break;
    }

    if (vector_index >= 0) {
      bool keys_match;
      if (key_dim == 4) {
        keys_match = vec_equal_4d(&vector_keys[vector_index * 4], candidate);
      } else {
        keys_match = vec_equal(&vector_keys[vector_index * key_dim], candidate, key_dim);
      }
      if (keys_match) {
        found_existing = true;
        break;
      }
    }

    slot = (slot + 1) & capacity_mask;
    if (slot == initial_slot) {
      break;
    }
    attempts++;
  }

  if (found_existing) {
    return;  // Already present, skip insert
  }

  // --- Insert phase: always go to global memory, invalidate cache entry ---
  insert_candidate_if_absent<HashFuncT>(table_kvs,
                                        vector_keys,
                                        candidate,
                                        key_dim,
                                        table_capacity,
                                        vector_capacity,
                                        num_entries_ptr,
                                        status_ptr);

  // Invalidate cache entries that may have been modified by the insert.
  // The insert probes from the hash slot, so invalidate the initial slot's
  // cache entry to avoid stale reads by other threads in this block.
  int insert_slot = HashFuncT::hash(candidate, key_dim, table_capacity);
  int inv_cache_idx = insert_slot & CACHE_MASK;
  smem_cache_tags[inv_cache_idx] = -1;
}

// --- Warp-Cooperative Search Kernel ---
// Each WARP (32 threads) handles one query key, probing 32 consecutive slots
// simultaneously per iteration. This improves hash table search throughput by
// reducing the number of serial probe iterations.
template <typename HashFuncT>
__device__ void warp_cooperative_search_kernel_templated(const int* __restrict__ table_kvs,
                                                         const int* __restrict__ vector_keys,
                                                         const int* __restrict__ search_keys,
                                                         int* __restrict__ results,
                                                         int num_search_keys,
                                                         int key_dim,
                                                         int table_capacity) {
  // Each WARP handles one query key
  int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
  int lane = threadIdx.x % 32;

  if (warp_id >= num_search_keys) return;

  const int* query_key = &search_keys[warp_id * key_dim];
  int slot = HashFuncT::hash(query_key, key_dim, table_capacity);
  const int capacity_mask = table_capacity - 1;

  int result = -1;

  // Each iteration, 32 lanes probe 32 consecutive slots
  for (int attempt = 0; attempt < table_capacity; attempt += 32) {
    int my_slot = (slot + lane + attempt) & capacity_mask;

    // Load slot entry (64-bit load for both marker and index)
    long long pair = *reinterpret_cast<const long long*>(&table_kvs[my_slot * 2]);
    int slot_marker = (int)(pair & 0xFFFFFFFF);
    int vector_index = (int)(pair >> 32);

    // Check for empty slot — if any lane hits empty before finding key, key is absent
    unsigned empty_mask = __ballot_sync(0xFFFFFFFF, slot_marker == -1);

    // Check for key match
    bool match = false;
    if (slot_marker != -1 && vector_index >= 0) {
      if (key_dim == 4) {
        match = vec_equal_4d(&vector_keys[vector_index * 4], query_key);
      } else {
        match = vec_equal(&vector_keys[vector_index * key_dim], query_key, key_dim);
      }
    }
    unsigned match_mask = __ballot_sync(0xFFFFFFFF, match);

    if (match_mask != 0) {
      // Found! Get result from the matching lane
      int winner_lane = __ffs(match_mask) - 1;
      result = __shfl_sync(0xFFFFFFFF, vector_index, winner_lane);
      break;
    }

    if (empty_mask != 0) {
      // Found an empty slot in the probe sequence.
      // With linear probing, if we've checked all slots before the
      // first empty and found no match, the key is absent.
      //
      // Lanes are ordered 0..31 and all share the same attempt offset,
      // so lane ordering IS probe ordering within this batch.
      break;
    }
    // All 32 slots occupied with non-matching keys — continue to next batch
  }

  if (lane == 0) {
    results[warp_id] = result;
  }
}

// --- Extern "C" Wrappers for CuPy ---

// Insert Wrappers
extern "C" __global__ void insert_kernel_fnv1a(
    int* table_kvs, const int* vector_keys, int num_keys, int key_dim, int table_capacity) {
  insert_kernel_templated<FNV1AHash>(table_kvs, vector_keys, num_keys, key_dim, table_capacity);
}
extern "C" __global__ void insert_kernel_city(
    int* table_kvs, const int* vector_keys, int num_keys, int key_dim, int table_capacity) {
  insert_kernel_templated<CityHash>(table_kvs, vector_keys, num_keys, key_dim, table_capacity);
}
extern "C" __global__ void insert_kernel_murmur(
    int* table_kvs, const int* vector_keys, int num_keys, int key_dim, int table_capacity) {
  insert_kernel_templated<MurmurHash>(table_kvs, vector_keys, num_keys, key_dim, table_capacity);
}

// Search Wrappers
extern "C" __global__ void search_kernel_fnv1a(const int* table_kvs,
                                               const int* vector_keys,
                                               const int* search_keys,
                                               int* results,
                                               int num_search_keys,
                                               int key_dim,
                                               int table_capacity) {
  search_kernel_templated<FNV1AHash>(
      table_kvs, vector_keys, search_keys, results, num_search_keys, key_dim, table_capacity);
}
extern "C" __global__ void search_kernel_city(const int* table_kvs,
                                              const int* vector_keys,
                                              const int* search_keys,
                                              int* results,
                                              int num_search_keys,
                                              int key_dim,
                                              int table_capacity) {
  search_kernel_templated<CityHash>(
      table_kvs, vector_keys, search_keys, results, num_search_keys, key_dim, table_capacity);
}
extern "C" __global__ void search_kernel_murmur(const int* table_kvs,
                                                const int* vector_keys,
                                                const int* search_keys,
                                                int* results,
                                                int num_search_keys,
                                                int key_dim,
                                                int table_capacity) {
  search_kernel_templated<MurmurHash>(
      table_kvs, vector_keys, search_keys, results, num_search_keys, key_dim, table_capacity);
}

// Warp-Cooperative Search Wrappers
extern "C" __global__ void warp_search_kernel_fnv1a(const int* table_kvs,
                                                    const int* vector_keys,
                                                    const int* search_keys,
                                                    int* results,
                                                    int num_search_keys,
                                                    int key_dim,
                                                    int table_capacity) {
  warp_cooperative_search_kernel_templated<FNV1AHash>(
      table_kvs, vector_keys, search_keys, results, num_search_keys, key_dim, table_capacity);
}
extern "C" __global__ void warp_search_kernel_city(const int* table_kvs,
                                                   const int* vector_keys,
                                                   const int* search_keys,
                                                   int* results,
                                                   int num_search_keys,
                                                   int key_dim,
                                                   int table_capacity) {
  warp_cooperative_search_kernel_templated<CityHash>(
      table_kvs, vector_keys, search_keys, results, num_search_keys, key_dim, table_capacity);
}
extern "C" __global__ void warp_search_kernel_murmur(const int* table_kvs,
                                                     const int* vector_keys,
                                                     const int* search_keys,
                                                     int* results,
                                                     int num_search_keys,
                                                     int key_dim,
                                                     int table_capacity) {
  warp_cooperative_search_kernel_templated<MurmurHash>(
      table_kvs, vector_keys, search_keys, results, num_search_keys, key_dim, table_capacity);
}

extern "C" __global__ void expand_insert_kernel_fnv1a(int* table_kvs,
                                                      int* vector_keys,
                                                      const int* base_coords,
                                                      const int* offsets,
                                                      int num_base_coords,
                                                      int num_offsets,
                                                      int key_dim,
                                                      int table_capacity,
                                                      int vector_capacity,
                                                      int* num_entries_ptr,
                                                      int* status_ptr) {
  expand_insert_kernel_templated<FNV1AHash>(table_kvs,
                                            vector_keys,
                                            base_coords,
                                            offsets,
                                            num_base_coords,
                                            num_offsets,
                                            key_dim,
                                            table_capacity,
                                            vector_capacity,
                                            num_entries_ptr,
                                            status_ptr);
}

extern "C" __global__ void expand_insert_kernel_city(int* table_kvs,
                                                     int* vector_keys,
                                                     const int* base_coords,
                                                     const int* offsets,
                                                     int num_base_coords,
                                                     int num_offsets,
                                                     int key_dim,
                                                     int table_capacity,
                                                     int vector_capacity,
                                                     int* num_entries_ptr,
                                                     int* status_ptr) {
  expand_insert_kernel_templated<CityHash>(table_kvs,
                                           vector_keys,
                                           base_coords,
                                           offsets,
                                           num_base_coords,
                                           num_offsets,
                                           key_dim,
                                           table_capacity,
                                           vector_capacity,
                                           num_entries_ptr,
                                           status_ptr);
}

extern "C" __global__ void expand_insert_kernel_murmur(int* table_kvs,
                                                       int* vector_keys,
                                                       const int* base_coords,
                                                       const int* offsets,
                                                       int num_base_coords,
                                                       int num_offsets,
                                                       int key_dim,
                                                       int table_capacity,
                                                       int vector_capacity,
                                                       int* num_entries_ptr,
                                                       int* status_ptr) {
  expand_insert_kernel_templated<MurmurHash>(table_kvs,
                                             vector_keys,
                                             base_coords,
                                             offsets,
                                             num_base_coords,
                                             num_offsets,
                                             key_dim,
                                             table_capacity,
                                             vector_capacity,
                                             num_entries_ptr,
                                             status_ptr);
}
