bigdl-core-npu 2.6.0b20250114__cp311-cp311-win_amd64.whl

intel_npu_acceleration_library/external/openvino/__init__.py

@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)
+
+# Required for Windows OS platforms
+# Note: always top-level
+try:
+    from openvino.utils import _add_openvino_libs_to_search_path
+    _add_openvino_libs_to_search_path()
+except ImportError:
+    pass
+
+# #
+# # OpenVINO API
+# # This __init__.py forces checking of runtime modules to propagate errors.
+# # It is not compared with init files from openvino-dev package.
+# #
+# Import all public modules
+from openvino import runtime as runtime
+from openvino import frontend as frontend
+from openvino import helpers as helpers
+from openvino import experimental as experimental
+from openvino import preprocess as preprocess
+from openvino import utils as utils
+from openvino import properties as properties
+
+# Import most important classes and functions from openvino.runtime
+from openvino.runtime import Model
+from openvino.runtime import Core
+from openvino.runtime import CompiledModel
+from openvino.runtime import InferRequest
+from openvino.runtime import AsyncInferQueue
+
+from openvino.runtime import Symbol
+from openvino.runtime import Dimension
+from openvino.runtime import Strides
+from openvino.runtime import PartialShape
+from openvino.runtime import Shape
+from openvino.runtime import Layout
+from openvino.runtime import Type
+from openvino.runtime import Tensor
+from openvino.runtime import OVAny
+
+from openvino.runtime import compile_model
+from openvino.runtime import get_batch
+from openvino.runtime import set_batch
+from openvino.runtime import serialize
+from openvino.runtime import shutdown
+from openvino.runtime import tensor_from_file
+from openvino.runtime import save_model
+from openvino.runtime import layout_helpers
+
+from openvino._pyopenvino import RemoteContext
+from openvino._pyopenvino import RemoteTensor
+from openvino._pyopenvino import Op
+
+# libva related:
+from openvino._pyopenvino import VAContext
+from openvino._pyopenvino import VASurfaceTensor
+
+# Set version for openvino package
+from openvino.runtime import get_version
+__version__ = get_version()
+
+# Tools
+try:
+    # Model Conversion API - ovc should reside in the main namespace
+    from openvino.tools.ovc import convert_model
+except ImportError:
+    pass

intel_npu_acceleration_library/external/openvino/_offline_transformations/__init__.py

@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+
+from openvino._pyopenvino import get_version
+
+__version__ = get_version()
+
+from openvino._pyopenvino._offline_transformations import apply_fused_names_cleanup
+from openvino._pyopenvino._offline_transformations import apply_moc_transformations
+from openvino._pyopenvino._offline_transformations import apply_moc_legacy_transformations
+from openvino._pyopenvino._offline_transformations import apply_low_latency_transformation
+from openvino._pyopenvino._offline_transformations import apply_pruning_transformation
+from openvino._pyopenvino._offline_transformations import apply_make_stateful_transformation
+from openvino._pyopenvino._offline_transformations import compress_model_transformation
+from openvino._pyopenvino._offline_transformations import compress_quantize_weights_transformation
+from openvino._pyopenvino._offline_transformations import convert_sequence_to_tensor_iterator_transformation
+from openvino._pyopenvino._offline_transformations import paged_attention_transformation
+from openvino._pyopenvino._offline_transformations import stateful_to_stateless_transformation

