[POC] PyTorch Inlined Extension

[slyalin]

It is a seamless way to embed a Python portion of the original PyTorch model into OpenVINO model as a custom operation without the explicit creation of a custom operation class.

This feature combines two core technologies: Python OpenVINO custom ops and PyTorch opaque autograd.Function (the same tool that is used underneath the ModuleExtension).

The main target users are model enablers who are exporting complex models from PyTorch ecosystem to OpenVINO ecosystem. They can hide parts of the original PyTorch model in on-the-fly defined opaque custom operations if there are issues with model tracing/conversion to OpenVINO.

Run the resulting OpenVINO model as-is with vanilla OpenVINO API in the same Python process. Provide an optimized implementation of such operations in C++ OpenVINO custom operations for performance reasons and to be ready to deploy without Python/PyTorch dependency.

Example:

import openvino as ov
import torch

# ov.inlined_extension decorator instructs to auto-generate a custom operation for each `my_func` call.
# Alternatively, use @ov.inlined_extension(MyOp) to substitute `MyOp` Python custom operation class instead of auto-generation.
# In both cases, each call of `my_func` will appear as a single custom op node in the OpenVINO model graph.

@ov.inlined_extension
def my_func(tensor1, tensor2):
    # Arbitrary Python code that is connected with the caller context by means of using torch.Tensor objects in inputs/outputs.
    # Inputs and outputs can consist of tracible (torch.Tensor) and not tracible types.
    # All `list`s, `tuple`s and `dict`s in inputs/outputs arguments are unpacked recursively to extract `torch.Tensor` objects.
    return tensor1 + tensor2, tensor1 * tensor2

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
    def forward(self, tensor):
        t1, t2 = my_func(tensor + 1, tensor + 2)  # <-- here the wrapped function is called as a part of PyTorch model
        return t1 / t2

my_model = MyModel()
sample_data = torch.tensor([1.0, 2.0, 3.0])
print(f'Original model inference: {my_model(sample_data)}')

ov_model = ov.convert_model(my_model, example_input=sample_data)  # convert model as usual
core = ov.Core()
compiled = core.compile_model(ov_model, 'CPU')
print(f'Compiled model inference: {compiled(sample_data)}')

Resulting OpenVINO model (custom op is InlinedCustomOp, id attribute is unique for each op instance):

t0 = opset.Parameter({'shape': [-1], 'element_type': 'f32'})  #  -> f32[?]
t1 = opset.Constant(model, 1)                                 #  -> f32[1]([1.0])
t2 = opset.Add([t0, t1], {'auto_broadcast': 'numpy'})         # f32[?], f32[1] -> f32[?]
t3 = opset.Constant(model, 3)                                 #  -> f32[1]([2.0])
t4 = opset.Add([t0, t3], {'auto_broadcast': 'numpy'})         # f32[?], f32[1] -> f32[?]
t5, t6 = opset.InlinedCustomOp([t2, t4], {'id': 6})           # f32[?], f32[?] -> f32[?], f32[?]
t7 = opset.Divide([t5, t6], {'auto_broadcast': 'numpy', 'm_pythondiv': True})  # f32[?], f32[?] -> f32[?]
t8 = opset.Result([t7], {})                                   # f32[?] -> f32[?]

Discovered limitations

Thanks to @eaidova, it was found that when torch.Tensor objects is used inside inlined_extension decorated function and this object is not present among input arguments as traced object (that means it is hidden inside some opaque data structure in one of the arguments, or even not passed to a function as an argument), then it leads to the exception thrown from C++ torch implementation like this one:

  File "/home/developer/.local/lib/python3.10/site-packages/torch/autograd/function.py", line 575, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
RuntimeError: _Map_base::at