From dd3f953288ceb13406442164fc38a11aa7b942d5 Mon Sep 17 00:00:00 2001
From: Yishuo Wang <yishuo.wang@intel.com>
Date: Wed, 12 Jul 2023 10:11:15 +0800
Subject: [PATCH] Support vnni check (#8497)

---
 python/llm/setup.py                           |  2 +-
 python/llm/src/bigdl/llm/utils/isa_checker.py | 77 +++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 python/llm/src/bigdl/llm/utils/isa_checker.py

diff --git a/python/llm/setup.py b/python/llm/setup.py
index 485096e9..dcf37cf2 100644
--- a/python/llm/setup.py
+++ b/python/llm/setup.py
@@ -199,7 +199,7 @@ def setup_package():
     for url in lib_urls[platform_name]:
         download_libs(url, change_permission=change_permission)
 
-    all_requires = []
+    all_requires = ['py-cpuinfo']
     all_requires += CONVERT_DEP
 
     metadata = dict(
diff --git a/python/llm/src/bigdl/llm/utils/isa_checker.py b/python/llm/src/bigdl/llm/utils/isa_checker.py
new file mode 100644
index 00000000..397ebed9
--- /dev/null
+++ b/python/llm/src/bigdl/llm/utils/isa_checker.py
@@ -0,0 +1,77 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+from cpuinfo import CPUID
+
+
+class ISAChecker:
+    def __init__(self):
+        cpuid = CPUID()
+        self.flags = cpuid.get_flags(cpuid.get_max_extension_support())
+        if self._avx_vnni(cpuid):
+            self.flags.append('avxvnni')
+
+    def _avx_vnni(self, cpuid):
+        eax = cpuid._run_asm(
+            b"\xB9\x01\x00\x00\x00",   # mov ecx, 0x1
+            b"\xB8\x07\x00\x00\x00",   # mov eax, 0x7
+            b"\x0f\xa2",               # cpuid
+            b"\xC3"                    # ret
+        )
+        return ((0x10) & eax) != 0
+
+    def check_avx(self):
+        return 'avx' in self.flags
+
+    def check_avx2(self):
+        return 'avx2' in self.flags
+
+    def check_avx_vnni(self):
+        return 'avxvnni' in self.flags
+
+    def check_avx512(self):
+        return 'avx512f' in self.flags and \
+               'avx512bw' in self.flags and \
+               'avx512cd' in self.flags and \
+               'avx512dq' in self.flags and \
+               'avx512vl' in self.flags
+
+    def check_avx512_vnni(self):
+        return 'avx512vnni' in self.flags
+
+
+isa_checker = ISAChecker()
+
+
+def check_avx():
+    return isa_checker.check_avx()
+
+
+def check_avx2():
+    return isa_checker.check_avx2()
+
+
+def check_avx_vnni():
+    return isa_checker.check_avx_vnni() and isa_checker.check_avx2()
+
+
+def check_avx512():
+    return isa_checker.check_avx512()
+
+
+def check_avx512_vnni():
+    return isa_checker.check_avx512_vnni() and isa_checker.check_avx512()