1/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#ifndef TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_
17#define TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_
18
19#include <memory>
20#include <string>
21#include <unordered_map>
22
23#include "tensorflow/core/platform/env.h"
24
25namespace tensorflow {
26
27// A file system that uses a graph saved in memmapped format by
28// MemmappedEnvWriter as a file system.
29//
30// The format supports saved tensors and protos. Tensors are saved at aligned
31// offsets.
32//
33// Format specification:
34// - last 8 bytes of a package is encoded offset to the directory. The encoding
35// is always little endian, independently from the platform, done by functions
36// EncodeUint64LittleEndian/DecodeUint64LittleEndian
37// - the directory starts from the encoded offset and is saved proto
38// MemmappedFileSystemDirectory with names and offsets to the regions.
39// - at the offsets in the directory the file regions are stored. Tensor regions
40// are aligned such way that when the package mapped to RAM they have the right
41// offset to be used by ImmutableConst operator.
42//
43// Region naming:
44// Region naming is up to the application, all of them starts from
45// kMemmappedPackagePrefix. The default graph usually has name
46// kMemmappedPackageDefaultGraphDef;
47//
48// A "frozen" GraphDef can be converted into this format using
49// tensorflow/contrib/util/convert_graphdef_memmapped_format
50class MemmappedFileSystem : public FileSystem {
51 public:
52 // Memmapped regions use this prefix to distinguish from
53 // the filesystem.
54 static constexpr const char kMemmappedPackagePrefix[] =
55 "memmapped_package://";
56
57 // The default graphdef in the package.
58 static constexpr const char kMemmappedPackageDefaultGraphDef[] =
59 "memmapped_package://.";
60
61 MemmappedFileSystem();
62 ~MemmappedFileSystem() override = default;
63
64 TF_USE_FILESYSTEM_METHODS_WITH_NO_TRANSACTION_SUPPORT;
65
66 Status FileExists(const string& fname, TransactionToken* token) override;
67 Status NewRandomAccessFile(
68 const string& filename, TransactionToken* token,
69 std::unique_ptr<RandomAccessFile>* result) override;
70 Status NewReadOnlyMemoryRegionFromFile(
71 const string& filename, TransactionToken* token,
72 std::unique_ptr<ReadOnlyMemoryRegion>* result) override;
73
74 // All these functions return Unimplemented error, the memmapped storage is
75 // read only.
76 Status NewWritableFile(const string& fname, TransactionToken* token,
77 std::unique_ptr<WritableFile>* result) override;
78 Status NewAppendableFile(const string& fname, TransactionToken* token,
79 std::unique_ptr<WritableFile>* result) override;
80 Status GetChildren(const string& dir, TransactionToken* token,
81 std::vector<string>* r) override;
82 Status GetMatchingPaths(const string& pattern, TransactionToken* token,
83 std::vector<string>* results) override;
84 Status DeleteFile(const string& f, TransactionToken* token) override;
85 Status CreateDir(const string& d, TransactionToken* token) override;
86 Status DeleteDir(const string& d, TransactionToken* token) override;
87 Status RenameFile(const string& s, const string& t,
88 TransactionToken* token) override;
89
90 // These functions are implemented.
91 Status GetFileSize(const string& f, TransactionToken* token,
92 uint64* s) override;
93 // Currently just returns size.
94 Status Stat(const string& fname, TransactionToken* token,
95 FileStatistics* stat) override;
96
97 // Initializes filesystem from a file in memmapped format.
98 Status InitializeFromFile(Env* env, const string& filename);
99
100 // Checks if the filename has a correct prefix.
101 static bool IsMemmappedPackageFilename(const string& filename);
102
103 static bool IsWellFormedMemmappedPackageFilename(const string& filename);
104
105 private:
106 struct FileRegion {
107 FileRegion(uint64 o, uint64 l) : offset(o), length(l) {}
108
109 uint64 offset; // Offset from the beginning of the file.
110 uint64 length; // Length of the region.
111 };
112
113 using DirectoryType = std::unordered_map<string, FileRegion>;
114
115 const void* GetMemoryWithOffset(uint64 offset) const;
116
117 std::unique_ptr<ReadOnlyMemoryRegion> mapped_memory_;
118 DirectoryType directory_;
119
120 TF_DISALLOW_COPY_AND_ASSIGN(MemmappedFileSystem);
121};
122
123class MemmappedEnv : public EnvWrapper {
124 public:
125 explicit MemmappedEnv(Env* env);
126 ~MemmappedEnv() override = default;
127 Status GetFileSystemForFile(const string& fname,
128 FileSystem** result) override;
129 Status GetRegisteredFileSystemSchemes(std::vector<string>* schemes) override;
130 Status InitializeFromFile(const string& filename);
131
132 protected:
133 std::unique_ptr<MemmappedFileSystem> memmapped_file_system_;
134};
135
136} // namespace tensorflow
137
138#endif // TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_
139