[Pytorch Edge] Model Ops compatibility api (pytorch#57501)

Summary:
Pull Request resolved: pytorch#57501

Add an api _get_model_ops_and_info to get root operators and versioning info of a model in both cxx and python, and the input can be from a file path or buffer.
ghstack-source-id: 129620112

Test Plan: unit test.

Reviewed By: xcheng16, raziel

Differential Revision: D28162765

fbshipit-source-id: 4413c1e906b8a872e4a717d849da37347adbbea4

Author: JacobSzwejbka

test/cpp/jit/script_module_v6.ptl

@@ -8,12 +8,13 @@
 from torch.jit.mobile import (
     _load_for_lite_interpreter,
     _get_model_bytecode_version,
+    _get_model_ops_and_info,
     _backport_for_mobile_to_buffer,
     _backport_for_mobile)
 from torch.testing._internal.common_utils import TestCase, run_tests
 from pathlib import Path
 
-pytorch_test_dri = Path(__file__).resolve().parents[1]
+pytorch_test_dir = Path(__file__).resolve().parents[1]
 
 # script_module_v4.ptl and script_module_v5.ptl source code
 # class TestModule(torch.nn.Module):
@@ -97,6 +98,38 @@
     ((('name', ''), ('type', 'Tensor'), ('default_value', None)),)))))
         '''
 
+SCRIPT_MODULE_V6_BYTECODE_PKL = '''
+(6,
+ ('__torch__.*.TestModule.forward',
+  (('instructions',
+    (('STOREN', 1, 2),
+     ('DROPR', 1, 0),
+     ('LOADC', 0, 0),
+     ('LOADC', 1, 0),
+     ('MOVE', 2, 0),
+     ('OP', 0, 0),
+     ('OP', 1, 0),
+     ('RET', 0, 0))),
+   ('operators', (('aten::add', 'int', 2), ('aten::add', 'Scalar', 2))),
+   ('constants',
+    (torch._utils._rebuild_tensor_v2(pers.obj(('storage', torch.DoubleStorage, '0', 'cpu', 8),),
+       0,
+       (2, 4),
+       (4, 1),
+       False,
+       collections.OrderedDict()),
+     1)),
+   ('types', ()),
+   ('register_size', 2)),
+  (('arguments',
+    ((('name', 'self'),
+      ('type', '__torch__.*.TestModule'),
+      ('default_value', None)),
+     (('name', 'y'), ('type', 'int'), ('default_value', None)))),
+   ('returns',
+    ((('name', ''), ('type', 'Tensor'), ('default_value', None)),)))))
+    '''
+
 SCRIPT_MODULE_BYTECODE_PKL = {
     4: {
         "bytecode_pkl": SCRIPT_MODULE_V4_BYTECODE_PKL,
@@ -113,7 +146,7 @@ def check_model_version(model_path, expect_version):
             actual_version = _get_model_bytecode_version(model_path)
             assert(actual_version == expect_version)
         for version, model_info in SCRIPT_MODULE_BYTECODE_PKL.items():
-            model_path = pytorch_test_dri / "cpp" / "jit" / model_info["model_name"]
+            model_path = pytorch_test_dir / "cpp" / "jit" / model_info["model_name"]
             check_model_version(model_path, version)
 
     def test_bytecode_values_for_all_backport_functions(self):
@@ -130,7 +163,7 @@ def test_bytecode_values_for_all_backport_functions(self):
             while current_from_version > MINIMUM_TO_VERSION:
                 # Load model v5 and run forward method
                 model_name = SCRIPT_MODULE_BYTECODE_PKL[current_from_version]["model_name"]
-                input_model_path = pytorch_test_dri / "cpp" / "jit" / model_name
+                input_model_path = pytorch_test_dir / "cpp" / "jit" / model_name
 
                 # A temporary model file will be export to this path, and run through bytecode.pkl
                 # content check.
@@ -205,7 +238,7 @@ def forward(self, y: int):
     # Check just the test_backport_bytecode_from_file_to_file mechanism but not the function implementations
     def test_backport_bytecode_from_file_to_file(self):
         maximum_checked_in_model_version = max(SCRIPT_MODULE_BYTECODE_PKL.keys())
-        script_module_v5_path = pytorch_test_dri / "cpp" / "jit" / SCRIPT_MODULE_BYTECODE_PKL[
+        script_module_v5_path = pytorch_test_dir / "cpp" / "jit" / SCRIPT_MODULE_BYTECODE_PKL[
             maximum_checked_in_model_version]["model_name"]
 
         if (maximum_checked_in_model_version > MINIMUM_TO_VERSION):
@@ -241,7 +274,7 @@ def test_backport_bytecode_from_file_to_file(self):
     # Check just the _backport_for_mobile_to_buffer mechanism but not the function implementations
     def test_backport_bytecode_from_file_to_buffer(self):
         maximum_checked_in_model_version = max(SCRIPT_MODULE_BYTECODE_PKL.keys())
-        script_module_v5_path = pytorch_test_dri / "cpp" / "jit" / SCRIPT_MODULE_BYTECODE_PKL[
+        script_module_v5_path = pytorch_test_dir / "cpp" / "jit" / SCRIPT_MODULE_BYTECODE_PKL[
             maximum_checked_in_model_version]["model_name"]
 
         if (maximum_checked_in_model_version > MINIMUM_TO_VERSION):
@@ -264,5 +297,12 @@ def test_backport_bytecode_from_file_to_buffer(self):
             torch.testing.assert_allclose(mobile_module_result, expected_mobile_module_result)
 
 
+    def test_get_model_ops_and_info(self):
+        # TODO update this to be more in the style of the above tests after a backport from 6 -> 5 exists
+        script_module_v6 = pytorch_test_dir / "cpp" / "jit" / "script_module_v6.ptl"
+        ops_v6 = _get_model_ops_and_info(script_module_v6)
+        assert(ops_v6["aten::add.int"].num_schema_args == 2)
+        assert(ops_v6["aten::add.Scalar"].num_schema_args == 2)
+
 if __name__ == '__main__':
     run_tests()

test/mobile/test_bytecode.py

@@ -260,6 +260,8 @@ def _backport_for_mobile(filename_input: Union[str, Path], filename_output: Unio
 def _backport_for_mobile_from_buffer(buffer: BinaryIO, filename_output: Union[str, Path], to_version: _int) -> None: ...
 def _backport_for_mobile_to_buffer(filename_input: Union[str, Path], to_version: _int) -> bytes:...
 def _backport_for_mobile_from_buffer_to_buffer(buffer: BinaryIO, to_version: _int) -> bytes:...
+def _get_model_ops_and_info(filename: Union[str, Path]): ...
+def _get_model_ops_and_info_from_buffer(buffer: BinaryIO): ...
 def _logging_set_logger(logger: LoggerBase) -> LoggerBase: ...
 def _get_graph_executor_optimize() -> _bool: ...
 def _set_graph_executor_optimize(optimize: _bool): ...

torch/_C/__init__.pyi.in

@@ -149,7 +149,7 @@ bool backport_v5_to_v4(
     PyTorchStreamReader& reader,
     PyTorchStreamWriter& writer) {
   // 1) read from archive `bytecode` archive
-  std::vector<IValue> bytecode_values = get_bytecode_values(reader);
+  std::vector<IValue> bytecode_values = get_bytecode_ivalues(reader);
   if (!check_bytecode_version(bytecode_values, kBytecodeVersionV5)) {
     TORCH_WARN("Incorrect bytecode version for input model.");
     return false;

torch/csrc/jit/mobile/backport_manager.cpp

@@ -49,12 +49,14 @@ c10::IValue readArchive(
   return ivalues;
 }
 
-std::vector<IValue> get_bytecode_values(PyTorchStreamReader& reader) {
+std::vector<IValue> get_bytecode_ivalues(PyTorchStreamReader& reader) {
   std::vector<IValue> bytecode_values;
   bytecode_values = readArchive("bytecode", reader).toTuple()->elements();
   return bytecode_values;
 }
 
+/********************** Bytecode **********************/
+
 // Forward declare
 int64_t _get_model_bytecode_version(
     const std::vector<IValue>& bytecode_ivalues);
@@ -76,7 +78,7 @@ int64_t _get_model_bytecode_version(std::shared_ptr<ReadAdapterInterface> rai) {
     return -1;
   }
   PyTorchStreamReader reader(std::move(rai));
-  auto bytecode_values = get_bytecode_values(reader);
+  auto bytecode_values = get_bytecode_ivalues(reader);
   return _get_model_bytecode_version(bytecode_values);
 }
 
@@ -90,5 +92,87 @@ int64_t _get_model_bytecode_version(
   return -1;
 }
 
+/********************** Operators and Info **********************/
+
+// Forward declare
+std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    std::vector<IValue> bytecode_ivalues);
+
+std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    std::istream& in) {
+  std::unique_ptr<IStreamAdapter> rai = std::make_unique<IStreamAdapter>(&in);
+  return _get_model_ops_and_info(std::move(rai));
+}
+
+std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    const std::string& filename) {
+  std::unique_ptr<FileAdapter> rai = std::make_unique<FileAdapter>(filename);
+  return _get_model_ops_and_info(std::move(rai));
+}
+
+std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    std::shared_ptr<ReadAdapterInterface> rai) {
+  if (!check_zip_file(rai)) {
+    TORCH_WARN("Failed to open zip file for model ops.");
+    return std::unordered_map<std::string, OperatorInfo>{};
+  }
+  PyTorchStreamReader reader(std::move(rai));
+  auto bytecode_values = get_bytecode_ivalues(reader);
+  return _get_model_ops_and_info(bytecode_values);
+}
+
+/* A function to retrieve the root (top level) operators of a model and their
+ * corresponding compatibility info. These root operators can call other
+ * operators within them (traced ops), and a root op can call many different
+ * traced ops depending on internal code paths in the root op. These traced ops
+ * are not returned by this function. Those operators are abstracted into the
+ * runtime as an implementation detail (and the traced ops themselves can also
+ * call other operators) making retrieving them difficult and their value from
+ * this api negligible since they will differ between which runtime version the
+ * model is run on. Because of this, there is a false positive this api can't
+ * prevent in a compatibility usecase. All the root ops of a model are present
+ * in a target runtime, but not all the traced ops are which prevents a model
+ * from being able to run.
+ **/
+std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    std::vector<IValue> bytecode_ivalues) {
+  constexpr uint64_t min_version_with_schema = 6;
+  if (_get_model_bytecode_version(bytecode_ivalues) < min_version_with_schema) {
+    TORCH_WARN(
+        "Only models with bytecode version 6 and above contain operator schema information. Please re-export your model to generate it");
+  }
+  std::unordered_map<std::string, OperatorInfo> result;
+  if (bytecode_ivalues.empty()) {
+    TORCH_WARN("Failed to get model ops and info.");
+    return result;
+  }
+  // loop over all the functions in the bytecode
+  for (int i = 1; i < bytecode_ivalues.size(); i++) {
+    // descend to the operators list
+    auto method_tuple = bytecode_ivalues.at(i).toTuple()->elements();
+    auto operators_tuple = method_tuple.at(1).toTuple()->elements()[1];
+    auto operators = operators_tuple.toTuple()->elements()[1];
+    for (auto& op_tuple : operators.toTuple()->elements()) {
+      auto op = op_tuple.toTuple()->elements();
+
+      // grab name
+      std::string op_name = op.at(0).toStringRef();
+      std::string op_overload_name = op.at(1).toStringRef();
+      if (op_overload_name != "") {
+        op_name.append(".");
+        op_name.append(op_overload_name);
+      }
+
+      // grab schema size
+      if (op.size() > 2) {
+        result.emplace(op_name, OperatorInfo{(int)op.at(2).toInt()});
+      } else { // no schema information use default
+        result.emplace(op_name, OperatorInfo{});
+      }
+    }
+  }
+  return result;
+}
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/mobile/model_compatibility.cpp

@@ -1,7 +1,10 @@
 #pragma once
 
+#include <torch/csrc/jit/mobile/runtime_compatibility.h>
+
 #include <istream>
 #include <memory>
+#include <unordered_map>
 
 namespace caffe2 {
 namespace serialize {
@@ -24,7 +27,7 @@ TORCH_API int64_t _get_model_bytecode_version(
 int64_t _get_model_bytecode_version(
     const std::vector<c10::IValue>& bytecode_ivalues);
 
-std::vector<c10::IValue> get_bytecode_values(
+std::vector<c10::IValue> get_bytecode_ivalues(
     caffe2::serialize::PyTorchStreamReader& reader);
 
 c10::IValue readArchive(
@@ -34,5 +37,15 @@ c10::IValue readArchive(
 bool check_zip_file(
     std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);
 
+// The family of methods below to get the root ops and information from a model
+TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    std::istream& in);
+
+TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    const std::string& filename);
+
+TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/mobile/model_compatibility.h

@@ -1,12 +1,16 @@
 #pragma once
 
+#include <c10/util/Optional.h>
+
 #include <memory>
 #include <unordered_map>
 
 namespace torch {
 namespace jit {
 
+// Struct storing metadata of an operator that can be useful for versioning
 struct OperatorInfo {
+  // The number of arguments within the schema of the op
   c10::optional<int> num_schema_args;
 };
 

torch/csrc/jit/mobile/runtime_compatibility.h

@@ -1770,6 +1770,15 @@ void initJitScriptBindings(PyObject* module) {
         std::istringstream in(buffer);
         return _get_model_bytecode_version(in);
       });
+  py::class_<OperatorInfo>(m, "OperatorInfo")
+      .def_readonly("num_schema_args", &OperatorInfo::num_schema_args);
+  m.def("_get_model_ops_and_info", [](const std::string& filename) {
+    return _get_model_ops_and_info(filename);
+  });
+  m.def("_get_model_ops_and_info_from_buffer", [](const std::string& buffer) {
+    std::istringstream in(buffer);
+    return _get_model_ops_and_info(in);
+  });
   m.def("_export_operator_list", [](torch::jit::mobile::Module& sm) {
     return debugMakeSet(torch::jit::mobile::_export_operator_list(sm));
   });

torch/csrc/jit/python/script_init.cpp

@@ -145,3 +145,43 @@ def _backport_for_mobile_to_buffer(f_input, to_version):
         return torch._C._backport_for_mobile_to_buffer(str(f_input), to_version)  # type: ignore[attr-defined]
     else:
         return torch._C._backport_for_mobile_from_buffer_to_buffer(f_input.read(), to_version)  # type: ignore[attr-defined]
+
+def _get_model_ops_and_info(f_input):
+    r"""
+    A function to retrieve the root (top level) operators of a model and their corresponding
+    compatibility info. These root operators can call other operators within them (traced ops), and
+    a root op can call many different traced ops depending on internal code paths in the root op.
+    These traced ops are not returned by this function. Those operators are abstracted into the
+    runtime as an implementation detail (and the traced ops themselves can also call other operators)
+    making retrieving them difficult and their value from this api negligible since they will differ
+    between which runtime version the model is run on. Because of this, there is a false positive this
+    api can't prevent in a compatibility usecase. All the root ops of a model are present in a
+    target runtime, but not all the traced ops are which prevents a model from being able to run.
+    Args:
+        f_input: a file-like object (has to implement read, readline, tell, and seek),
+            or a string containing a file name
+
+    Returns:
+        Operators and info: A Dictionary mapping strings (the qualified names of the root operators)
+        of the model to their OperatorInfo structs.
+
+    Example:
+
+    .. testcode::
+
+        from torch.jit.mobile import _get_model_ops_and_info
+
+        # Get bytecode version from a saved file path
+        ops_and_info = _get_model_ops_and_info("path/to/model.ptl")
+
+    """
+    if isinstance(f_input, str):
+        if not os.path.exists(f_input):
+            raise ValueError(f"The provided filename {f_input} does not exist")
+        if os.path.isdir(f_input):
+            raise ValueError(f"The provided filename {f_input} is a directory")
+
+    if (isinstance(f_input, str) or isinstance(f_input, pathlib.Path)):
+        return torch._C._get_model_ops_and_info(str(f_input))  # type: ignore[attr-defined]
+    else:
+        return torch._C._get_model_ops_and_info(f_input.read())  # type: ignore[attr-defined]