Introduce SinkHolder to keep nodes that do not return any outputs

Author: slyalin

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

@@ -23,15 +23,22 @@ def __init__(self, *args):
             # TODO: What about attributes?
             super().__init__(self, args)
             self.attrs = {"id": global_counter_id}  # `id` attribute distinguishes different instances of the same class, we need it because different instances may have different behaviour
-            # print(f'Made custom op class with id = {self.attrs["id"]}')
-            # print(f"Input signature: {input_signature}")
-            # print(f"Output signature: {output_signature}")
+            #print('output_signature from ctro:', output_signature)
+            if output_signature == ():
+                # The operation doesn't have outputs, so we need to take extra care to avoid eliminating the op from the graph
+                #print('===================== MARKING AS SINK ========================')
+                self.get_rt_info()['__sink__'] = True
+            #print(f'Made custom op class with id = {self.attrs["id"]}')
+            #print(f"Input signature: {input_signature}")
+            #print(f"Output signature: {output_signature}")
             self.constructor_validate_and_infer_types()
 
         def evaluate(self, outputs, inputs):
             # print("called evaluate")
             inputs_torch = (torch.from_numpy(input.data) for input in inputs)   # TODO: Check memory sharing
             result = func(*inputs_torch)
+            if result is None:
+                result = ()
             if not isinstance(result, tuple):
                 result = (result,)
             for i, tensor in enumerate(result):
@@ -47,8 +54,13 @@ def visit_attributes(self, visitor):
 
         def validate_and_infer_types(self):
             #TODO: Validate input signature
-            for i, output in enumerate(output_signature):
-                self.set_output_type(i, output[0], output[1])
+            if output_signature == ():
+                # Even when the original wrapped function doesn't give any return value, we need to set some output type to avoid eliminating the op from the graph
+                # Data type and shape doesn't matter, so we set default empty tensor of type u8
+                self.set_output_type(0, ov.Type.u8, ov.PartialShape([0]))
+            else:
+                for i, output in enumerate(output_signature):
+                    self.set_output_type(i, output[0], output[1])
     global_counter_id += 1
     return InlinedCustomOp
 
@@ -71,6 +83,7 @@ def make_input_signature(args, kwargs):
 def make_output_signature(args):
     if args is None:
         # TODO: This case is not really supported by PT FE -- because we don't support ops that do not have outputs, they will be lost
+        #print('=================== None PROCESSING ======================')
         args = ()
     if not isinstance(args, tuple):
         args = (args,)
@@ -108,6 +121,7 @@ def forward(ctx, *call_args, **call_kwargs):  #TODO: what is `ctx`?
             result = func_target(*call_args, **call_kwargs)
             if not op:
                 output_signature = make_output_signature(result)
+                #print('about to make custom op class with output signature', output_signature)
                 __class__.op = make_custom_op_class(func_target, input_signature, output_signature)
             else:
                 __class__.op = op
@@ -129,7 +143,9 @@ def trampoline(*args, **kwargs):
             # It is required because `func` is fused inside Trampoline class and can have different behaviour from call to call in PyTorch world even if
             # the same op is specified to wrap multiple different functions.
             trampoline = make_trampoline_class(func, op, op_attrs)
-            return trampoline.apply(*args, **kwargs)
+            result = trampoline.apply(*args, **kwargs)
+            #print('just called trampoline with result:', result)
+            return result
         return trampoline
 
     if len(args) == 1 and callable(args[0]) and not (isinstance(args[0], type) and issubclass(args[0], ov.Op)):

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

@@ -635,6 +635,11 @@ def convert(self, node_context):
         if op_extension := self.get_op_extension():
             trampoline, target_extension = op_extension
             assert isinstance(target_extension, InlineConversionExtension)
-            result = trampoline.convert(node_context)
-            return result
+            try:
+                return trampoline.convert(node_context)
+            except Exception as e:
+                print('[ ERROR ] Exception happened during calling of custom converter for PyTorch operation')
+                print('          PyTorch Script code:', self.graph_element)
+                print('          Exception:', e)
+                raise
         assert False, "PyTorch FrontEnd Internal Error: `converter` method of TorchScriptPythonDecoder is called for node that has no custom converter"

src/frontends/pytorch/src/translate_session.cpp

@@ -45,6 +45,32 @@ std::shared_ptr<ov::Model> TranslateSession::get_converted_model() {
     return m_ov_model;
 }
 
+/// @brief Sink operation that preserves another node from removing from the graph, when it is not possible to derive that target node op from Sink class
+class SinkHolder : public Sink {
+public:
+    OPENVINO_OP("SinkHolder", "util", ov::op::Sink);
+    SinkHolder() = default;
+    explicit SinkHolder(const OutputVector& arguments) : Sink(arguments) {}
+
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override {
+        check_new_args_count(this, new_args);
+        return std::make_shared<SinkHolder>(new_args);
+    }
+
+    void validate_and_infer_types() override {
+        set_output_type(0, element::u8, PartialShape({0}));
+    }
+
+    bool evaluate(ov::TensorVector& output_values, const ov::TensorVector& input_values) const override {
+        // This op has a dummy output -- no need to evaluate
+        return true;
+    }
+
+    bool has_evaluate() const override {
+        return true;
+    }
+};
+
 std::shared_ptr<ov::Model> TranslateSession::translate_graph(const ov::frontend::InputModel::Ptr& input_model) {
     auto pytorch_model = std::dynamic_pointer_cast<pytorch::InputModel>(input_model);
     FRONT_END_GENERAL_CHECK(pytorch_model != nullptr, "Invalid input model");
@@ -88,6 +114,7 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
         auto tensor_map = std::make_shared<TensorMap>();  // tensor map of the current context
         auto mutated_tensors = std::make_shared<std::set<size_t>>();
         std::vector<size_t> inserted_params;
+        SinkVector sinks;
 
         if (input_model && input_model->m_requested_places.size() == 0) {
             // When we have input model we should use its inputs order to create Parameters
@@ -179,6 +206,12 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
 
             const bool has_inputs = !node->inputs().empty();
             const size_t in_tensor_id = has_inputs ? node->inputs().at(0) : 0;
+            if (converted_outputs.size() == 1) {
+                // Check if converted node is marked as a sink
+                if(converted_outputs.front().get_node()->get_rt_info().count("__sink__")) {
+                    sinks.push_back(std::make_shared<SinkHolder>(converted_outputs));
+                }
+            }
             for (size_t i = 0; i < fw_outputs.size(); ++i) {
                 size_t fw_tensor_id = node->output(i);
                 if (has_inputs && node->may_produce_alias(0, i)) {
@@ -303,7 +336,7 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
                                              }),
                               parameters->end());
         }
-        resulting_model = std::make_shared<Model>(results, *parameters);
+        resulting_model = std::make_shared<Model>(results, sinks, *parameters);
         // Did a conversion in a nested scope to automatically remove any holders of nodes except those in the graph
     }