bigdl-core-npu 2.5.0__cp310-cp310-win_amd64.whl

intel_npu_acceleration_library/external/openvino/frontend/pytorch/module_extension.py

@@ -0,0 +1,39 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+class ModuleExtension:
+    def __init__(self, module, target_op, evaluate=None, convert=None):
+        """
+        Creates an extension that replaces entire PyTorch module by a single operation.
+        This functionality works with PyTorch models only. A module can be identified by
+        module type (e.g. torch.nn.Linear), module instance in the model or module name.
+
+        Args:
+            module (str, torch.nn.Module, type(torch.nn.Module)): PyTorch module to replace
+
+            target_op (str): a target operation that will be used as a replacer for the module,
+                            could be a name of the extension operation or existing PyTorch operation
+                            (with prim:: or aten:: prefix following TorchScript syntax).
+
+            evaluate (callable with args module, *args, **kwargs): a callable that will replace a target
+                            module in model execution it is responsible for producing valid output for
+                            the module to allow correct model tracing. By default it calls original module
+                            forward with the same arguments. The provided code will not be a part of the final
+                            traced model, it is used only to produce valid results in the tracing.
+
+            convert (callable with args target_op, *args, **kwargs): a callable that will be traced and become
+                            a part of the final model instead of the target module. It accepts target_op as
+                            the first parameter, target_op is callable that will appear as a single node in the
+                            graph, the type of the node is target_op provided as another argument above.
+        """
+        self.module = module
+        self.target_op = target_op
+        self.evaluate = evaluate
+        if self.evaluate is None:
+            self.evaluate = lambda module, *args, **kwargs: module(*args, **kwargs)
+        self.convert = convert
+        if self.convert is None:
+            self.convert = lambda module, target_op, *args, **kwargs: target_op(*args, **kwargs)

intel_npu_acceleration_library/external/openvino/frontend/pytorch/patch_model.py

@@ -0,0 +1,98 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+import torch
+from openvino.frontend.pytorch import ModuleExtension
+
+
+class no_jit_trace:
+    def __enter__(self):
+        self.state = torch._C._get_tracing_state()
+        torch._C._set_tracing_state(None)
+
+    def __exit__(self, *args):
+        torch._C._set_tracing_state(self.state)
+        self.state = None
+
+
+def patch_model(model, module_extensions, orig_forward_name):
+    def module_patcher(m, name):
+        extension = None
+        if m in module_extensions:
+            extension = module_extensions[m]
+        elif m.__class__ in module_extensions:
+            extension = module_extensions[m.__class__]
+        elif name in module_extensions:
+            extension = module_extensions[name]
+
+        if extension:
+            # The Trampoline class is instantiated for every module replacement, so we can use class members individually for each module.
+            class Trampoline(torch.autograd.Function):
+                target_extension = extension
+                original_module = m
+                stashed_args = None
+                stashed_kwargs = None
+
+                @staticmethod
+                @torch.jit.ignore
+                def forward(*args, **kwargs):
+                    with no_jit_trace():
+                        # `module` is going to be passed to a user-defined function `evaluate`
+                        # `module` is patched: forward function was replaced, and we are actually in this patched function right in this code
+                        # if we pass `module` as-is to the user code below, and it happens to call forward it will lead to infinite recursion or fail
+                        # so we need to temporary patch the module back to the original forward and then return it back again
+                        # stash the current forward to be able to return it back
+                        patched_forward = m.forward
+                        # set original forward for the module
+                        m.forward = getattr(m, orig_forward_name)
+                        # call user code
+                        results = extension.evaluate(
+                            m, *Trampoline.stashed_args, **Trampoline.stashed_kwargs)  # call user code
+                        m.forward = patched_forward  # return patched forward back
+                        return results
+
+            def new_forward(*args, **kwargs):
+                Trampoline.stashed_args = args
+                Trampoline.stashed_kwargs = kwargs
+                return extension.convert(m, Trampoline.apply, *args, **kwargs)
+            setattr(m, orig_forward_name, m.forward)
+            m.forward = new_forward
+
+    for name, m in model.named_modules():
+        if hasattr(m, orig_forward_name):
+            # already patched, skipping with a warning because it is unexpected
+            print(f'[ WARNING ] Unexpectedly found already patched module {name} while applying ModuleExtension during PyTorch model conversion. '
+                  'Result of the conversion maybe broken. Depending on the exact issue it may lead to broken original model.')
+            continue
+        module_patcher(m, name)
+
+
+def unpatch_model(model, orig_forward_name):
+    for _, m in model.named_modules():
+        if hasattr(m, orig_forward_name):
+            try:
+                m.forward = getattr(m, orig_forward_name)
+                delattr(m, orig_forward_name)
+            except Exception as error:
+                print('[ WARNING ] Exception raised during model unpatching. Depending on the exact issue it may lead to broken original model.')
+                print('Original exception details:')
+                print(error)
+
+
+def __make_16bit_traceable(model: torch.nn.Module):
+    # Replace torch.nn.Linear with ModuleExtension and move other modules to fp32
+    extensions = {torch.nn.Linear: ModuleExtension(
+        torch.nn.Linear,
+        "aten::linear",
+        evaluate=lambda module, *args, **kwargs: torch.ones(
+            list(args[0].shape[:-1]) + [module.out_features], dtype=torch.float32) * 0.5,
+        convert=lambda module, target_op, *args, **kwargs: target_op(args[0], module.weight, module.bias))
+    }
+    patch_model(model, extensions,
+                "_openvino_module_extension_patch_orig_forward")
+    for _, module in model.named_modules():
+        if module.__class__ not in extensions and hasattr(module, "weight") and module.weight.dtype in [torch.float16, torch.bfloat16]:
+            module.float()

