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

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

@@ -0,0 +1,189 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# mypy: ignore-errors
+
+from copy import deepcopy
+from dataclasses import dataclass
+from functools import lru_cache
+from types import MappingProxyType
+from warnings import warn
+
+import torch
+import torch.overrides
+
+from torch.fx import GraphModule
+from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
+
+from openvino.frontend import FrontEndManager
+from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder
+from openvino.frontend.pytorch.torchdynamo.partition import Partitioner
+from openvino.frontend.pytorch.torchdynamo.compile import openvino_compile
+from openvino.runtime import Core, Type, PartialShape
+from openvino.frontend.pytorch.torchdynamo.backend_utils import _get_cache_dir, _get_device, _get_aot_autograd
+
+from typing import Callable, Optional, Any
+
+from torch.fx.experimental.proxy_tensor import make_fx, wrapper_and_args_for_make_fx
+
+import logging
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+DEFAULT_OPENVINO_PYTHON_CONFIG = MappingProxyType(
+    {
+        "use_python_fusion_cache": True,
+        "allow_single_op_fusion": True,
+    },
+)
+
+compiled_cache = {}
+req_cache = {}
+max_openvino_partitions = 0
+partitioned_modules = {}
+
+
+def execute(
+    gm: GraphModule,
+    *args,
+    executor: str = "openvino",
+    executor_parameters: Optional[dict] = None,
+    options: Optional[Any] = None,
+):
+    if executor == "openvino":
+        return openvino_execute_partitioned(gm, *args, executor_parameters=executor_parameters, options=options)
+    elif executor == "strictly_openvino":
+        return openvino_execute(gm, *args, executor_parameters=executor_parameters)
+
+    msg = "Received unexpected value for 'executor': {0}. Allowed values are: openvino, strictly_openvino.".format(executor)
+    raise ValueError(msg)
+
+
+import numpy as np
+
+
+def execute_cached(compiled_model, *args):
+    ov_inputs = [a.detach().cpu().numpy() for a in args]
+    ov_inputs.reverse()
+    res = compiled_model(ov_inputs)
+    result = [torch.from_numpy(res[out]) for out in compiled_model.outputs]
+    return result
+
+
+def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition_id, options):
+
+    executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG
+
+    use_cache = executor_parameters.get(
+        "use_python_fusion_cache",
+        DEFAULT_OPENVINO_PYTHON_CONFIG["use_python_fusion_cache"],
+    )
+    global compiled_cache
+
+    model_hash_str = executor_parameters.get("model_hash_str", None)
+    if model_hash_str is not None:
+        fully_supported = False
+        if len(model_hash_str) > 3 and model_hash_str[-3:] == "_fs":
+            fully_supported = True
+        if not fully_supported:
+            model_hash_str = model_hash_str + "_p" + str(partition_id)
+
+    if use_cache and (partition_id in compiled_cache):
+        compiled = compiled_cache[partition_id]
+        req = req_cache[partition_id]
+    else:
+        compiled = openvino_compile(gm, *args, model_hash_str=model_hash_str, options=options)
+        compiled_cache[partition_id] = compiled
+        req = compiled.create_infer_request()
+        req_cache[partition_id] = req
+
+    flat_args, _ = tree_flatten(args)
+    ov_inputs = []
+    for arg in flat_args:
+        ov_inputs.append((arg if isinstance(arg, int) else arg.detach().cpu().numpy()))
+
+    res = req.infer(ov_inputs, share_inputs=True, share_outputs=True)
+
+    results1 = [torch.from_numpy(res[out]) for out in compiled.outputs]
+    if len(results1) == 1:
+        return results1[0]
+    return results1
+
+
+class OpenVINOGraphModule(torch.nn.Module):
+    def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None, options=None):
+        super().__init__()
+        self.gm = gm
+        self.partition_id = partition_id
+        self.executor_parameters = {"use_python_fusion_cache": use_python_fusion_cache,
+                                    "model_hash_str": model_hash_str}
+        self.perm_fallback = False
+        self.options = options
+
+    def __call__(self, *args):
+        if self.perm_fallback:
+            return self.gm(*args)
+
+        try:
+            result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id, options=self.options)
+        except Exception:
+            logger.debug("OpenVINO execution failed. Falling back to native PyTorch execution.")
+            self.perm_fallback = True
+            return self.gm(*args)
+
+        return result
+
+
+def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_str: str = None, options=None):
+    global max_openvino_partitions
+    partition_id = max_openvino_partitions
+    for node in gm.graph.nodes:
+        # TODO: use a better way to identify fused submodule
+        if node.op == "call_module" and "fused_" in node.name:
+            openvino_submodule = getattr(gm, node.name)
+            gm.delete_submodule(node.target)
+            gm.add_submodule(
+                node.target,
+                OpenVINOGraphModule(openvino_submodule, partition_id, use_python_fusion_cache,
+                                    model_hash_str=model_hash_str, options=options),
+            )
+            partition_id = partition_id + 1
+
+    max_openvino_partitions = partition_id
+
+    return gm
+
+
+def openvino_execute_partitioned(gm: GraphModule, *args, executor_parameters=None, options=None):
+    executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG
+
+    global partitioned_modules
+
+    use_python_fusion_cache = executor_parameters.get(
+        "use_python_fusion_cache",
+        DEFAULT_OPENVINO_PYTHON_CONFIG["use_python_fusion_cache"],
+    )
+    model_hash_str = executor_parameters.get("model_hash_str", None)
+
+    signature = str(id(gm))
+    if (not _get_aot_autograd(options)):
+        for idx, input_data in enumerate(args):
+            if isinstance(input_data, torch.Tensor):
+                signature = signature + "_" + str(idx) + ":" + str(input_data.type())[6:] + ":" + str(input_data.size())[11:-1].replace(" ", "")
+            else:
+                signature = signature + "_" + str(idx) + ":" + type(input_data).__name__ + ":val(" + str(input_data) + ")"
+
+    if signature not in partitioned_modules:
+        partitioned_modules[signature] = partition_graph(gm, use_python_fusion_cache=use_python_fusion_cache,
+                                                         model_hash_str=model_hash_str, options=options)
+    return partitioned_modules[signature](*args)
+
+
+def clear_caches():
+    global partitioned_modules
+    global compiled_cache
+
+    compiled_cache.clear()
+    partitioned_modules.clear()

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