intel_npu_acceleration_library/external/openvino/experimental/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Package: openvino
+This module provides access to experimental functionality that is subject to change without prior notice.
+"""
+
+# flake8: noqa
+
+from openvino._pyopenvino.experimental import evaluate_as_partial_shape
+from openvino._pyopenvino.experimental import evaluate_both_bounds
+from openvino._pyopenvino.experimental import set_element_type
+from openvino._pyopenvino.experimental import set_tensor_type

intel_npu_acceleration_library/external/openvino/frontend/__init__.py

@@ -0,0 +1,34 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Package: openvino
+Low level wrappers for the FrontEnd C++ API.
+"""
+
+# flake8: noqa
+
+from openvino._pyopenvino import get_version
+
+__version__ = get_version()
+
+# main classes
+from openvino.frontend.frontend import FrontEndManager
+from openvino.frontend.frontend import FrontEnd
+from openvino._pyopenvino import InputModel
+from openvino._pyopenvino import NodeContext
+from openvino._pyopenvino import Place
+
+# extensions
+from openvino._pyopenvino import DecoderTransformationExtension
+from openvino._pyopenvino import ConversionExtension
+from openvino._pyopenvino import OpExtension
+from openvino._pyopenvino import ProgressReporterExtension
+from openvino._pyopenvino import TelemetryExtension
+
+# exceptions
+from openvino._pyopenvino import NotImplementedFailure
+from openvino._pyopenvino import InitializationFailure
+from openvino._pyopenvino import OpConversionFailure
+from openvino._pyopenvino import OpValidationFailure
+from openvino._pyopenvino import GeneralFailure

intel_npu_acceleration_library/external/openvino/frontend/frontend.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Union
+
+from openvino._pyopenvino import FrontEnd as FrontEndBase
+from openvino._pyopenvino import FrontEndManager as FrontEndManagerBase
+from openvino._pyopenvino import InputModel
+from openvino.runtime import Model
+
+
+class FrontEnd(FrontEndBase):
+    def __init__(self, fe: FrontEndBase) -> None:
+        super().__init__(fe)
+
+    def convert(self, model: Union[Model, InputModel]) -> Model:
+        converted_model = super().convert(model)
+        if isinstance(model, InputModel):
+            return Model(converted_model)
+        return converted_model
+
+    def convert_partially(self, model: InputModel) -> Model:
+        return Model(super().convert_partially(model))
+
+    def decode(self, model: InputModel) -> Model:
+        return Model(super().decode(model))
+
+    def normalize(self, model: Model) -> None:
+        super().normalize(model)
+
+
+class FrontEndManager(FrontEndManagerBase):
+    def load_by_framework(self, framework: str) -> Union[FrontEnd, None]:
+        fe = super().load_by_framework(framework)
+        if fe is not None:
+            return FrontEnd(fe)
+        return fe
+
+    def load_by_model(self, model: str) -> Union[FrontEnd, None]:
+        fe = super().load_by_model(model)
+        if fe is not None:
+            return FrontEnd(fe)
+        return fe