intel_npu_acceleration_library/external/openvino/frontend/pytorch/torchdynamo/backend.py

@@ -0,0 +1,119 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+import logging
+import os
+from functools import partial
+from hashlib import sha256
+
+import torch
+from torch._dynamo.backends.common import fake_tensor_unsupported, aot_autograd
+from torch._dynamo.backends.registry import register_backend
+from torch._inductor.compile_fx import compile_fx
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch._decomp import decomposition_table, get_decompositions
+
+from openvino.frontend import FrontEndManager
+from openvino.runtime import Core, Type, PartialShape
+from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
+from openvino.frontend.pytorch.torchdynamo import decompositions
+from openvino.frontend.pytorch.torchdynamo.decompositions import get_aot_decomposition_list, get_inf_decomposition_list
+from openvino.frontend.pytorch.torchdynamo.partition import Partitioner
+from openvino.frontend.pytorch.torchdynamo.execute import execute, execute_cached
+from openvino.frontend.pytorch.torchdynamo.compile import cached_model_name, openvino_compile_cached_model
+from openvino.frontend.pytorch.torchdynamo.backend_utils import _get_cache_dir, _get_device, _get_model_caching, _get_decompositions, _get_aot_autograd
+
+from openvino.runtime import Core, Type, PartialShape
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+"""
+    This is a preview feature in OpenVINO. This feature
+    enables users to compile PyTorch models using torch.compile
+    with OpenVINO as a target backend in PyTorch applications
+
+    Sample usage:
+    This sample code loads resnet50 torchvision model and compiles it using torch dynamo.
+    We can then use this model for inference. We only need to add two lines of code to
+    the Pytorch applications which are marked in the code below
+
+    1) import openvino.torch
+    model = torchvision.models.resnet50()
+    2) model = torch.compile(model, backend="openvino")
+"""
+
+openvino_options = {}
+
+@register_backend
+@fake_tensor_unsupported
+def openvino(subgraph, example_inputs, options=None):
+    if (_get_aot_autograd(options)):
+        global openvino_options
+        openvino_options = options
+        decompositions = _get_decompositions(options) + get_inf_decomposition_list()
+        decompositions = decompositions + get_aot_decomposition_list()
+        return aot_autograd(fw_compiler=fx_openvino, 
+                            bw_compiler=fx_openvino, 
+                            decompositions=get_decompositions(decompositions))(subgraph, example_inputs)
+    return fx_openvino(subgraph, example_inputs, options)
+
+def fx_openvino(subgraph, example_inputs, options=None):
+    try:
+        if len(openvino_options) != 0:
+            options = openvino_options
+        executor_parameters = None
+        inputs_reversed = False
+        openvino_model_caching = _get_model_caching(options)
+        if openvino_model_caching is not None and openvino_model_caching:
+            # Create a hash to be used for caching
+            model_hash_str = sha256(subgraph.code.encode('utf-8')).hexdigest()
+            executor_parameters = {"model_hash_str": model_hash_str}
+            # Check if the model was fully supported and already cached
+            example_inputs.reverse()
+            inputs_reversed = True
+            maybe_fs_cached_name = cached_model_name(model_hash_str + "_fs", _get_device(options), example_inputs, _get_cache_dir(options))
+            if os.path.isfile(maybe_fs_cached_name + ".xml") and os.path.isfile(maybe_fs_cached_name + ".bin"):
+                # Model is fully supported and already cached. Run the cached OV model directly.
+                compiled_model = openvino_compile_cached_model(maybe_fs_cached_name, options, *example_inputs)
+                def _call(*args):
+                    res = execute_cached(compiled_model, *args)
+                    return res
+                return _call
+        if inputs_reversed:
+            example_inputs.reverse()
+
+        if (_get_aot_autograd(options)):
+            model = subgraph
+        else:
+            from torch._subclasses.fake_tensor import FakeTensorMode
+            decompositions = _get_decompositions(options) + get_inf_decomposition_list()
+            with FakeTensorMode(allow_non_fake_inputs=True):
+                model = make_fx(subgraph, decomposition_table=get_decompositions(decompositions))(*example_inputs)
+
+            with torch.no_grad():
+                model.eval()
+
+        partitioner = Partitioner(options)
+        compiled_model = partitioner.make_partitions(model, options)
+
+        if executor_parameters is not None and 'model_hash_str' in executor_parameters:
+            # Check if the model is fully supported.
+            fully_supported = partitioner.check_fully_supported(compiled_model)
+            if fully_supported:
+                executor_parameters["model_hash_str"] += "_fs"
+
+        def _call(*args):
+            res = execute(compiled_model, *args, executor="openvino",
+                          executor_parameters=executor_parameters, options=options)
+            return res
+        return _call
+    except Exception as e:
+        logger.debug(f"Failed in OpenVINO execution: {e}")
+        return compile_fx(subgraph, example_inputs)
+
+def reset():
+    clear_caches()

