StdAllocator.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2023 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include "open3d/core/Device.h"
#include "open3d/core/MemoryManager.h"
 
namespace open3d {
namespace core {
 
template <typename T>
class StdAllocator {
public:
    using value_type = T;
 
    StdAllocator() = default;
 
    explicit StdAllocator(int device_id) : device_id_(device_id) {}
 
    StdAllocator(const StdAllocator&) = default;
 
    StdAllocator& operator=(const StdAllocator&) = default;
 
    StdAllocator(StdAllocator&&) = default;
 
    StdAllocator& operator=(StdAllocator&&) = default;
 
    template <typename U>
    StdAllocator(const StdAllocator<U>& other) : device_id_(other.device_id_) {}
 
    T* allocate(std::size_t n) {
        void* ptr = MemoryManager::Malloc(
                n * sizeof(T), Device(Device::DeviceType::CUDA, device_id_));
        return static_cast<T*>(ptr);
    }
 
    void deallocate(T* p, std::size_t n) {
        MemoryManager::Free(p, Device(Device::DeviceType::CUDA, device_id_));
    }
 
    bool operator==(const StdAllocator& other) const {
        return device_id_ == other.device_id_;
    }
 
    bool operator!=(const StdAllocator& other) const {
        return !operator==(other);
    }
 
    int GetDeviceID() const { return device_id_; }
 
private:
    // Allow access in rebind constructor.
    template <typename T2>
    friend class StdAllocator;
 
    int device_id_;
};
 
}  // namespace core
}  // namespace open3d