wehub-resource-sync
62 lines
1.8 KiB
C++
62 lines
1.8 KiB
C++
// 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);
|
|
std::free((void*)iter->first);
|
|
}
|
|
_locked_tensors.clear();
|
|
}
|
|
|
|
torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem,
|
|
const torch::TensorOptions& options)
|
|
{
|
|
const auto scalar_dtype = torch::typeMetaToScalarType(options.dtype());
|
|
const auto num_bytes = num_elem * torch::elementSize(scalar_dtype);
|
|
auto pinned_buffer = ds_page_aligned_alloc(num_bytes, true);
|
|
assert(nullptr != pinned_buffer);
|
|
|
|
_locked_tensors[pinned_buffer] = num_bytes;
|
|
|
|
return at::from_blob(pinned_buffer, static_cast<int64_t>(num_elem), options);
|
|
}
|
|
|
|
torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, const at::ScalarType& elem_type)
|
|
{
|
|
auto options = torch::TensorOptions().dtype(elem_type).device(torch::kCPU).requires_grad(false);
|
|
return alloc(num_elem, 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]);
|
|
std::free(addr);
|
|
_locked_tensors.erase(addr);
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
bool deepspeed_pin_tensor_t::is_managed(const torch::Tensor& buffer)
|
|
{
|
|
if (!buffer.is_cpu()) { return false; }
|
|
auto addr = buffer.data_ptr();
|
|
if (_locked_tensors.find(addr) != _locked_tensors.end()) { return true; }
|
|
return false;
|
|
};
|