Skip to content

Commit 67c0883

Browse files
glookkaglookka
authored
Improved HNSW performance; added AVX-512 distance functions
1 parent 975b294 commit 67c0883

13 files changed

Lines changed: 1039 additions & 159 deletions

‎.github/workflows/build_template.yml‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ on:
4949
artifact_list:
5050
required: false
5151
type: string
52-
default: "build/columnar/lib_manticore_columnar_avx2.so build/columnar/lib_manticore_columnar.so build/secondary/lib_manticore_secondary_avx2.so build/secondary/lib_manticore_secondary.so build/knn/lib_manticore_knn.so build/knn/lib_manticore_knn_avx2.so build/_deps/manticore-build/src/searchd build/_deps/manticore-build/src/indexer build/_deps/manticore-build/src/indextool build/_deps/manticore-build/config/*.c build/_deps/manticore-build/config/*.h"
52+
default: "build/columnar/lib_manticore_columnar_avx2.so build/columnar/lib_manticore_columnar_avx512.so build/columnar/lib_manticore_columnar.so build/secondary/lib_manticore_secondary_avx2.so build/secondary/lib_manticore_secondary_avx512.so build/secondary/lib_manticore_secondary.so build/knn/lib_manticore_knn.so build/knn/lib_manticore_knn_avx2.so build/knn/lib_manticore_knn_avx512.so build/_deps/manticore-build/src/searchd build/_deps/manticore-build/src/indexer build/_deps/manticore-build/src/indextool build/_deps/manticore-build/config/*.c build/_deps/manticore-build/config/*.h"
5353
HOMEBREW_PREFIX:
5454
required: false
5555
type: string
@@ -159,22 +159,44 @@ jobs:
159159
- name: 🚀 Build
160160
run: ${{ inputs.cmake_command }}
161161

162-
- name: Move AVX2-optimized libraries
162+
- name: Move AVX2/AVX-512-optimized libraries
163163
if: success() && inputs.skip_avx2_move != true
164164
run: |
165165
if [ -d build/avx2/share/manticore/modules ]; then
166-
echo "Move linux based libs"
166+
echo "Move linux based AVX2 libs"
167167
mv build/avx2/share/manticore/modules/lib_manticore_columnar_avx2.so build/columnar/
168168
mv build/avx2/share/manticore/modules/lib_manticore_secondary_avx2.so build/secondary/
169+
mv build/avx2/share/manticore/modules/lib_manticore_knn_avx2.so build/knn/
169170
ls build/columnar/
170171
ls build/secondary/
172+
ls build/knn/
173+
fi
174+
if [ -d build/avx512/share/manticore/modules ]; then
175+
echo "Move linux based AVX-512 libs"
176+
mv build/avx512/share/manticore/modules/lib_manticore_columnar_avx512.so build/columnar/
177+
mv build/avx512/share/manticore/modules/lib_manticore_secondary_avx512.so build/secondary/
178+
mv build/avx512/share/manticore/modules/lib_manticore_knn_avx512.so build/knn/
179+
ls build/columnar/
180+
ls build/secondary/
181+
ls build/knn/
171182
fi
172183
if [ -d build/avx2/share/modules ]; then
173-
echo "Move windows based libs"
184+
echo "Move windows based AVX2 libs"
174185
mv build/avx2/share/modules/lib_manticore_columnar_avx2.dll build/columnar/Debug/
175186
mv build/avx2/share/modules/lib_manticore_secondary_avx2.dll build/secondary/Debug/
187+
mv build/avx2/share/modules/lib_manticore_knn_avx2.dll build/knn/Debug/
188+
ls build/columnar/Debug/
189+
ls build/secondary/Debug/
190+
ls build/knn/Debug/
191+
fi
192+
if [ -d build/avx512/share/modules ]; then
193+
echo "Move windows based AVX-512 libs"
194+
mv build/avx512/share/modules/lib_manticore_columnar_avx512.dll build/columnar/Debug/
195+
mv build/avx512/share/modules/lib_manticore_secondary_avx512.dll build/secondary/Debug/
196+
mv build/avx512/share/modules/lib_manticore_knn_avx512.dll build/knn/Debug/
176197
ls build/columnar/Debug/
177198
ls build/secondary/Debug/
199+
ls build/knn/Debug/
178200
fi
179201
180202
- name: Set artifact name

