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#include "tensorflow/core/util/memmapped_file_system_writer.h"
16
17#include <algorithm>
18
19namespace tensorflow {
20
21Status 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
30Status 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
58Status 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
81namespace {
82
83StringPiece 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
92Status 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
112Status 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
127void 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