bigdl-core-npu 2.5.0__cp311-cp311-win_amd64.whl → 2.6.0__cp311-cp311-win_amd64.whl

intel_npu_acceleration_library/device.py

@@ -4,6 +4,7 @@
 #
 
 from intel_npu_acceleration_library.nn.module import convert_to_npu_module
+from intel_npu_acceleration_library.backend.tensor import RemoteTensor
 from torch.overrides import TorchFunctionMode
 from functools import lru_cache
 from typing import Any, MutableMapping
@@ -165,8 +166,7 @@ def to(super_fn: Any, self: Any, *args: Any, **kwargs: Any):
     """
     npu_device, args, kwargs = parse_to_arguments(*args, **kwargs)
     if npu_device:
-        # None for now, once the remote tensor feature lands, it can be converted to a remote tensor
-        pass
+        return super_fn(RemoteTensor.from_torch(self), *args, **kwargs)
     return super_fn(self, *args, **kwargs)
 
 

intel_npu_acceleration_library/dtypes.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass
 from typing import Union
 import numpy as np
 import torch
-
+import ctypes
 
 @dataclass(frozen=True)
 class NPUDtype:
@@ -81,6 +81,39 @@ class NPUDtype:
         return self.name
 
 
+def get_backend_dtype(dtype) -> ctypes.c_char_p:
+    """Get the string representation of the dtype.
+    Args:
+        dtype: numpy dtype
+    Raises:
+        RuntimeError: Unsupported datatype
+    Returns:
+        ctypes.c_char_p: string representation of the dtype
+    """
+    if dtype in [np.int8, torch.int8]:
+        str_dtype = "int8"
+    elif dtype in [np.uint8, int4, torch.uint8]:
+        # u8 represents packed i4 dtypes
+        str_dtype = "int4"
+    elif dtype in [np.int16, torch.int16]:
+        str_dtype = "int16"
+    elif dtype in [np.int32, torch.int32]:
+        str_dtype = "int32"
+    elif dtype in [np.int64, torch.int64]:
+        str_dtype = "int64"
+    elif dtype in [np.float16, torch.float16]:
+        str_dtype = "float16"
+    elif dtype in [np.float32, torch.float32]:
+        str_dtype = "float32"
+    elif dtype in [np.float64, torch.float64]:
+        str_dtype = "float64"
+    elif dtype in [bfloat16, torch.bfloat16]:
+        str_dtype = "bfloat16"
+    else:
+        raise RuntimeError(f"DType is not supported {dtype}")
+    return ctypes.c_char_p(str_dtype.encode())
+
+
 float16 = NPUDtype(
     "fp16",
     16,

intel_npu_acceleration_library/external/openvino/__init__.py

@@ -21,6 +21,7 @@ except ImportError:
 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

intel_npu_acceleration_library/external/openvino/_offline_transformations/__init__.py

@@ -18,3 +18,4 @@ from openvino._pyopenvino._offline_transformations import compress_model_transfo
 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/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]}