‎.github/workflows/test.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ jobs:
148148
CTEST_CMAKE_GENERATOR: "Ninja Multi-Config"
149149
CTEST_CONFIGURATION_TYPE: Debug
150150
cache_key: build_windows_x64
151-
artifact_list: "build/columnar/Debug/lib_manticore_columnar.dll build/secondary/Debug/lib_manticore_secondary.dll build/knn/Debug/lib_manticore_knn.dll build/columnar/Debug/lib_manticore_columnar_avx2.dll build/secondary/Debug/lib_manticore_secondary_avx2.dll build/knn/Debug/lib_manticore_knn_avx2.dll build/_deps/manticore-build/src/Debug/indexer.exe build/_deps/manticore-build/src/Debug/searchd.exe" #
151+
artifact_list: "build/columnar/Debug/lib_manticore_columnar.dll build/secondary/Debug/lib_manticore_secondary.dll build/knn/Debug/lib_manticore_knn.dll build/columnar/Debug/lib_manticore_columnar_avx2.dll build/secondary/Debug/lib_manticore_secondary_avx2.dll build/knn/Debug/lib_manticore_knn_avx2.dll build/columnar/Debug/lib_manticore_columnar_avx512.dll build/secondary/Debug/lib_manticore_secondary_avx512.dll build/knn/Debug/lib_manticore_knn_avx512.dll build/_deps/manticore-build/src/Debug/indexer.exe build/_deps/manticore-build/src/Debug/searchd.exe" #
152152
MANTICORE_LOCATOR: ${{ needs.check_branch.outputs.manticore_locator }}
153153
embeddings_artifact: embeddings_windows_x86_64
154154

‎.github/workflows/test_template.yml‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,10 @@ on:
3737
required: false
3838
type: boolean
3939
default: true
40+
USE_AVX512:
41+
required: false
42+
type: boolean
43+
default: false
4044

4145
jobs:
4246
test:
@@ -139,9 +143,20 @@ jobs:
139143
if: ${{ inputs.USE_AVX2 == true }}
140144
run: |
141145
if [ -d /__w/columnar/columnar/build/columnar ]; then
142-
echo "Move linux based libs"
146+
echo "Move linux based AVX2 libs"
143147
mv /__w/columnar/columnar/build/columnar/lib_manticore_columnar_avx2.so /__w/columnar/columnar/build/columnar/lib_manticore_columnar.so
144148
mv /__w/columnar/columnar/build/secondary/lib_manticore_secondary_avx2.so /__w/columnar/columnar/build/secondary/lib_manticore_secondary.so
149+
mv /__w/columnar/columnar/build/knn/lib_manticore_knn_avx2.so /__w/columnar/columnar/build/knn/lib_manticore_knn.so
150+
fi
151+
152+
- name: Replace standard libs to AVX-512
153+
if: ${{ inputs.USE_AVX512 == true }}
154+
run: |
155+
if [ -d /__w/columnar/columnar/build/columnar ]; then
156+
echo "Move linux based AVX-512 libs"
157+
mv /__w/columnar/columnar/build/columnar/lib_manticore_columnar_avx512.so /__w/columnar/columnar/build/columnar/lib_manticore_columnar.so
158+
mv /__w/columnar/columnar/build/secondary/lib_manticore_secondary_avx512.so /__w/columnar/columnar/build/secondary/lib_manticore_secondary.so
159+
mv /__w/columnar/columnar/build/knn/lib_manticore_knn_avx512.so /__w/columnar/columnar/build/knn/lib_manticore_knn.so
145160
fi
146161
147162
- name: 🚀 Test

‎.github/workflows/win_test_template.yml‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,10 @@ on:
2222
required: false
2323
type: boolean
2424
default: true
25+
USE_AVX512:
26+
required: false
27+
type: boolean
28+
default: false
2529

