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