HashmapCUDAImpl.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#include "open3d/core/hashmap/CUDA/HashmapBufferCUDA.cuh"
#include "open3d/core/hashmap/CUDA/InternalNodeManager.h"
#include "open3d/core/hashmap/CUDA/Macros.h"
#include "open3d/core/hashmap/CUDA/Traits.h"
#include "open3d/core/hashmap/DeviceHashmap.h"

namespace open3d {
namespace core {

template <typename Hash, typename KeyEq>
class CUDAHashmapImplContext {
public:
    CUDAHashmapImplContext();

    __host__ void Setup(int64_t init_buckets,
                        int64_t init_capacity,
                        int64_t dsize_key,
                        int64_t dsize_value,
                        const InternalNodeManagerContext& node_mgr_ctx,
                        const CUDAHashmapBufferContext& kv_mgr_ctx);

    __device__ bool Insert(bool lane_active,
                           uint32_t lane_id,
                           uint32_t bucket_id,
                           const void* key_ptr,
                           addr_t iterator_addr);

    __device__ Pair<addr_t, bool> Find(bool lane_active,
                                       uint32_t lane_id,
                                       uint32_t bucket_id,
                                       const void* key_ptr);

    __device__ Pair<addr_t, bool> Erase(bool lane_active,
                                        uint32_t lane_id,
                                        uint32_t bucket_id,
                                        const void* key_ptr);

    __device__ void WarpSyncKey(const void* key_ptr,
                                uint32_t lane_id,
                                void* ret_key_ptr);
    __device__ int32_t WarpFindKey(const void* src_key_ptr,
                                   uint32_t lane_id,
                                   addr_t ptr);
    __device__ int32_t WarpFindEmpty(addr_t unit_data);

    // Hash function.
    __device__ int64_t ComputeBucket(const void* key_ptr) const;

    // Node manager.
    __device__ addr_t AllocateSlab(uint32_t lane_id);
    __device__ void FreeSlab(addr_t slab_ptr);

    // Helpers.
    __device__ addr_t* get_unit_ptr_from_list_nodes(addr_t slab_ptr,
                                                    uint32_t lane_id) {
        return node_mgr_ctx_.get_unit_ptr_from_slab(slab_ptr, lane_id);
    }
    __device__ addr_t* get_unit_ptr_from_list_head(uint32_t bucket_id,
                                                   uint32_t lane_id) {
        return reinterpret_cast<uint32_t*>(bucket_list_head_) +
               bucket_id * kWarpSize + lane_id;
    }

public:
    Hash hash_fn_;
    KeyEq cmp_fn_;

    int64_t bucket_count_;
    int64_t capacity_;
    int64_t dsize_key_;
    int64_t dsize_value_;

