[PT FE] torch.export support (openvinotoolkit#22397)

* [PT FE] torch.export support

* Apply code style

* Fix build

* Support torch <2.2

* Support more operations

* Update tests/model_hub_tests/torch_tests/torch_utils.py

* Support for model operations

* Support is_causal as kwarg for SDPA

* Update src/frontends/pytorch/src/op/addcmul.cpp

* Update tests/model_hub_tests/torch_tests/test_timm.py

* Support only decoder passed to convert_model

* Fix tests

* Update src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py

* Update src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py

Co-authored-by: Ekaterina Aidova <ekaterina.aidova@intel.com>

* Apply suggestions from code review

* Apply suggestions from code review

* Improve testing with caching model

* Apply suggestions from code review

---------

Co-authored-by: Ekaterina Aidova <ekaterina.aidova@intel.com>

Author: mvafin

.github/workflows/job_python_unit_tests.yml

@@ -216,6 +216,15 @@ jobs:
           TEST_DEVICE: CPU
           TEST_PRECISION: FP32
 
+      - name: PyTorch torch.export Layer Tests
+        if: ${{ fromJSON(inputs.affected-components).PyTorch_FE.test && runner.arch != 'ARM64' }} # Ticket: 126287
+        run: |
+          python3 -m pytest ${LAYER_TESTS_INSTALL_DIR}/pytorch_tests -m precommit_torch_export --junitxml=${INSTALL_TEST_DIR}/TEST-pytorch.xml
+        env:
+          TEST_DEVICE: CPU
+          TEST_PRECISION: FP32
+          PYTORCH_TRACING_MODE: EXPORT
+
       - name: PyTorch torch.compile TORCHFX Layer Tests
         if: ${{ fromJSON(inputs.affected-components).PyTorch_FE.test && runner.os != 'macOS' }}
         run: |

src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py

@@ -406,6 +406,18 @@ def may_produce_alias(self, in_index: int, out_index: int) -> bool:
     def inlined_inputs(self, index):
         return []
 
+    def inlined_input(self, index):
+        return []
+
+    def is_input_inlined(self, index):
+        return False
+
+    def get_attribute(self, name):
+        return OVAny(None)
+
+    def get_named_input(self, name):
+        raise RuntimeError("There is no named inputs in TS graph")
+
     @staticmethod
     def _transform_tensor_list_constants_to_listconstruct(graph: torch.Graph):
         # Function replaces prim::Constant containing List of Tensors with

src/bindings/python/src/openvino/frontend/pytorch/ts_decoder.py

@@ -137,8 +137,13 @@ def graph_has_ops(graph, op_types:list) -> bool:
     "torch.bool": OVType.boolean,
     "torch.DoubleTensor": OVType.f64,
     "torch.FloatTensor": OVType.f32,
+    "torch.HalfTensor": OVType.f16,
+    "torch.BFloat16Tensor": OVType.bf16,
     "torch.IntTensor": OVType.i32,
     "torch.LongTensor": OVType.i64,
+    "torch.ShortTensor": OVType.i16,
+    "torch.CharTensor": OVType.i8,
+    "torch.ByteTensor": OVType.u8,
     "torch.BoolTensor": OVType.boolean,
     "torch.quint8": OVType.u8,
     "torch.qint8": OVType.i8,

src/bindings/python/src/openvino/frontend/pytorch/utils.py

@@ -110,8 +110,21 @@ class PyDecoder : public ov::frontend::pytorch::TorchDecoder {
         PYBIND11_OVERRIDE_PURE(bool, TorchDecoder, may_produce_alias, in_index, out_index);
     }
 
-    ov::OutputVector inlined_inputs(size_t start_index) const override {
-        PYBIND11_OVERRIDE_PURE(ov::OutputVector, TorchDecoder, inlined_inputs, start_index); }
+    ov::OutputVector inlined_input(size_t index) const override {
+        PYBIND11_OVERRIDE_PURE(ov::OutputVector, TorchDecoder, inlined_input, index);
+    }
+
+    bool is_input_inlined(size_t index) const override {
+        PYBIND11_OVERRIDE_PURE(bool, TorchDecoder, is_input_inlined, index);
+    }
+
+    ov::Any get_attribute(const std::string &name) const override{
+        PYBIND11_OVERRIDE_PURE(ov::Any, TorchDecoder, get_attribute, name);
+    }
+
+    size_t get_named_input(const std::string &name) const override{
+        PYBIND11_OVERRIDE_PURE(size_t, TorchDecoder, get_named_input, name);
+    }
 
     const std::string& decoder_type_name() const override {
         PYBIND11_OVERRIDE_PURE(const std::string&, TorchDecoder, decoder_type_name);

src/bindings/python/src/pyopenvino/frontend/pytorch/decoder.hpp

@@ -375,6 +375,8 @@ ov::Any py_object_to_any(const py::object& py_obj) {
         return py::cast<ov::Affinity>(py_obj);
     } else if (py::isinstance<ov::Tensor>(py_obj)) {
         return py::cast<ov::Tensor>(py_obj);
+    } else if (py::isinstance<ov::Output<ov::Node>>(py_obj)) {
+        return py::cast<ov::Output<ov::Node>>(py_obj);
         // FrontEnd Decoder
     } else if (py::isinstance<ov::frontend::IDecoder>(py_obj)) {
         return py::cast<std::shared_ptr<ov::frontend::IDecoder>>(py_obj);

src/bindings/python/src/pyopenvino/utils/utils.cpp

@@ -19,7 +19,7 @@ class TorchDecoder : public IDecoder {
     // fundamental types like int, float etc.
     virtual Any const_input(size_t index) const = 0;
 
-    // Using size_t for input/output unuque ids are in sync with torch code, see def in
+    // Using size_t for input/output unique ids are in sync with torch code, see def in
     // torch/include/torch/csrc/jit/ir/ir.h, Value::unique_
 
     // TODO: set of input and output methods are not aligned; also they are not aligned with the rest of FEs
@@ -89,7 +89,7 @@ class TorchDecoder : public IDecoder {
     virtual size_t output(size_t index) const = 0;
 
     // Embed mapping to/from the original node representation from/to node passed as a parameter
-    // the representation of this mapping is specific for particular decored type and may be NOP
+    // the representation of this mapping is specific for particular decorated type and may be NOP
     // returns the same node as syntactically convenient way to make nested sentences in code
     virtual std::shared_ptr<Node> mark_node(std::shared_ptr<Node> ov_node) const = 0;
 
@@ -109,9 +109,19 @@ class TorchDecoder : public IDecoder {
 
     /// Returns new nodes for inputs inlined in the op itself
     // Used in Torch.FX decoder
-    virtual OutputVector inlined_inputs(size_t start_index) const = 0;
+    virtual OutputVector inlined_input(size_t index) const = 0;
 
-    /// Returns the id of the deccoder type (0: TorchFX, 1: TorchScript)
+    /// Returns if input is inlined
+    // Used in Torch.FX decoder
+    virtual bool is_input_inlined(size_t index) const = 0;
+
+    /// Returns named attribute as Any. For example kwargs input for FX graph
+    virtual ov::Any get_attribute(const std::string& name) const = 0;
+
+    /// Returns index of named input. For example kwargs input for FX graph
+    virtual size_t get_named_input(const std::string& name) const = 0;
+
+    /// Returns the id of the decoder type ("fx": TorchFX, "ts": TorchScript)
     virtual const std::string& decoder_type_name() const = 0;
 };
 

src/frontends/pytorch/include/openvino/frontend/pytorch/decoder.hpp

@@ -51,12 +51,36 @@ class NodeContext : public frontend::NodeContext {
     // Search for input in tensor map and return an output port for already converted op
     // TODO: int due to base class uses it, but naturally it should be size_t for PT
     Output<Node> get_input(int index) const override {
-        FRONT_END_GENERAL_CHECK(!input_is_none(index), "Input is none with index: ", index);
+        size_t index_ = static_cast<size_t>(index);
+        FRONT_END_GENERAL_CHECK(!input_is_none(index_), "Input doesn't exist with index: ", index);
         auto input = m_decoder_inputs.at(index);
+        if (input == 0) {
+            // Case when input can be inlined (possible only for fx decoder)
+            if (m_decoder->is_input_inlined(index_)) {
+                auto inlined_input = m_decoder->inlined_input(index_);
+                FRONT_END_GENERAL_CHECK(inlined_input.size() == 1, "Incorrect inlined input with index:", index);
+                return inlined_input[0];
+            }
+        }
         FRONT_END_GENERAL_CHECK(m_tensor_map->count(input), "No tensor corresponding input: ", input, " exist.");
         return m_tensor_map->at(input);
     }
 
+    Output<Node> get_input(const std::string& name) const override {
+        FRONT_END_GENERAL_CHECK(has_attribute(name), "Input with name ", name, " doesn't exist");
+        auto attr = get_attribute_as_any(name);
+        if (attr.is<Output<Node>>()) {
+            // Case when input is constant value
+            return attr.as<Output<Node>>();
+        } else if (attr.is<type::PyNone>()) {
+            // None means input is unknown type, most likely a Node
+            auto input = m_decoder->get_named_input(name);
+            FRONT_END_GENERAL_CHECK(m_tensor_map->count(input), "No tensor corresponding input: ", input, " exist.");
+            return m_tensor_map->at(input);
+        }
+        FRONT_END_GENERAL_CHECK(false, "Input has type which can't be converted to ov::Node.");
+    }
+
     Any get_values_from_const_input(int index) const override;
 
     // TODO: upstream to base class
@@ -112,9 +136,8 @@ class NodeContext : public frontend::NodeContext {
         return ov_output;
     }
 
-    Any get_attribute_as_any(const std::string&) const override {
-        throw std::runtime_error(
-            "There is no any named attributes in PyTorch node, query by attribute name is not implemented");
+    Any get_attribute_as_any(const std::string& name) const override {
+        return m_decoder->get_attribute(name);
     }
 
     void mutate_input(size_t index, Output<Node> ov_output) const;

src/frontends/pytorch/include/openvino/frontend/pytorch/node_context.hpp

@@ -194,6 +194,7 @@ void FrontEnd::normalize(const std::shared_ptr<ov::Model>& model) const {
     manager.register_pass<ov::frontend::pytorch::pass::PrimListUnpackReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::AtenGetItemReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::ListConstructReplacer>();
+    // TODO: remove AtenIndexToSelect when problem with  dynamic input rank is gone.
     manager.register_pass<ov::frontend::pytorch::pass::AtenIndexToSelect>();
     manager.register_pass<ov::frontend::pytorch::pass::AtenIndexPutReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::PrimListConstructPadReplacer>();

src/frontends/pytorch/src/frontend.cpp

@@ -109,7 +109,7 @@ Output<Node> NodeContext::get_tensor_from_model_or_create_input(size_t index) co
 }
 
 Output<Node> NodeContext::get_input_from_visible_context(size_t index) const {
-    FRONT_END_GENERAL_CHECK(index < get_input_size(), "Index is lower then number of inputs.");
+    FRONT_END_GENERAL_CHECK(index < get_input_size(), "Index ", index, " is lower then number of inputs.");
     auto input_tensor = get_input(static_cast<int>(index));
     auto input_node = input_tensor.get_node_shared_ptr();
     if (std::dynamic_pointer_cast<v0::Parameter>(input_node)) {

src/frontends/pytorch/src/node_context.cpp

@@ -117,6 +117,21 @@ OutputVector translate_adaptive_max_pool1d(const NodeContext& context) {
     return translate_adaptive_max_pool_base(context, const_tile_params, const_neg_1);
 };
 
+OutputVector translate_adaptive_max_pool3d_fx(const NodeContext& context) {
+    auto outs = translate_adaptive_max_pool3d(context);
+    return {context.mark_node(make_list_construct(outs))};
+};
+
+OutputVector translate_adaptive_max_pool2d_fx(const NodeContext& context) {
+    auto outs = translate_adaptive_max_pool2d(context);
+    return {context.mark_node(make_list_construct(outs))};
+};
+
+OutputVector translate_adaptive_max_pool1d_fx(const NodeContext& context) {
+    auto outs = translate_adaptive_max_pool1d(context);
+    return {context.mark_node(make_list_construct(outs))};
+};
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend

src/frontends/pytorch/src/op/adaptive_poolnd.cpp

@@ -34,8 +34,14 @@ OutputVector translate_add_common(const NodeContext& context, bool inplace) {
     } else {
         align_eltwise_input_types(context, lhs, rhs, true);
     }
+    Output<Node> alpha;
     if (!context.input_is_none(2)) {
-        auto converted_alpha = context.mark_node(std::make_shared<v1::ConvertLike>(context.get_input(2), rhs));
+        alpha = context.get_input(2);
+    } else if (context.has_attribute("alpha")) {
+        alpha = context.get_attribute<Output<Node>>("alpha");
+    }
+    if (alpha.get_node_shared_ptr()) {
+        auto converted_alpha = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, rhs));
         rhs = context.mark_node(std::make_shared<v1::Multiply>(converted_alpha, rhs));
     }
     auto add = context.mark_node(std::make_shared<v1::Add>(lhs, rhs));

src/frontends/pytorch/src/op/add.cpp

@@ -17,15 +17,30 @@ namespace op {
 
 using namespace ov::op;
 
-OutputVector translate_addcmul(const NodeContext& context) {
-    num_inputs_check(context, 4, 4);
+namespace {
+OutputVector addcmul_common(const NodeContext& context, const Output<Node>& value) {
     const auto eltwise_mult = std::make_shared<v1::Multiply>(context.get_input(1), context.get_input(2));
-    const auto value = context.get_input(3);
     const auto converted_value = std::make_shared<v1::ConvertLike>(value, context.get_input(1));
     const auto scalar_mult = std::make_shared<v1::Multiply>(eltwise_mult, converted_value);
     context.mark_nodes({eltwise_mult, converted_value, scalar_mult});
     return {context.mark_node(std::make_shared<v1::Add>(context.get_input(0), scalar_mult))};
 };
+}  // namespace
+
+OutputVector translate_addcmul(const NodeContext& context) {
+    num_inputs_check(context, 4, 4);
+    const auto value = context.get_input(3);
+    return addcmul_common(context, value);
+};
+
+OutputVector translate_addcmul_fx(const NodeContext& context) {
+    num_inputs_check(context, 3, 3);
+    Output<Node> value = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1}));
+    if (context.has_attribute("value")) {
+        value = context.get_input("value");
+    }
+    return addcmul_common(context, value);
+};
 
 }  // namespace op
 }  // namespace pytorch

src/frontends/pytorch/src/op/addcmul.cpp

@@ -16,12 +16,22 @@ namespace op {
 
 using namespace ov::op;
 
-OutputVector translate_addmm(const NodeContext& context) {
-    num_inputs_check(context, 3, 5);
+namespace {
+OutputVector translate_addmm_common(const NodeContext& context, const Output<Node> beta, const Output<Node> alpha) {
     auto input = context.get_input(0);
     auto m1 = context.get_input(1);
     auto m2 = context.get_input(2);
     auto mm = context.mark_node(std::make_shared<v0::MatMul>(m1, m2));
+    auto beta_converted = context.mark_node(std::make_shared<v1::ConvertLike>(beta, input));
+    auto alpha_converted = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, mm));
+    auto input_beta = context.mark_node(std::make_shared<v1::Multiply>(input, beta_converted));
+    auto mm_alpha = context.mark_node(std::make_shared<v1::Multiply>(mm, alpha_converted));
+    return {context.mark_node(std::make_shared<v1::Add>(input_beta, mm_alpha))};
+};
+}  // namespace
+
+OutputVector translate_addmm(const NodeContext& context) {
+    num_inputs_check(context, 3, 5);
     auto one = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1}));
     ov::Output<Node> alpha = one;
     ov::Output<Node> beta = one;
@@ -31,11 +41,21 @@ OutputVector translate_addmm(const NodeContext& context) {
     if (!context.input_is_none(4)) {
         alpha = context.get_input(4);
     }
-    auto beta_converted = context.mark_node(std::make_shared<v1::ConvertLike>(beta, input));
-    auto alpha_converted = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, mm));
-    auto input_beta = context.mark_node(std::make_shared<v1::Multiply>(input, beta_converted));
-    auto mm_alpha = context.mark_node(std::make_shared<v1::Multiply>(mm, alpha_converted));
-    return {context.mark_node(std::make_shared<v1::Add>(input_beta, mm_alpha))};
+    return {translate_addmm_common(context, beta, alpha)};
+};
+
+OutputVector translate_addmm_fx(const NodeContext& context) {
+    num_inputs_check(context, 3, 3);
+    auto one = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1}));
+    ov::Output<Node> alpha = one;
+    ov::Output<Node> beta = one;
+    if (context.has_attribute("beta")) {
+        beta = context.get_input("beta");
+    }
+    if (context.has_attribute("alpha")) {
+        alpha = context.get_input("alpha");
+    }
+    return {translate_addmm_common(context, beta, alpha)};
 };
 
 }  // namespace op