[LLM] llm supports vnni link on windows (#8543)

* support win vnni link

* fix style

* fix style

* use isa_checker

* fix

* typo

* fix

* update

python/llm/src/bigdl/llm/ggml/model/bloom/bloom_cpp.py

@@ -64,31 +64,13 @@ from ctypes import (
     c_size_t,
 )
 import pathlib
-from bigdl.llm.utils import get_avx_flags
+from bigdl.llm.utils import get_shared_lib_info
 from bigdl.llm.utils.common import invalidInputError
 
 
 # Load the library
 def _load_shared_library(lib_base_name: str):
-    # Determine the file extension based on the platform
+    _base_path, _lib_paths = get_shared_lib_info(lib_base_name=lib_base_name)
-    if sys.platform.startswith("linux") or sys.platform == "darwin":
-        lib_ext = ".so"
-    elif sys.platform == "win32":
-        lib_ext = ".dll"
-    else:
-        invalidInputError(False, "Unsupported platform")
-
-    avx = get_avx_flags()
-
-    # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs)
-    _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve()
-    _base_path = _base_path / 'libs'
-    # Searching for the library in the current directory under the name "libbloom" (default name
-    # for bloomcpp) and "bloom" (default name for this repo)
-    _lib_paths = [
-        _base_path / f"lib{lib_base_name}{avx}{lib_ext}",
-        _base_path / f"{lib_base_name}{avx}{lib_ext}",
-    ]
 
     if "BLOOM_CPP_LIB" in os.environ:
         lib_base_name = os.environ["BLOOM_CPP_LIB"]

python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py

@@ -64,30 +64,12 @@ from ctypes import (
 )
 import pathlib
 from bigdl.llm.utils.common import invalidInputError
-from bigdl.llm.utils import get_avx_flags
+from bigdl.llm.utils import get_shared_lib_info
 
 
 # Load the library
 def _load_shared_library(lib_base_name: str):
-    # Determine the file extension based on the platform
+    _base_path, _lib_paths = get_shared_lib_info(lib_base_name=lib_base_name)
-    if sys.platform.startswith("linux") or sys.platform == "darwin":
-        lib_ext = ".so"
-    elif sys.platform == "win32":
-        lib_ext = ".dll"
-    else:
-        invalidInputError(False, "Unsupported platform.")
-
-    avx = get_avx_flags()
-
-    # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs)
-    _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve()
-    _base_path = _base_path / 'libs'
-    # Searching for the library in the current directory under the name "libgptneox" (default name
-    # for gptneoxcpp) and "gptneox" (default name for this repo)
-    _lib_paths = [
-        _base_path / f"lib{lib_base_name}{avx}{lib_ext}",
-        _base_path / f"{lib_base_name}{avx}{lib_ext}",
-    ]
 
     if "GPTNEOX_CPP_LIB" in os.environ:
         lib_base_name = os.environ["GPTNEOX_CPP_LIB"]

python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py

@@ -64,30 +64,12 @@ from ctypes import (
 )
 import pathlib
 from bigdl.llm.utils.common import invalidInputError
-from bigdl.llm.utils import get_avx_flags
+from bigdl.llm.utils import get_shared_lib_info
 
 
 # Load the library
 def _load_shared_library(lib_base_name: str):
-    # Determine the file extension based on the platform
+    _base_path, _lib_paths = get_shared_lib_info(lib_base_name=lib_base_name)
-    if sys.platform.startswith("linux") or sys.platform == "darwin":
-        lib_ext = ".so"
-    elif sys.platform == "win32":
-        lib_ext = ".dll"
-    else:
-        invalidInputError(False, "Unsupported platform.")
-
-    avx = get_avx_flags()
-
-    # Construct the paths to the possible shared library names
-    _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve()
-    _base_path = _base_path / 'libs'
-    # Searching for the library in the current directory under the name "libllama" (default name
-    # for llamacpp) and "llama" (default name for this repo)
-    _lib_paths = [
-        _base_path / f"lib{lib_base_name}{avx}{lib_ext}",
-        _base_path / f"{lib_base_name}{avx}{lib_ext}",
-    ]
 
     if "LLAMA_CPP_LIB" in os.environ:
         lib_base_name = os.environ["LLAMA_CPP_LIB"]

python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder_cpp.py

@@ -64,31 +64,13 @@ from ctypes import (
     c_size_t,
 )
 import pathlib
-from bigdl.llm.utils import get_avx_flags
+from bigdl.llm.utils import get_shared_lib_info
 from bigdl.llm.utils.common import invalidInputError
 
 
 # Load the library
 def _load_shared_library(lib_base_name: str):
-    # Determine the file extension based on the platform
+    _base_path, _lib_paths = get_shared_lib_info(lib_base_name=lib_base_name)
-    if sys.platform.startswith("linux") or sys.platform == "darwin":
-        lib_ext = ".so"
-    elif sys.platform == "win32":
-        lib_ext = ".dll"
-    else:
-        invalidInputError(False, "Unsupported platform")
-
-    avx = get_avx_flags()
-
-    # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs)
-    _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve()
-    _base_path = _base_path / 'libs'
-    # Searching for the library in the current directory under the name "libbloom" (default name
-    # for bloomcpp) and "bloom" (default name for this repo)
-    _lib_paths = [
-        _base_path / f"lib{lib_base_name}{avx}{lib_ext}",
-        _base_path / f"{lib_base_name}{avx}{lib_ext}",
-    ]
 
     if "STARCODER_CPP_LIB" in os.environ:
         lib_base_name = os.environ["STARCODER_CPP_LIB"]

python/llm/src/bigdl/llm/utils/utils.py

@@ -15,18 +15,44 @@
 #
 
 import sys
+import pathlib
+from bigdl.llm.utils.isa_checker import check_avx_vnni, check_avx2, check_avx512_vnni
 from bigdl.llm.utils.common import invalidInputError, invalidOperationError
 
 
-def get_avx_flags():
+def get_cpu_flags():
-    avx = ""
+    flags = ""
     if sys.platform != "win32":
-        import subprocess
+        if check_avx512_vnni():
-        msg = subprocess.check_output(["lscpu"]).decode("utf-8")
+            flags = "_avx512"
-        if "avx512_vnni" in msg:
+        elif check_avx_vnni():
-            avx = "_avx512"
+            flags = "_avx2"
-        elif "avx2" in msg:
-            avx = "_avx2"
         else:
             invalidOperationError(False, "Unsupported CPUFLAGS.")
-    return avx
+    else:
+        flags = "_vnni" if check_avx_vnni() else ""
+    return flags
+
+
+def get_shared_lib_info(lib_base_name: str):
+    # Determine the file extension based on the platform
+    if sys.platform.startswith("linux") or sys.platform == "darwin":
+        lib_ext = ".so"
+    elif sys.platform == "win32":
+        lib_ext = ".dll"
+    else:
+        invalidInputError(False, "Unsupported platform.")
+
+    cpuflags = get_cpu_flags()
+
+    # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs)
+    _base_path = pathlib.Path(__file__).parent.parent.resolve()
+    _base_path = _base_path / 'libs'
+    # Searching for the library in the current directory under the name "lib{lib_base_name}"
+    # (default name for llmcpp) and "{lib_base_name}" (default name for this repo)
+    _lib_paths = [
+        _base_path / f"lib{lib_base_name}{cpuflags}{lib_ext}",
+        _base_path / f"{lib_base_name}{cpuflags}{lib_ext}",
+    ]
+
+    return _base_path, _lib_paths

python/llm/test/inference/test_call_models.py

@@ -16,7 +16,6 @@
 
 
 from bigdl.llm.models import Llama, Bloom, Gptneox, Starcoder
-from bigdl.llm.utils import get_avx_flags
 import pytest
 from unittest import TestCase
 import os

python/llm/test/langchain/test_langchain.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 #
 
-from bigdl.llm.utils import get_avx_flags
 from bigdl.llm.langchain.embeddings import BigdlNativeEmbeddings
 from bigdl.llm.langchain.llms import BigdlNativeLLM
 import pytest