intel_npu_acceleration_library/external/openvino/frontend/pytorch/torchdynamo/backend_utils.py

@@ -0,0 +1,85 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+from typing import Optional, Any
+from openvino.runtime import Core
+
+
+def _get_device(options) -> Optional[Any]:
+    core = Core()
+    device = "CPU"
+
+    if options is not None and "device" in options:
+        device = options["device"]
+
+    if device is not None:
+        assert device in core.available_devices, (
+            "Specified device "
+            + device
+            + " is not in the list of OpenVINO Available Devices"
+        )
+    else:
+        device = "CPU"
+    return device
+
+def _is_cache_dir_in_config(options) -> Optional[Any]:
+    if options is not None and "config" in options:
+        cfg = options["config"]
+        if cfg is not None and "CACHE_DIR" in cfg:
+            return True
+    return False
+
+
+def _get_cache_dir(options) -> Optional[Any]:
+    cache_dir = "./cache"
+    if options is not None and "cache_dir" in options:
+        cache_dir = options["cache_dir"]
+    if _is_cache_dir_in_config(options):
+        cache_dir = options["config"]["CACHE_DIR"]
+    return cache_dir
+
+
+def _get_aot_autograd(options) -> Optional[Any]:
+    if options is not None and "aot_autograd" in options:
+        aot_autograd = options["aot_autograd"]
+        if bool(aot_autograd) and str(aot_autograd).lower() not in ["false", "0"]:
+            return True
+        else:
+            return False
+
+
+def _get_model_caching(options) -> Optional[Any]:
+    if options is not None and "model_caching" in options:
+        caching = options["model_caching"]
+        if bool(caching) and str(caching).lower() not in ["false", "0"]:
+            return True
+    return False
+
+
+def _get_config(options) -> Optional[Any]:
+    if options is not None and "config" in options:
+        return options["config"]
+    return {}
+
+def _get_decompositions(options) -> Optional[Any]:
+    decompositions = []
+    if options is not None and "decompositions" in options:
+        decompositions = options["decompositions"]
+    return decompositions
+
+def _get_disabled_ops(options) -> Optional[Any]:
+    disabled_ops = []
+    if options is not None and "disabled_ops" in options:
+        disabled_ops = options["disabled_ops"]
+    return disabled_ops
+
+def _is_testing(options) -> Optional[Any]:
+    if options is not None and "testing" in options:
+        is_testing = options["testing"]
+        if bool(is_testing) and str(is_testing).lower not in ["false", "0"]:
+            return True
+    return False
+

