[LLM] Fix linux load libs for NeoX and llama (#8257)

* init

* add lisence

* fix style

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

@@ -63,28 +63,29 @@ from ctypes import (
 )
 import pathlib
 from bigdl.llm.utils.common import invalidInputError
+from bigdl.llm.utils import get_avx_flags
 
 
 # Load the library
 def _load_shared_library(lib_base_name: str):
     # Determine the file extension based on the platform
-    if sys.platform.startswith("linux"):
+    if sys.platform.startswith("linux") or sys.platform == "darwin":
-        lib_ext = ".so"
-    elif 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}{lib_ext}",
+        _base_path / f"lib{lib_base_name}{avx}{lib_ext}",
-        _base_path / f"{lib_base_name}{lib_ext}",
+        _base_path / f"{lib_base_name}{avx}{lib_ext}",
     ]
 
     if "GPTNEOX_CPP_LIB" in os.environ:

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

@@ -63,28 +63,29 @@ from ctypes import (
 )
 import pathlib
 from bigdl.llm.utils.common import invalidInputError
+from bigdl.llm.utils import get_avx_flags
 
 
 # Load the library
 def _load_shared_library(lib_base_name: str):
     # Determine the file extension based on the platform
-    if sys.platform.startswith("linux"):
+    if sys.platform.startswith("linux") or sys.platform == "darwin":
-        lib_ext = ".so"
-    elif 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}{lib_ext}",
+        _base_path / f"lib{lib_base_name}{avx}{lib_ext}",
-        _base_path / f"{lib_base_name}{lib_ext}",
+        _base_path / f"{lib_base_name}{avx}{lib_ext}",
     ]
 
     if "LLAMA_CPP_LIB" in os.environ:

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

@@ -18,3 +18,5 @@
 # physically located elsewhere.
 # Otherwise there would be module not found error in non-pip's setting as Python would
 # only search the first bigdl package and end up finding only one sub-package.
+
+from .utils import *

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

@@ -19,4 +19,4 @@
 # Otherwise there would be module not found error in non-pip's setting as Python would
 # only search the first bigdl package and end up finding only one sub-package.
 
-from .log4Error import invalidInputError
+from .log4Error import invalidInputError, invalidOperationError

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

@@ -0,0 +1,32 @@
+#
+# 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.
+#
+
+import sys
+from bigdl.llm.utils.common import invalidInputError, invalidOperationError
+
+
+def get_avx_flags():
+    avx = ""
+    if sys.platform != "win32":
+        import subprocess
+        msg = subprocess.check_output(["lscpu"]).decode("utf-8")
+        if "avx512_vnni" in msg:
+            avx = "_avx512"
+        elif "avx2" in msg:
+            avx = "_avx2"
+        else:
+            invalidOperationError(False, "Unsupported CPUFLAGS.")
+    return avx