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
|
#pragma once
#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include <ATen/core/ivalue.h>
#include <c10/macros/Macros.h>
#include <torch/csrc/jit/mobile/module.h>
/**
* Defines the public API for serializing mobile modules to flatbuffer.
* Note that this header must not include or depend on flatbuffer-defined
* types, to avoid leaking those details to PyTorch clients.
*/
namespace torch::jit {
/// Maps file names to file contents.
using ExtraFilesMap = std::unordered_map<std::string, std::string>;
/**
* Represents a span of data. Typically owned by a UniqueDetachedBuffer.
*/
class TORCH_API DetachedBuffer final {
public:
/// Creates a new DetachedBuffer with an optional data owner. This interface
/// is provided to let users create objects of this type for testing.
DetachedBuffer(void* data, size_t size, void* internal_data_owner = nullptr)
: data_(data), size_(size), data_owner_(internal_data_owner) {}
/// Returns a pointer to the data.
[[nodiscard]] void* data() {
return data_;
}
/// Returns a pointer to the data.
[[nodiscard]] const void* data() const {
return data_;
}
/// Returns the size of the data, in bytes.
[[nodiscard]] size_t size() const {
return size_;
}
/// Wrapper type that typically owns data_owner_.
using UniqueDetachedBuffer =
std::unique_ptr<DetachedBuffer, std::function<void(DetachedBuffer*)>>;
private:
/// Deletes the owner, if present, and the buf itself.
/// Note: we could have provided a movable type with a destructor that did
/// this work, but the unique wrapper was easier in practice.
static void destroy(DetachedBuffer* buf);
/// Provides access to destroy() for implementation and testing.
friend struct DetachedBufferFriend;
friend struct DetachedBufferTestingFriend;
/// Pointer to the data. Not owned by this class.
void* data_;
/// The size of `data_`, in bytes.
size_t size_;
/// Opaque pointer to the underlying owner of `data_`. This class
/// (DetachedBuffer) does not own the owner or the data. It will typically be
/// owned by a UniqueDetachedBuffer that knows how to delete the owner along
/// with this class.
void* data_owner_;
};
TORCH_API void save_mobile_module(
const mobile::Module& module,
const std::string& filename,
const ExtraFilesMap& extra_files = ExtraFilesMap(),
const ExtraFilesMap& jit_sources = ExtraFilesMap(),
const std::vector<IValue>& jit_constants = {});
TORCH_API DetachedBuffer::UniqueDetachedBuffer save_mobile_module_to_bytes(
const mobile::Module& module,
const ExtraFilesMap& extra_files = ExtraFilesMap(),
const ExtraFilesMap& jit_sources = ExtraFilesMap(),
const std::vector<IValue>& jit_constants = {});
TORCH_API void save_mobile_module_to_func(
const mobile::Module& module,
const std::function<size_t(const void*, size_t)>& writer_func);
// TODO(qihan): delete
TORCH_API bool register_flatbuffer_serializer();
} // namespace torch::jit
|
