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);
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  AddToDirectoryElement(element_name);
73  const string encoded = message.SerializeAsString();
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 Status::OK();
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 Status::OK();
125}
126
127void MemmappedFileSystemWriter::AddToDirectoryElement(const string& name) {
128  MemmappedFileSystemDirectoryElement* new_directory_element =
129      directory_.add_element();
130  new_directory_element->set_offset(output_file_offset_);
131  new_directory_element->set_name(name);
132}
133
134}  // namespace tensorflow
135