1 | /* Copyright 2016 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | #include "tensorflow/core/util/memmapped_file_system_writer.h" |
16 | |
17 | #include <algorithm> |
18 | |
19 | namespace tensorflow { |
20 | |
21 | Status MemmappedFileSystemWriter::InitializeToFile(Env* env, |
22 | const string& filename) { |
23 | auto status = env->NewWritableFile(filename, &output_file_); |
24 | if (status.ok()) { |
25 | output_file_offset_ = 0; |
26 | } |
27 | return status; |
28 | } |
29 | |
30 | Status MemmappedFileSystemWriter::SaveTensor(const Tensor& tensor, |
31 | const string& element_name) { |
32 | if (!output_file_) { |
33 | return errors::FailedPrecondition( |
34 | "MemmappedEnvWritter: saving tensor into not opened file" ); |
35 | } |
36 | if (!MemmappedFileSystem::IsWellFormedMemmappedPackageFilename( |
37 | element_name)) { |
38 | return errors::InvalidArgument( |
39 | "MemmappedEnvWritter: element_name is invalid: must have memmapped " , |
40 | "package prefix " , MemmappedFileSystem::kMemmappedPackagePrefix, |
41 | " and include [A-Za-z0-9_.]" ); |
42 | } |
43 | const auto tensor_data = tensor.tensor_data(); |
44 | if (tensor_data.empty()) { |
45 | return errors::InvalidArgument( |
46 | "MemmappedEnvWritter: saving tensor with 0 size" ); |
47 | } |
48 | // Adds pad for correct alignment after memmapping. |
49 | TF_RETURN_IF_ERROR(AdjustAlignment(Allocator::kAllocatorAlignment)); |
50 | AddToDirectoryElement(element_name, tensor_data.size()); |
51 | const auto result = output_file_->Append(tensor_data); |
52 | if (result.ok()) { |
53 | output_file_offset_ += tensor_data.size(); |
54 | } |
55 | return result; |
56 | } |
57 | |
58 | Status MemmappedFileSystemWriter::SaveProtobuf( |
59 | const protobuf::MessageLite& message, const string& element_name) { |
60 | if (!output_file_) { |
61 | return errors::FailedPrecondition( |
62 | "MemmappedEnvWritter: saving protobuf into not opened file" ); |
63 | } |
64 | if (!MemmappedFileSystem::IsWellFormedMemmappedPackageFilename( |
65 | element_name)) { |
66 | return errors::InvalidArgument( |
67 | "MemmappedEnvWritter: element_name is invalid: must have memmapped " |
68 | "package prefix " , |
69 | MemmappedFileSystem::kMemmappedPackagePrefix, |
70 | " and include [A-Za-z0-9_.]" ); |
71 | } |
72 | const string encoded = message.SerializeAsString(); |
73 | AddToDirectoryElement(element_name, encoded.size()); |
74 | const auto res = output_file_->Append(encoded); |
75 | if (res.ok()) { |
76 | output_file_offset_ += encoded.size(); |
77 | } |
78 | return res; |
79 | } |
80 | |
81 | namespace { |
82 | |
83 | StringPiece EncodeUint64LittleEndian(uint64 val, char* output_buffer) { |
84 | for (unsigned int i = 0; i < sizeof(uint64); ++i) { |
85 | output_buffer[i] = (val >> i * 8); |
86 | } |
87 | return {output_buffer, sizeof(uint64)}; |
88 | } |
89 | |
90 | } // namespace |
91 | |
92 | Status MemmappedFileSystemWriter::FlushAndClose() { |
93 | if (!output_file_) { |
94 | return errors::FailedPrecondition( |
95 | "MemmappedEnvWritter: flushing into not opened file" ); |
96 | } |
97 | const string dir = directory_.SerializeAsString(); |
98 | TF_RETURN_IF_ERROR(output_file_->Append(dir)); |
99 | |
100 | // Write the directory offset. |
101 | char buffer[sizeof(uint64)]; |
102 | TF_RETURN_IF_ERROR(output_file_->Append( |
103 | EncodeUint64LittleEndian(output_file_offset_, buffer))); |
104 | |
105 | // Flush and close the file. |
106 | TF_RETURN_IF_ERROR(output_file_->Flush()); |
107 | TF_RETURN_IF_ERROR(output_file_->Close()); |
108 | output_file_.reset(); |
109 | return OkStatus(); |
110 | } |
111 | |
112 | Status MemmappedFileSystemWriter::AdjustAlignment(uint64 alignment) { |
113 | const uint64 alignment_rest = output_file_offset_ % alignment; |
114 | const uint64 to_write_for_alignment = |
115 | (alignment_rest == 0) ? 0 : alignment - (output_file_offset_ % alignment); |
116 | static constexpr uint64 kFillerBufferSize = 16; |
117 | const char kFillerBuffer[kFillerBufferSize] = {}; |
118 | for (uint64 rest = to_write_for_alignment; rest > 0;) { |
119 | StringPiece sp(kFillerBuffer, std::min(rest, kFillerBufferSize)); |
120 | TF_RETURN_IF_ERROR(output_file_->Append(sp)); |
121 | rest -= sp.size(); |
122 | output_file_offset_ += sp.size(); |
123 | } |
124 | return OkStatus(); |
125 | } |
126 | |
127 | void MemmappedFileSystemWriter::AddToDirectoryElement(const string& name, |
128 | uint64 length) { |
129 | MemmappedFileSystemDirectoryElement* new_directory_element = |
130 | directory_.add_element(); |
131 | new_directory_element->set_offset(output_file_offset_); |
132 | new_directory_element->set_name(name); |
133 | new_directory_element->set_length(length); |
134 | } |
135 | |
136 | } // namespace tensorflow |
137 | |