Skip to content

Commit 5b1be39

Browse files
[executorch][native] Map PTN packages and view constants in place (pytorch#23335)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: pytorch#23230 by @digantdesai ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/digantdesai/110/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/110/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/digantdesai/109/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/110/orig Differential Revision: [D121704319](https://our.internmc.facebook.com/intern/diff/D121704319/) @diff-train-skip-merge Co-authored-by: Digant Desai <digantdesai@meta.com>
1 parent d2d0c25 commit 5b1be39

9 files changed

Lines changed: 257 additions & 13 deletions

File tree

‎backends/native/runtime/deserialize/OwnedBytes.cpp‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#include <executorch/backends/native/runtime/deserialize/OwnedBytes.h>
88

99
#include <cerrno>
10+
#include <cstdint>
1011
#include <filesystem>
1112
#include <fstream>
1213
#include <limits>
@@ -43,6 +44,23 @@ OwnedBytes OwnedBytes::from_vector(std::vector<uint8_t> bytes) {
4344
return OwnedBytes(std::move(bytes));
4445
}
4546

47+
OwnedBytes OwnedBytes::view(
48+
std::shared_ptr<const OwnedBytes> owner,
49+
ByteSpan bytes) {
50+
if (owner == nullptr) {
51+
throw std::invalid_argument("OwnedBytes::view: owner is null");
52+
}
53+
const ByteSpan whole = owner->span();
54+
const auto begin = reinterpret_cast<uintptr_t>(whole.data());
55+
const auto first = reinterpret_cast<uintptr_t>(bytes.data());
56+
if (!bytes.empty() &&
57+
(first < begin || first - begin > whole.size() ||
58+
bytes.size() > whole.size() - (first - begin))) {
59+
throw std::invalid_argument("OwnedBytes::view: bytes lie outside owner");
60+
}
61+
return OwnedBytes(SharedView{std::move(owner), bytes});
62+
}
63+
4664
OwnedBytes OwnedBytes::from_file(
4765
const std::string& path,
4866
bool use_mmap,
@@ -160,11 +178,17 @@ ByteSpan OwnedBytes::span() const {
160178
return ByteSpan(
161179
static_cast<const uint8_t*>(mapped->get()), mapped->get_deleter().size);
162180
}
181+
if (const SharedView* view = std::get_if<SharedView>(&storage_)) {
182+
return view->bytes;
183+
}
163184
const std::vector<uint8_t>& bytes = std::get<std::vector<uint8_t>>(storage_);
164185
return ByteSpan(bytes.data(), bytes.size());
165186
}
166187

167188
bool OwnedBytes::is_mapped() const {
189+
if (const SharedView* view = std::get_if<SharedView>(&storage_)) {
190+
return view->owner->is_mapped();
191+
}
168192
return std::holds_alternative<MappedFile>(storage_);
169193
}
170194

‎backends/native/runtime/deserialize/OwnedBytes.h‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@
1717

1818
namespace ptn {
1919

20-
// Owning, read-only bytes: either a heap buffer or a read-only file mapping.
20+
// Owning, read-only bytes: a heap buffer, a read-only file mapping, or a view
21+
// that keeps another OwnedBytes alive.
2122
//
2223
// Hands out spans that alias the storage. Each alternative keeps its payload
2324
// address across a move, so spans taken before a move stay valid for as long as
@@ -45,13 +46,20 @@ class OwnedBytes {
4546
// A read-only mapping of an entire file.
4647
using MappedFile = std::unique_ptr<void, Unmap>;
4748

48-
std::variant<std::vector<uint8_t>, HeapBuffer, MappedFile> storage_;
49+
struct SharedView {
50+
std::shared_ptr<const OwnedBytes> owner;
51+
ByteSpan bytes;
52+
};
53+
54+
std::variant<std::vector<uint8_t>, HeapBuffer, MappedFile, SharedView>
55+
storage_;
4956

5057
explicit OwnedBytes(std::vector<uint8_t> bytes)
5158
: storage_(std::move(bytes)) {}
5259
explicit OwnedBytes(HeapBuffer buffer) : storage_(std::move(buffer)) {}
5360
explicit OwnedBytes(MappedFile mapped_file)
5461
: storage_(std::move(mapped_file)) {}
62+
explicit OwnedBytes(SharedView view) : storage_(std::move(view)) {}
5563

5664
public:
5765
// Empty, owning nothing.
@@ -66,12 +74,20 @@ class OwnedBytes {
6674
// The whole payload. Valid for this OwnedBytes' lifetime.
6775
ByteSpan span() const;
6876

69-
// True when these bytes are a file mapping rather than a heap buffer.
77+
// True when these bytes are, or view, a file mapping rather than a heap
78+
// buffer.
7079
bool is_mapped() const;
7180

7281
// Take ownership of a buffer the caller already has, without copying it.
7382
static OwnedBytes from_vector(std::vector<uint8_t> bytes);
7483

84+
// A view of `bytes`, which must lie within `owner` unless empty, that shares
85+
// ownership of `owner` rather than copying. Throws std::invalid_argument
86+
// otherwise.
87+
static OwnedBytes view(
88+
std::shared_ptr<const OwnedBytes> owner,
89+
ByteSpan bytes);
90+
7591
// Acquire the contents of `path`. Maps it read-only by default: nothing is
7692
// copied, pages arrive on demand, and they are shared with any other process
7793
// mapping the same file. Pass use_mmap=false to read it into the heap

‎backends/native/runtime/deserialize/Package.cpp‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -86,16 +86,28 @@ Package::Package()
8686

8787
Package Package::load(OwnedBytes bytes) {
8888
Package out;
89-
out.archive_bytes_ = std::move(bytes);
90-
if (out.archive_bytes_.span().size() > detail::kMaxPackageBytes) {
89+
out.archive_bytes_ = std::make_shared<const OwnedBytes>(std::move(bytes));
90+
if (out.archive_bytes_->span().size() > detail::kMaxPackageBytes) {
9191
throw ResourceLimitError("package: image exceeds size limit");
9292
}
93-
out.zip_ = ZipReader::open(out.archive_bytes_.span());
93+
out.zip_ = ZipReader::open(out.archive_bytes_->span());
9494
out.load_metadata();
9595
return out;
9696
}
9797

9898
Package Package::load(const std::string& path) {
99+
std::optional<OwnedBytes> mapped;
100+
try {
101+
mapped = OwnedBytes::from_file(path, true, detail::kMaxPackageBytes);
102+
} catch (const ResourceLimitError&) {
103+
throw;
104+
} catch (const std::runtime_error&) {
105+
// No mmap on this platform or filesystem. Opening through libzip reports
106+
// any error that is not about mapping.
107+
}
108+
if (mapped) {
109+
return load(std::move(*mapped));
110+
}
99111
Package out;
100112
out.zip_ = ZipReader::open(path);
101113
out.load_metadata();
@@ -202,6 +214,22 @@ std::optional<OwnedBytes> Package::acquire_constant(
202214
return OwnedBytes::from_vector(std::move(bytes));
203215
}
204216

217+
std::optional<OwnedBytes> Package::view_constant(const std::string& key) const {
218+
const std::optional<ConstantInfo> info = constant_info(key);
219+
if (!info) {
220+
return std::nullopt;
221+
}
222+
const std::optional<ByteSpan> member = zip_->member_bytes(kSafeTensorsEntry);
223+
if (!member) {
224+
return std::nullopt;
225+
}
226+
return OwnedBytes::view(
227+
archive_bytes_,
228+
member->subspan(
229+
tensor_data_offset_ + tensors_->find(info->owner)->offset,
230+
info->nbytes));
231+
}
232+
205233
bool Package::load_constant_into(
206234
const std::string& key,
207235
MutableByteSpan destination) const {

‎backends/native/runtime/deserialize/Package.h‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
#include <cstddef>
1010
#include <cstdint>
11+
#include <memory>
1112
#include <optional>
1213
#include <string>
1314
#include <unordered_map>
@@ -42,12 +43,14 @@ struct ConstantInfo {
4243
// A loaded .ptn package: the serialized native Program plus the constants it
4344
// references.
4445
//
45-
// Opening a file-backed package reads its directory and metadata, but leaves
46-
// weight payloads on disk until an engine requests them.
46+
// Opening a package from a path reads its directory and metadata, but leaves
47+
// weight payloads on disk until an engine requests them. Opening from bytes
48+
// keeps the caller's whole image resident.
4749
class Package {
4850
private:
4951
uint64_t id_ = 0;
50-
OwnedBytes archive_bytes_;
52+
// Null when the package is read through file I/O rather than memory.
53+
std::shared_ptr<const OwnedBytes> archive_bytes_;
5154
std::optional<ZipReader> zip_;
5255
std::vector<uint8_t> program_;
5356
// Absent when the program references no constants, in which case the package
@@ -66,6 +69,7 @@ class Package {
6669
Package& operator=(const Package&) = delete;
6770

6871
// Open and parse the .ptn at `path` without loading its weight payloads.
72+
// Maps the file where supported; otherwise opens it through libzip.
6973
static Package load(const std::string& path);
7074

7175
// Parse a .ptn image already in hand. Takes ownership rather than copying, so
@@ -100,6 +104,13 @@ class Package {
100104
// Load one constant into a new owning buffer. nullopt when absent.
101105
std::optional<OwnedBytes> acquire_constant(const std::string& key) const;
102106

107+
// One constant's bytes in place in the package's memory, keeping that memory
108+
// alive past the Package instead of copying. nullopt when absent, when the
109+
// package is read through file I/O, or when the archive's local headers are
110+
// not laid out back to back. Alignment is whatever the file gives, so
111+
// a caller with stricter needs checks it and falls back to acquire_constant.
112+
std::optional<OwnedBytes> view_constant(const std::string& key) const;
113+
103114
// Load one constant directly into an exact-sized destination. Returns false
104115
// when absent and throws when the destination has the wrong size.
105116
bool load_constant_into(const std::string& key, MutableByteSpan destination)

‎backends/native/runtime/deserialize/ZipReader.cpp‎

Lines changed: 79 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#include <algorithm>
1010
#include <array>
1111
#include <cstdio>
12+
#include <cstring>
1213
#include <limits>
1314
#include <mutex>
1415
#include <stdexcept>
@@ -48,6 +49,19 @@ using ZipFileHandle = std::unique_ptr<zip_file_t, ZipFileDeleter>;
4849
"zip: " + operation + ": " + zip_file_strerror(file));
4950
}
5051

52+
constexpr uint32_t kLocalHeaderSignature = 0x04034b50;
53+
constexpr size_t kLocalHeaderSize = 30;
54+
constexpr uint16_t kDataDescriptorFlag = 1 << 3;
55+
56+
uint16_t read_u16(const uint8_t* p) {
57+
return static_cast<uint16_t>(p[0] | (p[1] << 8));
58+
}
59+
60+
uint32_t read_u32(const uint8_t* p) {
61+
return static_cast<uint32_t>(read_u16(p)) |
62+
(static_cast<uint32_t>(read_u16(p + 2)) << 16);
63+
}
64+
5165
ZipHandle open_path(const std::string& path) {
5266
int error_code = 0;
5367
ZipHandle archive(
@@ -91,9 +105,12 @@ ZipHandle open_memory(ByteSpan bytes) {
91105
} // namespace
92106

93107
struct ZipReader::Impl {
94-
explicit Impl(ZipHandle handle) : archive(std::move(handle)) {}
108+
explicit Impl(ZipHandle handle, ByteSpan bytes = {})
109+
: archive(std::move(handle)), memory(bytes) {}
95110

96111
ZipHandle archive;
112+
// The archive itself when memory-backed; empty when file-backed.
113+
ByteSpan memory;
97114
mutable std::mutex mutex;
98115
};
99116

@@ -134,12 +151,55 @@ ZipReader::ZipReader(std::unique_ptr<Impl> impl) : impl_(std::move(impl)) {
134151
if (stat.size > std::numeric_limits<size_t>::max()) {
135152
throw std::runtime_error("zip: member is too large: " + name);
136153
}
137-
if (!entries_.emplace(name, Entry{index, static_cast<size_t>(stat.size)})
138-
.second) {
154+
const Entry entry{index, static_cast<size_t>(stat.size), std::nullopt};
155+
if (!entries_.emplace(name, entry).second) {
139156
throw std::runtime_error("zip: duplicate member name: " + name);
140157
}
141158
names_.push_back(name);
142159
}
160+
if (!impl_->memory.empty()) {
161+
locate_member_data();
162+
}
163+
}
164+
165+
// libzip does not expose where member data starts, so walk the local headers.
166+
// Accept only a simple layout: stored members contiguous from offset zero in
167+
// central-directory order. Otherwise leave offsets unset and read via libzip.
168+
void ZipReader::locate_member_data() {
169+
const ByteSpan archive = impl_->memory;
170+
std::vector<size_t> offsets;
171+
offsets.reserve(names_.size());
172+
size_t position = 0;
173+
for (const std::string& name : names_) {
174+
if (archive.size() - position < kLocalHeaderSize) {
175+
return;
176+
}
177+
const uint8_t* header = archive.data() + position;
178+
const size_t name_size = read_u16(header + 26);
179+
const size_t extra_size = read_u16(header + 28);
180+
if (read_u32(header) != kLocalHeaderSignature ||
181+
(read_u16(header + 6) & kDataDescriptorFlag) != 0 ||
182+
read_u16(header + 8) != ZIP_CM_STORE) {
183+
return;
184+
}
185+
const size_t name_start = position + kLocalHeaderSize;
186+
if (archive.size() - name_start < name_size + extra_size ||
187+
std::string_view(
188+
reinterpret_cast<const char*>(archive.data() + name_start),
189+
name_size) != name) {
190+
return;
191+
}
192+
const size_t data = name_start + name_size + extra_size;
193+
const size_t size = entries_.find(name)->second.size;
194+
if (archive.size() - data < size) {
195+
return;
196+
}
197+
offsets.push_back(data);
198+
position = data + size;
199+
}
200+
for (size_t i = 0; i < offsets.size(); ++i) {
201+
entries_.find(names_[i])->second.data_offset = offsets[i];
202+
}
143203
}
144204

145205
ZipReader::~ZipReader() = default;
@@ -151,14 +211,22 @@ ZipReader ZipReader::open(const std::string& path) {
151211
}
152212

153213
ZipReader ZipReader::open(ByteSpan archive) {
154-
return ZipReader(std::make_unique<Impl>(open_memory(archive)));
214+
return ZipReader(std::make_unique<Impl>(open_memory(archive), archive));
155215
}
156216

157217
std::optional<size_t> ZipReader::member_size(std::string_view name) const {
158218
const Entry* entry = find_entry(name);
159219
return entry == nullptr ? std::nullopt : std::optional<size_t>(entry->size);
160220
}
161221

222+
std::optional<ByteSpan> ZipReader::member_bytes(std::string_view name) const {
223+
const Entry* entry = find_entry(name);
224+
if (entry == nullptr || !entry->data_offset) {
225+
return std::nullopt;
226+
}
227+
return impl_->memory.subspan(*entry->data_offset, entry->size);
228+
}
229+
162230
std::vector<uint8_t> ZipReader::read(std::string_view name) const {
163231
const Entry* entry = find_entry(name);
164232
if (entry == nullptr) {
@@ -197,6 +265,13 @@ void ZipReader::read_entry_into(
197265
if (destination.empty()) {
198266
return;
199267
}
268+
if (entry.data_offset) {
269+
std::memcpy(
270+
destination.data(),
271+
impl_->memory.data() + *entry.data_offset + offset,
272+
destination.size());
273+
return;
274+
}
200275

201276
const std::lock_guard<std::mutex> lock(impl_->mutex);
202277
ZipFileHandle file(

‎backends/native/runtime/deserialize/ZipReader.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ class ZipReader {
2626
struct Entry {
2727
uint64_t index = 0;
2828
size_t size = 0;
29+
// Where the member's data starts in a memory-backed archive, when known.
30+
std::optional<size_t> data_offset;
2931
};
3032

3133
struct StringHash {
@@ -43,6 +45,7 @@ class ZipReader {
4345
std::vector<std::string> names_;
4446

4547
explicit ZipReader(std::unique_ptr<Impl> impl);
48+
void locate_member_data();
4649
const Entry* find_entry(std::string_view name) const;
4750
void read_entry_into(
4851
std::string_view name,
@@ -69,6 +72,11 @@ class ZipReader {
6972
// Member size, or nullopt when the member is absent.
7073
std::optional<size_t> member_size(std::string_view name) const;
7174

75+
// A member's bytes in place within a memory-backed archive. nullopt when the
76+
// member is absent, the reader is file-backed, or the archive's local headers
77+
// are not laid out back to back in central-directory order.
78+
std::optional<ByteSpan> member_bytes(std::string_view name) const;
79+
7280
// Copies one complete member.
7381
std::vector<uint8_t> read(std::string_view name) const;
7482

0 commit comments

Comments
 (0)