add GGML_NUMA_MIRROR macro and modify tensor def

Author: wkgcass

ggml/CMakeLists.txt

@@ -198,6 +198,8 @@ option(GGML_OPENCL_USE_ADRENO_KERNELS       "ggml: use optimized kernels for Adr
 set   (GGML_OPENCL_TARGET_VERSION "300" CACHE STRING
                                             "gmml: OpenCL API version to target")
 
+option(GGML_NUMA_MIRROR                     "ggml: support numa aware tensor data"            OFF)
+
 # toolchain for vulkan-shaders-gen
 set   (GGML_VULKAN_SHADERS_GEN_TOOLCHAIN "" CACHE FILEPATH "ggml: toolchain file for vulkan-shaders-gen")
 
@@ -317,6 +319,18 @@ set(variable_set_statements
 
 set(GGML_SHARED_LIB ${BUILD_SHARED_LIBS})
 
+if (GGML_NUMA_MIRROR)
+    message(STATUS
+            "-----------------\n"
+            "Enabling GGML_NUMA_MIRROR"
+    message(STATUS
+            "-----------------")
+
+    foreach(lib "ggml" "ggml-base")
+        target_compile_definitions(${lib} PUBLIC GGML_NUMA_MIRROR)
+    endforeach()
+endif()
+
 get_cmake_property(all_variables VARIABLES)
 foreach(variable_name IN LISTS all_variables)
     if(variable_name MATCHES "^GGML_")

ggml/include/ggml.h

@@ -598,21 +598,46 @@ extern "C" {
         struct ggml_tensor * view_src;
         size_t               view_offs;
 
+#ifdef GGML_NUMA_MIRROR
+        union {
+        #ifdef __NVCC__
+            void * data;
+        #endif
+            void * __data[2];
+        };
+#else
         void * data;
+#endif
 
         char name[GGML_MAX_NAME];
 
         void * extra; // extra things e.g. for ggml-cuda.cu
 
+#ifndef GGML_NUMA_MIRROR
         char padding[8];
+#endif
     };
 
+#ifdef GGML_NUMA_MIRROR
+    extern __thread int ggml_current_numa_node;
+#endif
+
     static inline void * tensor_data(const struct ggml_tensor * tensor) {
+#ifdef GGML_NUMA_MIRROR
+        int n = ggml_current_numa_node;
+        return tensor->__data[n];
+#else
         return tensor->data;
+#endif
     }
 
     static inline void tensor_set_data(struct ggml_tensor * tensor, void * data) {
+#ifdef GGML_NUMA_MIRROR
+        tensor->__data[0] = data;
+        tensor->__data[1] = data;
+#else
         tensor->data = data;
+#endif
     }
 
     static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor);

ggml/src/ggml.c

@@ -60,6 +60,10 @@
 #define m512i(p) (__m512i)(p)
 #endif
 
+#ifdef GGML_NUMA_MIRROR
+__thread int ggml_current_numa_node = 0;
+#endif
+
 // precomputed f32 table for f16 (256 KB) (ggml-impl.h)
 float ggml_table_f32_f16[1 << 16];
 
@@ -1609,10 +1613,16 @@ static struct ggml_tensor * ggml_new_tensor_impl(
         /*.src          =*/ { NULL },
         /*.view_src     =*/ view_src,
         /*.view_offs    =*/ view_offs,
+#ifdef GGML_NUMA_MIRROR
+        /*.data         =*/ { .__data = { NULL, NULL } },
+#else
         /*.data         =*/ NULL,
+#endif
         /*.name         =*/ { 0 },
         /*.extra        =*/ NULL,
+#ifndef GGML_NUMA_MIRROR
         /*.padding      =*/ { 0 },
+#endif
     };
     tensor_set_data(result, obj_alloc_size > 0 ? (void *)(result + 1) : data);