Add CPU backend check and i2 dequantize helper for GGML

include/bitnet-lut-kernels.h

@@ -1124,7 +1124,7 @@ void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT
 }
 
 void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {
-    if (!(is_type_supported(tensor->type) && tensor->extra == nullptr)) {
+    if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {
         return;
     }
 

src/CMakeLists.txt

@@ -1,8 +1,7 @@
-set(GGML_HEADERS_BITNET ../include/ggml-bitnet.h)
-set(GGML_SOURCES_BITNET ggml-bitnet-mad.cpp)
-set(GGML_SOURCES_BITNET ggml-bitnet-lut.cpp)
+set(GGML_HEADERS_BITNET ../include/ggml-bitnet.h ../3rdparty/llama.cpp/ggml/src/ggml/src/ggml-cpu/ggml-cpu-impl.h)
+set(GGML_SOURCES_BITNET ggml-bitnet-mad.cpp ggml-bitnet-lut.cpp)
 
-include_directories(3rdparty/llama.cpp/ggml/include)
+include_directories(3rdparty/llama.cpp/ggml/include 3rdparty/llama.cpp/ggml/src)
 
 if (NOT (CMAKE_C_COMPILER_ID MATCHES "Clang" OR CMAKE_C_COMPILER_ID STREQUAL "GNU") OR
     NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang" OR CMAKE_CXX_COMPILER_ID STREQUAL "GNU"))

src/ggml-bitnet-lut.cpp

@@ -8,10 +8,10 @@
 #ifdef __x86_64__
 #include <immintrin.h>
 #endif
+#include "../include/ggml-bitnet.h"
+#include "../3rdparty/llama.cpp/ggml/src/ggml-quants.h"
 
-#include "ggml-bitnet.h"
-#include "ggml-quants.h"
-#include "ggml-cpu-impl.h"
+#include "../3rdparty/llama.cpp/ggml/src/ggml-cpu/ggml-cpu-impl.h"
 
 #if defined(GGML_BITNET_ARM_TL1) || defined(GGML_BITNET_X86_TL2)
 #include "bitnet-lut-kernels.h"

src/ggml-bitnet-mad.cpp

@@ -1,10 +1,13 @@
+#include "../include/gemm-config.h"
 #include <vector>
 #include <type_traits>
 #include <assert.h>
-#include "ggml-bitnet.h"
-#include "ggml-quants.h"
 #include "gemm-config.h"
-#include "ggml-cpu-impl.h"
+#include "../include/ggml-bitnet.h"
+#include "../3rdparty/llama.cpp/ggml/src/ggml-quants.h"
+
+#include "../3rdparty/llama.cpp/ggml/src/ggml-cpu/ggml-cpu-impl.h"
+
 #include <cmath>
 #include <cstring>
 
@@ -195,6 +198,29 @@ size_t quantize_i2_s(const float * src, void * dst, int64_t nrow, int64_t n_per_
 #endif
 }
 
+void dequantize_row_i2_s(const uint8_t * x, float * y, int64_t n, const float i2_scale) {
+    static const float map2bit[4] = { -1.0f, 0.0f, 1.0f, 0.0f };
+    int64_t done = 0;
+    while (done < n) {
+        int64_t cols0 = MIN(32, n - done - 0*32);
+        int64_t cols1 = MIN(32, n - done - 1*32);
+        int64_t cols2 = MIN(32, n - done - 2*32);
+        int64_t cols3 = MIN(32, n - done - 3*32);
+        for (int gp = 0; gp < 32; gp++) {
+            uint8_t byte = x[(done/4) + gp];
+            uint8_t c0 = (byte >> 6) & 0x03;
+            uint8_t c1 = (byte >> 4) & 0x03;
+            uint8_t c2 = (byte >> 2) & 0x03;
+            uint8_t c3 = (byte >> 0) & 0x03;
+            if (gp < cols0) y[done + 0*32 + gp] = i2_scale * map2bit[c0];
+            if (gp < cols1) y[done + 1*32 + gp] = i2_scale * map2bit[c1];
+            if (gp < cols2) y[done + 2*32 + gp] = i2_scale * map2bit[c2];
+            if (gp < cols3) y[done + 3*32 + gp] = i2_scale * map2bit[c3];
+        }
+        done += 128;
+    }
+}
+
 void ggml_vec_dot_i2_i8_s_1x1(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) {
 #if defined(__AVX2__)
     const uint8_t *    x = (uint8_t *)vx;