AMX enable for FP32 innerproduct

nmslib · Jan 21, 2025 · 7e6fb30 · 7e6fb30
1 parent b2122ac
commit 7e6fb30
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -72,7 +72,7 @@ if(AMXEnable)
     if(HAS_AMX_SUPPORT)
         message(STATUS "Compiler supports AMX")
         set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS_BACKUP} -mamx-bf16")
-        #add_compile_definitions(USE_AMX=1)
+        add_compile_definitions(USE_AMX=1)
     else()
         message(STATUS "Compiler does NOT support AMX")
         set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS_BACKUP}")

diff --git a/hnswlib/space_ip.h b/hnswlib/space_ip.h
@@ -729,6 +729,9 @@ static float InnerProductDistanceBf16(const void* a, const void* b, const void *
 }
 class Bf16InnerProductSpace : public hnswlib::SpaceInterface<float> {
     DISTFUNC<float> fstdistfunc_;
+#ifdef USE_AMX
+    AMXDISTFUNC<float> amxdistfunc_;
+#endif
     size_t data_size_;
     size_t dim_;
  public: