1 | #pragma once |
2 | |
3 | #include <functional> |
4 | #include <memory> |
5 | #include <string> |
6 | #include <unordered_map> |
7 | #include <vector> |
8 | |
9 | #include <ATen/core/ivalue.h> |
10 | #include <c10/macros/Macros.h> |
11 | #include <torch/csrc/jit/mobile/module.h> |
12 | |
13 | /** |
14 | * Defines the public API for serializing mobile modules to flatbuffer. |
15 | * Note that this header must not include or depend on flatbuffer-defined |
16 | * types, to avoid leaking those details to PyTorch clients. |
17 | */ |
18 | |
19 | namespace torch { |
20 | namespace jit { |
21 | |
22 | /// Maps file names to file contents. |
23 | using = std::unordered_map<std::string, std::string>; |
24 | |
25 | /** |
26 | * Represents a span of data. Typically owned by a UniqueDetachedBuffer. |
27 | */ |
28 | class TORCH_API DetachedBuffer final { |
29 | public: |
30 | /// Creates a new DetachedBuffer with an optional data owner. This interface |
31 | /// is provided to let users create objects of this type for testing. |
32 | DetachedBuffer(void* data, size_t size, void* internal_data_owner = nullptr) |
33 | : data_(data), size_(size), data_owner_(internal_data_owner) {} |
34 | |
35 | /// Returns a pointer to the data. |
36 | C10_NODISCARD void* data() { |
37 | return data_; |
38 | } |
39 | /// Returns a pointer to the data. |
40 | C10_NODISCARD const void* data() const { |
41 | return data_; |
42 | } |
43 | /// Returns the size of the data, in bytes. |
44 | C10_NODISCARD size_t size() const { |
45 | return size_; |
46 | } |
47 | |
48 | /// Wrapper type that typically owns data_owner_. |
49 | using UniqueDetachedBuffer = |
50 | std::unique_ptr<DetachedBuffer, std::function<void(DetachedBuffer*)>>; |
51 | |
52 | private: |
53 | /// Deletes the owner, if present, and the buf itself. |
54 | /// Note: we could have provided a movable type with a destructor that did |
55 | /// this work, but the unique wrapper was easier in practice. |
56 | static void destroy(DetachedBuffer* buf); |
57 | |
58 | /// Provides access to destroy() for implementation and testing. |
59 | friend struct DetachedBufferFriend; |
60 | friend struct DetachedBufferTestingFriend; |
61 | |
62 | /// Pointer to the data. Not owned by this class. |
63 | void* data_; |
64 | /// The size of `data_`, in bytes. |
65 | size_t size_; |
66 | /// Opaque pointer to the underlying owner of `data_`. This class |
67 | /// (DetachedBuffer) does not own the owner or the data. It will typically be |
68 | /// owned by a UniqueDetachedBuffer that knows how to delete the owner along |
69 | /// with this class. |
70 | void* data_owner_; |
71 | }; |
72 | |
73 | TORCH_API void save_mobile_module( |
74 | const mobile::Module& module, |
75 | const std::string& filename, |
76 | const ExtraFilesMap& = ExtraFilesMap(), |
77 | const ExtraFilesMap& jit_sources = ExtraFilesMap(), |
78 | const std::vector<IValue>& jit_constants = {}); |
79 | |
80 | TORCH_API DetachedBuffer::UniqueDetachedBuffer save_mobile_module_to_bytes( |
81 | const mobile::Module& module, |
82 | const ExtraFilesMap& = ExtraFilesMap(), |
83 | const ExtraFilesMap& jit_sources = ExtraFilesMap(), |
84 | const std::vector<IValue>& jit_constants = {}); |
85 | |
86 | TORCH_API void save_mobile_module_to_func( |
87 | const mobile::Module& module, |
88 | const std::function<size_t(const void*, size_t)>& writer_func); |
89 | |
90 | // TODO(qihan): delete |
91 | TORCH_API bool register_flatbuffer_serializer(); |
92 | |
93 | } // namespace jit |
94 | } // namespace torch |
95 | |