// Copyright (c) Microsoft Corporation.
// SPDX-License-Identifier: Apache-2.0

// DeepSpeed Team

/*
Functionality for managing CPU tensors occupying page-locked memory.
*/

#include "deepspeed_pin_tensor.h"

using namespace std;

deepspeed_pin_tensor_t::~deepspeed_pin_tensor_t()
{
    for (auto iter = _locked_tensors.begin(); iter != _locked_tensors.end(); ++iter) {
        munlock(iter->first, iter->second);
    }
    _locked_tensors.clear();
}

torch::Tensor deepspeed_pin_tensor_t::alloc(const size_t num_elem, const at::ScalarType& elem_type)
{
    const auto num_bytes = num_elem * elementSize(elem_type);
    auto pinned_buffer = ds_page_aligned_alloc(num_bytes, true);
    assert(nullptr != pinned_buffer);

    _locked_tensors[pinned_buffer] = num_bytes;

    auto options = torch::TensorOptions().dtype(elem_type).device(torch::kCPU);

    return at::from_blob(pinned_buffer, static_cast<long int>(num_bytes), options);
}

bool deepspeed_pin_tensor_t::free(torch::Tensor& locked_tensor)
{
    auto addr = locked_tensor.data_ptr();
    if (_locked_tensors.find(addr) != _locked_tensors.end()) {
        munlock(addr, _locked_tensors[addr]);
        _locked_tensors.erase(addr);
        return true;
    }

    return false;
}