intel_npu_acceleration_library/external/openvino/frontend/jax/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Package: openvino
+Low level wrappers for the FrontEnd C++ API.
+"""
+
+# flake8: noqa
+
+try:
+    from openvino.frontend.jax.py_jax_frontend import _FrontEndJaxDecoder as Decoder
+except ImportError as err:
+    raise ImportError("OpenVINO JAX frontend is not available, please make sure the frontend is built."
+                      "{}".format(err))

intel_npu_acceleration_library/external/openvino/frontend/jax/jaxpr_decoder.py

@@ -0,0 +1,293 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+import jax.core
+from openvino.frontend.jax.py_jax_frontend import _FrontEndJaxDecoder as Decoder
+from openvino.runtime import PartialShape, Type as OVType, OVAny
+from openvino.frontend.jax.utils import jax_array_to_ov_const, get_ov_type_for_value, \
+    ivalue_to_constant, param_to_constants
+
+import jax
+import numpy as np
+
+from typing import List
+import logging
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+class JaxprPythonDecoder (Decoder):
+    '''
+    The jaxpr decoder uses Jaxpr to get graph information from a jax module.
+    It takes use of the following parts.
+    
+    - `ClosedJaxpr`: the jaxpr object that contains the jaxpr and literals.
+        - `Jaxpr`: the jaxpr object that contains the invars, outvars, and eqns.
+            - `JaxEqns`: A list of jaxpr equations, which contains the information of the operation.
+                - `Primitive`: the operation that is used in the equation.
+                - `invars`: the input variables of the equation.
+                    - `aval`: the abstract value.
+                - `outvars`: the output variables of the equation.
+                    - `aval`: the abstract value.
+                - `params`: the named params of this equation.
+            - `invars`: the inputs of the model (traced graph).
+                - `aval`: the abstract value.
+            - `outvars`: the outputs of the model (traced graph).
+                - `aval`: the abstract value.
+            - `constvars`: the constant variables used in this model.
+                - `aval`: the abstract value.
+        - `Literal`: the literal object that contains the value of the constants.
+    '''
+    
+    def __init__(self, jaxpr, name=None, literals=None):
+        '''
+        Inputs: 
+            - jaxpr: for users, `ClosedJaxpr` is expected here. See https://github.com/google/jax/blob/jaxlib-v0.4.29/jax/_src/core.py#L197
+            - name: the name for the model.
+            - literals: the literals (constants) that are used in the model.
+        '''
+        Decoder.__init__(self)
+        
+        if isinstance(jaxpr, (jax.core.JaxprEqn, jax.core.Jaxpr)):
+            self.jaxpr = jaxpr
+        elif isinstance(jaxpr, jax.core.ClosedJaxpr):
+            # Take the `Jaxpr` from `ClosedJaxpr`, see https://github.com/google/jax/blob/jaxlib-v0.4.29/jax/_src/core.py#L85
+            self.jaxpr = jaxpr.jaxpr
+            # Literal should be a `Jax.core.Var`, see https://github.com/google/jax/blob/jaxlib-v0.4.29/jax/_src/core.py#L85
+            self.literals = jaxpr.literals
+        else:
+            raise ValueError(f"Unexpected type of jaxpr: {type(jaxpr)}")
+        self.name = name
+        if self.name is None:
+            self.name = "jax_module"
+        if literals is not None:
+            self.literals = literals
+            
+        self.params = {}
+        if hasattr(self.jaxpr, 'params') and isinstance(self.jaxpr.params, dict):
+            for k in self.jaxpr.params.keys():
+                converted = self.convert_param_to_constant_node(self.jaxpr, k)
+                if converted is not None:
+                    self.params.update(converted)
+                
+        # TODO: this implementation may lead to memory increasing. Any better solution?
+        self.m_decoders = []
+        
+    def inputs(self) -> List[int]:
+        if isinstance(self.jaxpr, jax.core.JaxprEqn):
+            idx = 0
+            res = []
+            for inp in self.jaxpr.invars:
+                if isinstance(inp, jax.core.Literal):
+                    res.append(self.literals[idx].output(0))
+                    idx += 1
+                else:
+                    res.append(id(inp))
+            return res
+        else:
+            return [id(v) for v in self.jaxpr.invars]
+    
+    def input(self, idx: int) -> int:
+        return id(self.jaxpr.invars[idx])
+    
+    def get_input_shape(self, index):
+        return PartialShape(self.jaxpr.invars[index].aval.shape)
+    
+    def get_input_signature_name(self, index) -> str:
+        return "jaxpr_invar_" + str(index)
+    
+    def get_input_type(self, index) -> OVType:
+        return get_ov_type_for_value(self.jaxpr.invars[index])
+        
+    def get_named_param(self, name):
+        '''
+        Get the object id of the named parameter by the name.
+        '''
+        return self.params[name].output(0)
+    
+    def get_named_param_as_constant(self, name):
+        '''
+        The named parameter in JAX is a python object but we want to use its value in cpp.
+        Therefore this API is used to get the named parameter as a constant, which can be used
+        to extract the value of it in cpp-level.
+        '''
+        return self.params[name].as_constant()
+    
+    def get_param_names(self):
+        '''
+        In JAX, the named parameters may exist in `params` attribute of `JaxEqn`.
+        For example, the `jax.lax.cat` operation has a named parameter `dim`, 
+        which is used to indicate the dimension to concatenate the tensors.
+        
+        Here we return the names of all the named params that appear in the model for the current `JaxEqn`.
+        '''
+        return list(self.params.keys())
+    
+    def get_output_type(self, index) -> OVType:
+        return get_ov_type_for_value(self.jaxpr.outvars[index])
+        
+    def get_output_name(self, index) -> str:
+        return "jaxpr_outvar_" + str(index)
+    
+    def get_output_shape(self, index):
+        return PartialShape(self.jaxpr.outvars[index].aval.shape)
+    
+    def visit_subgraph(self, node_visitor) -> None:
+        if isinstance(self.jaxpr, jax.core.JaxprEqn):
+            return
+        for _, decoder in self.params.items():
+            self.m_decoders.append(decoder)
+            node_visitor(decoder)
+        for idx, node in enumerate(self.jaxpr.constvars):
+            decoder = self.convert_literal_to_constant_node(
+                literal=self.literals[idx], 
+                name=self.name + "/" + f"const({id(node)})", 
+                output_id=id(node)
+            )
+            self.m_decoders.append(decoder)
+            node_visitor(decoder)
+        # Visit every `JaxEqn` in the jaxpr, see https://github.com/google/jax/blob/jaxlib-v0.4.29/jax/_src/core.py#L285
+        for node in self.jaxpr.eqns:
+            literal_decoders = []
+            for inp in node.invars:
+                if isinstance(inp, jax.core.Literal):
+                    literal_decoder = self.convert_literal_to_constant_node(inp)
+                    literal_decoders.append(literal_decoder)
+                    node_visitor(literal_decoder)
+            decoder = JaxprPythonDecoder(node, name=self.name + "/" + node.primitive.name, literals=literal_decoders)
+            self.m_decoders.append(decoder)
+            node_visitor(decoder)
+            
+    def get_op_type(self) -> str:
+        if isinstance(self.jaxpr, jax.core.JaxprEqn):
+            return self.jaxpr.primitive.name
+        else:
+            return "root"
+    
+    def outputs(self) -> List[int]:
+        return [id(v) for v in self.jaxpr.outvars]
+    
+    def output(self, idx: int) -> int:
+        return id(self.jaxpr.outvars[idx])
+    
+    def num_inputs(self) -> int:
+        return len(self.jaxpr.invars)
+    
+    def num_outputs(self) -> int:
+        return len(self.jaxpr.outvars)
+    
+    def as_constant(self):
+        if self.get_op_type() == 'constant':
+            value = self.literals
+            # TODO: dig out how to share the memory.
+            # Currently, using shared_memory will raise `ValueError: array is not writeable``
+            ov_const = jax_array_to_ov_const(value, shared_memory=False)
+            return ov_const.outputs()
+        else:
+            raise ValueError("This is not a constant node so it cannot be converted to a constant.")
+        
+    @staticmethod
+    def convert_param_to_constant_node(jaxpr, param) -> dict:
+        assert hasattr(jaxpr, 'params'), "The jaxpr does not have params."
+        if hasattr(jaxpr, 'primitive'):
+            param_map = param_to_constants(jaxpr.primitive.name, param, jaxpr, shared_memory=False)
+            res = {}
+            for name, constant in param_map.items():
+                if constant is not None:
+                    res[name] = _JaxprPythonConstantDecoder(constant=constant)
+        else:
+            constant = ivalue_to_constant(jaxpr.params[param], shared_memory=False)
+            res = {param: _JaxprPythonConstantDecoder(constant=constant)} if constant is not None else {}
+        return res
+    
+    @staticmethod
+    def convert_literal_to_constant_node(literal, name=None, output_id=None):
+        if isinstance(literal, jax.core.Literal):
+            constant = ivalue_to_constant(literal.val, shared_memory=False)
+        elif isinstance(literal, (jax.Array, np.ndarray)):
+            constant = ivalue_to_constant(literal, shared_memory=False)
+        else:
+            raise TypeError( f"The input should be a literal or jax array, but got {type(literal)}.")
+        return _JaxprPythonConstantDecoder(constant=constant, name=name, output_id=output_id)
+        
+class _JaxprPythonConstantDecoder (Decoder):
+    def __init__(self, name=None, constant=None, output_id=None):
+        '''
+        A decoder specially for constants and named parameters.
+        
+        Inputs:
+            - name: the name for the model.
+            - literals: the literals (constants) that are used in the model.
+            - output_id: the id specified for this decoder's output. If none, use `id(self.constant)`.
+        '''
+        Decoder.__init__(self)
+        
+        self.name = name
+        self.constant = constant
+        self.output_id = id(self.constant) if output_id is None else output_id
+        
+    def inputs(self) -> List[int]:
+        return []
+    
+    def input(self, idx: int) -> int:
+        raise ValueError("This is a constant node so it does not have input.")
+    
+    def get_input_shape(self, index):
+        raise ValueError("This is a constant node so it does not have input shape.")
+    
+    def get_input_signature_name(self, index) -> str:
+        raise ValueError("This is a constant node so it does not have input signature name.")
+    
+    def get_input_type(self, index) -> OVType:
+        raise ValueError("This is a constant node so it does not have input type.")
+        
+    def get_named_param(self, name):
+        raise ValueError("This is a constant node so it does not have named param.")
+    
+    def get_named_param_as_constant(self, name):
+        raise ValueError("This is a constant node so it does not have named param.")
+    
+    def get_param_names(self):
+        '''
+        In JAX, the named parameters may exist in `params` attribute of `JaxEqn`.
+        For example, the `jax.lax.cat` operation has a named parameter `dim`, 
+        which is used to indicate the dimension to concatenate the tensors.
+        
+        However, `_JaxprPythonConstantDecoder` is already a named param or a constant.
+        So it will never have a named param.
+        '''
+        return []
+    
+    def get_output_type(self, index) -> OVType:
+        assert len(self.constant) == 1
+        return OVAny(self.constant[0].element_type)
+        
+    def get_output_name(self, index) -> str:
+        return "jaxpr_outvar_" + str(index)
+    
+    def get_output_shape(self, index):
+        assert len(self.constant) == 1
+        return PartialShape(self.constant[0].shape)
+    
+    def visit_subgraph(self, node_visitor) -> None:
+        return
+            
+    def get_op_type(self) -> str:
+        return "constant"
+    
+    def outputs(self) -> List[int]:
+        return [self.output_id]
+    
+    def output(self, idx: int) -> int:
+        return self.output_id
+    
+    def num_inputs(self) -> int:
+        return 0
+    
+    def num_outputs(self) -> int:
+        return 1
+    
+    def as_constant(self):
+        return self.constant

intel_npu_acceleration_library/external/openvino/frontend/jax/passes.py

@@ -0,0 +1,65 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+from enum import Enum
+from jax.lax import ConvDimensionNumbers
+
+def enum_values_pass(value):
+    if isinstance(value, Enum):
+        return value.value
+    return value
+
+
+def conv_dimension_numbers_pass(value):
+    if isinstance(value, ConvDimensionNumbers):
+        return [
+            list(value.lhs_spec),
+            list(value.rhs_spec),
+            list(value.out_spec)
+        ]
+    return value
+
+
+def filter_element(value):
+    passes = [enum_values_pass]
+    for pass_ in passes:
+        value = pass_(value)
+    return value
+
+
+def filter_ivalue(value):
+    passes = [conv_dimension_numbers_pass]
+    for pass_ in passes:
+        value = pass_(value)
+    return value
+
+
+def dot_general_param_pass(param_name: str, jax_eqn):
+    param = jax_eqn.params[param_name]
+    res = {}
+    if param_name == 'dimension_numbers':
+        contract_dimensions = param[0]
+        assert len(contract_dimensions) == 2
+        res['contract_dimensions'] = [list(contract_dimensions[0]), list(contract_dimensions[1])]
+        
+        batch_dimensions = param[1]
+        assert len(batch_dimensions) == 2
+        lhs_length = len(batch_dimensions[0])
+        rhs_length = len(batch_dimensions[1])
+        assert lhs_length == rhs_length
+        if lhs_length > 0:
+            res['batch_dimensions'] = [list(batch_dimensions[0]), list(batch_dimensions[1])]
+    return res
+
+# mapping from primitive to pass 
+param_passes = {
+    'dot_general': dot_general_param_pass,
+}
+
+def filter_param(primitive: str, param_name: str, jax_eqn):
+    if primitive in param_passes:
+        return param_passes[primitive](param_name, jax_eqn)
+    return {param_name: jax_eqn.params[param_name]}