Support arbitrary nested tuples, lists and dicts in inputs and outputs of inlined custom ops. Also supported kwargs for original functions.

Author: slyalin

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

@@ -9,10 +9,69 @@
 from openvino.frontend.pytorch.utils import pt_to_ov_type_map
 import openvino as ov
 
+
+class ConstWrap:
+    def __init__(self, value):
+        self.value = value
+    def __eq__(self, x):
+        return self.value == x
+
+
+def unpack(packed, types, index=0):
+    unpacked_result = ()
+    if isinstance(packed, tuple):
+        packer_result = ()
+        for el in packed:
+            unpacked, packer, index = unpack(el, types, index)
+            unpacked_result += unpacked
+            packer_result += (packer,)
+    elif isinstance(packed, list):
+        packer_result = []
+        for el in packed:
+            unpacked, packer, index = unpack(el, types, index)
+            packer_result.append(packer)
+    elif isinstance(packed, dict):
+        packer_result = {}
+        for k, v in packed.items():
+            unpacked, packer, index = unpack(v, types, index)
+            unpacked_result += unpacked
+            packer_result[k] = packer
+    elif isinstance(packed, types):
+        unpacked_result = (packed,)
+        packer_result = index
+        index += 1
+    else:
+        packer_result = ConstWrap(packed)
+    return unpacked_result, packer_result, index
+
+
+def pack(unpacked, packer):
+    if isinstance(packer, tuple):
+        packed_result = ()
+        for el in packer:
+            packed = pack(unpacked, el)
+            packed_result += (packed,)
+    elif isinstance(packer, list):
+        packed_result = []
+        for el in packer:
+            packed = pack(unpacked, el)
+            packed_result.append(packed)
+    elif isinstance(packer, dict):
+        packed_result = {}
+        for k, v in packer.items():
+            packed = pack(unpacked, v)
+            packed_result[k] = packed
+    elif isinstance(packer, ConstWrap):
+        packed_result = packer.value
+    else:
+        packed_result = unpacked[packer]
+    return packed_result
+
+
 global_counter_id = 0
 
 # makes a custom op class from a func and input/output signatures
-def make_custom_op_class(func, input_signature, output_signature):
+def make_custom_op_class(func, input_signature, output_signature, input_packer, output_packer):
     import torch, numpy
     global global_counter_id
     # print('make_custom_op_class, id =', global_counter_id)
@@ -35,12 +94,11 @@ def __init__(self, *args):
 
         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,)
+            inputs_torch = tuple(torch.from_numpy(input.data) for input in inputs)   # TODO: Check memory sharing
+            args, kwargs = pack(inputs_torch, input_packer)
+            result = func(*args, **kwargs)
+            result, result_packer, _ = unpack(result, torch.Tensor)
+            assert result_packer == output_packer
             for i, tensor in enumerate(result):
                 ov.Tensor(numpy.array(tensor), shared_memory=True).copy_to(outputs[i])    # TODO: set the output tensor directly without copying
             return True
@@ -73,11 +131,10 @@ def make_signature(args):
 
 # Returns a tuple of tuples (element_type, partial_shape) for each argument, flattening nested structures if needed, setting all dimensions dynamic preserving rank
 # Currently assumes that all input arguments are torch.Tensor objects
-def make_input_signature(args, kwargs):
+def make_input_signature(args):
     # TODO: Avoid the current limitation: kwargs parameters should be passed in the same order as the function signature without gaps
     # flatten kwargs relying on the order of the keys
-    assert not kwargs, "Keyword arguments are not supported yet"
-    return make_signature(args + tuple(kwargs.values()))
+    return make_signature(args)
 
 
 def make_output_signature(args):
@@ -90,43 +147,46 @@ def make_output_signature(args):
     return make_signature(args)
 
 
-def is_class_method(obj):
-    if not inspect.isfunction(obj) and not inspect.ismethod(obj):
-        return False
-    argspec = inspect.getfullargspec(obj)
-    if argspec.args and argspec.args[0] == 'self':
-        return True
-    else:
-        return False
-
-
 def make_trampoline_class(func, op, op_attrs):
     import torch
     class Trampoline(torch.autograd.Function):
         target_extension = InlineConversionExtension()  # this is a marker for this type of extension
 
         # This function defines how the operation behaves when called as a part of PyTorch model code in eager execution or while jit.trace
         @staticmethod
-        def forward(ctx, *call_args, **call_kwargs):  #TODO: what is `ctx`?
+        def forward(ctx, *call_args):  #TODO: what is `ctx`?
             # print('Called through the trampoline')
-            func_target = func
             if not op:
-                if is_class_method(func):
-                    self_obj = call_args[0]
-                    call_args = call_args[1:]
-                    wrapped = lambda *distil_args, **distil_kwargs: func(self_obj, *distil_args, **distil_kwargs)
-                    func_target = wrapped
-                input_signature = make_input_signature(call_args, call_kwargs)
+                input_signature = make_input_signature(call_args)
             # TODO: Try to trace `func` with the hope to obtain tracable shapes to build more precise `validate_and_infer_types` automatically (unlikely possible)
-            result = func_target(*call_args, **call_kwargs)
+            print('about to call func_target with call_args:', call_args)
+            packed_args, packed_kwargs = pack(call_args, __class__.input_packer)
+            assert isinstance(packed_args, tuple)
+            assert isinstance(packed_kwargs, dict)
+            result = func(*packed_args, **packed_kwargs)
+            result, __class__.output_packer, _ = unpack(result, torch.Tensor)
+            print('output_packer:', __class__.output_packer)
             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)
+                __class__.op = make_custom_op_class(func, input_signature, output_signature, __class__.input_packer, __class__.output_packer)
             else:
                 __class__.op = op
             return result
 
+        # Unpack each element that is a tuple, a list, or a dict to a tuple of their values and concatenate together
+        # Build `packer` function to pack the result back to the original nested data types, save this `packer` to
+        # class member to use it later in `forward`.
+        @staticmethod
+        def unpack_inputs(args, kwargs):
+            unpacked, __class__.input_packer, _ = unpack((args, kwargs), torch.Tensor)
+            print('input_packer:', __class__.input_packer)
+            return unpacked
+
+        @staticmethod
+        def pack_outputs(result):
+            return pack(result, __class__.output_packer)
+
         # This function defines how the operation is represented in OpenVINO model graph
         @staticmethod
         def convert(node_context):
@@ -143,7 +203,11 @@ 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)
-            result = trampoline.apply(*args, **kwargs)
+            print('calling trampoline.apply with args:', args, 'kwargs:', kwargs)
+            args = trampoline.unpack_inputs(args, kwargs)
+            result = trampoline.apply(*args)
+            result = trampoline.pack_outputs(result)
+            # pack to the expected nested data types here
             #print('just called trampoline with result:', result)
             return result
         return trampoline