runtime.py

# -*- coding: utf-8 -*-
# Copyright (C) 2018-2023 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

"""Provide a layer of abstraction for an OpenVINO runtime environment."""

import logging
from typing import Dict, List, Union

import numpy as np

from openvino.runtime import Core

from openvino.runtime.exceptions import UserInputError
from openvino.runtime import Model, Node, Tensor, Type
from openvino.runtime.utils.types import NumericData, get_shape, get_dtype

import tests

log = logging.getLogger(__name__)


def runtime(backend_name: str = "CPU") -> "Runtime":
    """Create a Runtime object (helper factory)."""
    return Runtime(backend_name)


def get_runtime():
    """Return runtime object."""
    if tests.BACKEND_NAME is not None:
        return runtime(backend_name=tests.BACKEND_NAME)
    else:
        return runtime()


class Runtime(object):
    """Represents a graph runtime environment."""

    def __init__(self, backend_name: str) -> None:
        self.backend_name = backend_name
        log.debug(f"Creating runtime for {backend_name}")
        self.backend = Core()
        assert backend_name in self.backend.available_devices, 'The requested device "' + backend_name + '" is not supported!'

    def set_config(self, config: Dict[str, str]) -> None:
        """Set the runtime configuration."""
        self.backend.set_property(device_name=self.backend_name, properties=config)

    def computation(self, node_or_model: Union[Node, Model], *inputs: Node) -> "Computation":
        """Return a callable Computation object."""
        if isinstance(node_or_model, Node):
            model = Model(node_or_model, inputs, node_or_model.name)
            return Computation(self, model)
        elif isinstance(node_or_model, Model):
            return Computation(self, node_or_model)
        else:
            raise TypeError(
                "Runtime.computation must be called with an OpenVINO Model object "
                "or an OpenVINO node object an optionally Parameter node objects. "
                "Called with: %s",
                node_or_model,
            )

    def __repr__(self) -> str:
        return f"<Runtime: Backend='{self.backend_name}'>"


class Computation(object):
    """Graph callable computation object."""

    def __init__(self, runtime: Runtime, model: Model) -> None:
        self.runtime = runtime
        self.model = model
        self.parameters = model.get_parameters()
        self.results = model.get_results()
        self.network_cache = {}

    def convert_buffers(self, source_buffers, target_dtypes):
        converted_buffers = []
        for i in range(len(source_buffers)):
            key = list(source_buffers)[i]
            target_dtype = target_dtypes[i]
            # custom conversion for bf16
            if self.results[i].get_output_element_type(0) == Type.bf16:
                converted_buffers.append((source_buffers[key].view(target_dtype)).astype(target_dtype))
            else:
                converted_buffers.append(source_buffers[key].astype(target_dtype))
        return converted_buffers

    def convert_to_tensors(self, input_values):
        input_tensors = []
        for parameter, input_val in zip(self.parameters, input_values):
            if not isinstance(input_val, (np.ndarray)):
                input_val = np.ndarray([], type(input_val), np.array(input_val))
            if parameter.get_output_element_type(0) == Type.bf16:
                input_tensors.append(Tensor(Type.bf16, input_val.shape))
                input_tensors[-1].data[:] = input_val.view(np.float16)
            else:
                input_tensors.append(Tensor(input_val))
        return input_tensors

    def __repr__(self) -> str:
        params_string = ", ".join([param.name for param in self.parameters])
        return f"<Computation: {self.model.get_name()}({params_string})>"

    def __call__(self, *input_values: NumericData) -> List[NumericData]:
        """Run computation on input values and return result."""
        # Input validation
        if len(input_values) < len(self.parameters):
            raise UserInputError(
                "Expected %s params, received not enough %s values.",
                len(self.parameters),
                len(input_values),
            )

        param_names = [param.friendly_name for param in self.parameters]
        input_shapes = [get_shape(input_value) for input_value in input_values]
        if self.network_cache.get(str(input_shapes)) is None:
            model = self.model
            self.network_cache[str(input_shapes)] = model
        else:
            model = self.network_cache[str(input_shapes)]

        compiled_model = self.runtime.backend.compile_model(model, self.runtime.backend_name)
        is_bfloat16 = any(parameter.get_output_element_type(0) == Type.bf16 for parameter in self.parameters)
        if is_bfloat16:
            input_values = self.convert_to_tensors(input_values)
        request = compiled_model.create_infer_request()
        result_buffers = request.infer(dict(zip(param_names, input_values)))
        """Note: other methods to get result_buffers from request
           First call infer with no return value:
            request.infer(dict(zip(param_names, input_values)))
           Now use any of following options:
            result_buffers = [request.get_tensor(n).data for n in request.outputs]
            result_buffers = [request.get_output_tensor(i).data for i in range(len(request.outputs))]
            result_buffers = [t.data for t in request.output_tensors]
        """
        # # Since OV overwrite result data type we have to convert results to the original one.
        original_dtypes = [get_dtype(result.get_output_element_type(0)) for result in self.results]
        converted_buffers = self.convert_buffers(result_buffers, original_dtypes)
        return converted_buffers