add memkind as hbm allocator

Author: jikunshang

CMakeLists.txt

@@ -502,6 +502,11 @@ endif()
 # Build libraries
 #
 
+if (GGML_USE_HBM)
+    add_definitions(-DGGML_USE_HBM)
+    find_library(memkind memkind REQUIRED)
+endif()
+
 add_library(ggml OBJECT
             ggml.c
             ggml.h
@@ -517,6 +522,9 @@ add_library(ggml OBJECT
 target_include_directories(ggml PUBLIC . ${LLAMA_EXTRA_INCLUDES})
 target_compile_features(ggml PUBLIC c_std_11) # don't bump
 target_link_libraries(ggml PUBLIC Threads::Threads ${LLAMA_EXTRA_LIBS})
+if (GGML_USE_HBM)
+    target_link_libraries(ggml PUBLIC memkind)
+endif()
 
 add_library(ggml_static STATIC $<TARGET_OBJECTS:ggml>)
 if (BUILD_SHARED_LIBS)

ggml.c

@@ -103,6 +103,9 @@ typedef void * thread_ret_t;
 #include <sys/stat.h>
 #include <unistd.h>
 
+#endif
+#ifdef GGML_USE_HBM
+#include <hbwmalloc.h>
 #endif
 
 // __FMA__ and __F16C__ are not defined in MSVC, however they are implied with AVX2/AVX512
@@ -197,7 +200,9 @@ typedef void * thread_ret_t;
 #else
 inline static void * ggml_aligned_malloc(size_t size) {
     void * aligned_memory = NULL;
-#ifdef GGML_USE_METAL
+#ifdef GGML_USE_HBM
+    int result = hbw_posix_memalign(&aligned_memory, 16, size);
+#elif GGML_USE_METAL
     int result = posix_memalign(&aligned_memory, getpagesize(), size);
 #else
     int result = posix_memalign(&aligned_memory, GGML_MEM_ALIGN, size);
@@ -220,8 +225,12 @@ inline static void * ggml_aligned_malloc(size_t size) {
     return aligned_memory;
 }
 #define GGML_ALIGNED_MALLOC(size)  ggml_aligned_malloc(size)
+#ifdef GGML_USE_HBM
+#define GGML_ALIGNED_FREE(ptr)     hbw_free(ptr)
+#else
 #define GGML_ALIGNED_FREE(ptr)     free(ptr)
 #endif
+#endif
 
 #define UNUSED GGML_UNUSED
 #define SWAP(x, y, T) do { T SWAP = x; x = y; y = SWAP; } while (0)

llama.cpp

@@ -71,6 +71,9 @@ static void llama_log_callback_default(llama_log_level level, const char * text,
 #define LLAMA_MAX_SCRATCH_BUFFERS 16
 #endif
 
+#ifdef GGML_USE_HBM
+#include <hbwmalloc.h>
+#endif
 
 // available llama models
 enum e_model {
@@ -789,7 +792,11 @@ struct llama_model_loader {
             // allocate temp buffer if not using mmap
             if (!use_mmap && lt.data == NULL) {
                 GGML_ASSERT(lt.ggml_tensor->backend != GGML_BACKEND_CPU);
+                #ifdef GGML_USE_HBM
+                lt.data = (uint8_t*)hbw_malloc(ggml_nbytes(lt.ggml_tensor));
+                #else
                 lt.data = (uint8_t*)malloc(ggml_nbytes(lt.ggml_tensor));
+                #endif
             }
 
             load_data_for(lt);