fix: add build for armv82 and v84 and sync llama.cpp (#67)

* fix: add build for armv82 and v84 and sync llama.cpp * chore(android): format * chore: revert unnecessary change --------- Co-authored-by: jhen <[email protected]>
mybigday · Jul 22, 2024 · 1942159 · 1942159
1 parent 39d1e05
commit 1942159
Show file tree

Hide file tree

Showing 36 changed files with 19,428 additions and 6,810 deletions.
diff --git a/android/src/main/CMakeLists.txt b/android/src/main/CMakeLists.txt
@@ -22,13 +22,14 @@ set(
     ${RNLLAMA_LIB_DIR}/unicode.cpp
     ${RNLLAMA_LIB_DIR}/llama.cpp
     ${RNLLAMA_LIB_DIR}/sgemm.cpp
+    ${RNLLAMA_LIB_DIR}/ggml-aarch64.c
     ${RNLLAMA_LIB_DIR}/rn-llama.hpp
     ${CMAKE_SOURCE_DIR}/jni.cpp
 )
 
 find_library(LOG_LIB log)
 
-function(build_library target_name)
+function(build_library target_name cpu_flags)
     add_library(
         ${target_name}
         SHARED
@@ -37,32 +38,34 @@ function(build_library target_name)
 
     target_link_libraries(${target_name} ${LOG_LIB} android)
 
-    target_compile_options(${target_name} PRIVATE -pthread)
-
-    if (${target_name} STREQUAL "rnllama_v8fp16_va")
-        target_compile_options(${target_name} PRIVATE -march=armv8.4-a+fp16+dotprod)
-    endif ()
+    target_compile_options(${target_name} PRIVATE -pthread ${cpu_flags})
 
     if (${CMAKE_BUILD_TYPE} STREQUAL "Debug")
         target_compile_options(${target_name} PRIVATE -DRNLLAMA_ANDROID_ENABLE_LOGGING)
     endif ()
 
-    # NOTE: If you want to debug the native code, you can uncomment if and endif
-    # if (NOT ${CMAKE_BUILD_TYPE} STREQUAL "Debug")
-
-    target_compile_options(${target_name} PRIVATE -O3 -DNDEBUG)
-    target_compile_options(${target_name} PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)
-    target_compile_options(${target_name} PRIVATE -ffunction-sections -fdata-sections)
+    if (NOT ${CMAKE_BUILD_TYPE} STREQUAL "Debug")
+        target_compile_options(${target_name} PRIVATE -O3 -DNDEBUG)
+        target_compile_options(${target_name} PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)
+        target_compile_options(${target_name} PRIVATE -ffunction-sections -fdata-sections)
 
-    target_link_options(${target_name} PRIVATE -Wl,--gc-sections)
-    target_link_options(${target_name} PRIVATE -Wl,--exclude-libs,ALL)
-    target_link_options(${target_name} PRIVATE -flto)
-
-    # endif ()
+        target_link_options(${target_name} PRIVATE -Wl,--gc-sections)
+        target_link_options(${target_name} PRIVATE -Wl,--exclude-libs,ALL)
+        target_link_options(${target_name} PRIVATE -flto)
+    endif ()
 endfunction()
 
-build_library("rnllama") # Default target
+# Default target (no specific CPU features)
+build_library("rnllama" "")
 
 if (${ANDROID_ABI} STREQUAL "arm64-v8a")
-    build_library("rnllama_v8fp16_va")
+    # ARM64 targets
+    build_library("rnllama_v8_4_fp16_dotprod" "-march=armv8.4-a+fp16+dotprod")
+    build_library("rnllama_v8_2_fp16_dotprod" "-march=armv8.2-a+fp16+dotprod")
+    build_library("rnllama_v8_2_fp16" "-march=armv8.2-a+fp16")
+    build_library("rnllama_v8" "-march=armv8-a")
+elseif (${ANDROID_ABI} STREQUAL "x86_64")
+    # x86_64 target
+    build_library("rnllama_x86_64" "-march=x86-64" "-mtune=intel" "-msse4.2" "-mpopcnt")
+
 endif ()
diff --git a/android/src/main/java/com/rnllama/LlamaContext.java b/android/src/main/java/com/rnllama/LlamaContext.java
@@ -237,28 +237,32 @@ public void release() {
   static {
     Log.d(NAME, "Primary ABI: " + Build.SUPPORTED_ABIS[0]);
     if (LlamaContext.isArm64V8a()) {
-      boolean loadV8fp16 = false;
-      if (LlamaContext.isArm64V8a()) {
-        // ARMv8.2a needs runtime detection support
-        String cpuInfo = LlamaContext.cpuInfo();
-        if (cpuInfo != null) {
-          Log.d(NAME, "CPU info: " + cpuInfo);
-          if (cpuInfo.contains("fphp")) {
-            Log.d(NAME, "CPU supports fp16 arithmetic");
-            loadV8fp16 = true;
-          }
-        }
-      }
+      String cpuFeatures = LlamaContext.getCpuFeatures();
+      Log.d(NAME, "CPU features: " + cpuFeatures);
 
-      if (loadV8fp16) {
-        Log.d(NAME, "Loading librnllama_v8fp16_va.so");
-        System.loadLibrary("rnllama_v8fp16_va");
+      boolean hasFp16 = cpuFeatures.contains("fp16") || cpuFeatures.contains("fphp");
+      boolean hasDotProd = cpuFeatures.contains("dotprod") || cpuFeatures.contains("asimddp");
+      boolean isAtLeastArmV82 = cpuFeatures.contains("asimd") && cpuFeatures.contains("crc32") && cpuFeatures.contains("aes");
+      boolean isAtLeastArmV84 = cpuFeatures.contains("dcpop") && cpuFeatures.contains("uscat");
+
+      if (isAtLeastArmV84 && hasFp16 && hasDotProd) {
+        Log.d(NAME, "Loading librnllama_v8_4_fp16_dotprod.so");
+        System.loadLibrary("rnllama_v8_4_fp16_dotprod");
+      } else if (isAtLeastArmV82 && hasFp16 && hasDotProd) {
+        Log.d(NAME, "Loading librnllama_v8_2_fp16_dotprod.so");
+        System.loadLibrary("rnllama_v8_2_fp16_dotprod");
+      } else if (isAtLeastArmV82 && hasFp16) {
+        Log.d(NAME, "Loading librnllama_v8_2_fp16.so");
+        System.loadLibrary("rnllama_v8_2_fp16");
       } else {
-        Log.d(NAME, "Loading librnllama.so");
-        System.loadLibrary("rnllama");
+        Log.d(NAME, "Loading librnllama_v8.so");
+        System.loadLibrary("rnllama_v8");
       }
     } else if (LlamaContext.isX86_64()) {
-      Log.d(NAME, "Loading librnllama.so");
+      Log.d(NAME, "Loading librnllama_x86_64.so");
+      System.loadLibrary("rnllama_x86_64");
+    } else {
+      Log.d(NAME, "Loading default librnllama.so");
       System.loadLibrary("rnllama");
     }
   }
@@ -271,20 +275,23 @@ private static boolean isX86_64() {
     return Build.SUPPORTED_ABIS[0].equals("x86_64");
   }
 
-  private static String cpuInfo() {
+  private static String getCpuFeatures() {
     File file = new File("/proc/cpuinfo");
     StringBuilder stringBuilder = new StringBuilder();
     try {
       BufferedReader bufferedReader = new BufferedReader(new FileReader(file));
       String line;
       while ((line = bufferedReader.readLine()) != null) {
+        if (line.startsWith("Features")) {
           stringBuilder.append(line);
+          break;
+        }
       }
       bufferedReader.close();
       return stringBuilder.toString();
     } catch (IOException e) {
       Log.w(NAME, "Couldn't read /proc/cpuinfo", e);
-      return null;
+      return "";
     }
   }