257 lines
8.4 KiB
C++
257 lines
8.4 KiB
C++
/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
#include "tensorflow/lite/delegates/xnnpack/mmap_handle.h"
|
|
|
|
#if defined(_WIN32)
|
|
#include <io.h>
|
|
#include <windows.h>
|
|
#else
|
|
#include <sys/mman.h>
|
|
#include <unistd.h>
|
|
#endif
|
|
|
|
#if defined(__linux__) && !defined(MADV_PAGEOUT)
|
|
#define MADV_PAGEOUT 21
|
|
#endif
|
|
|
|
#include <fcntl.h>
|
|
#include <sys/stat.h>
|
|
|
|
#include <algorithm>
|
|
#include <cerrno>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <memory>
|
|
|
|
#include "tensorflow/lite/delegates/xnnpack/file_util.h"
|
|
#include "tensorflow/lite/delegates/xnnpack/macros.h"
|
|
#include "tensorflow/lite/delegates/xnnpack/windows_util.h"
|
|
|
|
namespace tflite::xnnpack {
|
|
|
|
// NOLINTBEGIN(whitespace/line_length)
|
|
#ifdef XNNPACK_CACHE_NO_MMAP_FOR_TEST
|
|
#pragma message( \
|
|
"XNNPACK_CACHE_NO_MMAP_FOR_TEST has been deprecated and doesn't do anything anymore. If you were using it, consider whether filing a bug to get your usecase fixed wouldn't be a better solution. If you really need that, the new define is XNNPACK_CACHE_NO_FILE_MAPPING_FOR_DEBUG.")
|
|
#endif
|
|
|
|
#ifdef XNNPACK_CACHE_NO_FILE_MAPPING_FOR_DEBUG
|
|
#pragma message( \
|
|
"XNNPACK_CACHE_NO_FILE_MAPPING_FOR_DEBUG is defined. This can lead to highly degraded performance.")
|
|
#endif
|
|
// NOLINTEND(whitespace/line_length)
|
|
|
|
#ifdef _WIN32
|
|
// Helper to split a value in high/low parts to pass to Windows APIs.
|
|
struct HighLow {
|
|
DWORD high;
|
|
DWORD low;
|
|
static HighLow From(uint64_t val) {
|
|
static_assert(sizeof(val) <= 2 * sizeof(DWORD),
|
|
"Value type doesn't fit in two DWORDs.");
|
|
return {static_cast<DWORD>(val >> CHAR_BIT * sizeof(DWORD)),
|
|
static_cast<DWORD>(val)};
|
|
}
|
|
};
|
|
#endif
|
|
|
|
void swap(MMapHandle& a, MMapHandle& b) {
|
|
using std::swap;
|
|
swap(a.size_, b.size_);
|
|
swap(a.offset_, b.offset_);
|
|
swap(a.offset_page_adjustment_, b.offset_page_adjustment_);
|
|
swap(a.data_, b.data_);
|
|
}
|
|
|
|
MMapHandle::~MMapHandle() { UnMap(); }
|
|
|
|
MMapHandle::MMapHandle(MMapHandle&& other) { swap(*this, other); }
|
|
|
|
MMapHandle& MMapHandle::operator=(MMapHandle&& other) {
|
|
swap(*this, other);
|
|
return *this;
|
|
}
|
|
|
|
bool MMapHandle::Map(const char* path, const size_t offset) {
|
|
return this->Map(FileDescriptor::Open(path, O_RDONLY), offset, path);
|
|
}
|
|
|
|
bool MMapHandle::Map(const FileDescriptorView& fd, const size_t offset,
|
|
const char* const path) {
|
|
this->UnMap();
|
|
const char* const safe_path = path != nullptr ? path : "[unspecified]";
|
|
|
|
XNNPACK_RETURN_CHECK(fd.IsValid(),
|
|
"cannot mmap invalid file descriptor %d ('%s').",
|
|
fd.Value(), path);
|
|
|
|
#if defined(_WIN32)
|
|
struct _stat64 file_stats;
|
|
XNNPACK_RETURN_CHECK(_fstat64(fd.Value(), &file_stats) == 0,
|
|
"could not access file stats to get size ('%s'): %s.",
|
|
safe_path, strerror(errno));
|
|
#else
|
|
struct stat file_stats;
|
|
XNNPACK_RETURN_CHECK(
|
|
fstat(fd.Value(), &file_stats) == 0,
|
|
"could not access file descriptor %d stats to get size ('%s'): %s.",
|
|
fd.Value(), safe_path, strerror(errno));
|
|
#endif
|
|
|
|
// This will reset data_ and size_ on return until it is deactivated.
|
|
ScopeGuard unmap_on_error([this] { UnMap(); });
|
|
size_ = file_stats.st_size - offset;
|
|
offset_ = offset;
|
|
#if defined(XNNPACK_CACHE_NO_FILE_MAPPING_FOR_DEBUG)
|
|
// This allocation is freed in UnMap and in the destructor.
|
|
data_ = new uint8_t[size_];
|
|
fd.SetPos(offset);
|
|
XNNPACK_RETURN_CHECK(fd.Read(data_, size_), "could not read file ('%s'): %s.",
|
|
safe_path, strerror(errno));
|
|
#elif defined(_WIN32)
|
|
HANDLE osf_handle = reinterpret_cast<HANDLE>(_get_osfhandle(fd.Value()));
|
|
XNNPACK_RETURN_CHECK(osf_handle != INVALID_HANDLE_VALUE,
|
|
"could not convert file descriptor to file handle: %s.",
|
|
strerror(errno));
|
|
|
|
file_mapping_ =
|
|
CreateFileMappingA(osf_handle, /*lpFileMappingAttributes=*/nullptr,
|
|
/*flProtect=*/PAGE_READONLY, /*dwMaximumSizeHigh=*/0,
|
|
/*dwMaximumSizeLow=*/0, /*lpName=*/nullptr);
|
|
XNNPACK_RETURN_CHECK(file_mapping_ != NULL,
|
|
"could not create a file mapping: %s",
|
|
GetLastErrorString().c_str());
|
|
|
|
SYSTEM_INFO sys_info;
|
|
GetSystemInfo(&sys_info);
|
|
|
|
offset_page_adjustment_ = offset_ % sys_info.dwAllocationGranularity;
|
|
|
|
const size_t adjusted_offset = offset_ - offset_page_adjustment_;
|
|
const size_t adjusted_size = size_ + offset_page_adjustment_;
|
|
HighLow file_offset = HighLow::From(adjusted_offset);
|
|
|
|
data_ = static_cast<uint8_t*>(MapViewOfFile(
|
|
file_mapping_, FILE_MAP_READ, file_offset.high, file_offset.low,
|
|
/*dwNumberOfBytesToMap=*/adjusted_size));
|
|
|
|
XNNPACK_RETURN_CHECK(data_ != nullptr, "could not map file (%s): %s",
|
|
safe_path, GetLastErrorString().c_str());
|
|
#else
|
|
offset_page_adjustment_ = offset_ % getpagesize();
|
|
data_ = static_cast<uint8_t*>(
|
|
mmap(/*addr=*/nullptr, size_ + offset_page_adjustment_, PROT_READ,
|
|
MAP_SHARED, fd.Value(), offset_ - offset_page_adjustment_));
|
|
XNNPACK_RETURN_CHECK(data_ != MAP_FAILED,
|
|
"could not mmap file descriptor %d (%s): %s.",
|
|
fd.Value(), safe_path, strerror(errno));
|
|
#endif
|
|
unmap_on_error.Deactivate();
|
|
return true;
|
|
}
|
|
|
|
bool MMapHandle::Resize(size_t new_size) {
|
|
#if (defined(__linux__) || defined(__ANDROID__)) && \
|
|
!defined(XNNPACK_CACHE_NO_FILE_MAPPING_FOR_DEBUG)
|
|
void* const remapped_data =
|
|
mremap(data_, size_ + offset_page_adjustment_,
|
|
new_size + offset_page_adjustment_, /*flags=*/0);
|
|
if (remapped_data == MAP_FAILED) {
|
|
XNNPACK_RETURN_CHECK(errno == ENOMEM, "remap failed: %s", strerror(errno));
|
|
return false;
|
|
}
|
|
size_ = new_size;
|
|
return true;
|
|
#else
|
|
// The current implementation uses new/delete which doesn't provide a way to
|
|
// modify an allocation size. Changing to malloc/realloc/free doesn't ensure
|
|
// that a memory allocation will not be moved when reallocating
|
|
return false;
|
|
#endif
|
|
}
|
|
|
|
void MMapHandle::UnMap() {
|
|
if (data_) {
|
|
#if defined(XNNPACK_CACHE_NO_FILE_MAPPING_FOR_DEBUG)
|
|
delete[] data_;
|
|
#elif defined(_WIN32)
|
|
UnmapViewOfFile(data_);
|
|
CloseHandle(file_mapping_);
|
|
#else
|
|
munmap(data_, size_ + offset_page_adjustment_);
|
|
#endif
|
|
}
|
|
data_ = nullptr;
|
|
offset_ = 0;
|
|
offset_page_adjustment_ = 0;
|
|
size_ = 0;
|
|
}
|
|
|
|
bool MMapHandle::LockMemory() {
|
|
#if defined(XNNPACK_CACHE_NO_FILE_MAPPING_FOR_DEBUG)
|
|
return true;
|
|
#elif defined(_WIN32)
|
|
return VirtualLock(data_, size_) != 0;
|
|
#else
|
|
return mlock(data_, size_ + offset_page_adjustment_) == 0;
|
|
#endif
|
|
}
|
|
|
|
bool MMapHandle::UnlockMemory() {
|
|
#if defined(XNNPACK_CACHE_NO_FILE_MAPPING_FOR_DEBUG)
|
|
return true;
|
|
#elif defined(_WIN32)
|
|
return VirtualUnlock(data_, size_) != 0;
|
|
#else
|
|
return munlock(data_, size_ + offset_page_adjustment_) == 0;
|
|
#endif
|
|
}
|
|
|
|
bool MarkMemoryNotNeeded(void* data, size_t size) {
|
|
#if defined(_WIN32)
|
|
SYSTEM_INFO sysInfo;
|
|
GetSystemInfo(&sysInfo);
|
|
DWORD page_size = sysInfo.dwPageSize;
|
|
// We don't mark memory chunks that are too small.
|
|
// TODO: b/510899567 - This threshold should be re-evaluated on Windows.
|
|
constexpr const size_t size_threshold = 1024 * 1024;
|
|
if (size <= size_threshold) {
|
|
return true;
|
|
}
|
|
#else
|
|
size_t page_size = getpagesize();
|
|
#endif
|
|
|
|
// We align the data buffer to the next page boundary and the size to be a
|
|
// multiple of the page size.
|
|
//
|
|
// - Windows will unlock all pages that contain at least a byte of the given
|
|
// range, which we want to avoid.
|
|
// - Linux requires the address to be on a page boundary.
|
|
void* aligned_data = std::align(page_size, page_size, data, size);
|
|
size -= size % page_size;
|
|
if (aligned_data) {
|
|
#if defined(_WIN32)
|
|
return VirtualUnlock(aligned_data, size);
|
|
#elif defined(__ANDROID__) || defined(__linux__)
|
|
return madvise(aligned_data, size, MADV_PAGEOUT) == 0;
|
|
#endif
|
|
}
|
|
return true;
|
|
}
|
|
|
|
} // namespace tflite::xnnpack
|