2630
jobs:
2731
test_windows:
@@ -160,9 +164,22 @@ jobs:
160164
if: ${{ inputs.USE_AVX2 == true }}
161165
run: |
162166
if (Test-Path -Path "D:\a\columnar\columnar\build\columnar\Debug\") {
163-
Write-Host "Moving Windows-based libs"
167+
Write-Host "Moving Windows-based AVX2 libs"
164168
Move-Item -Path "D:\a\columnar\columnar\build\columnar\Debug\lib_manticore_columnar_avx2.dll" -Destination "D:\a\columnar\columnar\build\columnar\Debug\lib_manticore_columnar.dll" -Force
165169
Move-Item -Path "D:\a\columnar\columnar\build\secondary\Debug\lib_manticore_secondary_avx2.dll" -Destination "D:\a\columnar\columnar\build\secondary\Debug\lib_manticore_secondary.dll" -Force
170+
Move-Item -Path "D:\a\columnar\columnar\build\knn\Debug\lib_manticore_knn_avx2.dll" -Destination "D:\a\columnar\columnar\build\knn\Debug\lib_manticore_knn.dll" -Force
171+
} else {
172+
Write-Host "Directory D:\a\columnar\columnar\build\columnar\Debug\ does not exist"
173+
}
174+
175+
- name: Replace standard libs to AVX-512
176+
if: ${{ inputs.USE_AVX512 == true }}
177+
run: |
178+
if (Test-Path -Path "D:\a\columnar\columnar\build\columnar\Debug\") {
179+
Write-Host "Moving Windows-based AVX-512 libs"
180+
Move-Item -Path "D:\a\columnar\columnar\build\columnar\Debug\lib_manticore_columnar_avx512.dll" -Destination "D:\a\columnar\columnar\build\columnar\Debug\lib_manticore_columnar.dll" -Force
181+
Move-Item -Path "D:\a\columnar\columnar\build\secondary\Debug\lib_manticore_secondary_avx512.dll" -Destination "D:\a\columnar\columnar\build\secondary\Debug\lib_manticore_secondary.dll" -Force
182+
Move-Item -Path "D:\a\columnar\columnar\build\knn\Debug\lib_manticore_knn_avx512.dll" -Destination "D:\a\columnar\columnar\build\knn\Debug\lib_manticore_knn.dll" -Force
166183
} else {
167184
Write-Host "Directory D:\a\columnar\columnar\build\columnar\Debug\ does not exist"
168185
}

‎CMakeLists.txt‎

Lines changed: 41 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ if (EXTERNAL_PACK)
4444
set ( BUILD_TESTING OFF )
4545
endif ()
4646

47+
option ( USE_AVX512 "Use AVX-512 instruction set" OFF )
48+
4749
set (EXTERNAL_OPTION "EXTERNAL_LIB")
4850

4951
# check what we're build. api just packs headers and build nothing
@@ -210,7 +212,10 @@ if ( FLAGS )
210212
target_compile_options ( columnar_root INTERFACE ${FLAGS} )
211213
endif()
212214

