Make flatbuffer loads faster if loading as mobile module. (pytorch#78998)

BCFC check: verified that flatbuffer file created in this commit can
be loaded in HEAD and file created in HEAD can be loaded in this commit

Fixes #ISSUE_NUMBER

Pull Request resolved: pytorch#78998
Approved by: https://github.com/zhxchen17

Author: qihqi

torch/csrc/jit/mobile/flatbuffer_loader.cpp

@@ -177,6 +177,18 @@ void parseExtraFiles(
   parseExtraFilesFromVector(extra_files_offsets, &extra_files);
 }
 
+void FlatbufferLoader::parseAndPopulate(
+    uint32_t i,
+    const mobile::serialization::IValue* ivalue) {
+  if (const auto* func = ivalue->val_as_Function()) {
+    auto func_ptr = parseFunction(func);
+    all_functions_[i] = func_ptr.get();
+    mcu_->register_function(std::move(func_ptr));
+  } else {
+    all_ivalues_[i] = parseIValue(ivalue);
+  }
+}
+
 mobile::Module FlatbufferLoader::parseModule(
     mobile::serialization::Module* module) {
   module_ = module;
@@ -192,15 +204,14 @@ mobile::Module FlatbufferLoader::parseModule(
   storages_.resize(module->storage_data_size());
   storage_loaded_.resize(module->storage_data_size(), false);
 
-  for (uint32_t i = 0; i < ivalues->size(); i++) {
+  mobile_ivalue_size_ = module_->mobile_ivalue_size();
+  if (mobile_ivalue_size_ == 0) {
+    mobile_ivalue_size_ = ivalues->size();
+  }
+
+  for (uint32_t i = 0; i < mobile_ivalue_size_; i++) {
     const auto* ival = ivalues->Get(i);
-    if (const auto* func = ival->val_as_Function()) {
-      auto func_ptr = parseFunction(func);
-      all_functions_[i] = func_ptr.get();
-      mcu_->register_function(std::move(func_ptr));
-    } else {
-      all_ivalues_[i] = parseIValue(ival);
-    }
+    parseAndPopulate(i, ival);
   }
   IValue& module_ivalue = getIValue(module->state_obj());
 
@@ -660,6 +671,21 @@ void FlatbufferLoader::extractJitSourceAndConstants(
   AT_ASSERT(
       module_parsed_,
       "Need to first parse a flatbuffer file before extracing jit_sources");
+
+  const auto* ivalues = module_->ivalues();
+  for (uint32_t i = mobile_ivalue_size_; i < ivalues->size(); i++) {
+    const auto* ival = ivalues->Get(i);
+    parseAndPopulate(i, ival);
+  }
+  // register functions
+  for (const auto& f : all_functions_) {
+    if (f.first >= mobile_ivalue_size_) {
+      uint32_t class_index =
+          ivalues->Get(f.first)->val_as_Function()->class_type();
+      ClassTypePtr class_type = all_types_[class_index];
+      class_type->addMethod(f.second);
+    }
+  }
   const auto* jit_constants = module_->jit_constants();
   for (auto i = 0; i < jit_constants->size(); ++i) {
     constants->emplace_back(getIValue(jit_constants->Get(i)));

torch/csrc/jit/mobile/flatbuffer_loader.h

@@ -141,6 +141,9 @@ class TORCH_API FlatbufferLoader {
   IValue parseIValue(const mobile::serialization::IValue* ivalue);
   std::unique_ptr<mobile::Function> parseFunction(
       const mobile::serialization::Function* method);
+  void parseAndPopulate(
+      uint32_t i,
+      const mobile::serialization::IValue* ivalue);
 
   std::unordered_map<uint32_t, mobile::Function*> all_functions_;
   std::vector<ClassTypePtr> all_types_;
@@ -158,6 +161,8 @@ class TORCH_API FlatbufferLoader {
   bool module_parsed_ = false;
   bool should_copy_tensor_memory_ = false;
   bool should_load_operators_ = true;
+  // 0 -> mobile_ivalue_size_ elements are from the mobile module.
+  uint32_t mobile_ivalue_size_ = 0;
 };
 
 } // namespace jit

torch/csrc/jit/serialization/flatbuffer_serializer.cpp

@@ -358,6 +358,7 @@ flatbuffers::DetachedBuffer FlatbufferSerializer::serializeModule(
   auto jit_source_offset = storeExtraFilesAndGetOffset(fbb, jit_sources);
   std::vector<uint32_t> jit_constants_indexes;
   jit_constants_indexes.reserve(jit_constants.size());
+  const uint32_t mobile_ivalue_size = ivalue_offsets_.size();
   for (const auto& ival : jit_constants) {
     jit_constants_indexes.emplace_back(storeIValueAndGetIndex(fbb, ival));
   }
@@ -408,7 +409,8 @@ flatbuffers::DetachedBuffer FlatbufferSerializer::serializeModule(
       fbb.CreateVector(obj_types_offset_),
       jit_source_offset,
       fbb.CreateVector(jit_constants_indexes),
-      operator_version);
+      operator_version,
+      mobile_ivalue_size);
   FinishModuleBuffer(fbb, mod);
   return fbb.Release();
 }

torch/csrc/jit/serialization/mobile_bytecode.fbs

@@ -211,6 +211,12 @@ table Module {
   // To read more:
   // https://github.com/pytorch/rfcs/blob/master/RFC-0017-PyTorch-Operator-Versioning.md
   operator_version:uint;
+
+  // Size of ivalue that comes from the mobile module.
+  // Because the ivalues array above can also have ivalues that cames from
+  // the jit::Module that got it's source attached to flatbuffer file.
+  // this should be smaller than ivalues.size()
+  mobile_ivalue_size:uint;
 }
 
 root_type Module;

torch/csrc/jit/serialization/mobile_bytecode_generated.h

@@ -2228,7 +2228,8 @@ struct Module FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
     VT_OBJECT_TYPES = 18,
     VT_JIT_SOURCES = 20,
     VT_JIT_CONSTANTS = 22,
-    VT_OPERATOR_VERSION = 24
+    VT_OPERATOR_VERSION = 24,
+    VT_MOBILE_IVALUE_SIZE = 26
   };
   uint32_t bytecode_version() const {
     return GetField<uint32_t>(VT_BYTECODE_VERSION, 0);
@@ -2296,6 +2297,12 @@ struct Module FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   bool mutate_operator_version(uint32_t _operator_version = 0) {
     return SetField<uint32_t>(VT_OPERATOR_VERSION, _operator_version, 0);
   }
+  uint32_t mobile_ivalue_size() const {
+    return GetField<uint32_t>(VT_MOBILE_IVALUE_SIZE, 0);
+  }
+  bool mutate_mobile_ivalue_size(uint32_t _mobile_ivalue_size = 0) {
+    return SetField<uint32_t>(VT_MOBILE_IVALUE_SIZE, _mobile_ivalue_size, 0);
+  }
   bool Verify(flatbuffers::Verifier &verifier) const {
     return VerifyTableStart(verifier) &&
            VerifyField<uint32_t>(verifier, VT_BYTECODE_VERSION) &&
@@ -2321,6 +2328,7 @@ struct Module FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
            VerifyOffset(verifier, VT_JIT_CONSTANTS) &&
            verifier.VerifyVector(jit_constants()) &&
            VerifyField<uint32_t>(verifier, VT_OPERATOR_VERSION) &&
+           VerifyField<uint32_t>(verifier, VT_MOBILE_IVALUE_SIZE) &&
            verifier.EndTable();
   }
 };
@@ -2362,6 +2370,9 @@ struct ModuleBuilder {
   void add_operator_version(uint32_t operator_version) {
     fbb_.AddElement<uint32_t>(Module::VT_OPERATOR_VERSION, operator_version, 0);
   }
+  void add_mobile_ivalue_size(uint32_t mobile_ivalue_size) {
+    fbb_.AddElement<uint32_t>(Module::VT_MOBILE_IVALUE_SIZE, mobile_ivalue_size, 0);
+  }
   explicit ModuleBuilder(flatbuffers::FlatBufferBuilder &_fbb)
         : fbb_(_fbb) {
     start_ = fbb_.StartTable();
@@ -2385,8 +2396,10 @@ inline flatbuffers::Offset<Module> CreateModule(
     flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>>> object_types = 0,
     flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>>> jit_sources = 0,
     flatbuffers::Offset<flatbuffers::Vector<uint32_t>> jit_constants = 0,
-    uint32_t operator_version = 0) {
+    uint32_t operator_version = 0,
+    uint32_t mobile_ivalue_size = 0) {
   ModuleBuilder builder_(_fbb);
+  builder_.add_mobile_ivalue_size(mobile_ivalue_size);
   builder_.add_operator_version(operator_version);
   builder_.add_jit_constants(jit_constants);
   builder_.add_jit_sources(jit_sources);
@@ -2413,7 +2426,8 @@ inline flatbuffers::Offset<Module> CreateModuleDirect(
     const std::vector<flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>> *object_types = nullptr,
     const std::vector<flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *jit_sources = nullptr,
     const std::vector<uint32_t> *jit_constants = nullptr,
-    uint32_t operator_version = 0) {
+    uint32_t operator_version = 0,
+    uint32_t mobile_ivalue_size = 0) {
   auto extra_files__ = extra_files ? _fbb.CreateVector<flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>>(*extra_files) : 0;
   auto methods__ = methods ? _fbb.CreateVector<uint32_t>(*methods) : 0;
   auto ivalues__ = ivalues ? _fbb.CreateVector<flatbuffers::Offset<torch::jit::mobile::serialization::IValue>>(*ivalues) : 0;
@@ -2433,7 +2447,8 @@ inline flatbuffers::Offset<Module> CreateModuleDirect(
       object_types__,
       jit_sources__,
       jit_constants__,
-      operator_version);
+      operator_version,
+      mobile_ivalue_size);
 }
 
 inline bool VerifyIValueUnion(flatbuffers::Verifier &verifier, const void *obj, IValueUnion type) {