/* Copyright 2016 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.
|
==============================================================================*/
|
|
#ifndef TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_
|
#define TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_
|
|
#include <memory>
|
#include <string>
|
#include <unordered_map>
|
|
#include "tensorflow/core/platform/env.h"
|
|
namespace tensorflow {
|
|
// A file system that uses a graph saved in memmapped format by
|
// MemmappedEnvWriter as a file system.
|
//
|
// The format supports saved tensors and protos. Tensors are saved at aligned
|
// offsets.
|
//
|
// Format specification:
|
// - last 8 bytes of a package is encoded offset to the directory. The encoding
|
// is always little endian, independently from the platform, done by functions
|
// EncodeUint64LittleEndian/DecodeUint64LittleEndian
|
// - the directory starts from the encoded offset and is saved proto
|
// MemmappedFileSystemDirectory with names and offsets to the regions.
|
// - at the offsets in the directory the file regions are stored. Tensor regions
|
// are aligned such way that when the package mapped to RAM they have the right
|
// offset to be used by ImmutableConst operator.
|
//
|
// Region naming:
|
// Region naming is up to the application, all of them starts from
|
// kMemmappedPackagePrefix. The default graph usually has name
|
// kMemmappedPackageDefaultGraphDef; for more details see the conversion
|
// utility
|
// third_party/tensorflow/contrib/util/convert_graphdef_memmapped_format.cc
|
//
|
// A "frozen" GraphDef can be converted into this format using
|
// tensorflow/contrib/util/convert_graphdef_memmapped_format
|
class MemmappedFileSystem : public FileSystem {
|
public:
|
// Memmapped regions use this prefix to distinguish from
|
// the filesystem.
|
#if defined(_MSC_VER)
|
static constexpr char* kMemmappedPackagePrefix =
|
#else
|
static constexpr char kMemmappedPackagePrefix[] =
|
#endif
|
"memmapped_package://";
|
|
// The default graphdef in the package.
|
#if defined(_MSC_VER)
|
static constexpr char* kMemmappedPackageDefaultGraphDef =
|
#else
|
static constexpr char kMemmappedPackageDefaultGraphDef[] =
|
#endif
|
"memmapped_package://.";
|
|
MemmappedFileSystem();
|
~MemmappedFileSystem() override = default;
|
Status FileExists(const string& fname) override;
|
Status NewRandomAccessFile(
|
const string& filename,
|
std::unique_ptr<RandomAccessFile>* result) override;
|
Status NewReadOnlyMemoryRegionFromFile(
|
const string& filename,
|
std::unique_ptr<ReadOnlyMemoryRegion>* result) override;
|
|
// All these functions return Unimplemented error, the memmapped storage is
|
// read only.
|
Status NewWritableFile(const string& fname,
|
std::unique_ptr<WritableFile>* result) override;
|
Status NewAppendableFile(const string& fname,
|
std::unique_ptr<WritableFile>* result) override;
|
Status GetChildren(const string& dir, std::vector<string>* r) override;
|
Status GetMatchingPaths(const string& pattern,
|
std::vector<string>* results) override;
|
Status DeleteFile(const string& f) override;
|
Status CreateDir(const string& d) override;
|
Status DeleteDir(const string& d) override;
|
Status RenameFile(const string& s, const string& t) override;
|
|
// These functions are implemented.
|
Status GetFileSize(const string& f, uint64* s) override;
|
// Currently just returns size.
|
Status Stat(const string& fname, FileStatistics* stat) override;
|
|
// Initializes filesystem from a file in memmapped format.
|
Status InitializeFromFile(Env* env, const string& filename);
|
|
// Checks if the filename has a correct prefix.
|
static bool IsMemmappedPackageFilename(const string& filename);
|
|
static bool IsWellFormedMemmappedPackageFilename(const string& filename);
|
|
private:
|
struct FileRegion {
|
FileRegion(uint64 o, uint64 l) : offset(o), length(l) {}
|
|
uint64 offset; // Offset from the beginning of the file.
|
uint64 length; // Length of the region.
|
};
|
|
using DirectoryType = std::unordered_map<string, FileRegion>;
|
|
const void* GetMemoryWithOffset(uint64 offset) const;
|
|
std::unique_ptr<ReadOnlyMemoryRegion> mapped_memory_;
|
DirectoryType directory_;
|
|
TF_DISALLOW_COPY_AND_ASSIGN(MemmappedFileSystem);
|
};
|
|
class MemmappedEnv : public EnvWrapper {
|
public:
|
explicit MemmappedEnv(Env* env);
|
~MemmappedEnv() override = default;
|
Status GetFileSystemForFile(const string& fname,
|
FileSystem** result) override;
|
Status GetRegisteredFileSystemSchemes(std::vector<string>* schemes) override;
|
Status InitializeFromFile(const string& filename);
|
|
protected:
|
std::unique_ptr<MemmappedFileSystem> memmapped_file_system_;
|
};
|
|
} // namespace tensorflow
|
|
#endif // TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_
|