    Slab* bucket_list_head_;
    InternalNodeManagerContext node_mgr_ctx_;
    CUDAHashmapBufferContext kv_mgr_ctx_;
};

template <typename Hash, typename KeyEq>
__global__ void InsertKernelPass0(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                  const void* input_keys,
                                  addr_t* output_addrs,
                                  int heap_counter_prev,
                                  int64_t count);

template <typename Hash, typename KeyEq>
__global__ void InsertKernelPass1(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                  const void* input_keys,
                                  addr_t* output_addrs,
                                  bool* output_masks,
                                  int64_t count);

template <typename Hash, typename KeyEq>
__global__ void InsertKernelPass2(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                  const void* input_values,
                                  addr_t* output_addrs,
                                  bool* output_masks,
                                  int64_t count);

template <typename Hash, typename KeyEq>
__global__ void FindKernel(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                           const void* input_keys,
                           addr_t* output_addrs,
                           bool* output_masks,
                           int64_t count);

template <typename Hash, typename KeyEq>
__global__ void EraseKernelPass0(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                 const void* input_keys,
                                 addr_t* output_addrs,
                                 bool* output_masks,
                                 int64_t count);

template <typename Hash, typename KeyEq>
__global__ void EraseKernelPass1(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                 addr_t* output_addrs,
                                 bool* output_masks,
                                 int64_t count);

template <typename Hash, typename KeyEq>
__global__ void GetActiveIndicesKernel(
        CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
        addr_t* output_addrs,
        uint32_t* output_iterator_count);

template <typename Hash, typename KeyEq>
__global__ void CountElemsPerBucketKernel(
        CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
        int64_t* bucket_elem_counts);

template <typename Hash, typename KeyEq>
CUDAHashmapImplContext<Hash, KeyEq>::CUDAHashmapImplContext()
    : bucket_count_(0), bucket_list_head_(nullptr) {}

template <typename Hash, typename KeyEq>
void CUDAHashmapImplContext<Hash, KeyEq>::Setup(
        int64_t init_buckets,
        int64_t init_capacity,
        int64_t dsize_key,
        int64_t dsize_value,
        const InternalNodeManagerContext& allocator_ctx,
        const CUDAHashmapBufferContext& pair_allocator_ctx) {
    bucket_count_ = init_buckets;
    capacity_ = init_capacity;
    dsize_key_ = dsize_key;
    dsize_value_ = dsize_value;

    node_mgr_ctx_ = allocator_ctx;
    kv_mgr_ctx_ = pair_allocator_ctx;

    hash_fn_.key_size_in_int_ = dsize_key / sizeof(int);
    cmp_fn_.key_size_in_int_ = dsize_key / sizeof(int);
}

template <typename Hash, typename KeyEq>
__device__ bool CUDAHashmapImplContext<Hash, KeyEq>::Insert(
        bool lane_active,
        uint32_t lane_id,
        uint32_t bucket_id,
        const void* key,
        addr_t iterator_addr) {
    uint32_t work_queue = 0;
    uint32_t prev_work_queue = 0;
    uint32_t curr_slab_ptr = kHeadSlabAddr;
    uint8_t src_key[kMaxKeyByteSize];

    bool mask = false;

    // > Loop when we have active lanes
    while ((work_queue = __ballot_sync(kSyncLanesMask, lane_active))) {
        // 0. Restart from linked list head if last insertion is finished
        curr_slab_ptr =
                (prev_work_queue != work_queue) ? kHeadSlabAddr : curr_slab_ptr;
        uint32_t src_lane = __ffs(work_queue) - 1;
        uint32_t src_bucket =
                __shfl_sync(kSyncLanesMask, bucket_id, src_lane, kWarpSize);

        WarpSyncKey(key, src_lane, src_key);

        // Each lane in the warp reads a unit in the slab
        uint32_t unit_data =
                (curr_slab_ptr == kHeadSlabAddr)
                        ? *(get_unit_ptr_from_list_head(src_bucket, lane_id))
                        : *(get_unit_ptr_from_list_nodes(curr_slab_ptr,
                                                         lane_id));

        int32_t lane_found = WarpFindKey(src_key, lane_id, unit_data);
        int32_t lane_empty = WarpFindEmpty(unit_data);

        // Branch 1: key already existing, ABORT
        if (lane_found >= 0) {
            if (lane_id == src_lane) {
                // free memory heap
                lane_active = false;
            }
        }

        // Branch 2: empty slot available, try to insert
        else if (lane_empty >= 0) {
            if (lane_id == src_lane) {
                // TODO: check why we cannot put malloc here
                const uint32_t* unit_data_ptr =
                        (curr_slab_ptr == kHeadSlabAddr)
                                ? get_unit_ptr_from_list_head(src_bucket,
                                                              lane_empty)
                                : get_unit_ptr_from_list_nodes(curr_slab_ptr,
                                                               lane_empty);

                addr_t old_iterator_addr =
                        atomicCAS((unsigned int*)unit_data_ptr, kEmptyNodeAddr,
                                  iterator_addr);

                // Remember to clean up in another pass
                // Branch 2.1: SUCCEED
                if (old_iterator_addr == kEmptyNodeAddr) {
                    lane_active = false;
                    mask = true;
                }
                // Branch 2.2: failed: RESTART
                // In the consequent attempt,
                // > if the same key was inserted in this slot,
                //   we fall back to Branch 1;
                // > if a different key was inserted,
                //   we go to Branch 2 or 3.
            }
        }

        // Branch 3: nothing found in this slab, goto next slab
        else {
            // broadcast next slab
            addr_t next_slab_ptr = __shfl_sync(kSyncLanesMask, unit_data,
                                               kNextSlabPtrLaneId, kWarpSize);

            // Branch 3.1: next slab existing, RESTART this lane
            if (next_slab_ptr != kEmptySlabAddr) {
                curr_slab_ptr = next_slab_ptr;
            }

            // Branch 3.2: next slab empty, try to allocate one
            else {
                addr_t new_next_slab_ptr = AllocateSlab(lane_id);

                if (lane_id == kNextSlabPtrLaneId) {
                    const uint32_t* unit_data_ptr =
                            (curr_slab_ptr == kHeadSlabAddr)
                                    ? get_unit_ptr_from_list_head(
                                              src_bucket, kNextSlabPtrLaneId)
                                    : get_unit_ptr_from_list_nodes(
                                              curr_slab_ptr,
                                              kNextSlabPtrLaneId);

                    addr_t old_next_slab_ptr =
                            atomicCAS((unsigned int*)unit_data_ptr,
                                      kEmptySlabAddr, new_next_slab_ptr);

                    // Branch 3.2.1: other thread allocated, RESTART lane. In
                    // the consequent attempt, goto Branch 2'
                    if (old_next_slab_ptr != kEmptySlabAddr) {
                        FreeSlab(new_next_slab_ptr);
                    }
                    // Branch 3.2.2: this thread allocated, RESTART lane, 'goto
                    // Branch 2'
                }
            }
        }

        prev_work_queue = work_queue;
    }

    return mask;
}

template <typename Hash, typename KeyEq>
__device__ Pair<addr_t, bool> CUDAHashmapImplContext<Hash, KeyEq>::Find(
        bool lane_active,
        uint32_t lane_id,
        uint32_t bucket_id,
        const void* query_key) {
    uint32_t work_queue = 0;
    uint32_t prev_work_queue = work_queue;
    uint32_t curr_slab_ptr = kHeadSlabAddr;

    addr_t iterator = kNullAddr;
    bool mask = false;

    // > Loop when we have active lanes.
    while ((work_queue = __ballot_sync(kSyncLanesMask, lane_active))) {
        // 0. Restart from linked list head if the last query is finished.
        curr_slab_ptr =
                (prev_work_queue != work_queue) ? kHeadSlabAddr : curr_slab_ptr;
        uint32_t src_lane = __ffs(work_queue) - 1;
        uint32_t src_bucket =
                __shfl_sync(kSyncLanesMask, bucket_id, src_lane, kWarpSize);

        uint8_t src_key[kMaxKeyByteSize];
        WarpSyncKey(query_key, src_lane, src_key);

        // Each lane in the warp reads a unit in the slab in parallel.
        const uint32_t unit_data =
                (curr_slab_ptr == kHeadSlabAddr)
                        ? *(get_unit_ptr_from_list_head(src_bucket, lane_id))
                        : *(get_unit_ptr_from_list_nodes(curr_slab_ptr,
                                                         lane_id));

        int32_t lane_found = WarpFindKey(src_key, lane_id, unit_data);

        // 1. Found in this slab, SUCCEED.
        if (lane_found >= 0) {
            // broadcast found value
            addr_t found_pair_internal_ptr = __shfl_sync(
                    kSyncLanesMask, unit_data, lane_found, kWarpSize);

            if (lane_id == src_lane) {
                lane_active = false;

                // Actually iterator_addr
                iterator = found_pair_internal_ptr;
                mask = true;
            }
        }

        // 2. Not found in this slab.
        else {
            // Broadcast next slab: lane 31 reads 'next'.
            addr_t next_slab_ptr = __shfl_sync(kSyncLanesMask, unit_data,
                                               kNextSlabPtrLaneId, kWarpSize);

            // 2.1. Next slab is empty, ABORT.
            if (next_slab_ptr == kEmptySlabAddr) {
                if (lane_id == src_lane) {
                    lane_active = false;
                }
            }
            // 2.2. Next slab exists, RESTART.
            else {
                curr_slab_ptr = next_slab_ptr;
            }
        }

        prev_work_queue = work_queue;
    }

    return make_pair(iterator, mask);
}

template <typename Hash, typename KeyEq>
__device__ Pair<addr_t, bool> CUDAHashmapImplContext<Hash, KeyEq>::Erase(
        bool lane_active,
        uint32_t lane_id,
        uint32_t bucket_id,
        const void* key) {
    uint32_t work_queue = 0;
    uint32_t prev_work_queue = 0;
    uint32_t curr_slab_ptr = kHeadSlabAddr;
    uint8_t src_key[kMaxKeyByteSize];

    addr_t iterator_addr = 0;
    bool mask = false;

    // > Loop when we have active lanes.
    while ((work_queue = __ballot_sync(kSyncLanesMask, lane_active))) {
        // 0. Restart from linked list head if last insertion is finished.
        curr_slab_ptr =
                (prev_work_queue != work_queue) ? kHeadSlabAddr : curr_slab_ptr;
        uint32_t src_lane = __ffs(work_queue) - 1;
        uint32_t src_bucket =
                __shfl_sync(kSyncLanesMask, bucket_id, src_lane, kWarpSize);

        WarpSyncKey(key, src_lane, src_key);

        const uint32_t unit_data =
                (curr_slab_ptr == kHeadSlabAddr)
                        ? *(get_unit_ptr_from_list_head(src_bucket, lane_id))
                        : *(get_unit_ptr_from_list_nodes(curr_slab_ptr,
                                                         lane_id));

        int32_t lane_found = WarpFindKey(src_key, lane_id, unit_data);

        // Branch 1: key found.
        if (lane_found >= 0) {
            if (lane_id == src_lane) {
                uint32_t* unit_data_ptr =
                        (curr_slab_ptr == kHeadSlabAddr)
                                ? get_unit_ptr_from_list_head(src_bucket,
                                                              lane_found)
                                : get_unit_ptr_from_list_nodes(curr_slab_ptr,
                                                               lane_found);

                uint32_t pair_to_delete = atomicExch(
                        (unsigned int*)unit_data_ptr, kEmptyNodeAddr);
                mask = pair_to_delete != kEmptyNodeAddr;
                iterator_addr = pair_to_delete;
                // Branch 1.2: other thread did the job, avoid double free
            }
        } else {  // no matching slot found:
            addr_t next_slab_ptr = __shfl_sync(kSyncLanesMask, unit_data,
                                               kNextSlabPtrLaneId, kWarpSize);
            if (next_slab_ptr == kEmptySlabAddr) {
                // not found:
                if (lane_id == src_lane) {
                    lane_active = false;
                }
            } else {
                curr_slab_ptr = next_slab_ptr;
            }
        }
        prev_work_queue = work_queue;
    }

    return make_pair(iterator_addr, mask);
}

template <typename Hash, typename KeyEq>
__device__ void CUDAHashmapImplContext<Hash, KeyEq>::WarpSyncKey(
        const void* key_ptr, uint32_t lane_id, void* ret_key_ptr) {
    auto dst_key_ptr = static_cast<int*>(ret_key_ptr);
    auto src_key_ptr = static_cast<const int*>(key_ptr);
    for (int i = 0; i < hash_fn_.key_size_in_int_; ++i) {
        dst_key_ptr[i] =
                __shfl_sync(kSyncLanesMask, src_key_ptr[i], lane_id, kWarpSize);
    }
}

template <typename Hash, typename KeyEq>
__device__ int32_t CUDAHashmapImplContext<Hash, KeyEq>::WarpFindKey(
        const void* key_ptr, uint32_t lane_id, addr_t ptr) {
    bool is_lane_found =
            // Select key lanes.
            ((1 << lane_id) & kNodePtrLanesMask)
            // Validate key addrs.
            && (ptr != kEmptyNodeAddr)
            // Find keys in memory heap.
            && cmp_fn_(kv_mgr_ctx_.ExtractIterator(ptr).first, key_ptr);

    return __ffs(__ballot_sync(kNodePtrLanesMask, is_lane_found)) - 1;
}

template <typename Hash, typename KeyEq>
__device__ int32_t
CUDAHashmapImplContext<Hash, KeyEq>::WarpFindEmpty(addr_t ptr) {
    bool is_lane_empty = (ptr == kEmptyNodeAddr);
    return __ffs(__ballot_sync(kNodePtrLanesMask, is_lane_empty)) - 1;
}

template <typename Hash, typename KeyEq>
__device__ int64_t
CUDAHashmapImplContext<Hash, KeyEq>::ComputeBucket(const void* key) const {
    return hash_fn_(key) % bucket_count_;
}

template <typename Hash, typename KeyEq>
__device__ addr_t
CUDAHashmapImplContext<Hash, KeyEq>::AllocateSlab(uint32_t lane_id) {
    return node_mgr_ctx_.WarpAllocate(lane_id);
}

template <typename Hash, typename KeyEq>
__device__ __forceinline__ void CUDAHashmapImplContext<Hash, KeyEq>::FreeSlab(
        addr_t slab_ptr) {
    node_mgr_ctx_.FreeUntouched(slab_ptr);
}

template <typename Hash, typename KeyEq>
__global__ void InsertKernelPass0(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                  const void* input_keys,
                                  addr_t* output_addrs,
                                  int heap_counter_prev,
                                  int64_t count) {
    uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;

    if (tid < count) {
        // First write ALL input_keys to avoid potential thread conflicts.
        addr_t iterator_addr =
                hash_ctx.kv_mgr_ctx_.heap_[heap_counter_prev + tid];
        iterator_t iterator =
                hash_ctx.kv_mgr_ctx_.ExtractIterator(iterator_addr);

        MEMCPY_AS_INTS(iterator.first,
                       static_cast<const uint8_t*>(input_keys) +
                               tid * hash_ctx.dsize_key_,
                       hash_ctx.dsize_key_);
        output_addrs[tid] = iterator_addr;
    }
}

template <typename Hash, typename KeyEq>
__global__ void InsertKernelPass1(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                  const void* input_keys,
                                  addr_t* output_addrs,
                                  bool* output_masks,
                                  int64_t count) {
    uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
    uint32_t lane_id = tid & 0x1F;

    if (tid - lane_id >= count) {
        return;
    }

    hash_ctx.node_mgr_ctx_.Init(tid, lane_id);

    bool lane_active = false;
    uint32_t bucket_id = 0;
    addr_t iterator_addr = 0;

    // Dummy.
    uint8_t dummy_key[kMaxKeyByteSize];
    const void* key = reinterpret_cast<const void*>(dummy_key);

    if (tid < count) {
        lane_active = true;
        key = static_cast<const uint8_t*>(input_keys) +
              tid * hash_ctx.dsize_key_;
        iterator_addr = output_addrs[tid];
        bucket_id = hash_ctx.ComputeBucket(key);
    }

    // Index out-of-bound threads still have to run for warp synchronization.
    bool mask = hash_ctx.Insert(lane_active, lane_id, bucket_id, key,
                                iterator_addr);

    if (tid < count) {
        output_masks[tid] = mask;
    }
}

template <typename Hash, typename KeyEq>
__global__ void InsertKernelPass2(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                  const void* input_values,
                                  addr_t* output_addrs,
                                  bool* output_masks,
                                  int64_t count) {
    uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;

    if (tid < count) {
        addr_t iterator_addr = output_addrs[tid];

        if (output_masks[tid]) {
            iterator_t iterator =
                    hash_ctx.kv_mgr_ctx_.ExtractIterator(iterator_addr);

            // Success: copy remaining input_values
            if (input_values != nullptr) {
                MEMCPY_AS_INTS(iterator.second,
                               static_cast<const uint8_t*>(input_values) +
                                       tid * hash_ctx.dsize_value_,
                               hash_ctx.dsize_value_);
            }

        } else {
            hash_ctx.kv_mgr_ctx_.DeviceFree(iterator_addr);
        }
    }
}

template <typename Hash, typename KeyEq>
__global__ void FindKernel(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                           const void* input_keys,
                           addr_t* output_addrs,
                           bool* output_masks,
                           int64_t count) {
    uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
    uint32_t lane_id = threadIdx.x & 0x1F;

    // This warp is idle.
    if ((tid - lane_id) >= count) {
        return;
    }

    // Initialize the memory allocator on each warp.
    hash_ctx.node_mgr_ctx_.Init(tid, lane_id);

    bool lane_active = false;
    uint32_t bucket_id = 0;

    // Dummy.
    uint8_t dummy_key[kMaxKeyByteSize];
    const void* key = reinterpret_cast<const void*>(dummy_key);
    Pair<addr_t, bool> result;

    if (tid < count) {
        lane_active = true;
        key = static_cast<const uint8_t*>(input_keys) +
              tid * hash_ctx.dsize_key_;
        bucket_id = hash_ctx.ComputeBucket(key);
    }

    result = hash_ctx.Find(lane_active, lane_id, bucket_id, key);

    if (tid < count) {
        output_addrs[tid] = result.first;
        output_masks[tid] = result.second;
    }
}

template <typename Hash, typename KeyEq>
__global__ void EraseKernelPass0(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                 const void* input_keys,
                                 addr_t* output_addrs,
                                 bool* output_masks,
                                 int64_t count) {
    uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
    uint32_t lane_id = threadIdx.x & 0x1F;

    if (tid - lane_id >= count) {
        return;
    }

    hash_ctx.node_mgr_ctx_.Init(tid, lane_id);

    bool lane_active = false;
    uint32_t bucket_id = 0;

    uint8_t dummy_key[kMaxKeyByteSize];
    const void* key = reinterpret_cast<const void*>(dummy_key);

    if (tid < count) {
        lane_active = true;
        key = static_cast<const uint8_t*>(input_keys) +
              tid * hash_ctx.dsize_key_;
        bucket_id = hash_ctx.ComputeBucket(key);
    }

    auto result = hash_ctx.Erase(lane_active, lane_id, bucket_id, key);

    if (tid < count) {
        output_addrs[tid] = result.first;
        output_masks[tid] = result.second;
    }
}

template <typename Hash, typename KeyEq>
__global__ void EraseKernelPass1(CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
                                 addr_t* output_addrs,
                                 bool* output_masks,
                                 int64_t count) {
    uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid < count && output_masks[tid]) {
        hash_ctx.kv_mgr_ctx_.DeviceFree(output_addrs[tid]);
    }
}

template <typename Hash, typename KeyEq>
__global__ void GetActiveIndicesKernel(
        CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
        addr_t* output_addrs,
        uint32_t* output_iterator_count) {
    uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
    uint32_t lane_id = threadIdx.x & 0x1F;

    // Assigning a warp per bucket.
    uint32_t bucket_id = tid >> 5;
    if (bucket_id >= hash_ctx.bucket_count_) {
        return;
    }

    hash_ctx.node_mgr_ctx_.Init(tid, lane_id);

    uint32_t src_unit_data =
            *hash_ctx.get_unit_ptr_from_list_head(bucket_id, lane_id);
    bool is_active = src_unit_data != kEmptyNodeAddr;

    if (is_active && ((1 << lane_id) & kNodePtrLanesMask)) {
        uint32_t index = atomicAdd(output_iterator_count, 1);
        output_addrs[index] = src_unit_data;
    }

    addr_t next = __shfl_sync(kSyncLanesMask, src_unit_data, kNextSlabPtrLaneId,
                              kWarpSize);

    // Count following nodes,
    while (next != kEmptySlabAddr) {
        src_unit_data = *hash_ctx.get_unit_ptr_from_list_nodes(next, lane_id);
        is_active = (src_unit_data != kEmptyNodeAddr);

        if (is_active && ((1 << lane_id) & kNodePtrLanesMask)) {
            uint32_t index = atomicAdd(output_iterator_count, 1);
            output_addrs[index] = src_unit_data;
        }
        next = __shfl_sync(kSyncLanesMask, src_unit_data, kNextSlabPtrLaneId,
                           kWarpSize);
    }
}

template <typename Hash, typename KeyEq>
__global__ void CountElemsPerBucketKernel(
        CUDAHashmapImplContext<Hash, KeyEq> hash_ctx,
        int64_t* bucket_elem_counts) {
    uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
    uint32_t lane_id = threadIdx.x & 0x1F;

    // Assigning a warp per bucket.
    uint32_t bucket_id = tid >> 5;
    if (bucket_id >= hash_ctx.bucket_count_) {
        return;
    }

    hash_ctx.node_mgr_ctx_.Init(tid, lane_id);

    uint32_t count = 0;

    // Count head node.
    uint32_t src_unit_data =
            *hash_ctx.get_unit_ptr_from_list_head(bucket_id, lane_id);
    count += __popc(
            __ballot_sync(kNodePtrLanesMask, src_unit_data != kEmptyNodeAddr));
    addr_t next = __shfl_sync(kSyncLanesMask, src_unit_data, kNextSlabPtrLaneId,
                              kWarpSize);

    // Count following nodes.
    while (next != kEmptySlabAddr) {
        src_unit_data = *hash_ctx.get_unit_ptr_from_list_nodes(next, lane_id);
        count += __popc(__ballot_sync(kNodePtrLanesMask,
                                      src_unit_data != kEmptyNodeAddr));
        next = __shfl_sync(kSyncLanesMask, src_unit_data, kNextSlabPtrLaneId,
                           kWarpSize);
    }

    // Write back the results.
    if (lane_id == 0) {
        bucket_elem_counts[bucket_id] = count;
    }
}

}  // namespace core
}  // namespace open3d