intel_npu_acceleration_library/external/openvino/frontend/pytorch/torchdynamo/compile.py

@@ -0,0 +1,141 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+import os
+import torch
+import torch.overrides
+
+from hashlib import sha256
+from torch.fx import GraphModule
+
+from openvino.frontend import FrontEndManager
+from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder
+from openvino.runtime import Core, Type, PartialShape, serialize
+from openvino.frontend.pytorch.torchdynamo.backend_utils import _get_cache_dir, _get_device, _get_config, _is_cache_dir_in_config
+
+from typing import Callable, Optional
+
+import logging
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+def cached_model_name(model_hash_str, device, args, cache_root, reversed = False):
+    if model_hash_str is None:
+        return None
+
+    model_cache_dir = cache_root + "/model/"
+
+    try:
+        os.makedirs(model_cache_dir, exist_ok=True)
+        file_name = model_cache_dir + model_hash_str + "_" + device
+    except OSError as error:
+        logger.warning(f"Cache directory {cache_root} cannot be created. Model caching is disabled. Error: {error }")
+        return None
+
+    inputs_str = ""
+    for idx, input_data in enumerate(args):
+        if reversed:
+            inputs_str = "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") + inputs_str
+        else:
+            inputs_str += "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "")
+    inputs_str = sha256(inputs_str.encode('utf-8')).hexdigest()
+    file_name += inputs_str
+
+    return file_name
+
+def openvino_compile_cached_model(cached_model_path, options, *example_inputs):
+    core = Core()
+    om = core.read_model(cached_model_path + ".xml")
+
+    dtype_mapping = {
+        torch.float32: Type.f32,
+        torch.float64: Type.f64,
+        torch.float16: Type.f16,
+        torch.int64: Type.i64,
+        torch.int32: Type.i32,
+        torch.uint8: Type.u8,
+        torch.int8: Type.i8,
+        torch.bool: Type.boolean
+    }
+
+    for idx, input_data in enumerate(example_inputs):
+        om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
+        om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
+    om.validate_nodes_and_infer_types()
+
+    config = {}
+
+    if _is_cache_dir_in_config(options):
+        config = _get_config(options)
+    else:
+        config["CACHE_DIR"] = _get_cache_dir(options)
+
+    compiled_model = core.compile_model(om, _get_device(options), config)
+
+    return compiled_model
+
+def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None, options=None):
+    core = Core()
+
+    device = _get_device(options)
+    cache_root = _get_cache_dir(options)
+    file_name = cached_model_name(model_hash_str, device, args, cache_root)
+
+    if file_name is not None and os.path.isfile(file_name + ".xml") and os.path.isfile(file_name + ".bin"):
+        om = core.read_model(file_name + ".xml")
+    else:
+        fe_manager = FrontEndManager()
+        fe = fe_manager.load_by_framework("pytorch")
+
+        input_shapes = []
+        input_types = []
+        for idx, input_data in enumerate(args):
+            if isinstance(input_data, int):
+                input_types.append(torch.int64)
+                input_shapes.append(torch.Size([1]))
+            else:
+                input_types.append(input_data.type())
+                input_shapes.append(input_data.size())
+
+        decoder = TorchFXPythonDecoder(gm)
+
+        im = fe.load(decoder)
+
+        om = fe.convert(im)
+
+        if file_name is not None:
+            serialize(om, file_name + ".xml", file_name + ".bin")
+
+    dtype_mapping = {
+        torch.float32: Type.f32,
+        torch.float64: Type.f64,
+        torch.float16: Type.f16,
+        torch.int64: Type.i64,
+        torch.int32: Type.i32,
+        torch.uint8: Type.u8,
+        torch.int8: Type.i8,
+        torch.bool: Type.boolean
+    }
+
+    for idx, input_data in enumerate(args):
+        if isinstance(input_data, int):
+            om.inputs[idx].get_node().set_element_type(dtype_mapping[torch.int64])
+            om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([1]))))
+        else:
+            om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
+            om.inputs[idx].get_node().set_partial_shape(PartialShape(list(decoder.input_shapes[idx])))
+
+    om.validate_nodes_and_infer_types()
+
+    config = _get_config(options)
+
+    if model_hash_str is not None:
+        if not _is_cache_dir_in_config(options):
+            config["CACHE_DIR"] = cache_root
+
+    compiled = core.compile_model(om, device, config)
+    return compiled

