File size: 5,085 Bytes
5f923cd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
// Copyright 2024 The ODML Authors.
//
// 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.

#include <fcntl.h>
#include <sys/mman.h>

#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <memory>

#include "absl/status/status.h"  // from @com_google_absl
#include "absl/status/statusor.h"  // from @com_google_absl
#include "absl/strings/string_view.h"  // from @com_google_absl
#include "runtime/util/memory_mapped_file.h"
#include "runtime/util/scoped_file.h"
#include "runtime/util/status_macros.h"

namespace litert::lm {
namespace {

class MemoryMappedFilePosix : public MemoryMappedFile {
 public:
  MemoryMappedFilePosix(uint64_t length, void* data)
      : length_(length), data_(data) {}
  ~MemoryMappedFilePosix() override {
    if (data_) {
      munmap(data_, length_);
    }
  }

  // Move constructor
  MemoryMappedFilePosix(MemoryMappedFilePosix&& other) noexcept
      : length_(other.length_), data_(other.data_) {
    // After transferring ownership of the data pointer and length,
    // we must reset the other object so its destructor doesn't free
    // the memory we just took ownership of.
    other.length_ = 0;
    other.data_ = nullptr;
  }

  // Move assignment
  MemoryMappedFilePosix& operator=(MemoryMappedFilePosix&& other) noexcept {
    if (this != &other) {
      // Free existing resource before taking ownership of the new one
      if (data_ != nullptr) {
        munmap(data_, length_);
      }

      // Transfer ownership from the other object
      length_ = other.length_;
      data_ = other.data_;

      // Reset the other object
      other.length_ = 0;
      other.data_ = nullptr;
    }
    return *this;
  }

  // Disable copy operations.
  MemoryMappedFilePosix(const MemoryMappedFilePosix&) = delete;
  MemoryMappedFilePosix& operator=(const MemoryMappedFilePosix&) = delete;

  uint64_t length() override { return length_; }

  void* data() override { return data_; }

 private:
  uint64_t length_;
  void* data_;
};

}  // namespace

// static
size_t MemoryMappedFile::GetOffsetAlignment() { return getpagesize(); }

// static
absl::StatusOr<std::unique_ptr<MemoryMappedFile>> MemoryMappedFile::Create(
    absl::string_view path) {
  ASSIGN_OR_RETURN(auto scoped_file, ScopedFile::Open(path));
  return Create(scoped_file.file());
}

// static
absl::StatusOr<std::unique_ptr<MemoryMappedFile>> MemoryMappedFile::Create(
    int file, uint64_t offset, uint64_t length, absl::string_view key) {
  RET_CHECK_EQ(offset % GetOffsetAlignment(), 0)
      << "Offset must be a multiple of page size : " << offset << ", "
      << GetOffsetAlignment();

  ASSIGN_OR_RETURN(size_t file_size, ScopedFile::GetSize(file));
  RET_CHECK_GE(file_size, length + offset) << "Length and offset too large.";
  if (length == 0) {
    length = file_size - offset;
  }

  void* data =
      mmap(nullptr, length, PROT_READ | PROT_WRITE, MAP_PRIVATE, file, offset);
  RET_CHECK_NE(data, MAP_FAILED) << "Failed to map, error: " << strerror(errno);
  RET_CHECK_NE(data, nullptr) << "Failed to map.";
#ifdef __APPLE__
  // Mark it not needed to avoid unnecessary page loading on MacOS or iOS.
  RET_CHECK_EQ(madvise(data, length, MADV_DONTNEED), 0) << "madvise failed.";
#else
  RET_CHECK_EQ(madvise(data, length, MADV_WILLNEED), 0) << "madvise failed.";
#endif

  return std::make_unique<MemoryMappedFilePosix>(length, data);
}

absl::StatusOr<std::unique_ptr<MemoryMappedFile>>
MemoryMappedFile::CreateMutable(absl::string_view path) {
  ASSIGN_OR_RETURN(auto scoped_file, ScopedFile::OpenWritable(path));
  return CreateMutable(scoped_file.file());
}

absl::StatusOr<std::unique_ptr<MemoryMappedFile>>
MemoryMappedFile::CreateMutable(int file, uint64_t offset, uint64_t length,
                                absl::string_view key) {
  RET_CHECK_EQ(offset % GetOffsetAlignment(), 0)
      << "Offset must be a multiple of page size : " << offset << ", "
      << GetOffsetAlignment();

  ASSIGN_OR_RETURN(size_t file_size, ScopedFile::GetSize(file));
  RET_CHECK_GE(file_size, length + offset) << "Length and offset too large.";
  if (length == 0) {
    length = file_size - offset;
  }
  if (length == 0) {
    return absl::InvalidArgumentError("Cannot mmap empty file.");
  }

  void* data =
      mmap(nullptr, length, PROT_READ | PROT_WRITE, MAP_SHARED, file, offset);
  RET_CHECK_NE(data, MAP_FAILED) << "Failed to map, error: " << strerror(errno);
  RET_CHECK_NE(data, nullptr) << "Failed to map.";

  return std::make_unique<MemoryMappedFilePosix>(length, data);
}

}  // namespace litert::lm