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
19namespace torch {
20namespace jit {
21
22/// Maps file names to file contents.
23using ExtraFilesMap = std::unordered_map<std::string, std::string>;
24
25/**
26 * Represents a span of data. Typically owned by a UniqueDetachedBuffer.
27 */
28class 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
73TORCH_API void save_mobile_module(
74 const mobile::Module& module,
75 const std::string& filename,
76 const ExtraFilesMap& extra_files = ExtraFilesMap(),
77 const ExtraFilesMap& jit_sources = ExtraFilesMap(),
78 const std::vector<IValue>& jit_constants = {});
79
80TORCH_API DetachedBuffer::UniqueDetachedBuffer save_mobile_module_to_bytes(
81 const mobile::Module& module,
82 const ExtraFilesMap& extra_files = ExtraFilesMap(),
83 const ExtraFilesMap& jit_sources = ExtraFilesMap(),
84 const std::vector<IValue>& jit_constants = {});
85
86TORCH_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
91TORCH_API bool register_flatbuffer_serializer();
92
93} // namespace jit
94} // namespace torch
95