[jit][edge] Do not reuse mobile type parser for all unpicklers. (pytorch#70338)

zhxchen17 · facebook-github-bot · commit 0a921ba0d066 · 2022-01-07T18:35:32.000-08:00
Summary: Pull Request resolved: pytorch#70338 Today Unpickler is used by both server and mobile for deserializing model, and it always fallback to mobile parser when there's no type resolver provided by user. However this is not intended as server and mobile type parser supports different things. In this diff we provide a default fallback using script parser and opt it out for all mobile cases. ghstack-source-id: 146727330 (Note: this ignores all push blocking failures!) Test Plan: CI Reviewed By: iseeyuan Differential Revision: D33284352 fbshipit-source-id: 997c4f110b36eee6596e8f23f6a87bf91a4197ed
diff --git a/test/cpp/jit/test_mobile_type_parser.cpp b/test/cpp/jit/test_mobile_type_parser.cpp
@@ -2,11 +2,7 @@
 #include <test/cpp/jit/test_utils.h>
 
 #include <ATen/core/jit_type.h>
-
-namespace c10 {
-TypePtr parseType(const std::string& pythonStr);
-std::vector<TypePtr> parseType(std::vector<std::string>& pythonStr);
-} // namespace c10
+#include <torch/csrc/jit/mobile/type_parser.h>
 
 namespace torch {
 namespace jit {
diff --git a/test/mobile/test_upgrader_bytecode_table_example.cpp b/test/mobile/test_upgrader_bytecode_table_example.cpp
@@ -5,13 +5,11 @@
  * cd ~/pytorch && python torch/csrc/jit/mobile/upgrader_mobile.cpp
  */
 
-#include <caffe2/serialize/versions.h>
 #include <torch/csrc/jit/mobile/upgrader_mobile.h>
-#include <ATen/core/ivalue.h>
 
-namespace c10 {
-TypePtr parseType(const std::string& pythonStr);
-} // namespace c10
+#include <ATen/core/ivalue.h>
+#include <caffe2/serialize/versions.h>
+#include <torch/csrc/jit/mobile/type_parser.h>
 
 namespace torch {
 namespace jit {
diff --git a/tools/codegen/operator_versions/gen_mobile_upgraders.py b/tools/codegen/operator_versions/gen_mobile_upgraders.py
@@ -101,13 +101,11 @@ class ByteCode(Enum):
  * cd ~/pytorch && python torch/csrc/jit/mobile/upgrader_mobile.cpp
  */
 
-#include <caffe2/serialize/versions.h>
 #include <torch/csrc/jit/mobile/upgrader_mobile.h>
-#include <ATen/core/ivalue.h>
 
-namespace c10 {
-TypePtr parseType(const std::string& pythonStr);
-} // namespace c10
+#include <ATen/core/ivalue.h>
+#include <caffe2/serialize/versions.h>
+#include <torch/csrc/jit/mobile/type_parser.h>
 
 namespace torch {
 namespace jit {
diff --git a/torch/csrc/jit/frontend/tree.h b/torch/csrc/jit/frontend/tree.h
@@ -29,8 +29,6 @@ struct Tree;
 using TreeRef = c10::intrusive_ptr<Tree>;
 using TreeList = at::SmallVector<TreeRef, 4>;
 
-static const TreeList empty_trees = {};
-
 struct Tree : c10::intrusive_ptr_target {
   Tree(int kind_) : kind_(kind_) {}
   int kind() const {
@@ -46,6 +44,7 @@ struct Tree : c10::intrusive_ptr_target {
     throw std::runtime_error("stringValue can only be called on TK_STRING");
   }
   virtual const TreeList& trees() const {
+    static const TreeList empty_trees = {};
     return empty_trees;
   }
   const TreeRef& tree(size_t i) const {
@@ -149,11 +148,11 @@ struct Compound : public Tree {
     return false;
   }
   TreeRef map(const std::function<TreeRef(TreeRef)>& fn) override {
-    TreeList trees_;
+    TreeList ret;
     for (auto& t : trees()) {
-      trees_.push_back(fn(t));
+      ret.push_back(fn(t));
     }
-    return Compound::create(kind(), range(), std::move(trees_));
+    return Compound::create(kind(), range(), std::move(ret));
   }
 
   const SourceRange& range() const override {
diff --git a/torch/csrc/jit/mobile/debug_info.cpp b/torch/csrc/jit/mobile/debug_info.cpp
@@ -1,5 +1,6 @@
 #include <torch/csrc/jit/frontend/source_range.h>
 #include <torch/csrc/jit/mobile/debug_info.h>
+#include <torch/csrc/jit/mobile/type_parser.h>
 #include <torch/csrc/jit/serialization/callstack_debug_info_serialization.h>
 #include <torch/csrc/jit/serialization/source_range_serialization.h>
 
@@ -122,10 +123,13 @@ MobileDebugTable::MobileDebugTable(
       size_t debug_size{0};
       std::tie(debug_data, debug_size) = reader->getRecord(record_name);
       auto ivalues =
-          std::move(
-              *jit::unpickle(
-                   reinterpret_cast<const char*>(debug_data.get()), debug_size)
-                   .toTuple())
+          std::move(*jit::unpickle(
+                         reinterpret_cast<const char*>(debug_data.get()),
+                         debug_size,
+                         nullptr,
+                         {},
+                         c10::parseType)
+                         .toTuple())
               .elements();
       SourceRangeDeserializer deserializer;
       for (auto& val : ivalues) {
diff --git a/torch/csrc/jit/mobile/import.cpp b/torch/csrc/jit/mobile/import.cpp
@@ -12,6 +12,7 @@
 #include <torch/csrc/jit/api/compilation_unit.h>
 #include <torch/csrc/jit/mobile/interpreter.h>
 #include <torch/csrc/jit/mobile/observer.h>
+#include <torch/csrc/jit/mobile/type_parser.h>
 #include <torch/csrc/jit/mobile/upgrader_mobile.h>
 #include <torch/csrc/jit/runtime/instruction.h>
 #include <torch/csrc/jit/serialization/import_export_constants.h>
@@ -78,11 +79,6 @@
 //  - Argument::{known_length_,kwarg_only_}
 //  - FunctionSchema::{overload_name_, is_vararg_, is_varret_}
 
-namespace c10 {
-// std::string serializeType(const Type &t);
-TypePtr parseType(const std::string& pythonStr);
-} // namespace c10
-
 namespace torch {
 namespace jit {
 using caffe2::serialize::IStreamAdapter;
@@ -502,7 +498,8 @@ c10::IValue BytecodeDeserializer::readArchive(
       type_resolver,
       obj_loader,
       device_,
-      *reader_.get());
+      *reader_.get(),
+      nullptr);
   return ivalues;
 }
 
diff --git a/torch/csrc/jit/mobile/import_data.cpp b/torch/csrc/jit/mobile/import_data.cpp
@@ -5,6 +5,7 @@
 #include <caffe2/serialize/inline_container.h>
 #include <torch/csrc/jit/api/compilation_unit.h>
 #include <torch/csrc/jit/mobile/observer.h>
+#include <torch/csrc/jit/mobile/type_parser.h>
 #include <torch/csrc/jit/runtime/instruction.h>
 #include <torch/csrc/jit/serialization/unpickler.h>
 #include <torch/custom_class.h>
@@ -14,11 +15,6 @@
 #include <string>
 #include <vector>
 
-namespace c10 {
-// std::string serializeType(const Type &t);
-TypePtr parseType(const std::string& pythonStr);
-} // namespace c10
-
 namespace torch {
 namespace jit {
 using caffe2::serialize::IStreamAdapter;
@@ -151,7 +147,9 @@ c10::IValue BytecodeDeserializer::readArchive(
       std::move(obj_loader),
       std::move(read_record),
       // NOLINTNEXTLINE(performance-move-const-arg)
-      std::move(device));
+      std::move(device),
+      false,
+      nullptr);
   return unpickler.parse_ivalue();
 }
 
diff --git a/torch/csrc/jit/mobile/model_compatibility.cpp b/torch/csrc/jit/mobile/model_compatibility.cpp
@@ -53,7 +53,8 @@ c10::IValue readArchive(
       type_resolver,
       obj_loader,
       device,
-      stream_reader);
+      stream_reader,
+      nullptr);
   return ivalues;
 }
 
diff --git a/torch/csrc/jit/mobile/type_parser.cpp b/torch/csrc/jit/mobile/type_parser.cpp
@@ -1,3 +1,5 @@
+#include <torch/csrc/jit/mobile/type_parser.h>
+
 #include <ATen/core/jit_type.h>
 #include <c10/util/string_view.h>
 #include <torch/csrc/jit/frontend/parser_constants.h>
diff --git a/torch/csrc/jit/mobile/type_parser.h b/torch/csrc/jit/mobile/type_parser.h
@@ -1,3 +1,5 @@
+#pragma once
+
 #include <ATen/core/dynamic_type.h>
 #include <ATen/core/jit_type.h>
 
diff --git a/torch/csrc/jit/runtime/register_ops_utils.cpp b/torch/csrc/jit/runtime/register_ops_utils.cpp
@@ -384,8 +384,8 @@ void listAdd(Stack& stack) {
 }
 
 void listInplaceAdd(Stack& stack) {
-  c10::List<IValue> b = pop(stack).to<List<IValue>>();
-  c10::List<IValue> a = pop(stack).to<List<IValue>>();
+  c10::List<IValue> b = pop(stack).to<c10::List<IValue>>();
+  c10::List<IValue> a = pop(stack).to<c10::List<IValue>>();
   a.append(std::move(b));
   push(stack, std::move(a));
 }
diff --git a/torch/csrc/jit/runtime/register_prim_ops.cpp b/torch/csrc/jit/runtime/register_prim_ops.cpp
@@ -973,7 +973,7 @@ static const std::vector<OperatorGeneratorArgs> opGenArgs{
         TORCH_SELECTIVE_SCHEMA(
             "aten::index.Tensor_hacked_twin(Tensor self, Tensor[] indices) -> Tensor"),
         [](Stack& stack) {
-          auto indices = pop(stack).to<List<c10::optional<at::Tensor>>>();
+          auto indices = pop(stack).to<c10::List<c10::optional<at::Tensor>>>();
           auto self = pop(stack).toTensor();
           auto result = at::index(self, indices);
           push(stack, std::move(result));
@@ -986,7 +986,7 @@ static const std::vector<OperatorGeneratorArgs> opGenArgs{
           auto unsafe = pop(stack).toBool();
           auto accumulate = pop(stack).toBool();
           auto values = pop(stack).toTensor();
-          auto indices = pop(stack).to<List<c10::optional<at::Tensor>>>();
+          auto indices = pop(stack).to<c10::List<c10::optional<at::Tensor>>>();
           auto self = pop(stack).toTensor();
           auto result =
               at::_index_put_impl_(self, indices, values, accumulate, unsafe);
@@ -999,7 +999,7 @@ static const std::vector<OperatorGeneratorArgs> opGenArgs{
         [](Stack& stack) {
           auto accumulate = pop(stack).toBool();
           auto values = pop(stack).toTensor();
-          auto indices = pop(stack).to<List<c10::optional<at::Tensor>>>();
+          auto indices = pop(stack).to<c10::List<c10::optional<at::Tensor>>>();
           auto self = pop(stack).toTensor();
           auto result = at::index_put_(self, indices, values, accumulate);
           push(stack, std::move(result));
@@ -1011,7 +1011,7 @@ static const std::vector<OperatorGeneratorArgs> opGenArgs{
         [](Stack& stack) {
           auto accumulate = pop(stack).toBool();
           auto values = pop(stack).toTensor();
-          auto indices = pop(stack).to<List<c10::optional<at::Tensor>>>();
+          auto indices = pop(stack).to<c10::List<c10::optional<at::Tensor>>>();
           auto self = pop(stack).toTensor();
           auto result = at::index_put_(self, indices, values, accumulate);
           push(stack, std::move(result));
diff --git a/torch/csrc/jit/serialization/callstack_debug_info_serialization.cpp b/torch/csrc/jit/serialization/callstack_debug_info_serialization.cpp
@@ -1,4 +1,5 @@
 #include <torch/csrc/jit/api/compilation_unit.h>
+#include <torch/csrc/jit/mobile/type_parser.h>
 #include <torch/csrc/jit/serialization/callstack_debug_info_serialization.h>
 #include <torch/csrc/jit/serialization/pickle.h>
 
@@ -214,7 +215,12 @@ ska::flat_hash_map<int64_t, DebugInfoTuple> CallStackDebugInfoUnpickler::
         size_t size,
         const ska::flat_hash_map<int64_t, SourceRange>& source_range_map,
         const std::shared_ptr<CompilationUnit>& cu) {
-  auto ival = jit::unpickle(reinterpret_cast<const char*>(data.get()), size);
+  auto ival = jit::unpickle(
+      reinterpret_cast<const char*>(data.get()),
+      size,
+      nullptr,
+      {},
+      c10::parseType);
   ska::flat_hash_map<int64_t, DebugInfoTuple> callstack_ptrs;
   auto ivalues = std::move(*std::move(ival).toTuple()).elements();
   for (auto& val : ivalues) {
diff --git a/torch/csrc/jit/serialization/import.cpp b/torch/csrc/jit/serialization/import.cpp
@@ -176,6 +176,7 @@ IValue ScriptModuleDeserializer::readArchive(const std::string& archive_name) {
       obj_loader,
       device_,
       *reader_.get(),
+      nullptr,
       storage_context_);
 }
 
diff --git a/torch/csrc/jit/serialization/import_read.cpp b/torch/csrc/jit/serialization/import_read.cpp
@@ -12,6 +12,7 @@ IValue readArchiveAndTensors(
     c10::optional<ObjLoader> obj_loader,
     c10::optional<at::Device> device,
     caffe2::serialize::PyTorchStreamReader& stream_reader,
+    c10::TypePtr (*type_parser)(const std::string&),
     std::shared_ptr<DeserializationStorageContext> storage_context) {
   std::string picklename = pickle_prefix + archive_name + ".pkl";
   at::DataPtr pickle_ptr;
@@ -47,6 +48,7 @@ IValue readArchiveAndTensors(
       std::move(read_record),
       device,
       false,
+      type_parser,
       storage_context);
   unpickler.set_version(stream_reader.version());
   return unpickler.parse_ivalue();
diff --git a/torch/csrc/jit/serialization/import_read.h b/torch/csrc/jit/serialization/import_read.h
@@ -20,6 +20,8 @@ TORCH_API IValue readArchiveAndTensors(
     c10::optional<ObjLoader> obj_loader,
     c10::optional<at::Device> device,
     caffe2::serialize::PyTorchStreamReader& stream_reader,
+    c10::TypePtr (*type_parser)(const std::string&) =
+        Unpickler::defaultTypeParser,
     std::shared_ptr<DeserializationStorageContext> storage_context = nullptr);
 
 bool check_zip_file(
diff --git a/torch/csrc/jit/serialization/pickle.cpp b/torch/csrc/jit/serialization/pickle.cpp
@@ -120,17 +120,19 @@ IValue pickle_load(const std::vector<char>& data) {
 IValue unpickle(
     std::function<size_t(char*, size_t)> reader,
     TypeResolver type_resolver,
-    c10::ArrayRef<at::Tensor> tensor_table) {
+    c10::ArrayRef<at::Tensor> tensor_table,
+    c10::TypePtr (*type_parser)(const std::string&)) {
   Unpickler unpickler(
-      std::move(reader), std::move(type_resolver), tensor_table);
+      std::move(reader), std::move(type_resolver), tensor_table, type_parser);
   return unpickler.parse_ivalue();
 }
 
 IValue unpickle(
     const char* data,
     size_t size,
     TypeResolver type_resolver,
-    c10::ArrayRef<at::Tensor> tensor_table) {
+    c10::ArrayRef<at::Tensor> tensor_table,
+    c10::TypePtr (*type_parser)(const std::string&)) {
   size_t bytes_read = 0;
   return unpickle(
       [&](char* buffer, size_t len) -> size_t {
@@ -145,7 +147,8 @@ IValue unpickle(
         return len;
       },
       std::move(type_resolver),
-      tensor_table);
+      tensor_table,
+      type_parser);
 }
 
 } // namespace jit
diff --git a/torch/csrc/jit/serialization/pickle.h b/torch/csrc/jit/serialization/pickle.h
@@ -69,7 +69,9 @@ TORCH_API IValue pickle_load(const std::vector<char>& data);
 TORCH_API IValue unpickle(
     std::function<size_t(char*, size_t)> reader,
     TypeResolver type_resolver,
-    c10::ArrayRef<at::Tensor> tensor_table);
+    c10::ArrayRef<at::Tensor> tensor_table,
+    c10::TypePtr (*type_parser)(const std::string&) =
+        Unpickler::defaultTypeParser);
 
 /// Decode a chunk of memory containing pickled data into its `torch::IValue`s.
 ///
@@ -81,7 +83,9 @@ TORCH_API IValue unpickle(
     const char* data,
     size_t size,
     TypeResolver type_resolver = nullptr,
-    c10::ArrayRef<at::Tensor> tensor_table = {});
+    c10::ArrayRef<at::Tensor> tensor_table = {},
+    c10::TypePtr (*type_parser)(const std::string&) =
+        Unpickler::defaultTypeParser);
 
 } // namespace jit
 } // namespace torch
diff --git a/torch/csrc/jit/serialization/unpickler.cpp b/torch/csrc/jit/serialization/unpickler.cpp
@@ -565,7 +565,7 @@ void Unpickler::readGlobal(
           if (type_resolver_ == nullptr) {
             // If we haven't injected a custom way of retrieving types from
             // names, use a barebones type parser.
-            type = c10::parseType(type_str);
+            type = type_parser_(type_str);
           } else {
             type = type_resolver_(type_str).type_;
           }
diff --git a/torch/csrc/jit/serialization/unpickler.h b/torch/csrc/jit/serialization/unpickler.h

Original file line number	Diff line number	Diff line change
`@@ -29,8 +29,6 @@ struct Tree;`
`29`	`29`	`using TreeRef = c10::intrusive_ptr<Tree>;`
`30`	`30`	`using TreeList = at::SmallVector<TreeRef, 4>;`
`31`	`31`
`32`		`-static const TreeList empty_trees = {};`
`33`		`-`
`34`	`32`	`struct Tree : c10::intrusive_ptr_target {`
`35`	`33`	`Tree(int kind_) : kind_(kind_) {}`
`36`	`34`	`int kind() const {`
`@@ -46,6 +44,7 @@ struct Tree : c10::intrusive_ptr_target {`
`46`	`44`	`throw std::runtime_error("stringValue can only be called on TK_STRING");`
`47`	`45`	`}`
`48`	`46`	`virtual const TreeList& trees() const {`
	`47`	`+ static const TreeList empty_trees = {};`
`49`	`48`	`return empty_trees;`
`50`	`49`	`}`
`51`	`50`	`const TreeRef& tree(size_t i) const {`
`@@ -149,11 +148,11 @@ struct Compound : public Tree {`
`149`	`148`	`return false;`
`150`	`149`	`}`
`151`	`150`	`TreeRef map(const std::function<TreeRef(TreeRef)>& fn) override {`
`152`		`- TreeList trees_;`
	`151`	`+ TreeList ret;`
`153`	`152`	`for (auto& t : trees()) {`
`154`		`- trees_.push_back(fn(t));`
	`153`	`+ ret.push_back(fn(t));`
`155`	`154`	`}`
`156`		`- return Compound::create(kind(), range(), std::move(trees_));`
	`155`	`+ return Compound::create(kind(), range(), std::move(ret));`
`157`	`156`	`}`
`158`	`157`
`159`	`158`	`const SourceRange& range() const override {`
Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,8 @@ c10::IValue readArchive(`
`53`	`53`	`type_resolver,`
`54`	`54`	`obj_loader,`
`55`	`55`	`device,`
`56`		`- stream_reader);`
	`56`	`+ stream_reader,`
	`57`	`+ nullptr);`
`57`	`58`	`return ivalues;`
`58`	`59`	`}`
`59`	`60`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+#include <torch/csrc/jit/mobile/type_parser.h>`
	`2`	`+`
`1`	`3`	`#include <ATen/core/jit_type.h>`
`2`	`4`	`#include <c10/util/string_view.h>`
`3`	`5`	`#include <torch/csrc/jit/frontend/parser_constants.h>`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+#pragma once`
	`2`	`+`
`1`	`3`	`#include <ATen/core/dynamic_type.h>`
`2`	`4`	`#include <ATen/core/jit_type.h>`
`3`	`5`