chore: import upstream snapshot with attribution

csrc/aio/py_lib/deepspeed_pin_tensor.cpp

@@ -0,0 +1,61 @@
+// 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;
+};