intel_npu_acceleration_library/external/openvino/frontend/pytorch/torchdynamo/decompositions.py

@@ -0,0 +1,116 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+import torch
+from torch._decomp.decompositions import aten, pw_cast_for_opmath
+from torch._decomp import register_decomposition, get_decompositions
+
+
+@register_decomposition(aten.convolution_backward)
+@pw_cast_for_opmath
+def convolution_backward(
+    grad_output,
+    inp,
+    weight,
+    bias,
+    stride,
+    padding,
+    dilation,
+    transposed,
+    output_padding,
+    groups,
+    output_mask,
+):
+    if stride == [2, 2]:
+        output_padding = [1, 1]
+
+    # Compute the gradient of the input tensor
+    grad_input = torch.nn.functional.conv_transpose2d(
+        grad_output, weight, stride=stride, padding=padding, dilation=dilation, groups=groups, output_padding=output_padding
+    )
+
+    # Compute the gradient of the weight tensor
+    grad_weight = torch.nn.functional.conv_transpose2d(
+        inp, weight.transpose(0, 1), stride=stride, padding=padding, dilation=dilation, groups=groups, output_padding=output_padding
+    )
+
+    # Compute the gradient of the bias tensor
+    if bias is not None:
+        grad_bias = grad_output.sum([0, 2, 3], keepdim=True)
+    else:
+        grad_bias = None
+
+    return grad_input, grad_weight, grad_bias
+
+if len(get_decompositions([aten._scaled_dot_product_flash_attention.default])) == 0:
+    @register_decomposition(aten._scaled_dot_product_flash_attention.default)
+    def scaled_dot_product_flash_attention(
+        query,
+        key,
+        value,
+        dropout_p=0.0,
+        is_causal=False,
+        *,
+        return_debug_mask=False,
+        scale=None,
+    ):
+        batch_size, num_head, q_size, head_size = (
+            query.shape[0],
+            query.shape[1],
+            query.shape[2],
+            query.shape[3],
+        )
+
+        logsumexp = torch.empty([batch_size, q_size, num_head, head_size], dtype=torch.float)
+        cum_seq_q, cum_seq_k = torch.empty([], dtype=torch.long), torch.empty(
+            [], dtype=torch.long
+        )
+        max_q, max_k = 0, 0
+        philox_seed, philox_offset = torch.empty([], dtype=torch.long), torch.empty(
+            [], dtype=torch.long
+        )
+        debug_attn_mask = torch.empty(
+            [],
+            dtype=query.dtype,
+            device=query.device,
+            requires_grad=query.requires_grad,
+        )
+        output, _ = aten._scaled_dot_product_attention_math.default(
+            query, key, value, None, dropout_p, is_causal, None, scale=scale
+        )
+
+        scores = torch.matmul(query, key.transpose(-2, -1)) / (key.size(-1) ** 0.5)
+        logsumexp = torch.logsumexp(scores, dim=-1)
+
+        output = output.transpose(1, 2).contiguous(memory_format=torch.contiguous_format)
+        return (
+            output.transpose(1, 2),
+            logsumexp,
+            cum_seq_q,
+            cum_seq_k,
+            max_q,
+            max_k,
+            philox_seed,
+            philox_offset,
+            debug_attn_mask,
+        )
+
+
+def get_aot_decomposition_list():
+    return ([torch.ops.aten._scaled_dot_product_flash_attention.default,
+             torch.ops.aten._softmax.default,
+             torch.ops.aten._softmax_backward_data.default,
+             torch.ops.aten.convolution_backward.default,
+             torch.ops.aten.gelu_backward.default,
+             torch.ops.aten.native_group_norm.default,
+             torch.ops.aten.native_group_norm_backward.default,
+             torch.ops.aten.native_layer_norm.default,
+             torch.ops.aten.native_layer_norm_backward.default,
+             torch.ops.aten.slice_backward.default])
+
+def get_inf_decomposition_list():
+    return ([torch.ops.aten.nll_loss_forward.default])