@@ -0,0 +1,289 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+from typing import Dict
+
+import torch
+from torch.nn import Module
+from torch._ops import OpOverload
+
+from torch.fx.node import Node, _get_qualified_name
+from torch.fx.passes.operator_support import OperatorSupport
+from torch.fx.passes.tools_common import CALLABLE_NODE_OPS
+from openvino.frontend.pytorch.torchdynamo.backend_utils import _get_disabled_ops
+
+import typing as t
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+class OperatorSupport(OperatorSupport):
+    """
+    Operator support for OpenVINO backend.
+    """
+
+    def __init__(self, options):
+        support_dict = {
+            "_operator.add": None,
+            "_operator.floordiv": None,
+            "_operator.getitem": None,
+            "_operator.mul": None,
+            "_operator.sub": None,
+            "torch.ops.aten.sym_size.int": None,
+            "torch.ops.aten._adaptive_avg_pool1d.default": None,
+            "torch.ops.aten._adaptive_avg_pool2d.default": None,
+            "torch.ops.aten._adaptive_avg_pool3d.default": None,
+            "torch.ops.aten._convolution.default": None,
+            "torch.ops.aten._embedding_bag.default": None,
+            "torch.ops.aten._fake_quantize_per_tensor_affine_cachemask_tensor_qparams.default": None,
+            "torch.ops.aten._local_scalar_dense.default": None,
+            "torch.ops.aten._log_softmax.default": None,
+            "torch.ops.aten._native_batch_norm_legit.default": None,
+            "torch.ops.aten._native_batch_norm_legit.no_stats": None,
+            "torch.ops.aten._native_batch_norm_legit_functional.default": None,
+            "torch.ops.aten._native_batch_norm_legit_no_training.default": None,
+            "torch.ops.aten._scaled_dot_product_flash_attention.default": None,
+            "torch.ops.aten._scaled_dot_product_flash_attention_for_cpu.default": None,
+            "torch.ops.aten._softmax.default": None,
+            "torch.ops.aten._to_copy.default": None,
+            "torch.ops.aten._unsafe_view.default": None,
+            "torch.ops.aten.abs.default": None,
+            "torch.ops.aten.acos.default": None,
+            "torch.ops.aten.acosh.default": None,
+            "torch.ops.aten.adaptive_max_pool1d.default": None,
+            "torch.ops.aten.adaptive_max_pool2d.default": None,
+            "torch.ops.aten.adaptive_max_pool3d.default": None,
+            "torch.ops.aten.add.Scalar": None,
+            "torch.ops.aten.add.Tensor": None,
+            "torch.ops.aten.add_.Tensor": None,
+            "torch.ops.aten.addcmul.default": None,
+            "torch.ops.aten.addmm.default": None,
+            "torch.ops.aten.alias.default": None,
+            "torch.ops.aten.all.default": None,
+            "torch.ops.aten.amax.default": None,
+            "torch.ops.aten.amin.default": None,
+            "torch.ops.aten.any.default": None,
+            "torch.ops.aten.any.dim": None,
+            "torch.ops.aten.arange.default": None,
+            "torch.ops.aten.arange.start": None,
+            "torch.ops.aten.arange.start_step": None,
+            "torch.ops.aten.argmax.default": None,
+            "torch.ops.aten.argmin.default": None,
+            "torch.ops.aten.as_strided.default": None,
+            "torch.ops.aten.as_strided_.default": None,
+            "torch.ops.aten.asin.default": None,
+            "torch.ops.aten.asinh.default": None,
+            "torch.ops.aten.asinh.default": None,
+            "torch.ops.aten.atanh.default": None,
+            "torch.ops.aten.avg_pool2d.default": None,
+            "torch.ops.aten.avg_pool3d.default": None,
+            "torch.ops.aten.baddbmm.default": None,
+            "torch.ops.aten.bitwise_and.Scalar": None,
+            "torch.ops.aten.bitwise_and.Tensor": None,
+            "torch.ops.aten.bitwise_not.default": None,
+            "torch.ops.aten.bitwise_or.Tensor": None,
+            "torch.ops.aten.bitwise_xor.Tensor": None,
+            "torch.ops.aten.bmm.default": None,
+            "torch.ops.aten.cat.default": None,
+            "torch.ops.aten.ceil.default": None,
+            "torch.ops.aten.clamp.default": None,
+            "torch.ops.aten.clamp_max.default": None,
+            "torch.ops.aten.clamp_max.Tensor": None,
+            "torch.ops.aten.clamp_min.default": None,
+            "torch.ops.aten.clamp_min.Tensor": None,
+            "torch.ops.aten.clone.default": None,
+            "torch.ops.aten.constant_pad_nd.default": None,
+            "torch.ops.aten.convolution.default": None,
+            "torch.ops.aten.copy.default": None,
+            "torch.ops.aten.copy_.default": None,
+            "torch.ops.aten.cos.default": None,
+            "torch.ops.aten.cosh.default": None,
+            "torch.ops.aten.cumsum.default": None,
+            "torch.ops.aten.detach.default": None,
+            "torch.ops.aten.detach_.default": None,
+            "torch.ops.aten.div.Scalar": None,
+            "torch.ops.aten.div.Tensor": None,
+            "torch.ops.aten.div.Tensor_mode": None,
+            "torch.ops.aten.div_.Tensor": None,
+            "torch.ops.aten.elu.default": None,
+            "torch.ops.aten.elu_.default": None,
+            "torch.ops.aten.embedding.default": None,
+            "torch.ops.aten.empty.memory_format": None,
+            "torch.ops.aten.eq.Scalar": None,
+            "torch.ops.aten.eq.Tensor": None,
+            "torch.ops.aten.erf.default": None,
+            "torch.ops.aten.exp.default": None,
+            "torch.ops.aten.expand.default": None,
+            "torch.ops.aten.fake_quantize_per_channel_affine_cachemask.default": None,
+            "torch.ops.aten.fill.Scalar": None,
+            "torch.ops.aten.fill_.Scalar": None,
+            "torch.ops.aten.fill.Tensor": None,
+            "torch.ops.aten.fill_.Tensor": None,
+            "torch.ops.aten.flip.default": None,
+            "torch.ops.aten.floor.default": None,
+            "torch.ops.aten.floor.default": None,
+            "torch.ops.aten.fmod.Scalar": None,
+            "torch.ops.aten.fmod.Tensor": None,
+            "torch.ops.aten.full.default": None,
+            "torch.ops.aten.full.names": None,
+            "torch.ops.aten.full_like.default": None,
+            "torch.ops.aten.gather.default": None,
+            "torch.ops.aten.ge.Scalar": None,
+            "torch.ops.aten.ge.Tensor": None,
+            "torch.ops.aten.gelu.default": None,
+            "torch.ops.aten.glu.default": None,
+            "torch.ops.aten.grid_sampler_2d.default": None,
+            "torch.ops.aten.gt.Scalar": None,
+            "torch.ops.aten.gt.Tensor": None,
+            "torch.ops.aten.hardsigmoid.default": None,
+            "torch.ops.aten.hardswish.default": None,
+            "torch.ops.aten.hardswish_.default": None,
+            "torch.ops.aten.hardtanh.default": None,
+            "torch.ops.aten.hardtanh_.default": None,
+            "torch.ops.aten.index.Tensor": None,
+            "torch.ops.aten.index_select.default": None,
+            "torch.ops.aten.isfinite.default": None,
+            "torch.ops.aten.isinf.default": None,
+            "torch.ops.aten.isnan.default": None,
+            "torch.ops.aten.le.Scalar": None,
+            "torch.ops.aten.le.Tensor": None,
+            "torch.ops.aten.leaky_relu.default": None,
+            "torch.ops.aten.leaky_relu_.default": None,
+            "torch.ops.aten.lift_fresh_copy.default": None,
+            "torch.ops.aten.linalg_vector_norm.default": None,
+            "torch.ops.aten.log.default": None,
+            "torch.ops.aten.log_sigmoid_forward.default": None,
+            "torch.ops.aten.log10.default": None,
+            "torch.ops.aten.log1p.default": None,
+            "torch.ops.aten.log2.default": None,
+            "torch.ops.aten.logical_not.default": None,
+            "torch.ops.aten.logsumexp.default": None,
+            "torch.ops.aten.lt.Scalar": None,
+            "torch.ops.aten.lt.Tensor": None,
+            "torch.ops.aten.masked_fill.Scalar": None,
+            "torch.ops.aten.masked_fill.Tensor": None,
+            "torch.ops.aten.masked_fill_.Scalar": None,
+            "torch.ops.aten.masked_fill_.Tensor": None,
+            "torch.ops.aten.max.default": None,
+            "torch.ops.aten.max.dim": None,
+            "torch.ops.aten.max_pool2d_with_indices.default": None,
+            "torch.ops.aten.max_pool3d_with_indices.default": None,
+            "torch.ops.aten.maximum.default": None,
+            "torch.ops.aten.mean.default": None,
+            "torch.ops.aten.mean.dim": None,
+            "torch.ops.aten.min.default": None,
+            "torch.ops.aten.min.dim": None,
+            "torch.ops.aten.minimum.default": None,
+            "torch.ops.aten.mm.default": None,
+            "torch.ops.aten.mul.Scalar": None,
+            "torch.ops.aten.mul.Tensor": None,
+            "torch.ops.aten.mul_.Tensor": None,
+            "torch.ops.aten.native_batch_norm.default": None,
+            "torch.ops.aten.native_dropout.default": None,
+            "torch.ops.aten.native_group_norm.default": None,
+            "torch.ops.aten.native_layer_norm.default": None,
+            "torch.ops.aten.ne.Scalar": None,
+            "torch.ops.aten.ne.Tensor": None,
+            "torch.ops.aten.neg.default": None,
+            "torch.ops.aten.new_full.default": None,
+            "torch.ops.aten.new_ones.default": None,
+            "torch.ops.aten.ones_like.default": None,
+            "torch.ops.aten.new_zeros.default": None,
+            "torch.ops.aten.ones.default": None,
+            "torch.ops.aten.permute.default": None,
+            "torch.ops.aten.pow.Scalar": None,
+            "torch.ops.aten.pow.Tensor_Scalar": None,
+            "torch.ops.aten.pow.Tensor_Tensor": None,
+            "torch.ops.aten.rand.default": None,
+            "torch.ops.aten.reflection_pad2d.default": None,
+            "torch.ops.aten.reciprocal.default": None,
+            "torch.ops.aten.relu.default": None,
+            "torch.ops.aten.relu_.default": None,
+            "torch.ops.aten.repeat.default": None,
+            "torch.ops.aten.roll.default": None,
+            "torch.ops.aten.rsqrt.default": None,
+            "torch.ops.aten.rsub.Scalar": None,
+            "torch.ops.aten.rsub.Tensor": None,
+            "torch.ops.aten.scalar_tensor.default": None,
+            "torch.ops.aten.scatter.src": None,
+            "torch.ops.aten.scatter.value": None,
+            "torch.ops.aten.select.int": None,
+            "torch.ops.aten.select_scatter.default": None,
+            "torch.ops.aten.sigmoid.default": None,
+            "torch.ops.aten.sigmoid_.default": None,
+            "torch.ops.aten.sign.default": None,
+            "torch.ops.aten.silu.default": None,
+            "torch.ops.aten.silu_.default": None,
+            "torch.ops.aten.sin.default": None,
+            "torch.ops.aten.sinh.default": None,
+            "torch.ops.aten.slice.Tensor": None,
+            "torch.ops.aten.slice_scatter.default": None,
+            "torch.ops.aten.sort.default": None,
+            "torch.ops.aten.split.Tensor": None,
+            "torch.ops.aten.split_with_sizes.default": None,
+            "torch.ops.aten.sqrt.default": None,
+            "torch.ops.aten.squeeze.dim": None,
+            "torch.ops.aten.squeeze.dims": None,
+            "torch.ops.aten.stack.default": None,
+            "torch.ops.aten.std.correction": None,
+            "torch.ops.aten.sub.default": None,
+            "torch.ops.aten.sub.Tensor": None,
+            "torch.ops.aten.sum.default": None,
+            "torch.ops.aten.sum.dim_IntList": None,
+            "torch.ops.aten.t.default": None,
+            "torch.ops.aten.tan.default": None,
+            "torch.ops.aten.tanh.default": None,
+            "torch.ops.aten.topk.default": None,
+            "torch.ops.aten.transpose.int": None,
+            "torch.ops.aten.tril.default": None,
+            "torch.ops.aten.tril_.default": None,
+            "torch.ops.aten.unbind.int": None,
+            "torch.ops.aten.unfold.default": None,
+            "torch.ops.aten.unsqueeze.default": None,
+            "torch.ops.aten.upsample_nearest2d.default": None,
+            "torch.ops.aten.var.correction": None,
+            "torch.ops.aten.var_mean.correction": None,
+            "torch.ops.aten.view.default": None,
+            "torch.ops.aten.where.self": None,
+            "torch.ops.aten.zeros.default": None,
+            "torch.ops.aten.zeros_like.default": None,
+            "torch.ops.torchvision.deform_conv2d.default": None,
+            "torch.ops.torchvision.roi_align.default": None,
+            "torch.ops.quantized_decomposed.quantize_per_tensor.default": None,
+            "torch.ops.quantized_decomposed.quantize_per_channel.default": None,
+            "torch.ops.quantized_decomposed.dequantize_per_tensor.default": None,
+            "torch.ops.quantized_decomposed.dequantize_per_channel.default": None
+
+        }
+            
+        self.enabled_op_names = []
+
+        for op in _get_disabled_ops(options):
+            del support_dict[op]
+
+        super().__init__(support_dict)
+
+    def enable_by_name(self, node: Node):
+        self.enabled_op_names.append(node.name)
+
+    def is_node_supported(self, submodules: t.Mapping[str, Module], node: Node) -> bool:
+        # OpenVINO FX subgraph should be purely functional
+        if node.op not in CALLABLE_NODE_OPS:
+            return False
+
+        # ops in supported_dict doesn't have overload name
+        # use overloadpacket's qualified_name for OpOverload
+        if isinstance(node.target, OpOverload):
+            target = _get_qualified_name(node.target.overloadpacket)
+
+            if target in self._support_dict:
+                return True
+
+        if node.name in self.enabled_op_names:
+            return True
+
+        return super().is_node_supported(submodules, node)

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

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# mypy: ignore-errors
+
+from typing import Dict
+
+import torch
+from torch.nn import Module
+from torch.fx import GraphModule, Node
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
+
+from torch.fx.experimental.proxy_tensor import DecompositionInterpreter
+from torch._decomp import decomposition_table
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
+from openvino.frontend.pytorch.torchdynamo.op_support import OperatorSupport
+from openvino.frontend.pytorch.torchdynamo.backend_utils import _is_testing
+
+import typing as t
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+class Partitioner:
+    def __init__(self, options):
+        self.supported_ops = OperatorSupport(options)
+
+    def fx_serialize(self, graph_module: GraphModule, *args, **kwargs):
+        fx_gm = make_fx(graph_module)(*args)
+        return fx_gm
+
+    def add_get_attr_inputs(self, partitions: t.List[Partition]):
+        # TODO: Find a more efficient way to include input
+        # "get_attr" nodes to the partitions.
+        getattr_to_merge: Dict[Node, Node] = {}
+        for partition in partitions:
+            for pnode in partition.nodes:
+                for pnode_input in pnode.all_input_nodes:
+                    if pnode_input.op in ["get_attr"] and pnode_input.op not in getattr_to_merge:
+                        getattr_to_merge[pnode_input] = partition
+        for getattr_node, getattr_part in getattr_to_merge.items():
+            getattr_part.add_node(getattr_node)
+
+    def check_fully_supported(self, graph_module: GraphModule) -> bool:
+        num_fused = 0
+        for node in graph_module.graph.nodes:
+            if node.op == "call_module" and "fused_" in node.name:
+                num_fused += 1
+            elif node.op != "placeholder" and node.op != "output":
+                return False
+        if num_fused == 1:
+            return True
+        return False
+
+    def capture_gptq_patterns(self, graph_module: GraphModule) -> bool:
+        for node in graph_module.graph.nodes:
+            if str(node.op) == "call_function" and str(node.target) == "aten.bitwise_and.Scalar":
+                bitwise_and_in_nodes = node.all_input_nodes
+                if len(bitwise_and_in_nodes) != 1:
+                    continue
+                to_copy_node = bitwise_and_in_nodes[0]
+                if str(to_copy_node.op) != "call_function" or str(to_copy_node.target) != "aten._to_copy.default":
+                    continue
+                to_copy_in_nodes = to_copy_node.all_input_nodes
+                if len(to_copy_in_nodes) != 1:
+                    continue
+                bitwise_right_shift_node = to_copy_in_nodes[0]
+                if str(bitwise_right_shift_node.op) != "call_function" or str(bitwise_right_shift_node.target) != "aten.bitwise_right_shift.Tensor":
+                    continue
+                bitwise_right_shift_in_nodes = bitwise_right_shift_node.all_input_nodes
+                if len(bitwise_right_shift_in_nodes) != 2:
+                    continue
+                expand_node = bitwise_right_shift_in_nodes[0]
+                if str(expand_node.op) != "call_function" or str(expand_node.target) != "aten.expand.default":
+                    continue
+                expand_in_nodes = expand_node.all_input_nodes
+                if len(expand_in_nodes) != 1:
+                    continue
+                unsqueeze_0_node = expand_in_nodes[0]
+                if str(unsqueeze_0_node.op) != "call_function" or str(unsqueeze_0_node.target) != "aten.unsqueeze.default":
+                    continue
+                unsqueeze_0_in_nodes = unsqueeze_0_node.all_input_nodes
+                if len(unsqueeze_0_in_nodes) != 1:
+                    continue
+                const_0_node = unsqueeze_0_in_nodes[0]
+                if str(const_0_node.op) != "get_attr":
+                    continue
+                unsqueeze_1_node = bitwise_right_shift_in_nodes[1]
+                if str(unsqueeze_1_node.op) != "call_function" or str(unsqueeze_1_node.target) != "aten.unsqueeze.default":
+                    continue
+                unsqueeze_1_in_nodes = unsqueeze_1_node.all_input_nodes
+                if len(unsqueeze_1_in_nodes) != 1:
+                    continue
+                const_1_node = unsqueeze_1_in_nodes[0]
+                if str(const_1_node.op) != "get_attr":
+                    continue
+
+                self.supported_ops.enable_by_name(node)
+                self.supported_ops.enable_by_name(to_copy_node)
+                self.supported_ops.enable_by_name(bitwise_right_shift_node)
+                self.supported_ops.enable_by_name(expand_node)
+                self.supported_ops.enable_by_name(unsqueeze_0_node)
+                self.supported_ops.enable_by_name(unsqueeze_1_node)
+
+    def make_partitions(self, graph_module: GraphModule, options) -> GraphModule:
+        allow_single_node_partition = _is_testing(options)
+        self.capture_gptq_patterns(graph_module)
+        partitioner = CapabilityBasedPartitioner(
+            graph_module, self.supported_ops, allows_single_node_partition=allow_single_node_partition)
+        partitions = partitioner.propose_partitions()
+        self.add_get_attr_inputs(partitions)
+        fused_graph_module = partitioner.fuse_partitions(partitions)
+
+        return fused_graph_module