From 843340d3b3cf51ef5c0b9876fcb6598d811e54b0 Mon Sep 17 00:00:00 2001
From: Disty0 <semihgulec2005@gmail.com>
Date: Sat, 23 Dec 2023 15:30:33 +0300
Subject: [PATCH] OpenVINO 8 bit support for CPUs

---
 CHANGELOG.md                       |  1 +
 installer.py                       |  1 +
 modules/intel/openvino/__init__.py | 32 ++++++++++++++++++++----------
 modules/shared.py                  |  1 +
 4 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8a2d44611..cdac8d8e4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -75,6 +75,7 @@
     - add `DISABLE_IPEXRUN` and `DISABLE_IPEX_1024_WA` environment variables  
     - compatibility improvements  
   - **OpenVINO**, thanks @disty0  
+    - **8 bit support for CPUs**  
     - add *Directory for OpenVINO cache* option to *System Paths*  
     - remove Intel ARC specific 1024x1024 workaround  
   - **HDR controls**  
diff --git a/installer.py b/installer.py
index 6c2798c99..29d65bb5d 100644
--- a/installer.py
+++ b/installer.py
@@ -523,6 +523,7 @@ def check_torch():
         install('hidet', 'hidet')
     if args.use_openvino or opts.get('cuda_compile_backend', '') == 'openvino_fx':
         install(os.environ.get('OPENVINO_PACKAGE', 'openvino==2023.2.0'), 'openvino')
+        install('nncf==2.7.0', 'nncf')
         install('onnxruntime-openvino', 'onnxruntime-openvino', ignore=True) # TODO openvino: numpy version conflicts with tensorflow and doesn't support Python 3.11
         os.environ.setdefault('PYTORCH_TRACING_MODE', 'TORCHFX')
         os.environ.setdefault('NEOReadDebugKeys', '1')
diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py
index 6bd090819..d9650ed20 100644
--- a/modules/intel/openvino/__init__.py
+++ b/modules/intel/openvino/__init__.py
@@ -1,20 +1,35 @@
 import os
 import sys
 import torch
+import nncf
+
 from openvino.frontend import FrontEndManager
 from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder
 from openvino.frontend.pytorch.torchdynamo.partition import Partitioner
 from openvino.runtime import Core, Type, PartialShape, serialize
+
 from torch._dynamo.backends.common import fake_tensor_unsupported
 from torch._dynamo.backends.registry import register_backend
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.fx import GraphModule
 from torch.utils._pytree import tree_flatten
+
 from types import MappingProxyType
 from hashlib import sha256
 import functools
+
 from modules import shared, devices
 
+NNCFNodeName = str
+def get_node_by_name(self, name: NNCFNodeName) ->  nncf.common.graph.NNCFNode:
+    node_ids = self._node_name_to_node_id_map.get(name, None)
+    if node_ids is None:
+        raise RuntimeError("Could not find a node {} in NNCFGraph!".format(name))
+
+    node_key = f"{node_ids[0]} {name}"
+    return self._nodes[node_key]
+nncf.common.graph.NNCFGraph.get_node_by_name = get_node_by_name
+
 def BUILD_MAP_UNPACK(self, inst):
         items = self.popn(inst.argval)
         # ensure everything is a dict
@@ -50,18 +65,9 @@ class OpenVINOGraphModule(torch.nn.Module):
         self.executor_parameters = {"use_python_fusion_cache": use_python_fusion_cache,
                                     "model_hash_str": model_hash_str}
         self.file_name = file_name
-        self.perm_fallback = False
 
     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, file_name=self.file_name)
-        #except Exception:
-        #    self.perm_fallback = True
-        #    return self.gm(*args)
-
         return result
 
 def get_device():
@@ -223,12 +229,14 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None, file_na
         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()
+    if shared.opts.openvino_compress_weights:
+        om = nncf.compress_weights(om)
 
     if model_hash_str is not None:
         core.set_property({'CACHE_DIR': cache_root + '/blob'})
 
-    compiled = core.compile_model(om, device)
-    return compiled
+    compiled_model = core.compile_model(om, device)
+    return compiled_model
 
 def openvino_compile_cached_model(cached_model_path, *example_inputs):
     core = Core()
@@ -249,6 +257,8 @@ def openvino_compile_cached_model(cached_model_path, *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()
+    if shared.opts.openvino_compress_weights:
+        om = nncf.compress_weights(om)
 
     core.set_property({'CACHE_DIR': shared.opts.openvino_cache_path + '/blob'})
 
diff --git a/modules/shared.py b/modules/shared.py
index b8068cbff..05cdcc9fb 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -329,6 +329,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "openvino_hetero_gpu": OptionInfo(False, "OpenVINO use Hetero Device for single inference with multiple devices"),
     "openvino_remove_cpu_from_hetero": OptionInfo(False, "OpenVINO remove CPU from Hetero Device"),
     "openvino_remove_igpu_from_hetero": OptionInfo(False, "OpenVINO remove iGPU from Hetero Device"),
+    "openvino_compress_weights": OptionInfo(False, "OpenVINO compress weights to 8 bit (CPU Only)"),
 }))
 
 options_templates.update(options_section(('advanced', "Inference Settings"), {