HashBackendBuffer.cpp

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        CloudViewer: www.cloudViewer.org                  -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2024 www.cloudViewer.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#include "cloudViewer/core/hashmap/HashBackendBuffer.h"
 
#include "cloudViewer/core/CUDAUtils.h"
 
namespace cloudViewer {
namespace core {
 
HashBackendBuffer::HashBackendBuffer(int64_t capacity,
                                     int64_t key_dsize,
                                     std::vector<int64_t> value_dsizes,
                                     const Device &device) {
    // First compute common bytesize divisor for fast copying values.
    const std::vector<int64_t> kDivisors = {16, 12, 8, 4, 2, 1};
 
    for (const auto &divisor : kDivisors) {
        bool valid = true;
        blocks_per_element_.clear();
        for (size_t i = 0; i < value_dsizes.size(); ++i) {
            int64_t bytesize = value_dsizes[i];
            valid = valid && (bytesize % divisor == 0);
            blocks_per_element_.push_back(bytesize / divisor);
        }
        if (valid) {
            common_block_size_ = divisor;
            break;
        }
    }
 
    heap_ = Tensor({capacity}, core::UInt32, device);
 
    key_buffer_ = Tensor({capacity},
                         Dtype(Dtype::DtypeCode::Object, key_dsize, "_hash_k"),
                         device);
 
    value_buffers_.clear();
    for (size_t i = 0; i < value_dsizes.size(); ++i) {
        int64_t dsize_value = value_dsizes[i];
        Tensor value_buffer_i({capacity},
                              Dtype(Dtype::DtypeCode::Object, dsize_value,
                                    "_hash_v_" + std::to_string(i)),
                              device);
        value_buffers_.push_back(value_buffer_i);
    }
 
    // Heap top is device specific
    if (device.IsCUDA()) {
        heap_top_.cuda = Tensor({1}, Dtype::Int32, device);
    }
 
    ResetHeap();
}
 
void HashBackendBuffer::ResetHeap() {
    Device device = GetDevice();
 
    Tensor heap = GetIndexHeap();
    if (device.IsCPU()) {
        CPUResetHeap(heap);
        heap_top_.cpu = 0;
    } else if (device.IsCUDA()) {
        CUDA_CALL(CUDAResetHeap, heap);
        heap_top_.cuda.Fill<int>(0);
    }
}
 
Device HashBackendBuffer::GetDevice() const { return heap_.GetDevice(); }
 
int64_t HashBackendBuffer::GetCapacity() const { return heap_.GetLength(); }
 
int64_t HashBackendBuffer::GetKeyDsize() const {
    return key_buffer_.GetDtype().ByteSize();
}
 
std::vector<int64_t> HashBackendBuffer::GetValueDsizes() const {
    std::vector<int64_t> value_dsizes;
    for (auto &value_buffer : value_buffers_) {
        value_dsizes.push_back(value_buffer.GetDtype().ByteSize());
    }
    return value_dsizes;
}
 
int64_t HashBackendBuffer::GetCommonBlockSize() const {
    return common_block_size_;
}
 
std::vector<int64_t> HashBackendBuffer::GetValueBlocksPerElement() const {
    return blocks_per_element_;
}
 
Tensor HashBackendBuffer::GetIndexHeap() const { return heap_; }
 
HashBackendBuffer::HeapTop &HashBackendBuffer::GetHeapTop() {
    return heap_top_;
}
 
int HashBackendBuffer::GetHeapTopIndex() const {
    if (heap_.IsCUDA()) {
        return heap_top_.cuda[0].Item<int>();
    }
    return heap_top_.cpu.load();
}
 
Tensor HashBackendBuffer::GetKeyBuffer() const { return key_buffer_; }
 
std::vector<Tensor> HashBackendBuffer::GetValueBuffers() const {
    return value_buffers_;
}
 
Tensor HashBackendBuffer::GetValueBuffer(size_t i) const {
    if (i >= value_buffers_.size()) {
        utility::LogError("Value buffer index out-of-bound ({} >= {}).", i,
                          value_buffers_.size());
    }
    return value_buffers_[i];
}
 
}  // namespace core
}  // namespace cloudViewer