213-
if ( USE_AVX2 )
215+
if ( USE_AVX512 )
216+
target_compile_definitions ( columnar_root INTERFACE USE_AVX512 )
217+
set ( lib_arch_suffix "_avx512" )
218+
elseif ( USE_AVX2 )
214219
target_compile_definitions ( columnar_root INTERFACE USE_AVX2 )
215220
set ( lib_arch_suffix "_avx2" )
216221
endif ()
@@ -242,17 +247,27 @@ if ((INSTALL_COLUMNAR OR INSTALL_SECONDARY OR INSTALL_KNN ) AND AMD AND NOT USE_
242247
endif()
243248

244249
if (ADD_AVX_BUILDS)
245-
message ( STATUS "Will build extra libraries with AVX2 support" )
250+
message ( STATUS "Will build extra libraries with AVX2 and AVX-512 support" )
246251
include ( ExternalProject )
247-
set ( EXTERNAL_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/avx2" ) # here build artefacts located
252+
set ( EXTERNAL_PREFIX_AVX2 "${CMAKE_CURRENT_BINARY_DIR}/avx2" ) # here build artefacts located
248253
ExternalProject_Add ( columnar_avx2
249254
SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
250255
PREFIX "${CMAKE_CURRENT_BINARY_DIR}/avx2"
251-
CMAKE_ARGS -D${EXTERNAL_OPTION}=1 -DCMAKE_INSTALL_PREFIX=${EXTERNAL_PREFIX}
256+
CMAKE_ARGS -D${EXTERNAL_OPTION}=1 -DUSE_AVX2=ON -DCMAKE_INSTALL_PREFIX=${EXTERNAL_PREFIX_AVX2}
252257
CONFIGURE_HANDLED_BY_BUILD TRUE
253258
TEST_COMMAND ""
254259
)
255-
set ( EXTERNAL_MODULES "${EXTERNAL_PREFIX}/share/${REL_MODULES_DIR}" )
260+
set ( EXTERNAL_MODULES_AVX2 "${EXTERNAL_PREFIX_AVX2}/share/${REL_MODULES_DIR}" )
261+
262+
set ( EXTERNAL_PREFIX_AVX512 "${CMAKE_CURRENT_BINARY_DIR}/avx512" )
263+
ExternalProject_Add ( columnar_avx512
264+
SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
265+
PREFIX "${CMAKE_CURRENT_BINARY_DIR}/avx512"
266+
CMAKE_ARGS -D${EXTERNAL_OPTION}=1 -DUSE_AVX512=ON -DCMAKE_INSTALL_PREFIX=${EXTERNAL_PREFIX_AVX512}
267+
CONFIGURE_HANDLED_BY_BUILD TRUE
268+
TEST_COMMAND ""
269+
)
270+
set ( EXTERNAL_MODULES_AVX512 "${EXTERNAL_PREFIX_AVX512}/share/${REL_MODULES_DIR}" )
256271
endif ()
257272

258273
if (INSTALL_COLUMNAR)
@@ -261,9 +276,14 @@ if (INSTALL_COLUMNAR)
261276
install ( FILES "$<TARGET_FILE_DIR:columnar_lib>/lib_manticore_columnar${lib_arch_suffix}.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
262277
if (ADD_AVX_BUILDS)
263278
add_library ( columnar_lib_avx2 MODULE IMPORTED )
264-
set_target_properties ( columnar_lib_avx2 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES}/lib_manticore_columnar_avx2.${MODULE_EXT}" )
279+
set_target_properties ( columnar_lib_avx2 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES_AVX2}/lib_manticore_columnar_avx2.${MODULE_EXT}" )
265280
install ( IMPORTED_RUNTIME_ARTIFACTS columnar_lib_avx2 RUNTIME DESTINATION ${MODULES_DIR} LIBRARY DESTINATION ${MODULES_DIR} COMPONENT columnar )
266-
install ( FILES "${EXTERNAL_MODULES}/lib_manticore_columnar_avx2.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
281+
install ( FILES "${EXTERNAL_MODULES_AVX2}/lib_manticore_columnar_avx2.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
282+
283+
add_library ( columnar_lib_avx512 MODULE IMPORTED )
284+
set_target_properties ( columnar_lib_avx512 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES_AVX512}/lib_manticore_columnar_avx512.${MODULE_EXT}" )
285+
install ( IMPORTED_RUNTIME_ARTIFACTS columnar_lib_avx512 RUNTIME DESTINATION ${MODULES_DIR} LIBRARY DESTINATION ${MODULES_DIR} COMPONENT columnar )
286+
install ( FILES "${EXTERNAL_MODULES_AVX512}/lib_manticore_columnar_avx512.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
267287
endif ()
268288
endif()
269289

@@ -273,9 +293,14 @@ if (INSTALL_SECONDARY)
273293
install ( FILES "$<TARGET_FILE_DIR:secondary_index>/lib_manticore_secondary${lib_arch_suffix}.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
274294
if (ADD_AVX_BUILDS)
275295
add_library ( secondary_index_avx2 MODULE IMPORTED )
276-
set_target_properties ( secondary_index_avx2 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES}/lib_manticore_secondary_avx2.${MODULE_EXT}" )
296+
set_target_properties ( secondary_index_avx2 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES_AVX2}/lib_manticore_secondary_avx2.${MODULE_EXT}" )
277297
install ( IMPORTED_RUNTIME_ARTIFACTS secondary_index_avx2 RUNTIME DESTINATION ${MODULES_DIR} LIBRARY DESTINATION ${MODULES_DIR} COMPONENT secondary )
278-
install ( FILES "${EXTERNAL_MODULES}/lib_manticore_secondary_avx2.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
298+
install ( FILES "${EXTERNAL_MODULES_AVX2}/lib_manticore_secondary_avx2.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
299+
300+
add_library ( secondary_index_avx512 MODULE IMPORTED )
301+
set_target_properties ( secondary_index_avx512 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES_AVX512}/lib_manticore_secondary_avx512.${MODULE_EXT}" )
302+
install ( IMPORTED_RUNTIME_ARTIFACTS secondary_index_avx512 RUNTIME DESTINATION ${MODULES_DIR} LIBRARY DESTINATION ${MODULES_DIR} COMPONENT secondary )
303+
install ( FILES "${EXTERNAL_MODULES_AVX512}/lib_manticore_secondary_avx512.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
279304
endif ()
280305
endif ()
281306

@@ -286,9 +311,14 @@ if (INSTALL_KNN)
286311

287312
if (ADD_AVX_BUILDS)
288313
add_library ( knn_lib_avx2 MODULE IMPORTED )
289-
set_target_properties ( knn_lib_avx2 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES}/lib_manticore_knn_avx2.${MODULE_EXT}" )
314+
set_target_properties ( knn_lib_avx2 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES_AVX2}/lib_manticore_knn_avx2.${MODULE_EXT}" )
290315
install ( IMPORTED_RUNTIME_ARTIFACTS knn_lib_avx2 RUNTIME DESTINATION ${MODULES_DIR} LIBRARY DESTINATION ${MODULES_DIR} COMPONENT knn )
291-
install ( FILES "${EXTERNAL_MODULES}/lib_manticore_knn_avx2.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
316+
install ( FILES "${EXTERNAL_MODULES_AVX2}/lib_manticore_knn_avx2.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
317+
318+
add_library ( knn_lib_avx512 MODULE IMPORTED )
319+
set_target_properties ( knn_lib_avx512 PROPERTIES IMPORTED_LOCATION "${EXTERNAL_MODULES_AVX512}/lib_manticore_knn_avx512.${MODULE_EXT}" )
320+
install ( IMPORTED_RUNTIME_ARTIFACTS knn_lib_avx512 RUNTIME DESTINATION ${MODULES_DIR} LIBRARY DESTINATION ${MODULES_DIR} COMPONENT knn )
321+
install ( FILES "${EXTERNAL_MODULES_AVX512}/lib_manticore_knn_avx512.pdb" DESTINATION ${MODULES_DIR} COMPONENT dbgsymbols OPTIONAL )
292322
endif ()
293323

294324
# Add component dependency - KNN depends on embeddings

‎cmake/GetHNSW.cmake‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1414
# See the License for the specific language governing permissions and
1515
# limitations under the License.
16-
set ( HNSW_GITHUB "https://github.com/manticoresoftware/hnswlib/archive/4f8bea5.zip" )
16+
set ( HNSW_GITHUB "https://github.com/manticoresoftware/hnswlib/archive/6568d3b.zip" )
1717
set ( HNSW_BUNDLEZIP "${LIBS_BUNDLE}/hnswlib-0.7.0.tar.gz" )
1818

1919
cmake_minimum_required ( VERSION 3.17 FATAL_ERROR )

‎cmake/update_bundle.cmake‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,16 @@ function ( get_avx_flags RESULT )
122122
return()
123123
endif()
124124

125-
if (USE_AVX2)
125+
if (USE_AVX512)
126+
message ( STATUS "Add AVX512 flags to compiler flags for ${CMAKE_SYSTEM_PROCESSOR} arch" )
127+
if (CLANG_CL)
128+
set ( ${RESULT} "/arch:AVX512;-mavx512vpopcntdq" PARENT_SCOPE )
129+
elseif (MSVC)
130+
set ( ${RESULT} "/arch:AVX512" PARENT_SCOPE )
131+
else()
132+
set ( ${RESULT} "-march=x86-64-v4;-mavx512vpopcntdq" PARENT_SCOPE )
133+
endif()
134+
elseif (USE_AVX2)
126135
message ( STATUS "Add AVX2 flags to compiler flags for ${CMAKE_SYSTEM_PROCESSOR} arch" )
127136
if (MSVC OR CLANG_CL)
128137
set ( ${RESULT} "/arch:AVX2" PARENT_SCOPE )
@@ -142,8 +151,9 @@ endfunction ( )
142151
function ( external_build module MODULE_SRC_NAME MODULE_BUILD_NAME )
143152
get_avx_flags ( FLAGS )
144153
if (FLAGS)
145-
set ( ENV{CXXFLAGS} ${FLAGS} )
146-
set ( ENV{CFLAGS} ${FLAGS} )
154+
string ( REPLACE ";" " " FLAGS_STR "${FLAGS}" )
155+
set ( ENV{CXXFLAGS} "${FLAGS_STR}" )
156+
set ( ENV{CFLAGS} "${FLAGS_STR}" )
147157
endif ()
148158

149159
set ( CMAKE_ARGS "" )

0 commit comments

Comments
 (0)