From be3ef883e104e6e259d681b996bf142f5921c12f Mon Sep 17 00:00:00 2001 From: wehub-resource-sync Date: Mon, 13 Jul 2026 12:47:42 +0800 Subject: [PATCH] chore: import upstream snapshot with attribution --- .clang-format | 31 + .clang-tidy | 10 + .gitattributes | 3 + .github/CODEOWNERS | 38 + .github/ISSUE_TEMPLATE/benchmark.yml | 49 + .github/ISSUE_TEMPLATE/bug_report.yml | 73 + .github/ISSUE_TEMPLATE/config.yml | 9 + .github/ISSUE_TEMPLATE/enhancement.yml | 41 + .github/ISSUE_TEMPLATE/feature_request.yml | 55 + .github/ISSUE_TEMPLATE/integration.yml | 36 + .github/ISSUE_TEMPLATE/profiling.yml | 41 + .github/auto-assign-config.yml | 24 + .../subproject-integration/CMakeLists.txt | 26 + .github/codecov.yml | 32 + .github/dependabot.yml | 17 + .github/workflows/01-ci-pipeline.yml | 116 + .github/workflows/02-lint-check.yml | 52 + .github/workflows/03-macos-linux-build.yml | 185 + .github/workflows/04-android-build.yml | 437 + .github/workflows/05-windows-build.yml | 157 + .github/workflows/06-ios-build.yml | 163 + .github/workflows/07-linux-riscv-build.yml | 342 + .../08-cmake-subproject-integration.yml | 102 + .github/workflows/_build_wheel_job.yml | 104 + .github/workflows/auto-assign.yml | 25 + .github/workflows/clang_tidy.yml | 197 + .github/workflows/community-pr-handler.yml | 148 + .github/workflows/continuous_bench.yml | 33 + .../docker/Dockerfile.linux_x64_glibc228 | 83 + .github/workflows/issue-auto-assign.yml | 99 + .github/workflows/nightly_coverage.yml | 102 + .github/workflows/scripts/run_vdb.sh | 88 + .gitignore | 68 + .gitmodules | 61 + CMakeLists.txt | 211 + CODE_OF_CONDUCT.md | 41 + CONTRIBUTING.md | 116 + LICENSE | 201 + NOTICE | 96 + README.md | 137 + README.wehub.md | 7 + README_CN.md | 133 + cmake/bazel.cmake | 2521 ++ cmake/iOSBundleInfo.plist.in | 22 + cmake/option.cmake | 321 + cmake/utils.cmake | 47 + examples/c++/CMakeLists.txt | 128 + examples/c++/ailego/main.cc | 11 + examples/c++/core/external_vector_example.cc | 219 + examples/c++/core/main.cc | 88 + examples/c++/db/main.cc | 264 + examples/c/CMakeLists.txt | 137 + examples/c/basic_example.c | 267 + examples/c/collection_schema_example.c | 255 + examples/c/diskann_example.c | 346 + examples/c/doc_example.c | 539 + examples/c/field_schema_example.c | 312 + examples/c/index_example.c | 398 + examples/c/optimized_example.c | 310 + pyproject.toml | 299 + python/tests/detail/distance_helper.py | 392 + python/tests/detail/doc_helper.py | 464 + python/tests/detail/fixture_helper.py | 652 + python/tests/detail/params_helper.py | 209 + python/tests/detail/support_helper.py | 126 + .../detail/test_collection_concurrency.py | 429 + .../detail/test_collection_create_and_open.py | 791 + python/tests/detail/test_collection_ddl.py | 1720 ++ python/tests/detail/test_collection_dml.py | 1111 + python/tests/detail/test_collection_dql.py | 1434 ++ .../tests/detail/test_collection_exception.py | 328 + python/tests/detail/test_collection_open.py | 967 + python/tests/detail/test_collection_recall.py | 740 + python/tests/detail/test_db_config.py | 307 + python/tests/test_collection.py | 1218 + python/tests/test_collection_diskann.py | 600 + python/tests/test_collection_fts.py | 188 + .../test_collection_fts_vector_hybrid.py | 391 + python/tests/test_collection_hnsw_rabitq.py | 574 + python/tests/test_convert.py | 584 + python/tests/test_doc.py | 269 + python/tests/test_embedding.py | 2099 ++ python/tests/test_fts_query.py | 158 + python/tests/test_gil_release.py | 273 + python/tests/test_hnsw_contiguous_memory.py | 415 + python/tests/test_jieba_default_dict.py | 201 + python/tests/test_params.py | 540 + python/tests/test_query_executor.py | 329 + python/tests/test_reranker.py | 948 + python/tests/test_schema.py | 246 + python/tests/test_typing.py | 141 + python/tests/test_util.py | 89 + python/tests/test_vamana.py | 501 + python/zvec/__init__.py | 238 + python/zvec/__init__.pyi | 202 + python/zvec/common/__init__.py | 18 + python/zvec/common/constants.py | 33 + python/zvec/executor/__init__.py | 24 + python/zvec/executor/query_executor.py | 266 + python/zvec/extension/__init__.py | 58 + .../zvec/extension/bm25_embedding_function.py | 375 + python/zvec/extension/embedding_function.py | 147 + .../zvec/extension/http_embedding_function.py | 162 + .../zvec/extension/jina_embedding_function.py | 240 + python/zvec/extension/jina_function.py | 182 + .../zvec/extension/multi_vector_reranker.py | 197 + .../extension/openai_embedding_function.py | 238 + python/zvec/extension/openai_function.py | 149 + .../zvec/extension/qwen_embedding_function.py | 537 + python/zvec/extension/qwen_function.py | 186 + python/zvec/extension/qwen_rerank_function.py | 177 + python/zvec/extension/rerank_function.py | 56 + ...sentence_transformer_embedding_function.py | 839 + .../sentence_transformer_function.py | 150 + .../sentence_transformer_rerank_function.py | 396 + python/zvec/model/__init__.py | 30 + python/zvec/model/collection.py | 439 + python/zvec/model/convert.py | 54 + python/zvec/model/doc.py | 178 + python/zvec/model/param/__init__.py | 60 + python/zvec/model/param/__init__.pyi | 1206 + python/zvec/model/param/query.py | 143 + python/zvec/model/schema/__init__.py | 21 + python/zvec/model/schema/__init__.pyi | 109 + python/zvec/model/schema/collection_schema.py | 215 + python/zvec/model/schema/field_schema.py | 310 + python/zvec/py.typed | 0 python/zvec/tool/__init__.py | 18 + python/zvec/tool/util.py | 63 + python/zvec/typing/__init__.py | 32 + python/zvec/typing/__init__.pyi | 404 + python/zvec/typing/enum.py | 62 + python/zvec/zvec.py | 246 + scripts/README.md | 0 scripts/build_android.sh | 395 + scripts/build_ios.sh | 178 + scripts/gcov.sh | 41 + .../generate_standard_tokenizer_unicode.py | 177 + src/CMakeLists.txt | 171 + src/ailego/CMakeLists.txt | 125 + src/ailego/algorithm/binary_quantizer.cc | 77 + src/ailego/algorithm/binary_quantizer.h | 69 + src/ailego/algorithm/integer_quantizer.cc | 466 + src/ailego/algorithm/integer_quantizer.h | 238 + src/ailego/algorithm/kmeans.h | 1055 + src/ailego/algorithm/lloyd_cluster.h | 364 + src/ailego/buffer/block_eviction_queue.cc | 139 + src/ailego/buffer/vector_page_table.cc | 621 + src/ailego/container/bitmap.cc | 443 + src/ailego/container/bitmap.h | 504 + src/ailego/container/bloom_filter.h | 279 + src/ailego/container/params.cc | 139 + src/ailego/container/reservoir.h | 140 + src/ailego/container/vector_array.h | 556 + src/ailego/encoding/json/mod_json.c | 3591 +++ src/ailego/hash/crc32c.cc | 637 + src/ailego/internal/cpu_features.cc | 427 + src/ailego/internal/cpu_features.h | 370 + src/ailego/io/file.cc | 780 + src/ailego/io/file_lock.cc | 80 + src/ailego/io/file_lock.h | 84 + src/ailego/io/file_writer.h | 127 + src/ailego/logger/logger.cc | 76 + src/ailego/math/cosine_distance_matrix.h | 90 + src/ailego/math/distance.h | 254 + src/ailego/math/distance_matrix.h | 20 + src/ailego/math/distance_matrix_accum_fp16.i | 879 + src/ailego/math/distance_matrix_accum_fp32.i | 1406 + src/ailego/math/distance_matrix_accum_int4.i | 1084 + src/ailego/math/distance_matrix_accum_int8.i | 1084 + .../math/distance_matrix_euclidean_utility.i | 253 + src/ailego/math/distance_matrix_fp16.i | 861 + src/ailego/math/distance_matrix_fp32.i | 1599 ++ .../distance_matrix_inner_product_utility.i | 208 + src/ailego/math/distance_matrix_int32.i | 859 + src/ailego/math/distance_matrix_int64.i | 472 + .../math/distance_matrix_mips_utility.i | 160 + src/ailego/math/distance_matrix_popcnt.i | 2565 ++ src/ailego/math/distance_utility.h | 65 + src/ailego/math/euclidean_distance_matrix.h | 649 + .../euclidean_distance_matrix_fp16_avx.cc | 36 + .../euclidean_distance_matrix_fp16_avx512.cc | 33 + ...clidean_distance_matrix_fp16_avx512fp16.cc | 82 + ...euclidean_distance_matrix_fp16_dispatch.cc | 85 + .../euclidean_distance_matrix_fp16_neon.cc | 34 + .../euclidean_distance_matrix_fp32_avx.cc | 106 + .../euclidean_distance_matrix_fp32_avx512.cc | 100 + ...euclidean_distance_matrix_fp32_dispatch.cc | 90 + .../euclidean_distance_matrix_fp32_neon.cc | 62 + .../euclidean_distance_matrix_fp32_sse.cc | 83 + .../euclidean_distance_matrix_int4_avx2.cc | 130 + ...euclidean_distance_matrix_int4_dispatch.cc | 65 + .../euclidean_distance_matrix_int4_sse.cc | 102 + .../euclidean_distance_matrix_int8_avx2.cc | 192 + ...euclidean_distance_matrix_int8_dispatch.cc | 65 + .../euclidean_distance_matrix_int8_sse.cc | 170 + .../math/euclidean_distance_matrix_scalar.cc | 114 + src/ailego/math/fht.h | 38 + src/ailego/math/fht_avx2.cc | 125 + src/ailego/math/fht_avx512.cc | 130 + src/ailego/math/fht_dispatch.cc | 150 + src/ailego/math/fht_neon.cc | 107 + src/ailego/math/fht_scalar.cc | 80 + src/ailego/math/fht_sse.cc | 106 + src/ailego/math/inner_product_matrix.h | 995 + src/ailego/math/inner_product_matrix_fp16.cc | 0 .../math/inner_product_matrix_fp16_avx.cc | 718 + .../math/inner_product_matrix_fp16_avx512.cc | 43 + .../inner_product_matrix_fp16_avx512fp16.cc | 757 + .../inner_product_matrix_fp16_dispatch.cc | 179 + .../math/inner_product_matrix_fp16_neon.cc | 43 + src/ailego/math/inner_product_matrix_fp32.cc | 0 .../math/inner_product_matrix_fp32_avx.cc | 114 + .../math/inner_product_matrix_fp32_avx512.cc | 102 + .../inner_product_matrix_fp32_dispatch.cc | 155 + .../math/inner_product_matrix_fp32_neon.cc | 65 + .../math/inner_product_matrix_fp32_sse.cc | 323 + .../math/inner_product_matrix_int4_avx2.cc | 138 + .../inner_product_matrix_int4_dispatch.cc | 80 + .../math/inner_product_matrix_int4_sse.cc | 107 + .../math/inner_product_matrix_int8_avx2.cc | 204 + .../inner_product_matrix_int8_dispatch.cc | 82 + .../math/inner_product_matrix_int8_sse.cc | 166 + .../math/inner_product_matrix_scalar.cc | 299 + src/ailego/math/matrix_define.i | 1850 ++ src/ailego/math/matrix_utility.i | 265 + .../math/mips_euclidean_distance_matrix.h | 1115 + ...mips_euclidean_distance_matrix_fp16_avx.cc | 149 + ...s_euclidean_distance_matrix_fp16_avx512.cc | 168 + ...euclidean_distance_matrix_fp16_dispatch.cc | 103 + ...ips_euclidean_distance_matrix_fp16_neon.cc | 159 + .../mips_euclidean_distance_matrix_fp32.cc | 0 ...mips_euclidean_distance_matrix_fp32_avx.cc | 159 + ...s_euclidean_distance_matrix_fp32_avx512.cc | 155 + ...euclidean_distance_matrix_fp32_dispatch.cc | 176 + ...ips_euclidean_distance_matrix_fp32_neon.cc | 77 + ...mips_euclidean_distance_matrix_fp32_sse.cc | 368 + ...ips_euclidean_distance_matrix_int4_avx2.cc | 172 + ...euclidean_distance_matrix_int4_dispatch.cc | 90 + ...mips_euclidean_distance_matrix_int4_sse.cc | 137 + ...ips_euclidean_distance_matrix_int8_avx2.cc | 191 + ...euclidean_distance_matrix_int8_dispatch.cc | 87 + ...mips_euclidean_distance_matrix_int8_sse.cc | 170 + .../mips_euclidean_distance_matrix_scalar.cc | 172 + src/ailego/math/norm1_matrix.h | 147 + src/ailego/math/norm1_matrix_fp16.cc | 90 + src/ailego/math/norm1_matrix_fp32.cc | 84 + src/ailego/math/norm2_matrix.h | 424 + src/ailego/math/norm2_matrix_fp16.cc | 91 + src/ailego/math/norm2_matrix_fp32.cc | 93 + src/ailego/math/norm_matrix.h | 18 + src/ailego/math/norm_matrix_fp16.i | 219 + src/ailego/math/norm_matrix_fp32.i | 185 + src/ailego/math/normalizer.cc | 437 + src/ailego/math/normalizer.h | 114 + src/ailego/math_batch/cosine_distance_batch.h | 58 + src/ailego/math_batch/distance_batch.h | 65 + src/ailego/math_batch/distance_batch_math.h | 30 + .../math_batch/euclidean_distance_batch.h | 143 + .../euclidean_distance_batch_dispatch.cc | 191 + ...euclidean_distance_batch_impl_fp16_avx2.cc | 114 + ...clidean_distance_batch_impl_fp16_avx512.cc | 137 + ...ean_distance_batch_impl_fp16_avx512fp16.cc | 99 + ...euclidean_distance_batch_impl_fp32_avx2.cc | 104 + ...clidean_distance_batch_impl_fp32_avx512.cc | 95 + .../math_batch/inner_product_distance_batch.h | 143 + .../inner_product_distance_batch_dispatch.cc | 228 + ...r_product_distance_batch_impl_fp16_avx2.cc | 109 + ...product_distance_batch_impl_fp16_avx512.cc | 128 + ...uct_distance_batch_impl_fp16_avx512fp16.cc | 92 + ...r_product_distance_batch_impl_fp32_avx2.cc | 134 + ...r_product_distance_batch_impl_int8_avx2.cc | 104 + ...uct_distance_batch_impl_int8_avx512fp16.cc | 173 + src/ailego/parallel/lock.h | 294 + src/ailego/parallel/multi_thread_list.h | 102 + src/ailego/parallel/semaphore.h | 181 + src/ailego/parallel/thread_pool.cc | 133 + src/ailego/pattern/defer.h | 24 + src/ailego/pattern/scope_guard.h | 177 + src/ailego/utility/bit_string_helper.h | 118 + src/ailego/utility/bitset_helper.cc | 2049 ++ src/ailego/utility/bitset_helper.h | 174 + src/ailego/utility/concurrency_helper.cc | 70 + src/ailego/utility/concurrency_helper.h | 34 + src/ailego/utility/dl_helper.cc | 74 + src/ailego/utility/dl_helper.h | 47 + src/ailego/utility/file_helper.cc | 404 + src/ailego/utility/float_helper.cc | 1246 + src/ailego/utility/math_helper.h | 109 + src/ailego/utility/matrix_helper.h | 70 + src/ailego/utility/memory_helper.cc | 492 + src/ailego/utility/memory_helper.h | 113 + src/ailego/utility/string_helper.cc | 121 + src/ailego/utility/time_helper.cc | 215 + src/ailego/version.cc | 48 + src/ailego/version.h | 33 + src/ailego/version.i | 293 + src/binding/CMakeLists.txt | 14 + src/binding/c/CMakeLists.txt | 256 + src/binding/c/c_api.cc | 7223 ++++++ src/binding/python/CMakeLists.txt | 133 + src/binding/python/binding.cc | 72 + src/binding/python/exports.mac | 1 + src/binding/python/exports.map | 34 + .../python/include/python_collection.h | 36 + src/binding/python/include/python_config.h | 30 + src/binding/python/include/python_doc.h | 34 + src/binding/python/include/python_param.h | 37 + src/binding/python/include/python_reranker.h | 31 + src/binding/python/include/python_schema.h | 35 + src/binding/python/include/python_type.h | 38 + .../python/model/common/python_config.cc | 223 + .../python/model/param/python_param.cc | 2269 ++ src/binding/python/model/python_collection.cc | 316 + src/binding/python/model/python_doc.cc | 452 + src/binding/python/model/python_reranker.cc | 91 + .../python/model/schema/python_schema.cc | 162 + src/binding/python/typing/python_type.cc | 229 + src/core/CMakeLists.txt | 87 + src/core/algorithm/CMakeLists.txt | 67 + src/core/algorithm/cluster/CMakeLists.txt | 19 + src/core/algorithm/cluster/cluster_params.h | 213 + src/core/algorithm/cluster/kmeans_cluster.cc | 917 + src/core/algorithm/cluster/linear_seeker.cc | 45 + src/core/algorithm/cluster/linear_seeker.h | 98 + .../algorithm/cluster/multi_chunk_cluster.cc | 440 + .../algorithm/cluster/multi_chunk_cluster.h | 482 + .../algorithm/cluster/opt_kmeans_cluster.cc | 1194 + src/core/algorithm/cluster/seeker.h | 72 + .../algorithm/cluster/stratified_cluster.cc | 502 + .../cluster/stratified_cluster_trainer.cc | 341 + .../cluster/stratified_cluster_trainer.h | 84 + src/core/algorithm/cluster/vector_mean.h | 791 + src/core/algorithm/diskann/CMakeLists.txt | 57 + .../algorithm/diskann/diskann_algorithm.cc | 324 + .../algorithm/diskann/diskann_algorithm.h | 72 + src/core/algorithm/diskann/diskann_builder.cc | 638 + src/core/algorithm/diskann/diskann_builder.h | 131 + .../diskann/diskann_builder_entity.cc | 643 + .../diskann/diskann_builder_entity.h | 107 + src/core/algorithm/diskann/diskann_context.cc | 142 + src/core/algorithm/diskann/diskann_context.h | 378 + .../diskann/diskann_dist_calculator.h | 150 + src/core/algorithm/diskann/diskann_entity.cc | 32 + src/core/algorithm/diskann/diskann_entity.h | 244 + .../algorithm/diskann/diskann_file_reader.cc | 299 + .../algorithm/diskann/diskann_file_reader.h | 103 + src/core/algorithm/diskann/diskann_holder.h | 325 + .../diskann/diskann_index_provider.h | 120 + src/core/algorithm/diskann/diskann_indexer.cc | 1205 + src/core/algorithm/diskann/diskann_indexer.h | 109 + src/core/algorithm/diskann/diskann_params.h | 53 + .../algorithm/diskann/diskann_pq_table.cc | 117 + src/core/algorithm/diskann/diskann_pq_table.h | 154 + .../algorithm/diskann/diskann_pq_trainer.cc | 417 + .../algorithm/diskann/diskann_pq_trainer.h | 87 + src/core/algorithm/diskann/diskann_reducer.cc | 190 + src/core/algorithm/diskann/diskann_reducer.h | 85 + .../diskann/diskann_reducer_entity.cc | 215 + .../diskann/diskann_reducer_entity.h | 69 + .../algorithm/diskann/diskann_searcher.cc | 312 + src/core/algorithm/diskann/diskann_searcher.h | 166 + .../diskann/diskann_searcher_entity.cc | 441 + .../diskann/diskann_searcher_entity.h | 125 + .../algorithm/diskann/diskann_streamer.cc | 345 + src/core/algorithm/diskann/diskann_streamer.h | 180 + src/core/algorithm/diskann/diskann_util.h | 221 + .../algorithm/diskann/diskann_vecs_reader.h | 320 + .../algorithm/diskann/diskann_visit_filter.h | 16 + src/core/algorithm/flat/CMakeLists.txt | 18 + src/core/algorithm/flat/flat_builder.cc | 283 + src/core/algorithm/flat/flat_builder.h | 102 + .../algorithm/flat/flat_distance_matrix.h | 184 + src/core/algorithm/flat/flat_index_format.h | 187 + src/core/algorithm/flat/flat_searcher.cc | 252 + src/core/algorithm/flat/flat_searcher.h | 182 + .../algorithm/flat/flat_searcher_context.h | 1114 + .../algorithm/flat/flat_searcher_provider.h | 220 + src/core/algorithm/flat/flat_streamer.cc | 491 + src/core/algorithm/flat/flat_streamer.h | 190 + .../algorithm/flat/flat_streamer_context.h | 259 + .../algorithm/flat/flat_streamer_dumper.h | 244 + .../algorithm/flat/flat_streamer_entity.cc | 1004 + .../algorithm/flat/flat_streamer_entity.h | 407 + .../algorithm/flat/flat_streamer_provider.h | 100 + src/core/algorithm/flat/flat_utility.h | 190 + src/core/algorithm/flat_sparse/CMakeLists.txt | 20 + .../flat_sparse/flat_sparse_builder.cc | 338 + .../flat_sparse/flat_sparse_builder.h | 105 + .../flat_sparse/flat_sparse_context.cc | 30 + .../flat_sparse/flat_sparse_context.h | 225 + .../flat_sparse/flat_sparse_entity.h | 246 + .../flat_sparse/flat_sparse_index_format.h | 45 + .../flat_sparse/flat_sparse_provider.h | 193 + .../flat_sparse/flat_sparse_search.h | 174 + .../flat_sparse/flat_sparse_searcher.cc | 182 + .../flat_sparse/flat_sparse_searcher.h | 185 + .../flat_sparse_searcher_entity.cc | 241 + .../flat_sparse/flat_sparse_searcher_entity.h | 126 + .../flat_sparse/flat_sparse_streamer.cc | 411 + .../flat_sparse/flat_sparse_streamer.h | 172 + .../flat_sparse_streamer_entity.cc | 1083 + .../flat_sparse/flat_sparse_streamer_entity.h | 265 + .../flat_sparse/flat_sparse_utility.h | 58 + src/core/algorithm/hnsw/CMakeLists.txt | 17 + src/core/algorithm/hnsw/hnsw_algorithm.cc | 699 + src/core/algorithm/hnsw/hnsw_algorithm.h | 156 + src/core/algorithm/hnsw/hnsw_chunk.cc | 219 + src/core/algorithm/hnsw/hnsw_chunk.h | 153 + src/core/algorithm/hnsw/hnsw_context.cc | 309 + src/core/algorithm/hnsw/hnsw_context.h | 596 + .../algorithm/hnsw/hnsw_dist_calculator.h | 247 + src/core/algorithm/hnsw/hnsw_entity.cc | 364 + src/core/algorithm/hnsw/hnsw_entity.h | 696 + src/core/algorithm/hnsw/hnsw_index_hash.h | 226 + src/core/algorithm/hnsw/hnsw_index_provider.h | 133 + src/core/algorithm/hnsw/hnsw_params.h | 122 + src/core/algorithm/hnsw/hnsw_streamer.cc | 927 + src/core/algorithm/hnsw/hnsw_streamer.h | 244 + .../algorithm/hnsw/hnsw_streamer_entity.cc | 1061 + .../algorithm/hnsw/hnsw_streamer_entity.h | 1195 + src/core/algorithm/hnsw_rabitq/CMakeLists.txt | 27 + .../hnsw_rabitq/hnsw_rabitq_algorithm.cc | 382 + .../hnsw_rabitq/hnsw_rabitq_algorithm.h | 124 + .../hnsw_rabitq/hnsw_rabitq_chunk.cc | 221 + .../algorithm/hnsw_rabitq/hnsw_rabitq_chunk.h | 140 + .../hnsw_rabitq/hnsw_rabitq_context.cc | 297 + .../hnsw_rabitq/hnsw_rabitq_context.h | 548 + .../hnsw_rabitq_dist_calculator.cc | 39 + .../hnsw_rabitq/hnsw_rabitq_dist_calculator.h | 239 + .../hnsw_rabitq/hnsw_rabitq_entity.cc | 382 + .../hnsw_rabitq/hnsw_rabitq_entity.h | 614 + .../hnsw_rabitq/hnsw_rabitq_index_hash.h | 236 + .../hnsw_rabitq/hnsw_rabitq_index_provider.h | 134 + .../hnsw_rabitq/hnsw_rabitq_params.h | 121 + .../hnsw_rabitq_query_algorithm.cc | 354 + .../hnsw_rabitq/hnsw_rabitq_query_algorithm.h | 136 + .../hnsw_rabitq/hnsw_rabitq_query_entity.h | 28 + .../hnsw_rabitq/hnsw_rabitq_register.cc | 24 + .../hnsw_rabitq/hnsw_rabitq_streamer.cc | 995 + .../hnsw_rabitq/hnsw_rabitq_streamer.h | 245 + .../hnsw_rabitq_streamer_entity.cc | 711 + .../hnsw_rabitq/hnsw_rabitq_streamer_entity.h | 559 + .../algorithm/hnsw_rabitq/rabitq_converter.cc | 315 + .../algorithm/hnsw_rabitq/rabitq_converter.h | 101 + .../algorithm/hnsw_rabitq/rabitq_params.h | 49 + .../algorithm/hnsw_rabitq/rabitq_reformer.cc | 471 + .../algorithm/hnsw_rabitq/rabitq_reformer.h | 73 + .../algorithm/hnsw_rabitq/rabitq_utils.cc | 115 + src/core/algorithm/hnsw_rabitq/rabitq_utils.h | 53 + src/core/algorithm/hnsw_sparse/CMakeLists.txt | 17 + .../hnsw_sparse/hnsw_sparse_algorithm.cc | 524 + .../hnsw_sparse/hnsw_sparse_algorithm.h | 127 + .../hnsw_sparse/hnsw_sparse_builder.cc | 528 + .../hnsw_sparse/hnsw_sparse_builder.h | 103 + .../hnsw_sparse/hnsw_sparse_builder_entity.cc | 291 + .../hnsw_sparse/hnsw_sparse_builder_entity.h | 159 + .../hnsw_sparse/hnsw_sparse_chunk.cc | 217 + .../algorithm/hnsw_sparse/hnsw_sparse_chunk.h | 142 + .../hnsw_sparse/hnsw_sparse_context.cc | 294 + .../hnsw_sparse/hnsw_sparse_context.h | 538 + .../hnsw_sparse/hnsw_sparse_dist_calculator.h | 184 + .../hnsw_sparse/hnsw_sparse_entity.cc | 443 + .../hnsw_sparse/hnsw_sparse_entity.h | 625 + .../hnsw_sparse/hnsw_sparse_index_hash.h | 232 + .../hnsw_sparse/hnsw_sparse_index_provider.h | 168 + .../hnsw_sparse/hnsw_sparse_params.h | 118 + .../hnsw_sparse/hnsw_sparse_searcher.cc | 564 + .../hnsw_sparse/hnsw_sparse_searcher.h | 159 + .../hnsw_sparse_searcher_entity.cc | 630 + .../hnsw_sparse/hnsw_sparse_searcher_entity.h | 181 + .../hnsw_sparse/hnsw_sparse_streamer.cc | 960 + .../hnsw_sparse/hnsw_sparse_streamer.h | 226 + .../hnsw_sparse_streamer_entity.cc | 1001 + .../hnsw_sparse/hnsw_sparse_streamer_entity.h | 540 + src/core/algorithm/ivf/CMakeLists.txt | 16 + src/core/algorithm/ivf/ivf_builder.cc | 829 + src/core/algorithm/ivf/ivf_builder.h | 315 + src/core/algorithm/ivf/ivf_centroid_index.cc | 617 + src/core/algorithm/ivf/ivf_centroid_index.h | 247 + .../algorithm/ivf/ivf_distance_calculator.cc | 46 + .../algorithm/ivf/ivf_distance_calculator.h | 185 + src/core/algorithm/ivf/ivf_dumper.cc | 409 + src/core/algorithm/ivf/ivf_dumper.h | 260 + src/core/algorithm/ivf/ivf_entity.cc | 989 + src/core/algorithm/ivf/ivf_entity.h | 367 + src/core/algorithm/ivf/ivf_index_format.h | 174 + src/core/algorithm/ivf/ivf_index_provider.h | 162 + src/core/algorithm/ivf/ivf_params.h | 90 + src/core/algorithm/ivf/ivf_searcher.cc | 324 + src/core/algorithm/ivf/ivf_searcher.h | 112 + src/core/algorithm/ivf/ivf_searcher_context.h | 239 + src/core/algorithm/ivf/ivf_streamer.cc | 324 + src/core/algorithm/ivf/ivf_streamer.h | 123 + src/core/algorithm/ivf/ivf_utility.h | 177 + src/core/algorithm/vamana/CMakeLists.txt | 17 + src/core/algorithm/vamana/vamana_algorithm.cc | 620 + src/core/algorithm/vamana/vamana_algorithm.h | 116 + src/core/algorithm/vamana/vamana_context.cc | 176 + src/core/algorithm/vamana/vamana_context.h | 350 + .../algorithm/vamana/vamana_dist_calculator.h | 174 + src/core/algorithm/vamana/vamana_entity.cc | 203 + src/core/algorithm/vamana/vamana_entity.h | 366 + .../algorithm/vamana/vamana_index_provider.h | 117 + src/core/algorithm/vamana/vamana_params.h | 81 + src/core/algorithm/vamana/vamana_streamer.cc | 715 + src/core/algorithm/vamana/vamana_streamer.h | 183 + .../vamana/vamana_streamer_entity.cc | 926 + .../algorithm/vamana/vamana_streamer_entity.h | 713 + src/core/framework/CMakeLists.txt | 11 + src/core/framework/index_cluster.cc | 159 + src/core/framework/index_context.cc | 28 + src/core/framework/index_converter.cc | 43 + src/core/framework/index_error.cc | 73 + src/core/framework/index_factory.cc | 261 + src/core/framework/index_flow.cc | 882 + src/core/framework/index_helper.cc | 255 + src/core/framework/index_logger.cc | 72 + src/core/framework/index_mapping.cc | 617 + src/core/framework/index_meta.cc | 208 + src/core/framework/index_plugin.cc | 59 + src/core/framework/index_version.cc | 34 + src/core/interface/CMakeLists.txt | 10 + src/core/interface/index.cc | 1069 + src/core/interface/index_factory.cc | 415 + src/core/interface/index_param.cc | 343 + src/core/interface/indexes/diskann_index.cc | 323 + src/core/interface/indexes/flat_index.cc | 72 + src/core/interface/indexes/hnsw_index.cc | 185 + .../interface/indexes/hnsw_rabitq_index.cc | 145 + src/core/interface/indexes/holder_builder.h | 87 + src/core/interface/indexes/ivf_index.cc | 273 + src/core/interface/indexes/vamana_index.cc | 116 + src/core/interface/utils/utils.h | 104 + src/core/metric/CMakeLists.txt | 17 + src/core/metric/cosine_metric.cc | 235 + src/core/metric/euclidean_metric.cc | 839 + src/core/metric/inner_product_metric.cc | 534 + src/core/metric/metric_params.h | 42 + src/core/metric/mips_euclidean_metric.cc | 683 + src/core/metric/quantized_integer_metric.cc | 372 + .../metric/quantized_integer_metric_batch.h | 367 + .../metric/quantized_integer_metric_matrix.h | 254 + src/core/metric/uniform_int8_metric.cc | 158 + src/core/mixed_reducer/CMakeLists.txt | 16 + src/core/mixed_reducer/mixed_reducer_params.h | 42 + .../mixed_reducer/mixed_streamer_reducer.cc | 590 + .../mixed_reducer/mixed_streamer_reducer.h | 114 + src/core/plugin/CMakeLists.txt | 29 + src/core/plugin/diskann_plugin.cc | 322 + src/core/quantizer/CMakeLists.txt | 17 + src/core/quantizer/binary_converter.cc | 231 + src/core/quantizer/binary_reformer.cc | 108 + src/core/quantizer/cosine_converter.cc | 452 + src/core/quantizer/cosine_reformer.cc | 322 + src/core/quantizer/half_float_converter.cc | 370 + src/core/quantizer/half_float_reformer.cc | 238 + .../quantizer/integer_quantizer_converter.cc | 659 + .../quantizer/integer_quantizer_reformer.cc | 542 + src/core/quantizer/mips_converter.cc | 660 + src/core/quantizer/mips_reformer.cc | 463 + src/core/quantizer/quantizer_params.h | 143 + src/core/quantizer/record_quantizer.h | 136 + src/core/quantizer/rotator/fht_rotator.cc | 188 + src/core/quantizer/rotator/fht_rotator.h | 60 + src/core/quantizer/rotator/matrix_rotator.cc | 182 + src/core/quantizer/rotator/matrix_rotator.h | 52 + src/core/quantizer/rotator/rotator.cc | 335 + src/core/quantizer/rotator/rotator.h | 153 + src/core/quantizer/uniform_int8_converter.cc | 336 + src/core/quantizer/uniform_int8_reformer.cc | 224 + src/core/utility/CMakeLists.txt | 37 + src/core/utility/basic_refiner.cc | 270 + src/core/utility/block_heap.cc | 184 + src/core/utility/block_heap.h | 117 + src/core/utility/buffer_storage.cc | 1553 ++ src/core/utility/file_dumper.cc | 135 + src/core/utility/file_read_storage.cc | 520 + src/core/utility/linear_pool.h | 237 + src/core/utility/memory_dumper.cc | 115 + src/core/utility/memory_read_storage.cc | 249 + src/core/utility/mmap_file_read_storage.cc | 299 + src/core/utility/mmap_file_storage.cc | 354 + src/core/utility/sparse_utility.h | 420 + src/core/utility/topk_result_builder.h | 144 + src/core/utility/utility_params.h | 234 + src/core/utility/visit_filter.h | 428 + src/db/CMakeLists.txt | 57 + src/db/collection.cc | 2148 ++ src/db/common/CMakeLists.txt | 14 + src/db/common/cgroup_util.cc | 494 + src/db/common/cgroup_util.h | 104 + src/db/common/concurrent_roaring_bitmap.cc | 211 + src/db/common/concurrent_roaring_bitmap.h | 323 + src/db/common/config.cc | 168 + src/db/common/constants.h | 90 + src/db/common/error_code.cc | 157 + src/db/common/error_code.h | 210 + src/db/common/file_helper.cc | 87 + src/db/common/file_helper.h | 298 + src/db/common/global_resource.cc | 33 + src/db/common/global_resource.h | 41 + src/db/common/glogger.h | 99 + src/db/common/logger.h | 77 + src/db/common/profiler.h | 230 + src/db/common/rocksdb_context.cc | 580 + src/db/common/rocksdb_context.h | 150 + src/db/common/status.cc | 56 + src/db/common/typedef.h | 198 + src/db/common/utils.cc | 23 + src/db/common/utils.h | 21 + src/db/index/CMakeLists.txt | 37 + src/db/index/column/column_indexer.h | 25 + src/db/index/column/common/index_results.h | 68 + src/db/index/column/fts_column/FtsLexer.g4 | 59 + src/db/index/column/fts_column/FtsParser.g4 | 92 + src/db/index/column/fts_column/bm25_scorer.cc | 196 + src/db/index/column/fts_column/bm25_scorer.h | 217 + .../column/fts_column/fts_ast_rewriter.cc | 410 + .../column/fts_column/fts_ast_rewriter.h | 43 + .../column/fts_column/fts_column_indexer.cc | 924 + .../column/fts_column/fts_column_indexer.h | 238 + .../column/fts_column/fts_index_results.h | 85 + src/db/index/column/fts_column/fts_indexer.cc | 331 + src/db/index/column/fts_column/fts_indexer.h | 93 + src/db/index/column/fts_column/fts_pipeline.h | 37 + .../index/column/fts_column/fts_query_ast.h | 191 + .../column/fts_column/fts_rocksdb_merge.cc | 181 + .../column/fts_column/fts_rocksdb_merge.h | 59 + .../column/fts_column/fts_rocksdb_reducer.cc | 494 + .../column/fts_column/fts_rocksdb_reducer.h | 157 + src/db/index/column/fts_column/fts_types.h | 61 + src/db/index/column/fts_column/fts_utils.cc | 38 + src/db/index/column/fts_column/fts_utils.h | 99 + .../index/column/fts_column/gen/FtsLexer.cc | 257 + src/db/index/column/fts_column/gen/FtsLexer.h | 73 + .../column/fts_column/gen/FtsLexer.interp | 67 + .../column/fts_column/gen/FtsLexer.tokens | 21 + .../index/column/fts_column/gen/FtsParser.cc | 1116 + .../index/column/fts_column/gen/FtsParser.h | 303 + .../column/fts_column/gen/FtsParser.interp | 53 + .../column/fts_column/gen/FtsParser.tokens | 21 + .../fts_column/gen/FtsParserBaseListener.cc | 8 + .../fts_column/gen/FtsParserBaseListener.h | 89 + .../fts_column/gen/FtsParserListener.cc | 8 + .../column/fts_column/gen/FtsParserListener.h | 66 + src/db/index/column/fts_column/gen_parser.sh | 9 + .../iterator/fts_candidate_iterator.cc | 53 + .../iterator/fts_candidate_iterator.h | 55 + .../iterator/fts_conjunction_iterator.cc | 295 + .../iterator/fts_conjunction_iterator.h | 84 + .../iterator/fts_disjunction_iterator.cc | 280 + .../iterator/fts_disjunction_iterator.h | 63 + .../fts_column/iterator/fts_doc_iterator.h | 123 + .../iterator/fts_phrase_iterator.cc | 210 + .../fts_column/iterator/fts_phrase_iterator.h | 83 + .../fts_column/iterator/fts_term_iterator.cc | 206 + .../fts_column/iterator/fts_term_iterator.h | 134 + .../fts_column/parser/fts_query_parser.cc | 415 + .../fts_column/parser/fts_query_parser.h | 67 + .../posting/bitpacked_posting_list.cc | 704 + .../posting/bitpacked_posting_list.h | 237 + .../fts_column/posting/bitpacked_simd_avx2.cc | 216 + .../fts_column/posting/bitpacked_simd_avx2.h | 49 + .../posting/bitpacked_simd_dispatch.cc | 60 + .../posting/bitpacked_simd_dispatch.h | 44 + .../posting/bitpacked_simd_scalar.cc | 115 + .../posting/bitpacked_simd_scalar.h | 48 + .../posting/bitpacked_simd_sse41.cc | 202 + .../fts_column/posting/bitpacked_simd_sse41.h | 50 + .../tokenizer/ascii_folding_token_filter.cc | 291 + .../tokenizer/ascii_folding_token_filter.h | 36 + .../fts_column/tokenizer/jieba_tokenizer.cc | 172 + .../fts_column/tokenizer/jieba_tokenizer.h | 86 + .../tokenizer/standard_tokenizer.cc | 1100 + .../fts_column/tokenizer/standard_tokenizer.h | 51 + .../tokenizer/standard_tokenizer_unicode.inc | 1398 + .../tokenizer/stemmer_token_filter.cc | 80 + .../tokenizer/stemmer_token_filter.h | 46 + .../fts_column/tokenizer/token_filter.cc | 51 + .../fts_column/tokenizer/token_filter.h | 67 + .../column/fts_column/tokenizer/tokenizer.h | 64 + .../fts_column/tokenizer/tokenizer_factory.cc | 117 + .../fts_column/tokenizer/tokenizer_factory.h | 65 + .../tokenizer/tokenizer_pipeline_manager.cc | 124 + .../tokenizer/tokenizer_pipeline_manager.h | 88 + .../tokenizer/whitespace_tokenizer.cc | 56 + .../tokenizer/whitespace_tokenizer.h | 39 + .../column/inverted_column/inverted_codec.h | 382 + .../inverted_column/inverted_column_indexer.h | 435 + .../inverted_column_indexer_search.cc | 1015 + .../inverted_column_indexer_util.cc | 207 + .../inverted_column_indexer_write.cc | 466 + .../inverted_column/inverted_doc_range.h | 127 + .../inverted_column/inverted_indexer.cc | 269 + .../column/inverted_column/inverted_indexer.h | 119 + .../inverted_column/inverted_rocksdb_merger.h | 114 + .../inverted_column/inverted_search_result.h | 199 + .../combined_vector_column_indexer.cc | 450 + .../combined_vector_column_indexer.h | 83 + .../column/vector_column/engine_helper.hpp | 516 + .../vector_column/vector_column_indexer.cc | 217 + .../vector_column/vector_column_indexer.h | 144 + .../vector_column/vector_column_params.h | 111 + .../vector_column/vector_index_results.h | 292 + src/db/index/common/delete_store.h | 152 + src/db/index/common/doc.cc | 1254 + src/db/index/common/id_map.cc | 184 + src/db/index/common/id_map.h | 97 + src/db/index/common/index_filter.h | 79 + src/db/index/common/index_params.cc | 136 + src/db/index/common/meta.h | 445 + src/db/index/common/proto_converter.cc | 439 + src/db/index/common/proto_converter.h | 81 + src/db/index/common/query.cc | 250 + src/db/index/common/schema.cc | 696 + src/db/index/common/stats.cc | 57 + src/db/index/common/type_helper.cc | 98 + src/db/index/common/type_helper.h | 574 + src/db/index/common/version_manager.cc | 304 + src/db/index/common/version_manager.h | 240 + src/db/index/segment/column_merging_reader.cc | 133 + src/db/index/segment/column_merging_reader.h | 50 + src/db/index/segment/segment.cc | 4821 ++++ src/db/index/segment/segment.h | 206 + src/db/index/segment/segment_helper.cc | 1146 + src/db/index/segment/segment_helper.h | 307 + src/db/index/segment/segment_manager.cc | 159 + src/db/index/segment/segment_manager.h | 55 + src/db/index/segment/sql_expr_parser.cc | 352 + src/db/index/segment/sql_expr_parser.h | 30 + src/db/index/storage/arrow_ipc_writer.cc | 196 + src/db/index/storage/arrow_ipc_writer.h | 55 + src/db/index/storage/base_forward_store.h | 59 + .../index/storage/bufferpool_forward_store.cc | 382 + .../index/storage/bufferpool_forward_store.h | 127 + src/db/index/storage/chunked_file_writer.cc | 138 + src/db/index/storage/chunked_file_writer.h | 49 + src/db/index/storage/forward_writer.cc | 32 + src/db/index/storage/forward_writer.h | 45 + .../index/storage/lazy_record_batch_reader.h | 186 + src/db/index/storage/memory_forward_store.cc | 478 + src/db/index/storage/memory_forward_store.h | 193 + src/db/index/storage/mmap_forward_store.cc | 571 + src/db/index/storage/mmap_forward_store.h | 204 + src/db/index/storage/parquet_buffer_pool.cc | 142 + src/db/index/storage/parquet_buffer_pool.h | 149 + src/db/index/storage/parquet_writer.cc | 185 + src/db/index/storage/parquet_writer.h | 55 + src/db/index/storage/store_helper.h | 949 + src/db/index/storage/wal/local_wal_file.cc | 247 + src/db/index/storage/wal/local_wal_file.h | 99 + src/db/index/storage/wal/wal_file.cc | 33 + src/db/index/storage/wal/wal_file.h | 67 + src/db/proto/zvec.proto | 239 + src/db/reranker/reranker.cc | 167 + src/db/sqlengine/CMakeLists.txt | 13 + src/db/sqlengine/analyzer/query_analyzer.cc | 548 + src/db/sqlengine/analyzer/query_analyzer.h | 60 + src/db/sqlengine/analyzer/query_field_info.cc | 37 + src/db/sqlengine/analyzer/query_field_info.h | 103 + src/db/sqlengine/analyzer/query_info.cc | 122 + src/db/sqlengine/analyzer/query_info.h | 387 + .../sqlengine/analyzer/query_info_helper.cc | 193 + src/db/sqlengine/analyzer/query_info_helper.h | 56 + src/db/sqlengine/analyzer/query_node.cc | 331 + src/db/sqlengine/analyzer/query_node.h | 373 + .../sqlengine/analyzer/query_node_walker.cc | 705 + src/db/sqlengine/analyzer/query_node_walker.h | 97 + .../sqlengine/analyzer/query_orderby_info.cc | 31 + .../sqlengine/analyzer/query_orderby_info.h | 62 + src/db/sqlengine/analyzer/simple_rewriter.cc | 198 + src/db/sqlengine/analyzer/simple_rewriter.h | 61 + src/db/sqlengine/antlr/SQLLexer.g4 | 95 + src/db/sqlengine/antlr/SQLParser.g4 | 210 + src/db/sqlengine/antlr/gen/SQLLexer.cc | 694 + src/db/sqlengine/antlr/gen/SQLLexer.h | 116 + src/db/sqlengine/antlr/gen/SQLLexer.interp | 194 + src/db/sqlengine/antlr/gen/SQLLexer.tokens | 103 + src/db/sqlengine/antlr/gen/SQLParser.cc | 3674 +++ src/db/sqlengine/antlr/gen/SQLParser.h | 811 + src/db/sqlengine/antlr/gen/SQLParser.interp | 163 + src/db/sqlengine/antlr/gen/SQLParser.tokens | 103 + .../antlr/gen/SQLParserBaseListener.cc | 8 + .../antlr/gen/SQLParserBaseListener.h | 211 + .../sqlengine/antlr/gen/SQLParserListener.cc | 8 + .../sqlengine/antlr/gen/SQLParserListener.h | 156 + src/db/sqlengine/antlr/gen_parser.sh | 9 + src/db/sqlengine/common/fts_cond_info.h | 43 + src/db/sqlengine/common/generic_node.h | 90 + src/db/sqlengine/common/group_by.h | 46 + src/db/sqlengine/common/util.cc | 93 + src/db/sqlengine/common/util.h | 54 + src/db/sqlengine/parser/base_info.h | 64 + .../parser/case_changing_charstream.h | 88 + .../sqlengine/parser/error_verbose_listener.h | 51 + src/db/sqlengine/parser/node.cc | 249 + src/db/sqlengine/parser/node.h | 285 + src/db/sqlengine/parser/orderby_elem_info.h | 55 + src/db/sqlengine/parser/query_parser.cc | 31 + src/db/sqlengine/parser/query_parser.h | 34 + src/db/sqlengine/parser/select_info.cc | 207 + src/db/sqlengine/parser/select_info.h | 101 + src/db/sqlengine/parser/selected_elem_info.cc | 39 + src/db/sqlengine/parser/selected_elem_info.h | 104 + src/db/sqlengine/parser/sql_info.cc | 85 + src/db/sqlengine/parser/sql_info.h | 65 + src/db/sqlengine/parser/sql_info_helper.cc | 117 + src/db/sqlengine/parser/sql_info_helper.h | 32 + .../parser/zvec_cached_sql_parser.cc | 234 + .../sqlengine/parser/zvec_cached_sql_parser.h | 57 + src/db/sqlengine/parser/zvec_parser.cc | 157 + src/db/sqlengine/parser/zvec_parser.h | 60 + src/db/sqlengine/parser/zvec_sql_parser.cc | 566 + src/db/sqlengine/parser/zvec_sql_parser.h | 54 + src/db/sqlengine/planner/doc_filter.cc | 168 + src/db/sqlengine/planner/doc_filter.h | 74 + src/db/sqlengine/planner/fts_recall_node.cc | 143 + src/db/sqlengine/planner/fts_recall_node.h | 59 + .../sqlengine/planner/invert_recall_node.cc | 87 + src/db/sqlengine/planner/invert_recall_node.h | 64 + src/db/sqlengine/planner/invert_search.cc | 177 + src/db/sqlengine/planner/invert_search.h | 42 + src/db/sqlengine/planner/op_register.cc | 30 + src/db/sqlengine/planner/op_register.h | 26 + .../planner/ops/check_not_filtered_op.cc | 72 + .../planner/ops/check_not_filtered_op.h | 114 + src/db/sqlengine/planner/ops/contain_op.cc | 220 + src/db/sqlengine/planner/ops/contain_op.h | 122 + .../sqlengine/planner/ops/fetch_vector_op.cc | 206 + .../sqlengine/planner/ops/fetch_vector_op.h | 183 + src/db/sqlengine/planner/optimizer.cc | 293 + src/db/sqlengine/planner/optimizer.h | 62 + src/db/sqlengine/planner/plan_info.cc | 33 + src/db/sqlengine/planner/plan_info.h | 45 + src/db/sqlengine/planner/query_planner.cc | 672 + src/db/sqlengine/planner/query_planner.h | 79 + src/db/sqlengine/planner/segment_node.cc | 85 + src/db/sqlengine/planner/segment_node.h | 55 + .../sqlengine/planner/vector_recall_node.cc | 288 + src/db/sqlengine/planner/vector_recall_node.h | 68 + src/db/sqlengine/sqlengine.cc | 24 + src/db/sqlengine/sqlengine.h | 42 + src/db/sqlengine/sqlengine_impl.cc | 747 + src/db/sqlengine/sqlengine_impl.h | 82 + .../zvec/ailego/buffer/block_eviction_queue.h | 153 + .../zvec/ailego/buffer/concurrentqueue.h | 4410 ++++ .../zvec/ailego/buffer/external_cache.h | 213 + .../zvec/ailego/buffer/vector_page_table.h | 301 + src/include/zvec/ailego/container/blob.h | 131 + src/include/zvec/ailego/container/cube.h | 414 + src/include/zvec/ailego/container/heap.h | 237 + src/include/zvec/ailego/container/hypercube.h | 268 + src/include/zvec/ailego/container/params.h | 776 + src/include/zvec/ailego/container/vector.h | 1055 + src/include/zvec/ailego/encoding/json.h | 17 + .../zvec/ailego/encoding/json/mod_json.h | 1386 + .../zvec/ailego/encoding/json/mod_json_plus.h | 3448 +++ src/include/zvec/ailego/hash/crc32c.h | 35 + src/include/zvec/ailego/hash/jump_hash.h | 34 + src/include/zvec/ailego/internal/platform.h | 524 + src/include/zvec/ailego/io/file.h | 310 + src/include/zvec/ailego/io/mmap_file.h | 256 + src/include/zvec/ailego/logger/logger.h | 184 + src/include/zvec/ailego/math_batch/utils.h | 22 + .../zvec/ailego/parallel/thread_pool.h | 402 + .../zvec/ailego/parallel/thread_queue.h | 291 + src/include/zvec/ailego/pattern/closure.h | 530 + src/include/zvec/ailego/pattern/expected.hpp | 2605 ++ src/include/zvec/ailego/pattern/factory.h | 230 + src/include/zvec/ailego/pattern/singleton.h | 51 + .../zvec/ailego/string/string_concat_helper.h | 120 + src/include/zvec/ailego/string/string_view.h | 65 + src/include/zvec/ailego/utility/file_helper.h | 154 + .../zvec/ailego/utility/float_helper.h | 237 + .../zvec/ailego/utility/string_helper.h | 321 + .../zvec/ailego/utility/string_helper_impl.h | 174 + src/include/zvec/ailego/utility/time_helper.h | 206 + src/include/zvec/ailego/utility/type_helper.h | 128 + src/include/zvec/c_api.h | 4143 +++ .../zvec/core/framework/index_builder.h | 61 + .../zvec/core/framework/index_bundle.h | 276 + .../zvec/core/framework/index_cluster.h | 291 + .../zvec/core/framework/index_context.h | 275 + .../zvec/core/framework/index_converter.h | 239 + .../zvec/core/framework/index_document.h | 342 + .../zvec/core/framework/index_dumper.h | 168 + src/include/zvec/core/framework/index_error.h | 170 + .../zvec/core/framework/index_factory.h | 287 + .../zvec/core/framework/index_features.h | 604 + .../zvec/core/framework/index_filter.h | 74 + src/include/zvec/core/framework/index_flow.h | 647 + .../zvec/core/framework/index_format.h | 185 + .../zvec/core/framework/index_framework.h | 32 + .../zvec/core/framework/index_groupby.h | 54 + .../zvec/core/framework/index_helper.h | 60 + .../zvec/core/framework/index_holder.h | 1777 ++ .../zvec/core/framework/index_logger.h | 169 + .../zvec/core/framework/index_mapping.h | 222 + .../zvec/core/framework/index_memory.h | 266 + src/include/zvec/core/framework/index_meta.h | 714 + .../zvec/core/framework/index_metric.h | 156 + .../zvec/core/framework/index_module.h | 67 + .../zvec/core/framework/index_packer.h | 223 + .../zvec/core/framework/index_plugin.h | 117 + .../zvec/core/framework/index_provider.h | 470 + .../zvec/core/framework/index_reducer.h | 229 + .../zvec/core/framework/index_refiner.h | 92 + .../zvec/core/framework/index_reformer.h | 134 + .../zvec/core/framework/index_runner.h | 749 + .../zvec/core/framework/index_searcher.h | 57 + .../core/framework/index_segment_storage.h | 238 + src/include/zvec/core/framework/index_stats.h | 69 + .../zvec/core/framework/index_storage.h | 405 + .../zvec/core/framework/index_streamer.h | 55 + .../zvec/core/framework/index_threads.h | 168 + .../zvec/core/framework/index_trainer.h | 112 + .../zvec/core/framework/index_unpacker.h | 317 + .../zvec/core/framework/index_version.h | 31 + src/include/zvec/core/interface/constants.h | 42 + src/include/zvec/core/interface/index.h | 406 + .../zvec/core/interface/index_factory.h | 54 + src/include/zvec/core/interface/index_param.h | 467 + .../core/interface/index_param_builders.h | 514 + .../zvec/core/interface/vector_source.h | 37 + src/include/zvec/db/collection.h | 122 + src/include/zvec/db/config.h | 203 + src/include/zvec/db/doc.h | 366 + src/include/zvec/db/index_params.h | 821 + src/include/zvec/db/options.h | 69 + src/include/zvec/db/query.h | 206 + src/include/zvec/db/query_params.h | 287 + src/include/zvec/db/reranker.h | 75 + src/include/zvec/db/schema.h | 424 + src/include/zvec/db/stats.h | 35 + src/include/zvec/db/status.h | 181 + src/include/zvec/db/type.h | 146 + src/include/zvec/plugin/diskann_plugin.h | 84 + src/include/zvec/turbo/turbo.h | 72 + src/include/zvec/version.h.in | 16 + src/turbo/CMakeLists.txt | 36 + .../record_quantized_int8/common.h | 315 + .../record_quantized_int8/cosine.cc | 144 + .../record_quantized_int8/cosine.h | 39 + .../squared_euclidean.cc | 173 + .../record_quantized_int8/squared_euclidean.h | 41 + .../avx512_vnni/uniform_int8/quantize.cc | 83 + src/turbo/avx512_vnni/uniform_int8/quantize.h | 33 + .../uniform_int8/squared_euclidean.cc | 222 + .../uniform_int8/squared_euclidean.h | 35 + src/turbo/turbo.cc | 103 + tests/CMakeLists.txt | 12 + tests/ailego/CMakeLists.txt | 15 + .../algorithm/integer_quantizer_test.cc | 564 + tests/ailego/algorithm/kmeans_test.cc | 343 + tests/ailego/container/bitmap_test.cc | 1522 ++ tests/ailego/container/blob_test.cc | 99 + tests/ailego/container/bloom_filter_test.cc | 69 + tests/ailego/container/cube_test.cc | 457 + tests/ailego/container/heap_test.cc | 468 + tests/ailego/container/hypercube_test.cc | 48 + tests/ailego/container/params_test.cc | 391 + tests/ailego/container/reservoir_test.cc | 52 + tests/ailego/container/vector_array_test.cc | 190 + tests/ailego/container/vector_test.cc | 410 + tests/ailego/encoding/json_parse_test.cc | 1796 ++ tests/ailego/hash/crc32c_test.cc | 141 + tests/ailego/hash/jump_hash_test.cc | 48 + tests/ailego/internal/cpu_features_test.cc | 286 + tests/ailego/io/file_lock_test.cc | 43 + tests/ailego/io/file_test.cc | 404 + tests/ailego/io/mmap_file_test.cc | 161 + tests/ailego/io/utf8_path_test.cc | 287 + tests/ailego/logger/logger_test.cc | 67 + .../math/cosine_distance_matrix_fp16_test.cc | 472 + .../math/cosine_distance_matrix_fp32_test.cc | 408 + .../math/cosine_distance_matrix_int8_test.cc | 428 + .../euclidean_distance_matrix_fp16_test.cc | 796 + .../euclidean_distance_matrix_fp32_test.cc | 786 + .../euclidean_distance_matrix_int4_test.cc | 682 + .../euclidean_distance_matrix_int8_test.cc | 960 + .../math/inner_product_matrix_fp16_test.cc | 1658 ++ .../math/inner_product_matrix_fp32_test.cc | 1628 ++ .../math/inner_product_matrix_int4_test.cc | 570 + .../math/inner_product_matrix_int8_test.cc | 730 + ...ips_euclidean_distance_matrix_fp16_test.cc | 988 + ...ips_euclidean_distance_matrix_fp32_test.cc | 962 + ...ips_euclidean_distance_matrix_int4_test.cc | 1000 + ...ips_euclidean_distance_matrix_int8_test.cc | 967 + tests/ailego/math/norm_matrix_fp16_test.cc | 349 + tests/ailego/math/norm_matrix_fp32_test.cc | 345 + tests/ailego/math/norm_matrix_int4_test.cc | 188 + tests/ailego/math/norm_matrix_int8_test.cc | 298 + tests/ailego/math/normalizer_test.cc | 109 + tests/ailego/parallel/lock_test.cc | 151 + .../ailego/parallel/multi_thread_list_test.cc | 259 + tests/ailego/parallel/semaphore_test.cc | 103 + tests/ailego/parallel/thread_pool_test.cc | 162 + tests/ailego/parallel/thread_queue_test.cc | 121 + tests/ailego/pattern/closure_test.cc | 1095 + tests/ailego/pattern/factory_test.cc | 48 + tests/ailego/pattern/scope_guard_test.cc | 102 + tests/ailego/pattern/singleton_test.cc | 45 + .../ailego/utility/bit_string_helper_test.cc | 49 + tests/ailego/utility/bitset_helper_test.cc | 47 + tests/ailego/utility/dl_helper_test.cc | 32 + tests/ailego/utility/float_helper_test.cc | 133 + tests/ailego/utility/matrix_helper_test.cc | 56 + tests/ailego/utility/memory_helper_test.cc | 37 + tests/ailego/utility/string_helper_test.cc | 480 + tests/ailego/utility/time_helper_test.cc | 79 + tests/ailego/utility/type_helper_test.cc | 106 + tests/ailego/version_test.cc | 23 + tests/android_gmock_main.cc | 27 + tests/android_gtest_main.cc | 26 + tests/c/CMakeLists.txt | 42 + tests/c/c_api_test.c | 6472 +++++ tests/c/utils.c | 993 + tests/c/utils.h | 248 + tests/core/CMakeLists.txt | 9 + tests/core/algorithm/CMakeLists.txt | 21 + tests/core/algorithm/cluster/CMakeLists.txt | 14 + .../algorithm/cluster/kmeans_cluster_test.cc | 106 + .../cluster/multi_chunk_cluster_test.cc | 221 + .../cluster/opt_kmeans_cluster_test.cc | 502 + tests/core/algorithm/diskann/CMakeLists.txt | 14 + .../algorithm/diskann/diskann_builder_test.cc | 169 + .../diskann/diskann_searcher_test.cc | 816 + tests/core/algorithm/flat/CMakeLists.txt | 14 + .../core/algorithm/flat/flat_builder_test.cc | 334 + .../algorithm/flat/flat_searcher_test.cpp | 1492 ++ .../flat/flat_streamer_buffer_test.cc | 623 + .../flat/flat_streamer_buffer_time_test.cc | 231 + .../flat/flat_streamer_buffer_time_test.cpp | 0 .../core/algorithm/flat/flat_streamer_test.cc | 1308 + .../core/algorithm/flat_sparse/CMakeLists.txt | 14 + .../flat_sparse/flat_sparse_builder_test.cc | 300 + .../flat_sparse/flat_sparse_searcher_test.cc | 819 + .../flat_sparse_streamer_buffer_test.cc | 230 + .../flat_sparse/flat_sparse_streamer_test.cc | 1726 ++ tests/core/algorithm/hnsw/CMakeLists.txt | 14 + .../hnsw/hnsw_streamer_buffer_test.cc | 517 + .../core/algorithm/hnsw/hnsw_streamer_test.cc | 3886 +++ .../core/algorithm/hnsw_rabitq/CMakeLists.txt | 34 + .../hnsw_rabitq/hnsw_rabitq_streamer_test.cc | 660 + .../core/algorithm/hnsw_sparse/CMakeLists.txt | 14 + .../hnsw_sparse/hnsw_sparse_builder_test.cc | 469 + .../hnsw_sparse/hnsw_sparse_searcher_test.cpp | 1780 ++ .../hnsw_sparse_streamer_buffer_test.cpp | 360 + .../hnsw_sparse/hnsw_sparse_streamer_test.cc | 2490 ++ tests/core/algorithm/ivf/CMakeLists.txt | 14 + tests/core/algorithm/ivf/ivf_builder_test.cc | 528 + tests/core/algorithm/ivf/ivf_searcher_test.cc | 3832 +++ tests/core/algorithm/vamana/CMakeLists.txt | 14 + .../algorithm/vamana/vamana_streamer_test.cc | 892 + tests/core/framework/CMakeLists.txt | 14 + tests/core/interface/CMakeLists.txt | 15 + tests/core/interface/index_group_by_test.cc | 539 + tests/core/interface/index_interface_test.cc | 2430 ++ tests/core/metric/CMakeLists.txt | 14 + tests/core/metric/cosine_metric_test.cc | 606 + tests/core/metric/euclidean_metric_test.cc | 134 + .../core/metric/inner_product_metric_test.cc | 77 + .../metric/quantized_integer_metric_test.cc | 1170 + tests/core/quantizer/CMakeLists.txt | 14 + .../quantizer/half_float_reformer_test.cc | 119 + .../integer_quantizer_reformer_test.cc | 1038 + .../quantizer/uniform_int8_reformer_test.cc | 539 + tests/core/utility/CMakeLists.txt | 16 + tests/core/utility/buffer_storage_test.cpp | 70 + .../core/utility/buffer_storage_write_test.cc | 1180 + tests/core/utility/file_dumper_test.cc | 130 + tests/core/utility/memory_dumper_test.cc | 62 + .../core/utility/mmap_file_container_test.cc | 109 + tests/core/utility/mmap_file_storage_test.cpp | 99 + tests/db/CMakeLists.txt | 49 + tests/db/collection_test.cc | 6547 +++++ tests/db/common/CMakeLists.txt | 32 + tests/db/common/config_test.cc | 246 + tests/db/common/status_test.cc | 160 + tests/db/crash_recovery/CMakeLists.txt | 85 + .../db/crash_recovery/collection_optimizer.cc | 139 + tests/db/crash_recovery/data_generator.cc | 257 + .../crash_recovery/optimize_recovery_test.cc | 264 + tests/db/crash_recovery/utility.h | 205 + .../db/crash_recovery/write_recovery_test.cc | 418 + tests/db/fts_query_test.cc | 230 + tests/db/index/CMakeLists.txt | 64 + .../ascii_folding_token_filter_test.cc | 265 + .../fts_column/bitpacked_posting_list_test.cc | 866 + .../fts_column/fts_ast_rewriter_test.cc | 560 + .../fts_column/fts_candidate_iterator_test.cc | 98 + .../fts_column/fts_column_indexer_test.cc | 1920 ++ .../fts_column/fts_conjunction_opt_test.cc | 305 + .../fts_column/fts_rocksdb_reducer_test.cc | 1137 + .../fts_column/lowercase_token_filter_test.cc | 124 + .../fts_column/standard_tokenizer_test.cc | 485 + .../fts_column/stemmer_token_filter_test.cc | 167 + .../tokenizer_pipeline_manager_test.cc | 271 + ...erted_column_indexer_array_numbers_test.cc | 348 + .../inverted_column_indexer_bool_test.cc | 366 + ...rted_column_indexer_cyclic_numbers_test.cc | 589 + ..._column_indexer_sequential_numbers_test.cc | 777 + .../inverted_column_indexer_string_test.cc | 378 + .../inverted_indexer_util_test.cc | 256 + .../column/vector_column_group_by_test.cc | 378 + .../column/vector_column_indexer_test.cc | 2687 ++ .../index/common/db_proto_converter_test.cc | 550 + tests/db/index/common/db_type_helper_test.cc | 291 + tests/db/index/common/doc_test.cc | 1593 ++ tests/db/index/common/index_params_test.cc | 281 + tests/db/index/common/meta_test.cc | 252 + tests/db/index/common/query_params_test.cc | 80 + tests/db/index/common/schema_test.cc | 1023 + tests/db/index/common/version_manager_test.cc | 303 + .../segment/column_merging_reader_test.cc | 437 + tests/db/index/segment/segment_helper_test.cc | 805 + tests/db/index/segment/segment_row_id_test.cc | 234 + tests/db/index/segment/segment_test.cc | 2298 ++ tests/db/index/segment/segment_test_fixture.h | 105 + .../db/index/segment/sql_expr_parser_test.cc | 290 + .../index/segment/sql_expr_validator_test.cc | 145 + .../db/index/storage/arrow_ipc_writer_test.cc | 130 + .../db/index/storage/bufferpool_store_test.cc | 441 + tests/db/index/storage/mem_store_test.cc | 1158 + tests/db/index/storage/mmap_store_test.cc | 676 + tests/db/index/storage/parquet_writer_test.cc | 102 + tests/db/index/storage/wal_file_test.cc | 737 + tests/db/index/utils/utils.cc | 804 + tests/db/index/utils/utils.h | 219 + tests/db/reranker_test.cc | 315 + tests/db/sqlengine/CMakeLists.txt | 43 + tests/db/sqlengine/contain_test.cc | 484 + tests/db/sqlengine/forward_recall_test.cc | 938 + tests/db/sqlengine/fts_multi_segment_test.cc | 233 + tests/db/sqlengine/fts_parser_test.cc | 889 + tests/db/sqlengine/fts_recall_test.cc | 824 + tests/db/sqlengine/invert_recall_test.cc | 789 + tests/db/sqlengine/like_test.cc | 372 + tests/db/sqlengine/mock_segment.h | 545 + tests/db/sqlengine/optimizer_test.cc | 322 + tests/db/sqlengine/query_info_test.cc | 697 + tests/db/sqlengine/recall_base.h | 270 + tests/db/sqlengine/simple_rewriter_test.cc | 650 + tests/db/sqlengine/sqlengine_test.cc | 171 + tests/db/sqlengine/test_helper.h | 112 + tests/db/sqlengine/vector_recall_test.cc | 291 + tests/db/utf8_collection_test.cc | 211 + tests/ios_test_sandbox.cc | 45 + tests/test_util.h | 86 + thirdparty/CMakeLists.txt | 34 + thirdparty/CRoaring/CMakeLists.txt | 62 + thirdparty/CRoaring/CRoaring-2.0.4/LICENSE | 48 + .../CRoaring-2.0.4/include/roaring/roaring.h | 2020 ++ thirdparty/CRoaring/CRoaring-2.0.4/roaring.c | 21180 ++++++++++++++++ thirdparty/CRoaring/CRoaring-2.0.4/roaring.h | 2020 ++ thirdparty/CRoaring/CRoaring-2.0.4/roaring.hh | 2881 +++ .../CRoaring/CRoaring-2.0.4/roaring64map.hh | 1864 ++ thirdparty/FastPFOR/CMakeLists.txt | 54 + thirdparty/FastPFOR/fastpfor.riscv.patch | 13 + thirdparty/RaBitQ-Library/CMakeLists.txt | 4 + thirdparty/antlr/CMakeLists.txt | 28 + thirdparty/antlr/antlr4.patch | 48 + thirdparty/antlr/antlr4.windows.patch | 19 + thirdparty/arrow/CMakeLists.txt | 235 + thirdparty/arrow/arrow.android.patch | 82 + thirdparty/arrow/arrow.ios.patch | 39 + thirdparty/arrow/arrow.patch | 26 + thirdparty/arrow/arrow.windows.patch | 25 + thirdparty/cppjieba/CMakeLists.txt | 17 + thirdparty/gflags/CMakeLists.txt | 57 + thirdparty/glog/CMakeLists.txt | 48 + thirdparty/glog/glog.android.patch | 78 + thirdparty/glog/glog.patch | 51 + thirdparty/googletest/CMakeLists.txt | 48 + .../googletest/googletest.windows.patch | 14 + thirdparty/limonp/CMakeLists.txt | 12 + thirdparty/lz4/CMakeLists.txt | 154 + thirdparty/magic_enum/CMakeLists.txt | 4 + thirdparty/protobuf/CMakeLists.txt | 76 + thirdparty/protobuf/protobuf.windows.patch | 15 + thirdparty/rocksdb/CMakeLists.txt | 83 + thirdparty/rocksdb/rocksdb.android.patch | 13 + thirdparty/rocksdb/rocksdb.ios.patch | 29 + thirdparty/rocksdb/rocksdb.windows.patch | 13 + thirdparty/snowball/CMakeLists.txt | 100 + thirdparty/sparsehash/CMakeLists.txt | 16 + thirdparty/sparsehash/sparseconfig.h | 46 + .../sparsehash/sparsehash.windows.patch | 13 + thirdparty/utf8proc/CMakeLists.txt | 4 + thirdparty/yaml-cpp/CMakeLists.txt | 17 + tools/CMakeLists.txt | 9 + tools/core/CMakeLists.txt | 59 + tools/core/README.md | 150 + tools/core/bench.cc | 911 + tools/core/bench_original.cc | 1066 + tools/core/bench_result.h | 147 + tools/core/convert_cohere_parquet.py | 187 + tools/core/filter_result_cache.h | 77 + tools/core/flow.h | 555 + tools/core/helper.h | 297 + tools/core/index_meta_helper.h | 121 + tools/core/local_builder.cc | 1357 + tools/core/local_builder_original.cc | 1254 + tools/core/meta_segment_common.h | 34 + tools/core/recall.cc | 1787 ++ tools/core/recall_original.cc | 2093 ++ tools/core/txt2vecs.cc | 723 + tools/core/txt_input_reader.h | 549 + tools/core/vecs_common.h | 56 + tools/core/vecs_index_holder.h | 433 + tools/core/vecs_reader.h | 378 + tools/db/CMakeLists.txt | 13 + tools/db/fts_bench_main.cc | 1859 ++ 1214 files changed, 431743 insertions(+) create mode 100644 .clang-format create mode 100644 .clang-tidy create mode 100644 .gitattributes create mode 100644 .github/CODEOWNERS create mode 100644 .github/ISSUE_TEMPLATE/benchmark.yml create mode 100644 .github/ISSUE_TEMPLATE/bug_report.yml create mode 100644 .github/ISSUE_TEMPLATE/config.yml create mode 100644 .github/ISSUE_TEMPLATE/enhancement.yml create mode 100644 .github/ISSUE_TEMPLATE/feature_request.yml create mode 100644 .github/ISSUE_TEMPLATE/integration.yml create mode 100644 .github/ISSUE_TEMPLATE/profiling.yml create mode 100644 .github/auto-assign-config.yml create mode 100644 .github/cmake/subproject-integration/CMakeLists.txt create mode 100644 .github/codecov.yml create mode 100644 .github/dependabot.yml create mode 100644 .github/workflows/01-ci-pipeline.yml create mode 100644 .github/workflows/02-lint-check.yml create mode 100644 .github/workflows/03-macos-linux-build.yml create mode 100644 .github/workflows/04-android-build.yml create mode 100644 .github/workflows/05-windows-build.yml create mode 100644 .github/workflows/06-ios-build.yml create mode 100644 .github/workflows/07-linux-riscv-build.yml create mode 100644 .github/workflows/08-cmake-subproject-integration.yml create mode 100644 .github/workflows/_build_wheel_job.yml create mode 100644 .github/workflows/auto-assign.yml create mode 100644 .github/workflows/clang_tidy.yml create mode 100644 .github/workflows/community-pr-handler.yml create mode 100644 .github/workflows/continuous_bench.yml create mode 100644 .github/workflows/docker/Dockerfile.linux_x64_glibc228 create mode 100644 .github/workflows/issue-auto-assign.yml create mode 100644 .github/workflows/nightly_coverage.yml create mode 100644 .github/workflows/scripts/run_vdb.sh create mode 100644 .gitignore create mode 100644 .gitmodules create mode 100644 CMakeLists.txt create mode 100644 CODE_OF_CONDUCT.md create mode 100644 CONTRIBUTING.md create mode 100644 LICENSE create mode 100644 NOTICE create mode 100644 README.md create mode 100644 README.wehub.md create mode 100644 README_CN.md create mode 100644 cmake/bazel.cmake create mode 100644 cmake/iOSBundleInfo.plist.in create mode 100644 cmake/option.cmake create mode 100644 cmake/utils.cmake create mode 100644 examples/c++/CMakeLists.txt create mode 100644 examples/c++/ailego/main.cc create mode 100644 examples/c++/core/external_vector_example.cc create mode 100644 examples/c++/core/main.cc create mode 100644 examples/c++/db/main.cc create mode 100644 examples/c/CMakeLists.txt create mode 100644 examples/c/basic_example.c create mode 100644 examples/c/collection_schema_example.c create mode 100644 examples/c/diskann_example.c create mode 100644 examples/c/doc_example.c create mode 100644 examples/c/field_schema_example.c create mode 100644 examples/c/index_example.c create mode 100644 examples/c/optimized_example.c create mode 100644 pyproject.toml create mode 100644 python/tests/detail/distance_helper.py create mode 100644 python/tests/detail/doc_helper.py create mode 100644 python/tests/detail/fixture_helper.py create mode 100644 python/tests/detail/params_helper.py create mode 100644 python/tests/detail/support_helper.py create mode 100644 python/tests/detail/test_collection_concurrency.py create mode 100644 python/tests/detail/test_collection_create_and_open.py create mode 100644 python/tests/detail/test_collection_ddl.py create mode 100644 python/tests/detail/test_collection_dml.py create mode 100644 python/tests/detail/test_collection_dql.py create mode 100644 python/tests/detail/test_collection_exception.py create mode 100644 python/tests/detail/test_collection_open.py create mode 100644 python/tests/detail/test_collection_recall.py create mode 100644 python/tests/detail/test_db_config.py create mode 100644 python/tests/test_collection.py create mode 100644 python/tests/test_collection_diskann.py create mode 100644 python/tests/test_collection_fts.py create mode 100644 python/tests/test_collection_fts_vector_hybrid.py create mode 100644 python/tests/test_collection_hnsw_rabitq.py create mode 100644 python/tests/test_convert.py create mode 100644 python/tests/test_doc.py create mode 100644 python/tests/test_embedding.py create mode 100644 python/tests/test_fts_query.py create mode 100644 python/tests/test_gil_release.py create mode 100644 python/tests/test_hnsw_contiguous_memory.py create mode 100644 python/tests/test_jieba_default_dict.py create mode 100644 python/tests/test_params.py create mode 100644 python/tests/test_query_executor.py create mode 100644 python/tests/test_reranker.py create mode 100644 python/tests/test_schema.py create mode 100644 python/tests/test_typing.py create mode 100644 python/tests/test_util.py create mode 100644 python/tests/test_vamana.py create mode 100644 python/zvec/__init__.py create mode 100644 python/zvec/__init__.pyi create mode 100644 python/zvec/common/__init__.py create mode 100644 python/zvec/common/constants.py create mode 100644 python/zvec/executor/__init__.py create mode 100644 python/zvec/executor/query_executor.py create mode 100644 python/zvec/extension/__init__.py create mode 100644 python/zvec/extension/bm25_embedding_function.py create mode 100644 python/zvec/extension/embedding_function.py create mode 100644 python/zvec/extension/http_embedding_function.py create mode 100644 python/zvec/extension/jina_embedding_function.py create mode 100644 python/zvec/extension/jina_function.py create mode 100644 python/zvec/extension/multi_vector_reranker.py create mode 100644 python/zvec/extension/openai_embedding_function.py create mode 100644 python/zvec/extension/openai_function.py create mode 100644 python/zvec/extension/qwen_embedding_function.py create mode 100644 python/zvec/extension/qwen_function.py create mode 100644 python/zvec/extension/qwen_rerank_function.py create mode 100644 python/zvec/extension/rerank_function.py create mode 100644 python/zvec/extension/sentence_transformer_embedding_function.py create mode 100644 python/zvec/extension/sentence_transformer_function.py create mode 100644 python/zvec/extension/sentence_transformer_rerank_function.py create mode 100644 python/zvec/model/__init__.py create mode 100644 python/zvec/model/collection.py create mode 100644 python/zvec/model/convert.py create mode 100644 python/zvec/model/doc.py create mode 100644 python/zvec/model/param/__init__.py create mode 100644 python/zvec/model/param/__init__.pyi create mode 100644 python/zvec/model/param/query.py create mode 100644 python/zvec/model/schema/__init__.py create mode 100644 python/zvec/model/schema/__init__.pyi create mode 100644 python/zvec/model/schema/collection_schema.py create mode 100644 python/zvec/model/schema/field_schema.py create mode 100644 python/zvec/py.typed create mode 100644 python/zvec/tool/__init__.py create mode 100644 python/zvec/tool/util.py create mode 100644 python/zvec/typing/__init__.py create mode 100644 python/zvec/typing/__init__.pyi create mode 100644 python/zvec/typing/enum.py create mode 100644 python/zvec/zvec.py create mode 100644 scripts/README.md create mode 100755 scripts/build_android.sh create mode 100755 scripts/build_ios.sh create mode 100755 scripts/gcov.sh create mode 100755 scripts/generate_standard_tokenizer_unicode.py create mode 100644 src/CMakeLists.txt create mode 100644 src/ailego/CMakeLists.txt create mode 100644 src/ailego/algorithm/binary_quantizer.cc create mode 100644 src/ailego/algorithm/binary_quantizer.h create mode 100644 src/ailego/algorithm/integer_quantizer.cc create mode 100644 src/ailego/algorithm/integer_quantizer.h create mode 100644 src/ailego/algorithm/kmeans.h create mode 100644 src/ailego/algorithm/lloyd_cluster.h create mode 100644 src/ailego/buffer/block_eviction_queue.cc create mode 100644 src/ailego/buffer/vector_page_table.cc create mode 100644 src/ailego/container/bitmap.cc create mode 100644 src/ailego/container/bitmap.h create mode 100644 src/ailego/container/bloom_filter.h create mode 100644 src/ailego/container/params.cc create mode 100644 src/ailego/container/reservoir.h create mode 100644 src/ailego/container/vector_array.h create mode 100644 src/ailego/encoding/json/mod_json.c create mode 100644 src/ailego/hash/crc32c.cc create mode 100644 src/ailego/internal/cpu_features.cc create mode 100644 src/ailego/internal/cpu_features.h create mode 100644 src/ailego/io/file.cc create mode 100644 src/ailego/io/file_lock.cc create mode 100644 src/ailego/io/file_lock.h create mode 100644 src/ailego/io/file_writer.h create mode 100644 src/ailego/logger/logger.cc create mode 100644 src/ailego/math/cosine_distance_matrix.h create mode 100644 src/ailego/math/distance.h create mode 100644 src/ailego/math/distance_matrix.h create mode 100644 src/ailego/math/distance_matrix_accum_fp16.i create mode 100644 src/ailego/math/distance_matrix_accum_fp32.i create mode 100644 src/ailego/math/distance_matrix_accum_int4.i create mode 100644 src/ailego/math/distance_matrix_accum_int8.i create mode 100644 src/ailego/math/distance_matrix_euclidean_utility.i create mode 100644 src/ailego/math/distance_matrix_fp16.i create mode 100644 src/ailego/math/distance_matrix_fp32.i create mode 100644 src/ailego/math/distance_matrix_inner_product_utility.i create mode 100644 src/ailego/math/distance_matrix_int32.i create mode 100644 src/ailego/math/distance_matrix_int64.i create mode 100644 src/ailego/math/distance_matrix_mips_utility.i create mode 100644 src/ailego/math/distance_matrix_popcnt.i create mode 100644 src/ailego/math/distance_utility.h create mode 100644 src/ailego/math/euclidean_distance_matrix.h create mode 100644 src/ailego/math/euclidean_distance_matrix_fp16_avx.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp16_avx512.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp16_avx512fp16.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp16_dispatch.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp16_neon.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp32_avx.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp32_avx512.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp32_dispatch.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp32_neon.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp32_sse.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_int4_avx2.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_int4_dispatch.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_int4_sse.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_int8_avx2.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_int8_dispatch.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_int8_sse.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_scalar.cc create mode 100644 src/ailego/math/fht.h create mode 100644 src/ailego/math/fht_avx2.cc create mode 100644 src/ailego/math/fht_avx512.cc create mode 100644 src/ailego/math/fht_dispatch.cc create mode 100644 src/ailego/math/fht_neon.cc create mode 100644 src/ailego/math/fht_scalar.cc create mode 100644 src/ailego/math/fht_sse.cc create mode 100644 src/ailego/math/inner_product_matrix.h create mode 100644 src/ailego/math/inner_product_matrix_fp16.cc create mode 100644 src/ailego/math/inner_product_matrix_fp16_avx.cc create mode 100644 src/ailego/math/inner_product_matrix_fp16_avx512.cc create mode 100644 src/ailego/math/inner_product_matrix_fp16_avx512fp16.cc create mode 100644 src/ailego/math/inner_product_matrix_fp16_dispatch.cc create mode 100644 src/ailego/math/inner_product_matrix_fp16_neon.cc create mode 100644 src/ailego/math/inner_product_matrix_fp32.cc create mode 100644 src/ailego/math/inner_product_matrix_fp32_avx.cc create mode 100644 src/ailego/math/inner_product_matrix_fp32_avx512.cc create mode 100644 src/ailego/math/inner_product_matrix_fp32_dispatch.cc create mode 100644 src/ailego/math/inner_product_matrix_fp32_neon.cc create mode 100644 src/ailego/math/inner_product_matrix_fp32_sse.cc create mode 100644 src/ailego/math/inner_product_matrix_int4_avx2.cc create mode 100644 src/ailego/math/inner_product_matrix_int4_dispatch.cc create mode 100644 src/ailego/math/inner_product_matrix_int4_sse.cc create mode 100644 src/ailego/math/inner_product_matrix_int8_avx2.cc create mode 100644 src/ailego/math/inner_product_matrix_int8_dispatch.cc create mode 100644 src/ailego/math/inner_product_matrix_int8_sse.cc create mode 100644 src/ailego/math/inner_product_matrix_scalar.cc create mode 100644 src/ailego/math/matrix_define.i create mode 100644 src/ailego/math/matrix_utility.i create mode 100644 src/ailego/math/mips_euclidean_distance_matrix.h create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp32.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_scalar.cc create mode 100644 src/ailego/math/norm1_matrix.h create mode 100644 src/ailego/math/norm1_matrix_fp16.cc create mode 100644 src/ailego/math/norm1_matrix_fp32.cc create mode 100644 src/ailego/math/norm2_matrix.h create mode 100644 src/ailego/math/norm2_matrix_fp16.cc create mode 100644 src/ailego/math/norm2_matrix_fp32.cc create mode 100644 src/ailego/math/norm_matrix.h create mode 100644 src/ailego/math/norm_matrix_fp16.i create mode 100644 src/ailego/math/norm_matrix_fp32.i create mode 100644 src/ailego/math/normalizer.cc create mode 100644 src/ailego/math/normalizer.h create mode 100644 src/ailego/math_batch/cosine_distance_batch.h create mode 100644 src/ailego/math_batch/distance_batch.h create mode 100644 src/ailego/math_batch/distance_batch_math.h create mode 100644 src/ailego/math_batch/euclidean_distance_batch.h create mode 100644 src/ailego/math_batch/euclidean_distance_batch_dispatch.cc create mode 100644 src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx2.cc create mode 100644 src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512.cc create mode 100644 src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512fp16.cc create mode 100644 src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx2.cc create mode 100644 src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx512.cc create mode 100644 src/ailego/math_batch/inner_product_distance_batch.h create mode 100644 src/ailego/math_batch/inner_product_distance_batch_dispatch.cc create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512fp16.cc create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512fp16.cc create mode 100644 src/ailego/parallel/lock.h create mode 100644 src/ailego/parallel/multi_thread_list.h create mode 100644 src/ailego/parallel/semaphore.h create mode 100644 src/ailego/parallel/thread_pool.cc create mode 100644 src/ailego/pattern/defer.h create mode 100644 src/ailego/pattern/scope_guard.h create mode 100644 src/ailego/utility/bit_string_helper.h create mode 100644 src/ailego/utility/bitset_helper.cc create mode 100644 src/ailego/utility/bitset_helper.h create mode 100644 src/ailego/utility/concurrency_helper.cc create mode 100644 src/ailego/utility/concurrency_helper.h create mode 100644 src/ailego/utility/dl_helper.cc create mode 100644 src/ailego/utility/dl_helper.h create mode 100644 src/ailego/utility/file_helper.cc create mode 100644 src/ailego/utility/float_helper.cc create mode 100644 src/ailego/utility/math_helper.h create mode 100644 src/ailego/utility/matrix_helper.h create mode 100644 src/ailego/utility/memory_helper.cc create mode 100644 src/ailego/utility/memory_helper.h create mode 100644 src/ailego/utility/string_helper.cc create mode 100644 src/ailego/utility/time_helper.cc create mode 100644 src/ailego/version.cc create mode 100644 src/ailego/version.h create mode 100644 src/ailego/version.i create mode 100644 src/binding/CMakeLists.txt create mode 100644 src/binding/c/CMakeLists.txt create mode 100644 src/binding/c/c_api.cc create mode 100644 src/binding/python/CMakeLists.txt create mode 100644 src/binding/python/binding.cc create mode 100644 src/binding/python/exports.mac create mode 100644 src/binding/python/exports.map create mode 100644 src/binding/python/include/python_collection.h create mode 100644 src/binding/python/include/python_config.h create mode 100644 src/binding/python/include/python_doc.h create mode 100644 src/binding/python/include/python_param.h create mode 100644 src/binding/python/include/python_reranker.h create mode 100644 src/binding/python/include/python_schema.h create mode 100644 src/binding/python/include/python_type.h create mode 100644 src/binding/python/model/common/python_config.cc create mode 100644 src/binding/python/model/param/python_param.cc create mode 100644 src/binding/python/model/python_collection.cc create mode 100644 src/binding/python/model/python_doc.cc create mode 100644 src/binding/python/model/python_reranker.cc create mode 100644 src/binding/python/model/schema/python_schema.cc create mode 100644 src/binding/python/typing/python_type.cc create mode 100644 src/core/CMakeLists.txt create mode 100644 src/core/algorithm/CMakeLists.txt create mode 100644 src/core/algorithm/cluster/CMakeLists.txt create mode 100644 src/core/algorithm/cluster/cluster_params.h create mode 100644 src/core/algorithm/cluster/kmeans_cluster.cc create mode 100644 src/core/algorithm/cluster/linear_seeker.cc create mode 100644 src/core/algorithm/cluster/linear_seeker.h create mode 100644 src/core/algorithm/cluster/multi_chunk_cluster.cc create mode 100644 src/core/algorithm/cluster/multi_chunk_cluster.h create mode 100644 src/core/algorithm/cluster/opt_kmeans_cluster.cc create mode 100644 src/core/algorithm/cluster/seeker.h create mode 100644 src/core/algorithm/cluster/stratified_cluster.cc create mode 100644 src/core/algorithm/cluster/stratified_cluster_trainer.cc create mode 100644 src/core/algorithm/cluster/stratified_cluster_trainer.h create mode 100644 src/core/algorithm/cluster/vector_mean.h create mode 100644 src/core/algorithm/diskann/CMakeLists.txt create mode 100644 src/core/algorithm/diskann/diskann_algorithm.cc create mode 100644 src/core/algorithm/diskann/diskann_algorithm.h create mode 100644 src/core/algorithm/diskann/diskann_builder.cc create mode 100644 src/core/algorithm/diskann/diskann_builder.h create mode 100644 src/core/algorithm/diskann/diskann_builder_entity.cc create mode 100644 src/core/algorithm/diskann/diskann_builder_entity.h create mode 100644 src/core/algorithm/diskann/diskann_context.cc create mode 100644 src/core/algorithm/diskann/diskann_context.h create mode 100644 src/core/algorithm/diskann/diskann_dist_calculator.h create mode 100644 src/core/algorithm/diskann/diskann_entity.cc create mode 100644 src/core/algorithm/diskann/diskann_entity.h create mode 100644 src/core/algorithm/diskann/diskann_file_reader.cc create mode 100644 src/core/algorithm/diskann/diskann_file_reader.h create mode 100644 src/core/algorithm/diskann/diskann_holder.h create mode 100644 src/core/algorithm/diskann/diskann_index_provider.h create mode 100644 src/core/algorithm/diskann/diskann_indexer.cc create mode 100644 src/core/algorithm/diskann/diskann_indexer.h create mode 100644 src/core/algorithm/diskann/diskann_params.h create mode 100644 src/core/algorithm/diskann/diskann_pq_table.cc create mode 100644 src/core/algorithm/diskann/diskann_pq_table.h create mode 100644 src/core/algorithm/diskann/diskann_pq_trainer.cc create mode 100644 src/core/algorithm/diskann/diskann_pq_trainer.h create mode 100644 src/core/algorithm/diskann/diskann_reducer.cc create mode 100644 src/core/algorithm/diskann/diskann_reducer.h create mode 100644 src/core/algorithm/diskann/diskann_reducer_entity.cc create mode 100644 src/core/algorithm/diskann/diskann_reducer_entity.h create mode 100644 src/core/algorithm/diskann/diskann_searcher.cc create mode 100644 src/core/algorithm/diskann/diskann_searcher.h create mode 100644 src/core/algorithm/diskann/diskann_searcher_entity.cc create mode 100644 src/core/algorithm/diskann/diskann_searcher_entity.h create mode 100644 src/core/algorithm/diskann/diskann_streamer.cc create mode 100644 src/core/algorithm/diskann/diskann_streamer.h create mode 100644 src/core/algorithm/diskann/diskann_util.h create mode 100644 src/core/algorithm/diskann/diskann_vecs_reader.h create mode 100644 src/core/algorithm/diskann/diskann_visit_filter.h create mode 100644 src/core/algorithm/flat/CMakeLists.txt create mode 100644 src/core/algorithm/flat/flat_builder.cc create mode 100644 src/core/algorithm/flat/flat_builder.h create mode 100644 src/core/algorithm/flat/flat_distance_matrix.h create mode 100644 src/core/algorithm/flat/flat_index_format.h create mode 100644 src/core/algorithm/flat/flat_searcher.cc create mode 100644 src/core/algorithm/flat/flat_searcher.h create mode 100644 src/core/algorithm/flat/flat_searcher_context.h create mode 100644 src/core/algorithm/flat/flat_searcher_provider.h create mode 100644 src/core/algorithm/flat/flat_streamer.cc create mode 100644 src/core/algorithm/flat/flat_streamer.h create mode 100644 src/core/algorithm/flat/flat_streamer_context.h create mode 100644 src/core/algorithm/flat/flat_streamer_dumper.h create mode 100644 src/core/algorithm/flat/flat_streamer_entity.cc create mode 100644 src/core/algorithm/flat/flat_streamer_entity.h create mode 100644 src/core/algorithm/flat/flat_streamer_provider.h create mode 100644 src/core/algorithm/flat/flat_utility.h create mode 100644 src/core/algorithm/flat_sparse/CMakeLists.txt create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_builder.cc create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_builder.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_context.cc create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_context.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_entity.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_index_format.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_provider.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_search.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_searcher.cc create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_searcher.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.cc create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_streamer.cc create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_streamer.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.cc create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.h create mode 100644 src/core/algorithm/flat_sparse/flat_sparse_utility.h create mode 100644 src/core/algorithm/hnsw/CMakeLists.txt create mode 100644 src/core/algorithm/hnsw/hnsw_algorithm.cc create mode 100644 src/core/algorithm/hnsw/hnsw_algorithm.h create mode 100644 src/core/algorithm/hnsw/hnsw_chunk.cc create mode 100644 src/core/algorithm/hnsw/hnsw_chunk.h create mode 100644 src/core/algorithm/hnsw/hnsw_context.cc create mode 100644 src/core/algorithm/hnsw/hnsw_context.h create mode 100644 src/core/algorithm/hnsw/hnsw_dist_calculator.h create mode 100644 src/core/algorithm/hnsw/hnsw_entity.cc create mode 100644 src/core/algorithm/hnsw/hnsw_entity.h create mode 100644 src/core/algorithm/hnsw/hnsw_index_hash.h create mode 100644 src/core/algorithm/hnsw/hnsw_index_provider.h create mode 100644 src/core/algorithm/hnsw/hnsw_params.h create mode 100644 src/core/algorithm/hnsw/hnsw_streamer.cc create mode 100644 src/core/algorithm/hnsw/hnsw_streamer.h create mode 100644 src/core/algorithm/hnsw/hnsw_streamer_entity.cc create mode 100644 src/core/algorithm/hnsw/hnsw_streamer_entity.h create mode 100644 src/core/algorithm/hnsw_rabitq/CMakeLists.txt create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_hash.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_provider.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_params.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_entity.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.h create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.cc create mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.h create mode 100644 src/core/algorithm/hnsw_rabitq/rabitq_converter.cc create mode 100644 src/core/algorithm/hnsw_rabitq/rabitq_converter.h create mode 100644 src/core/algorithm/hnsw_rabitq/rabitq_params.h create mode 100644 src/core/algorithm/hnsw_rabitq/rabitq_reformer.cc create mode 100644 src/core/algorithm/hnsw_rabitq/rabitq_reformer.h create mode 100644 src/core/algorithm/hnsw_rabitq/rabitq_utils.cc create mode 100644 src/core/algorithm/hnsw_rabitq/rabitq_utils.h create mode 100644 src/core/algorithm/hnsw_sparse/CMakeLists.txt create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_context.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_dist_calculator.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_index_hash.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_index_provider.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_params.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.h create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.cc create mode 100644 src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.h create mode 100644 src/core/algorithm/ivf/CMakeLists.txt create mode 100644 src/core/algorithm/ivf/ivf_builder.cc create mode 100644 src/core/algorithm/ivf/ivf_builder.h create mode 100644 src/core/algorithm/ivf/ivf_centroid_index.cc create mode 100644 src/core/algorithm/ivf/ivf_centroid_index.h create mode 100644 src/core/algorithm/ivf/ivf_distance_calculator.cc create mode 100644 src/core/algorithm/ivf/ivf_distance_calculator.h create mode 100644 src/core/algorithm/ivf/ivf_dumper.cc create mode 100644 src/core/algorithm/ivf/ivf_dumper.h create mode 100644 src/core/algorithm/ivf/ivf_entity.cc create mode 100644 src/core/algorithm/ivf/ivf_entity.h create mode 100644 src/core/algorithm/ivf/ivf_index_format.h create mode 100644 src/core/algorithm/ivf/ivf_index_provider.h create mode 100644 src/core/algorithm/ivf/ivf_params.h create mode 100644 src/core/algorithm/ivf/ivf_searcher.cc create mode 100644 src/core/algorithm/ivf/ivf_searcher.h create mode 100644 src/core/algorithm/ivf/ivf_searcher_context.h create mode 100644 src/core/algorithm/ivf/ivf_streamer.cc create mode 100644 src/core/algorithm/ivf/ivf_streamer.h create mode 100644 src/core/algorithm/ivf/ivf_utility.h create mode 100644 src/core/algorithm/vamana/CMakeLists.txt create mode 100644 src/core/algorithm/vamana/vamana_algorithm.cc create mode 100644 src/core/algorithm/vamana/vamana_algorithm.h create mode 100644 src/core/algorithm/vamana/vamana_context.cc create mode 100644 src/core/algorithm/vamana/vamana_context.h create mode 100644 src/core/algorithm/vamana/vamana_dist_calculator.h create mode 100644 src/core/algorithm/vamana/vamana_entity.cc create mode 100644 src/core/algorithm/vamana/vamana_entity.h create mode 100644 src/core/algorithm/vamana/vamana_index_provider.h create mode 100644 src/core/algorithm/vamana/vamana_params.h create mode 100644 src/core/algorithm/vamana/vamana_streamer.cc create mode 100644 src/core/algorithm/vamana/vamana_streamer.h create mode 100644 src/core/algorithm/vamana/vamana_streamer_entity.cc create mode 100644 src/core/algorithm/vamana/vamana_streamer_entity.h create mode 100644 src/core/framework/CMakeLists.txt create mode 100644 src/core/framework/index_cluster.cc create mode 100644 src/core/framework/index_context.cc create mode 100644 src/core/framework/index_converter.cc create mode 100644 src/core/framework/index_error.cc create mode 100644 src/core/framework/index_factory.cc create mode 100644 src/core/framework/index_flow.cc create mode 100644 src/core/framework/index_helper.cc create mode 100644 src/core/framework/index_logger.cc create mode 100644 src/core/framework/index_mapping.cc create mode 100644 src/core/framework/index_meta.cc create mode 100644 src/core/framework/index_plugin.cc create mode 100644 src/core/framework/index_version.cc create mode 100644 src/core/interface/CMakeLists.txt create mode 100644 src/core/interface/index.cc create mode 100644 src/core/interface/index_factory.cc create mode 100644 src/core/interface/index_param.cc create mode 100644 src/core/interface/indexes/diskann_index.cc create mode 100644 src/core/interface/indexes/flat_index.cc create mode 100644 src/core/interface/indexes/hnsw_index.cc create mode 100644 src/core/interface/indexes/hnsw_rabitq_index.cc create mode 100644 src/core/interface/indexes/holder_builder.h create mode 100644 src/core/interface/indexes/ivf_index.cc create mode 100644 src/core/interface/indexes/vamana_index.cc create mode 100644 src/core/interface/utils/utils.h create mode 100644 src/core/metric/CMakeLists.txt create mode 100644 src/core/metric/cosine_metric.cc create mode 100644 src/core/metric/euclidean_metric.cc create mode 100644 src/core/metric/inner_product_metric.cc create mode 100644 src/core/metric/metric_params.h create mode 100644 src/core/metric/mips_euclidean_metric.cc create mode 100644 src/core/metric/quantized_integer_metric.cc create mode 100644 src/core/metric/quantized_integer_metric_batch.h create mode 100644 src/core/metric/quantized_integer_metric_matrix.h create mode 100644 src/core/metric/uniform_int8_metric.cc create mode 100644 src/core/mixed_reducer/CMakeLists.txt create mode 100644 src/core/mixed_reducer/mixed_reducer_params.h create mode 100644 src/core/mixed_reducer/mixed_streamer_reducer.cc create mode 100644 src/core/mixed_reducer/mixed_streamer_reducer.h create mode 100644 src/core/plugin/CMakeLists.txt create mode 100644 src/core/plugin/diskann_plugin.cc create mode 100644 src/core/quantizer/CMakeLists.txt create mode 100644 src/core/quantizer/binary_converter.cc create mode 100644 src/core/quantizer/binary_reformer.cc create mode 100644 src/core/quantizer/cosine_converter.cc create mode 100644 src/core/quantizer/cosine_reformer.cc create mode 100644 src/core/quantizer/half_float_converter.cc create mode 100644 src/core/quantizer/half_float_reformer.cc create mode 100644 src/core/quantizer/integer_quantizer_converter.cc create mode 100644 src/core/quantizer/integer_quantizer_reformer.cc create mode 100644 src/core/quantizer/mips_converter.cc create mode 100644 src/core/quantizer/mips_reformer.cc create mode 100644 src/core/quantizer/quantizer_params.h create mode 100644 src/core/quantizer/record_quantizer.h create mode 100644 src/core/quantizer/rotator/fht_rotator.cc create mode 100644 src/core/quantizer/rotator/fht_rotator.h create mode 100644 src/core/quantizer/rotator/matrix_rotator.cc create mode 100644 src/core/quantizer/rotator/matrix_rotator.h create mode 100644 src/core/quantizer/rotator/rotator.cc create mode 100644 src/core/quantizer/rotator/rotator.h create mode 100644 src/core/quantizer/uniform_int8_converter.cc create mode 100644 src/core/quantizer/uniform_int8_reformer.cc create mode 100644 src/core/utility/CMakeLists.txt create mode 100644 src/core/utility/basic_refiner.cc create mode 100644 src/core/utility/block_heap.cc create mode 100644 src/core/utility/block_heap.h create mode 100644 src/core/utility/buffer_storage.cc create mode 100644 src/core/utility/file_dumper.cc create mode 100644 src/core/utility/file_read_storage.cc create mode 100644 src/core/utility/linear_pool.h create mode 100644 src/core/utility/memory_dumper.cc create mode 100644 src/core/utility/memory_read_storage.cc create mode 100644 src/core/utility/mmap_file_read_storage.cc create mode 100644 src/core/utility/mmap_file_storage.cc create mode 100644 src/core/utility/sparse_utility.h create mode 100644 src/core/utility/topk_result_builder.h create mode 100644 src/core/utility/utility_params.h create mode 100644 src/core/utility/visit_filter.h create mode 100644 src/db/CMakeLists.txt create mode 100644 src/db/collection.cc create mode 100644 src/db/common/CMakeLists.txt create mode 100644 src/db/common/cgroup_util.cc create mode 100644 src/db/common/cgroup_util.h create mode 100644 src/db/common/concurrent_roaring_bitmap.cc create mode 100644 src/db/common/concurrent_roaring_bitmap.h create mode 100644 src/db/common/config.cc create mode 100644 src/db/common/constants.h create mode 100644 src/db/common/error_code.cc create mode 100644 src/db/common/error_code.h create mode 100644 src/db/common/file_helper.cc create mode 100644 src/db/common/file_helper.h create mode 100644 src/db/common/global_resource.cc create mode 100644 src/db/common/global_resource.h create mode 100644 src/db/common/glogger.h create mode 100644 src/db/common/logger.h create mode 100644 src/db/common/profiler.h create mode 100644 src/db/common/rocksdb_context.cc create mode 100644 src/db/common/rocksdb_context.h create mode 100644 src/db/common/status.cc create mode 100644 src/db/common/typedef.h create mode 100644 src/db/common/utils.cc create mode 100644 src/db/common/utils.h create mode 100644 src/db/index/CMakeLists.txt create mode 100644 src/db/index/column/column_indexer.h create mode 100644 src/db/index/column/common/index_results.h create mode 100644 src/db/index/column/fts_column/FtsLexer.g4 create mode 100644 src/db/index/column/fts_column/FtsParser.g4 create mode 100644 src/db/index/column/fts_column/bm25_scorer.cc create mode 100644 src/db/index/column/fts_column/bm25_scorer.h create mode 100644 src/db/index/column/fts_column/fts_ast_rewriter.cc create mode 100644 src/db/index/column/fts_column/fts_ast_rewriter.h create mode 100644 src/db/index/column/fts_column/fts_column_indexer.cc create mode 100644 src/db/index/column/fts_column/fts_column_indexer.h create mode 100644 src/db/index/column/fts_column/fts_index_results.h create mode 100644 src/db/index/column/fts_column/fts_indexer.cc create mode 100644 src/db/index/column/fts_column/fts_indexer.h create mode 100644 src/db/index/column/fts_column/fts_pipeline.h create mode 100644 src/db/index/column/fts_column/fts_query_ast.h create mode 100644 src/db/index/column/fts_column/fts_rocksdb_merge.cc create mode 100644 src/db/index/column/fts_column/fts_rocksdb_merge.h create mode 100644 src/db/index/column/fts_column/fts_rocksdb_reducer.cc create mode 100644 src/db/index/column/fts_column/fts_rocksdb_reducer.h create mode 100644 src/db/index/column/fts_column/fts_types.h create mode 100644 src/db/index/column/fts_column/fts_utils.cc create mode 100644 src/db/index/column/fts_column/fts_utils.h create mode 100644 src/db/index/column/fts_column/gen/FtsLexer.cc create mode 100644 src/db/index/column/fts_column/gen/FtsLexer.h create mode 100644 src/db/index/column/fts_column/gen/FtsLexer.interp create mode 100644 src/db/index/column/fts_column/gen/FtsLexer.tokens create mode 100644 src/db/index/column/fts_column/gen/FtsParser.cc create mode 100644 src/db/index/column/fts_column/gen/FtsParser.h create mode 100644 src/db/index/column/fts_column/gen/FtsParser.interp create mode 100644 src/db/index/column/fts_column/gen/FtsParser.tokens create mode 100644 src/db/index/column/fts_column/gen/FtsParserBaseListener.cc create mode 100644 src/db/index/column/fts_column/gen/FtsParserBaseListener.h create mode 100644 src/db/index/column/fts_column/gen/FtsParserListener.cc create mode 100644 src/db/index/column/fts_column/gen/FtsParserListener.h create mode 100644 src/db/index/column/fts_column/gen_parser.sh create mode 100644 src/db/index/column/fts_column/iterator/fts_candidate_iterator.cc create mode 100644 src/db/index/column/fts_column/iterator/fts_candidate_iterator.h create mode 100644 src/db/index/column/fts_column/iterator/fts_conjunction_iterator.cc create mode 100644 src/db/index/column/fts_column/iterator/fts_conjunction_iterator.h create mode 100644 src/db/index/column/fts_column/iterator/fts_disjunction_iterator.cc create mode 100644 src/db/index/column/fts_column/iterator/fts_disjunction_iterator.h create mode 100644 src/db/index/column/fts_column/iterator/fts_doc_iterator.h create mode 100644 src/db/index/column/fts_column/iterator/fts_phrase_iterator.cc create mode 100644 src/db/index/column/fts_column/iterator/fts_phrase_iterator.h create mode 100644 src/db/index/column/fts_column/iterator/fts_term_iterator.cc create mode 100644 src/db/index/column/fts_column/iterator/fts_term_iterator.h create mode 100644 src/db/index/column/fts_column/parser/fts_query_parser.cc create mode 100644 src/db/index/column/fts_column/parser/fts_query_parser.h create mode 100644 src/db/index/column/fts_column/posting/bitpacked_posting_list.cc create mode 100644 src/db/index/column/fts_column/posting/bitpacked_posting_list.h create mode 100644 src/db/index/column/fts_column/posting/bitpacked_simd_avx2.cc create mode 100644 src/db/index/column/fts_column/posting/bitpacked_simd_avx2.h create mode 100644 src/db/index/column/fts_column/posting/bitpacked_simd_dispatch.cc create mode 100644 src/db/index/column/fts_column/posting/bitpacked_simd_dispatch.h create mode 100644 src/db/index/column/fts_column/posting/bitpacked_simd_scalar.cc create mode 100644 src/db/index/column/fts_column/posting/bitpacked_simd_scalar.h create mode 100644 src/db/index/column/fts_column/posting/bitpacked_simd_sse41.cc create mode 100644 src/db/index/column/fts_column/posting/bitpacked_simd_sse41.h create mode 100644 src/db/index/column/fts_column/tokenizer/ascii_folding_token_filter.cc create mode 100644 src/db/index/column/fts_column/tokenizer/ascii_folding_token_filter.h create mode 100644 src/db/index/column/fts_column/tokenizer/jieba_tokenizer.cc create mode 100644 src/db/index/column/fts_column/tokenizer/jieba_tokenizer.h create mode 100644 src/db/index/column/fts_column/tokenizer/standard_tokenizer.cc create mode 100644 src/db/index/column/fts_column/tokenizer/standard_tokenizer.h create mode 100644 src/db/index/column/fts_column/tokenizer/standard_tokenizer_unicode.inc create mode 100644 src/db/index/column/fts_column/tokenizer/stemmer_token_filter.cc create mode 100644 src/db/index/column/fts_column/tokenizer/stemmer_token_filter.h create mode 100644 src/db/index/column/fts_column/tokenizer/token_filter.cc create mode 100644 src/db/index/column/fts_column/tokenizer/token_filter.h create mode 100644 src/db/index/column/fts_column/tokenizer/tokenizer.h create mode 100644 src/db/index/column/fts_column/tokenizer/tokenizer_factory.cc create mode 100644 src/db/index/column/fts_column/tokenizer/tokenizer_factory.h create mode 100644 src/db/index/column/fts_column/tokenizer/tokenizer_pipeline_manager.cc create mode 100644 src/db/index/column/fts_column/tokenizer/tokenizer_pipeline_manager.h create mode 100644 src/db/index/column/fts_column/tokenizer/whitespace_tokenizer.cc create mode 100644 src/db/index/column/fts_column/tokenizer/whitespace_tokenizer.h create mode 100644 src/db/index/column/inverted_column/inverted_codec.h create mode 100644 src/db/index/column/inverted_column/inverted_column_indexer.h create mode 100644 src/db/index/column/inverted_column/inverted_column_indexer_search.cc create mode 100644 src/db/index/column/inverted_column/inverted_column_indexer_util.cc create mode 100644 src/db/index/column/inverted_column/inverted_column_indexer_write.cc create mode 100644 src/db/index/column/inverted_column/inverted_doc_range.h create mode 100644 src/db/index/column/inverted_column/inverted_indexer.cc create mode 100644 src/db/index/column/inverted_column/inverted_indexer.h create mode 100644 src/db/index/column/inverted_column/inverted_rocksdb_merger.h create mode 100644 src/db/index/column/inverted_column/inverted_search_result.h create mode 100644 src/db/index/column/vector_column/combined_vector_column_indexer.cc create mode 100644 src/db/index/column/vector_column/combined_vector_column_indexer.h create mode 100644 src/db/index/column/vector_column/engine_helper.hpp create mode 100644 src/db/index/column/vector_column/vector_column_indexer.cc create mode 100644 src/db/index/column/vector_column/vector_column_indexer.h create mode 100644 src/db/index/column/vector_column/vector_column_params.h create mode 100644 src/db/index/column/vector_column/vector_index_results.h create mode 100644 src/db/index/common/delete_store.h create mode 100644 src/db/index/common/doc.cc create mode 100644 src/db/index/common/id_map.cc create mode 100644 src/db/index/common/id_map.h create mode 100644 src/db/index/common/index_filter.h create mode 100644 src/db/index/common/index_params.cc create mode 100644 src/db/index/common/meta.h create mode 100644 src/db/index/common/proto_converter.cc create mode 100644 src/db/index/common/proto_converter.h create mode 100644 src/db/index/common/query.cc create mode 100644 src/db/index/common/schema.cc create mode 100644 src/db/index/common/stats.cc create mode 100644 src/db/index/common/type_helper.cc create mode 100644 src/db/index/common/type_helper.h create mode 100644 src/db/index/common/version_manager.cc create mode 100644 src/db/index/common/version_manager.h create mode 100644 src/db/index/segment/column_merging_reader.cc create mode 100644 src/db/index/segment/column_merging_reader.h create mode 100644 src/db/index/segment/segment.cc create mode 100644 src/db/index/segment/segment.h create mode 100644 src/db/index/segment/segment_helper.cc create mode 100644 src/db/index/segment/segment_helper.h create mode 100644 src/db/index/segment/segment_manager.cc create mode 100644 src/db/index/segment/segment_manager.h create mode 100644 src/db/index/segment/sql_expr_parser.cc create mode 100644 src/db/index/segment/sql_expr_parser.h create mode 100644 src/db/index/storage/arrow_ipc_writer.cc create mode 100644 src/db/index/storage/arrow_ipc_writer.h create mode 100644 src/db/index/storage/base_forward_store.h create mode 100644 src/db/index/storage/bufferpool_forward_store.cc create mode 100644 src/db/index/storage/bufferpool_forward_store.h create mode 100644 src/db/index/storage/chunked_file_writer.cc create mode 100644 src/db/index/storage/chunked_file_writer.h create mode 100644 src/db/index/storage/forward_writer.cc create mode 100644 src/db/index/storage/forward_writer.h create mode 100644 src/db/index/storage/lazy_record_batch_reader.h create mode 100644 src/db/index/storage/memory_forward_store.cc create mode 100644 src/db/index/storage/memory_forward_store.h create mode 100644 src/db/index/storage/mmap_forward_store.cc create mode 100644 src/db/index/storage/mmap_forward_store.h create mode 100644 src/db/index/storage/parquet_buffer_pool.cc create mode 100644 src/db/index/storage/parquet_buffer_pool.h create mode 100644 src/db/index/storage/parquet_writer.cc create mode 100644 src/db/index/storage/parquet_writer.h create mode 100644 src/db/index/storage/store_helper.h create mode 100644 src/db/index/storage/wal/local_wal_file.cc create mode 100644 src/db/index/storage/wal/local_wal_file.h create mode 100644 src/db/index/storage/wal/wal_file.cc create mode 100644 src/db/index/storage/wal/wal_file.h create mode 100644 src/db/proto/zvec.proto create mode 100644 src/db/reranker/reranker.cc create mode 100644 src/db/sqlengine/CMakeLists.txt create mode 100644 src/db/sqlengine/analyzer/query_analyzer.cc create mode 100644 src/db/sqlengine/analyzer/query_analyzer.h create mode 100644 src/db/sqlengine/analyzer/query_field_info.cc create mode 100644 src/db/sqlengine/analyzer/query_field_info.h create mode 100644 src/db/sqlengine/analyzer/query_info.cc create mode 100644 src/db/sqlengine/analyzer/query_info.h create mode 100644 src/db/sqlengine/analyzer/query_info_helper.cc create mode 100644 src/db/sqlengine/analyzer/query_info_helper.h create mode 100644 src/db/sqlengine/analyzer/query_node.cc create mode 100644 src/db/sqlengine/analyzer/query_node.h create mode 100644 src/db/sqlengine/analyzer/query_node_walker.cc create mode 100644 src/db/sqlengine/analyzer/query_node_walker.h create mode 100644 src/db/sqlengine/analyzer/query_orderby_info.cc create mode 100644 src/db/sqlengine/analyzer/query_orderby_info.h create mode 100644 src/db/sqlengine/analyzer/simple_rewriter.cc create mode 100644 src/db/sqlengine/analyzer/simple_rewriter.h create mode 100644 src/db/sqlengine/antlr/SQLLexer.g4 create mode 100644 src/db/sqlengine/antlr/SQLParser.g4 create mode 100644 src/db/sqlengine/antlr/gen/SQLLexer.cc create mode 100644 src/db/sqlengine/antlr/gen/SQLLexer.h create mode 100644 src/db/sqlengine/antlr/gen/SQLLexer.interp create mode 100644 src/db/sqlengine/antlr/gen/SQLLexer.tokens create mode 100644 src/db/sqlengine/antlr/gen/SQLParser.cc create mode 100644 src/db/sqlengine/antlr/gen/SQLParser.h create mode 100644 src/db/sqlengine/antlr/gen/SQLParser.interp create mode 100644 src/db/sqlengine/antlr/gen/SQLParser.tokens create mode 100644 src/db/sqlengine/antlr/gen/SQLParserBaseListener.cc create mode 100644 src/db/sqlengine/antlr/gen/SQLParserBaseListener.h create mode 100644 src/db/sqlengine/antlr/gen/SQLParserListener.cc create mode 100644 src/db/sqlengine/antlr/gen/SQLParserListener.h create mode 100755 src/db/sqlengine/antlr/gen_parser.sh create mode 100644 src/db/sqlengine/common/fts_cond_info.h create mode 100644 src/db/sqlengine/common/generic_node.h create mode 100644 src/db/sqlengine/common/group_by.h create mode 100644 src/db/sqlengine/common/util.cc create mode 100644 src/db/sqlengine/common/util.h create mode 100644 src/db/sqlengine/parser/base_info.h create mode 100644 src/db/sqlengine/parser/case_changing_charstream.h create mode 100644 src/db/sqlengine/parser/error_verbose_listener.h create mode 100644 src/db/sqlengine/parser/node.cc create mode 100644 src/db/sqlengine/parser/node.h create mode 100644 src/db/sqlengine/parser/orderby_elem_info.h create mode 100644 src/db/sqlengine/parser/query_parser.cc create mode 100644 src/db/sqlengine/parser/query_parser.h create mode 100644 src/db/sqlengine/parser/select_info.cc create mode 100644 src/db/sqlengine/parser/select_info.h create mode 100644 src/db/sqlengine/parser/selected_elem_info.cc create mode 100644 src/db/sqlengine/parser/selected_elem_info.h create mode 100644 src/db/sqlengine/parser/sql_info.cc create mode 100644 src/db/sqlengine/parser/sql_info.h create mode 100644 src/db/sqlengine/parser/sql_info_helper.cc create mode 100644 src/db/sqlengine/parser/sql_info_helper.h create mode 100644 src/db/sqlengine/parser/zvec_cached_sql_parser.cc create mode 100644 src/db/sqlengine/parser/zvec_cached_sql_parser.h create mode 100644 src/db/sqlengine/parser/zvec_parser.cc create mode 100644 src/db/sqlengine/parser/zvec_parser.h create mode 100644 src/db/sqlengine/parser/zvec_sql_parser.cc create mode 100644 src/db/sqlengine/parser/zvec_sql_parser.h create mode 100644 src/db/sqlengine/planner/doc_filter.cc create mode 100644 src/db/sqlengine/planner/doc_filter.h create mode 100644 src/db/sqlengine/planner/fts_recall_node.cc create mode 100644 src/db/sqlengine/planner/fts_recall_node.h create mode 100644 src/db/sqlengine/planner/invert_recall_node.cc create mode 100644 src/db/sqlengine/planner/invert_recall_node.h create mode 100644 src/db/sqlengine/planner/invert_search.cc create mode 100644 src/db/sqlengine/planner/invert_search.h create mode 100644 src/db/sqlengine/planner/op_register.cc create mode 100644 src/db/sqlengine/planner/op_register.h create mode 100644 src/db/sqlengine/planner/ops/check_not_filtered_op.cc create mode 100644 src/db/sqlengine/planner/ops/check_not_filtered_op.h create mode 100644 src/db/sqlengine/planner/ops/contain_op.cc create mode 100644 src/db/sqlengine/planner/ops/contain_op.h create mode 100644 src/db/sqlengine/planner/ops/fetch_vector_op.cc create mode 100644 src/db/sqlengine/planner/ops/fetch_vector_op.h create mode 100644 src/db/sqlengine/planner/optimizer.cc create mode 100644 src/db/sqlengine/planner/optimizer.h create mode 100644 src/db/sqlengine/planner/plan_info.cc create mode 100644 src/db/sqlengine/planner/plan_info.h create mode 100644 src/db/sqlengine/planner/query_planner.cc create mode 100644 src/db/sqlengine/planner/query_planner.h create mode 100644 src/db/sqlengine/planner/segment_node.cc create mode 100644 src/db/sqlengine/planner/segment_node.h create mode 100644 src/db/sqlengine/planner/vector_recall_node.cc create mode 100644 src/db/sqlengine/planner/vector_recall_node.h create mode 100644 src/db/sqlengine/sqlengine.cc create mode 100644 src/db/sqlengine/sqlengine.h create mode 100644 src/db/sqlengine/sqlengine_impl.cc create mode 100644 src/db/sqlengine/sqlengine_impl.h create mode 100644 src/include/zvec/ailego/buffer/block_eviction_queue.h create mode 100644 src/include/zvec/ailego/buffer/concurrentqueue.h create mode 100644 src/include/zvec/ailego/buffer/external_cache.h create mode 100644 src/include/zvec/ailego/buffer/vector_page_table.h create mode 100644 src/include/zvec/ailego/container/blob.h create mode 100644 src/include/zvec/ailego/container/cube.h create mode 100644 src/include/zvec/ailego/container/heap.h create mode 100644 src/include/zvec/ailego/container/hypercube.h create mode 100644 src/include/zvec/ailego/container/params.h create mode 100644 src/include/zvec/ailego/container/vector.h create mode 100644 src/include/zvec/ailego/encoding/json.h create mode 100644 src/include/zvec/ailego/encoding/json/mod_json.h create mode 100644 src/include/zvec/ailego/encoding/json/mod_json_plus.h create mode 100644 src/include/zvec/ailego/hash/crc32c.h create mode 100644 src/include/zvec/ailego/hash/jump_hash.h create mode 100644 src/include/zvec/ailego/internal/platform.h create mode 100644 src/include/zvec/ailego/io/file.h create mode 100644 src/include/zvec/ailego/io/mmap_file.h create mode 100644 src/include/zvec/ailego/logger/logger.h create mode 100644 src/include/zvec/ailego/math_batch/utils.h create mode 100644 src/include/zvec/ailego/parallel/thread_pool.h create mode 100644 src/include/zvec/ailego/parallel/thread_queue.h create mode 100644 src/include/zvec/ailego/pattern/closure.h create mode 100644 src/include/zvec/ailego/pattern/expected.hpp create mode 100644 src/include/zvec/ailego/pattern/factory.h create mode 100644 src/include/zvec/ailego/pattern/singleton.h create mode 100644 src/include/zvec/ailego/string/string_concat_helper.h create mode 100644 src/include/zvec/ailego/string/string_view.h create mode 100644 src/include/zvec/ailego/utility/file_helper.h create mode 100644 src/include/zvec/ailego/utility/float_helper.h create mode 100644 src/include/zvec/ailego/utility/string_helper.h create mode 100644 src/include/zvec/ailego/utility/string_helper_impl.h create mode 100644 src/include/zvec/ailego/utility/time_helper.h create mode 100644 src/include/zvec/ailego/utility/type_helper.h create mode 100644 src/include/zvec/c_api.h create mode 100644 src/include/zvec/core/framework/index_builder.h create mode 100644 src/include/zvec/core/framework/index_bundle.h create mode 100644 src/include/zvec/core/framework/index_cluster.h create mode 100644 src/include/zvec/core/framework/index_context.h create mode 100644 src/include/zvec/core/framework/index_converter.h create mode 100644 src/include/zvec/core/framework/index_document.h create mode 100644 src/include/zvec/core/framework/index_dumper.h create mode 100644 src/include/zvec/core/framework/index_error.h create mode 100644 src/include/zvec/core/framework/index_factory.h create mode 100644 src/include/zvec/core/framework/index_features.h create mode 100644 src/include/zvec/core/framework/index_filter.h create mode 100644 src/include/zvec/core/framework/index_flow.h create mode 100644 src/include/zvec/core/framework/index_format.h create mode 100644 src/include/zvec/core/framework/index_framework.h create mode 100644 src/include/zvec/core/framework/index_groupby.h create mode 100644 src/include/zvec/core/framework/index_helper.h create mode 100644 src/include/zvec/core/framework/index_holder.h create mode 100644 src/include/zvec/core/framework/index_logger.h create mode 100644 src/include/zvec/core/framework/index_mapping.h create mode 100644 src/include/zvec/core/framework/index_memory.h create mode 100644 src/include/zvec/core/framework/index_meta.h create mode 100644 src/include/zvec/core/framework/index_metric.h create mode 100644 src/include/zvec/core/framework/index_module.h create mode 100644 src/include/zvec/core/framework/index_packer.h create mode 100644 src/include/zvec/core/framework/index_plugin.h create mode 100644 src/include/zvec/core/framework/index_provider.h create mode 100644 src/include/zvec/core/framework/index_reducer.h create mode 100644 src/include/zvec/core/framework/index_refiner.h create mode 100644 src/include/zvec/core/framework/index_reformer.h create mode 100644 src/include/zvec/core/framework/index_runner.h create mode 100644 src/include/zvec/core/framework/index_searcher.h create mode 100644 src/include/zvec/core/framework/index_segment_storage.h create mode 100644 src/include/zvec/core/framework/index_stats.h create mode 100644 src/include/zvec/core/framework/index_storage.h create mode 100644 src/include/zvec/core/framework/index_streamer.h create mode 100644 src/include/zvec/core/framework/index_threads.h create mode 100644 src/include/zvec/core/framework/index_trainer.h create mode 100644 src/include/zvec/core/framework/index_unpacker.h create mode 100644 src/include/zvec/core/framework/index_version.h create mode 100644 src/include/zvec/core/interface/constants.h create mode 100644 src/include/zvec/core/interface/index.h create mode 100644 src/include/zvec/core/interface/index_factory.h create mode 100644 src/include/zvec/core/interface/index_param.h create mode 100644 src/include/zvec/core/interface/index_param_builders.h create mode 100644 src/include/zvec/core/interface/vector_source.h create mode 100644 src/include/zvec/db/collection.h create mode 100644 src/include/zvec/db/config.h create mode 100644 src/include/zvec/db/doc.h create mode 100644 src/include/zvec/db/index_params.h create mode 100644 src/include/zvec/db/options.h create mode 100644 src/include/zvec/db/query.h create mode 100644 src/include/zvec/db/query_params.h create mode 100644 src/include/zvec/db/reranker.h create mode 100644 src/include/zvec/db/schema.h create mode 100644 src/include/zvec/db/stats.h create mode 100644 src/include/zvec/db/status.h create mode 100644 src/include/zvec/db/type.h create mode 100644 src/include/zvec/plugin/diskann_plugin.h create mode 100644 src/include/zvec/turbo/turbo.h create mode 100644 src/include/zvec/version.h.in create mode 100644 src/turbo/CMakeLists.txt create mode 100644 src/turbo/avx512_vnni/record_quantized_int8/common.h create mode 100644 src/turbo/avx512_vnni/record_quantized_int8/cosine.cc create mode 100644 src/turbo/avx512_vnni/record_quantized_int8/cosine.h create mode 100644 src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc create mode 100644 src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.h create mode 100644 src/turbo/avx512_vnni/uniform_int8/quantize.cc create mode 100644 src/turbo/avx512_vnni/uniform_int8/quantize.h create mode 100644 src/turbo/avx512_vnni/uniform_int8/squared_euclidean.cc create mode 100644 src/turbo/avx512_vnni/uniform_int8/squared_euclidean.h create mode 100644 src/turbo/turbo.cc create mode 100644 tests/CMakeLists.txt create mode 100644 tests/ailego/CMakeLists.txt create mode 100644 tests/ailego/algorithm/integer_quantizer_test.cc create mode 100644 tests/ailego/algorithm/kmeans_test.cc create mode 100644 tests/ailego/container/bitmap_test.cc create mode 100644 tests/ailego/container/blob_test.cc create mode 100644 tests/ailego/container/bloom_filter_test.cc create mode 100644 tests/ailego/container/cube_test.cc create mode 100644 tests/ailego/container/heap_test.cc create mode 100644 tests/ailego/container/hypercube_test.cc create mode 100644 tests/ailego/container/params_test.cc create mode 100644 tests/ailego/container/reservoir_test.cc create mode 100644 tests/ailego/container/vector_array_test.cc create mode 100644 tests/ailego/container/vector_test.cc create mode 100644 tests/ailego/encoding/json_parse_test.cc create mode 100644 tests/ailego/hash/crc32c_test.cc create mode 100644 tests/ailego/hash/jump_hash_test.cc create mode 100644 tests/ailego/internal/cpu_features_test.cc create mode 100644 tests/ailego/io/file_lock_test.cc create mode 100644 tests/ailego/io/file_test.cc create mode 100644 tests/ailego/io/mmap_file_test.cc create mode 100644 tests/ailego/io/utf8_path_test.cc create mode 100644 tests/ailego/logger/logger_test.cc create mode 100644 tests/ailego/math/cosine_distance_matrix_fp16_test.cc create mode 100644 tests/ailego/math/cosine_distance_matrix_fp32_test.cc create mode 100644 tests/ailego/math/cosine_distance_matrix_int8_test.cc create mode 100644 tests/ailego/math/euclidean_distance_matrix_fp16_test.cc create mode 100644 tests/ailego/math/euclidean_distance_matrix_fp32_test.cc create mode 100644 tests/ailego/math/euclidean_distance_matrix_int4_test.cc create mode 100644 tests/ailego/math/euclidean_distance_matrix_int8_test.cc create mode 100644 tests/ailego/math/inner_product_matrix_fp16_test.cc create mode 100644 tests/ailego/math/inner_product_matrix_fp32_test.cc create mode 100644 tests/ailego/math/inner_product_matrix_int4_test.cc create mode 100644 tests/ailego/math/inner_product_matrix_int8_test.cc create mode 100644 tests/ailego/math/mips_euclidean_distance_matrix_fp16_test.cc create mode 100644 tests/ailego/math/mips_euclidean_distance_matrix_fp32_test.cc create mode 100644 tests/ailego/math/mips_euclidean_distance_matrix_int4_test.cc create mode 100644 tests/ailego/math/mips_euclidean_distance_matrix_int8_test.cc create mode 100644 tests/ailego/math/norm_matrix_fp16_test.cc create mode 100644 tests/ailego/math/norm_matrix_fp32_test.cc create mode 100644 tests/ailego/math/norm_matrix_int4_test.cc create mode 100644 tests/ailego/math/norm_matrix_int8_test.cc create mode 100644 tests/ailego/math/normalizer_test.cc create mode 100644 tests/ailego/parallel/lock_test.cc create mode 100644 tests/ailego/parallel/multi_thread_list_test.cc create mode 100644 tests/ailego/parallel/semaphore_test.cc create mode 100644 tests/ailego/parallel/thread_pool_test.cc create mode 100644 tests/ailego/parallel/thread_queue_test.cc create mode 100644 tests/ailego/pattern/closure_test.cc create mode 100644 tests/ailego/pattern/factory_test.cc create mode 100644 tests/ailego/pattern/scope_guard_test.cc create mode 100644 tests/ailego/pattern/singleton_test.cc create mode 100644 tests/ailego/utility/bit_string_helper_test.cc create mode 100644 tests/ailego/utility/bitset_helper_test.cc create mode 100644 tests/ailego/utility/dl_helper_test.cc create mode 100644 tests/ailego/utility/float_helper_test.cc create mode 100644 tests/ailego/utility/matrix_helper_test.cc create mode 100644 tests/ailego/utility/memory_helper_test.cc create mode 100644 tests/ailego/utility/string_helper_test.cc create mode 100644 tests/ailego/utility/time_helper_test.cc create mode 100644 tests/ailego/utility/type_helper_test.cc create mode 100644 tests/ailego/version_test.cc create mode 100644 tests/android_gmock_main.cc create mode 100644 tests/android_gtest_main.cc create mode 100644 tests/c/CMakeLists.txt create mode 100644 tests/c/c_api_test.c create mode 100644 tests/c/utils.c create mode 100644 tests/c/utils.h create mode 100644 tests/core/CMakeLists.txt create mode 100644 tests/core/algorithm/CMakeLists.txt create mode 100644 tests/core/algorithm/cluster/CMakeLists.txt create mode 100644 tests/core/algorithm/cluster/kmeans_cluster_test.cc create mode 100644 tests/core/algorithm/cluster/multi_chunk_cluster_test.cc create mode 100644 tests/core/algorithm/cluster/opt_kmeans_cluster_test.cc create mode 100644 tests/core/algorithm/diskann/CMakeLists.txt create mode 100644 tests/core/algorithm/diskann/diskann_builder_test.cc create mode 100644 tests/core/algorithm/diskann/diskann_searcher_test.cc create mode 100644 tests/core/algorithm/flat/CMakeLists.txt create mode 100644 tests/core/algorithm/flat/flat_builder_test.cc create mode 100644 tests/core/algorithm/flat/flat_searcher_test.cpp create mode 100644 tests/core/algorithm/flat/flat_streamer_buffer_test.cc create mode 100644 tests/core/algorithm/flat/flat_streamer_buffer_time_test.cc create mode 100644 tests/core/algorithm/flat/flat_streamer_buffer_time_test.cpp create mode 100644 tests/core/algorithm/flat/flat_streamer_test.cc create mode 100644 tests/core/algorithm/flat_sparse/CMakeLists.txt create mode 100644 tests/core/algorithm/flat_sparse/flat_sparse_builder_test.cc create mode 100644 tests/core/algorithm/flat_sparse/flat_sparse_searcher_test.cc create mode 100644 tests/core/algorithm/flat_sparse/flat_sparse_streamer_buffer_test.cc create mode 100644 tests/core/algorithm/flat_sparse/flat_sparse_streamer_test.cc create mode 100644 tests/core/algorithm/hnsw/CMakeLists.txt create mode 100644 tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc create mode 100644 tests/core/algorithm/hnsw/hnsw_streamer_test.cc create mode 100644 tests/core/algorithm/hnsw_rabitq/CMakeLists.txt create mode 100644 tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_test.cc create mode 100644 tests/core/algorithm/hnsw_sparse/CMakeLists.txt create mode 100644 tests/core/algorithm/hnsw_sparse/hnsw_sparse_builder_test.cc create mode 100644 tests/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_test.cpp create mode 100644 tests/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_buffer_test.cpp create mode 100644 tests/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_test.cc create mode 100644 tests/core/algorithm/ivf/CMakeLists.txt create mode 100644 tests/core/algorithm/ivf/ivf_builder_test.cc create mode 100644 tests/core/algorithm/ivf/ivf_searcher_test.cc create mode 100644 tests/core/algorithm/vamana/CMakeLists.txt create mode 100644 tests/core/algorithm/vamana/vamana_streamer_test.cc create mode 100644 tests/core/framework/CMakeLists.txt create mode 100644 tests/core/interface/CMakeLists.txt create mode 100644 tests/core/interface/index_group_by_test.cc create mode 100644 tests/core/interface/index_interface_test.cc create mode 100644 tests/core/metric/CMakeLists.txt create mode 100644 tests/core/metric/cosine_metric_test.cc create mode 100644 tests/core/metric/euclidean_metric_test.cc create mode 100644 tests/core/metric/inner_product_metric_test.cc create mode 100644 tests/core/metric/quantized_integer_metric_test.cc create mode 100644 tests/core/quantizer/CMakeLists.txt create mode 100644 tests/core/quantizer/half_float_reformer_test.cc create mode 100644 tests/core/quantizer/integer_quantizer_reformer_test.cc create mode 100644 tests/core/quantizer/uniform_int8_reformer_test.cc create mode 100644 tests/core/utility/CMakeLists.txt create mode 100644 tests/core/utility/buffer_storage_test.cpp create mode 100644 tests/core/utility/buffer_storage_write_test.cc create mode 100644 tests/core/utility/file_dumper_test.cc create mode 100644 tests/core/utility/memory_dumper_test.cc create mode 100644 tests/core/utility/mmap_file_container_test.cc create mode 100644 tests/core/utility/mmap_file_storage_test.cpp create mode 100644 tests/db/CMakeLists.txt create mode 100644 tests/db/collection_test.cc create mode 100644 tests/db/common/CMakeLists.txt create mode 100644 tests/db/common/config_test.cc create mode 100644 tests/db/common/status_test.cc create mode 100644 tests/db/crash_recovery/CMakeLists.txt create mode 100644 tests/db/crash_recovery/collection_optimizer.cc create mode 100644 tests/db/crash_recovery/data_generator.cc create mode 100644 tests/db/crash_recovery/optimize_recovery_test.cc create mode 100644 tests/db/crash_recovery/utility.h create mode 100644 tests/db/crash_recovery/write_recovery_test.cc create mode 100644 tests/db/fts_query_test.cc create mode 100644 tests/db/index/CMakeLists.txt create mode 100644 tests/db/index/column/fts_column/ascii_folding_token_filter_test.cc create mode 100644 tests/db/index/column/fts_column/bitpacked_posting_list_test.cc create mode 100644 tests/db/index/column/fts_column/fts_ast_rewriter_test.cc create mode 100644 tests/db/index/column/fts_column/fts_candidate_iterator_test.cc create mode 100644 tests/db/index/column/fts_column/fts_column_indexer_test.cc create mode 100644 tests/db/index/column/fts_column/fts_conjunction_opt_test.cc create mode 100644 tests/db/index/column/fts_column/fts_rocksdb_reducer_test.cc create mode 100644 tests/db/index/column/fts_column/lowercase_token_filter_test.cc create mode 100644 tests/db/index/column/fts_column/standard_tokenizer_test.cc create mode 100644 tests/db/index/column/fts_column/stemmer_token_filter_test.cc create mode 100644 tests/db/index/column/fts_column/tokenizer_pipeline_manager_test.cc create mode 100644 tests/db/index/column/inverted_column/inverted_column_indexer_array_numbers_test.cc create mode 100644 tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc create mode 100644 tests/db/index/column/inverted_column/inverted_column_indexer_cyclic_numbers_test.cc create mode 100644 tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc create mode 100644 tests/db/index/column/inverted_column/inverted_column_indexer_string_test.cc create mode 100644 tests/db/index/column/inverted_column/inverted_indexer_util_test.cc create mode 100644 tests/db/index/column/vector_column_group_by_test.cc create mode 100644 tests/db/index/column/vector_column_indexer_test.cc create mode 100644 tests/db/index/common/db_proto_converter_test.cc create mode 100644 tests/db/index/common/db_type_helper_test.cc create mode 100644 tests/db/index/common/doc_test.cc create mode 100644 tests/db/index/common/index_params_test.cc create mode 100644 tests/db/index/common/meta_test.cc create mode 100644 tests/db/index/common/query_params_test.cc create mode 100644 tests/db/index/common/schema_test.cc create mode 100644 tests/db/index/common/version_manager_test.cc create mode 100644 tests/db/index/segment/column_merging_reader_test.cc create mode 100644 tests/db/index/segment/segment_helper_test.cc create mode 100644 tests/db/index/segment/segment_row_id_test.cc create mode 100644 tests/db/index/segment/segment_test.cc create mode 100644 tests/db/index/segment/segment_test_fixture.h create mode 100644 tests/db/index/segment/sql_expr_parser_test.cc create mode 100644 tests/db/index/segment/sql_expr_validator_test.cc create mode 100644 tests/db/index/storage/arrow_ipc_writer_test.cc create mode 100644 tests/db/index/storage/bufferpool_store_test.cc create mode 100644 tests/db/index/storage/mem_store_test.cc create mode 100644 tests/db/index/storage/mmap_store_test.cc create mode 100644 tests/db/index/storage/parquet_writer_test.cc create mode 100644 tests/db/index/storage/wal_file_test.cc create mode 100644 tests/db/index/utils/utils.cc create mode 100644 tests/db/index/utils/utils.h create mode 100644 tests/db/reranker_test.cc create mode 100644 tests/db/sqlengine/CMakeLists.txt create mode 100644 tests/db/sqlengine/contain_test.cc create mode 100644 tests/db/sqlengine/forward_recall_test.cc create mode 100644 tests/db/sqlengine/fts_multi_segment_test.cc create mode 100644 tests/db/sqlengine/fts_parser_test.cc create mode 100644 tests/db/sqlengine/fts_recall_test.cc create mode 100644 tests/db/sqlengine/invert_recall_test.cc create mode 100644 tests/db/sqlengine/like_test.cc create mode 100644 tests/db/sqlengine/mock_segment.h create mode 100644 tests/db/sqlengine/optimizer_test.cc create mode 100644 tests/db/sqlengine/query_info_test.cc create mode 100644 tests/db/sqlengine/recall_base.h create mode 100644 tests/db/sqlengine/simple_rewriter_test.cc create mode 100644 tests/db/sqlengine/sqlengine_test.cc create mode 100644 tests/db/sqlengine/test_helper.h create mode 100644 tests/db/sqlengine/vector_recall_test.cc create mode 100644 tests/db/utf8_collection_test.cc create mode 100644 tests/ios_test_sandbox.cc create mode 100644 tests/test_util.h create mode 100644 thirdparty/CMakeLists.txt create mode 100644 thirdparty/CRoaring/CMakeLists.txt create mode 100644 thirdparty/CRoaring/CRoaring-2.0.4/LICENSE create mode 100644 thirdparty/CRoaring/CRoaring-2.0.4/include/roaring/roaring.h create mode 100644 thirdparty/CRoaring/CRoaring-2.0.4/roaring.c create mode 100644 thirdparty/CRoaring/CRoaring-2.0.4/roaring.h create mode 100644 thirdparty/CRoaring/CRoaring-2.0.4/roaring.hh create mode 100644 thirdparty/CRoaring/CRoaring-2.0.4/roaring64map.hh create mode 100644 thirdparty/FastPFOR/CMakeLists.txt create mode 100644 thirdparty/FastPFOR/fastpfor.riscv.patch create mode 100644 thirdparty/RaBitQ-Library/CMakeLists.txt create mode 100644 thirdparty/antlr/CMakeLists.txt create mode 100644 thirdparty/antlr/antlr4.patch create mode 100644 thirdparty/antlr/antlr4.windows.patch create mode 100644 thirdparty/arrow/CMakeLists.txt create mode 100644 thirdparty/arrow/arrow.android.patch create mode 100644 thirdparty/arrow/arrow.ios.patch create mode 100644 thirdparty/arrow/arrow.patch create mode 100644 thirdparty/arrow/arrow.windows.patch create mode 100644 thirdparty/cppjieba/CMakeLists.txt create mode 100644 thirdparty/gflags/CMakeLists.txt create mode 100644 thirdparty/glog/CMakeLists.txt create mode 100644 thirdparty/glog/glog.android.patch create mode 100644 thirdparty/glog/glog.patch create mode 100644 thirdparty/googletest/CMakeLists.txt create mode 100644 thirdparty/googletest/googletest.windows.patch create mode 100644 thirdparty/limonp/CMakeLists.txt create mode 100644 thirdparty/lz4/CMakeLists.txt create mode 100644 thirdparty/magic_enum/CMakeLists.txt create mode 100644 thirdparty/protobuf/CMakeLists.txt create mode 100644 thirdparty/protobuf/protobuf.windows.patch create mode 100644 thirdparty/rocksdb/CMakeLists.txt create mode 100644 thirdparty/rocksdb/rocksdb.android.patch create mode 100644 thirdparty/rocksdb/rocksdb.ios.patch create mode 100644 thirdparty/rocksdb/rocksdb.windows.patch create mode 100644 thirdparty/snowball/CMakeLists.txt create mode 100644 thirdparty/sparsehash/CMakeLists.txt create mode 100644 thirdparty/sparsehash/sparseconfig.h create mode 100644 thirdparty/sparsehash/sparsehash.windows.patch create mode 100644 thirdparty/utf8proc/CMakeLists.txt create mode 100644 thirdparty/yaml-cpp/CMakeLists.txt create mode 100644 tools/CMakeLists.txt create mode 100644 tools/core/CMakeLists.txt create mode 100644 tools/core/README.md create mode 100644 tools/core/bench.cc create mode 100644 tools/core/bench_original.cc create mode 100644 tools/core/bench_result.h create mode 100644 tools/core/convert_cohere_parquet.py create mode 100644 tools/core/filter_result_cache.h create mode 100644 tools/core/flow.h create mode 100644 tools/core/helper.h create mode 100644 tools/core/index_meta_helper.h create mode 100644 tools/core/local_builder.cc create mode 100644 tools/core/local_builder_original.cc create mode 100644 tools/core/meta_segment_common.h create mode 100644 tools/core/recall.cc create mode 100644 tools/core/recall_original.cc create mode 100644 tools/core/txt2vecs.cc create mode 100644 tools/core/txt_input_reader.h create mode 100644 tools/core/vecs_common.h create mode 100644 tools/core/vecs_index_holder.h create mode 100644 tools/core/vecs_reader.h create mode 100644 tools/db/CMakeLists.txt create mode 100644 tools/db/fts_bench_main.cc diff --git a/.clang-format b/.clang-format new file mode 100644 index 0000000..2574e7f --- /dev/null +++ b/.clang-format @@ -0,0 +1,31 @@ +# http://clang.llvm.org/docs/ClangFormatStyleOptions.html +# Defines the Google C++ style for automatic reformatting. +BasedOnStyle: Google +MaxEmptyLinesToKeep: 2 +DerivePointerAlignment: false +PointerAlignment: Right +AllowShortFunctionsOnASingleLine: Empty +IncludeBlocks: Merge +IncludeCategories: + - Regex: '^$' + Priority: 200 + - Regex: '^<[0-9A-Za-z_]+>$' + Priority: 201 + - Regex: '^<[0-9A-Za-z_]+\.[0-9A-Za-z]+>$' + Priority: 202 + - Regex: '^<[0-9A-Za-z_]+/[0-9A-Za-z]+' + Priority: 203 + - Regex: '^\"[0-9A-Za-z_]+/[0-9A-Za-z]+' + Priority: 300 + - Regex: '^\"[0-9A-Za-z_]+\.[0-9A-Za-z]+\"$' + Priority: 301 + - Regex: '.*' + Priority: 1000 diff --git a/.clang-tidy b/.clang-tidy new file mode 100644 index 0000000..929a369 --- /dev/null +++ b/.clang-tidy @@ -0,0 +1,10 @@ +Checks: > + -*, + bugprone-multiple-statement-macro, + bugprone-sizeof-expression, + modernize-use-nullptr, + modernize-use-override, +WarningsAsErrors: "*" +HeaderFilterRegex: "^(src|tests|tools)/(?!db/sqlengine/antlr/gen/|db/index/column/fts_column/gen/|include/zvec/ailego/encoding/json/mod_json\\.h).*" +FormatStyle: none +SystemHeaders: false diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..bb17898 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,3 @@ +# Auto-generated files — collapsed in GitHub PR diffs +src/db/index/column/fts_column/gen/** linguist-generated=true +src/db/sqlengine/antlr/gen/** linguist-generated=true diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000..97671f0 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,38 @@ +# CODEOWNERS - Auto-assign reviewers based on file paths +# Documentation: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners + +# Default: feihongxu0824 reviews everything not covered below +* @feihongxu0824 + +# .github directory +/.github/ @Cuiyus +/python/ @Cuiyus + +# cmake, examples, scripts +/cmake/ @egolearner +/examples/ @egolearner +/scripts/ @egolearner + +# src/db, tests +/src/db/ @zhourrr +/tests/db @zhourrr + +# src/core, src/turbo +/src/core/ @richyreachy +/src/turbo/ @richyreachy + +# src/ailego +/src/ailego/ @iaojnh +/tests/ailego @iaojnh +/tests/core @iaojnh + +# tools +/tools/ @JalinWang + +# src/binding +/src/binding/ @Cuiyus + +# src/include, thirdparty, .gitmodules +/src/include/ @chinaux +/thirdparty/ @chinaux +.gitmodules @chinaux diff --git a/.github/ISSUE_TEMPLATE/benchmark.yml b/.github/ISSUE_TEMPLATE/benchmark.yml new file mode 100644 index 0000000..0755be8 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/benchmark.yml @@ -0,0 +1,49 @@ +name: Benchmarking +description: Add, update, or fix benchmark cases for zvec +title: "[Benchmark]: " +labels: ["benchmark"] +body: + - type: markdown + attributes: + value: | + Use this for benchmark-related work: new test cases, CI integration, or performance regression tracking. + + - type: input + id: benchmark_type + attributes: + label: Benchmark Type + description: e.g., filtered search, batch insert, recall@k, ARM64 vs x86 + validations: + required: true + + - type: textarea + id: goal + attributes: + label: Goal + description: What performance aspect are you measuring or improving? + validations: + required: true + + - type: textarea + id: methodology + attributes: + label: Methodology + description: Dataset, query size, hardware, metrics (latency, throughput, memory) + validations: + required: true + + - type: textarea + id: baseline + attributes: + label: Baseline (if applicable) + description: Current performance numbers or competing systems for comparison. + validations: + required: false + + - type: textarea + id: ci_integration + attributes: + label: CI Integration Plan + description: Should this run in CI? How often? + validations: + required: false \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml new file mode 100644 index 0000000..8a8448e --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -0,0 +1,73 @@ +name: Bug Report +description: Report a bug or unexpected behavior (e.g., crash, incorrect vector query, memory leak) +title: "[Bug]: " +labels: ["bug", "triage"] +body: + - type: markdown + attributes: + value: | + Thank you for reporting! Please provide detailed info so we can reproduce and fix it quickly. + + - type: textarea + id: description + attributes: + label: Description + description: What happened? What did you expect? + placeholder: | + e.g. "Query with vector field crashes when using Zvec Python API" + validations: + required: true + + - type: textarea + id: steps_to_reproduce + attributes: + label: Steps to Reproduce + description: Exact steps to trigger the issue (code snippets welcome) + placeholder: | + 1. Build Zvec with CMake (Debug/Release) + 2. Run Python script: `python test.py` + 3. Call `collection.query(VectorQuery())` + 4. Process segfaults / hangs / returns wrong results + render: python + validations: + required: true + + - type: textarea + id: logs_or_trace + attributes: + label: Logs / Stack Trace + description: Paste relevant logs, LLDB/GDB backtrace, or CI failures + placeholder: | + Thread 1 "python" received signal SIGSEGV, Segmentation fault. + 0x0000000104a2c3f0 in std::__1::shared_ptr<...>::... + render: shell + validations: + required: false + + - type: input + id: os + attributes: + label: Operating System + placeholder: macOS 14 (M1), Ubuntu 22.04, Windows 11 (WSL2) + validations: + required: true + + - type: input + id: build_env + attributes: + label: Build & Runtime Environment + description: Compiler, CMake, Python, key dependencies + placeholder: | + clang 15.0.0, CMake 4.1.2, Python 3.11.9, magic_enum v0.9.7 (via git submodule) + validations: + required: true + + - type: checkboxes + id: additional_context + attributes: + label: Additional Context + options: + - label: I've checked `git status` — no uncommitted submodule changes + - label: I built with `CMAKE_BUILD_TYPE=Debug` + - label: This occurs with or without `COVERAGE=ON` + - label: The issue involves Python ↔ C++ integration (pybind11) \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000..6a2d7c5 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,9 @@ +blank_issues_enabled: false +contact_links: + - name: Documentation + url: https://zvec.org/en/ + about: Check the quickstart, build guide, and API docs first. + + - name: Python API Examples + url: https://zvec.org/en/docs/quickstart/ + about: See working usage examples. diff --git a/.github/ISSUE_TEMPLATE/enhancement.yml b/.github/ISSUE_TEMPLATE/enhancement.yml new file mode 100644 index 0000000..6c8af7d --- /dev/null +++ b/.github/ISSUE_TEMPLATE/enhancement.yml @@ -0,0 +1,41 @@ +name: Enhancement +description: Improve an existing feature or component +title: "[Enhance]: " +labels: ["enhancement"] +body: + - type: markdown + attributes: + value: | + This template is for improving existing functionality (e.g., performance, usability, robustness). + + - type: input + id: component + attributes: + label: Affected Component + description: e.g., HNSW index, buffer manager, Python API + validations: + required: true + + - type: textarea + id: current + attributes: + label: Current Behavior + description: What is the current state and its limitations? + validations: + required: true + + - type: textarea + id: desired + attributes: + label: Desired Improvement + description: What should be improved and how? + validations: + required: true + + - type: textarea + id: impact + attributes: + label: Impact + description: How will this benefit users? (e.g., faster queries, lower memory, easier integration) + validations: + required: true \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml new file mode 100644 index 0000000..1b97e50 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -0,0 +1,55 @@ +name: Feature Request +description: Suggest a new feature or improvement (e.g., better memory control, new query option) +title: "[Feature]: " +labels: ["feature"] +body: + - type: markdown + attributes: + value: | + Thanks for your idea! Help us understand the motivation and scope. + + - type: textarea + id: problem_or_motivation + attributes: + label: Problem / Motivation + description: What problem does this solve? Why is it needed? + placeholder: | + e.g. "Current vector queries don't allow filtering by metadata + distance threshold at once" + validations: + required: true + + - type: textarea + id: proposed_solution + attributes: + label: Proposed Solution + description: How should it work? API sketch or pseudocode welcome. + placeholder: | + Add `filter=` and `max_distance=` args to `Zvec.query()`: + ```python + results = db.query(vector, filter="category == 'A'", max_distance=0.5) + ``` + render: python + validations: + required: false + + - type: textarea + id: alternatives + attributes: + label: Alternatives Considered + description: Are there workarounds? Why not use them? + validations: + required: false + + - type: dropdown + id: impact_area + attributes: + label: Affected Area + multiple: true + options: + - label: C++ Core (storage, indexing) + - label: Python API / Bindings + - label: Build System (CMake, Homebrew pkg) + - label: Testing / CI / Coverage + - label: Documentation + validations: + required: false \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/integration.yml b/.github/ISSUE_TEMPLATE/integration.yml new file mode 100644 index 0000000..8ed39f3 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/integration.yml @@ -0,0 +1,36 @@ +name: Ecosystem Integration +description: Integrate zvec with external frameworks (e.g., LangChain, LlamaIndex) +title: "[Integration]: " +labels: ["integration"] +body: + - type: input + id: framework + attributes: + label: Target Framework + description: e.g., LangChain, LlamaIndex, Haystack + validations: + required: true + + - type: textarea + id: motivation + attributes: + label: Motivation + description: Why integrate with this framework? Who benefits? + validations: + required: true + + - type: textarea + id: interface + attributes: + label: Required Interface + description: What adapter or interface must be implemented? (e.g., VectorStore base class) + validations: + required: true + + - type: textarea + id: reference + attributes: + label: Reference Implementations + description: Links to similar integrations in other vector DBs. + validations: + required: false \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/profiling.yml b/.github/ISSUE_TEMPLATE/profiling.yml new file mode 100644 index 0000000..f881828 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/profiling.yml @@ -0,0 +1,41 @@ +name: Profiling / Investigation +description: Profile performance, compatibility, or behavior in a specific scenario +title: "[Profile]: " +labels: ["profile"] +body: + - type: markdown + attributes: + value: | + Use this for tasks like performance profiling, architecture compatibility checks, or feasibility studies. + + - type: input + id: scenario + attributes: + label: Target Scenario + description: e.g., ARM64 deployment, high-concurrency load, large dataset ingestion + validations: + required: true + + - type: textarea + id: objective + attributes: + label: Objective + description: What do you want to learn or validate? + validations: + required: true + + - type: textarea + id: methodology + attributes: + label: Proposed Methodology + description: How will you conduct the investigation? (tools, metrics, test data) + validations: + required: true + + - type: textarea + id: expected_outcome + attributes: + label: Expected Outcome + description: What deliverables are expected? (e.g., report, optimization PR, benchmark results) + validations: + required: true \ No newline at end of file diff --git a/.github/auto-assign-config.yml b/.github/auto-assign-config.yml new file mode 100644 index 0000000..5f76f4a --- /dev/null +++ b/.github/auto-assign-config.yml @@ -0,0 +1,24 @@ +# Auto-assign configuration +# Documentation: https://github.com/kentaro-m/auto-assign-action +# +# NOTE: Reviewers are assigned via CODEOWNERS file based on file paths. +# This config only handles assignee assignment. + +# Automatically add PR author as assignee +addAssignees: author + +# Number of assignees to add (0 = add all from the list) +numberOfAssignees: 0 + +# Reviewers are handled by CODEOWNERS, not here +addReviewers: false + +# Skip draft PRs +skipDraft: true + +# Skip keywords in PR title (won't assign if title contains these) +skipKeywords: + - wip + - WIP + - draft + - DRAFT diff --git a/.github/cmake/subproject-integration/CMakeLists.txt b/.github/cmake/subproject-integration/CMakeLists.txt new file mode 100644 index 0000000..a6cb349 --- /dev/null +++ b/.github/cmake/subproject-integration/CMakeLists.txt @@ -0,0 +1,26 @@ +cmake_minimum_required(VERSION 3.13) +project(zvec_subproject_integration LANGUAGES C CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +set(ZVEC_SOURCE_DIR "" CACHE PATH "Path to the zvec source tree") +if(NOT ZVEC_SOURCE_DIR) + message(FATAL_ERROR "ZVEC_SOURCE_DIR is required") +endif() + +add_subdirectory("${ZVEC_SOURCE_DIR}" zvec) + +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin") + +add_executable(zvec_embed_db_example + "${ZVEC_SOURCE_DIR}/examples/c++/db/main.cc") +target_link_libraries(zvec_embed_db_example PRIVATE zvec_shared) + +add_executable(zvec_embed_core_example + "${ZVEC_SOURCE_DIR}/examples/c++/core/main.cc") +target_link_libraries(zvec_embed_core_example PRIVATE zvec_core_shared) + +add_executable(zvec_embed_ailego_example + "${ZVEC_SOURCE_DIR}/examples/c++/ailego/main.cc") +target_link_libraries(zvec_embed_ailego_example PRIVATE zvec_ailego_shared) diff --git a/.github/codecov.yml b/.github/codecov.yml new file mode 100644 index 0000000..10ccdea --- /dev/null +++ b/.github/codecov.yml @@ -0,0 +1,32 @@ +codecov: + require_ci_to_pass: true + +coverage: + precision: 2 + round: down + range: "60...75" + + status: + project: + default: false + patch: + default: false + + +parsers: + gcov: + branch_detection: + conditional: true + loop: true + method: false + macro: false + +comment: + require_changes: false + layout: "reach,diff,flags,tree" + behavior: default + + +ignore: + - "thirdparty/" + - "tests/" \ No newline at end of file diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..30c1e44 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,17 @@ +version: 2 +updates: + # GitHub Actions dependencies + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + day: "monday" + time: "02:00" + timezone: "Asia/Shanghai" + labels: + - "dependencies" + - "github-actions" + commit-message: + prefix: "ci" + include: "scope" + open-pull-requests-limit: 5 diff --git a/.github/workflows/01-ci-pipeline.yml b/.github/workflows/01-ci-pipeline.yml new file mode 100644 index 0000000..ce9f8e2 --- /dev/null +++ b/.github/workflows/01-ci-pipeline.yml @@ -0,0 +1,116 @@ +name: Main + +on: + push: + branches: [ "main" ] + paths-ignore: &paths-ignore + - '**.md' + - '.github/CODEOWNERS' + - '.github/ISSUE_TEMPLATE/**' + - '.github/auto-assign-config.yml' + - '.github/codecov.yml' + - '.github/dependabot.yml' + - '.github/workflows/auto-assign.yml' + - '.github/workflows/community-pr-handler.yml' + - '.github/workflows/issue-auto-assign.yml' + - '.github/workflows/build_wheel.yml' + - '.github/workflows/build_test_wheel.yml' + - '.github/workflows/_build_wheel_job.yml' + - '.github/workflows/continuous_bench.yml' + - '.github/workflows/nightly_coverage.yml' + - '.github/workflows/07-linux-riscv-build.yml' + - '.github/workflows/docker/**' + - '.github/workflows/scripts/**' + merge_group: + pull_request: + branches: [ "main" ] + paths-ignore: *paths-ignore + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref || '' }}-${{ github.base_ref || '' }}-${{ github.ref != 'refs/heads/main' || github.sha }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + # Code quality checks (fast, run first) + lint: + if: >- + github.event_name != 'pull_request' || + (!github.event.pull_request.draft && + !contains(github.event.pull_request.title, 'wip')) + uses: ./.github/workflows/02-lint-check.yml + + # Static analysis: clang-tidy on changed C/C++ files + clang-tidy: + if: >- + github.event_name != 'pull_request' || + (!github.event.pull_request.draft && + !contains(github.event.pull_request.title, 'wip')) + uses: ./.github/workflows/clang_tidy.yml + + # Main build and test matrix + build-and-test-macos-arm64: + if: github.event_name != 'push' || github.ref != 'refs/heads/main' + name: Build & Test (macos-arm64) + needs: [lint, clang-tidy] + uses: ./.github/workflows/03-macos-linux-build.yml + with: + platform: macos-arm64 + os: macos-15 + + build-and-test-macos-26-arm64: + if: github.event_name != 'push' || github.ref != 'refs/heads/main' + name: Build & Test (macos-26-arm64) + needs: [lint, clang-tidy] + uses: ./.github/workflows/03-macos-linux-build.yml + with: + platform: macos-arm64 + os: macos-26 + + build-and-test-linux-arm64: + if: github.event_name != 'push' || github.ref != 'refs/heads/main' + name: Build & Test (linux-arm64) + needs: [lint, clang-tidy] + uses: ./.github/workflows/03-macos-linux-build.yml + with: + platform: linux-arm64 + os: ubuntu-24.04-arm + + build-and-test-linux-x64: + name: Build & Test (linux-x64) + needs: [lint, clang-tidy] + uses: ./.github/workflows/03-macos-linux-build.yml + with: + platform: linux-x64 + os: ubuntu-24.04 + + build-and-test-linux-x64-clang: + if: github.event_name != 'push' || github.ref != 'refs/heads/main' + name: Build & Test (linux-x64-clang) + needs: [lint, clang-tidy] + uses: ./.github/workflows/03-macos-linux-build.yml + with: + platform: linux-x64-clang + os: ubuntu-24.04 + compiler: clang + + build-android: + if: github.event_name != 'push' || github.ref != 'refs/heads/main' + name: Build & Test (android) + needs: [lint, clang-tidy] + uses: ./.github/workflows/04-android-build.yml + + build-and-test-on-windows: + if: github.event_name != 'push' || github.ref != 'refs/heads/main' + name: Build & Test (Windows) + needs: [lint, clang-tidy] + uses: ./.github/workflows/05-windows-build.yml + + build-ios: + if: github.event_name != 'push' || github.ref != 'refs/heads/main' + name: Build & Test (iOS) + needs: [lint, clang-tidy] + uses: ./.github/workflows/06-ios-build.yml diff --git a/.github/workflows/02-lint-check.yml b/.github/workflows/02-lint-check.yml new file mode 100644 index 0000000..038d73e --- /dev/null +++ b/.github/workflows/02-lint-check.yml @@ -0,0 +1,52 @@ +name: Lint + +on: + workflow_call: + +jobs: + lint: + name: Code Quality Checks + runs-on: ubuntu-24.04 + steps: + - name: Checkout code + uses: actions/checkout@v7 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.10' + cache: 'pip' + cache-dependency-path: 'pyproject.toml' + + - name: Install linting tools + run: | + python -m pip install --upgrade pip \ + ruff==v0.14.4 \ + clang-format==18.1.8 + shell: bash + + - name: Run Ruff Linter + run: python -m ruff check . + shell: bash + + - name: Run Ruff Formatter Check + run: python -m ruff format --check . + shell: bash + + - name: Run clang-format Check + run: | + CPP_FILES=$(find . -type f \( -name "*.cpp" -o -name "*.h" -o -name "*.hpp" -o -name "*.cc" -o -name "*.cxx" \) \ + ! -path "./build/*" \ + ! -path "./tests/*" \ + ! -path "./scripts/*" \ + ! -path "./python/*" \ + ! -path "./thirdparty/*" \ + ! -path "./.git/*") + + if [ -z "$CPP_FILES" ]; then + echo "No C++ files found to check." + exit 0 + fi + + clang-format --dry-run --Werror $CPP_FILES + shell: bash diff --git a/.github/workflows/03-macos-linux-build.yml b/.github/workflows/03-macos-linux-build.yml new file mode 100644 index 0000000..4e4f162 --- /dev/null +++ b/.github/workflows/03-macos-linux-build.yml @@ -0,0 +1,185 @@ +name: MacOS & Linux Build + +on: + workflow_call: + inputs: + platform: + description: 'Platform identifier' + required: true + type: string + os: + description: 'GitHub Actions runner OS' + required: true + type: string + compiler: + description: 'C++ compiler to use (gcc or clang)' + required: false + type: string + default: '' + +permissions: + contents: read + +jobs: + # Build and test matrix (parallel execution) + build-and-test: + name: Build & Test (${{ inputs.platform }}) + runs-on: ${{ inputs.os }} + + strategy: + fail-fast: false + matrix: + include: + - os: ${{ inputs.os }} + platform: ${{ inputs.platform }} + arch_flag: "" # Use appropriate architecture + + steps: + - name: Checkout code + uses: actions/checkout@v7 + with: + submodules: recursive + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: ${{ inputs.platform }}-${{ inputs.os }}-${{ inputs.compiler || 'gcc' }} + max-size: 150M + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.10' + cache: 'pip' + cache-dependency-path: 'pyproject.toml' + + - name: Install Clang + if: inputs.compiler == 'clang' && runner.os == 'Linux' + run: | + sudo apt-get update -y + sudo apt-get install -y clang libomp-dev + shell: bash + + - name: Install AIO + if: runner.os == 'Linux' && runner.arch == 'X64' + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libaio-dev + shell: bash + + - name: Print CPU info + if: runner.os == 'Linux' + run: lscpu + shell: bash + + - name: Set up environment variables + run: | + # Set number of processors for parallel builds + if [[ "${{ matrix.platform }}" == "macos-arm64" ]]; then + NPROC=$(sysctl -n hw.ncpu 2>/dev/null || echo 2) + else + NPROC=$(nproc 2>/dev/null || echo 2) + fi + echo "NPROC=$NPROC" >> $GITHUB_ENV + echo "Using $NPROC parallel jobs for builds" + + # Set compiler when clang is requested + if [[ "${{ inputs.compiler }}" == "clang" ]]; then + echo "CC=clang" >> $GITHUB_ENV + echo "CXX=clang++" >> $GITHUB_ENV + fi + + # Add Python user base bin to PATH for pip-installed CLI tools + echo "$(python -c 'import site; print(site.USER_BASE)')/bin" >> $GITHUB_PATH + shell: bash + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install \ + pybind11==3.0 \ + cmake==3.30.0 \ + ninja==1.11.1 \ + pytest \ + pytest-xdist \ + scikit-build-core \ + setuptools_scm + shell: bash + + - name: Check CMake subproject integration + if: matrix.platform == 'linux-x64' + run: | + SMOKE_SOURCE="$RUNNER_TEMP/zvec-subproject-smoke" + SMOKE_BUILD="$RUNNER_TEMP/zvec-subproject-build" + + mkdir -p "$SMOKE_SOURCE" "$SMOKE_BUILD" + cp "$GITHUB_WORKSPACE/.github/cmake/subproject-integration/CMakeLists.txt" \ + "$SMOKE_SOURCE/CMakeLists.txt" + + cmake -S "$SMOKE_SOURCE" -B "$SMOKE_BUILD" -G Ninja \ + -DZVEC_SOURCE_DIR="$GITHUB_WORKSPACE" \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_TOOLS=OFF \ + -DBUILD_C_BINDINGS=OFF \ + -DBUILD_PYTHON_BINDINGS=OFF \ + -DENABLE_WERROR=ON \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + shell: bash + + - name: Build from source + run: | + cd "$GITHUB_WORKSPACE" + + CMAKE_GENERATOR="Ninja" \ + CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \ + python -m pip install -v . \ + --no-build-isolation \ + --config-settings='cmake.define.BUILD_TOOLS=ON' \ + --config-settings='cmake.define.ENABLE_WERROR=ON' \ + --config-settings='cmake.define.CMAKE_C_COMPILER_LAUNCHER=ccache' \ + --config-settings='cmake.define.CMAKE_CXX_COMPILER_LAUNCHER=ccache' \ + ${{ matrix.arch_flag }} + shell: bash + + - name: Run C++ Tests + run: | + cd "$GITHUB_WORKSPACE/build" + cmake --build . --target unittest --parallel $NPROC + shell: bash + + - name: Run Python Tests + run: | + cd "$GITHUB_WORKSPACE" + python -m pytest python/tests/ + shell: bash + + - name: Run C++ Examples + run: | + cd "$GITHUB_WORKSPACE/examples/c++" + mkdir build && cd build + cmake .. -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + cmake --build . --parallel $NPROC + ./db-example + ./core-example + ./ailego-example + shell: bash + + - name: Run C Examples + run: | + cd "$GITHUB_WORKSPACE/examples/c" + mkdir build && cd build + cmake .. -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + cmake --build . --parallel $NPROC + ./c_api_basic_example + ./c_api_collection_schema_example + ./c_api_doc_example + ./c_api_field_schema_example + ./c_api_index_example + ./c_api_optimized_example + shell: bash diff --git a/.github/workflows/04-android-build.yml b/.github/workflows/04-android-build.yml new file mode 100644 index 0000000..aa44084 --- /dev/null +++ b/.github/workflows/04-android-build.yml @@ -0,0 +1,437 @@ +name: Android Cross Build & Test + +on: + workflow_call: + workflow_dispatch: + inputs: + api: + description: 'Android API level' + required: false + default: '34' + type: string + +permissions: + contents: read + +env: + NDK_VERSION: '26.1.10909125' + +jobs: + build-and-test: + runs-on: ubuntu-24.04 + timeout-minutes: 120 + strategy: + fail-fast: false + matrix: + abi: [x86_64] + api: ${{ github.event.inputs.api && fromJSON(format('["{0}"]', github.event.inputs.api)) || fromJSON('["34"]') }} + steps: + # ── Environment setup ────────────────────────────────────────────── + - name: Checkout + uses: actions/checkout@v7 + with: + submodules: recursive + + - name: Install dependencies + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + cmake ninja-build git ca-certificates python3 \ + build-essential make unzip curl + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: android-${{ matrix.abi }} + max-size: 300M + + - name: Setup Java 17 + uses: actions/setup-java@v5 + with: + distribution: temurin + java-version: '17' + + - name: Setup Android SDK + uses: android-actions/setup-android@v4 + + - name: Enable KVM + run: sudo chmod 666 /dev/kvm || true + + - name: Install NDK, emulator and system image + shell: bash + run: | + sdkmanager --install \ + "ndk;$NDK_VERSION" \ + "platform-tools" \ + "platforms;android-${{ matrix.api }}" \ + "emulator" + + # Install x86_64 system image (try variants in order of availability) + sdkmanager --install "system-images;android-${{ matrix.api }};google_apis;x86_64" 2>/dev/null || \ + sdkmanager --install "system-images;android-${{ matrix.api }};google_apis_playstore;x86_64" 2>/dev/null || \ + sdkmanager --install "system-images;android-${{ matrix.api }};default;x86_64" + + # ── Step 1: build host protoc (using HOST compiler, NOT NDK) ─────── + - name: Cache host protoc + uses: actions/cache@v6 + with: + path: build_host + key: ${{ runner.os }}-host-protoc-${{ hashFiles('thirdparty/protobuf/**', 'CMakeLists.txt') }} + restore-keys: | + ${{ runner.os }}-host-protoc- + + - name: 'Step 1: Build host protoc' + shell: bash + run: | + if [ ! -f "build_host/bin/protoc" ]; then + git submodule foreach --recursive 'git stash --include-untracked' 2>/dev/null || true + cmake -S . -B build_host \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_TOOLCHAIN_FILE="" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \ + -G Ninja + cmake --build build_host --target protoc --parallel + else + echo "Using cached host protoc" + fi + + # ── Step 2: cross-compile zvec + tests for Android ───────────────── + - name: 'Step 2: Cross-compile zvec and tests' + shell: bash + env: + BUILD_DIR: build_android_${{ matrix.abi }} + run: | + ANDROID_NDK_HOME="$ANDROID_HOME/ndk/$NDK_VERSION" + + # Reset thirdparty so the cross toolchain can patch cleanly + git submodule foreach --recursive 'git stash --include-untracked' 2>/dev/null || true + + # Force reconfigure to pick up any cmake changes + rm -f "$BUILD_DIR/CMakeCache.txt" + + cmake -S . -B "$BUILD_DIR" -G Ninja \ + -DANDROID_NDK="$ANDROID_NDK_HOME" \ + -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake" \ + -DANDROID_ABI=${{ matrix.abi }} \ + -DANDROID_NATIVE_API_LEVEL=${{ matrix.api }} \ + -DANDROID_STL=c++_static \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_PYTHON_BINDINGS=OFF \ + -DBUILD_TOOLS=OFF \ + -DENABLE_NATIVE=OFF \ + -DAUTO_DETECT_ARCH=OFF \ + -DENABLE_WERROR=ON \ + -DCMAKE_INSTALL_PREFIX="$BUILD_DIR/install" \ + -DGLOBAL_CC_PROTOBUF_PROTOC="$GITHUB_WORKSPACE/build_host/bin/protoc" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + + echo "Building all targets..." + cmake --build "$BUILD_DIR" --parallel + + # Discover test targets from ctest metadata + echo "Discovering test targets..." + TEST_NAMES=() + while IFS= read -r line; do + name=$(echo "$line" | sed -n 's/.*Test[[:space:]]*#[0-9]*:[[:space:]]*//p') + [ -n "$name" ] && TEST_NAMES+=("$name") + done < <(cd "$BUILD_DIR" && ctest --show-only 2>/dev/null || true) + + # Fallback: scan ninja targets for *_test + if [ ${#TEST_NAMES[@]} -eq 0 ]; then + echo "ctest unavailable, scanning ninja targets..." + while IFS= read -r line; do + name=$(echo "$line" | sed -n 's/^\([a-zA-Z0-9_]*_test\): .*/\1/p') + [ -n "$name" ] && TEST_NAMES+=("$name") + done < <(ninja -C "$BUILD_DIR" -t targets all 2>/dev/null || true) + fi + + echo "Building ${#TEST_NAMES[@]} test executables..." + ninja -C "$BUILD_DIR" -j$(nproc) "${TEST_NAMES[@]}" + + # ── Step 3: start emulator ───────────────────────────────────────── + - name: 'Step 3: Start Android emulator' + shell: bash + run: | + AVD_NAME="zvec_test_avd" + AVD_DIR="$HOME/.android/avd/${AVD_NAME}.avd" + AVD_INI="$HOME/.android/avd/${AVD_NAME}.ini" + + # Find installed system image (same priority as build_android.sh) + SYS_IMG="" + for variant in google_apis google_apis_playstore default; do + candidate="$ANDROID_HOME/system-images/android-${{ matrix.api }}/$variant/x86_64" + if [ -d "$candidate" ]; then + SYS_IMG="$candidate" + break + fi + done + if [ -z "$SYS_IMG" ]; then + SYS_IMG=$(find "$ANDROID_HOME/system-images" -type d -name "x86_64" 2>/dev/null | head -1) + fi + if [ -z "$SYS_IMG" ]; then + echo "ERROR: No x86_64 system image found" + exit 1 + fi + echo "System image: $SYS_IMG" + + # Extract tag from path (e.g. .../google_apis/x86_64 -> google_apis) + SYS_TAG=$(basename "$(dirname "$SYS_IMG")") + + # Create AVD via INI files (faster and more reliable than avdmanager) + mkdir -p "$AVD_DIR" + + cat > "$AVD_INI" << EOAVD + avd.ini.encoding=UTF-8 + path=${AVD_DIR} + path.rel=avd/${AVD_NAME}.avd + target=android-${{ matrix.api }} + EOAVD + + cat > "$AVD_DIR/config.ini" << EOCFG + AvdId=${AVD_NAME} + PlayStore.enabled=false + abi.type=x86_64 + avd.ini.displayname=${AVD_NAME} + avd.ini.encoding=UTF-8 + disk.dataPartition.size=8G + hw.accelerator.isConfigured=true + hw.cpu.arch=x86_64 + hw.cpu.ncore=4 + hw.lcd.density=420 + hw.lcd.height=1920 + hw.lcd.width=1080 + hw.ramSize=4096 + image.sysdir.1=${SYS_IMG}/ + tag.display=${SYS_TAG} + tag.id=${SYS_TAG} + EOCFG + + echo "Created AVD: $AVD_NAME" + + # Launch emulator in background + $ANDROID_HOME/emulator/emulator -avd "$AVD_NAME" \ + -no-window -no-audio -no-boot-anim \ + -gpu swiftshader_indirect \ + -netdelay none -netspeed full \ + -memory 4096 \ + -no-snapshot \ + -wipe-data & + echo "EMULATOR_PID=$!" >> "$GITHUB_ENV" + + # Wait for device to be reachable + adb wait-for-device + + # Poll for boot completion with timeout + echo "Waiting for boot to complete..." + TIMEOUT=300 + ELAPSED=0 + while true; do + BOOTED=$(adb shell getprop sys.boot_completed 2>/dev/null | tr -d '\r\n ' || true) + if [ "$BOOTED" = "1" ]; then + echo "Emulator ready (took ${ELAPSED}s)" + break + fi + if [ $ELAPSED -ge $TIMEOUT ]; then + echo "ERROR: Emulator failed to boot within ${TIMEOUT}s" + exit 1 + fi + sleep 3 + ELAPSED=$((ELAPSED + 3)) + done + + echo "Device ABI: $(adb shell getprop ro.product.cpu.abi | tr -d '\r')" + echo "ABI list : $(adb shell getprop ro.product.cpu.abilist | tr -d '\r')" + + # ── Step 4: run unit tests on emulator ───────────────────────────── + - name: 'Step 4: Run unit tests on emulator' + shell: bash + env: + BUILD_DIR: build_android_${{ matrix.abi }} + run: | + DEVICE_TEST_DIR="/data/local/tmp/zvec_tests" + DEVICE_LIB_DIR="$DEVICE_TEST_DIR/lib" + adb shell "mkdir -p $DEVICE_TEST_DIR $DEVICE_LIB_DIR" + + # Push shared libraries + echo "Pushing shared libraries..." + SO_COUNT=0 + while IFS= read -r so_file; do + adb push "$so_file" "$DEVICE_LIB_DIR/$(basename "$so_file")" > /dev/null 2>&1 + SO_COUNT=$((SO_COUNT + 1)) + done < <(find "$BUILD_DIR/lib" -name "*.so" -type f 2>/dev/null) + echo "Pushed $SO_COUNT shared libraries" + + # Push helper binaries (needed by crash_recovery tests which fork+exec them) + echo "Pushing helper binaries..." + for helper_name in data_generator collection_optimizer; do + helper_path=$(find "$BUILD_DIR" -name "$helper_name" -type f -executable ! -name "*_test" 2>/dev/null | head -1) + if [ -n "$helper_path" ]; then + adb push "$helper_path" "$DEVICE_TEST_DIR/$helper_name" > /dev/null 2>&1 + adb shell "chmod 755 $DEVICE_TEST_DIR/$helper_name" + echo " Pushed $helper_name" + fi + done + + # Discover test targets from ctest metadata + TEST_NAMES=() + while IFS= read -r line; do + name=$(echo "$line" | sed -n 's/.*Test[[:space:]]*#[0-9]*:[[:space:]]*//p') + [ -n "$name" ] && TEST_NAMES+=("$name") + done < <(cd "$BUILD_DIR" && ctest --show-only 2>/dev/null || true) + + if [ ${#TEST_NAMES[@]} -eq 0 ]; then + while IFS= read -r line; do + name=$(echo "$line" | sed -n 's/^\([a-zA-Z0-9_]*_test\): .*/\1/p') + [ -n "$name" ] && TEST_NAMES+=("$name") + done < <(ninja -C "$BUILD_DIR" -t targets all 2>/dev/null || true) + fi + + # Collect test binaries + TEST_BINS=() + for name in "${TEST_NAMES[@]}"; do + bin_path=$(find "$BUILD_DIR" -name "$name" -type f -executable 2>/dev/null | head -1) + if [ -n "$bin_path" ]; then + TEST_BINS+=("$bin_path") + else + echo "WARNING: binary not found for '$name'" + fi + done + + TOTAL=${#TEST_BINS[@]} + PASSED=0 + FAILED=0 + FAILED_NAMES=() + IDX=0 + + echo "Running $TOTAL unit tests on emulator..." + + for test_bin in "${TEST_BINS[@]}"; do + IDX=$((IDX + 1)) + test_name=$(basename "$test_bin") + device_path="$DEVICE_TEST_DIR/$test_name" + # Give each test its own working directory to avoid name collisions + WORK_DIR="$DEVICE_TEST_DIR/workdir_${test_name}" + + echo "" + echo "────────────────────────────────────────" + echo " [$IDX/$TOTAL] $test_name" + echo "────────────────────────────────────────" + + set +e + # Create isolated working directory + adb shell "mkdir -p $WORK_DIR" 2>/dev/null + + # Copy helper binaries into working directory so crash_recovery tests + # (which fork+exec data_generator / collection_optimizer) can find them + adb shell "for h in $DEVICE_TEST_DIR/data_generator $DEVICE_TEST_DIR/collection_optimizer; do [ -f \$h ] && cp \$h $WORK_DIR/; done" 2>/dev/null + + # Push test binary + adb push "$test_bin" "$device_path" > /dev/null 2>&1 + adb shell "chmod 755 $device_path" 2>/dev/null + + # Run test from its own working directory with LD_LIBRARY_PATH + OUTPUT=$(adb shell "cd $WORK_DIR && LD_LIBRARY_PATH=$DEVICE_LIB_DIR $device_path 2>&1; echo EXIT_CODE=\$?" 2>&1) + + # Extract exit code from the output + EXIT_CODE=$(echo "$OUTPUT" | grep -o 'EXIT_CODE=[0-9]*' | tail -1 | cut -d= -f2) + set -e + + # Print test output (without the EXIT_CODE marker) + echo "$OUTPUT" | grep -v 'EXIT_CODE=' | sed 's/^/ /' || true + + if [ "$EXIT_CODE" = "0" ]; then + echo " >>> PASSED" + PASSED=$((PASSED + 1)) + else + # Detect "crash-on-exit" pattern: all gtest assertions passed but + # process crashed during static destructor teardown (common with c++_static STL) + GTEST_PASSED_LINE=$(echo "$OUTPUT" | grep '\[ PASSED \]' | tail -1 || true) + GTEST_FAILED_LINE=$(echo "$OUTPUT" | grep '\[ FAILED \]' | head -1 || true) + if [ -n "$GTEST_PASSED_LINE" ] && [ -z "$GTEST_FAILED_LINE" ] && \ + { [ "$EXIT_CODE" = "139" ] || [ "$EXIT_CODE" = "134" ] || [ "$EXIT_CODE" = "135" ]; }; then + echo " >>> PASSED (crash-on-exit ignored, exit=$EXIT_CODE)" + PASSED=$((PASSED + 1)) + else + echo " >>> FAILED (exit=$EXIT_CODE)" + FAILED=$((FAILED + 1)) + FAILED_NAMES+=("$test_name") + fi + fi + + # Clean up binary and working directory to reclaim disk space + adb shell "rm -rf $device_path $WORK_DIR" 2>/dev/null || true + done + + echo "" + echo "============================================================" + echo " Test Summary" + echo "============================================================" + echo " Total : $TOTAL" + echo " Passed : $PASSED" + echo " Failed : $FAILED" + if [ $FAILED -gt 0 ]; then + echo "" + echo " Failed tests:" + for name in "${FAILED_NAMES[@]}"; do + echo " - $name" + done + fi + echo "============================================================" + + if [ $FAILED -gt 0 ]; then + exit 1 + fi + echo "All tests passed!" + + # ── Step 5: build and run examples ───────────────────────────────── + - name: 'Step 5: Build and run examples' + shell: bash + env: + BUILD_DIR: build_android_${{ matrix.abi }} + run: | + ANDROID_NDK_HOME="$ANDROID_HOME/ndk/$NDK_VERSION" + EXAMPLES_BUILD="examples/c++/build-android-examples-${{ matrix.abi }}" + + cmake -S examples/c++ -B "$EXAMPLES_BUILD" -G Ninja \ + -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake" \ + -DANDROID_ABI=${{ matrix.abi }} \ + -DANDROID_PLATFORM=android-${{ matrix.api }} \ + -DANDROID_STL=c++_static \ + -DCMAKE_BUILD_TYPE=Release \ + -DHOST_BUILD_DIR="$BUILD_DIR" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + cmake --build "$EXAMPLES_BUILD" --parallel + + # Reuse the shared-library directory from Step 4; push again in + # case Step 4 was skipped or the directory was cleaned. + DEVICE_LIB_DIR="/data/local/tmp/zvec_tests/lib" + adb shell "mkdir -p $DEVICE_LIB_DIR" 2>/dev/null || true + SO_COUNT=0 + while IFS= read -r so_file; do + adb push "$so_file" "$DEVICE_LIB_DIR/$(basename "$so_file")" > /dev/null 2>&1 + SO_COUNT=$((SO_COUNT + 1)) + done < <(find "$BUILD_DIR/lib" -name "*.so" -type f 2>/dev/null) + echo "Pushed $SO_COUNT shared libraries to $DEVICE_LIB_DIR" + + for example in ailego-example core-example db-example; do + if [ -f "$EXAMPLES_BUILD/$example" ]; then + echo "=== Running $example ===" + adb push "$EXAMPLES_BUILD/$example" "/data/local/tmp/$example" > /dev/null 2>&1 + adb shell "chmod 755 /data/local/tmp/$example && cd /data/local/tmp && LD_LIBRARY_PATH=$DEVICE_LIB_DIR ./$example" + adb shell "rm -f /data/local/tmp/$example" + fi + done + + # ── Cleanup ──────────────────────────────────────────────────────── + - name: Stop emulator + if: always() + shell: bash + run: | + adb emu kill 2>/dev/null || true + sleep 2 + if [ -n "$EMULATOR_PID" ]; then + kill "$EMULATOR_PID" 2>/dev/null || true + fi diff --git a/.github/workflows/05-windows-build.yml b/.github/workflows/05-windows-build.yml new file mode 100644 index 0000000..dcdbda3 --- /dev/null +++ b/.github/workflows/05-windows-build.yml @@ -0,0 +1,157 @@ +name: Windows Build + +on: + workflow_call: + workflow_dispatch: + +permissions: + contents: read + +jobs: + # Windows build and test matrix + build-and-test-windows: + name: Build & Test (${{ matrix.platform }}) + runs-on: ${{ matrix.platform }} + + strategy: + fail-fast: false + matrix: + include: + - platform: windows-2022 + - platform: windows-2025 + + env: + SCCACHE_GHA_ENABLED: "true" + + steps: + - name: Show env info + run: | + Get-CimInstance -ClassName Win32_Processor + where cl + & "C:\Program Files (x86)\Microsoft Visual Studio\Installer\vswhere.exe" -all -products * -prerelease -format json + shell: powershell + + - name: Checkout code + uses: actions/checkout@v7 + with: + submodules: recursive + + - name: Cleanup left marker files + run: | + git submodule foreach --recursive 'git reset --hard && git clean -ffdx' + shell: powershell + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.10 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.10' + cache: 'pip' + cache-dependency-path: 'pyproject.toml' + + - name: Set up MSVC environment + uses: ilammy/msvc-dev-cmd@v1.13.0 + with: + arch: x64 + + - name: Set up environment variables + run: | + $nproc = (Get-CimInstance Win32_ComputerSystem).NumberOfLogicalProcessors + echo "NPROC=$nproc" >> $env:GITHUB_ENV + echo "Using $nproc parallel jobs for builds" + shell: powershell + + - name: Install dependencies + run: | + python -m pip install --upgrade pip ` + pybind11==3.0 ` + cmake==3.30.0 ` + ninja==1.11.1 ` + pytest ` + pytest-xdist ` + scikit-build-core ` + setuptools_scm + shell: powershell + + - name: Build from source + run: | + cd "$env:GITHUB_WORKSPACE" + $env:CMAKE_GENERATOR = "Ninja" + $env:CMAKE_BUILD_PARALLEL_LEVEL = "$env:NPROC" + python -m pip install -v . ` + --no-build-isolation ` + --config-settings='cmake.define.BUILD_TOOLS=ON' ` + --config-settings='cmake.define.ENABLE_WERROR=ON' ` + --config-settings='cmake.define.CMAKE_C_COMPILER_LAUNCHER=sccache' ` + --config-settings='cmake.define.CMAKE_CXX_COMPILER_LAUNCHER=sccache' + shell: powershell + + - name: Show sccache statistics after pip install + if: always() + run: sccache --show-stats + shell: powershell + + - name: Run C++ Tests + run: | + cd "$env:GITHUB_WORKSPACE\build" + cmake --build . --target unittest --config Release --parallel $env:NPROC + shell: powershell + + - name: Show sccache statistics after tests + if: always() + run: sccache --show-stats + shell: powershell + + - name: Run Python Tests + run: | + cd "$env:GITHUB_WORKSPACE" + python -m pytest python/tests/ --basetemp=./.pytest_tmp + shell: powershell + + - name: Run C++ Examples + run: | + cd "$env:GITHUB_WORKSPACE\examples\c++" + mkdir build + cd build + cmake .. -G Ninja -DCMAKE_BUILD_TYPE=Release ` + -DCMAKE_C_COMPILER_LAUNCHER=sccache ` + -DCMAKE_CXX_COMPILER_LAUNCHER=sccache + cmake --build . --config Release --parallel $env:NPROC + + # Copy zvec DLLs next to the example executables so Windows can find them. + # CMake places DLLs in bin/ (RUNTIME output) and import libs in lib/. + $buildDir = "$env:GITHUB_WORKSPACE\build" + foreach ($dllName in @("zvec.dll", "zvec_core.dll", "zvec_ailego.dll")) { + $found = $false + foreach ($sub in @("$buildDir\bin", "$buildDir\bin\Release", "$buildDir\lib", "$buildDir\lib\Release")) { + $dllPath = Join-Path $sub $dllName + if (Test-Path $dllPath) { + Copy-Item $dllPath -Destination . -Force + Write-Host "Copied $dllName from $sub" + $found = $true + break + } + } + if (-not $found) { + Write-Host "WARNING: $dllName not found, searching recursively..." + $dll = Get-ChildItem -Path $buildDir -Filter $dllName -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1 + if ($dll) { + Copy-Item $dll.FullName -Destination . -Force + Write-Host "Copied $dllName from $($dll.DirectoryName)" + } else { + Write-Error "$dllName not found anywhere under $buildDir" + } + } + } + + .\db-example.exe + .\core-example.exe + .\ailego-example.exe + shell: powershell + + - name: Show sccache statistics + if: always() + run: sccache --show-stats + shell: powershell diff --git a/.github/workflows/06-ios-build.yml b/.github/workflows/06-ios-build.yml new file mode 100644 index 0000000..98a444f --- /dev/null +++ b/.github/workflows/06-ios-build.yml @@ -0,0 +1,163 @@ +name: iOS Cross Build + +on: + workflow_call: + workflow_dispatch: + +permissions: + contents: read + +jobs: + build-ios: + runs-on: macos-15 + strategy: + fail-fast: false + matrix: + include: + - platform: SIMULATORARM64 + arch: arm64 + sdk: iphonesimulator + test_on_simulator: true + - platform: OS + arch: arm64 + sdk: iphoneos + test_on_simulator: false + + name: iOS (${{ matrix.platform }}) + + steps: + - name: Checkout + uses: actions/checkout@v7 + with: + submodules: recursive + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: ios-${{ matrix.platform }} + max-size: 150M + + - name: Cache host protoc build + uses: actions/cache@v6 + with: + path: build_host + key: macos-host-protoc-${{ hashFiles('thirdparty/protobuf/**', 'CMakeLists.txt') }} + restore-keys: | + macos-host-protoc- + + - name: Build host protoc + run: | + if [ ! -f "build_host/bin/protoc" ]; then + cmake -S . -B build_host -DCMAKE_BUILD_TYPE=Release -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + cmake --build build_host --target protoc --parallel $(sysctl -n hw.ncpu) + else + echo "Using cached host protoc" + fi + + - name: Configure and Build + run: | + git submodule foreach --recursive 'git stash --include-untracked' || true + + SDK_PATH=$(xcrun --sdk ${{ matrix.sdk }} --show-sdk-path) + NPROC=$(sysctl -n hw.ncpu) + + cmake -S . -B build_ios_${{ matrix.platform }} \ + -DCMAKE_SYSTEM_NAME=iOS \ + -DCMAKE_OSX_DEPLOYMENT_TARGET="13.0" \ + -DCMAKE_OSX_ARCHITECTURES="${{ matrix.arch }}" \ + -DCMAKE_OSX_SYSROOT="$SDK_PATH" \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_PYTHON_BINDINGS=OFF \ + -DBUILD_TOOLS=OFF \ + -DENABLE_WERROR=ON \ + -DCMAKE_INSTALL_PREFIX="./install" \ + -DGLOBAL_CC_PROTOBUF_PROTOC="$GITHUB_WORKSPACE/build_host/bin/protoc" \ + -DIOS=ON \ + -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + + cmake --build build_ios_${{ matrix.platform }} --parallel $NPROC + + - name: Build test targets + if: matrix.test_on_simulator + run: | + NPROC=$(sysctl -n hw.ncpu) + cmake --build build_ios_${{ matrix.platform }} --target unittest --parallel $NPROC + + - name: Boot iOS Simulator + if: matrix.test_on_simulator + run: | + DEVICE_ID=$(xcrun simctl list devices available -j \ + | python3 -c " + import json, sys + data = json.load(sys.stdin) + for runtime, devices in data['devices'].items(): + if 'iOS' in runtime: + for d in devices: + if 'iPhone' in d['name'] and d['isAvailable']: + print(d['udid']) + sys.exit(0) + sys.exit(1) + ") + echo "DEVICE_ID=$DEVICE_ID" >> $GITHUB_ENV + xcrun simctl boot "$DEVICE_ID" + echo "Booted simulator: $DEVICE_ID" + + - name: Run all tests on simulator + if: matrix.test_on_simulator + run: | + FAILED_TESTS="" + PASSED=0 + TOTAL=0 + + for APP in build_ios_${{ matrix.platform }}/bin/*_test.app; do + [ -d "$APP" ] || continue + TEST_NAME=$(basename "$APP" .app) + BUNDLE_ID="com.zvec.${TEST_NAME}" + TOTAL=$((TOTAL + 1)) + + echo "::group::Running ${TEST_NAME}" + xcrun simctl install "$DEVICE_ID" "$APP" + set +eo pipefail + for attempt in 1 2 3; do + xcrun simctl launch --console "$DEVICE_ID" "$BUNDLE_ID" 2>&1 | tee /tmp/${TEST_NAME}.log + LAUNCH_EXIT=${PIPESTATUS[0]} + if ! grep -q "unknown to FrontBoard" /tmp/${TEST_NAME}.log; then + break + fi + echo "::warning::Attempt ${attempt}/3: FrontBoard has not registered ${TEST_NAME} yet, retrying in 3s..." + sleep 3 + done + set -eo pipefail + + if grep -q '\[ FAILED \]' /tmp/${TEST_NAME}.log; then + echo "::error::${TEST_NAME} has failing tests" + FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" + elif grep -q '\[ PASSED \]' /tmp/${TEST_NAME}.log; then + PASSED=$((PASSED + 1)) + elif grep -qE 'Failed: 0$' /tmp/${TEST_NAME}.log; then + # c_api_test uses a custom test framework (not GTest) + PASSED=$((PASSED + 1)) + elif [ "$LAUNCH_EXIT" -eq 0 ]; then + echo "::warning::${TEST_NAME} exited 0 but produced no recognisable test summary" + PASSED=$((PASSED + 1)) + else + echo "::error::${TEST_NAME} exited ${LAUNCH_EXIT} with no test summary" + FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" + fi + echo "::endgroup::" + done + + echo "Test summary: ${PASSED}/${TOTAL} passed" + if [ -n "$FAILED_TESTS" ]; then + echo "::error::Failed tests:${FAILED_TESTS}" + exit 1 + fi + + - name: Shutdown Simulator + if: matrix.test_on_simulator && always() + run: | + xcrun simctl shutdown "$DEVICE_ID" || true diff --git a/.github/workflows/07-linux-riscv-build.yml b/.github/workflows/07-linux-riscv-build.yml new file mode 100644 index 0000000..9f3be60 --- /dev/null +++ b/.github/workflows/07-linux-riscv-build.yml @@ -0,0 +1,342 @@ +name: Weekly Linux RISC-V Build + +on: + schedule: + # Saturday 00:00 UTC+8 (Friday 16:00 UTC) + - cron: '0 16 * * 5' + + workflow_dispatch: + +permissions: + contents: read + +env: + RISE_PYPI: https://gitlab.com/api/v4/projects/56254198/packages/pypi/simple + PIP_BREAK_SYSTEM_PACKAGES: 1 + +jobs: + build: + name: Build (linux-riscv64) + runs-on: ubuntu-24.04-riscv + + steps: + - name: Checkout code + uses: actions/checkout@v7 + with: + submodules: recursive + + - name: Install build dependencies + run: | + sudo mkdir -p /var/lib/dpkg/updates + sudo mkdir -p /var/lib/apt/lists/ + sudo mkdir -p /var/cache/apt/archives/ + sudo touch /var/lib/dpkg/status + sudo apt-get purge -y byobu || true + sudo apt-get update -o Dpkg::Lock::Timeout=300 + sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \ + -o Dpkg::Options::="--force-confdef" \ + -o Dpkg::Options::="--force-confold" \ + -o Dpkg::Lock::Timeout=300 \ + ccache python3-pybind11 pybind11-dev + shell: bash + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: linux-riscv64-${{ runner.os }}-gcc + max-size: 150M + + - name: Build from source + run: | + cd "$GITHUB_WORKSPACE" + NPROC=$(nproc 2>/dev/null || echo 2) + echo "Using $NPROC parallel jobs for builds" + cmake -S . -B build \ + -G "Unix Makefiles" \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_TOOLS=ON \ + -DBUILD_PYTHON_BINDINGS=ON \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + make -C build -j"$NPROC" + shell: bash + + - name: Archive entire workspace + run: | + cd "$GITHUB_WORKSPACE" + tar -cf linux-riscv64-workspace.tar . + shell: bash + + - name: Upload workspace artifacts + uses: actions/upload-artifact@v7 + with: + name: linux-riscv64-workspace + path: ${{ github.workspace }}/linux-riscv64-workspace.tar + if-no-files-found: error + + cpp-tests: + name: C++ Tests + runs-on: ubuntu-24.04-riscv + needs: build + + steps: + - name: Download workspace artifacts + uses: actions/download-artifact@v8 + with: + name: linux-riscv64-workspace + path: ${{ github.workspace }} + + - name: Extract workspace + run: | + cd "$GITHUB_WORKSPACE" + tar -xf linux-riscv64-workspace.tar + shell: bash + + - name: Install test dependencies + run: | + sudo mkdir -p /var/lib/dpkg/updates + sudo mkdir -p /var/lib/apt/lists/ + sudo mkdir -p /var/cache/apt/archives/ + sudo touch /var/lib/dpkg/status + sudo apt-get purge -y byobu || true + sudo apt-get update -o Dpkg::Lock::Timeout=300 + sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \ + -o Dpkg::Options::="--force-confdef" \ + -o Dpkg::Options::="--force-confold" \ + -o Dpkg::Lock::Timeout=300 \ + ccache python3-pybind11 pybind11-dev libgtest-dev liburing-dev + shell: bash + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: linux-riscv64-${{ runner.os }}-gcc + max-size: 150M + + - name: Reconfigure build directory + run: | + cd "$GITHUB_WORKSPACE" + cmake -S . -B build \ + -G "Unix Makefiles" \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_TOOLS=ON \ + -DBUILD_PYTHON_BINDINGS=ON \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + shell: bash + + - name: Run C++ Tests + run: | + cd "$GITHUB_WORKSPACE/build" + NPROC=$(nproc 2>/dev/null || echo 2) + make unittest -j"$NPROC" + shell: bash + + python-tests: + name: Python Tests + runs-on: ubuntu-24.04-riscv + needs: build + + steps: + - name: Select Python + run: | + if command -v python3 >/dev/null 2>&1; then + PYTHON_BIN=python3 + elif command -v python >/dev/null 2>&1; then + PYTHON_BIN=python + else + echo "No local Python interpreter found on PATH" + exit 1 + fi + "$PYTHON_BIN" --version + echo "PYTHON=$PYTHON_BIN" >> "$GITHUB_ENV" + shell: bash + + - name: Download workspace artifacts + uses: actions/download-artifact@v8 + with: + name: linux-riscv64-workspace + path: ${{ github.workspace }} + + - name: Extract workspace + run: | + cd "$GITHUB_WORKSPACE" + tar -xf linux-riscv64-workspace.tar + echo "$($PYTHON -c 'import site; print(site.USER_BASE)')/bin" >> "$GITHUB_PATH" + shell: bash + + - name: Cache pip + uses: actions/cache@v6 + with: + path: ~/.cache/pip + key: linux-riscv64-${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }} + restore-keys: | + linux-riscv64-${{ runner.os }}-pip- + + - name: Install dependencies + run: | + sudo mkdir -p /var/lib/dpkg/updates + sudo mkdir -p /var/lib/apt/lists/ + sudo mkdir -p /var/cache/apt/archives/ + sudo touch /var/lib/dpkg/status + sudo apt-get purge -y byobu || true + sudo apt-get update -o Dpkg::Lock::Timeout=300 + sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \ + -o Dpkg::Options::="--force-confdef" \ + -o Dpkg::Options::="--force-confold" \ + -o Dpkg::Lock::Timeout=300 \ + ccache libgtest-dev liburing-dev + + $PYTHON -m pip install --upgrade pip + $PYTHON -m pip install numpy==2.2.2 cmake==3.30.0 ninja==1.11.1.1 --index-url "$RISE_PYPI" + $PYTHON -m pip install pybind11==3.0 pytest scikit-build-core setuptools_scm pytest-xdist + shell: bash + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: linux-riscv64-${{ runner.os }}-gcc + max-size: 150M + + - name: Reconfigure build directory + run: | + cd "$GITHUB_WORKSPACE" + cmake -S . -B build \ + -G "Unix Makefiles" \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_TOOLS=ON \ + -DBUILD_PYTHON_BINDINGS=ON \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \ + -Dpybind11_DIR="$($PYTHON -c 'import pybind11; print(pybind11.get_cmake_dir())')" + shell: bash + + - name: Install from existing build directory + run: | + cd "$GITHUB_WORKSPACE" + NPROC=$(nproc 2>/dev/null || echo 2) + export SKBUILD_BUILD_DIR="$GITHUB_WORKSPACE/build" + CMAKE_GENERATOR="Unix Makefiles" \ + CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \ + $PYTHON -m pip install -v . \ + --no-build-isolation \ + --config-settings='cmake.define.BUILD_TOOLS="ON"' \ + --config-settings='cmake.define.ENABLE_WERROR=ON' \ + --config-settings='cmake.define.CMAKE_C_COMPILER_LAUNCHER=ccache' \ + --config-settings='cmake.define.CMAKE_CXX_COMPILER_LAUNCHER=ccache' + shell: bash + + - name: Run Python Tests + run: | + cd "$GITHUB_WORKSPACE" + $PYTHON -m pytest python/tests/ + shell: bash + + cpp-examples: + name: C++ Examples + runs-on: ubuntu-24.04-riscv + needs: build + + steps: + - name: Download workspace artifacts + uses: actions/download-artifact@v8 + with: + name: linux-riscv64-workspace + path: ${{ github.workspace }} + + - name: Extract workspace + run: | + cd "$GITHUB_WORKSPACE" + tar -xf linux-riscv64-workspace.tar + shell: bash + + - name: Install ccache + run: | + sudo mkdir -p /var/lib/dpkg/updates + sudo mkdir -p /var/lib/apt/lists/ + sudo mkdir -p /var/cache/apt/archives/ + sudo touch /var/lib/dpkg/status + sudo apt-get purge -y byobu || true + sudo apt-get update -o Dpkg::Lock::Timeout=300 + sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \ + -o Dpkg::Options::="--force-confdef" \ + -o Dpkg::Options::="--force-confold" \ + -o Dpkg::Lock::Timeout=300 \ + ccache + shell: bash + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: linux-riscv64-${{ runner.os }}-gcc + max-size: 150M + + - name: Run C++ Examples + run: | + cd "$GITHUB_WORKSPACE/examples/c++" + NPROC=$(nproc 2>/dev/null || echo 2) + mkdir -p build && cd build + cmake .. -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + make -j "$NPROC" + ./db-example + ./core-example + ./ailego-example + shell: bash + + c-examples: + name: C Examples + runs-on: ubuntu-24.04-riscv + needs: build + + steps: + - name: Download workspace artifacts + uses: actions/download-artifact@v8 + with: + name: linux-riscv64-workspace + path: ${{ github.workspace }} + + - name: Extract workspace + run: | + cd "$GITHUB_WORKSPACE" + tar -xf linux-riscv64-workspace.tar + shell: bash + + - name: Install ccache + run: | + sudo mkdir -p /var/lib/dpkg/updates + sudo mkdir -p /var/lib/apt/lists/ + sudo mkdir -p /var/cache/apt/archives/ + sudo touch /var/lib/dpkg/status + sudo apt-get purge -y byobu || true + sudo apt-get update -o Dpkg::Lock::Timeout=300 + sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \ + -o Dpkg::Options::="--force-confdef" \ + -o Dpkg::Options::="--force-confold" \ + -o Dpkg::Lock::Timeout=300 \ + ccache + shell: bash + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: linux-riscv64-${{ runner.os }}-gcc + max-size: 150M + + - name: Run C Examples + run: | + cd "$GITHUB_WORKSPACE/examples/c" + NPROC=$(nproc 2>/dev/null || echo 2) + mkdir -p build && cd build + cmake .. -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + make -j "$NPROC" + ./c_api_basic_example + ./c_api_collection_schema_example + ./c_api_doc_example + ./c_api_field_schema_example + ./c_api_index_example + ./c_api_optimized_example + shell: bash diff --git a/.github/workflows/08-cmake-subproject-integration.yml b/.github/workflows/08-cmake-subproject-integration.yml new file mode 100644 index 0000000..564a304 --- /dev/null +++ b/.github/workflows/08-cmake-subproject-integration.yml @@ -0,0 +1,102 @@ +name: Nightly CMake Subproject Integration + +on: + schedule: + # Run after the other nightly jobs to avoid starting every heavy job at once. + - cron: '30 16 * * *' + workflow_dispatch: + +permissions: + contents: read + +jobs: + subproject-integration: + name: Subproject Build & Example + runs-on: ubuntu-24.04 + timeout-minutes: 90 + + steps: + - name: Checkout code + uses: actions/checkout@v7 + with: + submodules: recursive + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: cmake-subproject-linux-x64 + max-size: 150M + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.10' + cache: 'pip' + cache-dependency-path: 'pyproject.toml' + + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libaio-dev + shell: bash + + - name: Set up environment variables + run: | + NPROC=$(nproc 2>/dev/null || echo 2) + echo "NPROC=$NPROC" >> "$GITHUB_ENV" + echo "$(python -c 'import site; print(site.USER_BASE)')/bin" >> "$GITHUB_PATH" + shell: bash + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install \ + cmake==3.30.0 \ + ninja==1.11.1 + shell: bash + + - name: Configure subproject integration + run: | + SMOKE_SOURCE="$RUNNER_TEMP/zvec-subproject-nightly" + SMOKE_BUILD="$RUNNER_TEMP/zvec-subproject-nightly-build" + + mkdir -p "$SMOKE_SOURCE" "$SMOKE_BUILD" + cp "$GITHUB_WORKSPACE/.github/cmake/subproject-integration/CMakeLists.txt" \ + "$SMOKE_SOURCE/CMakeLists.txt" + + CMAKE_ARGS=( + -DZVEC_SOURCE_DIR="$GITHUB_WORKSPACE" + -DCMAKE_BUILD_TYPE=Release + -DENABLE_WERROR=ON + -DUSE_OSS_MIRROR=ON + -DCMAKE_C_COMPILER_LAUNCHER=ccache + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + -DBUILD_TOOLS=OFF + -DBUILD_C_BINDINGS=OFF + -DBUILD_PYTHON_BINDINGS=OFF + ) + + cmake -S "$SMOKE_SOURCE" -B "$SMOKE_BUILD" -G Ninja \ + "${CMAKE_ARGS[@]}" + + echo "SMOKE_BUILD=$SMOKE_BUILD" >> "$GITHUB_ENV" + shell: bash + + - name: Build subproject targets + run: | + cmake --build "$SMOKE_BUILD" --target \ + core_knn_diskann \ + zvec_embed_db_example \ + zvec_embed_core_example \ + zvec_embed_ailego_example \ + --parallel "$NPROC" + shell: bash + + - name: Run subproject examples + run: | + cd "$SMOKE_BUILD" + ./bin/zvec_embed_db_example + ./bin/zvec_embed_core_example + ./bin/zvec_embed_ailego_example + shell: bash diff --git a/.github/workflows/_build_wheel_job.yml b/.github/workflows/_build_wheel_job.yml new file mode 100644 index 0000000..5eec8a1 --- /dev/null +++ b/.github/workflows/_build_wheel_job.yml @@ -0,0 +1,104 @@ +name: "(Reusable) Build, Publish and Smoke-test a Wheel" + +on: + workflow_call: + inputs: + runner: + description: "GitHub Actions runner label" + required: true + type: string + pypi_repository_url: + description: "PyPI repository URL (empty string means official PyPI)" + required: false + type: string + default: "" + secrets: + PYPI_API_TOKEN: + required: true + +jobs: + build_publish_test: + name: Build / publish / smoke-test on ${{ inputs.runner }} + runs-on: ${{ inputs.runner }} + permissions: + contents: read + + steps: + - name: Checkout code + uses: actions/checkout@v7 + with: + submodules: recursive + fetch-tags: true + fetch-depth: 0 + + - name: Set up Python (for cibuildwheel controller) + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install cibuildwheel + run: | + pip install --upgrade pip + pip install cibuildwheel==3.4.0 + + - name: Build wheels using cibuildwheel + run: | + python -m cibuildwheel --output-dir wheelhouse + # Save list of built wheels for publishing + ls wheelhouse/*.whl | tee $GITHUB_STEP_SUMMARY + echo "wheels=$(ls wheelhouse/*.whl | tr '\n' ' ')" >> $GITHUB_ENV + + - name: Publish to PyPI + if: success() && github.event_name == 'workflow_dispatch' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + TWINE_REPOSITORY_URL: ${{ inputs.pypi_repository_url }} + run: | + pip install twine + twine upload --skip-existing --verbose wheelhouse/*.whl + + - name: Smoke test from PyPI + if: success() && github.event_name == 'workflow_dispatch' + shell: bash + env: + PYPI_REPOSITORY_URL: ${{ inputs.pypi_repository_url }} + run: | + # Extract version from wheel filename (e.g. zvec-0.2.1.dev24-cp311-...whl -> 0.2.1.dev24) + WHEEL_FILE=$(ls wheelhouse/zvec-*.whl | head -1) + ZVEC_VERSION=$(basename "$WHEEL_FILE" | sed 's/zvec-\([^-]*\)-.*/\1/') + + # Build index-url flags: use TestPyPI when repository URL is set, otherwise official PyPI + if [ -n "$PYPI_REPOSITORY_URL" ]; then + INDEX_FLAGS="--index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/" + echo "Waiting for zvec==$ZVEC_VERSION to become available on TestPyPI..." + else + INDEX_FLAGS="" + echo "Waiting for zvec==$ZVEC_VERSION to become available on PyPI..." + fi + # Poll until the version is available (max 5 minutes) + FOUND=0 + for i in $(seq 1 30); do + if pip install $INDEX_FLAGS --dry-run "zvec==$ZVEC_VERSION" > /dev/null 2>&1; then + echo "Version $ZVEC_VERSION is available." + FOUND=1 + break + fi + echo "Attempt $i/30: not yet available, retrying in 10s..." + sleep 10 + done + + if [ "$FOUND" -eq 0 ]; then + echo "ERROR: Timed out (5 min) waiting for zvec==$ZVEC_VERSION on PyPI. Aborting smoke test." + exit 1 + fi + + # Create a clean venv and install + python -m venv test_env + source test_env/bin/activate + pip install --upgrade pip + pip install $INDEX_FLAGS "zvec==$ZVEC_VERSION" + pip install --upgrade pip + pip install $INDEX_FLAGS "zvec==$ZVEC_VERSION" + # Run a simple smoke test + python -c "import zvec; print('Import OK:', zvec.__version__)" diff --git a/.github/workflows/auto-assign.yml b/.github/workflows/auto-assign.yml new file mode 100644 index 0000000..7b73f1a --- /dev/null +++ b/.github/workflows/auto-assign.yml @@ -0,0 +1,25 @@ +name: Auto Assign PR + +on: + pull_request_target: + types: [opened, reopened, ready_for_review] + +permissions: + contents: read + pull-requests: write + +jobs: + auto-assign: + runs-on: ubuntu-latest + steps: + - name: Checkout for config + uses: actions/checkout@v7 + with: + ref: ${{ github.event.pull_request.base.ref }} + sparse-checkout: .github/auto-assign-config.yml + sparse-checkout-cone-mode: false + + - name: Auto Assign PR Author + uses: kentaro-m/auto-assign-action@v2.0.2 + with: + configuration-path: '.github/auto-assign-config.yml' diff --git a/.github/workflows/clang_tidy.yml b/.github/workflows/clang_tidy.yml new file mode 100644 index 0000000..20a302d --- /dev/null +++ b/.github/workflows/clang_tidy.yml @@ -0,0 +1,197 @@ +name: Clang-Tidy + +on: + workflow_call: + workflow_dispatch: + +permissions: + contents: read + +jobs: + clang_tidy: + name: Clang-Tidy Checks + runs-on: ubuntu-24.04 + steps: + - name: Shallow checkout (no submodules) + uses: actions/checkout@v7 + with: + submodules: false + fetch-depth: 0 + + - name: Collect changed C/C++ files + id: changed_files + uses: tj-actions/changed-files@v47 + with: + files: | + **/*.c + **/*.cc + **/*.cpp + **/*.cxx + **/*.h + **/*.hpp + files_ignore: | + thirdparty/** + build/** + src/db/sqlengine/antlr/gen/** + src/db/index/column/fts_column/gen/** + + - name: No C/C++ files changed - skip + if: steps.changed_files.outputs.any_changed != 'true' + run: echo "No C/C++ files changed. Skipping clang-tidy." + + - name: Checkout submodules + if: steps.changed_files.outputs.any_changed == 'true' + run: git submodule update --init --recursive --depth 1 + + - name: Install dependencies + if: steps.changed_files.outputs.any_changed == 'true' + run: | + sudo apt-get update + sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev libaio-dev + + - name: Setup ccache + if: steps.changed_files.outputs.any_changed == 'true' + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: clang-tidy + max-size: 500M + + - name: Configure CMake and export compile commands + if: steps.changed_files.outputs.any_changed == 'true' + run: | + cmake -S . -B build -G Ninja \ + -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \ + -DBUILD_TOOLS=ON \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + + - name: Filter changed files against compile_commands.json + id: tidy_files + if: steps.changed_files.outputs.any_changed == 'true' + run: | + python3 - <<'PY' + import json + import os + from pathlib import Path + + raw = os.environ.get("ALL_CHANGED_FILES", "") + changed = [f for f in raw.split() if f] + + compile_db_path = Path("build/compile_commands.json") + with compile_db_path.open("r", encoding="utf-8") as fh: + compile_db = json.load(fh) + + compile_entries = set() + for entry in compile_db: + file_field = entry.get("file", "") + if file_field: + compile_entries.add(os.path.normpath(file_field)) + + cwd = Path.cwd().resolve() + selected = [] + skipped = [] + + for rel_path in changed: + abs_path = os.path.normpath(str((cwd / rel_path).resolve())) + if abs_path in compile_entries: + selected.append(rel_path) + elif not Path(rel_path).is_file(): + skipped.append(f"{rel_path} (file not found)") + else: + skipped.append(f"{rel_path} (not in compile_commands.json)") + + github_output = os.environ["GITHUB_OUTPUT"] + with open(github_output, "a", encoding="utf-8") as out: + out.write(f"any_tidy_files={'true' if selected else 'false'}\n") + out.write("all_tidy_files<> "$GITHUB_OUTPUT" + + - name: Restore generated headers cache + id: cache_headers + if: steps.changed_files.outputs.any_changed == 'true' + uses: actions/cache/restore@v6 + with: + path: | + build/external + build/thirdparty + build/src/db/proto + key: clang-tidy-headers-${{ runner.os }}-${{ hashFiles('thirdparty/**/*.cmake', 'thirdparty/**/CMakeLists.txt', 'thirdparty/**/*.patch', 'src/db/proto/*.proto') }}-${{ steps.submodule_hash.outputs.value }} + + - name: Build generated headers only + if: steps.tidy_files.outputs.any_tidy_files == 'true' && steps.cache_headers.outputs.cache-hit != 'true' + run: | + ninja -C build clang_tidy_deps + + - name: Save generated headers cache + if: steps.tidy_files.outputs.any_tidy_files == 'true' && steps.cache_headers.outputs.cache-hit != 'true' + uses: actions/cache/save@v6 + with: + path: | + build/external + build/thirdparty + build/src/db/proto + key: clang-tidy-headers-${{ runner.os }}-${{ hashFiles('thirdparty/**/*.cmake', 'thirdparty/**/CMakeLists.txt', 'thirdparty/**/*.patch', 'src/db/proto/*.proto') }}-${{ steps.submodule_hash.outputs.value }} + + - name: Run clang-tidy on changed files (parallel) + if: steps.tidy_files.outputs.any_tidy_files == 'true' + run: | + mapfile -t files_to_check <<'TIDY_EOF' + ${{ steps.tidy_files.outputs.all_tidy_files }} + TIDY_EOF + + log_dir=$(mktemp -d) + printf '%s\n' "${files_to_check[@]}" \ + | grep -v '^\s*$' \ + | xargs -I{} -P "$(nproc)" sh -c ' + file="{}" + if [ -f "$file" ]; then + log="'"$log_dir"'/$$.log" + echo "$file" > "$log" + if clang-tidy -p build --quiet --warnings-as-errors="*" "$file" >> "$log" 2>&1; then + echo "PASS: $file" + rm -f "$log" + else + echo "FAIL: $file" + fi + fi' + + failed=0 + for f in "$log_dir"/*.log; do + [ -e "$f" ] || break + failed=1 + src=$(head -1 "$f") + echo "" + echo "::group::clang-tidy errors: $src" + tail -n +2 "$f" + echo "::endgroup::" + done + + rm -rf "$log_dir" + if [ "$failed" -eq 1 ]; then + echo "::error::clang-tidy found issues in one or more files" + exit 1 + fi + + - name: No files to analyse + if: steps.changed_files.outputs.any_changed == 'true' && steps.tidy_files.outputs.any_tidy_files != 'true' + run: echo "Changed C/C++ files not in compile_commands.json. Nothing to analyse." diff --git a/.github/workflows/community-pr-handler.yml b/.github/workflows/community-pr-handler.yml new file mode 100644 index 0000000..2296c87 --- /dev/null +++ b/.github/workflows/community-pr-handler.yml @@ -0,0 +1,148 @@ +name: Community PR Handler + +on: + pull_request_target: + types: [opened, reopened, ready_for_review] + +permissions: + contents: read + pull-requests: write + +jobs: + handle-community-pr: + runs-on: ubuntu-latest + if: github.event.pull_request.draft == false + steps: + - name: Check if community contributor + id: check + uses: actions/github-script@v9 + with: + script: | + const pr = context.payload.pull_request; + const prAuthor = pr.user.login; + const association = pr.author_association; + + console.log('========== PR Info =========='); + console.log(`PR Number: #${pr.number}`); + console.log(`PR Title: ${pr.title}`); + console.log(`PR Author: ${prAuthor}`); + console.log(`Author Association: ${association}`); + console.log(`Base Branch: ${pr.base.ref}`); + console.log(`Head Branch: ${pr.head.ref}`); + console.log(`Head Repo: ${pr.head.repo?.full_name || 'same repo'}`); + + // OWNER, MEMBER, COLLABORATOR are maintainers; others are community contributors + const maintainerRoles = ['OWNER', 'MEMBER', 'COLLABORATOR']; + const isCommunity = !maintainerRoles.includes(association); + + console.log('========== Decision =========='); + console.log(`Maintainer Roles: ${maintainerRoles.join(', ')}`); + console.log(`Is Community Contributor: ${isCommunity}`); + + if (!isCommunity) { + console.log('Skipping: PR author is a maintainer, no action needed.'); + } + + core.setOutput('is_community', isCommunity); + + - name: Checkout base branch for CODEOWNERS + if: steps.check.outputs.is_community == 'true' + uses: actions/checkout@v7 + with: + ref: ${{ github.event.pull_request.base.ref }} + sparse-checkout: .github/CODEOWNERS + sparse-checkout-cone-mode: false + + - name: Assign reviewer as assignee + if: steps.check.outputs.is_community == 'true' + uses: actions/github-script@v9 + with: + script: | + const fs = require('fs'); + const prNumber = context.payload.pull_request.number; + + console.log('========== Fetching Changed Files =========='); + // Get changed files + const { data: files } = await github.rest.pulls.listFiles({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: prNumber + }); + const changedFiles = files.map(f => f.filename); + console.log(`Total changed files: ${changedFiles.length}`); + changedFiles.forEach((f, i) => console.log(` [${i + 1}] ${f}`)); + + console.log('========== Parsing CODEOWNERS =========='); + // Parse CODEOWNERS + let codeowners = []; + try { + const content = fs.readFileSync('.github/CODEOWNERS', 'utf8'); + const lines = content.split('\n').filter(l => l.trim() && !l.trim().startsWith('#')); + console.log(`Found ${lines.length} rules in CODEOWNERS`); + for (const line of lines) { + const parts = line.trim().split(/\s+/); + if (parts.length >= 2) { + const pattern = parts[0]; + const owners = parts.slice(1).map(o => o.replace('@', '')); + codeowners.push({ pattern, owners }); + console.log(` Rule: "${pattern}" -> [${owners.join(', ')}]`); + } + } + } catch (e) { + console.log('ERROR: Could not read CODEOWNERS:', e.message); + return; + } + + console.log('========== Matching Files to Owners =========='); + // Find matching owners for changed files + const matchedOwners = new Set(); + for (const file of changedFiles) { + let matchedOwner = null; + let matchedPattern = null; + for (const rule of codeowners) { + const pattern = rule.pattern; + if (pattern === '*') { + matchedOwner = rule.owners[0]; + matchedPattern = pattern; + } else if (pattern.endsWith('/')) { + const dir = pattern.replace(/^\//, '').replace(/\/$/, ''); + if (file.startsWith(dir + '/') || file.startsWith(dir)) { + matchedOwner = rule.owners[0]; + matchedPattern = pattern; + } + } else if (file === pattern || file === pattern.replace(/^\//, '')) { + matchedOwner = rule.owners[0]; + matchedPattern = pattern; + } + } + if (matchedOwner) { + matchedOwners.add(matchedOwner); + console.log(` "${file}" -> matched "${matchedPattern}" -> @${matchedOwner}`); + } else { + console.log(` "${file}" -> NO MATCH`); + } + } + + console.log('========== Setting Assignees =========='); + // Set assignees + const assignees = Array.from(matchedOwners); + console.log(`Assignees to set: [${assignees.join(', ')}]`); + + if (assignees.length > 0) { + try { + await github.rest.issues.addAssignees({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: prNumber, + assignees: assignees + }); + console.log('SUCCESS: Assignees set successfully!'); + } catch (e) { + console.log('ERROR: Failed to set assignees:', e.message); + console.log('Error details:', JSON.stringify(e, null, 2)); + } + } else { + console.log('WARNING: No assignees to set - no matching owners found.'); + } + + console.log('========== Done =========='); diff --git a/.github/workflows/continuous_bench.yml b/.github/workflows/continuous_bench.yml new file mode 100644 index 0000000..c37b2aa --- /dev/null +++ b/.github/workflows/continuous_bench.yml @@ -0,0 +1,33 @@ +name: Continuous Benchmark +on: + # Temporarily disabled. Restore original triggers below to re-enable. + # push: + # branches: [ "main"] + # paths-ignore: + # - '**.md' + # workflow_dispatch: + workflow_dispatch: + inputs: + __disabled__: + description: 'This workflow is temporarily disabled.' + required: false + +concurrency: + group: cb-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + benchmark: + if: false + runs-on: vdbbench + steps: + - uses: actions/checkout@v7 + + - name: Run VectorDBBench + env: + DATABASE_URL: ${{ secrets.DATABASE_URL }} + run: | + bash .github/workflows/scripts/run_vdb.sh \ No newline at end of file diff --git a/.github/workflows/docker/Dockerfile.linux_x64_glibc228 b/.github/workflows/docker/Dockerfile.linux_x64_glibc228 new file mode 100644 index 0000000..503829e --- /dev/null +++ b/.github/workflows/docker/Dockerfile.linux_x64_glibc228 @@ -0,0 +1,83 @@ +# ============================================================================= +# Dockerfile.linux_x64_glibc228 +# Purpose: Ubuntu 18.10 gcc-9 + glibc 2.28 + CMake 3.30.0 + PyBind11 build environment +# Warning: ubuntu:18.10 is EOL; use only for glibc 2.28 compatibility testing. +# ============================================================================= + +# Use official Ubuntu 18.10 (Cosmic Cuttlefish) +# glibc version: 2.28 (confirmed via `ldd --version`) +FROM ubuntu:18.10 + +# Replace Ubuntu mirror with old-releases.ubuntu.com for older glibc compatibility +RUN sed -i 's|http://\(.*\)/ubuntu|http://old-releases.ubuntu.com/ubuntu|g' /etc/apt/sources.list && \ + sed -i 's|http://security.ubuntu.com/ubuntu|http://old-releases.ubuntu.com/ubuntu|g' /etc/apt/sources.list + +# Add Ubuntu 20.04 (focal) repo for GCC 9 ONLY +RUN echo "deb http://archive.ubuntu.com/ubuntu/ focal main universe" >> /etc/apt/sources.list && \ + echo "deb http://security.ubuntu.com/ubuntu/ focal-security main universe" >> /etc/apt/sources.list + +# Prevent interactive prompts & set non-root user +ENV DEBIAN_FRONTEND=noninteractive \ + TZ=Etc/UTC + +# Create non-root user for safety (optional but recommended) +RUN useradd -m -u 1000 builder && \ + mkdir -p /workspace && chown builder:builder /workspace + +# Install base system dependencies +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + build-essential \ + gcc-9 g++-9 \ + ninja-build git curl ca-certificates vim wget lcov gnupg clang-format-18\ + rsync lsb-release \ + uuid-dev zlib1g-dev libssl-dev libffi-dev \ + pybind11-dev && \ + update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 90 \ + --slave /usr/bin/g++ g++ /usr/bin/g++-9 && \ + rm -rf /var/lib/apt/lists/* + +# Install Miniforge (Conda) as root, then assign to builder +ENV MINIFORGE_VERSION="latest" +ENV MINIFORGE_HOME="/opt/miniforge3" + +RUN curl -sSL "https://github.com/conda-forge/miniforge/releases/${MINIFORGE_VERSION}/download/Miniforge3-Linux-x86_64.sh" -o miniforge.sh && \ + bash miniforge.sh -b -p ${MINIFORGE_HOME} && \ + rm miniforge.sh && \ + chown -R builder:builder ${MINIFORGE_HOME} + +# Switch to non-root user +USER builder +ENV PATH="${MINIFORGE_HOME}/bin:${PATH}" +WORKDIR /workspace + +# Create conda envs for supported Python versions +RUN conda create -n py310 python=3.10 -y && \ + conda create -n py311 python=3.11 -y && \ + conda create -n py312 python=3.12 -y +RUN conda clean --all -f -y + +# Install CMake 3.30.0 from Kitware official binary +# Ref: https://github.com/Kitware/CMake/releases/tag/v3.30.0 +RUN mkdir -p /tmp/cmake && cd /tmp/cmake && \ + curl -sSL -o cmake.tar.gz \ + "https://github.com/Kitware/CMake/releases/download/v3.30.0/cmake-3.30.0-linux-x86_64.tar.gz" && \ + tar -xzf cmake.tar.gz --strip-components=1 -C /tmp/cmake && \ + mkdir -p /home/builder/.local && \ + mv * /home/builder/.local/ && \ + chown -R builder:builder /home/builder/.local && \ + rm -rf /tmp/cmake + +# Add CMake to PATH +ENV PATH="/home/builder/.local/bin:${PATH}" + +# Verify installations +RUN cmake --version && \ + conda info && \ + conda env list && \ + python --version && \ + gcc --version && \ + ldd --version | head -n1 + +# Final setup +WORKDIR /workspace \ No newline at end of file diff --git a/.github/workflows/issue-auto-assign.yml b/.github/workflows/issue-auto-assign.yml new file mode 100644 index 0000000..bbd1334 --- /dev/null +++ b/.github/workflows/issue-auto-assign.yml @@ -0,0 +1,99 @@ +name: Issue Auto Assign + +on: + issues: + types: [opened] + +jobs: + auto-assign: + runs-on: ubuntu-latest + permissions: + issues: write + + steps: + - name: Parse issue and assign + uses: actions/github-script@v9 + with: + script: | + const issue = context.payload.issue; + const issueBody = issue.body || ''; + const issueLabels = (issue.labels || []).map(l => l.name); + + // Default assignees per template type (based on labels) + const defaultAssignees = { + 'bug': 'zhourrr', + 'feature': 'feihongxu0824', + 'benchmark': 'egolearner', + 'enhancement': 'feihongxu0824', + 'integration': 'chinaux', + 'profile': 'richyreachy' + }; + + // Global fallback assignee + const fallbackAssignee = 'feihongxu0824'; + + // Parse user-selected assignee from issue body + // The input field renders as: "### Preferred Assignee\n\n" + let selectedAssignee = null; + const assigneeMatch = issueBody.match(/### Preferred Assignee\s*\n+([^\n#]+)/); + if (assigneeMatch) { + const selection = assigneeMatch[1].trim(); + console.log(`Parsed assignee input: "${selection}"`); + // If user entered a valid GitHub username (not empty, not placeholder text) + if (selection && + selection !== '_No response_' && + selection !== 'None' && + !selection.toLowerCase().includes('leave empty') && + !selection.startsWith('e.g.,')) { + // Clean up the username (remove @ if present) + selectedAssignee = selection.replace(/^@/, '').trim(); + } + } + + // Determine final assignee + let finalAssignee = selectedAssignee; + + // If no user selection, use default based on label + if (!finalAssignee && issueLabels.length > 0) { + for (const [label, assignee] of Object.entries(defaultAssignees)) { + if (issueLabels.includes(label)) { + finalAssignee = assignee; + console.log(`Matched label "${label}" -> assignee "${assignee}"`); + break; + } + } + } + + // Fallback to default assignee if no match + if (!finalAssignee) { + finalAssignee = fallbackAssignee; + console.log(`No match found, using fallback assignee: ${fallbackAssignee}`); + } + + console.log(`Issue #${issue.number}: Labels = [${issueLabels.join(', ')}]`); + console.log(`User selected assignee: ${selectedAssignee || 'None (Auto)'}`); + console.log(`Final assignee: ${finalAssignee}`); + + // Assign the issue + try { + await github.rest.issues.addAssignees({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue.number, + assignees: [finalAssignee] + }); + console.log(`Successfully assigned issue #${issue.number} to ${finalAssignee}`); + } catch (error) { + console.error(`Failed to assign issue: ${error.message}`); + // If assignment fails (user may not have permission), add a comment + try { + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue.number, + body: `⚠️ Auto-assignment to \`${finalAssignee}\` failed. Please assign manually.` + }); + } catch (commentError) { + console.error(`Failed to create comment: ${commentError.message}`); + } + } diff --git a/.github/workflows/nightly_coverage.yml b/.github/workflows/nightly_coverage.yml new file mode 100644 index 0000000..81429b1 --- /dev/null +++ b/.github/workflows/nightly_coverage.yml @@ -0,0 +1,102 @@ +name: Nightly Coverage Report + +on: + schedule: + # Runs daily at 00:00 CST (China Standard Time) = 16:00 UTC + - cron: '0 16 * * *' + + workflow_dispatch: + +permissions: + contents: read + +jobs: + coverage: + name: Nightly Coverage Report + runs-on: ubuntu-24.04 + + strategy: + matrix: + python-version: ['3.10'] + fail-fast: false + + steps: + - name: Checkout code + uses: actions/checkout@v7 + with: + ref: main # Always use main for nightly + submodules: recursive + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + cache-dependency-path: 'pyproject.toml' + + - name: Set up environment variables + run: | + # Set number of processors for parallel builds + NPROC=$(nproc 2>/dev/null || echo 2) + echo "NPROC=$NPROC" >> $GITHUB_ENV + echo "Using $NPROC parallel jobs for builds" + + # Add Python user base bin to PATH for pip-installed CLI tools + echo "$(python -c 'import site; print(site.USER_BASE)')/bin" >> $GITHUB_PATH + shell: bash + + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + lcov libaio-dev + shell: bash + + - name: Install dependencies + run: | + python -m pip install --upgrade pip \ + pybind11==3.0 \ + cmake==3.30.0 \ + ninja==1.11.1 \ + pytest \ + pytest-cov \ + pytest-xdist \ + scikit-build-core \ + setuptools_scm + shell: bash + + - name: Build with COVERAGE config + run: | + cd "$GITHUB_WORKSPACE" + + CMAKE_GENERATOR="Unix Makefiles" \ + CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \ + python -m pip install -v . \ + --no-build-isolation \ + --config-settings="cmake.build-type=COVERAGE" + shell: bash + + - name: Run Python Tests with Coverage + run: | + cd "$GITHUB_WORKSPACE" + python -m pytest python/tests/ --cov=zvec --cov-report=xml \ + --deselect=python/tests/test_gil_release.py::TestGILRelease::test_gil_released_during_query + shell: bash + + - name: Run C++ Tests and Generate Coverage + run: | + cd "$GITHUB_WORKSPACE/build" + make unittest -j$NPROC + cd "$GITHUB_WORKSPACE" + # Ensure gcov.sh is executable + chmod +x scripts/gcov.sh + bash scripts/gcov.sh -k + shell: bash + + - name: Upload Coverage to Codecov + uses: codecov/codecov-action@v7 + with: + files: ./proxima-zvec-filtered.lcov.info,./coverage.xml + flags: python,cpp,nightly + name: nightly-linux-py${{ matrix.python-version }} + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/.github/workflows/scripts/run_vdb.sh b/.github/workflows/scripts/run_vdb.sh new file mode 100644 index 0000000..f153a59 --- /dev/null +++ b/.github/workflows/scripts/run_vdb.sh @@ -0,0 +1,88 @@ +set -e + +QUANTIZE_TYPE_LIST="int8 int4 fp16 fp32" +CASE_TYPE_LIST="Performance768D1M Performance768D10M Performance1536D500K" # respectively test cosine, ip # Performance960D1M l2 metrics +LOG_FILE="bench.log" +DATE=$(date +%Y-%m-%d_%H-%M-%S) +NPROC=$(nproc 2>/dev/null || getconf _NPROCESSORS_ONLN 2>/dev/null || echo 2) + +# COMMIT_ID = branch-date-sha +COMMIT_ID=${GITHUB_REF_NAME}-"$DATE"-$(echo ${GITHUB_WORKFLOW_SHA} | cut -c1-8) +COMMIT_ID=$(echo "$COMMIT_ID" | sed 's/\//_/g') +echo "COMMIT_ID: $COMMIT_ID" +echo "GITHUB_WORKFLOW_SHA: $GITHUB_WORKFLOW_SHA" +echo "workspace: $GITHUB_WORKSPACE" +DB_LABEL_PREFIX="Zvec16c64g-$COMMIT_ID" + +# install zvec +git submodule update --init + +# for debug +#cd .. +#export SKBUILD_BUILD_DIR="$GITHUB_WORKSPACE/../build" +pwd + +python3 -m venv .venv +source .venv/bin/activate +pip install cmake ninja psycopg2-binary loguru fire +pip install -e /opt/VectorDBBench + +CMAKE_GENERATOR="Unix Makefiles" \ +CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \ +pip install -v "$GITHUB_WORKSPACE" + +for CASE_TYPE in $CASE_TYPE_LIST; do + echo "Running VectorDBBench for $CASE_TYPE" + DATASET_DESC="" + if [ "$CASE_TYPE" == "Performance768D1M" ]; then + DATASET_DESC="Performance768D1M - Cohere Cosine" + elif [ "$CASE_TYPE" == "Performance768D10M" ]; then + DATASET_DESC="Performance768D10M - Cohere Cosine" + else + DATASET_DESC="Performance1536D500K - OpenAI IP" + fi + + for QUANTIZE_TYPE in $QUANTIZE_TYPE_LIST; do + DB_LABEL="$DB_LABEL_PREFIX-$CASE_TYPE-$QUANTIZE_TYPE" + echo "Running VectorDBBench for $DB_LABEL" + + VDB_PARAMS="--path ${DB_LABEL} --db-label ${DB_LABEL} --case-type ${CASE_TYPE} --num-concurrency 12,14,16,18,20" + if [ "$CASE_TYPE" == "Performance768D1M" ]; then + VDB_PARAMS="${VDB_PARAMS} --m 15 --ef-search 180" + elif [ "$CASE_TYPE" == "Performance768D10M" ]; then + VDB_PARAMS="${VDB_PARAMS} --m 50 --ef-search 118 --is-using-refiner" + else #Performance1536D500K using default params + refiner to monitor performance degradation + VDB_PARAMS="${VDB_PARAMS} --m 50 --ef-search 100 --is-using-refiner" + fi + + if [ "$QUANTIZE_TYPE" == "fp32" ]; then + vectordbbench zvec ${VDB_PARAMS} 2>&1 | tee $LOG_FILE + else + vectordbbench zvec ${VDB_PARAMS} --quantize-type "${QUANTIZE_TYPE}" 2>&1 | tee $LOG_FILE + fi + + RESULT_JSON_PATH=$(grep -o "/opt/VectorDBBench/.*\.json" $LOG_FILE) + QPS=$(jq -r '.results[0].metrics.qps' "$RESULT_JSON_PATH") + RECALL=$(jq -r '.results[0].metrics.recall' "$RESULT_JSON_PATH") + LATENCY_P99=$(jq -r '.results[0].metrics.serial_latency_p99' "$RESULT_JSON_PATH") + LOAD_DURATION=$(jq -r '.results[0].metrics.load_duration' "$RESULT_JSON_PATH") + + #quote the var to avoid space in the label + label_list="case_type=\"${CASE_TYPE}\",dataset_desc=\"${DATASET_DESC}\",db_label=\"${DB_LABEL}\",commit=\"${COMMIT_ID}\",date=\"${DATE}\",quantize_type=\"${QUANTIZE_TYPE}\"" + # replace `/` with `_` in label_list + label_list=$(echo "$label_list" | sed 's/\//_/g') + cat < prom_metrics.txt + # TYPE vdb_bench_qps gauge + vdb_bench_qps{$label_list} $QPS + # TYPE vdb_bench_recall gauge + vdb_bench_recall{$label_list} $RECALL + # TYPE vdb_bench_latency_p99 gauge + vdb_bench_latency_p99{$label_list} $LATENCY_P99 + # TYPE vdb_bench_load_duration gauge + vdb_bench_load_duration{$label_list} $LOAD_DURATION +EOF + echo "prom_metrics:" + cat prom_metrics.txt + curl --data-binary @prom_metrics.txt "http://47.93.34.27:9091/metrics/job/benchmarks-${CASE_TYPE}/case_type/${CASE_TYPE}/quantize_type/${QUANTIZE_TYPE}" -v + done +done \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..65859ca --- /dev/null +++ b/.gitignore @@ -0,0 +1,68 @@ +.* +*~ +bazel-* +build* +bin/* +lib/* +var/* +venv* +tests/integration/conf/* +tests/de_integration/conf/* +**/__pycache__/* +tests/bench/log/* +tests/integration/integration +tests/integration/log +tests/integration/*.log +tests/de_integration/log +tests/de_integration/*.log +!.git* +!.clang-format +!.clang-tidy +!.circleci +!.drone.yml +sdk/python/dist/ +compile_commands.json +dist +html +*.lcov.info + +# Compiled extension modules must never live in the source package: they are +# built by CMake and installed into the wheel. A stray copy left here gets +# bundled verbatim by scikit-build's wheel.packages and breaks the wheel +# (e.g. a stale binary with the wrong macOS deployment target). +python/zvec/*.so +python/zvec/*.dylib +python/zvec/*.pyd + +# Dependencies +/node_modules + +# Production +/build + +# Generated files +.docusaurus +.cache-loader + +# Misc +.DS_Store +.env.local +.env.development.local +.env.test.local +.env.production.local + +npm-debug.log* +yarn-debug.log* +yarn-error.log* + +allure-* + +!build_android.sh +!build_ios.sh + +# congfig +doc/ +config/ +examples/python/ +examples/c_api/ +logs/ \ No newline at end of file diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..f919c73 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,61 @@ +[submodule "thirdparty/googletest/googletest-1.10.0"] + path = thirdparty/googletest/googletest-1.10.0 + url = https://github.com/google/googletest.git + ignore = dirty +[submodule "thirdparty/sparsehash/sparsehash-2.0.4"] + path = thirdparty/sparsehash/sparsehash-2.0.4 + url = https://github.com/sparsehash/sparsehash.git + ignore = dirty +[submodule "thirdparty/gflags/gflags-2.2.2"] + path = thirdparty/gflags/gflags-2.2.2 + url = https://github.com/gflags/gflags.git +[submodule "thirdparty/rocksdb/rocksdb-8.1.1"] + path = thirdparty/rocksdb/rocksdb-8.1.1 + url = https://github.com/facebook/rocksdb.git + ignore = dirty +[submodule "thirdparty/yaml-cpp/yaml-cpp-0.6.3"] + path = thirdparty/yaml-cpp/yaml-cpp-0.6.3 + url = https://github.com/jbeder/yaml-cpp.git +[submodule "thirdparty/arrow/apache-arrow-21.0.0"] + path = thirdparty/arrow/apache-arrow-21.0.0 + url = https://github.com/apache/arrow.git + ignore = dirty +[submodule "thirdparty/glog/glog-0.5.0"] + path = thirdparty/glog/glog-0.5.0 + url = https://github.com/google/glog.git + ignore = dirty +[submodule "thirdparty/protobuf/protobuf-3.21.12"] + path = thirdparty/protobuf/protobuf-3.21.12 + url = https://github.com/protocolbuffers/protobuf.git +[submodule "thirdparty/lz4/lz4-1.9.4"] + path = thirdparty/lz4/lz4-1.9.4 + url = https://github.com/lz4/lz4.git +[submodule "thirdparty/antlr/antlr4"] + path = thirdparty/antlr/antlr4 + url = https://github.com/antlr/antlr4.git + ignore = dirty +[submodule "thirdparty/magic_enum/magic_enum-0.9.7"] + path = thirdparty/magic_enum/magic_enum-0.9.7 + url = https://github.com/Neargye/magic_enum.git + ignore = dirty +[submodule "thirdparty/RaBitQ-Library/RaBitQ-Library-0.1"] + path = thirdparty/RaBitQ-Library/RaBitQ-Library-0.1 + url = https://github.com/VectorDB-NTU/RaBitQ-Library.git +[submodule "thirdparty/aio/libaio-0.3"] + path = thirdparty/aio/libaio-0.3 + url = https://github.com/yugabyte/libaio.git +[submodule "thirdparty/cppjieba/cppjieba-5.6.7"] + path = thirdparty/cppjieba/cppjieba-5.6.7 + url = https://github.com/yanyiwu/cppjieba.git +[submodule "thirdparty/FastPFOR/FastPFOR-0.4.0"] + path = thirdparty/FastPFOR/FastPFOR-0.4.0 + url = https://github.com/fast-pack/FastPFOR.git +[submodule "thirdparty/limonp/limonp-v1.0.2"] + path = thirdparty/limonp/limonp-v1.0.2 + url = https://github.com/yanyiwu/limonp.git +[submodule "thirdparty/utf8proc/utf8proc-2.11.3"] + path = thirdparty/utf8proc/utf8proc-2.11.3 + url = https://github.com/JuliaStrings/utf8proc.git +[submodule "thirdparty/snowball/snowball-3.1.1"] + path = thirdparty/snowball/snowball-3.1.1 + url = https://github.com/snowballstem/snowball.git diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..d334b85 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,211 @@ +cmake_minimum_required(VERSION 3.26) +cmake_policy(SET CMP0077 NEW) +project(zvec) +set(CC_CXX_STANDARD 17) + +if(MSVC) + set(INTTYPES_FORMAT VC7) + # Use /Z7 (embed debug info in .obj) instead of /Zi (separate .pdb) to enable sccache/ccache. + # Use CACHE FORCE so it propagates into thirdparty subdirectories that call project(). + foreach(_cfg DEBUG RELWITHDEBINFO) + foreach(_lang C CXX) + string(REPLACE "/Zi" "/Z7" _flags "${CMAKE_${_lang}_FLAGS_${_cfg}}") + set(CMAKE_${_lang}_FLAGS_${_cfg} "${_flags}" CACHE STRING "" FORCE) + endforeach() + endforeach() + add_compile_options(/EHsc) # def c++ exception behavior + add_compile_options(/Zc:preprocessor /Zc:__cplusplus) + add_compile_options(/we4716) # -Werror=return-type + + ###### reduce output length to make vibe coding work better :) should be removed or solved later + # TODO(windows): fix&cleanup + add_compile_options(/wd4267 /wd4244 /wd4305 /wd4242) # conversion/truncation from 'size_t' to 'type', possible loss of data + add_compile_options(/wd4310) # cast truncates constant value + add_compile_options(/wd4146) # unary minus operator applied to unsigned type #usage: uint32_t seg_id_{-1U}; + add_compile_options(/wd4245) # conversion from 'type1' to 'type2', signed/unsigned mismatch + add_compile_options(/wd4251) # class 'type' needs to have dll-interface to be used by clients of class 'type2' + add_compile_options(/wd4099) # type name first seen using 'object_type1' now seen using 'object_type2' + add_compile_options(/wd4200) # nonstandard extension used : zero-sized array in struct/union + add_compile_options(/wd4324) # structure was padded due to alignment specifier + add_compile_options(/wd4702) # unreachable code + add_compile_options(/wd4996) # wstring_convert are deprecated in C++17 # usage: in antlr4 + add_compile_options(/wd4101) # unreferenced local variable + add_compile_options(/wd4996) # The POSIX name for this item is deprecated. Instead, use the ISO C and C++ conformant name: _lseek + # result of 32-bit shift implicitly converted to 64 bits (was 64-bit shift intended?) # usage: keeps |= (1 << k); + add_compile_options(/wd4334) + ###### + +else() + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Werror=return-type") + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Werror=return-type") +endif() + +if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND NOT IOS) + set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -Wl,--no-as-needed") + set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--no-as-needed") +endif() + +if(NOT DEFINED PROJECT_ROOT_DIR OR NOT PROJECT_ROOT_DIR) + set(PROJECT_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR} CACHE PATH "Root directory of the project" FORCE) +endif() + +message(STATUS "PROJECT_ROOT_DIR = ${PROJECT_ROOT_DIR}") + +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +# iOS platform detection +if(NOT ANDROID AND NOT IOS AND CMAKE_SYSTEM_NAME STREQUAL "iOS") + set(IOS TRUE) +endif() + +# iOS bundle properties for test executables +if(IOS) + set(MACOSX_BUNDLE_BUNDLE_VERSION "1") + set(MACOSX_BUNDLE_SHORT_VERSION_STRING "1.0") + set(CMAKE_MACOSX_BUNDLE_INFO_PLIST "${PROJECT_ROOT_DIR}/cmake/iOSBundleInfo.plist.in") +endif() + +if (NOT ANDROID AND NOT IOS AND AUTO_DETECT_ARCH AND HOST_ARCH MATCHES "^(x86|x64)$") + setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512 MATH_MARCH_FLAG_AVX512FP16) + message(STATUS "best compiler march, sse: " ${MATH_MARCH_FLAG_SSE} ", avx2: " ${MATH_MARCH_FLAG_AVX2} ", avx512: " ${MATH_MARCH_FLAG_AVX512} ", avx512fp16: " ${MATH_MARCH_FLAG_AVX512FP16}) +endif() + +include_directories(${PROJECT_ROOT_DIR}/src/include) +include_directories(${PROJECT_ROOT_DIR}/src) + +option(BUILD_ZVEC_SHARED "Build all-in-one C++ shared library libzvec" ON) +option(BUILD_ZVEC_AILEGO_SHARED "Build all-in-one zvec-ailego shared library libzvec_ailego" ON) +option(BUILD_ZVEC_CORE_SHARED "Build all-in-one zvec-core shared library libzvec_core" ON) + +option(BUILD_PYTHON_BINDINGS "Build Python bindings using pybind11" OFF) +option(BUILD_C_BINDINGS "Build C bindings" ON) +option(BUILD_TOOLS "Build tools" ON) + +message(STATUS "BUILD_ZVEC_SHARED:${BUILD_ZVEC_SHARED}") +message(STATUS "BUILD_ZVEC_AILEGO_SHARED:${BUILD_ZVEC_AILEGO_SHARED}") +message(STATUS "BUILD_ZVEC_CORE_SHARED:${BUILD_ZVEC_CORE_SHARED}") +message(STATUS "BUILD_PYTHON_BINDINGS:${BUILD_PYTHON_BINDINGS}") +message(STATUS "BUILD_C_BINDINGS:${BUILD_C_BINDINGS}") +message(STATUS "BUILD_TOOLS:${BUILD_TOOLS}") + +option(RABITQ_ENABLE_AVX512 "Compile RaBitQ with AVX-512 support" OFF) + +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64|AMD64" AND NOT ANDROID AND NOT IOS) + include(CheckCCompilerFlag) + + check_c_compiler_flag("-mavx2" COMPILER_SUPPORTS_AVX2) + check_c_compiler_flag("-mavx512f -mavx512bw -mavx512vl" COMPILER_SUPPORTS_AVX512) + + if(COMPILER_SUPPORTS_AVX2 OR COMPILER_SUPPORTS_AVX512) + set(RABITQ_SUPPORTED ON) + add_definitions(-DRABITQ_SUPPORTED=1) + if(RABITQ_ENABLE_AVX512 AND COMPILER_SUPPORTS_AVX512) + add_definitions(-DRABITQ_COMPILED_AVX512=1) + set(RABITQ_ARCH_FLAG "${MATH_MARCH_FLAG_AVX512}") + else() + set(RABITQ_ARCH_FLAG "${MATH_MARCH_FLAG_AVX2}") + endif() + else() + set(RABITQ_SUPPORTED OFF) + add_definitions(-DRABITQ_SUPPORTED=0) + message(STATUS "RaBitQ support disabled - compiler does not support AVX2 or AVX-512") + endif() +elseif(IOS) + set(RABITQ_SUPPORTED OFF) + add_definitions(-DRABITQ_SUPPORTED=0) + message(STATUS "RaBitQ support disabled - not supported on iOS") +else() + set(RABITQ_SUPPORTED OFF) + add_definitions(-DRABITQ_SUPPORTED=0) + message(STATUS "RaBitQ support disabled - only supported on Linux x86_64") +endif() +message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") + +# DiskAnn support (Linux x86_64 only, requires libaio) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS) + set(DISKANN_SUPPORTED ON) + add_definitions(-DDISKANN_SUPPORTED=1) +else() + set(DISKANN_SUPPORTED OFF) + add_definitions(-DDISKANN_SUPPORTED=0) + message(STATUS "DiskAnn support disabled - only supported on Linux x86_64") +endif() +message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}") + +option(USE_OSS_MIRROR "Use OSS mirror for faster third-party downloads" OFF) +if(DEFINED ENV{USE_OSS_MIRROR} AND NOT "$ENV{USE_OSS_MIRROR}" STREQUAL "") + set(USE_OSS_MIRROR "$ENV{USE_OSS_MIRROR}" CACHE BOOL "Use OSS mirror for faster third-party downloads" FORCE) +endif() +message(STATUS "USE_OSS_MIRROR:${USE_OSS_MIRROR}") + +cc_directory(thirdparty) +cc_directories(src) + +cc_directories(tests) + +add_custom_target(clang_tidy_deps DEPENDS zvec_proto ARROW.BUILD glog gflags Lz4.BUILD) + +if(BUILD_TOOLS) + cc_directories(tools) +endif() + +git_version(GIT_SRCS_VER ${PROJECT_ROOT_DIR}) +set(CPACK_PACKAGE_VERSION ${GIT_SRCS_VER}) +set(CPACK_PACKAGE_NAME zvec) +include(CPack) + +if(BUILD_PYTHON_BINDINGS) + if(APPLE) + set(CMAKE_STRIP "") + message(STATUS "Disabled strip on macOS to preserve code signature") + endif() + + include(GNUInstallDirs) + if(DEFINED SKBUILD_PLATLIB_DIR) + set(ZVEC_PY_INSTALL_DIR "${SKBUILD_PLATLIB_DIR}") + elseif(DEFINED Python_SITEARCH) + set(ZVEC_PY_INSTALL_DIR "${Python_SITEARCH}") + else() + set(ZVEC_PY_INSTALL_DIR "${CMAKE_INSTALL_LIBDIR}") + endif() + + message(STATUS "Zvec install path: ${ZVEC_PY_INSTALL_DIR}") + # COMPONENT python: only these runtime artifacts are pulled into the wheel + # (see install.components in pyproject.toml). The cc_library(PACKED) SDK + # install rules use the default component and are excluded from the wheel. + # Install the extension inside the zvec package (zvec/_zvec*.so) rather than + # at the site-packages root, so it does not pollute the top-level namespace. + # The Python code imports it as `zvec._zvec` accordingly. + install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec + COMPONENT python) + + # DiskAnn ships as a runtime-loaded shared module + # (libzvec_diskann_plugin.so) that is brought online implicitly the + # first time a DiskAnn index is created — users never call any load + # function. The Python extension resolves the module next to _zvec.so + # (see the $ORIGIN rpath in src/binding/python/CMakeLists.txt); the + # module must therefore be installed alongside _zvec.so, i.e. inside the + # zvec package directory as well. + # + # Gate on DISKANN_SUPPORTED, not on the target's existence: on unsupported + # platforms (e.g. macOS / ARM64) the core_knn_diskann target is still + # defined, but built from an empty stub (src/core/algorithm/CMakeLists.txt) + # with zero exported symbols and a runtime load path compiled out + # (#if DISKANN_SUPPORTED). Shipping that stub is pure dead weight, so it is + # only packaged where DiskAnn is real — currently Linux x86_64 with libaio. + if(DISKANN_SUPPORTED) + install(TARGETS core_knn_diskann LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec + COMPONENT python) + endif() + # Bundle cppjieba's dictionary files so the `jieba` FTS tokenizer works + # out of the box. python/zvec/__init__.py resolves this directory via + # importlib.resources and registers it with set_default_jieba_dict_dir(). + set(ZVEC_JIEBA_DICT_SRC + "${PROJECT_SOURCE_DIR}/thirdparty/cppjieba/cppjieba-5.6.7/dict") + install(FILES + "${ZVEC_JIEBA_DICT_SRC}/jieba.dict.utf8" + "${ZVEC_JIEBA_DICT_SRC}/hmm_model.utf8" + DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec/data/jieba_dict + COMPONENT python) +endif() diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..24d18ee --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,41 @@ +# Zvec Code of Conduct + +## Our Pledge + +We pledge to foster an open, respectful, and harassment-free environment for everyone, regardless of age, body size, disability, ethnicity, gender identity and expression, level of experience, education, socioeconomic status, nationality, personal appearance, race, religion, or sexual identity and orientation. + +## Expected Behavior + +- Use welcoming and inclusive language +- Respect differing viewpoints and experiences +- Gracefully accept constructive criticism +- Focus on what is best for the community +- Show empathy and kindness toward others + +## Unacceptable Behavior + +- Harassment, intimidation, or discriminatory conduct +- Trolling, insulting, or derogatory comments +- Public or private harassment +- Publishing others’ private information without consent +- Any conduct that would reasonably be considered inappropriate in a professional setting + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be reported by contacting the project team +at **zvec@alibaba-inc.com** (replace with your contact). All complaints will be reviewed and investigated promptly +and fairly. + +The project team is obligated to respect the privacy and security of the reporter. + +Consequences may include: +- A formal warning +- Temporary or permanent ban from project spaces +- Removal of contributions (e.g. comments, PRs) + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 2.1, available at +https://www.contributor-covenant.org/version/2/1/code_of_conduct.html + +[homepage]: https://www.contributor-covenant.org \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..4650d7c --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,116 @@ +# Contributing to Zvec + +First off, thank you for considering contributing to Zvec! 🙌 +Whether you're reporting a bug, proposing a feature, improving documentation, or submitting code — every contribution helps make Zvec better. + +## Code of Conduct + +By participating, you agree to abide by our [Code of Conduct](CODE_OF_CONDUCT.md). Please be respectful, collaborative, and inclusive. + +--- + +## Development Setup + +> [!TIP] +> **Linux** is the recommended environment for development and performance benchmarking. + +### Prerequisites + +- Python 3.10 - 3.14 (64-bit only) +- CMake ≥ 3.26, < 4.0 (`cmake --version`) +- A C++17-compatible compiler (e.g., `g++-11+`, `clang++`, Apple Clang on macOS) + +### Clone & Initialize + +```bash +git clone --recursive https://github.com/alibaba/zvec.git +cd zvec +``` + +> 💡 **Tip** +> - Forgot `--recursive`? Run: +> ```bash +> git submodule update --init --recursive +> ``` +> - Set up pre-commit hooks: +> ```bash +> pip install pre-commit && pre-commit install +> ``` + +### Build from Source (Editable Install) +```bash +pip install -e ".[dev]" +# This installs dev dependencies (pytest, ruff, etc.) and builds the C++ extension in-place +``` + +> ✅ Verify: +> ```bash +> python -c "import zvec; print('Success!')" +> ``` + +--- + +## Testing + +### Run All Tests +```bash +pytest python/tests/ -v +``` + +### Run with Coverage (Debug/CI) +```bash +pytest python/tests/ --cov=zvec --cov-report=term-missing +``` + +> 🔎 See full rules in `[tool.ruff]` section of `pyproject.toml`. + +--- + +## Build Customization + +You can control build behavior via environment variables or `pyproject.toml`: + +| Option | How to Set | Description | +|--------|------------|-------------| +| **Build Type** | `CMAKE_BUILD_TYPE=Debug` | `Debug`, `Release`, or `Coverage` (for gcov/lcov) | +| **Generator** | `CMAKE_GENERATOR="Unix Makefiles"` | Default: `Ninja`; use Make if preferred | +| **AVX-512** | `ENABLE_SKYLAKE_AVX512=ON` | Enable AVX-512 optimizations (x86_64 only) | + +Example (Debug + Make): +```bash +CMAKE_BUILD_TYPE=Debug CMAKE_GENERATOR="Unix Makefiles" pip install -v . +``` + +--- + +## Submitting Changes + +1. Fork the repo and create a feature branch (`feat/...`, `fix/...`, `docs/...`) +2. Write clear commit messages (e.g., `fix(query): handle null vector in dense_fp32`) +3. Ensure tests pass & linter is clean +4. Open a Pull Request to `main` +5. Link related issue (e.g., `Closes #123`) + +✅ **PRs should include**: +- Test coverage for new behavior +- Updates to documentation (if applicable) +- Reasoning behind non-obvious design choices + +--- + +## Documentation + +- User guides: `docs/` (built with MkDocs) +- API reference: generated from docstrings (follow [Google style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings)) +- Build & deploy: `mkdocs serve` / `mkdocs build` + +--- + +## Need Help + +- Browse [existing issues](https://github.com/alibaba/zvec/issues) +- For sensitive/security issues: email `zvec@alibaba-inc.com` + +--- + +✨ Thanks again for being part of Zvec! diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..f49a4e1 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..5e26964 --- /dev/null +++ b/NOTICE @@ -0,0 +1,96 @@ +zvec +Copyright 2025-present the zvec project + +This product is licensed under the Apache License, Version 2.0 (see the LICENSE +file). It includes third-party software components that are distributed under +their own licenses, as listed below. + +================================================================================ +Third-Party Components +================================================================================ + +-------------------------------------------------------------------------------- +Unicode Character Database +-------------------------------------------------------------------------------- +Project: Unicode Character Database +Homepage: https://www.unicode.org/ +License: Unicode License V3 +Used in: src/db/index/column/fts_column/tokenizer/standard_tokenizer_unicode.inc + +The generated standard tokenizer lookup tables are derived from Unicode 17.0.0 +data files: auxiliary/WordBreakProperty.txt, emoji/emoji-data.txt, +LineBreak.txt, and Scripts.txt. + +Unicode License V3 copyright and permission notice: + + UNICODE LICENSE V3 COPYRIGHT AND PERMISSION NOTICE + Copyright © 1991-2026 Unicode, Inc. + NOTICE TO USER: Carefully read the following legal agreement. BY DOWNLOADING, + INSTALLING, COPYING OR OTHERWISE USING DATA FILES, AND/OR SOFTWARE, YOU + UNEQUIVOCALLY ACCEPT, AND AGREE TO BE BOUND BY, ALL OF THE TERMS AND + CONDITIONS OF THIS AGREEMENT. IF YOU DO NOT AGREE, DO NOT DOWNLOAD, INSTALL, + COPY, DISTRIBUTE OR USE THE DATA FILES OR SOFTWARE. + + Permission is hereby granted, free of charge, to any person obtaining a copy + of data files and any associated documentation (the "Data Files") or software + and any associated documentation (the "Software") to deal in the Data Files + or Software without restriction, including without limitation the rights to + use, copy, modify, merge, publish, distribute, and/or sell copies of the Data + Files or Software, and to permit persons to whom the Data Files or Software + are furnished to do so, provided that either (a) this copyright and + permission notice appear with all copies of the Data Files or Software, or + (b) this copyright and permission notice appear in associated Documentation. + + THE DATA FILES AND SOFTWARE ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY + KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT OF + THIRD PARTY RIGHTS. + + IN NO EVENT SHALL THE COPYRIGHT HOLDER OR HOLDERS INCLUDED IN THIS NOTICE BE + LIABLE FOR ANY CLAIM, OR ANY SPECIAL INDIRECT OR CONSEQUENTIAL DAMAGES, OR + ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER + IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THE DATA FILES OR SOFTWARE. + + Except as contained in this notice, the name of a copyright holder shall not + be used in advertising or otherwise to promote the sale, use or other + dealings in these Data Files or Software without prior written authorization + of the copyright holder. + +-------------------------------------------------------------------------------- +pyglass +-------------------------------------------------------------------------------- +Project: pyglass — Graph Library for Approximate Similarity Search +Homepage: https://github.com/zilliztech/pyglass +License: MIT License +Used in: src/core/utility/linear_pool.h + +The LinearPool implementation (and the accompanying Neighbor / Bitset helpers) +in src/core/utility/linear_pool.h is adapted from pyglass, with modifications +(a BlockHeap-compatible reset()/push_block() interface and the use of +MemoryHelper for huge-page-backed allocation). The related BlockHeap design in +src/core/utility/block_heap.{h,cc} is also derived from pyglass. + +Original license text: + + MIT License + + Copyright (c) 2023 zh Wang + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..ed27ba2 --- /dev/null +++ b/README.md @@ -0,0 +1,137 @@ +

+ English | 中文 +

+ +
+ + + zvec logo + +
+ +

+ Code Coverage + Main + License + PyPI Release + Python Versions + npm Release +

+ +

+ alibaba%2Fzvec | Trendshift +

+ +

+ 🚀 Quickstart | + 🏠 Home | + 📚 Docs | + 📊 Benchmarks | + 🔎 DeepWiki | + 🎮 Discord | + 🐦 X (Twitter) +

+ +**Zvec** is an open-source, in-process vector database — lightweight, lightning-fast, and designed to embed directly into applications. Battle-tested within Alibaba Group, it delivers production-grade, low-latency and scalable similarity search with minimal setup. + +> [!Important] +> 🚀 **v0.5.0 (June 12, 2026)** +> +> - **Full-Text Search (FTS)**: Native full-text search — attach an FTS index to any string field and query it with natural-language or structured expressions, no external search engine required. +> - **Hybrid Retrieval**: Combine full-text and vector search in a single `MultiQuery` across dense vectors, sparse vectors, scalar filters, and text. +> - **DiskANN Index**: New on-disk index that keeps the bulk of the index on disk, drastically cutting memory usage for large-scale datasets. +> - **Ecosystem & Platforms**: New official [Go](https://github.com/zvec-ai/zvec-go) / [Rust](https://github.com/zvec-ai/zvec-rust) SDKs, the [Zvec Studio](https://github.com/zvec-ai/zvec-studio) visual tool, and RISC-V support. +> +> 👉 [Read the Release Notes](https://github.com/alibaba/zvec/releases/tag/v0.5.0) | [View Roadmap 📍](https://github.com/alibaba/zvec/issues/309) + +## 💫 Features + +- **Blazing Fast**: Searches billions of vectors in milliseconds. +- **Simple, Just Works**: [Install](#-installation) and start searching in seconds. Pure local, no servers, no config, no fuss. +- **Dense + Sparse Vectors**: Support dense and sparse embeddings, multi-vector queries, and a rich selection of [vector index types](https://zvec.org/en/docs/db/concepts/vector-index/#vector-index-types) that scale from memory to disk. +- **Full-Text Search (FTS)**: Native keyword-based full-text search — query string fields with natural-language or structured expressions. +- **Hybrid Search**: Fuse vector similarity, full-text search, and structured filters in a single query for precise results. +- **Durable Storage**: Write-ahead logging (WAL) guarantees persistence — data is never lost, even on process crash or power failure. +- **Concurrent Access**: Multiple processes can read the same collection simultaneously; writes are single-process exclusive. +- **Runs Anywhere**: As an in-process library, Zvec runs wherever your code runs — notebooks, servers, CLI tools, or even edge devices. + +## 📦 Installation + +Zvec offers official SDKs across multiple languages: + +- **[Python](https://pypi.org/project/zvec/)**: `pip install zvec` (requires Python 3.10–3.14) +- **[Node.js](https://www.npmjs.com/package/@zvec/zvec)**: `npm install @zvec/zvec` +- **[Go](https://github.com/zvec-ai/zvec-go)**: High-performance Go bindings. +- **[Rust](https://github.com/zvec-ai/zvec-rust)**: High-performance Rust bindings. +- **[Dart/Flutter](https://pub.dev/packages/zvec)**: `flutter pub add zvec` + +Prefer a visual tool? Try **[Zvec Studio](https://github.com/zvec-ai/zvec-studio)** to browse data and debug queries — no code required. + +### ✅ Supported Platforms + +- Linux (x86_64, ARM64) +- macOS (ARM64) +- Windows (x86_64) + +### 🛠️ Building from Source + +If you prefer to build Zvec from source, please check the [Building from Source](https://zvec.org/en/docs/db/build/) guide. + +## ⚡ One-Minute Example + +```python +import zvec + +# Define collection schema +schema = zvec.CollectionSchema( + name="example", + vectors=zvec.VectorSchema("embedding", zvec.DataType.VECTOR_FP32, 4), +) + +# Create collection +collection = zvec.create_and_open(path="./zvec_example", schema=schema) + +# Insert documents +collection.insert([ + zvec.Doc(id="doc_1", vectors={"embedding": [0.1, 0.2, 0.3, 0.4]}), + zvec.Doc(id="doc_2", vectors={"embedding": [0.2, 0.3, 0.4, 0.1]}), +]) + +# Search by vector similarity +results = collection.query( + zvec.Query(field_name="embedding", vector=[0.4, 0.3, 0.3, 0.1]), + topk=10 +) + +# Results: list of {'id': str, 'score': float, ...}, sorted by relevance +print(results) +``` + +## 📈 Performance at Scale + +Zvec delivers exceptional speed and efficiency, making it ideal for demanding production workloads. + +Zvec Performance Benchmarks + +For detailed benchmark methodology, configurations, and complete results, please see our [Benchmarks documentation](https://zvec.org/en/docs/db/benchmarks/). + +## 🤝 Join Our Community + +
+ +
+ +| 💬 DingTalk | 📱 WeChat | 🎮 Discord | X (Twitter) | +| :---: | :---: | :---: | :---: | +| DingTalk QR Code | WeChat QR Code | [![Discord](https://img.shields.io/badge/Discord-Join%20Server-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/rKddFBBu9z) | [![X (formerly Twitter) Follow](https://img.shields.io/twitter/follow/ZvecAI)]() | +| Scan to join | Scan to join | Click to join | Click to follow | + +
+ +
+ +## ❤️ Contributing + +We welcome and appreciate contributions from the community! Whether you're fixing a bug, adding a feature, or improving documentation, your help makes Zvec better for everyone. + +Check out our [Contributing Guide](./CONTRIBUTING.md) to get started! diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..04b9a1b --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`alibaba/zvec` +- 原始仓库：https://github.com/alibaba/zvec +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/README_CN.md b/README_CN.md new file mode 100644 index 0000000..4ad14c6 --- /dev/null +++ b/README_CN.md @@ -0,0 +1,133 @@ +

+ English | 中文 +

+ +
+ + + zvec logo + +
+ +

+ 代码覆盖率 + Main + 许可证 + PyPI 版本 + Python 版本 + npm 版本 +

+ +

+ alibaba%2Fzvec | Trendshift +

+ +

+ 🚀 快速开始 | + 🏠 主页 | + 📚 文档 | + 📊 性能报告 | + 🔎 DeepWiki | + 🎮 Discord | + 🐦 X (Twitter) +

+ +**Zvec** 是一款开源的嵌入式(进程内)向量数据库 — 轻量、极速，可直接嵌入应用程序。以极简的配置提供生产级、低延迟、可扩展的向量检索能力。 + +> [!IMPORTANT] +> 🚀 **v0.5.0（2026 年 6 月 12 日）** +> +> - **全文检索（FTS）**：原生全文检索能力——可为任意字符串字段挂载 FTS 索引，使用自然语言或结构化表达式检索，无需外接搜索引擎。 +> - **混合检索**：在单次 `MultiQuery` 中融合全文与向量检索，跨稠密向量、稀疏向量、标量过滤与文本。 +> - **DiskANN 索引**：全新磁盘索引，将索引主体存于磁盘，大幅降低大规模数据集的内存占用。 +> - **生态与平台**：全新官方 [Go](https://github.com/zvec-ai/zvec-go) / [Rust](https://github.com/zvec-ai/zvec-rust) SDK、可视化工具 [Zvec Studio](https://github.com/zvec-ai/zvec-studio)，以及 RISC-V 架构支持。 +> +> 👉 [查看更新日志](https://github.com/alibaba/zvec/releases/tag/v0.5.0) | [查看路线图 📍](https://github.com/alibaba/zvec/issues/309) + +## 💫 核心特性 + +- **极致性能**：毫秒级响应，轻松检索数十亿级向量。 +- **开箱即用**：[安装](#-安装)后即刻开始搜索，纯本地运行，无需服务器、无需配置、零门槛。 +- **稠密 + 稀疏向量**：支持稠密向量、稀疏向量与多向量查询，以及从内存到磁盘、丰富多样的[向量索引类型](https://zvec.org/zh/docs/db/concepts/vector-index/#向量索引类型)。 +- **全文检索（FTS）**：原生的基于关键词的全文检索——使用自然语言或结构化表达式检索字符串字段。 +- **混合检索**：在单次查询中融合向量语义、全文检索与标量过滤，获得精确结果。 +- **持久化存储**：WAL 预写日志保障数据持久性 — 即使进程崩溃或意外断电，数据也不会丢失。 +- **并发访问**：支持多进程同时读取同一个 Collection；写入为单进程独占模式。 +- **进程内运行**：无需单独部署服务，纯进程内运行。Notebook、高性能服务器、CLI 工具、边缘设备 — 随处可用。 + +## 📦 安装 + +Zvec 提供多语言官方 SDK： + +- **[Python](https://pypi.org/project/zvec/)**：`pip install zvec`（需 64 位 Python 3.10–3.14） +- **[Node.js](https://www.npmjs.com/package/@zvec/zvec)**：`npm install @zvec/zvec` +- **[Go](https://github.com/zvec-ai/zvec-go)**：高性能的 Go 绑定。 +- **[Rust](https://github.com/zvec-ai/zvec-rust)**：高性能的 Rust 绑定。 +- **[Dart/Flutter](https://pub.dev/packages/zvec)**：`flutter pub add zvec` + +想要图形界面？试试 **[Zvec Studio](https://github.com/zvec-ai/zvec-studio)**，零代码浏览数据与调试查询。 + +### ✅ 支持的平台 + +- Linux (x86_64, ARM64) +- macOS (ARM64) +- Windows (x86_64) + +### 🛠️ 源码构建 + +如需从源码构建 Zvec，请参考[源码构建指南](https://zvec.org/zh/docs/db/build/)。 + +## ⚡ 一分钟上手 + +```python +import zvec + +# 定义 collection schema +schema = zvec.CollectionSchema( + name="example", + vectors=zvec.VectorSchema("embedding", zvec.DataType.VECTOR_FP32, 4), +) + +# 创建 collection +collection = zvec.create_and_open(path="./zvec_example", schema=schema) + +# 插入 documents +collection.insert([ + zvec.Doc(id="doc_1", vectors={"embedding": [0.1, 0.2, 0.3, 0.4]}), + zvec.Doc(id="doc_2", vectors={"embedding": [0.2, 0.3, 0.4, 0.1]}), +]) + +# 向量相似度检索 +results = collection.query( + zvec.Query(field_name="embedding", vector=[0.4, 0.3, 0.3, 0.1]), + topk=10 +) + +# 查询结果：按相关性排序的 {'id': str, 'score': float, ...} 列表 +print(results) +``` + +## 📈 极致性能 + +Zvec 提供极致的速度和效率，能够轻松应对高要求的生产环境负载。 + +Zvec 性能基准测试 + +有关具体的测试方法、配置及完整结果，请参阅[性能报告](https://zvec.org/zh/docs/db/benchmarks/)。 + +## 🤝 加入社区 + +
+ +| 💬 钉钉群 | 📱 微信群 | 🎮 Discord | X (Twitter) | +| :---: | :---: | :---: | :---: | +| 钉钉二维码 | 微信二维码 | [![Discord](https://img.shields.io/badge/Discord-Join%20Server-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/rKddFBBu9z) | [![X (formerly Twitter) Follow](https://img.shields.io/twitter/follow/ZvecAI)]() | +| 扫码加入 | 扫码加入 | 点击加入 | 点击关注 | + +
+ +## ❤️ 参与贡献 + +非常欢迎来自社区的每一份贡献！无论是修复 Bug、新增功能，还是完善文档，都将让 Zvec 变得更好。 + +请查阅我们的[贡献指南](./CONTRIBUTING.md)开始参与！ diff --git a/cmake/bazel.cmake b/cmake/bazel.cmake new file mode 100644 index 0000000..b097340 --- /dev/null +++ b/cmake/bazel.cmake @@ -0,0 +1,2521 @@ +## +## The following functions used by user's CMakeLists.txt: +## + +## 1. Functions for C/C++ +## +## 1.1. Add a subdirectory to the build +## cc_directory( [binary_dir]) +## +## 1.2. Add subdirectories to the build +## cc_directories( [source_dir2 ...]) +## +## 1.3. Build a C/C++ static or shared library +## cc_library( +## NAME +## [STATIC] [SHARED] [STRICT] [ALWAYS_LINK] [EXCLUDE] [PACKED] [SRCS_NO_GLOB] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [PUBINCS public_dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [PACKED_EXCLUDES pattern1 ...] +## [VERSION ] +## ) +## +## 1.4. Build a C/C++ executable program +## cc_binary( +## NAME +## [STRICT] [PACKED] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [VERSION ] +## ) +## +## 1.5. Build a C/C++ executable test program +## cc_test( +## NAME +## [STRICT] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [ARGS args1 ...] +## [VERSION ] +## ) +## +## 1.6. Add existing test cases to a test suite +## cc_test_suite( [test_name ...]) +## +## 1.7. Import a C/C++ static or shared library +## cc_import( +## NAME +## [STATIC | SHARED] [PACKED] +## PATH +## [INCS dir1 ...] +## [PUBINCS public_dir1 ...] +## [DEPS target1 ...] +## [IMPLIB ] +## [PACKED_EXCLUDES pattern1 ...] +## ) +## +## 1.8. Import a C/C++ interface library +## cc_interface( +## NAME +## [PACKED] +## [INCS dir1 ...] +## [PUBINCS public_dir1 ...] +## [DEPS target1 ...] +## [PACKED_EXCLUDES pattern1 ...] +## ) +## +## 1.9. Build a C/C++ executable google test program +## cc_gtest( +## NAME +## [STRICT] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [ARGS args1 ...] +## [VERSION ] +## ) +## +## 1.10. Build a C/C++ executable google mock program +## cc_gmock( +## NAME +## [STRICT] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [ARGS args1 ...] +## [VERSION ] +## ) +## +## 1.11. Build a C++ protobuf static or shared library +## cc_proto_library( +## NAME +## [STATIC] [SHARED] [STRICT] [EXCLUDE] [PACKED] +## SRCS [file2.proto ...] +## [PROTOROOT path] +## [CXXFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [VERSION ] +## [PROTOBUF_VERSION ] +## ) +## + +## 2. Functions for CUDA +## +## 2.1. Add a subdirectory to the build +## cuda_directory( [binary_dir]) +## +## 2.2. Add subdirectories to the build +## cuda_directories( [source_dir2 ...]) +## +## 2.3. Build a CUDA static or shared library +## cuda_library( +## NAME +## [STATIC] [SHARED] [STRICT] [ALWAYS_LINK] [EXCLUDE] [PACKED] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [PUBINCS public_dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [CUDAFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [PACKED_EXCLUDES pattern1 ...] +## [VERSION ] +## ) +## +## 2.4. Build a CUDA executable program +## cuda_binary( +## NAME +## [STRICT] [PACKED] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [CUDAFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [VERSION ] +## ) +## +## 2.5. Build a CUDA executable test program +## cuda_test( +## NAME +## [STRICT] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [CUDAFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [ARGS args1 ...] +## [VERSION ] +## ) +## +## 2.6. Add existing test cases to a test suite +## cuda_test_suite( [test_name ...]) +## +## 2.7. Import a C/C++/CUDA static or shared library +## cuda_import( +## NAME +## [STATIC | SHARED] [PACKED] +## PATH +## [INCS dir1 ...] +## [PUBINCS public_dir1 ...] +## [DEPS target1 ...] +## [IMPLIB ] +## [PACKED_EXCLUDES pattern1 ...] +## ) +## +## 2.8. Import a C/C++/CUDA interface library +## cuda_interface( +## NAME +## [PACKED] +## [INCS dir1 ...] +## [PUBINCS public_dir1 ...] +## [DEPS target1 ...] +## [PACKED_EXCLUDES pattern1 ...] +## ) +## +## 2.9. Build a CUDA executable google test program +## cuda_gtest( +## NAME +## [STRICT] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [CUDAFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [ARGS args1 ...] +## [VERSION ] +## ) +## +## 2.10. Build a CUDA executable google mock program +## cuda_gmock( +## NAME +## [STRICT] +## SRCS [file2 ...] +## [INCS dir1 ...] +## [DEFS DEF1=1 ...] +## [LIBS lib1 ...] +## [CFLAGS flag1 ...] +## [CXXFLAGS flag1 ...] +## [CUDAFLAGS flag1 ...] +## [LDFLAGS flag1 ...] +## [DEPS target1 ...] +## [ARGS args1 ...] +## [VERSION ] +## ) +## + +## 3. Utility functions +## +## 3.1. Download a git repository +## git_repository( +## NAME +## URL +## [TAG ] +## [PATH ] +## ) +## +## 3.2. Download a hg repository +## hg_repository( +## NAME +## URL +## [TAG ] +## [PATH ] +## ) +## +## 3.3. Download a svn repository +## svn_repository( +## NAME +## URL +## [REV ] +## [PATH ] +## ) +## +## 3.4. Download a http archive +## http_archive( +## NAME +## URL +## [SHA256 | SHA1 | MD5 ] +## [PATH ] +## ) +## +## 3.5. Retrieve a version string from GIT +## git_version( +## +## +## ) +## +## 3.6. Retrieve a version string from HG +## hg_version( +## +## +## ) +## +## 3.7. Retrieve a version string from SVN +## svn_version( +## +## +## ) +## + +cmake_minimum_required(VERSION 3.13 FATAL_ERROR) +include(CMakeParseArguments) + +# Using AppleClang instead of Clang (Compiler id) +if(POLICY CMP0025) + cmake_policy(SET CMP0025 NEW) +endif() + +# Enable unit testing +enable_testing() + +# Add unittest target +if(NOT TARGET unittest) + if(IOS) + # iOS: build-only target; tests are run on simulator separately + add_custom_target(unittest) + else() + include(ProcessorCount) + ProcessorCount(NPROC) + if(NPROC EQUAL 0) + set(NPROC 1) + endif() + math(EXPR PARALLEL_JOBS "${NPROC} - 1") + if(PARALLEL_JOBS LESS 1) + set(PARALLEL_JOBS 1) + endif() + add_custom_target( + unittest + COMMAND ${CMAKE_CTEST_COMMAND} --output-on-failure + --build-config $ + --parallel ${PARALLEL_JOBS} + ) + endif() +endif() + +# Directories of target output +if(NOT CMAKE_ARCHIVE_OUTPUT_DIRECTORY) + set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/lib) +endif() +if(NOT CMAKE_LIBRARY_OUTPUT_DIRECTORY) + set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/lib) +endif() +if(NOT CMAKE_RUNTIME_OUTPUT_DIRECTORY) + set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/bin) +endif() + +# RPATH settings +set(CMAKE_MACOSX_RPATH ON) +if(NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin") + set(CMAKE_SKIP_BUILD_RPATH ON) + set(CMAKE_BUILD_WITH_INSTALL_RPATH ON) + if(${CMAKE_SIZEOF_VOID_P} EQUAL "8") + set(CMAKE_INSTALL_RPATH "$ORIGIN/../lib64:$ORIGIN/../lib:$ORIGIN") + else() + set(CMAKE_INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN") + endif() +else() + set(CMAKE_INSTALL_RPATH "@loader_path/../lib:@loader_path") +endif() + +# Define standard installation directories +if(NOT CMAKE_INSTALL_LIBDIR) + set(CMAKE_INSTALL_LIBDIR lib) +endif() +if(NOT CMAKE_INSTALL_BINDIR) + set(CMAKE_INSTALL_BINDIR bin) +endif() +if(NOT CMAKE_INSTALL_INCDIR) + set(CMAKE_INSTALL_INCDIR include) +endif() +if(NOT CMAKE_INSTALL_ETCDIR) + set(CMAKE_INSTALL_ETCDIR etc) +endif() + +# Generates a compile_commands.json +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) +set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON) + +if(APPLE OR ANDROID) + option(CLANG_USE_LIBCXX "Use libc++ instead of libstdc++" ON) +else() + option(CLANG_USE_LIBCXX "Use libc++ instead of libstdc++" OFF) +endif() + +set(CLANG_STDLIB_OPTION "") +if(CLANG_USE_LIBCXX) + set(CLANG_STDLIB_OPTION "-stdlib=libc++") +else() + set(CLANG_STDLIB_OPTION "-stdlib=libstdc++") +endif() + +if(NOT MSVC) + # Use color in diagnostics + set( + _C_FLAGS + "$<$:-fcolor-diagnostics>" + "$<$:-fcolor-diagnostics>" + "$<$:-fdiagnostics-color=always>" + ) + set( + _CXX_FLAGS + "$<$:-fcolor-diagnostics;${CLANG_STDLIB_OPTION}>" + "$<$:-fcolor-diagnostics>" + "$<$:-fdiagnostics-color=always>" + ) + add_compile_options( + "$<$:${_C_FLAGS}>" + "$<$:${_CXX_FLAGS}>" + ) + unset(_C_FLAGS) + unset(_CXX_FLAGS) +else() + option(ZVEC_USE_STATIC_CRT "Use static CRT (/MT) instead of dynamic CRT (/MD), default=ON" ON) + + if(ZVEC_USE_STATIC_CRT) + set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$:Debug>" CACHE STRING "" FORCE) + else() + set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreadedDLL$<$:Debug>" CACHE STRING "" FORCE) + endif() + + set( + _COMPILER_FLAGS + CMAKE_CXX_FLAGS + CMAKE_CXX_FLAGS_DEBUG + CMAKE_CXX_FLAGS_RELEASE + CMAKE_CXX_FLAGS_RELWITHDEBINFO + CMAKE_CXX_FLAGS_MINSIZEREL + CMAKE_C_FLAGS + CMAKE_C_FLAGS_DEBUG + CMAKE_C_FLAGS_RELEASE + CMAKE_C_FLAGS_RELWITHDEBINFO + CMAKE_C_FLAGS_MINSIZEREL + ) + if(ZVEC_USE_STATIC_CRT) + foreach(COMPILER_FLAG ${_COMPILER_FLAGS}) + string(REPLACE "/MD" "/MT" ${COMPILER_FLAG} "${${COMPILER_FLAG}}") + string(REGEX REPLACE "/W[0-9]" "" ${COMPILER_FLAG} "${${COMPILER_FLAG}}") + endforeach() + else() + foreach(COMPILER_FLAG ${_COMPILER_FLAGS}) + string(REPLACE "/MT" "/MD" ${COMPILER_FLAG} "${${COMPILER_FLAG}}") + string(REGEX REPLACE "/W[0-9]" "" ${COMPILER_FLAG} "${${COMPILER_FLAG}}") + endforeach() + endif() + unset(_COMPILER_FLAGS) + + add_definitions(-D_CRT_SECURE_NO_WARNINGS) + set(BUILD_SHARED_LIBS OFF) +endif() + +set(CMAKE_C_FLAGS_ASAN ${CMAKE_C_FLAGS_DEBUG}) +set(CMAKE_CXX_FLAGS_ASAN ${CMAKE_CXX_FLAGS_DEBUG}) +set(CMAKE_EXE_LINKER_FLAGS_ASAN ${CMAKE_EXE_LINKER_FLAGS_DEBUG}) +set(CMAKE_SHARED_LINKER_FLAGS_ASAN ${CMAKE_SHARED_LINKER_FLAGS_DEBUG}) +set(CMAKE_STATIC_LINKER_FLAGS_ASAN ${CMAKE_STATIC_LINKER_FLAGS_DEBUG}) +set(CMAKE_MODULE_LINKER_FLAGS_ASAN ${CMAKE_MODULE_LINKER_FLAGS_DEBUG}) +set(CMAKE_C_FLAGS_COVERAGE ${CMAKE_C_FLAGS_DEBUG}) +set(CMAKE_CXX_FLAGS_COVERAGE ${CMAKE_CXX_FLAGS_DEBUG}) +set(CMAKE_EXE_LINKER_FLAGS_COVERAGE ${CMAKE_EXE_LINKER_FLAGS_DEBUG}) +set(CMAKE_SHARED_LINKER_FLAGS_COVERAGE ${CMAKE_SHARED_LINKER_FLAGS_DEBUG}) +set(CMAKE_STATIC_LINKER_FLAGS_COVERAGE ${CMAKE_STATIC_LINKER_FLAGS_DEBUG}) +set(CMAKE_MODULE_LINKER_FLAGS_COVERAGE ${CMAKE_MODULE_LINKER_FLAGS_DEBUG}) + +# C/C++ ASAN compile flags +set( + BAZEL_CC_ASAN_COMPILE_FLAGS + "$<$:$<$:-fsanitize=address>>" + "$<$:$<$:-fsanitize=address>>" + "$<$:$<$:-fsanitize=address>>" + "$<$:$<$:/fsanitize=address>>" + ) + +# C/C++ COVERAGE compile flags +set( + BAZEL_CC_COVERAGE_COMPILE_FLAGS + "$<$:$<$:--coverage>>" + "$<$:$<$:--coverage>>" + "$<$:$<$:--coverage>>" + "$<$:-fprofile-update=atomic>" + ) + +# C/C++ strict compile flags +if(ENABLE_WERROR) + set(BAZEL_CC_WERROR_FLAGS + "$<$:-Werror>" + "$<$:-Werror>" + "$<$:-Werror>" + "$<$:/WX>" + ) +else() + set(BAZEL_CC_WERROR_FLAGS "") +endif() + +if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 7.0) + set( + BAZEL_CC_STRICT_COMPILE_FLAGS + "$<$:-Wall;-Wextra;-Wshadow>" + "$<$:-Wall;-Wextra;-Wshadow>" + "$<$:-Wall;-Wextra;-Wshadow-local;-Wno-misleading-indentation>" + "$<$:/W4>" + ${BAZEL_CC_WERROR_FLAGS} + ${BAZEL_CC_ASAN_COMPILE_FLAGS} + ${BAZEL_CC_COVERAGE_COMPILE_FLAGS} + ) +else() + set( + BAZEL_CC_STRICT_COMPILE_FLAGS + "$<$:-Wall;-Wextra;-Wshadow>" + "$<$:-Wall;-Wextra;-Wshadow>" + "$<$:-Wall;-Wextra;-Wshadow;-Wno-misleading-indentation>" + "$<$:/W4>" + ${BAZEL_CC_WERROR_FLAGS} + ${BAZEL_CC_ASAN_COMPILE_FLAGS} + ${BAZEL_CC_COVERAGE_COMPILE_FLAGS} + ) +endif() + + +# C/C++ strict link flags +set( + BAZEL_CC_STRICT_LINK_FLAGS + ${BAZEL_CC_ASAN_COMPILE_FLAGS} + ${BAZEL_CC_COVERAGE_COMPILE_FLAGS} + ) + +# C/C++ unstrict compile flags +set( + BAZEL_CC_UNSTRICT_COMPILE_FLAGS + "$<$:-Wall>" + "$<$:-Wall>" + "$<$:-Wall>" + "$<$:/W3>" + ${BAZEL_CC_ASAN_COMPILE_FLAGS} + ${BAZEL_CC_COVERAGE_COMPILE_FLAGS} + ) + +# C/C++ unstrict link flags +set( + BAZEL_CC_UNSTRICT_LINK_FLAGS + ${BAZEL_CC_ASAN_COMPILE_FLAGS} + ${BAZEL_CC_COVERAGE_COMPILE_FLAGS} + ) + +# CUDA strict compile flags +set( + BAZEL_CUDA_STRICT_COMPILE_FLAGS + "$<$:$<$:-Wall;-Wextra;-Wshadow>>" + "$<$:$<$:-Wall;-Wextra;-Wshadow>>" + "$<$:$<$:-Wall;-Wextra;-Wshadow>>" + "$<$:$<$:/W4>>" + "$<$:$<$:-Wall;-Wextra;-Wshadow>>" + "$<$:$<$:-Wall;-Wextra;-Wshadow>>" + "$<$:$<$:-Wall;-Wextra;-Wshadow>>" + "$<$:$<$:/W4>>" + "$<$:$<$:-G>>" + ) + +# CUDA strict link flags +set(BAZEL_CUDA_STRICT_LINK_FLAGS "") + +# CUDA unstrict compile flags +set( + BAZEL_CUDA_UNSTRICT_COMPILE_FLAGS + "$<$:$<$:-Wall>>" + "$<$:$<$:-Wall>>" + "$<$:$<$:-Wall>>" + "$<$:$<$:/W3>>" + "$<$:$<$:-Wall>>" + "$<$:$<$:-Wall>>" + "$<$:$<$:-Wall>>" + "$<$:$<$:/W3>>" + "$<$:$<$:-G>>" + ) + +# CUDA unstrict link flags +set(BAZEL_CUDA_UNSTRICT_LINK_FLAGS "") + +## Find workspace directory +function(_find_workspace_directory _RESULT) + # Find Workspace.cmake folder + set(_CURRENT_WORKSPACE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + get_filename_component( + _PARENT_WORKSPACE_DIR ${_CURRENT_WORKSPACE_DIR} DIRECTORY + ) + while(NOT ("${_CURRENT_WORKSPACE_DIR}" STREQUAL "${_PARENT_WORKSPACE_DIR}")) + if(EXISTS "${_CURRENT_WORKSPACE_DIR}/Workspace.cmake") + set(${_RESULT} ${_CURRENT_WORKSPACE_DIR} PARENT_SCOPE) + message(STATUS "Found workspace at ${${_RESULT}}") + break() + endif() + + # Find next parent folder + set(_CURRENT_WORKSPACE_DIR ${_PARENT_WORKSPACE_DIR}) + get_filename_component( + _PARENT_WORKSPACE_DIR ${_CURRENT_WORKSPACE_DIR} DIRECTORY + ) + endwhile() +endfunction() + +## Retrieve absolute paths +function(_absolute_paths _RESULT) + foreach(FILEPATH ${ARGN}) + if(NOT IS_ABSOLUTE ${FILEPATH}) + get_filename_component(FILEPATH ${FILEPATH} ABSOLUTE) + endif() + list(APPEND FILEPATHS ${FILEPATH}) + endforeach() + set(${_RESULT} "${FILEPATHS}" PARENT_SCOPE) +endfunction() + +## Add both shared and static library +macro(_add_library _NAME _OPTION) + add_library(${_NAME}_objects OBJECT ${_OPTION} ${ARGN}) + add_library( + ${_NAME}_static STATIC ${_OPTION} $ + ) + if(IOS) + # iOS: create the main target as static too (no shared libs on iOS) + add_library( + ${_NAME} STATIC ${_OPTION} $ + ) + else() + add_library( + ${_NAME} SHARED ${_OPTION} $ + ) + endif() + add_dependencies(${_NAME} ${_NAME}_static) + if(NOT MSVC) + set_property(TARGET ${_NAME}_static PROPERTY OUTPUT_NAME ${_NAME}) + endif() +endmacro() + +## Link dependencies +function(_targets_link_dependencies _NAME) + foreach(LIB ${ARGN}) + if(TARGET ${LIB}) + list(APPEND LIBS_DEPS ${LIB}) + list( + APPEND LIBS_INCS + "$" + ) + list( + APPEND LIBS_SYSTEM_INCS + "$" + ) + endif() + endforeach() + + if(LIBS_DEPS) + add_dependencies(${_NAME} ${LIBS_DEPS}) + target_include_directories(${_NAME} SYSTEM PRIVATE "${LIBS_SYSTEM_INCS}") + target_include_directories(${_NAME} PRIVATE "${LIBS_INCS}") + endif() +endfunction() + +## Link libraries +function(_target_link_libraries _NAME) + function(_collect_always_link_libs LIB_LIST RESULT_VAR) + if(NOT _COLLECT_ALWAYS_LINK_VISITED) + set(_COLLECT_ALWAYS_LINK_VISITED "" PARENT_SCOPE) + endif() + + set(LOCAL_RESULT "") + foreach(LIB ${LIB_LIST}) + if(NOT TARGET ${LIB}) + continue() + endif() + + list(FIND _COLLECT_ALWAYS_LINK_VISITED ${LIB} ALREADY_VISITED) + if(NOT ALREADY_VISITED EQUAL -1) + continue() + endif() + + list(APPEND _COLLECT_ALWAYS_LINK_VISITED ${LIB}) + set(_COLLECT_ALWAYS_LINK_VISITED "${_COLLECT_ALWAYS_LINK_VISITED}" PARENT_SCOPE) + + get_target_property(ALWAYS_LINK ${LIB} ALWAYS_LINK) + if(ALWAYS_LINK) + list(APPEND LOCAL_RESULT ${LIB}) + elseif(MSVC AND TARGET ${LIB}_static) + get_target_property(_SIBLING_AL ${LIB}_static ALWAYS_LINK) + if(_SIBLING_AL) + list(APPEND LOCAL_RESULT ${LIB}_static) + endif() + endif() + + get_target_property(DEP_LIBS ${LIB} INTERFACE_LINK_LIBRARIES) + if(DEP_LIBS) + _collect_always_link_libs("${DEP_LIBS}" DEP_ALWAYS_LINK_LIBS) + list(APPEND LOCAL_RESULT ${DEP_ALWAYS_LINK_LIBS}) + endif() + + get_target_property(LINK_LIBS ${LIB} LINK_LIBRARIES) + if(LINK_LIBS) + _collect_always_link_libs("${LINK_LIBS}" LINK_ALWAYS_LINK_LIBS) + list(APPEND LOCAL_RESULT ${LINK_ALWAYS_LINK_LIBS}) + endif() + endforeach() + + list(REMOVE_DUPLICATES LOCAL_RESULT) + set(${RESULT_VAR} "${LOCAL_RESULT}" PARENT_SCOPE) + endfunction() + + _collect_always_link_libs("${ARGN}" ALL_ALWAYS_LINK_LIBS) + + set(ALL_LIBS_TO_PROCESS ${ARGN}) + foreach(ALWAYS_LIB ${ALL_ALWAYS_LINK_LIBS}) + list(FIND ARGN ${ALWAYS_LIB} FOUND_INDEX) + if(FOUND_INDEX EQUAL -1) + list(APPEND ALL_LIBS_TO_PROCESS ${ALWAYS_LIB}) + endif() + endforeach() + + list(REMOVE_DUPLICATES ALL_LIBS_TO_PROCESS) + + # On MSVC, each DLL has its own copy of template statics (e.g. Factory + # singletons), so registrations inside a DLL are invisible to the exe. + # Substitute SHARED libs with their ALWAYS_LINK _static counterparts and + # use /WHOLEARCHIVE so all registration code lives in the same module. + if(MSVC) + set(_SUBSTITUTED_LIBS "") + foreach(LIB ${ALL_LIBS_TO_PROCESS}) + if(TARGET ${LIB} AND TARGET ${LIB}_static) + get_target_property(_LIB_TYPE ${LIB} TYPE) + get_target_property(_STATIC_AL ${LIB}_static ALWAYS_LINK) + if("${_LIB_TYPE}" STREQUAL "SHARED_LIBRARY" AND _STATIC_AL) + list(APPEND _SUBSTITUTED_LIBS ${LIB}_static) + list(APPEND ALL_ALWAYS_LINK_LIBS ${LIB}_static) + continue() + endif() + endif() + list(APPEND _SUBSTITUTED_LIBS ${LIB}) + endforeach() + set(ALL_LIBS_TO_PROCESS ${_SUBSTITUTED_LIBS}) + if(ALL_ALWAYS_LINK_LIBS) + list(REMOVE_DUPLICATES ALL_ALWAYS_LINK_LIBS) + endif() + endif() + + foreach(LIB ${ALL_LIBS_TO_PROCESS}) + if(NOT TARGET ${LIB}) + list(APPEND LINK_LIBS ${LIB}) + continue() + endif() + + list(FIND ALL_ALWAYS_LINK_LIBS ${LIB} IS_ALWAYS_LINK) + if(IS_ALWAYS_LINK EQUAL -1) + list(APPEND LINK_LIBS ${LIB}) + continue() + endif() + + if(NOT MSVC) + if(NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin" AND NOT ${CMAKE_SYSTEM_NAME} MATCHES "iOS") + list(APPEND LINK_LIBS -Wl,--whole-archive ${LIB} -Wl,--no-whole-archive) + else() + list(APPEND LINK_LIBS -Wl,-force_load ${LIB}) + endif() + else() + # TODO(windows): revert maybe + list(APPEND MSVC_WHOLEARCHIVE_OPTS /WHOLEARCHIVE:$) + get_target_property(OTHER_LINK_LIBS ${LIB} INTERFACE_LINK_LIBRARIES) + if(OTHER_LINK_LIBS) + foreach(OTHER_LIB ${OTHER_LINK_LIBS}) + list(FIND ALL_LIBS_TO_PROCESS ${OTHER_LIB} FOUND_INDEX) + if(FOUND_INDEX EQUAL -1) + list(APPEND LINK_LIBS ${OTHER_LIB}) + endif() + endforeach() + endif() + list(APPEND LIBS_DEPS ${LIB}) + list( + APPEND LIBS_INCS + "$" + ) + endif() + endforeach() + + target_link_libraries(${_NAME} ${LINK_LIBS}) + if(MSVC_WHOLEARCHIVE_OPTS) + target_link_options(${_NAME} PRIVATE ${MSVC_WHOLEARCHIVE_OPTS}) + endif() + if(LIBS_DEPS) + add_dependencies(${_NAME} ${LIBS_DEPS}) + target_include_directories(${_NAME} PRIVATE "${LIBS_INCS}") + endif() +endfunction() + +## Add a subdirectory to the build +function(cc_directory) + add_subdirectory(${ARGN}) +endfunction() + +## Add subdirectories to the build +function(cc_directories) + foreach(SRC_DIR ${ARGN}) + add_subdirectory(${SRC_DIR}) + endforeach() +endfunction() + +## Set the properties of target +function(_cc_target_properties) + cmake_parse_arguments( + CC_ARGS "STRICT;ALWAYS_LINK" "NAME;VERSION;C_STANDARD;CXX_STANDARD" + "INCS;PUBINCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS" ${ARGN} + ) + + if(NOT CC_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + get_target_property(TARGET_TYPE ${CC_ARGS_NAME} TYPE) + if(("${TARGET_TYPE}" STREQUAL "SHARED_LIBRARY") OR + ("${TARGET_TYPE}" STREQUAL "STATIC_LIBRARY") OR + ("${TARGET_TYPE}" STREQUAL "EXECUTABLE")) + set(TARGET_LINKABLE TRUE) + endif() + + if(CC_ARGS_ALWAYS_LINK) + if(("${TARGET_TYPE}" STREQUAL "STATIC_LIBRARY") OR + ("${TARGET_TYPE}" STREQUAL "OBJECT_LIBRARY")) + set_property(TARGET ${CC_ARGS_NAME} PROPERTY ALWAYS_LINK TRUE) + endif() + endif() + + # Set the warning level of compiling + if(CC_ARGS_STRICT) + target_compile_options( + ${CC_ARGS_NAME} PRIVATE "${BAZEL_CC_STRICT_COMPILE_FLAGS}" + ) + if(TARGET_LINKABLE) + target_link_libraries(${CC_ARGS_NAME} "${BAZEL_CC_STRICT_LINK_FLAGS}") + endif() + else() + target_compile_options( + ${CC_ARGS_NAME} PRIVATE "${BAZEL_CC_UNSTRICT_COMPILE_FLAGS}" + ) + if(TARGET_LINKABLE) + target_link_libraries(${CC_ARGS_NAME} "${BAZEL_CC_UNSTRICT_LINK_FLAGS}") + endif() + endif() + + if(CC_ARGS_DEFS) + target_compile_definitions(${CC_ARGS_NAME} PRIVATE "${CC_ARGS_DEFS}") + endif() + + if(CC_ARGS_CFLAGS OR CC_ARGS_CXXFLAGS) + target_compile_options( + ${CC_ARGS_NAME} PRIVATE + "$<$:${CC_ARGS_CFLAGS}>" + "$<$:${CC_ARGS_CXXFLAGS}>" + ) + endif() + + if(CC_ARGS_LDFLAGS) + string(REPLACE ";" " " CC_ARGS_LDFLAGS "${CC_ARGS_LDFLAGS}") + set_property( + TARGET ${CC_ARGS_NAME} PROPERTY LINK_FLAGS "${CC_ARGS_LDFLAGS}" + ) + endif() + + if(CC_ARGS_INCS) + _absolute_paths(INC_DIRS ${CC_ARGS_INCS}) + target_include_directories(${CC_ARGS_NAME} PRIVATE "${INC_DIRS}") + endif() + + if(BAZEL_WORKSPACE_DIR) + target_include_directories(${CC_ARGS_NAME} PRIVATE "${BAZEL_WORKSPACE_DIR}") + endif() + + if(CC_ARGS_PUBINCS) + _absolute_paths(INC_DIRS ${CC_ARGS_PUBINCS}) + target_include_directories(${CC_ARGS_NAME} PUBLIC "${INC_DIRS}") + endif() + + if(CC_ARGS_LIBS) + if(NOT TARGET_LINKABLE) + _targets_link_dependencies(${CC_ARGS_NAME} ${CC_ARGS_LIBS}) + else() + if ("${TARGET_TYPE}" STREQUAL "EXECUTABLE") + _target_link_libraries(${CC_ARGS_NAME} "${CC_ARGS_LIBS}") + else() + target_link_libraries(${CC_ARGS_NAME} "${CC_ARGS_LIBS}") + endif() + endif() + endif() + + if(CC_ARGS_DEPS) + add_dependencies(${CC_ARGS_NAME} "${CC_ARGS_DEPS}") + endif() + + if(CC_ARGS_VERSION) + set_property( + TARGET ${CC_ARGS_NAME} PROPERTY VERSION "${CC_ARGS_VERSION}" + ) + endif() + + if(NOT CC_C_STANDARD) + set(CC_C_STANDARD 99) + endif() + + if(NOT CC_CXX_STANDARD) + set(CC_CXX_STANDARD 11) + endif() + + set_target_properties( + ${CC_ARGS_NAME} PROPERTIES DEFINE_SYMBOL "" + C_STANDARD ${CC_C_STANDARD} CXX_STANDARD ${CC_CXX_STANDARD} + C_STANDARD_REQUIRED ON C_EXTENSIONS ON + CXX_STANDARD_REQUIRED ON CXX_EXTENSIONS OFF + WINDOWS_EXPORT_ALL_SYMBOLS ON + ) +endfunction() + +## Build a C/C++ static or shared library +function(cc_library) + cmake_parse_arguments( + CC_ARGS + "STATIC;SHARED;EXCLUDE;PACKED;SRCS_NO_GLOB" + "NAME;VERSION" + "SRCS;INCS;PUBINCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS;PACKED_EXCLUDES" + ${ARGN} + ) + + if(NOT CC_ARGS_NAME) + message(FATAL_ERROR "No target name provided.") + endif() + + if(CC_ARGS_SRCS_NO_GLOB) + set(SOURCE_FILES ${CC_ARGS_SRCS}) + if(NOT SOURCE_FILES) + message(FATAL_ERROR "No source files provided for ${CC_ARGS_NAME} (SRCS_NO_GLOB mode).") + endif() + else() + set(SOURCE_FILES "") + foreach(_src IN LISTS CC_ARGS_SRCS) + if(IS_ABSOLUTE "${_src}" OR NOT "${_src}" MATCHES "[*?]") + list(APPEND SOURCE_FILES "${_src}") + else() + file(GLOB _globbed_srcs ${_src}) + list(APPEND SOURCE_FILES ${_globbed_srcs}) + endif() + endforeach() + if(NOT SOURCE_FILES) + message(FATAL_ERROR "No source files found for ${CC_ARGS_NAME} after globbing.") + endif() + endif() + + if(CC_ARGS_VERSION) + string(REPLACE "-" "_" MACRO_PREFIX "${CC_ARGS_NAME}") + list(APPEND CC_ARGS_DEFS ${MACRO_PREFIX}_VERSION="${CC_ARGS_VERSION}") + endif() + + if(CC_ARGS_EXCLUDE) + set(EXCLUDE_OPTION EXCLUDE_FROM_ALL) + endif() + + if(CC_ARGS_SHARED AND CC_ARGS_STATIC) + _add_library(${CC_ARGS_NAME} "${EXCLUDE_OPTION}" ${SOURCE_FILES}) + elseif(CC_ARGS_SHARED) + add_library(${CC_ARGS_NAME} SHARED ${EXCLUDE_OPTION} ${SOURCE_FILES}) + elseif(CC_ARGS_STATIC) + add_library(${CC_ARGS_NAME} STATIC ${EXCLUDE_OPTION} ${SOURCE_FILES}) + else() + add_library(${CC_ARGS_NAME} ${EXCLUDE_OPTION} ${SOURCE_FILES}) + endif() + + if(TARGET ${CC_ARGS_NAME}_objects) + _cc_target_properties( + NAME "${CC_ARGS_NAME}_objects" + INCS "${CC_ARGS_INCS};${CC_ARGS_PUBINCS}" + DEFS "${CC_ARGS_DEFS}" + LIBS "${CC_ARGS_LIBS}" + CFLAGS "${CC_ARGS_CFLAGS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + "${CC_ARGS_UNPARSED_ARGUMENTS}" + ) + endif() + + if(TARGET ${CC_ARGS_NAME}_static) + _cc_target_properties( + NAME "${CC_ARGS_NAME}_static" + INCS "${CC_ARGS_INCS}" + PUBINCS "${CC_ARGS_PUBINCS}" + DEFS "${CC_ARGS_DEFS}" + LIBS "${CC_ARGS_LIBS}" + CFLAGS "${CC_ARGS_CFLAGS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + "${CC_ARGS_UNPARSED_ARGUMENTS}" + ) + if(CC_ARGS_PACKED) + install( + TARGETS ${CC_ARGS_NAME}_static + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) + endif() + endif() + + _cc_target_properties( + NAME "${CC_ARGS_NAME}" + INCS "${CC_ARGS_INCS}" + PUBINCS "${CC_ARGS_PUBINCS}" + DEFS "${CC_ARGS_DEFS}" + LIBS "${CC_ARGS_LIBS}" + CFLAGS "${CC_ARGS_CFLAGS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + VERSION "${CC_ARGS_VERSION}" + "${CC_ARGS_UNPARSED_ARGUMENTS}" + ) + if(CC_ARGS_PACKED) + install( + TARGETS ${CC_ARGS_NAME} + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" + LIBRARY DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) + if(CC_ARGS_PUBINCS) + foreach(PACKED_EXCLUDE ${CC_ARGS_PACKED_EXCLUDES}) + list(APPEND PATTERN_EXCLUDES "PATTERN;${PACKED_EXCLUDE};EXCLUDE") + endforeach() + install( + DIRECTORY ${CC_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR} + FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp" PATTERN "*.hxx" + ${PATTERN_EXCLUDES} + ) + endif() + endif() +endfunction() + +## Build a C/C++ executable program +function(cc_binary) + cmake_parse_arguments( + CC_ARGS "PACKED" "NAME;VERSION" + "SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS" ${ARGN} + ) + + if(NOT CC_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB CC_ARGS_SRCS ${CC_ARGS_SRCS}) + if(NOT CC_ARGS_SRCS) + message(FATAL_ERROR "No source files found of ${CC_ARGS_NAME}.") + endif() + + if(CC_ARGS_VERSION) + string(REPLACE "-" "_" MACRO_PREFIX "${CC_ARGS_NAME}") + list(APPEND CC_ARGS_DEFS ${MACRO_PREFIX}_VERSION="${CC_ARGS_VERSION}") + endif() + add_executable(${CC_ARGS_NAME} ${CC_ARGS_SRCS}) + + # iOS: set bundle properties for simulator/device installation + if(IOS) + set_target_properties(${CC_ARGS_NAME} PROPERTIES + MACOSX_BUNDLE_INFO_PLIST "${PROJECT_ROOT_DIR}/cmake/iOSBundleInfo.plist.in" + ) + endif() + + if(CC_ARGS_PACKED) + install( + TARGETS ${CC_ARGS_NAME} RUNTIME DESTINATION "${CMAKE_INSTALL_BINDIR}" + ) + endif() + + _cc_target_properties( + NAME "${CC_ARGS_NAME}" + INCS "${CC_ARGS_INCS}" + DEFS "${CC_ARGS_DEFS}" + LIBS "${CC_ARGS_LIBS}" + CFLAGS "${CC_ARGS_CFLAGS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + VERSION "${CC_ARGS_VERSION}" + "${CC_ARGS_UNPARSED_ARGUMENTS}" + ) +endfunction() + +## Build a C/C++ executable test program +function(cc_test) + cmake_parse_arguments( + CC_ARGS "" "NAME;VERSION" + "SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS;ARGS" ${ARGN} + ) + + if(NOT CC_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB CC_ARGS_SRCS ${CC_ARGS_SRCS}) + if(NOT CC_ARGS_SRCS) + message(FATAL_ERROR "No source files found of ${CC_ARGS_NAME}.") + endif() + + if(CC_ARGS_VERSION) + string(REPLACE "-" "_" MACRO_PREFIX "${CC_ARGS_NAME}") + list(APPEND CC_ARGS_DEFS ${MACRO_PREFIX}_VERSION="${CC_ARGS_VERSION}") + endif() + # iOS: add sandbox helper to redirect CWD to writable directory + if(IOS) + list(APPEND CC_ARGS_SRCS "${PROJECT_ROOT_DIR}/tests/ios_test_sandbox.cc") + # Arrow's iOS code references CoreFoundation symbols; link Apple frameworks + list(APPEND CC_ARGS_LDFLAGS + -framework CoreFoundation + -framework CoreGraphics + -framework CoreData + -framework CoreText + -framework Security + -framework Foundation + -Wl,-U,_MallocExtension_ReleaseFreeMemory + -Wl,-U,_ProfilerStart + -Wl,-U,_ProfilerStop + -Wl,-U,_RegisterThriftProtocol + ) + endif() + + add_executable(${CC_ARGS_NAME} EXCLUDE_FROM_ALL ${CC_ARGS_SRCS}) + + # iOS: set bundle properties for simulator/device installation + if(IOS) + set_target_properties(${CC_ARGS_NAME} PROPERTIES + MACOSX_BUNDLE_INFO_PLIST "${PROJECT_ROOT_DIR}/cmake/iOSBundleInfo.plist.in" + ) + endif() + + _cc_target_properties( + NAME "${CC_ARGS_NAME}" + INCS "${CC_ARGS_INCS}" + DEFS "${CC_ARGS_DEFS}" + LIBS "${CC_ARGS_LIBS}" + CFLAGS "${CC_ARGS_CFLAGS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + "${CC_ARGS_UNPARSED_ARGUMENTS}" + ) + add_dependencies(unittest ${CC_ARGS_NAME}) + set(TEST_WORKING_DIR "${CMAKE_BINARY_DIR}/test_tmp/${CC_ARGS_NAME}") + file(MAKE_DIRECTORY "${TEST_WORKING_DIR}") + add_custom_target( + unittest.${CC_ARGS_NAME} + COMMAND $ "${CC_ARGS_ARGS}" + WORKING_DIRECTORY ${TEST_WORKING_DIR} + DEPENDS ${CC_ARGS_NAME} + ) + add_test( + NAME ${CC_ARGS_NAME} + COMMAND $ "${CC_ARGS_ARGS}" + WORKING_DIRECTORY ${TEST_WORKING_DIR} + ) +endfunction() + +## Add existing test cases to a test suite +function(cc_test_suite _NAME) + if(NOT TARGET unittest.${_NAME}) + add_custom_target(unittest.${_NAME} COMMAND "") + endif() + foreach(TEST_TARGET ${ARGN}) + list(APPEND TEST_TARGETS unittest.${TEST_TARGET}) + endforeach() + if(TEST_TARGETS) + add_dependencies(unittest.${_NAME} ${TEST_TARGETS}) + endif() +endfunction() + +## Import a C/C++ static or shared library +function(cc_import) + cmake_parse_arguments( + CC_ARGS "STATIC;SHARED;PACKED" + "NAME;PATH;IMPLIB" "INCS;PUBINCS;DEPS;PACKED_EXCLUDES" ${ARGN} + ) + + if(NOT CC_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB CC_ARGS_PATH ${CC_ARGS_PATH}) + if(NOT CC_ARGS_PATH) + message(FATAL_ERROR "No imported target file found of ${CC_ARGS_NAME}.") + endif() + if(MSVC AND CC_ARGS_SHARED AND NOT CC_ARGS_IMPLIB) + string(REGEX REPLACE + ".[Dd][Ll][Ll]$" ".lib" CC_ARGS_IMPLIB ${CC_ARGS_PATH} + ) + endif() + + if(CC_ARGS_SHARED) + add_library(${CC_ARGS_NAME} SHARED IMPORTED GLOBAL) + elseif(CC_ARGS_STATIC) + add_library(${CC_ARGS_NAME} STATIC IMPORTED GLOBAL) + else() + add_library(${CC_ARGS_NAME} UNKNOWN IMPORTED GLOBAL) + endif() + + set_property( + TARGET ${CC_ARGS_NAME} PROPERTY IMPORTED_LOCATION ${CC_ARGS_PATH} + ) + if(MSVC AND CC_ARGS_SHARED) + set_property( + TARGET ${CC_ARGS_NAME} PROPERTY IMPORTED_IMPLIB ${CC_ARGS_IMPLIB} + ) + endif() + + if(CC_ARGS_INCS) + _absolute_paths(INC_DIRS ${CC_ARGS_INCS}) + foreach(INC_DIR ${INC_DIRS}) + set_property( + TARGET ${CC_ARGS_NAME} APPEND PROPERTY + INTERFACE_INCLUDE_DIRECTORIES "${INC_DIR}" + ) + endforeach() + endif() + + if(CC_ARGS_PUBINCS) + _absolute_paths(INC_DIRS ${CC_ARGS_PUBINCS}) + foreach(INC_DIR ${INC_DIRS}) + set_property( + TARGET ${CC_ARGS_NAME} APPEND PROPERTY + INTERFACE_INCLUDE_DIRECTORIES "${INC_DIR}" + ) + endforeach() + endif() + + if(CC_ARGS_DEPS) + add_dependencies(${CC_ARGS_NAME} "${CC_ARGS_DEPS}") + endif() + + if(CC_ARGS_PACKED) + install( + TARGETS ${CC_ARGS_NAME} + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" + LIBRARY DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) + if(CC_ARGS_PUBINCS) + foreach(PACKED_EXCLUDE ${CC_ARGS_PACKED_EXCLUDES}) + list(APPEND PATTERN_EXCLUDES "PATTERN;${PACKED_EXCLUDE};EXCLUDE") + endforeach() + install( + DIRECTORY ${CC_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR} + FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp" PATTERN "*.hxx" + ${PATTERN_EXCLUDES} + ) + endif() + endif() +endfunction() + +## Import a C/C++ interface library +function(cc_interface) + cmake_parse_arguments( + CC_ARGS "PACKED" "NAME" "INCS;PUBINCS;DEPS;PACKED_EXCLUDES" ${ARGN} + ) + + if(NOT CC_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + add_library(${CC_ARGS_NAME} INTERFACE GLOBAL) + if(CC_ARGS_INCS) + _absolute_paths(INC_DIRS ${CC_ARGS_INCS}) + target_include_directories(${CC_ARGS_NAME} INTERFACE "${INC_DIRS}") + endif() + + if(CC_ARGS_PUBINCS) + _absolute_paths(INC_DIRS ${CC_ARGS_PUBINCS}) + target_include_directories(${CC_ARGS_NAME} INTERFACE "${INC_DIRS}") + endif() + + if(CC_ARGS_DEPS) + add_dependencies(${CC_ARGS_NAME} "${CC_ARGS_DEPS}") + endif() + + if(CC_ARGS_PACKED AND CC_ARGS_PUBINCS) + foreach(PACKED_EXCLUDE ${CC_ARGS_PACKED_EXCLUDES}) + list(APPEND PATTERN_EXCLUDES "PATTERN;${PACKED_EXCLUDE};EXCLUDE") + endforeach() + install( + DIRECTORY ${CC_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR} + FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp" PATTERN "*.hxx" + ${PATTERN_EXCLUDES} + ) + endif() +endfunction() + +## Find gtest library +function(_find_gtest) + if(DEFINED FIND_GTEST_LIBS AND DEFINED FIND_GTEST_INCS) + return() + endif() + + if(NOT TARGET gtest OR NOT TARGET gtest_main) + # Find gtest using 'find_package' + find_package(GTest REQUIRED) + set( + FIND_GTEST_INCS "${GTEST_INCLUDE_DIRS}" + CACHE STRING "GTest includes" + ) + set( + FIND_GTEST_LIBS "${GTEST_BOTH_LIBRARIES}" + CACHE STRING "GTest libraries" + ) + else() + # Find gtest using target names + set(FIND_GTEST_INCS "" CACHE STRING "GTest includes") + if(ANDROID) + # On Android, use a custom main that calls _exit() to skip static + # destructors and avoid glog/gflags teardown crashes. + if(NOT TARGET zvec_gtest_main) + add_library(zvec_gtest_main STATIC + ${PROJECT_ROOT_DIR}/tests/android_gtest_main.cc) + target_link_libraries(zvec_gtest_main PUBLIC gtest) + endif() + set(FIND_GTEST_LIBS "gtest;zvec_gtest_main" CACHE STRING "GTest libraries") + else() + set(FIND_GTEST_LIBS "gtest;gtest_main" CACHE STRING "GTest libraries") + endif() + endif() +endfunction() + +## Build a C/C++ executable google test program +function(cc_gtest) + cmake_parse_arguments( + CC_ARGS "" "NAME;VERSION" + "SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS;ARGS" ${ARGN} + ) + _find_gtest() + cc_test( + NAME "${CC_ARGS_NAME}" + VERSION "${CC_ARGS_VERSION}" + SRCS "${CC_ARGS_SRCS}" + INCS "${CC_ARGS_INCS};${FIND_GTEST_INCS}" + DEFS "${CC_ARGS_DEFS}" + LIBS "${CC_ARGS_LIBS};${FIND_GTEST_LIBS}" + CFLAGS "${CC_ARGS_CFLAGS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + ARGS "${CC_ARGS_ARGS}" + ) +endfunction() + +## Find gmock library +function(_find_gmock) + if(DEFINED FIND_GMOCK_LIBS AND DEFINED FIND_GMOCK_INCS) + return() + endif() + + if(NOT TARGET gmock OR NOT TARGET gmock_main) + # Find gmock/gtest using 'find_package' + find_package(GMock REQUIRED) + find_package(GTest REQUIRED) + set( + FIND_GMOCK_INCS "${GMOCK_INCLUDE_DIRS};${GTEST_INCLUDE_DIRS}" + CACHE STRING "GMock includes" + ) + set( + FIND_GMOCK_LIBS "${GMOCK_BOTH_LIBRARIES};${GTEST_LIBRARIES}" + CACHE STRING "GMock libraries" + ) + else() + # Find gmock using target names + set(FIND_GMOCK_INCS "" CACHE STRING "GMock includes") + if(ANDROID) + # On Android, use a custom main that calls _exit() to skip static + # destructors and avoid glog/gflags teardown crashes. + if(NOT TARGET zvec_gmock_main) + add_library(zvec_gmock_main STATIC + ${PROJECT_ROOT_DIR}/tests/android_gmock_main.cc) + target_link_libraries(zvec_gmock_main PUBLIC gmock gtest) + endif() + set(FIND_GMOCK_LIBS "gmock;zvec_gmock_main" CACHE STRING "GMock libraries") + else() + set(FIND_GMOCK_LIBS "gmock;gmock_main" CACHE STRING "GMock libraries") + endif() + endif() +endfunction() + +## Build a C/C++ executable google mock program +function(cc_gmock) + cmake_parse_arguments( + CC_ARGS "" "NAME;VERSION" + "SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS;ARGS" ${ARGN} + ) + _find_gmock() + cc_test( + NAME "${CC_ARGS_NAME}" + VERSION "${CC_ARGS_VERSION}" + SRCS "${CC_ARGS_SRCS}" + INCS "${CC_ARGS_INCS};${FIND_GMOCK_INCS}" + DEFS "${CC_ARGS_DEFS}" + LIBS "${CC_ARGS_LIBS};${FIND_GMOCK_LIBS}" + CFLAGS "${CC_ARGS_CFLAGS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + ARGS "${CC_ARGS_ARGS}" + ) +endfunction() + +## Find protobuf library +function(_find_protobuf _VERSION) + if(DEFINED CC_PROTOBUF_PROTOC_${_VERSION}) + return() + endif() + + # Find protobuf using 'find_package' + if(NOT TARGET protoc OR NOT TARGET libprotobuf) + find_package(Protobuf ${_VERSION} REQUIRED) + set( + CC_PROTOBUF_PROTOC_${_VERSION} + "${PROTOBUF_PROTOC_EXECUTABLE}" CACHE PATH "Protobuf compiler" + ) + set( + CC_PROTOBUF_INCS_${_VERSION} + "${PROTOBUF_INCLUDE_DIRS}" CACHE STRING "Protobuf includes" + ) + set( + CC_PROTOBUF_LIBS_${_VERSION} + "${PROTOBUF_LIBRARIES}" CACHE STRING "Protobuf libraries" + ) + return() + endif() + + # Find protobuf using target names + get_target_property(protoc_VERSION protoc VERSION) + get_target_property(libprotobuf_VERSION libprotobuf VERSION) + if(_VERSION) + if(${protoc_VERSION} VERSION_LESS ${_VERSION}) + message( + FATAL_ERROR + "The 'protoc' version is ${protoc_VERSION}, less than ${_VERSION}." + ) + endif() + if(${libprotobuf_VERSION} VERSION_LESS ${_VERSION}) + message( + FATAL_ERROR + "The 'libprotobuf' version is ${libprotobuf_VERSION}, " + "less than ${_VERSION}." + ) + endif() + endif() + + message(STATUS "Found binary 'protoc ${protoc_VERSION}'") + message(STATUS "Found library 'libprotobuf ${libprotobuf_VERSION}'") + set( + CC_PROTOBUF_PROTOC_${_VERSION} + "$" CACHE PATH "Protobuf compiler" + ) + get_target_property(protoc_SOURCE_DIR protoc SOURCE_DIR) + get_filename_component(protoc_INCLUDE_DIR ${protoc_SOURCE_DIR}/../src ABSOLUTE) + set( + CC_PROTOBUF_INCS_${_VERSION} + "${protoc_INCLUDE_DIR}" CACHE STRING "Protobuf includes" + ) + set( + CC_PROTOBUF_LIBS_${_VERSION} libprotobuf CACHE STRING "Protobuf libraries" + ) +endfunction() + +## Build a C++ protobuf static or shared library +function(cc_proto_library) + cmake_parse_arguments( + CC_ARGS "STATIC;SHARED;EXCLUDE;PACKED" + "NAME;VERSION;PROTOROOT;PROTOBUF_VERSION" + "SRCS;CXXFLAGS;LDFLAGS;DEPS" ${ARGN} + ) + + _find_protobuf("${CC_ARGS_PROTOBUF_VERSION}") + set(CC_PROTOBUF_PROTOC ${CC_PROTOBUF_PROTOC_${CC_ARGS_PROTOBUF_VERSION}}) + if(DEFINED GLOBAL_CC_PROTOBUF_PROTOC) + set(CC_PROTOBUF_PROTOC ${GLOBAL_CC_PROTOBUF_PROTOC}) + endif() + set(CC_PROTOBUF_INCS ${CC_PROTOBUF_INCS_${CC_ARGS_PROTOBUF_VERSION}}) + set(CC_PROTOBUF_LIBS ${CC_PROTOBUF_LIBS_${CC_ARGS_PROTOBUF_VERSION}}) + + if(NOT CC_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB CC_ARGS_SRCS ${CC_ARGS_SRCS}) + if(NOT CC_ARGS_SRCS) + message(FATAL_ERROR "No source files found of ${CC_ARGS_NAME}.") + endif() + + if(CC_ARGS_VERSION) + string(REPLACE "-" "_" MACRO_PREFIX "${CC_ARGS_NAME}") + list(APPEND CC_ARGS_DEFS ${MACRO_PREFIX}_VERSION="${CC_ARGS_VERSION}") + endif() + + if(CC_ARGS_EXCLUDE) + set(EXCLUDE_OPTION EXCLUDE_FROM_ALL) + endif() + + set(PROTO_ROOT ${CMAKE_CURRENT_SOURCE_DIR}) + if(CC_ARGS_PROTOROOT) + get_filename_component(PROTO_ROOT ${CC_ARGS_PROTOROOT} ABSOLUTE) + endif() + + # Compile proto files to C++ sources + set(CPP_OUTPATH "${CMAKE_CURRENT_BINARY_DIR}") + foreach(PROTO_FILE ${CC_ARGS_SRCS}) + get_filename_component(PROTO_FILE ${PROTO_FILE} ABSOLUTE) + + if(NOT ${PROTO_FILE} MATCHES "\\.proto$$") + message(FATAL_ERROR "Unrecognized proto file ${PROTOFILE}") + endif() + if(NOT ${PROTO_FILE} MATCHES "^${PROTO_ROOT}") + message(FATAL_ERROR "'${PROTO_FILE}' NOT IN '${PROTO_ROOT}'") + endif() + + string( + REGEX REPLACE "^${PROTO_ROOT}(/?)" "" ROOT_CLEANED_FILE ${PROTO_FILE} + ) + string(REGEX REPLACE "\\.proto$$" "" EXT_CLEANED_FILE ${ROOT_CLEANED_FILE}) + set(CPP_FILE "${CPP_OUTPATH}/${EXT_CLEANED_FILE}.pb.cc") + set(HDR_FILE "${CPP_OUTPATH}/${EXT_CLEANED_FILE}.pb.h") + set(INJ_FILE "${CPP_OUTPATH}/${EXT_CLEANED_FILE}.pb.cmake") + file(RELATIVE_PATH REL_CPP_FILE ${CMAKE_BINARY_DIR} ${CPP_FILE}) + + set(INJECTED_SCRIPT + "foreach(SRC ${EXT_CLEANED_FILE}.pb.cc ${EXT_CLEANED_FILE}.pb.h)\n" + " file(READ \$\{SRC\} SRC_CODE)\n" + " file(REMOVE \$\{SRC\})\n" + " file(APPEND \$\{SRC\} \"#ifdef __GNUC__\\n\")\n" + " file(APPEND \$\{SRC\} \"#pragma GCC diagnostic push\\n\")\n" + " file(APPEND \$\{SRC\} \"#pragma GCC diagnostic ignored \\\"-Wshadow\\\"\\n\")\n" + " file(APPEND \$\{SRC\} \"#pragma GCC diagnostic ignored \\\"-Wunused-parameter\\\"\\n\")\n" + " file(APPEND \$\{SRC\} \"#endif\\n\\n\")\n" + " file(APPEND \$\{SRC\} \"\$\{SRC_CODE\}\")\n" + " file(APPEND \$\{SRC\} \"\\n#ifdef __GNUC__\\n\")\n" + " file(APPEND \$\{SRC\} \"#pragma GCC diagnostic pop\\n\")\n" + " file(APPEND \$\{SRC\} \"#endif\\n\")\n" + "endforeach()\n" + ) + file(WRITE "${INJ_FILE}" ${INJECTED_SCRIPT}) + + add_custom_command( + OUTPUT "${CPP_FILE}" "${HDR_FILE}" + # COMMAND ${CMAKE_COMMAND} -E make_directory ${CPP_OUTPATH} + COMMAND ${CC_PROTOBUF_PROTOC} + --cpp_out "${CPP_OUTPATH}" --python_out "${CPP_OUTPATH}" + --proto_path "${PROTO_ROOT}" --proto_path "${CC_PROTOBUF_INCS}" "${PROTO_FILE}" + + COMMAND ${CMAKE_COMMAND} -P "${INJ_FILE}" + DEPENDS "${PROTO_FILE}" + COMMENT "Generating CXX source ${REL_CPP_FILE}" + VERBATIM + ) + list(APPEND CC_SRCS "${CPP_FILE}" "${HDR_FILE}") + endforeach() + + # Compile C++ sources + if(CC_ARGS_SHARED AND CC_ARGS_STATIC) + _add_library(${CC_ARGS_NAME} "${EXCLUDE_OPTION}" "${CC_SRCS}") + elseif(CC_ARGS_SHARED) + add_library(${CC_ARGS_NAME} SHARED ${EXCLUDE_OPTION} ${CC_SRCS}) + elseif(CC_ARGS_STATIC) + add_library(${CC_ARGS_NAME} STATIC ${EXCLUDE_OPTION} ${CC_SRCS}) + else() + add_library(${CC_ARGS_NAME} ${EXCLUDE_OPTION} ${CC_SRCS}) + endif() + + if(TARGET ${CC_ARGS_NAME}_objects) + _cc_target_properties( + NAME "${CC_ARGS_NAME}_objects" + INCS "${CPP_OUTPATH};${CC_PROTOBUF_INCS}" + LIBS "${CC_PROTOBUF_LIBS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + "${CC_ARGS_UNPARSED_ARGUMENTS}" + ) + endif() + + if(TARGET ${CC_ARGS_NAME}_static) + _cc_target_properties( + NAME "${CC_ARGS_NAME}_static" + PUBINCS "${CPP_OUTPATH};${CC_PROTOBUF_INCS}" + LIBS "${CC_PROTOBUF_LIBS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + "${CC_ARGS_UNPARSED_ARGUMENTS}" + ) + if(CC_ARGS_PACKED) + install( + TARGETS ${CC_ARGS_NAME}_static + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) + endif() + endif() + + _cc_target_properties( + NAME "${CC_ARGS_NAME}" + PUBINCS "${CPP_OUTPATH};${CC_PROTOBUF_INCS}" + LIBS "${CC_PROTOBUF_LIBS}" + CXXFLAGS "${CC_ARGS_CXXFLAGS}" + LDFLAGS "${CC_ARGS_LDFLAGS}" + DEPS "${CC_ARGS_DEPS}" + VERSION "${CC_ARGS_VERSION}" + "${CC_ARGS_UNPARSED_ARGUMENTS}" + ) + if(CC_ARGS_PACKED) + install( + TARGETS ${CC_ARGS_NAME} + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" + LIBRARY DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) + endif() +endfunction() + +## Add a subdirectory to the build +function(cuda_directory) + if(NOT CMAKE_CUDA_COMPILER) + message(FATAL_ERROR "No CUDA language supported.") + endif() + cc_directory(${ARGN}) +endfunction() + +## Add subdirectories to the build +function(cuda_directories) + if(NOT CMAKE_CUDA_COMPILER) + message(FATAL_ERROR "No CUDA language supported.") + endif() + cc_directories(${ARGN}) +endfunction() + +## Set the properties of cuda target +function(_cuda_target_properties) + cmake_parse_arguments( + CUDA_ARGS "STRICT;ALWAYS_LINK" "NAME;VERSION;C_STANDARD;CXX_STANDARD" + "INCS;PUBINCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS" ${ARGN} + ) + + if(NOT CUDA_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + get_target_property(TARGET_TYPE ${CUDA_ARGS_NAME} TYPE) + if(("${TARGET_TYPE}" STREQUAL "SHARED_LIBRARY") OR + ("${TARGET_TYPE}" STREQUAL "STATIC_LIBRARY") OR + ("${TARGET_TYPE}" STREQUAL "EXECUTABLE")) + set(TARGET_LINKABLE TRUE) + endif() + + if(CUDA_ARGS_ALWAYS_LINK) + if(("${TARGET_TYPE}" STREQUAL "STATIC_LIBRARY") OR + ("${TARGET_TYPE}" STREQUAL "OBJECT_LIBRARY")) + set_property(TARGET ${CUDA_ARGS_NAME} PROPERTY ALWAYS_LINK TRUE) + endif() + endif() + + # Set the warning level of compiling + if(CUDA_ARGS_STRICT) + target_compile_options( + ${CUDA_ARGS_NAME} PRIVATE "${BAZEL_CUDA_STRICT_COMPILE_FLAGS}" + ) + if(TARGET_LINKABLE) + target_link_libraries( + ${CUDA_ARGS_NAME} "${BAZEL_CUDA_STRICT_LINK_FLAGS}" + ) + endif() + else() + target_compile_options( + ${CUDA_ARGS_NAME} PRIVATE "${BAZEL_CUDA_UNSTRICT_COMPILE_FLAGS}" + ) + if(TARGET_LINKABLE) + target_link_libraries( + ${CUDA_ARGS_NAME} "${BAZEL_CUDA_UNSTRICT_LINK_FLAGS}" + ) + endif() + endif() + + target_compile_options( + ${CUDA_ARGS_NAME} PRIVATE + "$<$:-ccbin=${CMAKE_CXX_COMPILER}>" + ) + + if(CUDA_ARGS_DEFS) + target_compile_definitions(${CUDA_ARGS_NAME} PRIVATE "${CUDA_ARGS_DEFS}") + endif() + + if(CUDA_ARGS_CFLAGS OR CUDA_ARGS_CXXFLAGS OR CUDA_ARGS_CUDAFLAGS) + target_compile_options( + ${CUDA_ARGS_NAME} PRIVATE + "$<$:${CUDA_ARGS_CFLAGS}>" + "$<$:${CUDA_ARGS_CXXFLAGS}>" + "$<$:${CUDA_ARGS_CUDAFLAGS}>" + ) + endif() + + if(CUDA_ARGS_LDFLAGS) + string(REPLACE ";" " " CUDA_ARGS_LDFLAGS "${CUDA_ARGS_LDFLAGS}") + set_property( + TARGET ${CUDA_ARGS_NAME} PROPERTY LINK_FLAGS "${CUDA_ARGS_LDFLAGS}" + ) + endif() + + if(CUDA_ARGS_INCS) + _absolute_paths(INC_DIRS ${CUDA_ARGS_INCS}) + target_include_directories(${CUDA_ARGS_NAME} PRIVATE "${INC_DIRS}") + endif() + + target_include_directories( + ${CUDA_ARGS_NAME} PRIVATE "${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}" + ) + + if(BAZEL_WORKSPACE_DIR) + target_include_directories( + ${CUDA_ARGS_NAME} PRIVATE "${BAZEL_WORKSPACE_DIR}" + ) + endif() + + if(CUDA_ARGS_PUBINCS) + _absolute_paths(INC_DIRS ${CUDA_ARGS_PUBINCS}) + target_include_directories(${CUDA_ARGS_NAME} PUBLIC "${INC_DIRS}") + endif() + + if(CUDA_ARGS_LIBS) + if(NOT TARGET_LINKABLE) + _targets_link_dependencies(${CUDA_ARGS_NAME} ${CUDA_ARGS_LIBS}) + else() + if ("${TARGET_TYPE}" STREQUAL "EXECUTABLE") + _target_link_libraries(${CUDA_ARGS_NAME} "${CUDA_ARGS_LIBS}") + else() + target_link_libraries(${CUDA_ARGS_NAME} "${CUDA_ARGS_LIBS}") + endif() + endif() + endif() + + if(CUDA_ARGS_DEPS) + add_dependencies(${CUDA_ARGS_NAME} "${CUDA_ARGS_DEPS}") + endif() + + if(CUDA_ARGS_VERSION) + set_property( + TARGET ${CUDA_ARGS_NAME} PROPERTY VERSION "${CUDA_ARGS_VERSION}" + ) + endif() + + if(NOT CUDA_C_STANDARD) + set(CUDA_C_STANDARD 99) + endif() + + if(NOT CUDA_CXX_STANDARD) + set(CUDA_CXX_STANDARD 11) + endif() + + set_target_properties( + ${CUDA_ARGS_NAME} PROPERTIES DEFINE_SYMBOL "" + C_STANDARD ${CUDA_C_STANDARD} CXX_STANDARD ${CUDA_CXX_STANDARD} + C_STANDARD_REQUIRED ON C_EXTENSIONS ON + CXX_STANDARD_REQUIRED ON CXX_EXTENSIONS OFF + CUDA_STANDARD 11 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF + WINDOWS_EXPORT_ALL_SYMBOLS ON + ) +endfunction() + +## Build a CUDA static or shared library +function(cuda_library) + if(NOT CMAKE_CUDA_COMPILER) + message(FATAL_ERROR "No CUDA language supported.") + endif() + + cmake_parse_arguments( + CUDA_ARGS "STATIC;SHARED;EXCLUDE;PACKED" "NAME;VERSION" + "SRCS;INCS;PUBINCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS;PACKED_EXCS" + ${ARGN} + ) + + if(NOT CUDA_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB CUDA_ARGS_SRCS ${CUDA_ARGS_SRCS}) + if(NOT CUDA_ARGS_SRCS) + message(FATAL_ERROR "No source files found of ${CUDA_ARGS_NAME}.") + endif() + + if(CUDA_ARGS_VERSION) + string(REPLACE "-" "_" MACRO_PREFIX "${CUDA_ARGS_NAME}") + list(APPEND CUDA_ARGS_DEFS ${MACRO_PREFIX}_VERSION="${CUDA_ARGS_VERSION}") + endif() + + if(CUDA_ARGS_EXCLUDE) + set(EXCLUDE_OPTION EXCLUDE_FROM_ALL) + endif() + + if(CUDA_ARGS_SHARED AND CUDA_ARGS_STATIC) + _add_library(${CUDA_ARGS_NAME} "${EXCLUDE_OPTION}" "${CUDA_ARGS_SRCS}") + elseif(CUDA_ARGS_SHARED) + add_library(${CUDA_ARGS_NAME} SHARED ${EXCLUDE_OPTION} ${CUDA_ARGS_SRCS}) + elseif(CUDA_ARGS_STATIC) + add_library(${CUDA_ARGS_NAME} STATIC ${EXCLUDE_OPTION} ${CUDA_ARGS_SRCS}) + else() + add_library(${CUDA_ARGS_NAME} ${EXCLUDE_OPTION} ${CUDA_ARGS_SRCS}) + endif() + + if(TARGET ${CUDA_ARGS_NAME}_objects) + _cuda_target_properties( + NAME "${CUDA_ARGS_NAME}_objects" + INCS "${CUDA_ARGS_INCS};${CUDA_ARGS_PUBINCS}" + DEFS "${CUDA_ARGS_DEFS}" + LIBS "${CUDA_ARGS_LIBS}" + CFLAGS "${CUDA_ARGS_CFLAGS}" + CXXFLAGS "${CUDA_ARGS_CXXFLAGS}" + CUDAFLAGS "${CUDA_ARGS_CUDAFLAGS}" + LDFLAGS "${CUDA_ARGS_LDFLAGS}" + DEPS "${CUDA_ARGS_DEPS}" + "${CUDA_ARGS_UNPARSED_ARGUMENTS}" + ) + endif() + + if(TARGET ${CUDA_ARGS_NAME}_static) + _cuda_target_properties( + NAME "${CUDA_ARGS_NAME}_static" + INCS "${CUDA_ARGS_INCS}" + PUBINCS "${CUDA_ARGS_PUBINCS}" + DEFS "${CUDA_ARGS_DEFS}" + LIBS "${CUDA_ARGS_LIBS}" + CFLAGS "${CUDA_ARGS_CFLAGS}" + CXXFLAGS "${CUDA_ARGS_CXXFLAGS}" + CUDAFLAGS "${CUDA_ARGS_CUDAFLAGS}" + LDFLAGS "${CUDA_ARGS_LDFLAGS}" + DEPS "${CUDA_ARGS_DEPS}" + "${CUDA_ARGS_UNPARSED_ARGUMENTS}" + ) + if(CUDA_ARGS_PACKED) + install( + TARGETS ${CUDA_ARGS_NAME}_static + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) + endif() + endif() + + _cuda_target_properties( + NAME "${CUDA_ARGS_NAME}" + INCS "${CUDA_ARGS_INCS}" + PUBINCS "${CUDA_ARGS_PUBINCS}" + DEFS "${CUDA_ARGS_DEFS}" + LIBS "${CUDA_ARGS_LIBS}" + CFLAGS "${CUDA_ARGS_CFLAGS}" + CXXFLAGS "${CUDA_ARGS_CXXFLAGS}" + CUDAFLAGS "${CUDA_ARGS_CUDAFLAGS}" + LDFLAGS "${CUDA_ARGS_LDFLAGS}" + DEPS "${CUDA_ARGS_DEPS}" + VERSION "${CUDA_ARGS_VERSION}" + "${CUDA_ARGS_UNPARSED_ARGUMENTS}" + ) + if(CUDA_ARGS_PACKED) + install( + TARGETS ${CUDA_ARGS_NAME} + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" + LIBRARY DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) + if(CUDA_ARGS_PUBINCS) + foreach(PACKED_EXCLUDE ${CUDA_ARGS_PACKED_IGORNES}) + list(APPEND PATTERN_EXCLUDES "PATTERN;${PACKED_EXCLUDE};EXCLUDE") + endforeach() + install( + DIRECTORY ${CUDA_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR} + FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp" + PATTERN "*.hxx" PATTERN "*.cuh" + ${PATTERN_EXCLUDES} + ) + endif() + endif() +endfunction() + +## Build a CUDA executable program +function(cuda_binary) + if(NOT CMAKE_CUDA_COMPILER) + message(FATAL_ERROR "No CUDA language supported.") + endif() + + cmake_parse_arguments( + CUDA_ARGS "PACKED" "NAME;VERSION" + "SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS" ${ARGN} + ) + + if(NOT CUDA_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB CUDA_ARGS_SRCS ${CUDA_ARGS_SRCS}) + if(NOT CUDA_ARGS_SRCS) + message(FATAL_ERROR "No source files found of ${CUDA_ARGS_NAME}.") + endif() + + if(CUDA_ARGS_VERSION) + string(REPLACE "-" "_" MACRO_PREFIX "${CUDA_ARGS_NAME}") + list(APPEND CUDA_ARGS_DEFS ${MACRO_PREFIX}_VERSION="${CUDA_ARGS_VERSION}") + endif() + add_executable(${CUDA_ARGS_NAME} ${CUDA_ARGS_SRCS}) + + if(CUDA_ARGS_PACKED) + install( + TARGETS ${CUDA_ARGS_NAME} RUNTIME DESTINATION "${CMAKE_INSTALL_BINDIR}" + ) + endif() + + _cuda_target_properties( + NAME "${CUDA_ARGS_NAME}" + INCS "${CUDA_ARGS_INCS}" + DEFS "${CUDA_ARGS_DEFS}" + LIBS "${CUDA_ARGS_LIBS}" + CFLAGS "${CUDA_ARGS_CFLAGS}" + CXXFLAGS "${CUDA_ARGS_CXXFLAGS}" + CUDAFLAGS "${CUDA_ARGS_CUDAFLAGS}" + LDFLAGS "${CUDA_ARGS_LDFLAGS}" + DEPS "${CUDA_ARGS_DEPS}" + VERSION "${CUDA_ARGS_VERSION}" + "${CUDA_ARGS_UNPARSED_ARGUMENTS}" + ) +endfunction() + +## Build a CUDA executable test program +function(cuda_test) + if(NOT CMAKE_CUDA_COMPILER) + message(FATAL_ERROR "No CUDA language supported.") + endif() + + cmake_parse_arguments( + CUDA_ARGS "" "NAME;VERSION" + "SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS;ARGS" ${ARGN} + ) + + if(NOT CUDA_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB CUDA_ARGS_SRCS ${CUDA_ARGS_SRCS}) + if(NOT CUDA_ARGS_SRCS) + message(FATAL_ERROR "No source files found of ${CUDA_ARGS_NAME}.") + endif() + + if(CUDA_ARGS_VERSION) + string(REPLACE "-" "_" MACRO_PREFIX "${CUDA_ARGS_NAME}") + list(APPEND CUDA_ARGS_DEFS ${MACRO_PREFIX}_VERSION="${CUDA_ARGS_VERSION}") + endif() + add_executable(${CUDA_ARGS_NAME} EXCLUDE_FROM_ALL ${CUDA_ARGS_SRCS}) + + _cuda_target_properties( + NAME "${CUDA_ARGS_NAME}" + INCS "${CUDA_ARGS_INCS}" + DEFS "${CUDA_ARGS_DEFS}" + LIBS "${CUDA_ARGS_LIBS}" + CFLAGS "${CUDA_ARGS_CFLAGS}" + CXXFLAGS "${CUDA_ARGS_CXXFLAGS}" + CUDAFLAGS "${CUDA_ARGS_CUDAFLAGS}" + LDFLAGS "${CUDA_ARGS_LDFLAGS}" + DEPS "${CUDA_ARGS_DEPS}" + "${CUDA_ARGS_UNPARSED_ARGUMENTS}" + ) + add_dependencies(unittest ${CUDA_ARGS_NAME}) + set(TEST_WORKING_DIR "${CMAKE_BINARY_DIR}/test_tmp/${CUDA_ARGS_NAME}") + file(MAKE_DIRECTORY "${TEST_WORKING_DIR}") + add_custom_target( + unittest.${CUDA_ARGS_NAME} + COMMAND $ "${CUDA_ARGS_ARGS}" + WORKING_DIRECTORY ${TEST_WORKING_DIR} + DEPENDS ${CUDA_ARGS_NAME} + ) + add_test( + NAME ${CUDA_ARGS_NAME} + COMMAND $ "${CUDA_ARGS_ARGS}" + WORKING_DIRECTORY ${TEST_WORKING_DIR} + ) +endfunction() + +## Add existing test cases to a test suite +function(cuda_test_suite) + if(NOT CMAKE_CUDA_COMPILER) + message(FATAL_ERROR "No CUDA language supported.") + endif() + cc_test_suite(${ARGN}) +endfunction() + +## Import a C/C++/CUDA static or shared library +function(cuda_import) + if(NOT CMAKE_CUDA_COMPILER) + message(FATAL_ERROR "No CUDA language supported.") + endif() + + cmake_parse_arguments( + CUDA_ARGS "STATIC;SHARED;PACKED" + "NAME;PATH;IMPLIB" "INCS;PUBINCS;DEPS;PACKED_EXCLUDES" ${ARGN} + ) + + if(NOT CUDA_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB CUDA_ARGS_PATH ${CUDA_ARGS_PATH}) + if(NOT CUDA_ARGS_PATH) + message(FATAL_ERROR "No imported target file found of ${CUDA_ARGS_NAME}.") + endif() + if(MSVC AND CUDA_ARGS_SHARED AND NOT CUDA_ARGS_IMPLIB) + string(REGEX REPLACE + ".[Dd][Ll][Ll]$" ".lib" CUDA_ARGS_IMPLIB ${CUDA_ARGS_PATH} + ) + endif() + + if(CUDA_ARGS_SHARED) + add_library(${CUDA_ARGS_NAME} SHARED IMPORTED GLOBAL) + elseif(CUDA_ARGS_STATIC) + add_library(${CUDA_ARGS_NAME} STATIC IMPORTED GLOBAL) + else() + add_library(${CUDA_ARGS_NAME} UNKNOWN IMPORTED GLOBAL) + endif() + + set_property( + TARGET ${CUDA_ARGS_NAME} PROPERTY IMPORTED_LOCATION ${CUDA_ARGS_PATH} + ) + if(MSVC AND CUDA_ARGS_SHARED) + set_property( + TARGET ${CUDA_ARGS_NAME} PROPERTY IMPORTED_IMPLIB ${CUDA_ARGS_IMPLIB} + ) + endif() + + if(CUDA_ARGS_INCS) + _absolute_paths(INC_DIRS ${CUDA_ARGS_INCS}) + foreach(INC_DIR ${INC_DIRS}) + set_property( + TARGET ${CUDA_ARGS_NAME} APPEND PROPERTY + INTERFACE_INCLUDE_DIRECTORIES "${INC_DIR}" + ) + endforeach() + endif() + + if(CUDA_ARGS_PUBINCS) + _absolute_paths(INC_DIRS ${CUDA_ARGS_PUBINCS}) + foreach(INC_DIR ${INC_DIRS}) + set_property( + TARGET ${CUDA_ARGS_NAME} APPEND PROPERTY + INTERFACE_INCLUDE_DIRECTORIES "${INC_DIR}" + ) + endforeach() + endif() + + if(CUDA_ARGS_DEPS) + add_dependencies(${CUDA_ARGS_NAME} "${CUDA_ARGS_DEPS}") + endif() + + if(CUDA_ARGS_PACKED) + install( + TARGETS ${CUDA_ARGS_NAME} + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" + LIBRARY DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) + if(CUDA_ARGS_PUBINCS) + foreach(PACKED_EXCLUDE ${CUDA_ARGS_PACKED_EXCLUDES}) + list(APPEND PATTERN_EXCLUDES "PATTERN;${PACKED_EXCLUDE};EXCLUDE") + endforeach() + install( + DIRECTORY ${CUDA_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR} + FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp" + PATTERN "*.hxx" PATTERN "*.cuh" + ${PATTERN_EXCLUDES} + ) + endif() + endif() +endfunction() + +## Import a C/C++/CUDA interface library +function(cuda_interface) + if(NOT CMAKE_CUDA_COMPILER) + message(FATAL_ERROR "No CUDA language supported.") + endif() + + cmake_parse_arguments( + CUDA_ARGS "PACKED" "NAME" "INCS;PUBINCS;DEPS;PACKED_EXCLUDES" ${ARGN} + ) + + if(NOT CUDA_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + add_library(${CUDA_ARGS_NAME} INTERFACE GLOBAL) + if(CUDA_ARGS_INCS) + _absolute_paths(INC_DIRS ${CUDA_ARGS_INCS}) + target_include_directories(${CUDA_ARGS_NAME} INTERFACE "${INC_DIRS}") + endif() + + if(CUDA_ARGS_PUBINCS) + _absolute_paths(INC_DIRS ${CUDA_ARGS_PUBINCS}) + target_include_directories(${CUDA_ARGS_NAME} INTERFACE "${INC_DIRS}") + endif() + + if(CUDA_ARGS_DEPS) + add_dependencies(${CUDA_ARGS_NAME} "${CUDA_ARGS_DEPS}") + endif() + + if(CUDA_ARGS_PACKED AND CUDA_ARGS_PUBINCS) + foreach(PACKED_EXCLUDE ${CUDA_ARGS_PACKED_EXCLUDES}) + list(APPEND PATTERN_EXCLUDES "PATTERN;${PACKED_EXCLUDE};EXCLUDE") + endforeach() + install( + DIRECTORY ${CUDA_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR} + FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp" + PATTERN "*.hxx" PATTERN "*.cuh" + ${PATTERN_EXCLUDES} + ) + endif() +endfunction() + +## Build a C/C++/CUDA executable google test program +function(cuda_gtest) + cmake_parse_arguments( + CUDA_ARGS "" "NAME;VERSION" + "SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS;ARGS" ${ARGN} + ) + _find_gtest() + cuda_test( + NAME "${CUDA_ARGS_NAME}" + VERSION "${CUDA_ARGS_VERSION}" + SRCS "${CUDA_ARGS_SRCS}" + INCS "${CUDA_ARGS_INCS};${FIND_GTEST_INCS}" + DEFS "${CUDA_ARGS_DEFS}" + LIBS "${CUDA_ARGS_LIBS};${FIND_GTEST_LIBS}" + CFLAGS "${CUDA_ARGS_CFLAGS}" + CXXFLAGS "${CUDA_ARGS_CXXFLAGS}" + CUDAFLAGS "${CUDA_ARGS_CUDAFLAGS}" + LDFLAGS "${CUDA_ARGS_LDFLAGS}" + DEPS "${CUDA_ARGS_DEPS}" + ARGS "${CUDA_ARGS_ARGS}" + ) +endfunction() + +## Build a C/C++/CUDA executable google mock program +function(cuda_gmock) + cmake_parse_arguments( + CUDA_ARGS "" "NAME;VERSION" + "SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS;ARGS" ${ARGN} + ) + _find_gmock() + cuda_test( + NAME "${CUDA_ARGS_NAME}" + VERSION "${CUDA_ARGS_VERSION}" + SRCS "${CUDA_ARGS_SRCS}" + INCS "${CUDA_ARGS_INCS};${FIND_GMOCK_INCS}" + DEFS "${CUDA_ARGS_DEFS}" + LIBS "${CUDA_ARGS_LIBS};${FIND_GMOCK_LIBS}" + CFLAGS "${CUDA_ARGS_CFLAGS}" + CXXFLAGS "${CUDA_ARGS_CXXFLAGS}" + CUDAFLAGS "${CUDA_ARGS_CUDAFLAGS}" + LDFLAGS "${CUDA_ARGS_LDFLAGS}" + DEPS "${CUDA_ARGS_DEPS}" + ARGS "${CUDA_ARGS_ARGS}" + ) +endfunction() + +## Add a subdirectory to the build +function(go_directory) + add_subdirectory(${ARGN}) +endfunction() + +## Add subdirectories to the build +function(go_directories) + foreach(SRC_DIR ${ARGN}) + add_subdirectory(${SRC_DIR}) + endforeach() +endfunction() + +## Build a go executable program +function(go_binary) + find_program( + GO_EXECUTABLE go PATHS $ENV{HOME}/go ENV GOROOT GOPATH PATH_SUFFIXES bin + ) + if(NOT GO_EXECUTABLE) + message(FATAL_ERROR "No go language compiler found.") + endif() + + cmake_parse_arguments( + GO_ARGS "PACKED" "NAME" + "GOPATH;SRCS;ASMFLAGS;GCFLAGS;LDFLAGS;DEPS" ${ARGN} + ) + if(NOT GO_ARGS_NAME) + message(FATAL_ERROR "No target name privated.") + endif() + + file(GLOB GO_ARGS_SRCS ${GO_ARGS_SRCS}) + if(NOT GO_ARGS_SRCS) + message(FATAL_ERROR "No source files/directories found of ${GO_ARGS_NAME}.") + endif() + + if(${CMAKE_SYSTEM_NAME} MATCHES "Windows") + string(REPLACE ";" "\;" GO_ARGS_GOPATH "${GO_ARGS_GOPATH}") + else() + string(REPLACE ";" ":" GO_ARGS_GOPATH "${GO_ARGS_GOPATH}") + endif() + + set( + GO_OUTPUT_FILE + ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${GO_ARGS_NAME}${CMAKE_EXECUTABLE_SUFFIX} + ) + file(RELATIVE_PATH GO_OUTPUT_REL_FILE ${CMAKE_BINARY_DIR} ${GO_OUTPUT_FILE}) + add_custom_target( + ${GO_ARGS_NAME} + COMMAND ${CMAKE_COMMAND} -E env GOPATH="${GO_ARGS_GOPATH}" + "${GO_EXECUTABLE}" build -v -buildmode=exe + -compiler=gc -gcflags="${GO_ARGS_GCFLAGS}" -asmflags="${GO_ARGS_ASMFLAGS}" + -ldflags="${GO_ARGS_LDFLAGS}" + -o "${GO_OUTPUT_FILE}" "${GO_ARGS_SRCS}" + WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" + DEPENDS "${GO_ARGS_DEPS}" + COMMENT "Building GO executable ${GO_OUTPUT_REL_FILE}" + ) + if(GO_ARGS_PACKED) + install(PROGRAMS ${GO_OUTPUT_FILE} DESTINATION "${CMAKE_INSTALL_BINDIR}") + endif() +endfunction() + +## Fetch content +function(_fetch_content) + cmake_parse_arguments( + DL_ARGS "" + "NAME;PATH;GIT_URL;GIT_TAG;HG_URL;HG_TAG;SVN_URL;SVN_REV;URL;URL_HASH" + "" ${ARGN} + ) + + if(NOT DL_ARGS_NAME) + message(FATAL_ERROR "No fetch name privated.") + endif() + + if(NOT DL_ARGS_PATH) + # Download to current source directory + set(DL_ARGS_PATH "${CMAKE_CURRENT_SOURCE_DIR}/${DL_ARGS_NAME}") + endif() + + set( + CMAKELISTS_CONTENT + "cmake_minimum_required(VERSION 3.13)\n" + "project(${DL_ARGS_NAME})\n" + "include(ExternalProject)\n" + "ExternalProject_Add(\n" + " ${DL_ARGS_NAME}\n" + " PREFIX \"external\"\n" + " GIT_REPOSITORY \"${DL_ARGS_GIT_URL}\"\n" + " GIT_TAG \"${DL_ARGS_GIT_TAG}\"\n" + " HG_REPOSITORY \"${DL_ARGS_HG_URL}\"\n" + " HG_TAG \"${DL_ARGS_HG_TAG}\"\n" + " SVN_REPOSITORY \"${DL_ARGS_SVN_URL}\"\n" + " SVN_REVISION \"${DL_ARGS_SVN_REV}\"\n" + " URL \"${DL_ARGS_URL}\"\n" + " URL_HASH \"${DL_ARGS_URL_HASH}\"\n" + " SOURCE_DIR \"${DL_ARGS_PATH}\"\n" + " BINARY_DIR \"\"\n" + " CONFIGURE_COMMAND \"\"\n" + " BUILD_COMMAND \"\"\n" + " INSTALL_COMMAND \"\"\n" + " TEST_COMMAND \"\"\n" + " LOG_DOWNLOAD ON\n" + " )\n" + ) + set( + CMAKELISTS_DIRECTORY + "${PROJECT_BINARY_DIR}/downloads/${DL_ARGS_NAME}" + ) + add_custom_target( + external.${DL_ARGS_NAME} + COMMAND "${CMAKE_COMMAND}" -G "${CMAKE_GENERATOR}" . && + "${CMAKE_COMMAND}" --build . + WORKING_DIRECTORY "${CMAKELISTS_DIRECTORY}" + ) + + # Write a cmake script into folder + file(WRITE "${CMAKELISTS_DIRECTORY}/CMakeLists.txt" ${CMAKELISTS_CONTENT}) + + execute_process( + COMMAND "${CMAKE_COMMAND}" -G "${CMAKE_GENERATOR}" . + WORKING_DIRECTORY "${CMAKELISTS_DIRECTORY}" + ) + execute_process( + COMMAND "${CMAKE_COMMAND}" --build . + WORKING_DIRECTORY "${CMAKELISTS_DIRECTORY}" + ) +endfunction() + +## Download a git repository +function(git_repository) + cmake_parse_arguments(GIT_ARGS "" "NAME;PATH;URL;TAG" "" ${ARGN}) + + if(NOT GIT_ARGS_NAME) + message(FATAL_ERROR "No repository name privated.") + endif() + if(NOT GIT_ARGS_URL) + message(FATAL_ERROR "No repository URL privated.") + endif() + + if(GIT_ARGS_PATH AND NOT IS_ABSOLUTE ${GIT_ARGS_PATH}) + get_filename_component(GIT_ARGS_PATH ${GIT_ARGS_PATH} ABSOLUTE) + endif() + + _fetch_content( + NAME "${GIT_ARGS_NAME}" + PATH "${GIT_ARGS_PATH}" + GIT_URL "${GIT_ARGS_URL}" + GIT_TAG "${GIT_ARGS_TAG}" + ) +endfunction() + +## Download a hg repository +function(hg_repository) + cmake_parse_arguments(HG_ARGS "" "NAME;PATH;URL;TAG" "" ${ARGN}) + + if(NOT HG_ARGS_NAME) + message(FATAL_ERROR "No repository name privated.") + endif() + if(NOT HG_ARGS_URL) + message(FATAL_ERROR "No repository URL privated.") + endif() + + if(HG_ARGS_PATH AND NOT IS_ABSOLUTE ${HG_ARGS_PATH}) + get_filename_component(HG_ARGS_PATH ${HG_ARGS_PATH} ABSOLUTE) + endif() + + _fetch_content( + NAME "${HG_ARGS_NAME}" + PATH "${HG_ARGS_PATH}" + HG_URL "${HG_ARGS_URL}" + HG_TAG "${HG_ARGS_TAG}" + ) +endfunction() + +## Download a svn repository +function(svn_repository) + cmake_parse_arguments(SVN_ARGS "" "NAME;PATH;URL;REV" "" ${ARGN}) + + if(NOT SVN_ARGS_NAME) + message(FATAL_ERROR "No repository name privated.") + endif() + if(NOT SVN_ARGS_URL) + message(FATAL_ERROR "No repository URL privated.") + endif() + + if(SVN_ARGS_PATH AND NOT IS_ABSOLUTE ${SVN_ARGS_PATH}) + get_filename_component(SVN_ARGS_PATH ${SVN_ARGS_PATH} ABSOLUTE) + endif() + + _fetch_content( + NAME "${SVN_ARGS_NAME}" + PATH "${SVN_ARGS_PATH}" + SVN_URL "${SVN_ARGS_URL}" + SVN_REV "${SVN_ARGS_REV}" + ) +endfunction() + +## Download a http archive +function(http_archive) + cmake_parse_arguments(HTTP_ARGS "" "NAME;PATH;URL;SHA256;SHA1;MD5" "" ${ARGN}) + + if(NOT HTTP_ARGS_NAME) + message(FATAL_ERROR "No archive name privated.") + endif() + if(NOT HTTP_ARGS_URL) + message(FATAL_ERROR "No archive URL privated.") + endif() + + if(HTTP_ARGS_PATH AND NOT IS_ABSOLUTE ${HTTP_ARGS_PATH}) + get_filename_component(HTTP_ARGS_PATH ${HTTP_ARGS_PATH} ABSOLUTE) + endif() + + if(HTTP_ARGS_SHA256) + set(HTTP_URL_HASH "SHA256=${HTTP_ARGS_SHA256}") + elseif(HTTP_ARGS_SHA1) + set(HTTP_URL_HASH "SHA1=${HTTP_ARGS_SHA1}") + elseif(HTTP_ARGS_MD5) + set(HTTP_URL_HASH "MD5=${HTTP_ARGS_MD5}") + else() + set(HTTP_URL_HASH "") + endif() + + _fetch_content( + NAME "${HTTP_ARGS_NAME}" + PATH "${HTTP_ARGS_PATH}" + URL "${HTTP_ARGS_URL}" + URL_HASH "${HTTP_URL_HASH}" + ) +endfunction() + +## Retrieve a version string from GIT +function(git_version _RESULT _SOURCES_DIR) + find_package(Git REQUIRED) + + if(NOT IS_ABSOLUTE ${_SOURCES_DIR}) + get_filename_component(_SOURCES_DIR ${_SOURCES_DIR} ABSOLUTE) + endif() + + # git describe --tags + execute_process( + COMMAND "${GIT_EXECUTABLE}" describe --tags + WORKING_DIRECTORY "${_SOURCES_DIR}" + RESULT_VARIABLE GIT_VER_RESULT + OUTPUT_VARIABLE GIT_VER_OUTPUT + ERROR_VARIABLE GIT_VER_ERROR + ) + if(GIT_VER_RESULT EQUAL 0) + string(STRIP ${GIT_VER_OUTPUT} GIT_VER_OUTPUT) + set(${_RESULT} "${GIT_VER_OUTPUT}" PARENT_SCOPE) + return() + endif() + + # git rev-parse --short HEAD + execute_process( + COMMAND "${GIT_EXECUTABLE}" rev-parse --short HEAD + WORKING_DIRECTORY "${_SOURCES_DIR}" + RESULT_VARIABLE GIT_VER_RESULT + OUTPUT_VARIABLE GIT_VER_OUTPUT + ERROR_VARIABLE GIT_VER_ERROR + ) + if(GIT_VER_RESULT EQUAL 0) + string(STRIP ${GIT_VER_OUTPUT} GIT_VER_OUTPUT) + set(${_RESULT} "g${GIT_VER_OUTPUT}" PARENT_SCOPE) + return() + endif() + + set(${_RESULT} "" PARENT_SCOPE) +endfunction() + +## Retrieve a version string from HG +function(hg_version _RESULT _SOURCES_DIR) + find_package(Hg REQUIRED) + + if(NOT IS_ABSOLUTE ${_SOURCES_DIR}) + get_filename_component(_SOURCES_DIR ${_SOURCES_DIR} ABSOLUTE) + endif() + + # hg log -T "{latesttagdistance}" -r . + execute_process( + COMMAND "${HG_EXECUTABLE}" log -T "{latesttagdistance}" -r . + WORKING_DIRECTORY "${_SOURCES_DIR}" + RESULT_VARIABLE HG_VER_RESULT + OUTPUT_VARIABLE HG_VER_OUTPUT + ERROR_VARIABLE HG_VER_ERROR + ) + if(HG_VER_RESULT EQUAL 0) + string(STRIP ${HG_VER_OUTPUT} HG_VER_OUTPUT) + if(HG_VER_OUTPUT STREQUAL "0") + # hg log -T "{latesttag}" -r . + execute_process( + COMMAND "${HG_EXECUTABLE}" log -T "{latesttag}" -r . + WORKING_DIRECTORY "${_SOURCES_DIR}" + RESULT_VARIABLE HG_VER_RESULT + OUTPUT_VARIABLE HG_VER_OUTPUT + ERROR_VARIABLE HG_VER_ERROR + ) + else() + # hg log -T "{latesttag}-{latesttagdistance}-h{node|short}" -r . + execute_process( + COMMAND "${HG_EXECUTABLE}" log + -T "{latesttag}-{latesttagdistance}-h{node|short}" -r . + WORKING_DIRECTORY "${_SOURCES_DIR}" + RESULT_VARIABLE HG_VER_RESULT + OUTPUT_VARIABLE HG_VER_OUTPUT + ERROR_VARIABLE HG_VER_ERROR + ) + endif() + + if(HG_VER_RESULT EQUAL 0) + string(STRIP ${HG_VER_OUTPUT} HG_VER_OUTPUT) + if(NOT HG_VER_OUTPUT MATCHES "^null.*") + set(${_RESULT} "${HG_VER_OUTPUT}" PARENT_SCOPE) + return() + endif() + endif() + endif() + + # hg log -T "h{node|short}" -r . + execute_process( + COMMAND "${HG_EXECUTABLE}" log -T "h{node|short}" -r . + WORKING_DIRECTORY "${_SOURCES_DIR}" + RESULT_VARIABLE HG_VER_RESULT + OUTPUT_VARIABLE HG_VER_OUTPUT + ERROR_VARIABLE HG_VER_ERROR + ) + if(HG_VER_RESULT EQUAL 0) + string(STRIP ${HG_VER_OUTPUT} HG_VER_OUTPUT) + set(${_RESULT} "${HG_VER_OUTPUT}" PARENT_SCOPE) + return() + endif() + + set(${_RESULT} "" PARENT_SCOPE) +endfunction() + +## Retrieve a version string from SVN +function(svn_version _RESULT _SOURCES_DIR) + find_package(Subversion REQUIRED) + + if(NOT IS_ABSOLUTE ${_SOURCES_DIR}) + get_filename_component(_SOURCES_DIR ${_SOURCES_DIR} ABSOLUTE) + endif() + + # svn info --show-item revision + execute_process( + COMMAND "${Subversion_SVN_EXECUTABLE}" info --show-item revision + WORKING_DIRECTORY "${_SOURCES_DIR}" + RESULT_VARIABLE SVN_VER_RESULT + OUTPUT_VARIABLE SVN_VER_OUTPUT + ERROR_VARIABLE SVN_VER_ERROR + ) + if(SVN_VER_RESULT EQUAL 0) + string(STRIP ${SVN_VER_OUTPUT} SVN_VER_OUTPUT) + set(${_RESULT} "r${SVN_VER_OUTPUT}" PARENT_SCOPE) + return() + endif() + + set(${_RESULT} "" PARENT_SCOPE) +endfunction() + +_find_workspace_directory(BAZEL_WORKSPACE_DIR) +if(BAZEL_WORKSPACE_DIR) + include("${BAZEL_WORKSPACE_DIR}/Workspace.cmake") +endif() diff --git a/cmake/iOSBundleInfo.plist.in b/cmake/iOSBundleInfo.plist.in new file mode 100644 index 0000000..ac6e41e --- /dev/null +++ b/cmake/iOSBundleInfo.plist.in @@ -0,0 +1,22 @@ + + + + + CFBundleDevelopmentRegion + English + CFBundleExecutable + ${MACOSX_BUNDLE_EXECUTABLE_NAME} + CFBundleIdentifier + com.zvec.${MACOSX_BUNDLE_EXECUTABLE_NAME} + CFBundleInfoDictionaryVersion + 6.0 + CFBundleName + ${MACOSX_BUNDLE_EXECUTABLE_NAME} + CFBundlePackageType + APPL + CFBundleShortVersionString + ${MACOSX_BUNDLE_SHORT_VERSION_STRING} + CFBundleVersion + ${MACOSX_BUNDLE_BUNDLE_VERSION} + + diff --git a/cmake/option.cmake b/cmake/option.cmake new file mode 100644 index 0000000..fe10857 --- /dev/null +++ b/cmake/option.cmake @@ -0,0 +1,321 @@ +## https://en.wikipedia.org/wiki/List_of_Intel_CPU_microarchitectures +## https://en.wikipedia.org/wiki/List_of_AMD_CPU_microarchitectures +## https://gcc.gnu.org/onlinedocs/gcc/x86-Options.html + +## Intel Microarchitectures +option(ENABLE_NEHALEM "Enable Intel Nehalem CPU microarchitecture" OFF) +option(ENABLE_SANDYBRIDGE "Enable Intel Sandy Bridge CPU microarchitecture" OFF) +option(ENABLE_HASWELL "Enable Intel Haswell CPU microarchitecture" OFF) +option(ENABLE_BROADWELL "Enable Intel Broadwell CPU microarchitecture" OFF) +option(ENABLE_SKYLAKE "Enable Intel Skylake CPU microarchitecture" OFF) +option(ENABLE_SKYLAKE_AVX512 "Enable Intel Skylake Server CPU microarchitecture" OFF) +option(ENABLE_ICELAKE "Enable Intel Icelake CPU microarchitecture" OFF) +option(ENABLE_SAPPHIRERAPIDS "Enable Intel Sapphire Rapids Server CPU microarchitecture" OFF) +option(ENABLE_EMERALDRAPIDS "Enable Intel Emerald Rapids Server CPU microarchitecture" OFF) +option(ENABLE_GRANITERAPIDS "Enable Intel Granite Rapids Server CPU microarchitecture" OFF) + +option(ENABLE_NATIVE "Enable native CPU microarchitecture" OFF) + +## AMD Microarchitectures +option(ENABLE_ZEN1 "Enable AMD Zen+ Family 17h CPU microarchitecture" OFF) +option(ENABLE_ZEN2 "Enable AMD Zen 2 Family 17h CPU microarchitecture" OFF) +option(ENABLE_ZEN3 "Enable AMD Zen 3 Family 19h CPU microarchitecture" OFF) + +## ARM architectures +option(ENABLE_ARMV8A "Enable ARMv8-a architecture" OFF) +option(ENABLE_ARMV8.1A "Enable ARMv8.1-a architecture" OFF) +option(ENABLE_ARMV8.2A "Enable ARMv8.2-a architecture" OFF) +option(ENABLE_ARMV8.3A "Enable ARMv8.3-a architecture" OFF) +option(ENABLE_ARMV8.4A "Enable ARMv8.4-a architecture" OFF) +option(ENABLE_ARMV8.5A "Enable ARMv8.5-a architecture" OFF) +option(ENABLE_ARMV8.6A "Enable ARMv8.6-a architecture" OFF) + +## OpenMP option +option(ENABLE_OPENMP "Enable OpenMP support" OFF) + +## Treat warnings as errors (for CI) +option(ENABLE_WERROR "Treat compiler warnings as errors in strict targets" OFF) + +set(ARCH_OPTIONS + ENABLE_NEHALEM ENABLE_SANDYBRIDGE ENABLE_HASWELL ENABLE_BROADWELL ENABLE_SKYLAKE + ENABLE_SKYLAKE_AVX512 ENABLE_ICELAKE ENABLE_SAPPHIRERAPIDS ENABLE_EMERALDRAPIDS ENABLE_GRANITERAPIDS + ENABLE_ZEN1 ENABLE_ZEN2 ENABLE_ZEN3 + ENABLE_ARMV8A ENABLE_ARMV8.1A ENABLE_ARMV8.2A ENABLE_ARMV8.3A ENABLE_ARMV8.4A + ENABLE_ARMV8.5A ENABLE_ARMV8.6A + ENABLE_NATIVE +) + +option(AUTO_DETECT_ARCH "Auto detect CPU microarchitecture" ON) +foreach(opt IN LISTS ARCH_OPTIONS) + if(${opt}) + message(WARNING "opt ${opt}") + set(AUTO_DETECT_ARCH OFF) + break() + endif() +endforeach() + +include(CheckCCompilerFlag) + +function(_AppendFlags _RESULT _FLAG) + if(${_RESULT} AND NOT "${${_RESULT}}" MATCHES "${_FLAG}") + set(${_RESULT} "${${_RESULT}} ${_FLAG}" PARENT_SCOPE) + else() + set(${_RESULT} "${_FLAG}" PARENT_SCOPE) + endif() +endfunction() + +macro(add_arch_flag FLAG VAR_NAME OPTION_NAME) + if(MSVC) + if(${ARGC} GREATER 3) + set(_effective_flag "${ARGV3}") + else() + set(_effective_flag "") + endif() + else() + set(_effective_flag "${FLAG}") + endif() + + if(NOT "${_effective_flag}" STREQUAL "") + check_c_compiler_flag("${_effective_flag}" COMPILER_SUPPORT_${VAR_NAME}) + if(COMPILER_SUPPORT_${VAR_NAME}) + _AppendFlags(CMAKE_C_FLAGS "${_effective_flag}") + _AppendFlags(CMAKE_CXX_FLAGS "${_effective_flag}") + set(${VAR_NAME}_ENABLED ON) + else() + if(${OPTION_NAME}) + message(FATAL_ERROR "Compiler does not support required flag: '${_effective_flag}' for ${OPTION_NAME}") + else() + set(${VAR_NAME}_ENABLED OFF) + endif() + endif() + else() + set(${VAR_NAME}_ENABLED ON) + endif() +endmacro() + +function(_setup_armv8_march) + if(MSVC) + return() + endif() + set(_arch "armv8") + check_c_compiler_flag("-march=${_arch}" _COMP_SUPP_${_arch}) + if(_COMP_SUPP_${_arch}) + _AppendFlags(CMAKE_C_FLAGS "-march=${_arch}") + _AppendFlags(CMAKE_CXX_FLAGS "-march=${_arch}") + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}" PARENT_SCOPE) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}" PARENT_SCOPE) + return() + else() + message(WARNING "No ARMv8 march flag supported by compiler.") + endif() +endfunction() + +function(_setup_x86_march) + if(MSVC) + return() + endif() + set(_arch "x86-64") + check_c_compiler_flag("-march=${_arch}" _COMP_SUPP_${_arch}) + if(_COMP_SUPP_${_arch}) + _AppendFlags(CMAKE_C_FLAGS "-march=${_arch}") + _AppendFlags(CMAKE_CXX_FLAGS "-march=${_arch}") + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}" PARENT_SCOPE) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}" PARENT_SCOPE) + return() + else() + message(WARNING "No known x86 march flag supported; falling back to generic.") + endif() +endfunction() + +function(setup_compiler_march_for_x86 VAR_NAME_SSE VAR_NAME_AVX2 VAR_NAME_AVX512 VAR_NAME_AVX512FP16) + if(MSVC) + #sse + set(${VAR_NAME_SSE} "" PARENT_SCOPE) + + #avx2 + check_c_compiler_flag("/arch:AVX2" _COMP_SUPP_MSVC_AVX2) + if(_COMP_SUPP_MSVC_AVX2) + set(${VAR_NAME_AVX2} "/arch:AVX2" PARENT_SCOPE) + else() + set(${VAR_NAME_AVX2} "" PARENT_SCOPE) + endif() + + #avx512 + check_c_compiler_flag("/arch:AVX512" _COMP_SUPP_MSVC_AVX512) + if(_COMP_SUPP_MSVC_AVX512) + set(${VAR_NAME_AVX512} "/arch:AVX512" PARENT_SCOPE) + elseif(_COMP_SUPP_MSVC_AVX2) + set(${VAR_NAME_AVX512} "/arch:AVX2" PARENT_SCOPE) + else() + set(${VAR_NAME_AVX512} "" PARENT_SCOPE) + endif() + + #avx512fp16 + if(_COMP_SUPP_MSVC_AVX512) + set(${VAR_NAME_AVX512FP16} "/arch:AVX512" PARENT_SCOPE) + elseif(_COMP_SUPP_MSVC_AVX2) + set(${VAR_NAME_AVX512FP16} "/arch:AVX2" PARENT_SCOPE) + else() + set(${VAR_NAME_AVX512FP16} "" PARENT_SCOPE) + endif() + + return() + endif() + + #sse + set(${VAR_NAME_SSE} "-march=corei7" PARENT_SCOPE) + + #avx 2 + set(${VAR_NAME_AVX2} "-march=core-avx2" PARENT_SCOPE) + + #avx512 + set(_x86_flags_avx512 "icelake-server" "skylake-avx512" "core-avx2" "x86-64") + foreach(_arch_avx512 IN LISTS _x86_flags_avx512) + check_c_compiler_flag("-march=${_arch_avx512}" _COMP_SUPP_${_arch_avx512}) + if(_COMP_SUPP_${_arch_avx512}) + set(${VAR_NAME_AVX512} "-march=${_arch_avx512}" PARENT_SCOPE) + break() + endif() + endforeach() + + #avx512fp16 + set(_x86_flags_avx512fp16 + "sapphirerapids" "icelake-server" "skylake-avx512" "core-avx2" "x86-64" + ) + foreach(_arch_avx512fp16 IN LISTS _x86_flags_avx512fp16) + check_c_compiler_flag("-march=${_arch_avx512fp16}" _COMP_SUPP_${_arch_avx512fp16}) + if(_COMP_SUPP_${_arch_avx512fp16}) + set(${VAR_NAME_AVX512FP16} "-march=${_arch_avx512fp16}" PARENT_SCOPE) + break() + endif() + endforeach() +endfunction() + +# iOS: Skip -march flags and OpenMP; architecture is controlled by CMAKE_OSX_ARCHITECTURES +if(IOS OR CMAKE_SYSTEM_NAME STREQUAL "iOS") + return() +endif() + +if(NOT AUTO_DETECT_ARCH) + if(ENABLE_NATIVE) + if (NOT MSVC) + add_arch_flag("-march=native" NATIVE ENABLE_NATIVE) + endif () + endif() + + if(ENABLE_ZEN3) + add_arch_flag("-march=znver3" ZNVER3 ENABLE_ZEN3 "/arch:AVX2") + endif() + + if(ENABLE_ZEN2) + add_arch_flag("-march=znver2" ZNVER2 ENABLE_ZEN2 "/arch:AVX2") + endif() + + if(ENABLE_ZEN1) + add_arch_flag("-march=znver1" ZNVER1 ENABLE_ZEN1 "/arch:AVX2") + endif() + + if(ENABLE_GRANITERAPIDS) + add_arch_flag("-march=graniterapids" GRANITERAPIDS ENABLE_GRANITERAPIDS "/arch:AVX512") + endif() + + if(ENABLE_EMERALDRAPIDS) + add_arch_flag("-march=emeraldrapids" EMERALDRAPIDS ENABLE_EMERALDRAPIDS "/arch:AVX512") + endif() + + if(ENABLE_SAPPHIRERAPIDS) + add_arch_flag("-march=sapphirerapids" SAPPHIRERAPIDS ENABLE_SAPPHIRERAPIDS "/arch:AVX512") + endif() + + if(ENABLE_ICELAKE) + add_arch_flag("-march=icelake-server" ICELAKE ENABLE_ICELAKE "/arch:AVX512") + endif() + + if(ENABLE_SKYLAKE_AVX512) + add_arch_flag("-march=skylake-avx512" SKYLAKE_AVX512 ENABLE_SKYLAKE_AVX512 "/arch:AVX512") + endif() + + if(ENABLE_SKYLAKE) + add_arch_flag("-march=skylake" SKYLAKE ENABLE_SKYLAKE "/arch:AVX2") + endif() + + if(ENABLE_BROADWELL) + add_arch_flag("-march=broadwell" BROADWELL ENABLE_BROADWELL "/arch:AVX2") + endif() + + if(ENABLE_HASWELL) + add_arch_flag("-march=haswell" HASWELL ENABLE_HASWELL "/arch:AVX2") + endif() + + if(ENABLE_SANDYBRIDGE) + add_arch_flag("-march=sandybridge" SANDYBRIDGE ENABLE_SANDYBRIDGE "/arch:AVX") + endif() + + if(ENABLE_NEHALEM) + add_arch_flag("-march=nehalem" NEHALEM ENABLE_NEHALEM) + endif() + + # ARM (newest first — allow multiple? usually only one) + # But GCC allows only one -march=, so honor highest enabled + if(ENABLE_ARMV8.6A) + add_arch_flag("-march=armv8.6-a" ARMV86A ENABLE_ARMV8.6A) + endif() + if(ENABLE_ARMV8.5A) + add_arch_flag("-march=armv8.5-a" ARMV85A ENABLE_ARMV8.5A) + endif() + if(ENABLE_ARMV8.4A) + add_arch_flag("-march=armv8.4-a" ARMV84A ENABLE_ARMV8.4A) + endif() + if(ENABLE_ARMV8.3A) + add_arch_flag("-march=armv8.3-a" ARMV83A ENABLE_ARMV8.3A) + endif() + if(ENABLE_ARMV8.2A) + add_arch_flag("-march=armv8.2-a" ARMV82A ENABLE_ARMV8.2A) + endif() + if(ENABLE_ARMV8.1A) + add_arch_flag("-march=armv8.1-a" ARMV81A ENABLE_ARMV8.1A) + endif() + if(ENABLE_ARMV8A) + add_arch_flag("-march=armv8-a" ARMV8A ENABLE_ARMV8A) + endif() + +else() + # AUTO DETECT + # Heuristic: detect host architecture and probe appropriate flags + string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" SYSTEM_PROC_LOWER) + + if(SYSTEM_PROC_LOWER MATCHES "^(x86_64|amd64|x64)$") + set(HOST_ARCH x64) + elseif(SYSTEM_PROC_LOWER MATCHES "^(i[3-6]86|x86)$") + set(HOST_ARCH x86) + elseif(SYSTEM_PROC_LOWER MATCHES "^(aarch64|arm64)$") + set(HOST_ARCH arm64) + elseif(SYSTEM_PROC_LOWER MATCHES "^(arm|armv7|armv7-a|armv7l)$") + set(HOST_ARCH arm) + else() + set(HOST_ARCH unknown) + message(WARNING "unknown host arch: ${CMAKE_SYSTEM_PROCESSOR}") + endif() + # message(STATUS "host arch: ${HOST_ARCH}") + + if (HOST_ARCH MATCHES "^(arm|arm64)$") + _setup_armv8_march() + elseif (HOST_ARCH MATCHES "^(x86|x64)$") + _setup_x86_march() + else () + message(WARNING "unknown host arch - no -march set") + endif () +endif() + +# ----------------------------- +# OpenMP +# ----------------------------- +if(ENABLE_OPENMP) + find_package(OpenMP REQUIRED) + if(OpenMP_C_FLAGS) + _AppendFlags(CMAKE_C_FLAGS "${OpenMP_C_FLAGS}") + endif() + if(OpenMP_CXX_FLAGS) + _AppendFlags(CMAKE_CXX_FLAGS "${OpenMP_CXX_FLAGS}") + endif() +endif() diff --git a/cmake/utils.cmake b/cmake/utils.cmake new file mode 100644 index 0000000..30195c6 --- /dev/null +++ b/cmake/utils.cmake @@ -0,0 +1,47 @@ +# Promote a target's INTERFACE_INCLUDE_DIRECTORIES to be treated as SYSTEM +# includes by consumers, suppressing warnings from third-party headers. +function(mark_target_includes_system) + foreach(_target ${ARGN}) + if(NOT TARGET ${_target}) + continue() + endif() + get_target_property(_aliased ${_target} ALIASED_TARGET) + if(_aliased) + set(_target ${_aliased}) + endif() + get_target_property(_inc ${_target} INTERFACE_INCLUDE_DIRECTORIES) + if(_inc) + set_target_properties(${_target} PROPERTIES + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "${_inc}") + endif() + endforeach() +endfunction() + +function(apply_patch_once patch_name target_dir patch_file) + set(mark_file "${target_dir}/.${patch_name}_patched") + + if(EXISTS "${mark_file}") + #message(STATUS "Patch '${patch_name}' already applied to ${target_dir}, skipping.") + return() + endif() + + if(NOT EXISTS "${patch_file}") + message(FATAL_ERROR "Patch file '${patch_file}' not found!") + endif() + + #message(STATUS "Applying patch '${patch_name}' to ${target_dir} ...") + execute_process( + COMMAND git apply --ignore-space-change --ignore-whitespace "${patch_file}" + WORKING_DIRECTORY "${target_dir}" + RESULT_VARIABLE patch_result + OUTPUT_VARIABLE patch_stdout + ERROR_VARIABLE patch_stderr + ) + + if(NOT patch_result EQUAL 0) + message(FATAL_ERROR "Failed to apply patch '${patch_name}' to ${target_dir}:\n${patch_stderr}") + else() + #message(STATUS "Patch '${patch_name}' applied successfully:\n${patch_stdout}") + file(WRITE "${mark_file}" "patched") + endif() +endfunction() diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt new file mode 100644 index 0000000..d58814d --- /dev/null +++ b/examples/c++/CMakeLists.txt @@ -0,0 +1,128 @@ +cmake_minimum_required(VERSION 3.13) +cmake_policy(SET CMP0077 NEW) +project(zvec-example-c++) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +# Enable compile_commands.json +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + +# --- Paths to Zvec --- +# Allow custom host build directory, default to "build" +if(NOT DEFINED HOST_BUILD_DIR) + set(HOST_BUILD_DIR "build") +endif() + +get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) +set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) +set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) + +# Add include and library search paths +include_directories(${ZVEC_INCLUDE_DIR}) +set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR}) + +# Support multi-config builds (MSVC puts libs in Debug/Release subdirectories) +if(CMAKE_BUILD_TYPE) + set(ZVEC_CONFIG_LIB_DIR ${ZVEC_LIB_DIR}/${CMAKE_BUILD_TYPE}) + if(EXISTS "${ZVEC_CONFIG_LIB_DIR}") + list(APPEND ZVEC_LIB_SEARCH_DIRS ${ZVEC_CONFIG_LIB_DIR}) + endif() +endif() +if(WIN32) + set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$:Debug>") +endif() + +function(zvec_find_shared_library OUT_VAR LIB_NAME) + unset(${OUT_VAR} CACHE) + if(WIN32) + find_library(${OUT_VAR} + NAMES ${LIB_NAME}_shared ${LIB_NAME} + PATHS ${ZVEC_LIB_SEARCH_DIRS} + NO_DEFAULT_PATH + NO_CMAKE_FIND_ROOT_PATH + ) + else() + set(ZVEC_ORIGINAL_LIBRARY_SUFFIXES ${CMAKE_FIND_LIBRARY_SUFFIXES}) + if(APPLE) + set(CMAKE_FIND_LIBRARY_SUFFIXES ".dylib") + else() + set(CMAKE_FIND_LIBRARY_SUFFIXES ".so") + endif() + find_library(${OUT_VAR} + NAMES ${LIB_NAME} + PATHS ${ZVEC_LIB_SEARCH_DIRS} + NO_DEFAULT_PATH + NO_CMAKE_FIND_ROOT_PATH + ) + set(CMAKE_FIND_LIBRARY_SUFFIXES "${ZVEC_ORIGINAL_LIBRARY_SUFFIXES}") + endif() + set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE) +endfunction() + +function(zvec_require_shared_library OUT_VAR LIB_NAME) + zvec_find_shared_library(${OUT_VAR} ${LIB_NAME}) + if(NOT ${OUT_VAR}) + message(FATAL_ERROR + "lib${LIB_NAME} shared library was not found in ${ZVEC_LIB_SEARCH_DIRS}. " + "Build zvec first, or pass -DHOST_BUILD_DIR=.") + endif() + set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE) +endfunction() + +zvec_require_shared_library(ZVEC_SHARED_LIBRARY zvec) +zvec_require_shared_library(ZVEC_AILEGO_SHARED_LIBRARY zvec_ailego) +zvec_require_shared_library(ZVEC_CORE_SHARED_LIBRARY zvec_core) + +# --- Create INTERFACE target for libzvec (all-in-one C++ shared library) --- +# libzvec.so/.dylib/.dll already bundles all zvec internal components +# (zvec, zvec_core, zvec_ailego, zvec_turbo), so no individual dependency +# libraries need to be specified by the consumer. +add_library(zvec-lib INTERFACE) +target_link_libraries(zvec-lib INTERFACE "${ZVEC_SHARED_LIBRARY}") + +# --- Create INTERFACE target for libzvec_ailego (ailego-only all-in-one library) --- +# The ailego example intentionally depends only on libzvec_ailego. +add_library(zvec-ailego-lib INTERFACE) +target_link_libraries(zvec-ailego-lib INTERFACE "${ZVEC_AILEGO_SHARED_LIBRARY}") + +# --- Create INTERFACE target for libzvec_core (core-only all-in-one library) --- +# The core example intentionally depends only on libzvec_core. +add_library(zvec-core-lib INTERFACE) +target_link_libraries(zvec-core-lib INTERFACE "${ZVEC_CORE_SHARED_LIBRARY}") + +# --- Executables --- +set(ZVEC_EXAMPLE_TARGETS) + +add_executable(db-example db/main.cc) +target_link_libraries(db-example PRIVATE zvec-lib) +if(ANDROID) + target_link_libraries(db-example PRIVATE log) +endif() +list(APPEND ZVEC_EXAMPLE_TARGETS db-example) + +add_executable(ailego-example ailego/main.cc) +target_link_libraries(ailego-example PRIVATE zvec-ailego-lib) +list(APPEND ZVEC_EXAMPLE_TARGETS ailego-example) + +add_executable(core-example core/main.cc) +target_link_libraries(core-example PRIVATE zvec-core-lib) +list(APPEND ZVEC_EXAMPLE_TARGETS core-example) + +add_executable(external-vector-example core/external_vector_example.cc) +target_link_libraries(external-vector-example PRIVATE zvec-core-lib) +list(APPEND ZVEC_EXAMPLE_TARGETS external-vector-example) + +# Strip symbols to reduce executable size +if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID) + foreach(ZVEC_EXAMPLE_TARGET ${ZVEC_EXAMPLE_TARGETS}) + add_custom_command(TARGET ${ZVEC_EXAMPLE_TARGET} POST_BUILD + COMMAND ${CMAKE_STRIP} "$" + COMMENT "Stripping symbols from ${ZVEC_EXAMPLE_TARGET}") + endforeach() +endif() + +# Optimize for size +if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID AND ZVEC_EXAMPLE_TARGETS) + set_property(TARGET ${ZVEC_EXAMPLE_TARGETS} PROPERTY COMPILE_FLAGS "-Os") + set_property(TARGET ${ZVEC_EXAMPLE_TARGETS} PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE) +endif() diff --git a/examples/c++/ailego/main.cc b/examples/c++/ailego/main.cc new file mode 100644 index 0000000..da68df3 --- /dev/null +++ b/examples/c++/ailego/main.cc @@ -0,0 +1,11 @@ +#include +#include +#include + +using namespace zvec; + +int main() { + std::string a{"hello world"}; + + std::cout << ailego::StringHelper::StartsWith(a, "hello") << std::endl; +} \ No newline at end of file diff --git a/examples/c++/core/external_vector_example.cc b/examples/c++/core/external_vector_example.cc new file mode 100644 index 0000000..be37474 --- /dev/null +++ b/examples/c++/core/external_vector_example.cc @@ -0,0 +1,219 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/// @file external_vector_example.cc +/// @brief Demonstrates using HNSW index in external-vector mode. +/// +/// In external-vector mode the index does NOT store raw vectors internally. +/// Instead, a user-provided VectorSource is passed on every Add/Search call +/// so the index can fetch vectors on demand. This is useful when vectors are +/// already stored elsewhere (e.g. a columnar store, mmap file, remote storage) +/// and you want to avoid duplicating them inside the index. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::core_interface; + +// --------------------------------------------------------------------------- +// A simple VectorSource backed by an in-memory float matrix. +// In production this could be backed by mmap, a database, or remote storage. +// --------------------------------------------------------------------------- +class InMemoryVectorSource : public zvec::core::VectorSource { + public: + /// @param base Pointer to a contiguous float array of shape [n, dim]. + /// @param dim Dimensionality of each vector. + InMemoryVectorSource(const float *base, uint32_t dim) + : base_(base), dim_(dim) {} + + const void *get_vector(uint32_t node_id) const override { + return base_ + static_cast(node_id) * dim_; + } + + private: + const float *base_; + uint32_t dim_; +}; + +// --------------------------------------------------------------------------- +// Helper: generate random float vectors in [0, 1) +// --------------------------------------------------------------------------- +static std::vector generate_random_vectors(uint32_t count, + uint32_t dim) { + std::vector data(static_cast(count) * dim); + for (auto &v : data) { + v = static_cast(std::rand()) / static_cast(RAND_MAX); + } + return data; +} + +// --------------------------------------------------------------------------- +// Helper: brute-force kNN (L2) for recall verification +// --------------------------------------------------------------------------- +static std::vector brute_force_knn(const float *base, uint32_t n, + uint32_t dim, const float *query, + uint32_t topk) { + std::vector> dists(n); + for (uint32_t i = 0; i < n; ++i) { + float dist = 0.0f; + for (uint32_t d = 0; d < dim; ++d) { + float diff = base[static_cast(i) * dim + d] - query[d]; + dist += diff * diff; + } + dists[i] = {dist, i}; + } + std::partial_sort(dists.begin(), dists.begin() + topk, dists.end()); + std::vector result(topk); + for (uint32_t i = 0; i < topk; ++i) { + result[i] = dists[i].second; + } + return result; +} + +int main() { + constexpr uint32_t kDimension = 32; + constexpr uint32_t kDocCount = 200; + constexpr uint32_t kTopK = 5; + const std::string index_path = "external_vector_example.index"; + + // Clean up any previous run + std::filesystem::remove_all(index_path); + + // ------ Step 1: Generate random vector data (simulating external storage) + std::srand(42); + auto vectors = generate_random_vectors(kDocCount, kDimension); + + // Wrap data in our VectorSource + InMemoryVectorSource source(vectors.data(), kDimension); + + // ------ Step 2: Build HNSW index with external-vector mode enabled + auto param = HNSWIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithUseExternalVector(true) // <-- key setting + .Build(); + + auto index = IndexFactory::CreateAndInitIndex(*param); + if (!index) { + std::cerr << "Failed to create index." << std::endl; + return 1; + } + + int ret = index->Open( + index_path, StorageOptions{StorageOptions::StorageType::kMMAP, true}); + if (ret != 0) { + std::cerr << "Failed to open index." << std::endl; + return 1; + } + + // ------ Step 3: Add vectors using AddWithSource + for (uint32_t i = 0; i < kDocCount; ++i) { + VectorData vd; + vd.vector = + DenseVector{vectors.data() + static_cast(i) * kDimension}; + ret = index->AddWithSource(vd, i, source); + if (ret != 0) { + std::cerr << "Failed to add doc " << i << std::endl; + return 1; + } + } + std::cout << "[OK] Added " << kDocCount << " vectors in external mode." + << std::endl; + + // ------ Step 4: Search using SearchWithSource + auto query_param = + HNSWQueryParamBuilder().with_topk(kTopK).with_ef_search(64).build(); + + // Use the first vector as query + const float *query_vec = vectors.data(); + VectorData query; + query.vector = DenseVector{query_vec}; + + SearchResult result; + ret = index->SearchWithSource(query, query_param, source, &result); + if (ret != 0) { + std::cerr << "Search failed." << std::endl; + return 1; + } + + std::cout << "[OK] Search returned " << result.doc_list_.size() << " results." + << std::endl; + + // The closest vector to vectors[0] should be itself (doc_id=0) + if (!result.doc_list_.empty() && result.doc_list_[0].key() == 0) { + std::cout << "[OK] Nearest neighbor is doc_id=0 (self), score=" + << result.doc_list_[0].score() << std::endl; + } + + // ------ Step 5: Verify recall against brute-force + auto gt = + brute_force_knn(vectors.data(), kDocCount, kDimension, query_vec, kTopK); + uint32_t hits = 0; + for (const auto &doc : result.doc_list_) { + if (std::find(gt.begin(), gt.end(), static_cast(doc.key())) != + gt.end()) { + ++hits; + } + } + float recall = static_cast(hits) / static_cast(kTopK); + std::cout << "[OK] Recall@" << kTopK << " = " << recall * 100.0f << "%" + << std::endl; + + // ------ Step 6: Reopen index and search again (persistence verification) + index->Close(); + std::cout << "[OK] Index closed." << std::endl; + + // Must re-create index instance with same params before reopening + index = IndexFactory::CreateAndInitIndex(*param); + if (!index) { + std::cerr << "Failed to re-create index for reopen." << std::endl; + return 1; + } + + ret = index->Open(index_path, + StorageOptions{StorageOptions::StorageType::kMMAP, false}); + if (ret != 0) { + std::cerr << "Failed to reopen index." << std::endl; + return 1; + } + + SearchResult result2; + ret = index->SearchWithSource(query, query_param, source, &result2); + if (ret != 0) { + std::cerr << "Search after reopen failed." << std::endl; + return 1; + } + + std::cout << "[OK] After reopen: search returned " << result2.doc_list_.size() + << " results, top1 doc_id=" << result2.doc_list_[0].key() + << std::endl; + + // Cleanup + index->Close(); + std::filesystem::remove_all(index_path); + std::cout << "\n=== External Vector Example Complete ===" << std::endl; + return 0; +} diff --git a/examples/c++/core/main.cc b/examples/c++/core/main.cc new file mode 100644 index 0000000..2aaa3f7 --- /dev/null +++ b/examples/c++/core/main.cc @@ -0,0 +1,88 @@ +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::core_interface; + +constexpr uint32_t kDimension = 64; +const std::string index_name{"test.index"}; + +Index::Pointer create_index(const BaseIndexParam::Pointer ¶m, + int doc_num = 10) { + auto index = IndexFactory::CreateAndInitIndex(*param); + if (!index) { + std::cout << "Failed to create index." << std::endl; + return nullptr; + } + + int ret = index->Open( + index_name, StorageOptions{StorageOptions::StorageType::kMMAP, true}); + if (ret != 0) { + std::cout << "Failed to open index." << std::endl; + return nullptr; + } + + for (int i = 0; i < doc_num; ++i) { + std::vector vector(kDimension, i / 10.0f + 0.1f); + VectorData vector_data; + vector_data.vector = DenseVector{vector.data()}; + ret = index->Add(vector_data, i); + if (ret != 0) { + std::cout << "Failed to add to index." << std::endl; + return nullptr; + } + } + + ret = index->Train(); + if (ret != 0) { + std::cout << "Failed to train index." << std::endl; + return nullptr; + } + + return index; +} + +int main() { + std::filesystem::remove(index_name); + + auto param = HNSWIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .Build(); + auto index = create_index(param, 1); + std::cout << "index stats: " << index->GetDocCount() << std::endl; + + // query + auto query_param = HNSWQueryParamBuilder() + .with_topk(10) + .with_fetch_vector(true) + .with_ef_search(20) + .build(); + + SearchResult result; + VectorData query; + std::vector vector(kDimension, 0.1f); + query.vector = DenseVector{vector.data()}; + int ret = index->Search(query, query_param, &result); + if (ret != 0) { + std::cout << "Failed to search index." << std::endl; + return -1; + } + + std::cout << "query results: " << result.doc_list_.size() << std::endl; + if (result.doc_list_.size() == 0) { + std::cout << "No results found." << std::endl; + return -1; + } + + std::cout << "key: " << result.doc_list_[0].key() + << ", score: " << result.doc_list_[0].score() << std::endl; + + return 0; +} \ No newline at end of file diff --git a/examples/c++/db/main.cc b/examples/c++/db/main.cc new file mode 100644 index 0000000..2fbd36d --- /dev/null +++ b/examples/c++/db/main.cc @@ -0,0 +1,264 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec; + +Doc create_doc(const uint64_t doc_id, const CollectionSchema &schema, + std::string pk = "") { + Doc new_doc; + if (pk.empty()) { + pk = "pk_" + std::to_string(doc_id); + } + new_doc.set_pk(pk); + + for (auto &field : schema.fields()) { + switch (field->data_type()) { + case DataType::BINARY: { + std::string binary_str("binary_" + std::to_string(doc_id)); + new_doc.set(field->name(), binary_str); + break; + } + case DataType::BOOL: + new_doc.set(field->name(), doc_id % 10 == 0); + break; + case DataType::INT32: + new_doc.set(field->name(), (int32_t)doc_id); + break; + case DataType::INT64: + new_doc.set(field->name(), (int64_t)doc_id); + break; + case DataType::UINT32: + new_doc.set(field->name(), (uint32_t)doc_id); + break; + case DataType::UINT64: + new_doc.set(field->name(), (uint64_t)doc_id); + break; + case DataType::FLOAT: + new_doc.set(field->name(), (float)doc_id); + break; + case DataType::DOUBLE: + new_doc.set(field->name(), (double)doc_id); + break; + case DataType::STRING: + new_doc.set(field->name(), + "value_" + std::to_string(doc_id)); + break; + case DataType::ARRAY_BINARY: { + std::vector bin_vec; + for (size_t i = 0; i < (doc_id % 10); i++) { + bin_vec.push_back("bin_" + std::to_string(i)); + } + new_doc.set>(field->name(), bin_vec); + break; + } + case DataType::ARRAY_BOOL: + new_doc.set>(field->name(), + std::vector(10, doc_id % 10 == 0)); + break; + case DataType::ARRAY_INT32: + new_doc.set>( + field->name(), std::vector(10, (int32_t)doc_id)); + break; + case DataType::ARRAY_INT64: + new_doc.set>( + field->name(), std::vector(10, (int64_t)doc_id)); + break; + case DataType::ARRAY_UINT32: + new_doc.set>( + field->name(), std::vector(10, (uint32_t)doc_id)); + break; + case DataType::ARRAY_UINT64: + new_doc.set>( + field->name(), std::vector(10, (uint64_t)doc_id)); + break; + case DataType::ARRAY_FLOAT: + new_doc.set>(field->name(), + std::vector(10, (float)doc_id)); + break; + case DataType::ARRAY_DOUBLE: + new_doc.set>( + field->name(), std::vector(10, (double)doc_id)); + break; + case DataType::ARRAY_STRING: + new_doc.set>( + field->name(), + std::vector(10, "value_" + std::to_string(doc_id))); + break; + case DataType::VECTOR_BINARY32: + new_doc.set>( + field->name(), + std::vector(field->dimension(), uint32_t(doc_id + 0.1))); + break; + case DataType::VECTOR_BINARY64: + new_doc.set>( + field->name(), + std::vector(field->dimension(), uint64_t(doc_id + 0.1))); + break; + case DataType::VECTOR_FP32: + new_doc.set>( + field->name(), + std::vector(field->dimension(), float(doc_id + 0.1))); + break; + case DataType::VECTOR_FP64: + new_doc.set>( + field->name(), + std::vector(field->dimension(), double(doc_id + 0.1))); + break; + case DataType::VECTOR_FP16: + new_doc.set>( + field->name(), std::vector( + field->dimension(), static_cast( + float(doc_id + 0.1)))); + break; + case DataType::VECTOR_INT8: + new_doc.set>( + field->name(), + std::vector(field->dimension(), (int8_t)doc_id)); + break; + case DataType::VECTOR_INT16: + new_doc.set>( + field->name(), + std::vector(field->dimension(), (int16_t)doc_id)); + break; + case DataType::SPARSE_VECTOR_FP16: { + std::vector indices; + std::vector values; + for (uint32_t i = 0; i < 100; i++) { + indices.push_back(i); + values.push_back(zvec::float16_t(float(doc_id + 0.1))); + } + std::pair, std::vector> + sparse_float_vec; + sparse_float_vec.first = indices; + sparse_float_vec.second = values; + new_doc.set< + std::pair, std::vector>>( + field->name(), sparse_float_vec); + break; + } + case DataType::SPARSE_VECTOR_FP32: { + std::vector indices; + std::vector values; + for (uint32_t i = 0; i < 100; i++) { + indices.push_back(i); + values.push_back(float(doc_id + 0.1)); + } + std::pair, std::vector> sparse_float_vec; + sparse_float_vec.first = indices; + sparse_float_vec.second = values; + new_doc.set, std::vector>>( + field->name(), sparse_float_vec); + break; + } + default: + std::cout << "Unsupported data type: " << field->name() << std::endl; + throw std::runtime_error("Unsupported vector data type"); + } + } + + return new_doc; +} + +CollectionSchema::Ptr create_schema() { + auto schema = std::make_shared("demo"); + schema->set_max_doc_count_per_segment(1000); + + schema->add_field(std::make_shared( + "id", DataType::INT64, false, std::make_shared(true))); + schema->add_field(std::make_shared( + "name", DataType::STRING, false, + std::make_shared(false))); + schema->add_field( + std::make_shared("weight", DataType::FLOAT, true)); + + schema->add_field(std::make_shared( + "dense", DataType::VECTOR_FP32, 128, false, + std::make_shared(MetricType::IP))); + schema->add_field(std::make_shared( + "sparse", DataType::SPARSE_VECTOR_FP32, 0, false, + std::make_shared(MetricType::IP))); + + return schema; +} + +int main() { + std::string path = "./demo"; + std::filesystem::remove_all(path); + + auto schema = create_schema(); + CollectionOptions options{false, true}; + + auto result = Collection::CreateAndOpen(path, *schema, options); + if (!result.has_value()) { + std::cout << result.error().message() << std::endl; + return -1; + } + + std::cout << "init stats: " << result.value()->Stats().value().to_string() + << std::endl; + + auto coll = std::move(result).value(); + + // insert docs + { + auto doc1 = create_doc(0, *schema); + std::vector docs{doc1}; + auto res = coll->Insert(docs); + if (!res.has_value()) { + std::cout << res.error().message() << std::endl; + return -1; + } + std::cout << "after insert stats " << coll->Stats().value().to_string() + << std::endl; + } + + // optimize + { + auto res = coll->Optimize(); + if (!res.ok()) { + std::cout << res.message() << std::endl; + return -1; + } + std::cout << "after optimize stats " << coll->Stats().value().to_string() + << std::endl; + } + + // query + { + SearchQuery query; + query.topk_ = 10; + query.target_.field_name_ = "dense"; + query.include_vector_ = true; + std::vector query_vector = std::vector(128, 0.1); + query.target_.set_vector(std::string((char *)query_vector.data(), + query_vector.size() * sizeof(float))); + auto res = coll->Query(query); + if (!res.has_value()) { + std::cout << res.error().message() << std::endl; + return -1; + } + std::cout << "query result: doc_count[" << res.value().size() << "]" + << std::endl; + std::cout << "first doc: " << res.value()[0]->to_detail_string() + << std::endl; + } + + // close and reopen + coll.reset(); + options.read_only_ = true; + result = Collection::Open(path, options); + if (!result.has_value()) { + std::cout << result.error().message() << std::endl; + return -1; + } + std::cout << "reopen stats: " << result.value()->Stats().value().to_string() + << std::endl; + + return 0; +} \ No newline at end of file diff --git a/examples/c/CMakeLists.txt b/examples/c/CMakeLists.txt new file mode 100644 index 0000000..736200c --- /dev/null +++ b/examples/c/CMakeLists.txt @@ -0,0 +1,137 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +cmake_minimum_required(VERSION 3.13) +cmake_policy(SET CMP0077 NEW) +project(zvec-example-c) +set(CMAKE_C_STANDARD 99) +set(CMAKE_C_STANDARD_REQUIRED ON) + +# Enable compile_commands.json +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + +# --- Paths to Zvec and dependencies --- +# Allow custom host build directory, default to "build" +if(NOT DEFINED HOST_BUILD_DIR) + set(HOST_BUILD_DIR "build") +endif() + +set(ZVEC_INCLUDE_DIR ${CMAKE_BINARY_DIR}/../../../src/include) +set(ZVEC_GENERATED_INCLUDE_DIR ${CMAKE_BINARY_DIR}/../../../${HOST_BUILD_DIR}/src/generated) +set(ZVEC_LIB_DIR ${CMAKE_BINARY_DIR}/../../../${HOST_BUILD_DIR}/lib) + +# Add include and library search paths +include_directories(${ZVEC_INCLUDE_DIR} ${ZVEC_GENERATED_INCLUDE_DIR}) +link_directories(${ZVEC_LIB_DIR}) + +# Support multi-config builds (MSVC puts libs in Debug/Release subdirectories) +if(CMAKE_BUILD_TYPE) + link_directories(${ZVEC_LIB_DIR}/${CMAKE_BUILD_TYPE}) +endif() + +# Find required packages +find_package(Threads REQUIRED) + +# Create INTERFACE target for zvec_c_api (fat shared library) +# No whole-archive flags needed — all symbols are already resolved in the .so/.dylib +add_library(zvec-c-api INTERFACE) +target_link_libraries(zvec-c-api INTERFACE + zvec_c_api + Threads::Threads +) +if(NOT WIN32) + target_link_libraries(zvec-c-api INTERFACE ${CMAKE_DL_LIBS}) +endif() +if(APPLE) + target_link_options(zvec-c-api INTERFACE -Wl,-rpath,${ZVEC_LIB_DIR}) +endif() + +# Basic example +add_executable(c_api_basic_example basic_example.c) +target_link_libraries(c_api_basic_example PRIVATE + zvec-c-api +) + +# Schema example +add_executable(c_api_collection_schema_example collection_schema_example.c) +target_link_libraries(c_api_collection_schema_example PRIVATE + zvec-c-api +) + +# Struct document example +add_executable(c_api_doc_example doc_example.c) +target_link_libraries(c_api_doc_example PRIVATE + zvec-c-api +) + +# Index example +add_executable(c_api_index_example index_example.c) +target_link_libraries(c_api_index_example PRIVATE + zvec-c-api +) + +# Field schema example +add_executable(c_api_field_schema_example field_schema_example.c) +target_link_libraries(c_api_field_schema_example PRIVATE + zvec-c-api +) + +# Optimized example +add_executable(c_api_optimized_example optimized_example.c) +target_link_libraries(c_api_optimized_example PRIVATE + zvec-c-api +) + +# DiskANN example +add_executable(c_api_diskann_example diskann_example.c) +target_link_libraries(c_api_diskann_example PRIVATE + zvec-c-api +) + +# Strip symbols to reduce executable size +if(CMAKE_BUILD_TYPE STREQUAL "Release" AND (ANDROID OR (CMAKE_SYSTEM_NAME STREQUAL "Linux"))) + add_custom_command(TARGET c_api_basic_example POST_BUILD + COMMAND ${CMAKE_STRIP} "$" + COMMENT "Stripping symbols from c_api_basic_example") + add_custom_command(TARGET c_api_collection_schema_example POST_BUILD + COMMAND ${CMAKE_STRIP} "$" + COMMENT "Stripping symbols from c_api_collection_schema_example") + add_custom_command(TARGET c_api_doc_example POST_BUILD + COMMAND ${CMAKE_STRIP} "$" + COMMENT "Stripping symbols from c_api_doc_example") + add_custom_command(TARGET c_api_index_example POST_BUILD + COMMAND ${CMAKE_STRIP} "$" + COMMENT "Stripping symbols from c_api_index_example") + add_custom_command(TARGET c_api_field_schema_example POST_BUILD + COMMAND ${CMAKE_STRIP} "$" + COMMENT "Stripping symbols from c_api_field_schema_example") + add_custom_command(TARGET c_api_optimized_example POST_BUILD + COMMAND ${CMAKE_STRIP} "$" + COMMENT "Stripping symbols from c_api_optimized_example") + add_custom_command(TARGET c_api_diskann_example POST_BUILD + COMMAND ${CMAKE_STRIP} "$" + COMMENT "Stripping symbols from c_api_diskann_example") +endif() + +# Optimize for size +if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID) + set_property(TARGET c_api_basic_example c_api_collection_schema_example c_api_doc_example + c_api_index_example c_api_field_schema_example c_api_optimized_example + c_api_diskann_example + PROPERTY COMPILE_FLAGS "-Os") + set_property(TARGET c_api_basic_example c_api_collection_schema_example c_api_doc_example + c_api_index_example c_api_field_schema_example c_api_optimized_example + c_api_diskann_example + PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE) +endif() diff --git a/examples/c/basic_example.c b/examples/c/basic_example.c new file mode 100644 index 0000000..efe1803 --- /dev/null +++ b/examples/c/basic_example.c @@ -0,0 +1,267 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "zvec/c_api.h" + +/** + * @brief Print error message and return error code + */ +static zvec_error_code_t handle_error(zvec_error_code_t error, + const char *context) { + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + fprintf(stderr, "Error in %s: %d - %s\n", context, error, + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + } + return error; +} + +/** + * @brief Create a simple test collection using CollectionSchema + */ +static zvec_error_code_t create_simple_test_collection( + zvec_collection_t **collection) { + // Create collection schema using C API + zvec_collection_schema_t *schema = + zvec_collection_schema_create("test_collection"); + if (!schema) { + return ZVEC_ERROR_INTERNAL_ERROR; + } + + zvec_error_code_t error = ZVEC_OK; + + // Create index parameters using new API + zvec_index_params_t *invert_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT); + if (!invert_params) { + zvec_collection_schema_destroy(schema); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + zvec_index_params_set_invert_params(invert_params, true, false); + + zvec_index_params_t *hnsw_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW); + if (!hnsw_params) { + zvec_index_params_destroy(invert_params); + zvec_collection_schema_destroy(schema); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_COSINE); + zvec_index_params_set_hnsw_params(hnsw_params, 16, 200); + + // Create and add ID field (primary key) + zvec_field_schema_t *id_field = + zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0); + zvec_field_schema_set_index_params(id_field, invert_params); + error = zvec_collection_schema_add_field(schema, id_field); + if (error != ZVEC_OK) { + zvec_index_params_destroy(invert_params); + zvec_index_params_destroy(hnsw_params); + zvec_collection_schema_destroy(schema); + return error; + } + + // Create text field (inverted index) + zvec_field_schema_t *text_field = + zvec_field_schema_create("text", ZVEC_DATA_TYPE_STRING, true, 0); + zvec_field_schema_set_index_params(text_field, invert_params); + error = zvec_collection_schema_add_field(schema, text_field); + if (error != ZVEC_OK) { + zvec_index_params_destroy(invert_params); + zvec_index_params_destroy(hnsw_params); + zvec_collection_schema_destroy(schema); + return error; + } + + // Create embedding field (HNSW index) + zvec_field_schema_t *embedding_field = zvec_field_schema_create( + "embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, 3); + zvec_field_schema_set_index_params(embedding_field, hnsw_params); + error = zvec_collection_schema_add_field(schema, embedding_field); + if (error != ZVEC_OK) { + zvec_index_params_destroy(invert_params); + zvec_index_params_destroy(hnsw_params); + zvec_collection_schema_destroy(schema); + return error; + } + + // Cleanup index parameters (they have been copied to the field schemas) + zvec_index_params_destroy(invert_params); + zvec_index_params_destroy(hnsw_params); + + // Use default options + zvec_collection_options_t *options = zvec_collection_options_create(); + if (!options) { + zvec_collection_schema_destroy(schema); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + + // Create collection using the new API + error = zvec_collection_create_and_open("./test_collection", schema, options, + collection); + + // Cleanup resources + zvec_collection_options_destroy(options); + zvec_collection_schema_destroy(schema); + + return error; +} + +/** + * @brief Basic C API usage example + */ +int main() { + printf("=== ZVec C API Basic Example ===\n\n"); + + zvec_error_code_t error; + + // Create collection using simplified function + zvec_collection_t *collection = NULL; + error = create_simple_test_collection(&collection); + if (handle_error(error, "creating collection") != ZVEC_OK) { + return 1; + } + printf("✓ Collection created successfully\n"); + + // Prepare test data + float vector1[] = {0.1f, 0.2f, 0.3f}; + float vector2[] = {0.4f, 0.5f, 0.6f}; + + zvec_doc_t *docs[2]; + for (int i = 0; i < 2; ++i) { + docs[i] = zvec_doc_create(); + if (!docs[i]) { + fprintf(stderr, "Failed to create document %d\n", i); + // Cleanup allocated resources + for (int j = 0; j < i; ++j) { + zvec_doc_destroy(docs[j]); + } + return ZVEC_ERROR_INTERNAL_ERROR; + } + } + + // Manually add fields to document 1 + zvec_doc_set_pk(docs[0], "doc1"); + zvec_doc_add_field_by_value(docs[0], "id", ZVEC_DATA_TYPE_STRING, "doc1", + strlen("doc1")); + zvec_doc_add_field_by_value(docs[0], "text", ZVEC_DATA_TYPE_STRING, + "First document", strlen("First document")); + zvec_doc_add_field_by_value(docs[0], "embedding", ZVEC_DATA_TYPE_VECTOR_FP32, + vector1, 3 * sizeof(float)); + + // Manually add fields to document 2 + zvec_doc_set_pk(docs[1], "doc2"); + zvec_doc_add_field_by_value(docs[1], "id", ZVEC_DATA_TYPE_STRING, "doc2", + strlen("doc2")); + zvec_doc_add_field_by_value(docs[1], "text", ZVEC_DATA_TYPE_STRING, + "Second document", strlen("Second document")); + zvec_doc_add_field_by_value(docs[1], "embedding", ZVEC_DATA_TYPE_VECTOR_FP32, + vector2, 3 * sizeof(float)); + + // Insert documents + size_t success_count = 0; + size_t error_count = 0; + error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, 2, + &success_count, &error_count); + if (handle_error(error, "inserting documents") != ZVEC_OK) { + zvec_collection_destroy(collection); + return 1; + } + printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count, + error_count); + for (int i = 0; i < 2; ++i) { + zvec_doc_destroy(docs[i]); + } + + // Flush collection + error = zvec_collection_flush(collection); + if (handle_error(error, "flushing collection") != ZVEC_OK) { + printf("Collection flush failed\n"); + } else { + printf("✓ Collection flushed successfully\n"); + } + + // Get collection statistics + zvec_collection_stats_t *stats = NULL; + error = zvec_collection_get_stats(collection, &stats); + if (handle_error(error, "getting collection stats") == ZVEC_OK) { + printf("✓ Collection stats - Document count: %llu\n", + (unsigned long long)zvec_collection_stats_get_doc_count(stats)); + // Free statistics memory + zvec_collection_stats_destroy(stats); + } + + printf("Testing vector query...\n"); + // Query documents + zvec_vector_query_t *query = zvec_vector_query_create(); + if (!query) { + fprintf(stderr, "Failed to create vector query\n"); + zvec_collection_destroy(collection); + return 1; + } + + zvec_vector_query_set_field_name(query, "embedding"); + zvec_vector_query_set_query_vector(query, vector1, 3 * sizeof(float)); + zvec_vector_query_set_topk(query, 10); + zvec_vector_query_set_filter(query, ""); + zvec_vector_query_set_include_vector(query, true); + zvec_vector_query_set_include_doc_id(query, true); + + zvec_doc_t **results = NULL; + size_t result_count = 0; + error = zvec_collection_query(collection, (const zvec_vector_query_t *)query, + &results, &result_count); + + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + printf("[ERROR] Query failed: %s\n", + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + zvec_vector_query_destroy(query); + goto cleanup; + } + + zvec_vector_query_destroy(query); + + printf("✓ Query successful - Returned %zu results\n", result_count); + + // Process query results + for (size_t i = 0; i < result_count && i < 5; ++i) { + const zvec_doc_t *doc = results[i]; + const char *pk = zvec_doc_get_pk_copy(doc); + + printf(" Result %zu: PK=%s, DocID=%llu, Score=%.4f\n", i + 1, + pk ? pk : "NULL", (unsigned long long)zvec_doc_get_doc_id(doc), + zvec_doc_get_score(doc)); + + if (pk) { + zvec_free((void *)pk); + } + } + + // Free query results memory + zvec_docs_free(results, result_count); + +cleanup: + // Cleanup resources + zvec_collection_destroy(collection); + printf("✓ Example completed\n"); + return 0; +} \ No newline at end of file diff --git a/examples/c/collection_schema_example.c b/examples/c/collection_schema_example.c new file mode 100644 index 0000000..a2299b7 --- /dev/null +++ b/examples/c/collection_schema_example.c @@ -0,0 +1,255 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "zvec/c_api.h" + +/** + * @brief Print error message and return error code + */ +static zvec_error_code_t handle_error(zvec_error_code_t error, + const char *context) { + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + fprintf(stderr, "Error in %s: %d - %s\n", context, error, + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + } + return error; +} + +/** + * @brief Collection schema creation and management example + */ +int main() { + printf("=== ZVec Collection Schema Example ===\n\n"); + + zvec_error_code_t error; + + // 1. Create collection schema + zvec_collection_schema_t *schema = + zvec_collection_schema_create("schema_example_collection"); + if (!schema) { + fprintf(stderr, "Failed to create collection schema\n"); + return 1; + } + printf("✓ Collection schema created successfully\n"); + + // 2. Set schema properties + zvec_collection_schema_set_max_doc_count_per_segment(schema, 1000000); + printf("✓ Set max documents per segment: %llu\n", + (unsigned long long) + zvec_collection_schema_get_max_doc_count_per_segment(schema)); + + // 3. Create index parameters + zvec_index_params_t *invert_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT); + if (!invert_params) { + fprintf(stderr, "Failed to create invert index parameters\n"); + zvec_collection_schema_destroy(schema); + return 1; + } + zvec_index_params_set_invert_params(invert_params, true, false); + + zvec_index_params_t *hnsw_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW); + if (!hnsw_params) { + fprintf(stderr, "Failed to create HNSW index parameters\n"); + zvec_index_params_destroy(invert_params); + zvec_collection_schema_destroy(schema); + return 1; + } + zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_L2); + zvec_index_params_set_hnsw_params(hnsw_params, 16, 200); + + // 4. Create and add ID field (primary key) + zvec_field_schema_t *id_field = + zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0); + if (!id_field) { + fprintf(stderr, "Failed to create ID field\n"); + zvec_collection_schema_destroy(schema); + return 1; + } + + error = zvec_collection_schema_add_field(schema, id_field); + if (handle_error(error, "adding ID field") != ZVEC_OK) { + zvec_collection_schema_destroy(schema); + return 1; + } + printf("✓ ID field added successfully\n"); + + // 5. Create and add text field with inverted index + zvec_field_schema_t *text_field = + zvec_field_schema_create("content", ZVEC_DATA_TYPE_STRING, true, 0); + if (!text_field) { + fprintf(stderr, "Failed to create text field\n"); + zvec_collection_schema_destroy(schema); + return 1; + } + + zvec_field_schema_set_index_params(text_field, invert_params); + error = zvec_collection_schema_add_field(schema, text_field); + if (handle_error(error, "adding text field") != ZVEC_OK) { + zvec_collection_schema_destroy(schema); + return 1; + } + printf("✓ Text field with inverted index added successfully\n"); + + // 6. Create and add vector field with HNSW index + zvec_field_schema_t *vector_field = zvec_field_schema_create( + "embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, 128); + if (!vector_field) { + fprintf(stderr, "Failed to create vector field\n"); + zvec_collection_schema_destroy(schema); + return 1; + } + + zvec_field_schema_set_index_params(vector_field, hnsw_params); + error = zvec_collection_schema_add_field(schema, vector_field); + if (handle_error(error, "adding vector field") != ZVEC_OK) { + zvec_collection_schema_destroy(schema); + return 1; + } + printf("✓ Vector field with HNSW index added successfully\n"); + + // 7. Check field count + // Note: This function may not exist in current API, commenting out for now + // size_t field_count = zvec_collection_schema_get_field_count(schema); + // printf("✓ Total field count: %zu\n", field_count); + + // 8. Create collection with schema + zvec_collection_options_t *options = zvec_collection_options_create(); + if (!options) { + fprintf(stderr, "Failed to create collection options\n"); + zvec_collection_schema_destroy(schema); + return 1; + } + zvec_collection_t *collection = NULL; + + error = zvec_collection_create_and_open("./schema_example_collection", schema, + options, &collection); + if (handle_error(error, "creating collection with schema") != ZVEC_OK) { + zvec_collection_options_destroy(options); + zvec_collection_schema_destroy(schema); + return 1; + } + zvec_collection_options_destroy(options); + printf("✓ Collection created successfully with schema\n"); + + // 9. Prepare test data + float vector1[128]; + float vector2[128]; + for (int i = 0; i < 128; i++) { + vector1[i] = (float)(i + 1) / 128.0f; + vector2[i] = (float)(i + 2) / 128.0f; + } + + // 10. Create documents + zvec_doc_t *docs[2]; + for (int i = 0; i < 2; i++) { + docs[i] = zvec_doc_create(); + if (!docs[i]) { + fprintf(stderr, "Failed to create document %d\n", i); + // Cleanup + for (int j = 0; j < i; j++) { + zvec_doc_destroy(docs[j]); + } + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + return 1; + } + } + + // Add fields to document 1 + zvec_doc_set_pk(docs[0], "doc1"); + zvec_doc_add_field_by_value(docs[0], "id", ZVEC_DATA_TYPE_STRING, "doc1", + strlen("doc1")); + zvec_doc_add_field_by_value(docs[0], "content", ZVEC_DATA_TYPE_STRING, + "First test document", + strlen("First test document")); + zvec_doc_add_field_by_value(docs[0], "embedding", ZVEC_DATA_TYPE_VECTOR_FP32, + vector1, 128 * sizeof(float)); + + // Add fields to document 2 + zvec_doc_set_pk(docs[1], "doc2"); + zvec_doc_add_field_by_value(docs[1], "id", ZVEC_DATA_TYPE_STRING, "doc2", + strlen("doc2")); + zvec_doc_add_field_by_value(docs[1], "content", ZVEC_DATA_TYPE_STRING, + "Second test document", + strlen("Second test document")); + zvec_doc_add_field_by_value(docs[1], "embedding", ZVEC_DATA_TYPE_VECTOR_FP32, + vector2, 128 * sizeof(float)); + + // 11. Insert documents + size_t success_count = 0, error_count = 0; + error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, 2, + &success_count, &error_count); + if (handle_error(error, "inserting documents") != ZVEC_OK) { + // Cleanup + for (int i = 0; i < 2; i++) { + zvec_doc_destroy(docs[i]); + } + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + return 1; + } + printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count, + error_count); + + // Cleanup documents + for (int i = 0; i < 2; i++) { + zvec_doc_destroy(docs[i]); + } + + // 12. Flush collection + error = zvec_collection_flush(collection); + if (handle_error(error, "flushing collection") == ZVEC_OK) { + printf("✓ Collection flushed successfully\n"); + } + + // 13. Query test + zvec_vector_query_t *query = zvec_vector_query_create(); + if (!query) { + fprintf(stderr, "Failed to create vector query\n"); + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + return 1; + } + zvec_vector_query_set_field_name(query, "embedding"); + zvec_vector_query_set_query_vector(query, vector1, 128 * sizeof(float)); + zvec_vector_query_set_topk(query, 5); + zvec_vector_query_set_filter(query, ""); + zvec_vector_query_set_include_vector(query, true); + zvec_vector_query_set_include_doc_id(query, true); + + zvec_doc_t **results = NULL; + size_t result_count = 0; + error = zvec_collection_query(collection, (const zvec_vector_query_t *)query, + &results, &result_count); + if (error == ZVEC_OK) { + printf("✓ Vector query successful - Returned %zu results\n", result_count); + zvec_docs_free(results, result_count); + } + zvec_vector_query_destroy(query); + + // 14. Cleanup resources + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + printf("✓ Schema example completed\n"); + + return 0; +} \ No newline at end of file diff --git a/examples/c/diskann_example.c b/examples/c/diskann_example.c new file mode 100644 index 0000000..5ef8056 --- /dev/null +++ b/examples/c/diskann_example.c @@ -0,0 +1,346 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/** + * @file diskann_example.c + * @brief End-to-end example demonstrating DiskANN index usage via the C API. + * + * DiskANN is a disk-based approximate nearest neighbor search algorithm + * optimized for large-scale datasets that exceed available memory. It uses + * a Vamana graph structure combined with product quantization (PQ) to + * achieve high recall with efficient disk I/O. + * + * NOTE: DiskANN requires Linux x86_64 with libaio. On other platforms the + * example will compile but the runtime plugin will fail to load. + * + * Workflow demonstrated: + * 1. Create collection schema with DiskANN-indexed vector field + * 2. Insert documents with high-dimensional vectors + * 3. Flush collection (triggers PQ training + graph build) + * 4. Search using DiskANN query parameters (list_size controls recall) + * 5. Clean up all resources + */ + +#include +#include +#include +#include "zvec/c_api.h" + +/* -------------------------------------------------------------------------- + * Helpers + * -------------------------------------------------------------------------- */ + +static zvec_error_code_t handle_error(zvec_error_code_t error, + const char *context) { + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + fprintf(stderr, "Error in %s: %d - %s\n", context, error, + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + } + return error; +} + +#define VECTOR_DIM 64 +#define NUM_DOCS 100 +#define COLLECTION_DIR "./diskann_example_collection" + +/* -------------------------------------------------------------------------- + * Main + * -------------------------------------------------------------------------- */ + +int main(void) { + printf("=== ZVec DiskANN Index Example ===\n\n"); + + zvec_error_code_t error; + int i; + + /* ------------------------------------------------------------------ + * Step 1: Create collection schema + * ------------------------------------------------------------------ */ + printf("[Step 1] Creating collection schema...\n"); + + zvec_collection_schema_t *schema = + zvec_collection_schema_create("diskann_example"); + if (!schema) { + fprintf(stderr, "Failed to create schema\n"); + return 1; + } + + /* Index params — declared up-front and NULL-initialized so the + * cleanup_schema path never touches an uninitialized pointer even if an + * early field addition fails. */ + zvec_index_params_t *invert_params = NULL; + zvec_index_params_t *diskann_params = NULL; + + /* Scalar field with inverted index (for primary key / filtering) */ + invert_params = zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT); + zvec_index_params_set_invert_params(invert_params, true, false); + + zvec_field_schema_t *id_field = + zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0); + zvec_field_schema_set_index_params(id_field, invert_params); + error = zvec_collection_schema_add_field(schema, id_field); + if (handle_error(error, "adding id field") != ZVEC_OK) { + goto cleanup_schema; + } + printf(" + id field (STRING, inverted index)\n"); + + /* Vector field with DiskANN index */ + diskann_params = zvec_index_params_create(ZVEC_INDEX_TYPE_DISKANN); + if (!diskann_params) { + fprintf(stderr, "Failed to create DiskANN index parameters\n"); + goto cleanup_schema; + } + zvec_index_params_set_metric_type(diskann_params, ZVEC_METRIC_TYPE_L2); + zvec_index_params_set_diskann_params( + diskann_params, 64, /* max_degree: graph connectivity */ + 100, /* list_size: build-time candidates */ + 8); /* pq_chunk_num: PQ chunks (0=auto) */ + + printf( + " DiskANN index params: max_degree=%d, list_size=%d, pq_chunk_num=%d\n", + zvec_index_params_get_diskann_max_degree(diskann_params), + zvec_index_params_get_diskann_list_size(diskann_params), + zvec_index_params_get_diskann_pq_chunk_num(diskann_params)); + + zvec_field_schema_t *embedding_field = zvec_field_schema_create( + "embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, VECTOR_DIM); + zvec_field_schema_set_index_params(embedding_field, diskann_params); + error = zvec_collection_schema_add_field(schema, embedding_field); + if (handle_error(error, "adding embedding field") != ZVEC_OK) { + goto cleanup_schema; + } + printf(" + embedding field (VECTOR_FP32, %dD, DiskANN index)\n", VECTOR_DIM); + + /* Index params are copied into field schemas; safe to destroy now */ + zvec_index_params_destroy(invert_params); + zvec_index_params_destroy(diskann_params); + invert_params = NULL; + diskann_params = NULL; + + /* ------------------------------------------------------------------ + * Step 2: Create and open collection + * ------------------------------------------------------------------ */ + printf("\n[Step 2] Creating collection...\n"); + + zvec_collection_options_t *options = zvec_collection_options_create(); + zvec_collection_t *collection = NULL; + error = zvec_collection_create_and_open(COLLECTION_DIR, schema, options, + &collection); + zvec_collection_options_destroy(options); + if (handle_error(error, "creating collection") != ZVEC_OK) { + goto cleanup_schema; + } + printf(" Collection created at %s\n", COLLECTION_DIR); + + /* ------------------------------------------------------------------ + * Step 3: Generate and insert documents + * ------------------------------------------------------------------ */ + printf("\n[Step 3] Inserting %d documents with %dD vectors...\n", NUM_DOCS, + VECTOR_DIM); + + /* Allocate vector storage */ + float(*vectors)[VECTOR_DIM] = + (float(*)[VECTOR_DIM])malloc(NUM_DOCS * VECTOR_DIM * sizeof(float)); + if (!vectors) { + fprintf(stderr, "Failed to allocate vector storage\n"); + goto cleanup_collection; + } + + /* Generate deterministic vector data */ + for (i = 0; i < NUM_DOCS; i++) { + for (int d = 0; d < VECTOR_DIM; d++) { + vectors[i][d] = (float)((i * VECTOR_DIM + d) % 1000) / 1000.0f; + } + } + + /* Insert in batches */ + int batch_size = 20; + size_t total_success = 0, total_error = 0; + + for (int batch_start = 0; batch_start < NUM_DOCS; batch_start += batch_size) { + int count = batch_start + batch_size > NUM_DOCS ? NUM_DOCS - batch_start + : batch_size; + + zvec_doc_t **docs = + (zvec_doc_t **)malloc((size_t)count * sizeof(zvec_doc_t *)); + for (i = 0; i < count; i++) { + int idx = batch_start + i; + docs[i] = zvec_doc_create(); + + char pk[32]; + snprintf(pk, sizeof(pk), "doc_%04d", idx); + zvec_doc_set_pk(docs[i], pk); + + zvec_doc_add_field_by_value(docs[i], "id", ZVEC_DATA_TYPE_STRING, pk, + strlen(pk)); + zvec_doc_add_field_by_value(docs[i], "embedding", + ZVEC_DATA_TYPE_VECTOR_FP32, vectors[idx], + VECTOR_DIM * sizeof(float)); + } + + size_t success_count = 0, error_count = 0; + error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, + (size_t)count, &success_count, &error_count); + if (error != ZVEC_OK) { + handle_error(error, "inserting batch"); + } + total_success += success_count; + total_error += error_count; + + for (i = 0; i < count; i++) { + zvec_doc_destroy(docs[i]); + } + free(docs); + } + printf(" Inserted: %zu succeeded, %zu failed\n", total_success, total_error); + + /* ------------------------------------------------------------------ + * Step 4: Flush to trigger index build (PQ training + graph construction) + * ------------------------------------------------------------------ */ + printf("\n[Step 4] Flushing collection (triggers DiskANN index build)...\n"); + + error = zvec_collection_flush(collection); + if (handle_error(error, "flushing collection") != ZVEC_OK) { + goto cleanup_vectors; + } + + zvec_collection_stats_t *stats = NULL; + error = zvec_collection_get_stats(collection, &stats); + if (error == ZVEC_OK && stats) { + printf(" Document count after flush: %llu\n", + (unsigned long long)zvec_collection_stats_get_doc_count(stats)); + zvec_collection_stats_destroy(stats); + } + + /* ------------------------------------------------------------------ + * Step 5: Search with DiskANN query parameters + * ------------------------------------------------------------------ */ + printf("\n[Step 5] Searching with DiskANN query parameters...\n"); + + /* Create DiskANN query params — list_size controls the search frontier + * (beam width). Larger values improve recall at the cost of latency. */ + zvec_diskann_query_params_t *da_qp = zvec_query_params_diskann_create(200); + if (!da_qp) { + fprintf(stderr, "Failed to create DiskANN query params\n"); + goto cleanup_vectors; + } + printf(" DiskANN query params: list_size=%d\n", + zvec_query_params_diskann_get_list_size(da_qp)); + + /* Build the vector query */ + zvec_vector_query_t *query = zvec_vector_query_create(); + zvec_vector_query_set_field_name(query, "embedding"); + zvec_vector_query_set_query_vector(query, vectors[0], + VECTOR_DIM * sizeof(float)); + zvec_vector_query_set_topk(query, 10); + zvec_vector_query_set_include_vector(query, false); + zvec_vector_query_set_include_doc_id(query, true); + + /* Attach DiskANN query params (ownership transfers to query) */ + error = zvec_vector_query_set_diskann_params(query, da_qp); + if (handle_error(error, "setting DiskANN query params") != ZVEC_OK) { + zvec_vector_query_destroy(query); + goto cleanup_vectors; + } + /* da_qp is now owned by query — do NOT call diskann_destroy on it */ + + /* Execute the query */ + zvec_doc_t **results = NULL; + size_t result_count = 0; + error = zvec_collection_query(collection, (const zvec_vector_query_t *)query, + &results, &result_count); + if (error != ZVEC_OK) { + handle_error(error, "executing DiskANN query"); + printf( + " (This is expected on non-Linux platforms — DiskANN requires " + "libaio)\n"); + } else { + printf(" Query returned %zu results:\n", result_count); + for (size_t r = 0; r < result_count && r < 5; r++) { + const char *pk = zvec_doc_get_pk_copy(results[r]); + printf(" [%zu] pk=%s doc_id=%llu score=%.6f\n", r + 1, + pk ? pk : "NULL", + (unsigned long long)zvec_doc_get_doc_id(results[r]), + zvec_doc_get_score(results[r])); + if (pk) { + zvec_free((void *)pk); + } + } + if (result_count > 5) { + printf(" ... and %zu more\n", result_count - 5); + } + zvec_docs_free(results, result_count); + } + zvec_vector_query_destroy(query); + + /* ------------------------------------------------------------------ + * Step 6: Demonstrate list_size tuning (higher recall vs. lower latency) + * ------------------------------------------------------------------ */ + printf("\n[Step 6] Tuning list_size for recall/latency trade-off...\n"); + + int list_sizes[] = {50, 100, 300}; + for (int li = 0; li < 3; li++) { + zvec_diskann_query_params_t *tune_qp = + zvec_query_params_diskann_create(list_sizes[li]); + + zvec_vector_query_t *tune_query = zvec_vector_query_create(); + zvec_vector_query_set_field_name(tune_query, "embedding"); + zvec_vector_query_set_query_vector(tune_query, vectors[0], + VECTOR_DIM * sizeof(float)); + zvec_vector_query_set_topk(tune_query, 10); + zvec_vector_query_set_include_doc_id(tune_query, true); + zvec_vector_query_set_diskann_params(tune_query, tune_qp); + + zvec_doc_t **tune_results = NULL; + size_t tune_count = 0; + error = zvec_collection_query(collection, + (const zvec_vector_query_t *)tune_query, + &tune_results, &tune_count); + if (error == ZVEC_OK) { + printf(" list_size=%3d -> %zu results returned\n", list_sizes[li], + tune_count); + zvec_docs_free(tune_results, tune_count); + } else { + printf(" list_size=%3d -> query failed (expected on non-Linux)\n", + list_sizes[li]); + } + zvec_vector_query_destroy(tune_query); + } + + /* ------------------------------------------------------------------ + * Cleanup + * ------------------------------------------------------------------ */ +cleanup_vectors: + free(vectors); + +cleanup_collection: + zvec_collection_destroy(collection); + +cleanup_schema: + zvec_collection_schema_destroy(schema); + if (invert_params) { + zvec_index_params_destroy(invert_params); + } + if (diskann_params) { + zvec_index_params_destroy(diskann_params); + } + + printf("\n DiskANN index type string: %s\n", + zvec_index_type_to_string(ZVEC_INDEX_TYPE_DISKANN)); + printf("=== Example completed ===\n"); + return 0; +} diff --git a/examples/c/doc_example.c b/examples/c/doc_example.c new file mode 100644 index 0000000..3ca8a2a --- /dev/null +++ b/examples/c/doc_example.c @@ -0,0 +1,539 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "zvec/c_api.h" + +/** + * @brief Print error message and return error code + */ +static zvec_error_code_t handle_error(zvec_error_code_t error, + const char *context) { + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + fprintf(stderr, "Error in %s: %d - %s\n", context, error, + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + } + return error; +} + +/** + * @brief Create a test document with all data types + * @param doc_index Document index for generating unique data + * @return zvec_doc_t* Created document pointer + */ +static zvec_doc_t *create_full_type_test_doc(int doc_index) { + zvec_doc_t *doc = zvec_doc_create(); + if (!doc) { + fprintf(stderr, "Failed to create document\n"); + return NULL; + } + + // Set primary key + char pk_buffer[32]; + snprintf(pk_buffer, sizeof(pk_buffer), "doc_%d", doc_index); + zvec_doc_set_pk(doc, pk_buffer); + + // Add Id field with inverted index + char id_buffer[32]; + snprintf(id_buffer, sizeof(id_buffer), "id_%d", doc_index); + zvec_doc_add_field_by_value(doc, "id", ZVEC_DATA_TYPE_STRING, id_buffer, + strlen(id_buffer)); + + // Add scalar fields with different data types + // String field + char string_value[64]; + snprintf(string_value, sizeof(string_value), "test_string_%d", doc_index); + zvec_doc_add_field_by_value(doc, "string_field", ZVEC_DATA_TYPE_STRING, + string_value, strlen(string_value)); + + // Boolean field + bool bool_value = (doc_index % 2 == 0); + zvec_doc_add_field_by_value(doc, "bool_field", ZVEC_DATA_TYPE_BOOL, + &bool_value, sizeof(bool_value)); + + // Integer fields + int32_t int32_value = doc_index * 1000; + zvec_doc_add_field_by_value(doc, "int32_field", ZVEC_DATA_TYPE_INT32, + &int32_value, sizeof(int32_value)); + + int64_t int64_value = (int64_t)doc_index * 1000000LL; + zvec_doc_add_field_by_value(doc, "int64_field", ZVEC_DATA_TYPE_INT64, + &int64_value, sizeof(int64_value)); + + // Floating point fields + float float_value = (float)doc_index * 1.5f; + zvec_doc_add_field_by_value(doc, "float_field", ZVEC_DATA_TYPE_FLOAT, + &float_value, sizeof(float_value)); + + double double_value = (double)doc_index * 2.718281828; + zvec_doc_add_field_by_value(doc, "double_field", ZVEC_DATA_TYPE_DOUBLE, + &double_value, sizeof(double_value)); + + // Vector fields with different dimensions + // FP32 vector (3D) + float fp32_vector[3] = {(float)doc_index, (float)doc_index * 2.0f, + (float)doc_index * 3.0f}; + zvec_doc_add_field_by_value(doc, "vector_fp32", ZVEC_DATA_TYPE_VECTOR_FP32, + fp32_vector, 3 * sizeof(float)); + + // Larger FP32 vector (16D) + float large_vector[16]; + for (int i = 0; i < 16; i++) { + large_vector[i] = (float)(doc_index * 16 + i) / 256.0f; + } + zvec_doc_add_field_by_value(doc, "large_vector", ZVEC_DATA_TYPE_VECTOR_FP32, + large_vector, 16 * sizeof(float)); + + return doc; +} + +/** + * @brief Compare two documents for equality + */ +static bool compare_documents(const zvec_doc_t *doc1, const zvec_doc_t *doc2) { + if (!doc1 || !doc2) return false; + + // Compare primary keys + const char *pk1 = zvec_doc_get_pk_pointer(doc1); + const char *pk2 = zvec_doc_get_pk_pointer(doc2); + + if (!pk1 || !pk2 || strcmp(pk1, pk2) != 0) { + return false; + } + + // TODO: Compare other fields and values + + return true; +} + +/** + * @brief Print document fields and their values + * @param doc The document to print + * @param doc_index Document index for identification + */ +static void print_doc(const zvec_doc_t *doc, int doc_index) { + if (!doc) { + printf("Document %d: NULL document\n", doc_index); + return; + } + + printf("\n=== Document %d ===\n", doc_index); + + // Print primary key + const char *pk = zvec_doc_get_pk_pointer(doc); + printf("Primary Key: %s\n", pk ? pk : "NULL"); + + // Print document ID + uint64_t doc_id = zvec_doc_get_doc_id(doc); + printf("Document ID: %llu\n", (unsigned long long)doc_id); + + // Print score + float score = zvec_doc_get_score(doc); + printf("Score: %.6f\n", score); + + // Print scalar fields + printf("\nScalar Fields:\n"); + + // ID field (using pointer function for strings) + const void *id_value = NULL; + size_t id_size = 0; + zvec_error_code_t error = zvec_doc_get_field_value_pointer( + doc, "id", ZVEC_DATA_TYPE_STRING, &id_value, &id_size); + if (error == ZVEC_OK && id_value) { + printf(" id: %.*s\n", (int)id_size, (const char *)id_value); + } + + // String field (using pointer function for strings) + const void *string_value = NULL; + size_t string_size = 0; + error = zvec_doc_get_field_value_pointer( + doc, "string_field", ZVEC_DATA_TYPE_STRING, &string_value, &string_size); + if (error == ZVEC_OK && string_value) { + printf(" string_field: %.*s\n", (int)string_size, + (const char *)string_value); + } + + // Boolean field + bool bool_value; + error = zvec_doc_get_field_value_basic(doc, "bool_field", ZVEC_DATA_TYPE_BOOL, + &bool_value, sizeof(bool_value)); + if (error == ZVEC_OK) { + printf(" bool_field: %s\n", bool_value ? "true" : "false"); + } + + // Int32 field + int32_t int32_value; + error = + zvec_doc_get_field_value_basic(doc, "int32_field", ZVEC_DATA_TYPE_INT32, + &int32_value, sizeof(int32_value)); + if (error == ZVEC_OK) { + printf(" int32_field: %d\n", int32_value); + } + + // Int64 field + int64_t int64_value; + error = + zvec_doc_get_field_value_basic(doc, "int64_field", ZVEC_DATA_TYPE_INT64, + &int64_value, sizeof(int64_value)); + if (error == ZVEC_OK) { + printf(" int64_field: %lld\n", (long long)int64_value); + } + + // Float field + float float_value; + error = + zvec_doc_get_field_value_basic(doc, "float_field", ZVEC_DATA_TYPE_FLOAT, + &float_value, sizeof(float_value)); + if (error == ZVEC_OK) { + printf(" float_field: %.6f\n", float_value); + } + + // Double field + double double_value; + error = + zvec_doc_get_field_value_basic(doc, "double_field", ZVEC_DATA_TYPE_DOUBLE, + &double_value, sizeof(double_value)); + if (error == ZVEC_OK) { + printf(" double_field: %.6f\n", double_value); + } + + // Print vector fields (using copy function for complex types) + printf("\nVector Fields:\n"); + + // FP32 vector (3D) + void *fp32_vector = NULL; + size_t fp32_size = 0; + error = zvec_doc_get_field_value_copy( + doc, "vector_fp32", ZVEC_DATA_TYPE_VECTOR_FP32, &fp32_vector, &fp32_size); + if (error == ZVEC_OK && fp32_vector) { + const float *vec = (const float *)fp32_vector; + size_t dim = fp32_size / sizeof(float); + printf(" vector_fp32 (%zuD): [", dim); + for (size_t i = 0; i < dim && i < 10; i++) { // Limit to first 10 elements + printf("%.3f", vec[i]); + if (i < dim - 1 && i < 9) printf(", "); + } + if (dim > 10) printf(", ..."); + printf("]\n"); + zvec_free(fp32_vector); // Free the allocated memory + } + + // Large vector (16D) + void *large_vector = NULL; + size_t large_size = 0; + error = zvec_doc_get_field_value_copy(doc, "large_vector", + ZVEC_DATA_TYPE_VECTOR_FP32, + &large_vector, &large_size); + if (error == ZVEC_OK && large_vector) { + const float *vec = (const float *)large_vector; + size_t dim = large_size / sizeof(float); + printf(" large_vector (%zuD): [", dim); + for (size_t i = 0; i < dim && i < 10; i++) { // Limit to first 10 elements + printf("%.3f", vec[i]); + if (i < dim - 1 && i < 9) printf(", "); + } + if (dim > 10) printf(", ..."); + printf("]\n"); + zvec_free(large_vector); // Free the allocated memory + } + + printf("==================\n\n"); +} + +/** + * @brief Document creation, manipulation, and query example + */ +int main() { + printf("=== ZVec Document Example ===\n\n"); + + zvec_error_code_t error; + + // 1. Create collection schema for document testing + zvec_collection_schema_t *schema = + zvec_collection_schema_create("doc_example_collection"); + if (!schema) { + fprintf(stderr, "Failed to create collection schema\n"); + return -1; + } + printf("✓ Collection schema created\n"); + + // 2. Create index parameters + zvec_index_params_t *invert_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT); + if (!invert_params) { + fprintf(stderr, "Failed to create invert index parameters\n"); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_invert_params(invert_params, true, false); + + zvec_index_params_t *hnsw_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW); + if (!hnsw_params) { + fprintf(stderr, "Failed to create HNSW index parameters\n"); + zvec_index_params_destroy(invert_params); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_L2); + zvec_index_params_set_hnsw_params(hnsw_params, 16, 200); + + // 3. Create fields for all data types + printf("Creating fields for all data types...\n"); + + // Id field with inverted index + zvec_field_schema_t *id_field = + zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0); + if (id_field) { + zvec_field_schema_set_index_params(id_field, invert_params); + error = zvec_collection_schema_add_field(schema, id_field); + if (handle_error(error, "adding ID field") == ZVEC_OK) { + printf("✓ ID field with inverted index added\n"); + } + } + + // Scalar fields + zvec_field_schema_t *string_field = + zvec_field_schema_create("string_field", ZVEC_DATA_TYPE_STRING, true, 0); + zvec_field_schema_t *bool_field = + zvec_field_schema_create("bool_field", ZVEC_DATA_TYPE_BOOL, true, 0); + zvec_field_schema_t *int32_field = + zvec_field_schema_create("int32_field", ZVEC_DATA_TYPE_INT32, true, 0); + zvec_field_schema_t *int64_field = + zvec_field_schema_create("int64_field", ZVEC_DATA_TYPE_INT64, true, 0); + zvec_field_schema_t *float_field = + zvec_field_schema_create("float_field", ZVEC_DATA_TYPE_FLOAT, true, 0); + zvec_field_schema_t *double_field = + zvec_field_schema_create("double_field", ZVEC_DATA_TYPE_DOUBLE, true, 0); + + if (string_field) zvec_collection_schema_add_field(schema, string_field); + if (bool_field) zvec_collection_schema_add_field(schema, bool_field); + if (int32_field) zvec_collection_schema_add_field(schema, int32_field); + if (int64_field) zvec_collection_schema_add_field(schema, int64_field); + if (float_field) zvec_collection_schema_add_field(schema, float_field); + if (double_field) zvec_collection_schema_add_field(schema, double_field); + + // Vector fields + zvec_field_schema_t *vector_fp32_field = zvec_field_schema_create( + "vector_fp32", ZVEC_DATA_TYPE_VECTOR_FP32, false, 3); + zvec_field_schema_t *large_vector_field = zvec_field_schema_create( + "large_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 16); + + if (vector_fp32_field) { + zvec_field_schema_set_index_params(vector_fp32_field, hnsw_params); + error = zvec_collection_schema_add_field(schema, vector_fp32_field); + if (handle_error(error, "adding vector FP32 field") == ZVEC_OK) { + printf("✓ Vector FP32 field with HNSW index added\n"); + } + } + + if (large_vector_field) { + zvec_field_schema_set_index_params(large_vector_field, hnsw_params); + error = zvec_collection_schema_add_field(schema, large_vector_field); + if (handle_error(error, "adding large vector field") == ZVEC_OK) { + printf("✓ Large vector field with HNSW index added\n"); + } + } + + // 4. Create collection + zvec_collection_options_t *options = zvec_collection_options_create(); + if (!options) { + fprintf(stderr, "Failed to create collection options\n"); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_collection_t *collection = NULL; + + error = zvec_collection_create_and_open("./doc_example_collection", schema, + options, &collection); + zvec_collection_options_destroy(options); + if (handle_error(error, "creating collection") != ZVEC_OK) { + zvec_collection_schema_destroy(schema); + return -1; + } + printf("✓ Collection created successfully\n"); + + // 5. Create and insert multiple test documents + printf("Creating and inserting test documents...\n"); + +#define DOC_COUNT 5 + // Use dynamic allocation for MSVC compatibility (no VLA support) + zvec_doc_t **test_docs = + (zvec_doc_t **)malloc(DOC_COUNT * sizeof(zvec_doc_t *)); + if (!test_docs) { + fprintf(stderr, "Failed to allocate test documents\n"); + goto cleanup; + } + + for (int i = 0; i < DOC_COUNT; i++) { + test_docs[i] = create_full_type_test_doc(i); + if (!test_docs[i]) { + fprintf(stderr, "Failed to create document %d\n", i); + // Cleanup + for (int j = 0; j < i; j++) { + zvec_doc_destroy(test_docs[j]); + } + goto cleanup; + } + printf("✓ Created document %d with PK: %s\n", i, + zvec_doc_get_pk_pointer(test_docs[i])); + } + + // Print all documents before insertion + printf("\nDocuments before insertion:\n"); + for (int i = 0; i < DOC_COUNT; i++) { + print_doc(test_docs[i], i); + } + + // Insert documents + size_t success_count = 0, error_count = 0; + error = zvec_collection_insert(collection, (const zvec_doc_t **)test_docs, + DOC_COUNT, &success_count, &error_count); + if (handle_error(error, "inserting documents") == ZVEC_OK) { + printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count, + error_count); + } + + // 6. Flush collection + error = zvec_collection_flush(collection); + if (handle_error(error, "flushing collection") != ZVEC_OK) { + printf("Warning: Collection flush failed\n"); + } else { + printf("✓ Collection flushed successfully\n"); + } + + // Use the first document's vector for querying + float query_vector[] = {0.0f, 0.0f, 0.0f}; + zvec_vector_query_t *query = zvec_vector_query_create(); + if (!query) { + fprintf(stderr, "Failed to create vector query\n"); + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_vector_query_set_field_name(query, "vector_fp32"); + zvec_vector_query_set_query_vector(query, query_vector, 3 * sizeof(float)); + zvec_vector_query_set_topk(query, 5); + zvec_vector_query_set_filter(query, ""); + zvec_vector_query_set_include_vector(query, true); + zvec_vector_query_set_include_doc_id(query, true); + + zvec_doc_t **query_results = NULL; + size_t result_count = 0; + + error = zvec_collection_query(collection, (const zvec_vector_query_t *)query, + &query_results, &result_count); + if (handle_error(error, "querying documents") != ZVEC_OK) { + query_results = NULL; + result_count = 0; + } + + printf("Query returned %zu results\n", result_count); + + // Print query results + printf("\nQuery Results:\n"); + for (size_t i = 0; i < result_count; i++) { + print_doc(query_results[i], i); + } + + // Compare query results + for (size_t i = 0; i < result_count && i < DOC_COUNT; i++) { + const char *result_pk = zvec_doc_get_pk_pointer(query_results[i]); + printf("Comparing query result[%zu]: %s\n", i, result_pk); + + // Find matching original document + bool found = false; + for (int j = 0; j < DOC_COUNT; j++) { + const char *original_pk = zvec_doc_get_pk_pointer(test_docs[j]); + if (strcmp(result_pk, original_pk) == 0) { + if (compare_documents(test_docs[j], query_results[i])) { + printf("✓ Query result %s matches original document\n", result_pk); + } else { + printf("✗ Query result %s does not match original document\n", + result_pk); + } + found = true; + break; + } + } + + if (!found) { + printf("⚠ Original document not found for: %s\n", result_pk); + } + } + + // 7. Filter query test + printf("\n=== Filter Query Test ===\n"); + + // Create filtered query + zvec_vector_query_set_filter(query, "string_field = 'string_field_0'"); + + zvec_doc_t **filtered_results = NULL; + size_t filtered_count = 0; + + error = zvec_collection_query(collection, (const zvec_vector_query_t *)query, + &filtered_results, &filtered_count); + if (handle_error(error, "filtered querying") == ZVEC_OK) { + printf("Filtered query returned %zu results\n", filtered_count); + + // Verify filter results + bool filter_correct = true; + for (size_t i = 0; i < filtered_count; i++) { + // Note: Field value access may require different API + // For now, we'll just check that we got results + const char *pk = zvec_doc_get_pk_pointer(filtered_results[i]); + if (strstr(pk, "doc_") == NULL) { + filter_correct = false; + break; + } + } + + if (filter_correct) { + printf("✓ Filter query results are correct\n"); + } else { + printf("✗ Filter query results are incorrect\n"); + } + + if (filtered_results) { + zvec_docs_free(filtered_results, filtered_count); + } + } + + // 8. Cleanup query results + if (query_results) { + zvec_docs_free(query_results, result_count); + } + + // 9. Cleanup documents + for (int i = 0; i < DOC_COUNT; i++) { + zvec_doc_destroy(test_docs[i]); + } + free(test_docs); // Free the array itself + + // 10. Final cleanup +cleanup: + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + + printf("✓ Document example completed\n"); + + return 0; +} \ No newline at end of file diff --git a/examples/c/field_schema_example.c b/examples/c/field_schema_example.c new file mode 100644 index 0000000..6279b42 --- /dev/null +++ b/examples/c/field_schema_example.c @@ -0,0 +1,312 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "zvec/c_api.h" + +/** + * @brief Print error message and return error code + */ +static zvec_error_code_t handle_error(zvec_error_code_t error, + const char *context) { + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + fprintf(stderr, "Error in %s: %d - %s\n", context, error, + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + } + return error; +} + +/** + * @brief Field schema creation and management example + */ +int main() { + printf("=== ZVec Field Schema Example ===\n\n"); + + zvec_error_code_t error; + + // 1. Create collection schema + zvec_collection_schema_t *schema = + zvec_collection_schema_create("field_example_collection"); + if (!schema) { + fprintf(stderr, "Failed to create collection schema\n"); + return -1; + } + printf("✓ Collection schema created successfully\n"); + + // 2. Create different types of index parameters + zvec_index_params_t *invert_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT); + if (!invert_params) { + fprintf(stderr, "Failed to create invert index parameters\n"); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_invert_params(invert_params, true, false); + + zvec_index_params_t *hnsw_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW); + if (!hnsw_params) { + fprintf(stderr, "Failed to create HNSW index parameters\n"); + zvec_index_params_destroy(invert_params); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_COSINE); + zvec_index_params_set_hnsw_params(hnsw_params, 16, 200); + + zvec_index_params_t *flat_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT); + if (!flat_params) { + fprintf(stderr, "Failed to create Flat index parameters\n"); + zvec_index_params_destroy(invert_params); + zvec_index_params_destroy(hnsw_params); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(flat_params, ZVEC_METRIC_TYPE_L2); + + if (!invert_params || !hnsw_params || !flat_params) { + fprintf(stderr, "Failed to create index parameters\n"); + zvec_index_params_destroy(invert_params); + zvec_index_params_destroy(hnsw_params); + zvec_index_params_destroy(flat_params); + zvec_collection_schema_destroy(schema); + return -1; + } + + // 3. Create scalar fields with different data types + printf("Creating scalar fields...\n"); + + // String field with inverted index + zvec_field_schema_t *name_field = + zvec_field_schema_create("name", ZVEC_DATA_TYPE_STRING, false, 0); + if (name_field) { + zvec_field_schema_set_index_params(name_field, invert_params); + error = zvec_collection_schema_add_field(schema, name_field); + if (handle_error(error, "adding name field") == ZVEC_OK) { + printf("✓ String field 'name' with inverted index added\n"); + } + } + + // Integer field + zvec_field_schema_t *age_field = + zvec_field_schema_create("age", ZVEC_DATA_TYPE_INT32, true, 0); + if (age_field) { + error = zvec_collection_schema_add_field(schema, age_field); + if (handle_error(error, "adding age field") == ZVEC_OK) { + printf("✓ Integer field 'age' added\n"); + } + } + + // Float field + zvec_field_schema_t *score_field = + zvec_field_schema_create("score", ZVEC_DATA_TYPE_FLOAT, true, 0); + if (score_field) { + error = zvec_collection_schema_add_field(schema, score_field); + if (handle_error(error, "adding score field") == ZVEC_OK) { + printf("✓ Float field 'score' added\n"); + } + } + + // Boolean field + zvec_field_schema_t *active_field = + zvec_field_schema_create("active", ZVEC_DATA_TYPE_BOOL, false, 0); + if (active_field) { + error = zvec_collection_schema_add_field(schema, active_field); + if (handle_error(error, "adding active field") == ZVEC_OK) { + printf("✓ Boolean field 'active' added\n"); + } + } + + // 4. Create vector fields with different dimensions and indexes + printf("Creating vector fields...\n"); + + // Small dimension vector with HNSW index + zvec_field_schema_t *small_vector_field = zvec_field_schema_create( + "small_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 32); + if (small_vector_field) { + zvec_field_schema_set_index_params(small_vector_field, hnsw_params); + error = zvec_collection_schema_add_field(schema, small_vector_field); + if (handle_error(error, "adding small vector field") == ZVEC_OK) { + printf( + "✓ Small vector field 'small_vector' (32D) with HNSW index added\n"); + } + } + + // Medium dimension vector with Flat index + zvec_field_schema_t *medium_vector_field = zvec_field_schema_create( + "medium_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 128); + if (medium_vector_field) { + zvec_field_schema_set_index_params(medium_vector_field, flat_params); + error = zvec_collection_schema_add_field(schema, medium_vector_field); + if (handle_error(error, "adding medium vector field") == ZVEC_OK) { + printf( + "✓ Medium vector field 'medium_vector' (128D) with Flat index " + "added\n"); + } + } + + // Large dimension vector with HNSW index + zvec_field_schema_t *large_vector_field = zvec_field_schema_create( + "large_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 512); + if (large_vector_field) { + zvec_field_schema_set_index_params(large_vector_field, hnsw_params); + error = zvec_collection_schema_add_field(schema, large_vector_field); + if (handle_error(error, "adding large vector field") == ZVEC_OK) { + printf( + "✓ Large vector field 'large_vector' (512D) with HNSW index added\n"); + } + } + + // 5. Create collection with the schema + zvec_collection_options_t *options = zvec_collection_options_create(); + if (!options) { + fprintf(stderr, "Failed to create collection options\n"); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_collection_t *collection = NULL; + + error = zvec_collection_create_and_open("./field_example_collection", schema, + options, &collection); + zvec_collection_options_destroy(options); + if (handle_error(error, "creating collection") != ZVEC_OK) { + zvec_collection_schema_destroy(schema); + return -1; + } + printf("✓ Collection created successfully\n"); + + // 6. Create test documents with various field types + printf("Creating test documents...\n"); + + zvec_doc_t *doc1 = zvec_doc_create(); + zvec_doc_t *doc2 = zvec_doc_create(); + + if (!doc1 || !doc2) { + fprintf(stderr, "Failed to create documents\n"); + goto cleanup; + } + + // Document 1 + zvec_doc_set_pk(doc1, "user1"); + zvec_doc_add_field_by_value(doc1, "name", ZVEC_DATA_TYPE_STRING, + "Alice Johnson", strlen("Alice Johnson")); + int32_t age1 = 28; + zvec_doc_add_field_by_value(doc1, "age", ZVEC_DATA_TYPE_INT32, &age1, + sizeof(age1)); + float score1 = 87.5f; + zvec_doc_add_field_by_value(doc1, "score", ZVEC_DATA_TYPE_FLOAT, &score1, + sizeof(score1)); + bool active1 = true; + zvec_doc_add_field_by_value(doc1, "active", ZVEC_DATA_TYPE_BOOL, &active1, + sizeof(active1)); + + // Add vector data + float small_vec1[32]; + float medium_vec1[128]; + float large_vec1[512]; + + for (int i = 0; i < 32; i++) small_vec1[i] = (float)i / 32.0f; + for (int i = 0; i < 128; i++) medium_vec1[i] = (float)i / 128.0f; + for (int i = 0; i < 512; i++) large_vec1[i] = (float)i / 512.0f; + + zvec_doc_add_field_by_value(doc1, "small_vector", ZVEC_DATA_TYPE_VECTOR_FP32, + small_vec1, 32 * sizeof(float)); + zvec_doc_add_field_by_value(doc1, "medium_vector", ZVEC_DATA_TYPE_VECTOR_FP32, + medium_vec1, 128 * sizeof(float)); + zvec_doc_add_field_by_value(doc1, "large_vector", ZVEC_DATA_TYPE_VECTOR_FP32, + large_vec1, 512 * sizeof(float)); + + // Document 2 + zvec_doc_set_pk(doc2, "user2"); + zvec_doc_add_field_by_value(doc2, "name", ZVEC_DATA_TYPE_STRING, "Bob Smith", + strlen("Bob Smith")); + int32_t age2 = 35; + zvec_doc_add_field_by_value(doc2, "age", ZVEC_DATA_TYPE_INT32, &age2, + sizeof(age2)); + float score2 = 92.0f; + zvec_doc_add_field_by_value(doc2, "score", ZVEC_DATA_TYPE_FLOAT, &score2, + sizeof(score2)); + bool active2 = false; + zvec_doc_add_field_by_value(doc2, "active", ZVEC_DATA_TYPE_BOOL, &active2, + sizeof(active2)); + + // Add vector data + float small_vec2[32]; + float medium_vec2[128]; + float large_vec2[512]; + + for (int i = 0; i < 32; i++) small_vec2[i] = (float)(32 - i) / 32.0f; + for (int i = 0; i < 128; i++) medium_vec2[i] = (float)(128 - i) / 128.0f; + for (int i = 0; i < 512; i++) large_vec2[i] = (float)(512 - i) / 512.0f; + + zvec_doc_add_field_by_value(doc2, "small_vector", ZVEC_DATA_TYPE_VECTOR_FP32, + small_vec2, 32 * sizeof(float)); + zvec_doc_add_field_by_value(doc2, "medium_vector", ZVEC_DATA_TYPE_VECTOR_FP32, + medium_vec2, 128 * sizeof(float)); + zvec_doc_add_field_by_value(doc2, "large_vector", ZVEC_DATA_TYPE_VECTOR_FP32, + large_vec2, 512 * sizeof(float)); + + // 7. Insert documents + zvec_doc_t *docs[] = {doc1, doc2}; + size_t success_count = 0, error_count = 0; + error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, 2, + &success_count, &error_count); + if (handle_error(error, "inserting documents") == ZVEC_OK) { + printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count, + error_count); + } + + // 8. Flush and test queries + zvec_collection_flush(collection); + printf("✓ Collection flushed\n"); + + // Test vector query on medium vector field + zvec_vector_query_t *query = zvec_vector_query_create(); + if (!query) { + fprintf(stderr, "Failed to create vector query\n"); + goto cleanup; + } + zvec_vector_query_set_field_name(query, "medium_vector"); + zvec_vector_query_set_query_vector(query, medium_vec1, 128 * sizeof(float)); + zvec_vector_query_set_topk(query, 2); + zvec_vector_query_set_filter(query, ""); + zvec_vector_query_set_include_vector(query, false); + zvec_vector_query_set_include_doc_id(query, true); + + zvec_doc_t **results = NULL; + size_t result_count = 0; + error = zvec_collection_query(collection, (const zvec_vector_query_t *)query, + &results, &result_count); + if (error == ZVEC_OK) { + printf("✓ Vector query successful - Found %zu results\n", result_count); + zvec_docs_free(results, result_count); + } + zvec_vector_query_destroy(query); + + // 9. Cleanup +cleanup: + if (doc1) zvec_doc_destroy(doc1); + if (doc2) zvec_doc_destroy(doc2); + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + + printf("✓ Field schema example completed\n"); + return 0; +} \ No newline at end of file diff --git a/examples/c/index_example.c b/examples/c/index_example.c new file mode 100644 index 0000000..5cce395 --- /dev/null +++ b/examples/c/index_example.c @@ -0,0 +1,398 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "zvec/c_api.h" + +/** + * @brief Print error message and return error code + */ +static zvec_error_code_t handle_error(zvec_error_code_t error, + const char *context) { + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + fprintf(stderr, "Error in %s: %d - %s\n", context, error, + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + } + return error; +} + +/** + * @brief Index creation and management example + */ +int main() { + printf("=== ZVec Index Example ===\n\n"); + + zvec_error_code_t error; + + // 1. Create collection schema + zvec_collection_schema_t *schema = + zvec_collection_schema_create("index_example_collection"); + if (!schema) { + fprintf(stderr, "Failed to create collection schema\n"); + return -1; + } + printf("✓ Collection schema created successfully\n"); + + // 2. Create different index parameter configurations + printf("Creating index parameters...\n"); + + // Inverted index parameters + zvec_index_params_t *invert_params_standard = + zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT); + if (!invert_params_standard) { + fprintf(stderr, "Failed to create invert index parameters (standard)\n"); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_invert_params(invert_params_standard, true, false); + + zvec_index_params_t *invert_params_extended = + zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT); + if (!invert_params_extended) { + fprintf(stderr, "Failed to create invert index parameters (extended)\n"); + zvec_index_params_destroy(invert_params_standard); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_invert_params(invert_params_extended, true, true); + + // HNSW index parameters with different configurations + zvec_index_params_t *hnsw_params_fast = + zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW); + if (!hnsw_params_fast) { + fprintf(stderr, "Failed to create HNSW index parameters (fast)\n"); + zvec_index_params_destroy(invert_params_standard); + zvec_index_params_destroy(invert_params_extended); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(hnsw_params_fast, ZVEC_METRIC_TYPE_L2); + zvec_index_params_set_hnsw_params(hnsw_params_fast, 16, 100); + + // Demonstrate INT8 quantization with random rotation preprocessing + // (enable_rotate rotates vectors before INT8 quantization to reduce error) + zvec_index_params_set_quantize_type(hnsw_params_fast, + ZVEC_QUANTIZE_TYPE_INT8); + zvec_index_params_set_quantizer_enable_rotate(hnsw_params_fast, true); + + zvec_index_params_t *hnsw_params_balanced = + zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW); + if (!hnsw_params_balanced) { + fprintf(stderr, "Failed to create HNSW index parameters (balanced)\n"); + zvec_index_params_destroy(invert_params_standard); + zvec_index_params_destroy(invert_params_extended); + zvec_index_params_destroy(hnsw_params_fast); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(hnsw_params_balanced, + ZVEC_METRIC_TYPE_COSINE); + zvec_index_params_set_hnsw_params(hnsw_params_balanced, 32, 200); + + zvec_index_params_t *hnsw_params_accurate = + zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW); + if (!hnsw_params_accurate) { + fprintf(stderr, "Failed to create HNSW index parameters (accurate)\n"); + zvec_index_params_destroy(invert_params_standard); + zvec_index_params_destroy(invert_params_extended); + zvec_index_params_destroy(hnsw_params_fast); + zvec_index_params_destroy(hnsw_params_balanced); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(hnsw_params_accurate, ZVEC_METRIC_TYPE_IP); + zvec_index_params_set_hnsw_params(hnsw_params_accurate, 64, 400); + + // Flat index parameters + zvec_index_params_t *flat_params_l2 = + zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT); + if (!flat_params_l2) { + fprintf(stderr, "Failed to create Flat index parameters (L2)\n"); + zvec_index_params_destroy(invert_params_standard); + zvec_index_params_destroy(invert_params_extended); + zvec_index_params_destroy(hnsw_params_fast); + zvec_index_params_destroy(hnsw_params_balanced); + zvec_index_params_destroy(hnsw_params_accurate); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(flat_params_l2, ZVEC_METRIC_TYPE_L2); + + zvec_index_params_t *flat_params_cosine = + zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT); + if (!flat_params_cosine) { + fprintf(stderr, "Failed to create Flat index parameters (cosine)\n"); + zvec_index_params_destroy(invert_params_standard); + zvec_index_params_destroy(invert_params_extended); + zvec_index_params_destroy(hnsw_params_fast); + zvec_index_params_destroy(hnsw_params_balanced); + zvec_index_params_destroy(hnsw_params_accurate); + zvec_index_params_destroy(flat_params_l2); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(flat_params_cosine, + ZVEC_METRIC_TYPE_COSINE); + + // 3. Create fields with different index types + printf("Creating fields with various index types...\n"); + + // Fields with inverted indexes + zvec_field_schema_t *id_field = + zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0); + if (id_field) { + zvec_field_schema_set_index_params(id_field, invert_params_standard); + error = zvec_collection_schema_add_field(schema, id_field); + if (handle_error(error, "adding ID field") == ZVEC_OK) { + printf("✓ ID field with standard inverted index added\n"); + } + } + + zvec_field_schema_t *category_field = + zvec_field_schema_create("category", ZVEC_DATA_TYPE_STRING, true, 0); + if (category_field) { + zvec_field_schema_set_index_params(category_field, invert_params_extended); + error = zvec_collection_schema_add_field(schema, category_field); + if (handle_error(error, "adding category field") == ZVEC_OK) { + printf("✓ Category field with extended inverted index added\n"); + } + } + + // Vector fields with HNSW indexes (different configurations) + zvec_field_schema_t *fast_search_field = zvec_field_schema_create( + "fast_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 64); + if (fast_search_field) { + zvec_field_schema_set_index_params(fast_search_field, hnsw_params_fast); + error = zvec_collection_schema_add_field(schema, fast_search_field); + if (handle_error(error, "adding fast search field") == ZVEC_OK) { + printf("✓ Fast search vector field (64D) with HNSW index added\n"); + } + } + + zvec_field_schema_t *balanced_field = zvec_field_schema_create( + "balanced_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 128); + if (balanced_field) { + zvec_field_schema_set_index_params(balanced_field, hnsw_params_balanced); + error = zvec_collection_schema_add_field(schema, balanced_field); + if (handle_error(error, "adding balanced field") == ZVEC_OK) { + printf("✓ Balanced vector field (128D) with HNSW index added\n"); + } + } + + zvec_field_schema_t *accurate_field = zvec_field_schema_create( + "accurate_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 256); + if (accurate_field) { + zvec_field_schema_set_index_params(accurate_field, hnsw_params_accurate); + error = zvec_collection_schema_add_field(schema, accurate_field); + if (handle_error(error, "adding accurate field") == ZVEC_OK) { + printf("✓ Accurate vector field (256D) with HNSW index added\n"); + } + } + + // Vector field with Flat index + zvec_field_schema_t *exact_field = zvec_field_schema_create( + "exact_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 32); + if (exact_field) { + zvec_field_schema_set_index_params(exact_field, flat_params_l2); + error = zvec_collection_schema_add_field(schema, exact_field); + if (handle_error(error, "adding exact field") == ZVEC_OK) { + printf("✓ Exact search vector field (32D) with Flat index added\n"); + } + } + + // 4. Create collection + zvec_collection_options_t *options = zvec_collection_options_create(); + if (!options) { + fprintf(stderr, "Failed to create collection options\n"); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_collection_t *collection = NULL; + + error = zvec_collection_create_and_open("./index_example_collection", schema, + options, &collection); + zvec_collection_options_destroy(options); + if (handle_error(error, "creating collection") != ZVEC_OK) { + zvec_collection_schema_destroy(schema); + return -1; + } + printf("✓ Collection created successfully\n"); + + // 5. Create test data + printf("Creating test documents...\n"); + + zvec_doc_t *docs[3]; + for (int i = 0; i < 3; i++) { + docs[i] = zvec_doc_create(); + if (!docs[i]) { + fprintf(stderr, "Failed to create document %d\n", i); + // Cleanup + for (int j = 0; j < i; j++) { + zvec_doc_destroy(docs[j]); + } + goto cleanup; + } + } + + // Prepare vector data + float fast_vec[3][64]; + float balanced_vec[3][128]; + float accurate_vec[3][256]; + float exact_vec[3][32]; + + // Generate different vector patterns for testing + for (int doc_idx = 0; doc_idx < 3; doc_idx++) { + for (int i = 0; i < 64; i++) { + fast_vec[doc_idx][i] = (float)(doc_idx * 64 + i) / (64.0f * 3.0f); + } + for (int i = 0; i < 128; i++) { + balanced_vec[doc_idx][i] = (float)(doc_idx * 128 + i) / (128.0f * 3.0f); + } + for (int i = 0; i < 256; i++) { + accurate_vec[doc_idx][i] = (float)(doc_idx * 256 + i) / (256.0f * 3.0f); + } + for (int i = 0; i < 32; i++) { + exact_vec[doc_idx][i] = (float)(doc_idx * 32 + i) / (32.0f * 3.0f); + } + } + + // Populate documents + for (int i = 0; i < 3; i++) { + char pk[16]; + snprintf(pk, sizeof(pk), "doc%d", i + 1); + zvec_doc_set_pk(docs[i], pk); + + char id_val[16]; + snprintf(id_val, sizeof(id_val), "ID_%d", i + 1); + zvec_doc_add_field_by_value(docs[i], "id", ZVEC_DATA_TYPE_STRING, id_val, + strlen(id_val)); + + char category_val[16]; + snprintf(category_val, sizeof(category_val), "cat_%d", (i % 2) + 1); + zvec_doc_add_field_by_value(docs[i], "category", ZVEC_DATA_TYPE_STRING, + category_val, strlen(category_val)); + + zvec_doc_add_field_by_value(docs[i], "fast_vector", + ZVEC_DATA_TYPE_VECTOR_FP32, fast_vec[i], + 64 * sizeof(float)); + zvec_doc_add_field_by_value(docs[i], "balanced_vector", + ZVEC_DATA_TYPE_VECTOR_FP32, balanced_vec[i], + 128 * sizeof(float)); + zvec_doc_add_field_by_value(docs[i], "accurate_vector", + ZVEC_DATA_TYPE_VECTOR_FP32, accurate_vec[i], + 256 * sizeof(float)); + zvec_doc_add_field_by_value(docs[i], "exact_vector", + ZVEC_DATA_TYPE_VECTOR_FP32, exact_vec[i], + 32 * sizeof(float)); + } + + // 6. Insert documents + size_t success_count = 0, error_count = 0; + error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, 3, + &success_count, &error_count); + if (handle_error(error, "inserting documents") == ZVEC_OK) { + printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count, + error_count); + } + + // Cleanup documents + for (int i = 0; i < 3; i++) { + zvec_doc_destroy(docs[i]); + } + + // 7. Flush collection to build indexes + error = zvec_collection_flush(collection); + if (handle_error(error, "flushing collection") == ZVEC_OK) { + printf("✓ Collection flushed - indexes built\n"); + } + + // 8. Test different query types + printf("Testing various index queries...\n"); + + // Test HNSW query (balanced) + zvec_vector_query_t *hnsw_query = zvec_vector_query_create(); + if (!hnsw_query) { + fprintf(stderr, "Failed to create HNSW query\n"); + goto cleanup; + } + zvec_vector_query_set_field_name(hnsw_query, "balanced_vector"); + zvec_vector_query_set_query_vector(hnsw_query, balanced_vec[0], + 128 * sizeof(float)); + zvec_vector_query_set_topk(hnsw_query, 2); + zvec_vector_query_set_filter(hnsw_query, ""); + zvec_vector_query_set_include_vector(hnsw_query, false); + zvec_vector_query_set_include_doc_id(hnsw_query, true); + + zvec_doc_t **hnsw_results = NULL; + size_t hnsw_result_count = 0; + error = + zvec_collection_query(collection, (const zvec_vector_query_t *)hnsw_query, + &hnsw_results, &hnsw_result_count); + if (error == ZVEC_OK) { + printf("✓ HNSW query successful - Found %zu results\n", hnsw_result_count); + zvec_docs_free(hnsw_results, hnsw_result_count); + } + zvec_vector_query_destroy(hnsw_query); + + // Test Flat query (exact) + zvec_vector_query_t *flat_query = zvec_vector_query_create(); + if (!flat_query) { + fprintf(stderr, "Failed to create Flat query\n"); + goto cleanup; + } + zvec_vector_query_set_field_name(flat_query, "exact_vector"); + zvec_vector_query_set_query_vector(flat_query, exact_vec[0], + 32 * sizeof(float)); + zvec_vector_query_set_topk(flat_query, 2); + zvec_vector_query_set_filter(flat_query, ""); + zvec_vector_query_set_include_vector(flat_query, false); + zvec_vector_query_set_include_doc_id(flat_query, true); + + zvec_doc_t **flat_results = NULL; + size_t flat_result_count = 0; + error = + zvec_collection_query(collection, (const zvec_vector_query_t *)flat_query, + &flat_results, &flat_result_count); + if (error == ZVEC_OK) { + printf("✓ Flat (exact) query successful - Found %zu results\n", + flat_result_count); + zvec_docs_free(flat_results, flat_result_count); + } + zvec_vector_query_destroy(flat_query); + + // 9. Performance comparison information + printf("\nIndex Performance Characteristics:\n"); + printf("- Inverted Index: Fast text search, supports filtering\n"); + printf( + "- HNSW Index: Approximate nearest neighbor search, good balance of " + "speed/accuracy\n"); + printf("- Flat Index: Exact search, slower but 100%% accurate\n"); + printf( + "- Trade-off: Speed vs Accuracy - choose based on your requirements\n"); + + // 10. Cleanup +cleanup: + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + + // Cleanup index parameters + + printf("✓ Index example completed\n"); + return 0; +} \ No newline at end of file diff --git a/examples/c/optimized_example.c b/examples/c/optimized_example.c new file mode 100644 index 0000000..28be5c2 --- /dev/null +++ b/examples/c/optimized_example.c @@ -0,0 +1,310 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include "zvec/c_api.h" + +/** + * @brief Print error message and return error code + */ +static zvec_error_code_t handle_error(zvec_error_code_t error, + const char *context) { + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + fprintf(stderr, "Error in %s: %d - %s\n", context, error, + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + } + return error; +} + +/** + * @brief Create test vector data + */ +static float *create_test_vector(size_t dimension) { + float *vector = malloc(dimension * sizeof(float)); + if (!vector) { + return NULL; + } + + for (size_t i = 0; i < dimension; i++) { + vector[i] = (float)rand() / RAND_MAX; + } + + return vector; +} + +/** + * @brief Optimized C API usage example with performance considerations + */ +int main() { + printf("=== ZVec Optimized C API Example ===\n\n"); + + // Get version information + const char *version = zvec_get_version(); + printf("ZVec Version: %s\n\n", version ? version : "Unknown"); + + zvec_error_code_t error; + + // 1. Create optimized collection schema + zvec_collection_schema_t *schema = + zvec_collection_schema_create("optimized_example_collection"); + if (!schema) { + fprintf(stderr, "Failed to create collection schema\n"); + return -1; + } + printf("✓ Collection schema created\n"); + + // 2. Create optimized index parameters + zvec_index_params_t *hnsw_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW); + if (!hnsw_params) { + fprintf(stderr, "Failed to create HNSW index parameters\n"); + zvec_collection_schema_destroy(schema); + return -1; + } + zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_L2); + zvec_index_params_set_hnsw_params(hnsw_params, 32, 200); + + // 3. Create fields with optimized configuration + zvec_field_schema_t *id_field = + zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0); + zvec_field_schema_t *text_field = + zvec_field_schema_create("text", ZVEC_DATA_TYPE_STRING, true, 0); + zvec_field_schema_t *embedding_field = zvec_field_schema_create( + "embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, 128); + + if (!id_field || !text_field || !embedding_field) { + fprintf(stderr, "Failed to create field schemas\n"); + goto cleanup_params; + } + + // Set indexes + zvec_field_schema_set_index_params(embedding_field, hnsw_params); + + // Add fields to schema + error = zvec_collection_schema_add_field(schema, id_field); + if (handle_error(error, "adding ID field") != ZVEC_OK) goto cleanup_fields; + + error = zvec_collection_schema_add_field(schema, text_field); + if (handle_error(error, "adding text field") != ZVEC_OK) goto cleanup_fields; + + error = zvec_collection_schema_add_field(schema, embedding_field); + if (handle_error(error, "adding embedding field") != ZVEC_OK) + goto cleanup_fields; + + printf("✓ Fields configured with indexes\n"); + + // 4. Create collection with optimized options + zvec_collection_options_t *options = zvec_collection_options_create(); + if (!options) { + fprintf(stderr, "Failed to create collection options\n"); + goto cleanup_fields; + } + zvec_collection_options_set_enable_mmap( + options, true); // Enable memory mapping for better performance + + zvec_collection_t *collection = NULL; + error = zvec_collection_create_and_open("./optimized_example_collection", + schema, options, &collection); + zvec_collection_options_destroy(options); + if (handle_error(error, "creating collection") != ZVEC_OK) { + goto cleanup_fields; + } + printf("✓ Collection created with optimized settings\n"); + + // 5. Bulk insert test data + const size_t DOC_COUNT = 1000; + const size_t BATCH_SIZE = 100; + + printf("Inserting %zu documents in batches of %zu...\n", DOC_COUNT, + BATCH_SIZE); + + clock_t start_time = clock(); + + for (size_t batch_start = 0; batch_start < DOC_COUNT; + batch_start += BATCH_SIZE) { + size_t current_batch_size = (batch_start + BATCH_SIZE > DOC_COUNT) + ? DOC_COUNT - batch_start + : BATCH_SIZE; + + zvec_doc_t **batch_docs = malloc(current_batch_size * sizeof(zvec_doc_t *)); + if (!batch_docs) { + fprintf(stderr, "Failed to allocate batch documents\n"); + break; + } + + // Create batch documents + for (size_t i = 0; i < current_batch_size; i++) { + batch_docs[i] = zvec_doc_create(); + if (!batch_docs[i]) { + fprintf(stderr, "Failed to create document\n"); + // Cleanup previous documents in batch + for (size_t j = 0; j < i; j++) { + zvec_doc_destroy(batch_docs[j]); + } + free(batch_docs); + goto cleanup_collection; + } + + size_t doc_id = batch_start + i; + char pk[32]; + snprintf(pk, sizeof(pk), "doc_%zu", doc_id); + zvec_doc_set_pk(batch_docs[i], pk); + + // Add ID field + char id_str[32]; + snprintf(id_str, sizeof(id_str), "ID_%zu", doc_id); + zvec_doc_add_field_by_value(batch_docs[i], "id", ZVEC_DATA_TYPE_STRING, + id_str, strlen(id_str)); + + // Add text field + char text_str[64]; + snprintf(text_str, sizeof(text_str), + "Document number %zu with sample text", doc_id); + zvec_doc_add_field_by_value(batch_docs[i], "text", ZVEC_DATA_TYPE_STRING, + text_str, strlen(text_str)); + + // Add vector field + float *vector = create_test_vector(128); + if (vector) { + zvec_doc_add_field_by_value(batch_docs[i], "embedding", + ZVEC_DATA_TYPE_VECTOR_FP32, vector, + 128 * sizeof(float)); + free(vector); + } + } + + // Insert batch + size_t success_count, error_count; + error = zvec_collection_insert(collection, (const zvec_doc_t **)batch_docs, + current_batch_size, &success_count, + &error_count); + if (handle_error(error, "inserting batch") != ZVEC_OK) { + // Cleanup batch documents + for (size_t i = 0; i < current_batch_size; i++) { + zvec_doc_destroy(batch_docs[i]); + } + free(batch_docs); + goto cleanup_collection; + } + + printf(" Batch %zu-%zu: %zu successful, %zu failed\n", batch_start, + batch_start + current_batch_size - 1, success_count, error_count); + + // Cleanup batch documents + for (size_t i = 0; i < current_batch_size; i++) { + zvec_doc_destroy(batch_docs[i]); + } + free(batch_docs); + } + + clock_t insert_end_time = clock(); + double insert_time = + ((double)(insert_end_time - start_time)) / CLOCKS_PER_SEC; + printf("✓ Bulk insertion completed in %.3f seconds (%.0f docs/sec)\n", + insert_time, DOC_COUNT / insert_time); + + // 6. Flush and optimize collection + printf("Flushing and optimizing collection...\n"); + zvec_collection_flush(collection); + zvec_collection_optimize(collection); + printf("✓ Collection optimized\n"); + + // 7. Performance query test + printf("Testing query performance...\n"); + + float *query_vector = create_test_vector(128); + if (!query_vector) { + fprintf(stderr, "Failed to create query vector\n"); + goto cleanup_collection; + } + + zvec_vector_query_t *query = zvec_vector_query_create(); + if (!query) { + fprintf(stderr, "Failed to create vector query\n"); + free(query_vector); + goto cleanup_collection; + } + zvec_vector_query_set_field_name(query, "embedding"); + zvec_vector_query_set_query_vector(query, query_vector, 128 * sizeof(float)); + zvec_vector_query_set_topk(query, 10); + zvec_vector_query_set_filter(query, ""); + zvec_vector_query_set_include_vector(query, false); + zvec_vector_query_set_include_doc_id(query, true); + + const int QUERY_COUNT = 100; + start_time = clock(); + + for (int q = 0; q < QUERY_COUNT; q++) { + zvec_doc_t **results = NULL; + size_t result_count = 0; + + error = + zvec_collection_query(collection, (const zvec_vector_query_t *)query, + &results, &result_count); + if (error != ZVEC_OK) { + char *error_msg = NULL; + zvec_get_last_error(&error_msg); + printf("Query %d failed: %s\n", q, + error_msg ? error_msg : "Unknown error"); + zvec_free(error_msg); + continue; + } + + if (results) { + zvec_docs_free(results, result_count); + } + } + + clock_t query_end_time = clock(); + double query_time = ((double)(query_end_time - start_time)) / CLOCKS_PER_SEC; + double avg_query_time = (query_time * 1000) / QUERY_COUNT; + + printf("✓ Performance test completed\n"); + printf(" Average query time: %.2f ms\n", avg_query_time); + printf(" Queries per second: %.0f\n", 1000.0 / avg_query_time); + + free(query_vector); + zvec_vector_query_destroy(query); + + // 8. Memory usage information + zvec_collection_stats_t *stats = NULL; + error = zvec_collection_get_stats(collection, &stats); + if (error == ZVEC_OK && stats) { + printf("Collection Statistics:\n"); + printf(" Document count: %llu\n", + (unsigned long long)zvec_collection_stats_get_doc_count(stats)); + zvec_collection_stats_destroy(stats); + } + + // 9. Cleanup +cleanup_collection: + zvec_collection_destroy(collection); + +cleanup_fields: + // Field schemas are managed by the collection schema, no need to destroy + // individually + +cleanup_params: + zvec_collection_schema_destroy(schema); + + printf("✓ Optimized example completed\n"); + + return 0; +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..dbbe26c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,299 @@ +###################################################################################################### +# Zvec: High-Performance Vector Database with PyBind11 & C++ Backend +###################################################################################################### +[project] +name = "zvec" +dynamic = ["version"] +description = "A high-performance vector database engine with native C++ backend and Python bindings" +readme = "README.md" +license = { text = "Apache-2.0" } +authors = [ + { name = "zvec", email = "zvec@alibaba-inc.com" }, +] +maintainers = [ + { name = "Zvec Core Team", email = "zvec@alibaba-inc.com" }, +] +requires-python = ">=3.9" # NOTE: Only 64-bit Python interpreters are supported. +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: Apache Software License", + "Operating System :: POSIX :: Linux", + "Operating System :: MacOS", + "Programming Language :: C++", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Database", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries :: Python Modules", +] +keywords = [ + "vector-database", "ann", "nearest-neighbor" +] +dependencies = [ + "numpy >=1.23", +] + +[project.urls] +Homepage = "https://github.com/alibaba/zvec" +Repository = "https://github.com/alibaba/zvec" +"Bug Tracker" = "https://github.com/alibaba/zvec/issues" +"Documentation" = "https://zvec.org" + +[project.optional-dependencies] +test = [ + "pytest >=8.0", + "pytest-cov >=4.1", + "pytest-mock >=3.12", + "pytest-xdist >=3.5", + "cibuildwheel == 3.4.0", +] +docs = [ + "mkdocs >=1.5", + "mkdocs-material >=9.5", + "mkdocstrings[python] >=0.24", +] +dev = [ + "ruff >=0.4", + "black >=24.0", + "mypy >=1.8", + "pre-commit >=3.6", + "build >=1.0", + "twine >=4.0", + "numpy >=1.23", + # Inherit test deps + "pytest >=8.0", + "pytest-cov >=4.1", + "pytest-mock >=3.12", + "cibuildwheel == 3.4.0", + # Inherit docs deps + "mkdocs >=1.5", + "mkdocs-material >=9.5", + "mkdocstrings[python] >=0.24", + "pybind11-stubgen>=2.5.5", + "pybind11 >=3.0", +] +###################################################################################################### +# BUILD SYSTEM CONFIGURATION (scikit-build-core) +###################################################################################################### +[build-system] +requires = [ + "scikit-build-core >=0.11", + "pybind11 >=3.0", + "setuptools_scm>=8.0", + "cmake>=3.26,<4.0", + "ninja>=1.11", +] +build-backend = "scikit_build_core.build" + +[tool.scikit-build] +# Core settings +minimum-version = "0.11" +metadata.version.provider = "scikit_build_core.metadata.setuptools_scm" + +# CMake configuration +cmake.version = ">=3.26,<4.0" +ninja.version = ">=1.11" +cmake.build-type = "Release" +install.strip = true # Strip symbols in release builds to reduce wheel size +# Only pull the runtime artifacts tagged COMPONENT "python" into the wheel. +# This excludes the C++ SDK install rules (static/shared libs under lib/ and +# public headers under include/) that cc_library(PACKED) registers for a +# standalone `make install`, which would otherwise bloat the wheel. +install.components = ["python"] + +# Build directory +build-dir = "build" + +# Platform support +wheel.expand-macos-universal-tags = true +wheel.packages = ["python/zvec"] + +# Source distribution +sdist.include = [ + "README.md", + "LICENSE", + "pyproject.toml", + "CMakeLists.txt", + "src/**/*", + "stub/zvec/**/*", + "python/zvec/py.typed", +] + +# CMake defines (env-overridable) +[tool.scikit-build.cmake.define] +BUILD_TOOLS = "OFF" +BUILD_PYTHON_BINDINGS = "ON" +#CMAKE_VERBOSE_MAKEFILE = "ON" + +# Setuptools config for test pypi +[tool.setuptools_scm] +local_scheme = "no-local-version" +version_scheme = "guess-next-dev" +fallback_version = "0.0.0" +###################################################################################################### +# TESTING & QUALITY +###################################################################################################### +[tool.pytest.ini_options] +minversion = "8.0" +addopts = [ + "-ra", + "--showlocals", + "--strict-markers", + "--strict-config", + "--tb=short", + "-n=auto", +] +xfail_strict = true +log_cli_level = "INFO" +filterwarnings = [ + "error", + "ignore::pytest.PytestCacheWarning", + # Ignore numpy deprecation warnings in tests (if any) + "ignore:.*numpy.*:DeprecationWarning", +] +testpaths = ["python/tests"] +markers = [ + "title: Custom marker for test title/description", + # "slow: marks tests as slow", +] + +###################################################################################################### +# BUILD WHEEL +###################################################################################################### +[tool.cibuildwheel] +build = [ + "cp310-*", + "cp311-*", + "cp312-*", + "cp313-*", + "cp314-*", +] +build-frontend = "build" +test-requires = ["pytest", "pytest-xdist", "numpy"] +test-command = "cd {project} && pytest python/tests -v --tb=short" +build-verbosity = 1 + +[tool.cibuildwheel.linux] +# libaio is required by the C++ backend; install it inside the manylinux +# container (manylinux_2_28 is AlmaLinux 8 based, so use dnf/libaio-devel). +# libaio-devel lives in BaseOS; disable EPEL so a flaky EPEL mirror cannot +# break metadata refresh. +before-all = "dnf install -y --disablerepo=epel libaio-devel" +archs = ["auto"] +environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16" } +manylinux-x86_64-image = "manylinux_2_28" +manylinux-aarch64-image = "manylinux_2_28" +# Skip 32-bit builds and musllinux +skip = ["*-manylinux_i686", "*-musllinux*"] + +[tool.cibuildwheel.macos] +archs = ["arm64"] +# Inherits CMAKE_GENERATOR and CMAKE_BUILD_PARALLEL_LEVEL from [tool.cibuildwheel] won't work; +# platform-level environment overrides the top-level entirely, so all vars must be listed here +environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16", MACOSX_DEPLOYMENT_TARGET = "11.0" } + +[tool.cibuildwheel.windows] +archs = ["auto64"] +test-command = "cd /d {project} && pytest python/tests -v --tb=short --basetemp=./.pytest_tmp" +environment = { CMAKE_GENERATOR = "Ninja" } +#config-settings = { build-dir = "build/{wheel_tag}" } +#&& if exist .pytest_tmp rmdir /s /q .pytest_tmp" +before-build = [ + "python -c \"import shutil, glob, os; [shutil.rmtree(p, ignore_errors=True) for p in ['build', '.pytest_tmp']];\"" +] + +###################################################################################################### +# CODE QUALITY & FORMATTING (Ruff) +###################################################################################################### +[tool.ruff] +target-version = "py310" +line-length = 88 +exclude = [ + "build/", + "dist/", + ".git/", + ".venv/", + "venv/", + "thirdparty", +] + +[tool.ruff.lint] +extend-select = [ + "B", # flake8-bugbear + "I", # isort + "ARG", # flake8-unused-arguments + "C4", # flake8-comprehensions + "EM", # flake8-errmsg + "ICN", # flake8-import-conventions + "G", # flake8-logging-format + "PGH", # pygrep-hooks + "PIE", # flake8-pie + "PL", # pylint + "PT", # flake8-pytest-style + "PTH", # flake8-use-pathlib + "RET", # flake8-return + "RUF", # Ruff-specific + "SIM", # flake8-simplify + "T20", # flake8-print + "UP", # pyupgrade + "YTT", # flake8-2020 + "EXE", # flake8-executable + "NPY", # NumPy-specific + "PD", # pandas-vet +] +ignore = [ + "PLR0913", # Too many arguments (common in bindings) + "PLR2004", # Magic value used in comparison + "UP045", "UP007", # Use list() instead of [] (breaks C++ init) + "EM101", "EM102", # Exception messages as literals (ok in tests/utils) + "B008", # Mutable default args (cautiously allowed in config) + "E731", # Lambda assignment (used in callbacks) + "B019", # `functools.lru_cache` on methods (handled manually) + "PLR0912", # Too many branches + "PLC0105", # Ignore contravariant + "RUF002", # Ignore Unicode +] +fixable = ["ALL"] +unfixable = [] + +# Ignore all errors in docstrings +[tool.ruff.lint.pydocstyle] +convention = "google" # or "numpy", "pep257" +ignore-decorators = ["typing.overload"] + +[tool.ruff.lint.flake8-type-checking] +# Don't check code examples in docstrings +quote-annotations = true + +[tool.ruff.lint.isort] +required-imports = ["from __future__ import annotations"] +known-first-party = ["zvec"] + +[tool.ruff.lint.per-file-ignores] +"python/tests/**" = ["ALL"] +"bench/core/**" = ["ALL"] +"python/zvec/__init__.py" = [ + "F401", # Unused import (for __all__) + "E402", # Module level import not at top (C++ module init order) + "PLE0605", # Invalid format for __all__ + "RUF022", # __all__ is not sorted +] +"python/zvec/model/doc.py" = [ + "RUF023", # Unused sort (for __slot__) +] +"python/zvec/extension/**" = [ + "PLC0415", # Import outside top-level (dynamic imports in _get_model) +] + +[tool.ruff.format] +indent-style = "space" +quote-style = "double" +line-ending = "lf" +skip-magic-trailing-comma = false diff --git a/python/tests/detail/distance_helper.py b/python/tests/detail/distance_helper.py new file mode 100644 index 0000000..cf2815c --- /dev/null +++ b/python/tests/detail/distance_helper.py @@ -0,0 +1,392 @@ +import logging +import math +import numpy as np + +from zvec import ( + MetricType, + DataType, + QuantizeType, + Doc, + CollectionSchema, + FieldSchema, + VectorSchema, +) + +from typing import Dict + + +def is_float_equal(actual, expected, rel_tol=1e-5, abs_tol=1e-8): + if actual is None and expected is None: + return True + return math.isclose(actual, expected, rel_tol=rel_tol, abs_tol=abs_tol) + + +def is_dense_vector_equal(vec1, vec2, rtol=1e-5, atol=1e-8): + """Compare two dense vectors with tolerance.""" + return np.allclose(vec1, vec2, rtol=rtol, atol=atol) + + +def is_sparse_vector_equal(vec1, vec2, rtol=1e-5, atol=1e-8): + """Compare two sparse vectors with tolerance.""" + # Check if they have the same keys + if set(vec1.keys()) != set(vec2.keys()): + return False + + # Check if all values are close + for key in vec1: + if not math.isclose(vec1[key], vec2[key], rel_tol=rtol, abs_tol=atol): + return False + + return True + + +def is_float_array_equal(arr1, arr2, rtol=1e-5, atol=1e-8): + """Compare two float arrays with tolerance.""" + return np.allclose(arr1, arr2, rtol=rtol, atol=atol) + + +def is_double_array_equal(arr1, arr2, rtol=1e-9, atol=1e-12): + """Compare two double arrays with tolerance.""" + return np.allclose(arr1, arr2, rtol=rtol, atol=atol) + + +def is_int_array_equal(arr1, arr2): + """Compare two integer arrays with exact equality.""" + return np.array_equal(arr1, arr2) + + +def cosine_distance_dense( + vec1, + vec2, + dtype: DataType = DataType.VECTOR_FP32, + quantize_type: QuantizeType = QuantizeType.UNDEFINED, +): + if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16: + # More stable conversion to float16 to avoid numerical issues + vec1 = [float(np.float16(a)) for a in vec1] + vec2 = [float(np.float16(b)) for b in vec2] + elif dtype == DataType.VECTOR_INT8: + # For INT8 vectors, convert to integers for proper calculation + vec1 = [ + int(round(min(max(val, -128), 127))) for val in vec1 + ] # Clamp to valid INT8 range + vec2 = [ + int(round(min(max(val, -128), 127))) for val in vec2 + ] # Clamp to valid INT8 range + + dot_product = sum(a * b for a, b in zip(vec1, vec2)) + + magnitude1 = math.sqrt(sum(a * a for a in vec1)) + magnitude2 = math.sqrt(sum(b * b for b in vec2)) + + if magnitude1 == 0 or magnitude2 == 0: + return 1.0 # Zero vector case - maximum distance + + cosine_similarity = dot_product / (magnitude1 * magnitude2) + + # Clamp to [-1, 1] range to handle floating-point precision errors + cosine_similarity = max(-1.0, min(1.0, cosine_similarity)) + + # For identical vectors (within floating point precision), ensure cosine distance is 0.0 + # This is especially important for low-precision types which have limited precision + if ( + dtype == DataType.VECTOR_FP16 + or quantize_type == QuantizeType.FP16 + or dtype == DataType.VECTOR_INT8 + ): + if ( + abs(cosine_similarity - 1.0) < 1e-3 + ): # Handle precision issues for low-precision types + cosine_similarity = 1.0 + + # Return cosine distance (1 - cosine similarity) to maintain compatibility + # with system internal processing and existing test expectations + return 1.0 - cosine_similarity + + +def dp_distance_dense( + vec1, + vec2, + dtype: DataType = DataType.VECTOR_FP32, + quantize_type: QuantizeType = QuantizeType.UNDEFINED, +): + if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16: + # More stable computation to avoid numerical issues + products = [ + float(np.float16(a)) * float(np.float16(b)) for a, b in zip(vec1, vec2) + ] + return sum(products) + elif dtype == DataType.VECTOR_INT8: + # For INT8 vectors, convert to integers for proper calculation + products = [ + int(round(min(max(a, -128), 127))) * int(round(min(max(b, -128), 127))) + for a, b in zip(vec1, vec2) + ] + return sum(products) + return sum(a * b for a, b in zip(vec1, vec2)) + + +def euclidean_distance_dense( + vec1, + vec2, + dtype: DataType = DataType.VECTOR_FP32, + quantize_type: QuantizeType = QuantizeType.UNDEFINED, +): + if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16: + # Convert to float16 and compute squared differences safely + # Use a more stable computation to avoid overflow + squared_diffs = [] + for a, b in zip(vec1, vec2): + diff = np.float16(a) - np.float16(b) + squared_diff = float(diff) * float( + diff + ) # Convert to float for multiplication + squared_diffs.append(squared_diff) + squared_distance = sum(squared_diffs) + elif dtype == DataType.VECTOR_INT8: + # For INT8 vectors, convert to integers and handle potential scaling + # INT8 values might be treated differently in the library implementation + vec1_int = [ + int(round(min(max(val, -128), 127))) for val in vec1 + ] # Clamp to valid INT8 range + vec2_int = [ + int(round(min(max(val, -128), 127))) for val in vec2 + ] # Clamp to valid INT8 range + # Use float type to prevent overflow when summing large squared differences + squared_distance = sum(float(a - b) ** 2 for a, b in zip(vec1_int, vec2_int)) + else: + squared_distance = sum((a - b) ** 2 for a, b in zip(vec1, vec2)) + + return squared_distance # Return squared distance for INT8 + + +def distance_dense( + vec1, + vec2, + metric: MetricType, + data_type: DataType = DataType.VECTOR_FP32, + quantize_type: QuantizeType = QuantizeType.UNDEFINED, +): + if metric == MetricType.COSINE: + return cosine_distance_dense(vec1, vec2, data_type, quantize_type) + elif metric == MetricType.L2: + return euclidean_distance_dense(vec1, vec2, data_type, quantize_type) + elif metric == MetricType.IP: + return dp_distance_dense(vec1, vec2, data_type, quantize_type) + else: + raise ValueError("Unsupported metric type") + + +def dp_distance_sparse( + vec1, + vec2, + data_type: DataType = DataType.SPARSE_VECTOR_FP32, + quantize_type: QuantizeType = QuantizeType.UNDEFINED, +): + dot_product = 0.0 + for dim in set(vec1.keys()) & set(vec2.keys()): + print("dim,vec1,vec2:\n") + print(dim, vec1, vec2) + if ( + data_type == DataType.SPARSE_VECTOR_FP16 + or quantize_type == QuantizeType.FP16 + ): + vec1[dim] = np.float16(vec1[dim]) + vec2[dim] = np.float16(vec2[dim]) + dot_product += vec1[dim] * vec2[dim] + return dot_product + + +def distance( + vec1, + vec2, + metric: MetricType, + data_type: DataType, + quantize_type: QuantizeType = QuantizeType.UNDEFINED, +): + is_sparse = ( + data_type == DataType.SPARSE_VECTOR_FP32 + or data_type == DataType.SPARSE_VECTOR_FP16 + ) + + if is_sparse: + if metric != MetricType.IP: + raise ValueError("Unsupported metric type for sparse vectors") + + if is_sparse: + return dp_distance_sparse(vec1, vec2, data_type, quantize_type) + else: + return distance_dense(vec1, vec2, metric, data_type, quantize_type) + + +def distance_recall( + vec1, + vec2, + metric: MetricType, + data_type: DataType, + quantize_type: QuantizeType = QuantizeType.UNDEFINED, +): + is_sparse = ( + data_type == DataType.SPARSE_VECTOR_FP32 + or data_type == DataType.SPARSE_VECTOR_FP16 + ) + + if is_sparse: + return dp_distance_sparse(vec1, vec2, data_type, quantize_type) + else: + if data_type in [DataType.VECTOR_FP32, DataType.VECTOR_FP16]: + return distance_dense(vec1, vec2, metric, data_type, quantize_type) + elif data_type in [DataType.VECTOR_INT8] and metric in [ + MetricType.L2, + MetricType.IP, + ]: + return distance_dense(vec1, vec2, metric, data_type, quantize_type) + else: + return dp_distance_dense(vec1, vec2, data_type, quantize_type) + + +def calculate_rrf_score(rank, k=60): + return 1.0 / (k + rank + 1) + + +def calculate_multi_vector_rrf_scores(query_results: Dict[str, Doc], k=60): + rrf_scores = {} + + for vector_name, docs in query_results.items(): + for rank, doc in enumerate(docs): + doc_id = doc.id + rrf_score = calculate_rrf_score(rank, k) + if doc_id in rrf_scores: + rrf_scores[doc_id] += rrf_score + else: + rrf_scores[doc_id] = rrf_score + + return rrf_scores + + +def calculate_multi_vector_weighted_scores( + query_results: Dict[str, Doc], weights: Dict[str, float], metric: MetricType +): + def _normalize_score(score: float, metric: MetricType) -> float: + if metric == MetricType.L2: + return 1.0 - 2 * math.atan(score) / math.pi + if metric == MetricType.IP: + return 0.5 + math.atan(score) / math.pi + if metric == MetricType.COSINE: + return 1.0 - score / 2.0 + raise ValueError("Unsupported metric type") + + weighted_scores = {} + + for vector_name, docs in query_results.items(): + weight = weights.get(vector_name, 1.0) + + for doc in docs: + doc_id = doc.id + weighted_score = (_normalize_score(doc.score, metric)) * weight + if doc_id in weighted_scores: + weighted_scores[doc_id] += weighted_score + else: + weighted_scores[doc_id] = weighted_score + + return weighted_scores + + +def is_field_equal(field1, field2, schema: FieldSchema) -> bool: + if field1 is None and field2 is None: + return True + if field1 is None or field2 is None: + return False + + if schema.data_type == DataType.ARRAY_FLOAT: + return is_float_array_equal(field1, field2) + elif schema.data_type == DataType.ARRAY_DOUBLE: + return is_double_array_equal(field1, field2) + elif schema.data_type in [ + DataType.ARRAY_INT32, + DataType.ARRAY_INT64, + DataType.ARRAY_BOOL, + DataType.ARRAY_STRING, + DataType.ARRAY_UINT32, + DataType.ARRAY_UINT64, + DataType.ARRAY_INT64, + ]: + return is_int_array_equal(field1, field2) + elif schema.data_type in [DataType.FLOAT, DataType.DOUBLE]: + return is_float_equal(field1, field2) + + return field1 == field2 + + +def is_vector_equal(vec1, vec2, schema: VectorSchema) -> bool: + if ( + schema.data_type == DataType.SPARSE_VECTOR_FP16 + or schema.data_type == DataType.VECTOR_FP16 + ): + # skip fp16 vector equal + return True + + is_sparse = ( + schema.data_type == DataType.SPARSE_VECTOR_FP32 + or schema.data_type == DataType.SPARSE_VECTOR_FP16 + ) + + if is_sparse: + return is_sparse_vector_equal(vec1, vec2) + else: + return is_dense_vector_equal(vec1, vec2) + + +def is_doc_equal( + doc1: Doc, + doc2: Doc, + schema: CollectionSchema, + except_score: bool = True, + include_vector: bool = True, +): + if doc1.id != doc2.id: + logging.error("doc ids are not equal") + return False + + reduce_field_names = set(doc1.field_names() + doc2.field_names()) + reduce_vector_names = set(doc1.vector_names() + doc2.vector_names()) + + is_doc1_fields_empty = doc1.fields is None or doc1.fields == {} + is_doc2_fields_empty = doc2.fields is None or doc2.fields == {} + + if is_doc1_fields_empty or is_doc2_fields_empty: + if is_doc1_fields_empty != is_doc2_fields_empty: + return False + else: + for field_name in reduce_field_names: + field_schema = schema.field(field_name) + if field_schema is None: + return False + if is_field_equal( + doc1.field(field_name), doc2.field(field_name), field_schema + ): + continue + else: + logging.error(f"{field_name} are not equal") + return False + + if include_vector: + is_doc1_vectors_empty = doc1.vectors is None or doc1.vectors == {} + is_doc2_vectors_empty = doc2.vectors is None or doc2.vectors == {} + + if is_doc1_vectors_empty or is_doc2_vectors_empty: + if is_doc1_fields_empty != is_doc2_vectors_empty: + return False + else: + for vector_name in reduce_vector_names: + vector_schema = schema.vector(vector_name) + if vector_schema is None: + return False + if is_vector_equal( + doc1.vector(vector_name), doc2.vector(vector_name), vector_schema + ): + continue + else: + return False + + return True diff --git a/python/tests/detail/doc_helper.py b/python/tests/detail/doc_helper.py new file mode 100644 index 0000000..09e7892 --- /dev/null +++ b/python/tests/detail/doc_helper.py @@ -0,0 +1,464 @@ +from zvec import CollectionSchema, Doc + +from support_helper import * + +import numpy as np +from typing import Literal, Optional, Union, Tuple + +import random +import string +import math + + +def generate_constant_vector( + i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32" +): + if dtype == "int8": + vec = [(i % 127)] * dimension + vec[i % dimension] = (i + 1) % 127 + else: + base_val = (i % 1000) / 256.0 + special_val = ((i + 1) % 1000) / 256.0 + vec = [base_val] * dimension + vec[i % dimension] = special_val + + return vec + + +def generate_constant_vector_recall( + i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32" +): + if dtype == "int8": + vec = [(i % 127)] * dimension + vec[i % dimension] = (i + 1) % 127 + else: + base_val = math.sin((i) * 1000) / 256.0 + special_val = math.sin((i + 1) * 1000) / 256.0 + vec = [base_val] * dimension + vec[i % dimension] = special_val + + return vec + + +def generate_sparse_vector(i: int): + return {i: i + 0.1} + + +def generate_vectordict(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields = {} + doc_vectors = {} + for field in schema.fields: + if field.data_type == DataType.BOOL: + doc_fields[field.name] = i % 2 == 0 + elif field.data_type == DataType.INT32: + doc_fields[field.name] = i + elif field.data_type == DataType.UINT32: + doc_fields[field.name] = i + elif field.data_type == DataType.INT64: + doc_fields[field.name] = i + elif field.data_type == DataType.UINT64: + doc_fields[field.name] = i + elif field.data_type == DataType.FLOAT: + doc_fields[field.name] = float(i) + 0.1 + elif field.data_type == DataType.DOUBLE: + doc_fields[field.name] = float(i) + 0.11 + elif field.data_type == DataType.STRING: + doc_fields[field.name] = f"test_{i}" + elif field.data_type == DataType.ARRAY_BOOL: + doc_fields[field.name] = [i % 2 == 0, i % 3 == 0] + elif field.data_type == DataType.ARRAY_INT32: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_UINT32: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_INT64: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_UINT64: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_FLOAT: + doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)] + elif field.data_type == DataType.ARRAY_DOUBLE: + doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)] + elif field.data_type == DataType.ARRAY_STRING: + doc_fields[field.name] = [f"test_{i}", f"test_{i + 1}", f"test_{i + 2}"] + else: + raise ValueError(f"Unsupported field type: {field.data_type}") + for vector in schema.vectors: + if vector.data_type == DataType.VECTOR_FP16: + doc_vectors[vector.name] = generate_constant_vector( + i, vector.dimension, "float16" + ) + elif vector.data_type == DataType.VECTOR_FP32: + doc_vectors[vector.name] = generate_constant_vector( + i, vector.dimension, "float32" + ) + elif vector.data_type == DataType.VECTOR_INT8: + doc_vectors[vector.name] = generate_constant_vector( + i, + vector.dimension, + "int8", + ) + elif vector.data_type == DataType.SPARSE_VECTOR_FP32: + doc_vectors[vector.name] = generate_sparse_vector(i) + elif vector.data_type == DataType.SPARSE_VECTOR_FP16: + doc_vectors[vector.name] = generate_sparse_vector(i) + else: + raise ValueError(f"Unsupported vector type: {vector.data_type}") + return doc_fields, doc_vectors + + +def generate_vectordict_recall(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields = {} + doc_vectors = {} + for field in schema.fields: + if field.data_type == DataType.BOOL: + doc_fields[field.name] = i % 2 == 0 + elif field.data_type == DataType.INT32: + doc_fields[field.name] = i + elif field.data_type == DataType.UINT32: + doc_fields[field.name] = i + elif field.data_type == DataType.INT64: + doc_fields[field.name] = i + elif field.data_type == DataType.UINT64: + doc_fields[field.name] = i + elif field.data_type == DataType.FLOAT: + doc_fields[field.name] = float(i) + 0.1 + elif field.data_type == DataType.DOUBLE: + doc_fields[field.name] = float(i) + 0.11 + elif field.data_type == DataType.STRING: + doc_fields[field.name] = f"test_{i}" + elif field.data_type == DataType.ARRAY_BOOL: + doc_fields[field.name] = [i % 2 == 0, i % 3 == 0] + elif field.data_type == DataType.ARRAY_INT32: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_UINT32: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_INT64: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_UINT64: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_FLOAT: + doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)] + elif field.data_type == DataType.ARRAY_DOUBLE: + doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)] + elif field.data_type == DataType.ARRAY_STRING: + doc_fields[field.name] = [f"test_{i}", f"test_{i + 1}", f"test_{i + 2}"] + else: + raise ValueError(f"Unsupported field type: {field.data_type}") + for vector in schema.vectors: + if vector.data_type == DataType.VECTOR_FP16: + doc_vectors[vector.name] = generate_constant_vector_recall( + i, vector.dimension, "float16" + ) + elif vector.data_type == DataType.VECTOR_FP32: + doc_vectors[vector.name] = generate_constant_vector_recall( + i, vector.dimension, "float32" + ) + elif vector.data_type == DataType.VECTOR_INT8: + doc_vectors[vector.name] = generate_constant_vector_recall( + i, + vector.dimension, + "int8", + ) + elif vector.data_type == DataType.SPARSE_VECTOR_FP32: + doc_vectors[vector.name] = generate_sparse_vector(i) + elif vector.data_type == DataType.SPARSE_VECTOR_FP16: + doc_vectors[vector.name] = generate_sparse_vector(i) + else: + raise ValueError(f"Unsupported vector type: {vector.data_type}") + return doc_fields, doc_vectors + + +def generate_vectordict_update(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields = {} + doc_vectors = {} + for field in schema.fields: + if field.data_type == DataType.BOOL: + doc_fields[field.name] = (i + 1) % 2 == 0 + elif field.data_type == DataType.INT32: + doc_fields[field.name] = i + 1 + elif field.data_type == DataType.UINT32: + doc_fields[field.name] = i + 1 + elif field.data_type == DataType.INT64: + doc_fields[field.name] = i + 1 + elif field.data_type == DataType.UINT64: + doc_fields[field.name] = i + 1 + elif field.data_type == DataType.FLOAT: + doc_fields[field.name] = float(i + 1) + 0.1 + elif field.data_type == DataType.DOUBLE: + doc_fields[field.name] = float(i + 1) + 0.11 + elif field.data_type == DataType.STRING: + doc_fields[field.name] = f"test_{i + 1}" + elif field.data_type == DataType.ARRAY_BOOL: + doc_fields[field.name] = [(i + 1) % 2 == 0, (i + 1) % 3 == 0] + elif field.data_type == DataType.ARRAY_INT32: + doc_fields[field.name] = [i + 1, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_UINT32: + doc_fields[field.name] = [i + 1, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_INT64: + doc_fields[field.name] = [i + 1, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_UINT64: + doc_fields[field.name] = [i + 1, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_FLOAT: + doc_fields[field.name] = [float(i + 1.1), float(i + 2.1), float(i + 3.1)] + elif field.data_type == DataType.ARRAY_DOUBLE: + doc_fields[field.name] = [float(i + 1.11), float(i + 2.11), float(i + 3.11)] + elif field.data_type == DataType.ARRAY_STRING: + doc_fields[field.name] = [f"test_{i + 1}", f"test_{i + 2}", f"test_{i + 3}"] + else: + raise ValueError(f"Unsupported field type: {field.data_type}") + for vector in schema.vectors: + if vector.data_type == DataType.VECTOR_FP16: + doc_vectors[vector.name] = generate_constant_vector( + i + 1, vector.dimension, "float16" + ) + elif vector.data_type == DataType.VECTOR_FP32: + doc_vectors[vector.name] = generate_constant_vector( + i + 1, vector.dimension, "float32" + ) + elif vector.data_type == DataType.VECTOR_INT8: + doc_vectors[vector.name] = generate_constant_vector( + i + 1, + vector.dimension, + "int8", + ) + elif vector.data_type == DataType.SPARSE_VECTOR_FP32: + doc_vectors[vector.name] = generate_sparse_vector(i + 1) + elif vector.data_type == DataType.SPARSE_VECTOR_FP16: + doc_vectors[vector.name] = generate_sparse_vector(i + 1) + else: + raise ValueError(f"Unsupported vector type: {vector.data_type}") + return doc_fields, doc_vectors + + +def generate_doc(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields, doc_vectors = generate_vectordict(i, schema) + doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + return doc + + +def generate_doc_recall(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields, doc_vectors = generate_vectordict_recall(i, schema) + doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + return doc + + +def generate_update_doc(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields, doc_vectors = generate_vectordict_update(i, schema) + doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + return doc + + +def generate_doc_random(i, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + + random.seed(i) + + for field in schema.fields: + if field.data_type == DataType.BOOL: + doc_fields[field.name] = random.choice([True, False]) + elif field.data_type == DataType.INT32: + doc_fields[field.name] = random.randint(-2147483648, 2147483647) + elif field.data_type == DataType.UINT32: + doc_fields[field.name] = random.randint(0, 4294967295) + elif field.data_type == DataType.INT64: + doc_fields[field.name] = random.randint( + -9223372036854775808, 9223372036854775807 + ) + elif field.data_type == DataType.UINT64: + doc_fields[field.name] = random.randint(0, 18446744073709551615) + elif field.data_type == DataType.FLOAT: + doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38) + elif field.data_type == DataType.DOUBLE: + doc_fields[field.name] = random.uniform( + -1.7976931348623157e308, 1.7976931348623157e308 + ) + elif field.data_type == DataType.STRING: + length = random.randint(1, 999) + doc_fields[field.name] = "".join( + random.choices(string.ascii_letters + string.digits, k=length) + ) + elif field.data_type == DataType.ARRAY_BOOL: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.choice([True, False]) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_INT32: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.randint(-2147483648, 2147483647) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_UINT32: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.randint(0, 4294967295) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_INT64: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.randint(-9223372036854775808, 9223372036854775807) + for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_UINT64: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.randint(0, 18446744073709551615) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_FLOAT: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_DOUBLE: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.uniform(-1.7976931348623157e308, 1.7976931348623157e308) + for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_STRING: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + "".join( + random.choices( + string.ascii_letters + string.digits, k=random.randint(1, 100) + ) + ) + for _ in range(array_length) + ] + else: + raise ValueError(f"Unsupported field type: {field.data_type}") + + for vector in schema.vectors: + if vector.data_type == DataType.VECTOR_FP16: + doc_vectors[vector.name] = generate_constant_vector( + random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "float16" + ) + elif vector.data_type == DataType.VECTOR_FP32: + doc_vectors[vector.name] = generate_constant_vector( + random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "float32" + ) + elif vector.data_type == DataType.VECTOR_INT8: + doc_vectors[vector.name] = generate_constant_vector( + random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "int8" + ) + elif vector.data_type == DataType.SPARSE_VECTOR_FP32: + doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100)) + elif vector.data_type == DataType.SPARSE_VECTOR_FP16: + doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100)) + else: + raise ValueError(f"Unsupported vector type: {vector.data_type}") + + doc = Doc(id=i, fields=doc_fields, vectors=doc_vectors) + return doc + + +def generate_vectordict_random(schema: CollectionSchema): + doc_fields = {} + doc_vectors = {} + for field in schema.fields: + if field.data_type == DataType.BOOL: + doc_fields[field.name] = random.choice([True, False]) + elif field.data_type == DataType.INT32: + doc_fields[field.name] = random.randint(-2147483648, 2147483647) + elif field.data_type == DataType.UINT32: + doc_fields[field.name] = random.randint(0, 4294967295) + elif field.data_type == DataType.INT64: + doc_fields[field.name] = random.randint( + -9223372036854775808, 9223372036854775807 + ) + elif field.data_type == DataType.UINT64: + doc_fields[field.name] = random.randint(0, 18446744073709551615) + elif field.data_type == DataType.FLOAT: + doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38) + elif field.data_type == DataType.DOUBLE: + doc_fields[field.name] = random.uniform( + -1.7976931348623157e308, 1.7976931348623157e308 + ) + elif field.data_type == DataType.STRING: + length = random.randint(1, 999) + doc_fields[field.name] = "".join( + random.choices(string.ascii_letters + string.digits, k=length) + ) + elif field.data_type == DataType.ARRAY_BOOL: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.choice([True, False]) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_INT32: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.randint(-2147483648, 2147483647) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_UINT32: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.randint(0, 4294967295) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_INT64: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.randint(-9223372036854775808, 9223372036854775807) + for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_UINT64: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.randint(0, 18446744073709551615) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_FLOAT: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_DOUBLE: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + random.uniform(-1.7976931348623157e308, 1.7976931348623157e308) + for _ in range(array_length) + ] + elif field.data_type == DataType.ARRAY_STRING: + array_length = random.randint(0, 10) + doc_fields[field.name] = [ + "".join( + random.choices( + string.ascii_letters + string.digits, k=random.randint(1, 100) + ) + ) + for _ in range(array_length) + ] + else: + raise ValueError(f"Unsupported field type: {field.data_type}") + + for vector in schema.vectors: + if vector.data_type == DataType.VECTOR_FP16: + doc_vectors[vector.name] = generate_constant_vector( + random.randint(1, 100), vector.dimension, "float16" + ) + elif vector.data_type == DataType.VECTOR_FP32: + doc_vectors[vector.name] = generate_constant_vector( + random.randint(1, 100), vector.dimension, "float32" + ) + elif vector.data_type == DataType.VECTOR_INT8: + doc_vectors[vector.name] = generate_constant_vector( + random.randint(1, 100), vector.dimension, "int8" + ) + elif vector.data_type == DataType.SPARSE_VECTOR_FP32: + doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100)) + elif vector.data_type == DataType.SPARSE_VECTOR_FP16: + doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100)) + else: + raise ValueError(f"Unsupported vector type: {vector.data_type}") + + return doc_fields, doc_vectors diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py new file mode 100644 index 0000000..63aac79 --- /dev/null +++ b/python/tests/detail/fixture_helper.py @@ -0,0 +1,652 @@ +import pytest +import logging +import platform + +DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in ( + "x86_64", + "AMD64", + "i686", + "i386", +) + +from typing import Any, Generator +from zvec.typing import DataType, StatusCode, MetricType, QuantizeType +import zvec + + +# Cache the DiskAnn plugin preload status so we pay the load cost once per +# test session. The plugin normally auto-loads on first DiskAnn use, but we +# preload it explicitly here so a missing libaio / misplaced plugin .so +# surfaces as a clear pytest skip instead of a confusing +# "Create vector column indexer failed" deep inside the collection code path. +_DISKANN_PRELOAD_REASON: str | None = None +_DISKANN_PRELOAD_DONE: bool = False + + +def _ensure_diskann_runtime_or_reason() -> str | None: + """Preload the DiskAnn plugin and return None on success or a human-readable + skip reason on failure. Idempotent across calls.""" + global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON + if _DISKANN_PRELOAD_DONE: + return _DISKANN_PRELOAD_REASON + _DISKANN_PRELOAD_DONE = True + + if not DISKANN_SUPPORTED: + _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" + return _DISKANN_PRELOAD_REASON + + if not zvec.is_libaio_available(): + _DISKANN_PRELOAD_REASON = ( + "libaio is not available on this host; DiskAnn cannot run. " + "Install libaio1 (or libaio1t64 on Ubuntu 24.04+) and retry." + ) + return _DISKANN_PRELOAD_REASON + + status = zvec.load_diskann_plugin() + if status != zvec.DISKANN_PLUGIN_OK: + _DISKANN_PRELOAD_REASON = ( + f"Failed to load DiskAnn plugin (status={status}); " + "check that libzvec_diskann_plugin.so is installed alongside " + "_zvec.so in the Python site-packages directory." + ) + return _DISKANN_PRELOAD_REASON + + _DISKANN_PRELOAD_REASON = None + return None + + +from zvec import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + FlatIndexParam, + IVFIndexParam, + FieldSchema, + VectorSchema, + CollectionSchema, + Collection, + Doc, + Query, +) + +from support_helper import * + + +@pytest.fixture(scope="session") +def basic_schema(collection_name="test_collection") -> CollectionSchema: + return CollectionSchema( + name=collection_name if len(collection_name) > 0 else "test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", DataType.STRING, nullable=False, index_param=InvertIndexParam() + ), + FieldSchema("weight", DataType.FLOAT, nullable=True), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + VectorSchema( + "sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() + ), + ], + ) + + +@pytest.fixture(scope="session") +def full_schema( + nullable: bool = False, + has_index: bool = False, +) -> CollectionSchema: + scalar_index_param = None + vector_index_param = None + if has_index: + scalar_index_param = InvertIndexParam(enable_range_optimization=True) + vector_index_param = HnswIndexParam() + + fields = [] + for k, v in DEFAULT_SCALAR_FIELD_NAME.items(): + fields.append( + FieldSchema( + v, + k, + nullable=nullable, + index_param=scalar_index_param, + ) + ) + vetors = [] + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + vetors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + + return CollectionSchema( + name="full_collection", + fields=fields, + vectors=vetors, + ) + + +@pytest.fixture(scope="function") +def full_schema_new(request) -> CollectionSchema: + if hasattr(request, "param"): + nullable, has_index, vector_index = request.param + else: + nullable, has_index, vector_index = True, False, HnswIndexParam() + + # Skip DiskAnn tests on unsupported platforms or when the runtime cannot + # be brought up (missing libaio, plugin .so not installed, etc.). + from zvec.model.param import DiskAnnIndexParam + + if isinstance(vector_index, DiskAnnIndexParam): + skip_reason = _ensure_diskann_runtime_or_reason() + if skip_reason is not None: + pytest.skip(skip_reason) + + scalar_index_param = None + vector_index_param = None + if has_index: + scalar_index_param = InvertIndexParam(enable_range_optimization=True) + vector_index_param = vector_index + + fields = [] + for k, v in DEFAULT_SCALAR_FIELD_NAME.items(): + fields.append( + FieldSchema( + v, + k, + nullable=nullable, + index_param=scalar_index_param, + ) + ) + vectors = [] + + if vector_index_param in [ + HnswIndexParam(), + FlatIndexParam(), + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + ), + FlatIndexParam( + metric_type=MetricType.IP, + ), + ]: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + elif vector_index_param in [ + IVFIndexParam(), + IVFIndexParam( + metric_type=MetricType.IP, + n_list=100, + n_iters=10, + use_soar=False, + ), + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + ( + IVFIndexParam( + metric_type=MetricType.COSINE, + n_list=150, + n_iters=15, + use_soar=False, + ) + ), + ( + HnswIndexParam( + metric_type=MetricType.COSINE, + m=24, + ef_construction=150, + ) + ), + ( + HnswIndexParam( + metric_type=MetricType.L2, + m=32, + ef_construction=200, + ) + ), + ( + FlatIndexParam( + metric_type=MetricType.COSINE, + ) + ), + ( + FlatIndexParam( + metric_type=MetricType.L2, + ) + ), + ]: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + elif v in ["vector_int8_field"] and vector_index_param in [ + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + ( + HnswIndexParam( + metric_type=MetricType.L2, + m=32, + ef_construction=200, + ) + ), + ( + FlatIndexParam( + metric_type=MetricType.L2, + ) + ), + ]: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + else: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=HnswIndexParam(), + ) + ) + else: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + else: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=HnswIndexParam(), + ) + ) + + return CollectionSchema( + name="full_collection_new", + fields=fields, + vectors=vectors, + ) + + +@pytest.fixture(scope="function") +def full_schema_ivf(request) -> CollectionSchema: + if hasattr(request, "param"): + nullable, has_index, vector_index = request.param + else: + nullable, has_index, vector_index = True, False, IVFIndexParam() + + scalar_index_param = None + vector_index_param = None + if has_index: + scalar_index_param = InvertIndexParam(enable_range_optimization=True) + vector_index_param = vector_index + + fields = [] + for k, v in DEFAULT_SCALAR_FIELD_NAME.items(): + fields.append( + FieldSchema( + v, + k, + nullable=nullable, + index_param=scalar_index_param, + ) + ) + vectors = [] + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + + return CollectionSchema( + name="full_collection_ivf", + fields=fields, + vectors=vectors, + ) + + +@pytest.fixture(scope="function") +def full_schema_1024(request) -> CollectionSchema: + if hasattr(request, "param"): + nullable, has_index, vector_index = request.param + else: + nullable, has_index, vector_index = True, False, HnswIndexParam() + + scalar_index_param = None + vector_index_param = None + if has_index: + scalar_index_param = InvertIndexParam(enable_range_optimization=True) + vector_index_param = vector_index + + fields = [] + for k, v in DEFAULT_SCALAR_FIELD_NAME.items(): + fields.append( + FieldSchema( + v, + k, + nullable=nullable, + index_param=scalar_index_param, + ) + ) + vectors = [] + + if vector_index_param in [ + HnswIndexParam(), + FlatIndexParam(), + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + ), + FlatIndexParam( + metric_type=MetricType.IP, + ), + ]: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=vector_index_param, + ) + ) + elif vector_index_param in [ + IVFIndexParam(), + IVFIndexParam( + metric_type=MetricType.IP, + n_list=100, + n_iters=10, + use_soar=False, + ), + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + IVFIndexParam( + metric_type=MetricType.COSINE, + n_list=150, + n_iters=15, + use_soar=False, + ), + ]: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=vector_index_param, + ) + ) + elif v in ["vector_int8_field"] and vector_index_param in [ + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + IVFIndexParam( + metric_type=MetricType.COSINE, + n_list=150, + n_iters=15, + use_soar=False, + ), + ]: + vectors.append( + VectorSchema( + v, + k, + dimension=DVECTOR_DIMENSION_1024, + index_param=vector_index_param, + ) + ) + else: + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=HnswIndexParam(), + ) + ) + else: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field", "vector_int8_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=vector_index_param, + ) + ) + else: + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=HnswIndexParam(), + ) + ) + + return CollectionSchema( + name="full_collection_new", + fields=fields, + vectors=vectors, + ) + + +@pytest.fixture(scope="function") +def single_vector_schema( + data_type: DataType, +) -> CollectionSchema: + vector_schema = [ + VectorSchema( + DEFAULT_VECTOR_FIELD_NAME[data_type], + data_type, + DEFAULT_VECTOR_DIMENSION, + ) + ] + + return CollectionSchema( + name="full_collection", + vectors=vector_schema, + ) + + +@pytest.fixture(scope="function") +def single_vector_schema_with_index_param( + data_type: DataType, index_param +) -> CollectionSchema: + vector_schema = [ + VectorSchema( + DEFAULT_VECTOR_FIELD_NAME[data_type], + data_type, + DEFAULT_VECTOR_DIMENSION, + index_param, + ) + ] + + return CollectionSchema( + name="full_collection", + vectors=vector_schema, + ) + + +def create_collection_fixture( + collection_temp_dir, schema: CollectionSchema, collection_option: CollectionOption +) -> Generator[Any, Any, Collection]: + """Common helper function to create and manage collection fixtures.""" + coll = zvec.create_and_open( + path=str(collection_temp_dir), + schema=schema, + option=collection_option, + ) + + assert coll is not None, "Failed to create and open collection" + assert coll.path == str(collection_temp_dir) + assert coll.schema.name == schema.name + assert list(coll.schema.fields) == list(schema.fields) + assert list(coll.schema.vectors) == list(schema.vectors) + assert coll.option.read_only == collection_option.read_only + assert coll.option.enable_mmap == collection_option.enable_mmap + + try: + yield coll + finally: + if hasattr(coll, "destroy") and coll is not None: + try: + coll.destroy() + except Exception as e: + logging.warning(f"Warning: failed to destroy collection: {e}") + + +@pytest.fixture(scope="function") +def basic_collection( + collection_temp_dir, basic_schema, collection_option +) -> Generator[Any, Any, Collection]: + yield from create_collection_fixture( + collection_temp_dir, basic_schema, collection_option + ) + + +@pytest.fixture(scope="function") +def collection_option(): + return CollectionOption(read_only=False, enable_mmap=True) + + +@pytest.fixture(scope="function") +def collection_temp_dir(tmp_path_factory): + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection_path" + return str(collection_path) + + +@pytest.fixture(scope="function") +def full_collection( + collection_temp_dir, + full_schema, + collection_option, + nullable: bool = True, + has_index: bool = False, +) -> Generator[Any, Any, Collection]: + yield from create_collection_fixture( + collection_temp_dir, full_schema, collection_option + ) + + +@pytest.fixture(scope="function") +def full_collection_new( + collection_temp_dir, full_schema_new, collection_option +) -> Generator[Any, Any, Collection]: + yield from create_collection_fixture( + collection_temp_dir, full_schema_new, collection_option + ) + + +@pytest.fixture(scope="function") +def full_collection_ivf( + collection_temp_dir, full_schema_ivf, collection_option +) -> Generator[Any, Any, Collection]: + yield from create_collection_fixture( + collection_temp_dir, full_schema_ivf, collection_option + ) + + +@pytest.fixture(scope="function") +def full_collection_1024( + collection_temp_dir, full_schema_1024, collection_option +) -> Generator[Any, Any, Collection]: + yield from create_collection_fixture( + collection_temp_dir, full_schema_1024, collection_option + ) + + +@pytest.fixture +def sample_field_list(nullable: bool = True, scalar_index_param=None, name_prefix=""): + field_list = [] + for k, v in DEFAULT_SCALAR_FIELD_NAME.items(): + field_list.append( + FieldSchema( + f"{name_prefix}_{v}" if len(name_prefix) > 0 else v, + k, + nullable=nullable, + index_param=scalar_index_param, + ) + ) + return field_list + + +@pytest.fixture +def sample_vector_list(vector_index_param=None, name_prefix=""): + vector_list = [] + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + vector_list.append( + VectorSchema( + f"{name_prefix}_{v}" if len(name_prefix) > 0 else v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + return vector_list diff --git a/python/tests/detail/params_helper.py b/python/tests/detail/params_helper.py new file mode 100644 index 0000000..e373005 --- /dev/null +++ b/python/tests/detail/params_helper.py @@ -0,0 +1,209 @@ +from zvec import ( + CollectionOption, + IndexOption, + OptimizeOption, + InvertIndexParam, + HnswIndexParam, + IVFIndexParam, + FlatIndexParam, + AlterColumnOption, + AddColumnOption, + DataType, + MetricType, + QuantizeType, +) + + +VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP = { + DataType.VECTOR_FP32: [ + HnswIndexParam(), + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + quantize_type=QuantizeType.INT8, + ), + HnswIndexParam( + metric_type=MetricType.COSINE, + m=24, + ef_construction=150, + quantize_type=QuantizeType.INT4, + ), + HnswIndexParam( + metric_type=MetricType.L2, + m=32, + ef_construction=200, + quantize_type=QuantizeType.FP16, + ), + FlatIndexParam(), + FlatIndexParam(metric_type=MetricType.IP, quantize_type=QuantizeType.INT4), + FlatIndexParam(metric_type=MetricType.L2, quantize_type=QuantizeType.INT8), + FlatIndexParam(metric_type=MetricType.COSINE, quantize_type=QuantizeType.FP16), + IVFIndexParam(), + IVFIndexParam( + metric_type=MetricType.IP, + quantize_type=QuantizeType.INT4, + n_list=100, + n_iters=10, + use_soar=False, + ), + IVFIndexParam( + metric_type=MetricType.L2, + quantize_type=QuantizeType.INT8, + n_list=200, + n_iters=20, + use_soar=True, + ), + IVFIndexParam( + metric_type=MetricType.COSINE, + quantize_type=QuantizeType.FP16, + n_list=150, + n_iters=15, + use_soar=False, + ), + ], + DataType.VECTOR_FP16: [ + HnswIndexParam(), + FlatIndexParam(), + # IVFIndexParam(), + ], + DataType.VECTOR_INT8: [ + HnswIndexParam(), + FlatIndexParam(), + # IVFIndexParam(), + ], + DataType.SPARSE_VECTOR_FP32: [ + HnswIndexParam(), + FlatIndexParam(), + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + quantize_type=QuantizeType.FP16, + ), + ], + DataType.SPARSE_VECTOR_FP16: [ + HnswIndexParam(), + FlatIndexParam(), + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + ), + ], +} + +VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS = [ + (data_type, param) + for data_type, params in VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP.items() + for param in params +] + +INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP = { + DataType.VECTOR_FP32: [ + InvertIndexParam(), + ], + DataType.VECTOR_FP16: [ + InvertIndexParam(), + ], + DataType.VECTOR_INT8: [ + InvertIndexParam(), + ], + DataType.SPARSE_VECTOR_FP32: [ + HnswIndexParam(metric_type=MetricType.L2), + FlatIndexParam(metric_type=MetricType.COSINE), + IVFIndexParam(), + InvertIndexParam(), + ], + DataType.SPARSE_VECTOR_FP16: [ + HnswIndexParam(metric_type=MetricType.L2), + FlatIndexParam(metric_type=MetricType.COSINE), + IVFIndexParam(), + InvertIndexParam(), + ], +} + +INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS = [ + (data_type, param) + for data_type, params in INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP.items() + for param in params +] + +COLLECTION_NAME_MAX_LENGTH = 64 + +COLLECTION_NAME_VALID_LIST = [ + "col", + "C0llECTION", + "Collection1", + "collection_2", + "123collection-", + "a" * COLLECTION_NAME_MAX_LENGTH, +] + +COLLECTION_NAME_INVALID_LIST = [ + "l", + "1C", + "", + " ", + None, + "abcdefghijklmnopqrstuvwxzy123456abcdefghijklmnopqrstuvwxzy1234561", + "test/", + "!@#$%^&*()test", +] + +FIELD_NAME_VALID_LIST = [ + "1", + "12", + "col", + "ID", + "name1", + "Weigt_12-", + "123age", + "name_with_underscores", + "123numeric_start", + "name-with-dashes", +] + +FIELD_NAME_INVALID_LIST = [ + "", + " ", + None, + "abcdefghijklmnopqrstuvwxzy1234561", + "test/", + "!@#$%^&*()test", + "name@with#special$chars", + "name with spaces", +] + +FIELD_LIST_MAX_LENGTH = 1024 +VECTOR_LIST_MAX_LENGTH = 5 +DENSE_VECTOR_MAX_DIMENSION = 20000 +SPARSE_VECTOR_MAX_DIMENSION = 4096 + +FIELD_VECTOR_LIST_DIMENSION_VALID_LIST = [ + # field_list_len, vector_list_len, dimension + (1, 1, 1), + (2, 2, 512), + (512, 3, 1024), + (1024, 4, 20000), +] + +FIELD_VECTOR_LIST_DIMENSION_INVALID_LIST = [ + # field_list_len, vector_list_len, dimension + (1, 1, 0), + (1, 1, -1), + (1, 1, "1"), + (1, 1, 20001), +] + + +INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG = "incompatible constructor arguments" +SCHEMA_VALIDATE_ERROR_MSG = "schema validate failed" +CREATE_READ_ONLY_ERROR_MSG = "Unable to create collection with read-only mode" +INCOMPATIBLE_FUNCTION_ERROR_MSG = "incompatible function arguments" +INVALID_PATH_ERROR_MSG = "path validate failed" +INDEX_NON_EXISTENT_COLUMN_ERROR_MSG = "not found in schema" +ACCESS_DESTROYED_COLLECTION_ERROR_MSG = "is already destroyed" +COLLECTION_PATH_NOT_EXIST_ERROR_MSG = "not exist" +NOT_SUPPORT_ADD_COLUMN_ERROR_MSG = "Only support basic numeric data type" +NOT_EXIST_COLUMN_TO_DROP_ERROR_MSG = "Column not exists" diff --git a/python/tests/detail/support_helper.py b/python/tests/detail/support_helper.py new file mode 100644 index 0000000..38d8074 --- /dev/null +++ b/python/tests/detail/support_helper.py @@ -0,0 +1,126 @@ +from zvec import ( + CollectionOption, + IndexOption, + OptimizeOption, + InvertIndexParam, + HnswIndexParam, + IVFIndexParam, + FlatIndexParam, + DataType, + IndexType, + QuantizeType, +) + +SUPPORT_SCALAR_DATA_TYPES = [ + DataType.BOOL, + DataType.FLOAT, + DataType.DOUBLE, + DataType.INT32, + DataType.INT64, + DataType.UINT32, + DataType.UINT64, + DataType.STRING, + DataType.ARRAY_BOOL, + DataType.ARRAY_FLOAT, + DataType.ARRAY_DOUBLE, + DataType.ARRAY_INT32, + DataType.ARRAY_INT64, + DataType.ARRAY_UINT32, + DataType.ARRAY_UINT64, + DataType.ARRAY_STRING, +] + +DEFAULT_SCALAR_FIELD_NAME = { + DataType.BOOL: "bool_field", + DataType.FLOAT: "float_field", + DataType.DOUBLE: "double_field", + DataType.INT32: "int32_field", + DataType.INT64: "int64_field", + DataType.UINT32: "uint32_field", + DataType.UINT64: "uint64_field", + DataType.STRING: "string_field", + DataType.ARRAY_BOOL: "array_bool_field", + DataType.ARRAY_FLOAT: "array_float_field", + DataType.ARRAY_DOUBLE: "array_double_field", + DataType.ARRAY_INT32: "array_int32_field", + DataType.ARRAY_INT64: "array_int64_field", + DataType.ARRAY_UINT32: "array_uint32_field", + DataType.ARRAY_UINT64: "array_uint64_field", + DataType.ARRAY_STRING: "array_string_field", +} + +SUPPORT_SCALAR_INDEX_TYPES = [ + IndexType.INVERT, +] + +SUPPORT_VECTOR_DATA_TYPES = [ + DataType.VECTOR_FP16, + DataType.VECTOR_FP32, + DataType.VECTOR_INT8, + DataType.SPARSE_VECTOR_FP32, + DataType.SPARSE_VECTOR_FP16, +] + +SUPPORT_VECTOR_INDEX_TYPES = [ + IndexType.FLAT, + IndexType.HNSW, + IndexType.IVF, +] + +DEFAULT_VECTOR_FIELD_NAME = { + DataType.VECTOR_FP16: "vector_fp16_field", + DataType.VECTOR_FP32: "vector_fp32_field", + DataType.VECTOR_INT8: "vector_int8_field", + DataType.SPARSE_VECTOR_FP32: "sparse_vector_fp32_field", + DataType.SPARSE_VECTOR_FP16: "sparse_vector_fp16_field", +} + +DEFAULT_VECTOR_DIMENSION = 128 +VECTOR_DIMENSION_1024 = 4 +SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP = { + DataType.VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF], + DataType.VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF], + DataType.VECTOR_INT8: [IndexType.FLAT, IndexType.HNSW], + DataType.SPARSE_VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW], + DataType.SPARSE_VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW], +} + +SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP_PARAMS = [ + (data_type, index_type) + for data_type, index_types in SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP.items() + for index_type in index_types +] + +DEFAULT_INDEX_PARAMS = { + IndexType.FLAT: FlatIndexParam(), + IndexType.HNSW: HnswIndexParam(), + IndexType.IVF: IVFIndexParam(), + IndexType.INVERT: InvertIndexParam(), +} + +SUPPORT_VECTOR_DATA_TYPE_QUANT_MAP = { + DataType.VECTOR_FP32: [QuantizeType.FP16, QuantizeType.INT8, QuantizeType.INT4], + DataType.SPARSE_VECTOR_FP32: [QuantizeType.FP16], +} + +SUPPORT_ADD_COLUMN_DATA_TYPE = [ + DataType.INT32, + DataType.UINT32, + DataType.INT64, + DataType.UINT64, + DataType.FLOAT, + DataType.DOUBLE, +] + +NOT_SUPPORT_ADD_COLUMN_DATA_TYPE = [ + DataType.BOOL, + DataType.STRING, + DataType.ARRAY_BOOL, + DataType.ARRAY_INT32, + DataType.ARRAY_INT64, + DataType.ARRAY_UINT32, + DataType.ARRAY_UINT64, + DataType.ARRAY_FLOAT, + DataType.ARRAY_DOUBLE, + DataType.ARRAY_STRING, +] diff --git a/python/tests/detail/test_collection_concurrency.py b/python/tests/detail/test_collection_concurrency.py new file mode 100644 index 0000000..7a4e774 --- /dev/null +++ b/python/tests/detail/test_collection_concurrency.py @@ -0,0 +1,429 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import pytest +import threading +import numpy as np +import zvec + +from zvec import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + Collection, + Doc, + DataType, + FieldSchema, + VectorSchema, +) + + +class TestCollectionConcurrency: + @pytest.fixture(scope="function") + def test_collection(self, tmp_path_factory): + """Fixture to create a test collection""" + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(), + ), + FieldSchema("weight", DataType.FLOAT, nullable=True), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + VectorSchema( + "sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() + ), + ], + ) + + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + coll = zvec.create_and_open( + path=str(collection_path), + schema=collection_schema, + option=collection_option, + ) + + assert coll is not None, "Failed to create and open collection" + + yield coll + + # Clean up + if hasattr(coll, "destroy") and coll is not None: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + def test_concurrent_read_write(self, test_collection: Collection): + results = [] + + def insert_docs(thread_id): + try: + docs = [ + Doc( + id=f"{thread_id}_{i}", + fields={ + "id": int(f"{thread_id}{i}"), + "name": f"thread_{thread_id}_doc_{i}", + "weight": float(i), + }, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(i), 2: float(i * 2)}, + }, + ) + for i in range(5) + ] + + result = test_collection.insert(docs) + results.append((thread_id, "insert", len(result))) + except Exception as e: + results.append((thread_id, "insert_exception", str(e))) + + def query_docs(thread_id): + try: + result = test_collection.query(filter="id > 0", topk=10) + results.append((thread_id, "query", len(result))) + except Exception as e: + results.append((thread_id, "query_exception", str(e))) + + # Create threads for concurrent operations + threads = [] + + # Start insert threads + for i in range(3): + thread = threading.Thread(target=insert_docs, args=(i,)) + threads.append(thread) + thread.start() + + # Start query threads + for i in range(3): + thread = threading.Thread(target=query_docs, args=(i,)) + threads.append(thread) + thread.start() + + # Wait for all threads to complete + for thread in threads: + thread.join() + + # Analyze results + insert_results = [r for r in results if r[1] == "insert"] + query_results = [r for r in results if r[1] == "query"] + + logging.info( + f"Concurrent read/write results - Inserts: {len(insert_results)}, Queries: {len(query_results)}" + ) + + # At least some operations should succeed + assert len(insert_results) + len(query_results) > 0 + + def test_concurrent_query(self, test_collection: Collection): + # First insert some data + docs = [ + Doc( + id=f"{i}", + fields={"id": i, "name": f"test_{i}", "weight": float(i)}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(i), 2: float(i * 2)}, + }, + ) + for i in range(20) + ] + + insert_result = test_collection.insert(docs) + assert len(insert_result) == 20 + + results = [] + + def query_operation(thread_id): + """Perform query operation from a thread""" + try: + result = test_collection.query(filter=f"id > {thread_id}", topk=5) + results.append((thread_id, "query", len(result))) + except Exception as e: + results.append((thread_id, "query_exception", str(e))) + + # Create multiple threads for concurrent queries + threads = [] + for i in range(5): + thread = threading.Thread(target=query_operation, args=(i,)) + threads.append(thread) + thread.start() + + # Wait for all threads to complete + for thread in threads: + thread.join() + + # Analyze results + query_results = [r for r in results if r[1] == "query"] + logging.info(f"Concurrent query results - Queries: {len(query_results)}") + + # All query operations should succeed + assert len(query_results) == 5 + + def test_concurrent_modifications(self, test_collection: Collection): + # First insert some data + docs = [ + Doc( + id=f"{i}", + fields={"id": i, "name": f"test_{i}", "weight": float(i)}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(i), 2: float(i * 2)}, + }, + ) + for i in range(10) + ] + + insert_result = test_collection.insert(docs) + assert len(insert_result) == 10 + + results = [] + + def update_operation(thread_id): + """Perform update operation from a thread""" + try: + # Each thread updates different documents + update_docs = [ + Doc( + id=f"{i}", + fields={ + "id": i, + "name": f"updated_by_thread_{thread_id}", + "weight": float(i + thread_id), + }, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(i) + 0.5, 2: float(i * 2) + 0.5}, + }, + ) + for i in range(thread_id * 2, thread_id * 2 + 2) + ] + + result = test_collection.update(update_docs) + results.append((thread_id, "update", len(result))) + except Exception as e: + results.append((thread_id, "update_exception", str(e))) + + def delete_operation(thread_id): + """Perform delete operation from a thread""" + try: + # Each thread deletes different documents + delete_ids = [f"{thread_id * 2 + 2}", f"{thread_id * 2 + 3}"] + result = test_collection.delete(delete_ids) + results.append((thread_id, "delete", len(result))) + except Exception as e: + results.append((thread_id, "delete_exception", str(e))) + + # Create threads for concurrent operations + threads = [] + + # Start update threads + for i in range(3): + thread = threading.Thread(target=update_operation, args=(i,)) + threads.append(thread) + thread.start() + + # Start delete threads + for i in range(2): + thread = threading.Thread(target=delete_operation, args=(i,)) + threads.append(thread) + thread.start() + + # Wait for all threads to complete + for thread in threads: + thread.join() + + # Analyze results + update_results = [r for r in results if r[1] == "update"] + delete_results = [r for r in results if r[1] == "delete"] + + logging.info( + f"Concurrent modification results - Updates: {len(update_results)}, Deletes: {len(delete_results)}" + ) + + # At least some operations should succeed + assert len(update_results) + len(delete_results) > 0 + + def test_read_write_locking(self, test_collection: Collection): + # Perform operations that should be thread-safe + docs = [ + Doc( + id=f"{i}", + fields={"id": i, "name": f"test_{i}", "weight": float(i)}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(i), 2: float(i * 2)}, + }, + ) + for i in range(5) + ] + + # Insert data + insert_result = test_collection.insert(docs) + assert len(insert_result) == 5 + + # Concurrent operations should not cause data corruption + results = [] + + def mixed_operation(thread_id): + """Perform mixed operations from a thread""" + try: + # Mix of read and write operations + if thread_id % 2 == 0: + # Read operation + result = test_collection.fetch([f"{thread_id % 5}"]) + results.append((thread_id, "read", len(result))) + else: + # Write operation + doc = Doc( + id=f"{thread_id % 5}", + fields={ + "id": thread_id % 5, + "name": f"mixed_op_{thread_id}", + "weight": float(thread_id), + }, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(thread_id), 2: float(thread_id * 2)}, + }, + ) + result = test_collection.upsert(doc) + results.append((thread_id, "write", len(result))) + except Exception as e: + results.append((thread_id, "exception", str(e))) + + # Create multiple threads + threads = [] + for i in range(10): + thread = threading.Thread(target=mixed_operation, args=(i,)) + threads.append(thread) + thread.start() + + # Wait for all threads to complete + for thread in threads: + thread.join() + + # Verify that the collection is still in a consistent state + final_result = test_collection.query() + assert len(final_result) >= 0 # Should not crash or return corrupted data + + def test_race_condition_detection(self, test_collection: Collection): + # Insert initial data + docs = [ + Doc( + id=f"{i}", + fields={"id": i, "name": f"initial_{i}", "weight": float(i)}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(i), 2: float(i * 2)}, + }, + ) + for i in range(10) + ] + + insert_result = test_collection.insert(docs) + assert len(insert_result) == 10 + + # Perform many rapid concurrent operations + operation_count = 100 + results = [] + + def rapid_operation(op_id): + """Perform rapid operations""" + try: + # Alternate between different types of operations + if op_id % 4 == 0: + # Insert + doc = Doc( + id=f"rapid_{op_id}", + fields={ + "id": op_id, + "name": f"rapid_{op_id}", + "weight": float(op_id), + }, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(op_id), 2: float(op_id * 2)}, + }, + ) + result = test_collection.insert(doc) + results.append(("insert", len(result))) + elif op_id % 4 == 1: + # Update + doc = Doc( + id=f"{op_id % 10}", + fields={ + "id": op_id % 10, + "name": f"rapid_update_{op_id}", + "weight": float(op_id), + }, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: float(op_id), 2: float(op_id * 2)}, + }, + ) + result = test_collection.update(doc) + results.append(("update", len(result))) + elif op_id % 4 == 2: + # Query + result = test_collection.query(filter=f"id > {op_id % 5}", topk=3) + results.append(("query", len(result))) + else: + # Fetch + result = test_collection.fetch([f"{op_id % 10}"]) + results.append(("fetch", len(result))) + except Exception as e: + results.append(("exception", str(e))) + + # Create many threads for rapid concurrent operations + threads = [] + for i in range(operation_count): + thread = threading.Thread(target=rapid_operation, args=(i,)) + threads.append(thread) + thread.start() + + # Wait for all threads to complete + for thread in threads: + thread.join() + + # Verify collection is still functional + final_query = test_collection.query() + assert len(final_query) >= 0 # Should not be corrupted + + logging.info( + f"Rapid concurrent operations completed - Total operations: {len(results)}" + ) diff --git a/python/tests/detail/test_collection_create_and_open.py b/python/tests/detail/test_collection_create_and_open.py new file mode 100644 index 0000000..f746f78 --- /dev/null +++ b/python/tests/detail/test_collection_create_and_open.py @@ -0,0 +1,791 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import threading +import os + +from distance_helper import * +from fixture_helper import * +from doc_helper import * +from params_helper import * + + +def check_collection_info( + coll: Collection, schema: CollectionSchema, option: CollectionOption, path: str +): + assert coll is not None, "Failed to create and open collection" + assert coll.path == path + assert coll.schema.name == schema.name + assert list(coll.schema.fields) == list(schema.fields) + assert list(coll.schema.vectors) == list(schema.vectors) + assert coll.option.read_only == option.read_only + assert coll.option.enable_mmap == option.enable_mmap + + +def check_collection_basic(coll: Collection, optimize: bool = False): + schema = coll.schema + + docs = [generate_doc(i, schema) for i in range(10)] + + results = coll.insert(docs=docs) + assert len(results) == len(docs) + for result in results: + assert result.ok() + + assert coll.stats.doc_count == len(docs) + + def check_fetch_query(): + results = coll.fetch([str(i) for i in range(len(docs))]) + assert len(results) == len(docs) + for i in range(len(docs)): + assert str(i) in results + + results = coll.query() + assert len(results) == len(docs) + + check_fetch_query() + + if optimize: + coll.optimize() + check_fetch_query() + + +def check_collection_full(coll: Collection): + test_doc = generate_doc(1, coll.schema) + + insert_result = coll.insert(test_doc) + assert insert_result.ok() + + stats = coll.stats + assert stats.doc_count == 1 + + fetched_docs = coll.fetch(ids=["1"]) + assert len(fetched_docs) == 1 + assert "1" in fetched_docs + assert fetched_docs["1"] is not None + assert is_doc_equal(fetched_docs["1"], test_doc, coll.schema) + + query_result = coll.query() + assert len(query_result) == 1 + + updated_doc = Doc( + id="1", + fields={"int32_field": 1}, + vectors={"vector_fp32_field": [0.2] * 128}, + ) + update_result = coll.update(updated_doc) + assert update_result.ok() + + upserted_doc = generate_doc(1, coll.schema) + upsert_result = coll.upsert(upserted_doc) + assert upsert_result.ok() + + # 8. Delete document + delete_result = coll.delete("1") + assert delete_result.ok() + + # Verify document was deleted + stats = coll.stats + assert stats.doc_count == 0 + + +valid_collection_options = [ + # (read_only, enable_mmap) + (False, True), + (False, False), +] +invalid_collection_options = [ + # (read_only, enable_mmap) + (True, True), + (True, False), +] +duplicate_names_test = [ + ("field1", "field1", "vector1", "vector2"), + ("field1", "field2", "vector1", "vector1"), + ( + "shared_name1", + "shared_name2", + "shared_name1", + "shared_name2", + ), +] +long_names = [ + "a" * 100, # 100 characters + "b" * 200, # 200 characters +] + +valid_path_list = [ + "/tmp/nonexistent/directory/test_collection", + "test/collection/with/slashes", + "test/collection/with/slashes/哈哈", +] +invalid_path_list = [ + "invalid\0path", + "", +] + + +class TestCreateAndOpen: + @pytest.mark.parametrize("collection_name", COLLECTION_NAME_VALID_LIST) + def test_valid_collection_name( + self, + collection_temp_dir, + collection_name, + collection_option, + sample_field_list, + sample_vector_list, + ): + collection_schema = zvec.CollectionSchema( + name=collection_name, + fields=sample_field_list, + vectors=sample_vector_list, + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + check_collection_info( + coll, collection_schema, collection_option, collection_temp_dir + ) + check_collection_basic(coll) + + coll.destroy() + + @pytest.mark.parametrize("collection_name", COLLECTION_NAME_INVALID_LIST) + def test_invalid_collection_name( + self, + collection_temp_dir, + collection_name, + collection_option, + sample_field_list, + sample_vector_list, + ): + with pytest.raises(Exception) as exc_info: + collection_schema = zvec.CollectionSchema( + name=collection_name, + fields=sample_field_list, + vectors=sample_vector_list, + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + @pytest.mark.parametrize("name_prefix", FIELD_NAME_VALID_LIST) + def test_valid_field_vector_name( + self, + collection_temp_dir, + collection_option, + name_prefix, + sample_field_list, + sample_vector_list, + ): + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=sample_field_list, + vectors=sample_vector_list, + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + check_collection_info( + coll, collection_schema, collection_option, collection_temp_dir + ) + check_collection_basic(coll) + + coll.destroy() + + @pytest.mark.parametrize("field_name", FIELD_NAME_INVALID_LIST) + def test_invalid_field_name( + self, collection_temp_dir, collection_option, field_name + ): + with pytest.raises(Exception) as exc_info: + field_list = [FieldSchema(field_name, DataType.STRING)] + vector_list = [ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ] + + collection_schema = zvec.CollectionSchema( + name="collection_name", fields=field_list, vectors=vector_list + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + @pytest.mark.parametrize("vector_name", FIELD_NAME_INVALID_LIST) + def test_invalid_vector_name( + self, collection_temp_dir, collection_option, vector_name + ): + with pytest.raises(Exception) as exc_info: + field_list = [ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ) + ] + vector_list = [ + VectorSchema(vector_name, DataType.VECTOR_FP32, dimension=128) + ] + + collection_schema = zvec.CollectionSchema( + name="collection_name", fields=field_list, vectors=vector_list + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + @pytest.mark.parametrize( + "field_list_len,vector_list_len,dimension", + FIELD_VECTOR_LIST_DIMENSION_VALID_LIST, + ) + def test_valid_field_vector_size_dimension( + self, + collection_temp_dir, + collection_option, + field_list_len, + vector_list_len, + dimension, + ): + field_list = [] + vector_list = [] + for i in range(0, field_list_len): + field_list.append( + FieldSchema("id_" + str(i), DataType.INT64, nullable=True) + ) + + for i in range(0, vector_list_len): + vector_list.append( + VectorSchema( + "dense_vector_" + str(i), + DataType.VECTOR_FP32, + dimension=dimension, + index_param=HnswIndexParam(), + ) + ) + + collection_schema = zvec.CollectionSchema( + name="test_dense_vector_list", fields=field_list, vectors=vector_list + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + check_collection_info( + coll, collection_schema, collection_option, collection_temp_dir + ) + check_collection_basic(coll) + + coll.destroy() + + @pytest.mark.parametrize( + "field_list_len,vector_list_len,dimension", + FIELD_VECTOR_LIST_DIMENSION_INVALID_LIST, + ) + def test_invalid_field_vector_size_dimension( + self, + collection_temp_dir, + collection_option, + vector_list_len, + field_list_len, + dimension, + ): + with pytest.raises(Exception) as exc_info: + field_list = [] + vector_list = [] + for i in range(0, field_list_len): + field_list.append( + FieldSchema( + "id_" + str(i), + DataType.INT64, + nullable=False, + ) + ) + + for i in range(0, vector_list_len): + vector_list.append( + VectorSchema( + "dense_vector_" + str(i), + DataType.VECTOR_FP32, + dimension=dimension, + index_param=HnswIndexParam(), + ) + ) + + collection_schema = zvec.CollectionSchema( + name="test_dense_vector_list", fields=field_list, vectors=vector_list + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + def test_valid_single_vector_field_construction( + self, collection_temp_dir, collection_option + ): + field = FieldSchema( + "id", + DataType.INT64, + nullable=True, + index_param=InvertIndexParam(enable_range_optimization=True), + ) + + vector = VectorSchema( + "dense_vector", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + + collection_schema = zvec.CollectionSchema( + name="test_single_dense_vector_non_list", + fields=field, + vectors=vector, # Non-list form + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + check_collection_info( + coll, collection_schema, collection_option, collection_temp_dir + ) + check_collection_basic(coll) + coll.destroy() + + def test_collection_concurrent_create( + self, collection_temp_dir, basic_schema, collection_option + ): + results = [] + errors = [] + lock = threading.Lock() + + # Function to be executed by each thread + def create_collection_thread(thread_id): + try: + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=basic_schema, + option=collection_option, + ) + with lock: + results.append((thread_id, coll)) + except Exception as e: + with lock: + errors.append((thread_id, str(e))) + + threads = [] + for i in range(5): + thread = threading.Thread(target=create_collection_thread, args=(i,)) + threads.append(thread) + thread.start() + + for thread in threads: + thread.join() + assert len(results) == 1, ( + f"Expected exactly one successful creation, but got {len(results)}" + ) + assert len(errors) == 4, ( + f"Expected exactly four failures, but got {len(errors)}" + ) + + successful_thread_id, successful_collection = results[0] + assert successful_collection is not None, ( + "Successful creation should return a valid collection" + ) + assert successful_collection.path == collection_temp_dir, ( + "Collection path mismatch" + ) + + def test_create_open_loop( + self, collection_temp_dir, collection_option, full_schema + ): + for cycle in range(10): + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=full_schema, + option=collection_option, + ) + assert coll is not None, ( + f"Failed to create and open collection in cycle {cycle}" + ) + assert coll.path == collection_temp_dir, ( + f"Collection path mismatch in cycle {cycle}" + ) + + del coll + + reopened_coll = zvec.open( + path=collection_temp_dir, option=collection_option + ) + assert reopened_coll is not None, ( + f"Failed to reopen collection in cycle {cycle}" + ) + assert reopened_coll.path == collection_temp_dir, ( + f"Reopened collection path mismatch in cycle {cycle}" + ) + + check_collection_full(reopened_coll) + + reopened_coll.destroy() + + @pytest.mark.parametrize( + "data_type, index_param", VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS + ) + def test_valid_vector_index_params( + self, + data_type, + index_param, + single_vector_schema_with_index_param, + collection_temp_dir, + collection_option, + ): + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=single_vector_schema_with_index_param, + option=collection_option, + ) + + check_collection_info( + coll, + single_vector_schema_with_index_param, + collection_option, + collection_temp_dir, + ) + + check_collection_basic(coll, True) + + @pytest.mark.parametrize( + "data_type, index_param", INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS + ) + def test_invalid_vector_index_params( + self, + data_type, + index_param, + single_vector_schema_with_index_param, + collection_temp_dir, + collection_option, + ): + with pytest.raises(Exception) as exc_info: + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=single_vector_schema_with_index_param, + option=collection_option, + ) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + def test_open_concurrent_same_path(self, tmp_path_factory, collection_option): + """Test concurrent opening of the same collection path. + + - Multi-threading concurrency: 5 threads simultaneously open the same collection + - Result verification: Verify that only one can open successfully, others must fail + """ + # Create a temporary directory and path for the collection + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "concurrent_open_test_collection" + + # First, create a collection that we'll try to open concurrently + field_list = [ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", DataType.STRING, nullable=False, index_param=InvertIndexParam() + ), + ] + + vector_list = [ + VectorSchema( + "dense_vector", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ] + + collection_schema = zvec.CollectionSchema( + name="concurrent_open_test_collection", + fields=field_list, + vectors=vector_list, + ) + + # Create the collection first + coll = zvec.create_and_open( + path=str(collection_path), + schema=collection_schema, + option=collection_option, + ) + + # Close the collection so we can test opening it + if hasattr(coll, "close") and coll is not None: + coll.close() + + # Shared variables to collect results from threads + results = [] + errors = [] + + # Lock for thread-safe operations + lock = threading.Lock() + # Clean up the created collection reference + del coll + + # Function to be executed by each thread + def open_collection_thread(thread_id): + try: + reopened_coll = zvec.open( + path=str(collection_path), option=collection_option + ) + with lock: + results.append((thread_id, reopened_coll)) + # Clean up the collection if opened successfully + if hasattr(reopened_coll, "close") and reopened_coll is not None: + reopened_coll.close() + except Exception as e: + with lock: + errors.append((thread_id, str(e))) + + # Create and start 5 threads + threads = [] + for i in range(5): + thread = threading.Thread(target=open_collection_thread, args=(i,)) + threads.append(thread) + thread.start() + + # Wait for all threads to complete + for thread in threads: + thread.join() + + # Verify results: + # 1. Only one open should succeed (exactly one collection in results) + # 2. Others should fail (4 errors in errors) + assert len(results) == 1, ( + f"Expected exactly one successful open, but got {len(results)}" + ) + assert len(errors) == 4, ( + f"Expected exactly four failures, but got {len(errors)}" + ) + + # Additional verification: check that the successful open has a valid collection + successful_thread_id, successful_collection = results[0] + assert successful_collection is not None, ( + "Successful open should return a valid collection" + ) + assert successful_collection.path == str(collection_path), ( + "Collection path mismatch" + ) + + @pytest.mark.parametrize("read_only,enable_mmap", valid_collection_options) + def test_valid_option( + self, collection_temp_dir, basic_schema, read_only, enable_mmap + ): + option = CollectionOption(read_only=read_only, enable_mmap=enable_mmap) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=basic_schema, + option=option, + ) + + check_collection_info(coll, basic_schema, option, collection_temp_dir) + check_collection_basic(coll) + + coll.destroy() + + def test_valid_none_option(self, collection_temp_dir, basic_schema): + zvec.create_and_open( + path=collection_temp_dir, + schema=basic_schema, + option=None, + ) + + @pytest.mark.parametrize("read_only,enable_mmap", invalid_collection_options) + def test_invalid_option( + self, collection_temp_dir, basic_schema, read_only, enable_mmap + ): + with pytest.raises(Exception) as exc_info: + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=basic_schema, + option=CollectionOption(read_only=read_only, enable_mmap=enable_mmap), + ) + + assert CREATE_READ_ONLY_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + @pytest.mark.parametrize( + "field_name1,field_name2,vector_name1,vector_name2", + duplicate_names_test, + ) + def test_duplicate_field_names( + self, + collection_temp_dir, + collection_option, + field_name1, + field_name2, + vector_name1, + vector_name2, + ): + with pytest.raises(Exception) as exc_info: + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + field_name1, + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + field_name2, + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + ], + vectors=[ + VectorSchema( + vector_name1, + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + VectorSchema( + vector_name2, + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + ], + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + @pytest.mark.parametrize("long_name", long_names) + def test_invalid_long_field_names( + self, collection_option, collection_temp_dir, long_name + ): + collection_schema = zvec.CollectionSchema( + name=long_name, + fields=[ + FieldSchema( + long_name + "_field", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + ], + vectors=[ + VectorSchema( + long_name + "_vector", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ], + ) + + with pytest.raises(Exception) as exc_info: + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + def test_invalid_empty_fields_and_vectors( + self, collection_temp_dir, collection_option + ): + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=[], # Empty fields + vectors=[], # Empty vectors + ) + + with pytest.raises(Exception) as exc_info: + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=collection_schema, + option=collection_option, + ) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value) + + @pytest.mark.parametrize("valid_path", valid_path_list) + def test_valid_path(self, basic_schema, collection_option, valid_path): + if os.path.exists(valid_path): + import shutil + + shutil.rmtree(valid_path) + + coll = zvec.create_and_open( + path=valid_path, schema=basic_schema, option=collection_option + ) + + check_collection_info(coll, basic_schema, collection_option, valid_path) + + coll.destroy() + + @pytest.mark.parametrize("invalid_path", invalid_path_list) + def test_invalid_path(self, basic_schema, collection_option, invalid_path): + with pytest.raises(Exception) as exc_info: + coll = zvec.create_and_open( + path=invalid_path, schema=basic_schema, option=collection_option + ) + + assert INVALID_PATH_ERROR_MSG in str(exc_info.value), str(exc_info.value) diff --git a/python/tests/detail/test_collection_ddl.py b/python/tests/detail/test_collection_ddl.py new file mode 100644 index 0000000..a2c682f --- /dev/null +++ b/python/tests/detail/test_collection_ddl.py @@ -0,0 +1,1720 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from distance_helper import * +from fixture_helper import * +from doc_helper import * +from params_helper import * + + +class TestDDL: + def test_collection_stats(self, basic_collection: Collection): + assert basic_collection.stats is not None + stats = basic_collection.stats + assert stats.doc_count == 0 + assert len(stats.index_completeness) == 2 + assert stats.index_completeness["dense"] == 1 + assert stats.index_completeness["sparse"] == 1 + + def test_collection_destroy( + self, basic_collection: Collection, collection_temp_dir, collection_option + ): + doc = generate_doc(1, basic_collection.schema) + + result = basic_collection.insert(doc) + assert bool(result) + assert result.ok() + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + basic_collection.destroy() + + with pytest.raises(Exception) as exc_info: + stats = basic_collection.stats + assert ACCESS_DESTROYED_COLLECTION_ERROR_MSG in str(exc_info.value) + + with pytest.raises(Exception) as exc_info: + zvec.open(path=collection_temp_dir, option=collection_option) + assert COLLECTION_PATH_NOT_EXIST_ERROR_MSG in str(exc_info.value) + + def test_collection_flush(self, basic_collection: Collection): + doc = generate_doc(1, basic_collection.schema) + + result = basic_collection.insert(doc) + assert bool(result) + assert result.ok() + + basic_collection.flush() + + fetched_docs = basic_collection.fetch(["1"]) + assert "1" in fetched_docs + assert fetched_docs["1"].id == "1" + + +class TestIndexDDL: + @pytest.mark.parametrize("field_name", DEFAULT_SCALAR_FIELD_NAME.values()) + @pytest.mark.parametrize("index_type", SUPPORT_SCALAR_INDEX_TYPES) + def test_scalar_index_operation( + self, + full_collection: Collection, + field_name: str, + index_type: IndexType, + ): + # INSERT 0~5 Doc + docs = [generate_doc(i, full_collection.schema) for i in range(5)] + + result = full_collection.insert(docs) + assert len(result) == 5 + for item in result: + assert item.ok() + + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 5 + + if field_name in ["bool_field"]: + query_filter = f"{field_name} = true" + elif field_name in ["double_field", "float_field"]: + query_filter = f"{field_name} >= 3.0" + elif field_name in [ + "int32_field", + "int64_field", + "uint32_field", + "uint64_field", + ]: + query_filter = f"{field_name} >= 30" + elif field_name in ["string_field"]: + query_filter = f"{field_name} >= 'test_3'" + elif field_name in ["array_bool_field"]: + query_filter = f"{field_name} contain_any (false)" + elif field_name in ["array_double_field", "array_float_field"]: + query_filter = f"{field_name} contain_any (3.0, 4.0)" + elif field_name in [ + "array_int64_field", + "array_int32_field", + "array_uint64_field", + "array_uint32_field", + ]: + query_filter = f"{field_name} contain_any (3, 4)" + elif field_name == "array_string_field": + query_filter = f"{field_name} contain_any ('test_3', 'test_4')" + else: + assert False, f"Unsupported field type for index creation: {field_name}" + + query_result_before = full_collection.query(filter=query_filter, topk=10) + + if index_type not in DEFAULT_INDEX_PARAMS: + pytest.fail(f"Unsupported index type for index creation: {index_type}") + index_param = DEFAULT_INDEX_PARAMS[index_type] + + full_collection.create_index( + field_name=field_name, index_param=index_param, option=IndexOption() + ) + stats_after_create = full_collection.stats + assert stats_after_create is not None + assert stats_after_create.doc_count == 5 + + query_result_after = full_collection.query(filter=query_filter, topk=10) + + assert len(query_result_before) == len(query_result_after), ( + f"Query result count mismatch for {field_name} with index type {index_type}: before={len(query_result_before)}, after={len(query_result_after)}" + ) + + before_ids = set(doc.id for doc in query_result_before) + after_ids = set(doc.id for doc in query_result_after) + assert before_ids == after_ids, ( + f"Query result IDs mismatch for {field_name} with index type {index_type}: before={before_ids}, after={after_ids}" + ) + + # INSERT 5~8 Doc + new_docs = [generate_doc(i, full_collection.schema) for i in range(5, 8)] + + result = full_collection.insert(new_docs) + assert len(result) == 3 + for item in result: + assert item.ok() + + stats_after_insert1 = full_collection.stats + assert stats_after_insert1 is not None + assert stats_after_insert1.doc_count == 8 + + fetched_docs = full_collection.fetch([f"{i}" for i in range(5, 8)]) + assert len(fetched_docs) == 3 + + for i in range(5, 8): + doc_id = f"{i}" + assert doc_id in fetched_docs + + query_result = full_collection.query(filter=query_filter, topk=20) + assert len(query_result) >= len(query_result_before) + + full_collection.drop_index(field_name=field_name) + + # Insert 8~10 Doc + more_docs = [generate_doc(i, full_collection.schema) for i in range(8, 10)] + + result = full_collection.insert(more_docs) + assert len(result) == 2 + for item in result: + assert item.ok() + + stats_after_insert2 = full_collection.stats + assert stats_after_insert2 is not None + assert stats_after_insert2.doc_count == 10 + + fetched_docs = full_collection.fetch([f"{i}" for i in range(8, 10)]) + assert len(fetched_docs) == 2 + + for i in range(8, 10): + doc_id = f"{i}" + assert doc_id in fetched_docs + + query_result = full_collection.query(filter=query_filter, topk=20) + assert len(query_result) >= len(query_result_before) + + final_stats = full_collection.stats + assert final_stats is not None + assert final_stats.doc_count == 10 + full_collection.destroy() + + @pytest.mark.parametrize("field_name", DEFAULT_SCALAR_FIELD_NAME.values()) + @pytest.mark.parametrize("index_type", SUPPORT_SCALAR_INDEX_TYPES) + def test_duplicate_create_index( + self, full_collection: Collection, field_name: str, index_type: IndexType + ): + docs = [generate_doc(i, full_collection.schema) for i in range(10)] + + result = full_collection.insert(docs) + assert bool(result) + for item in result: + assert item.ok() + + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 10 + + if field_name in ["bool_field"]: + query_filter = f"{field_name} = true" + elif field_name in ["double_field", "float_field"]: + query_filter = f"{field_name} >= 3.0" + elif field_name in [ + "int32_field", + "int64_field", + "uint32_field", + "uint64_field", + ]: + query_filter = f"{field_name} >= 30" + elif field_name in ["string_field"]: + query_filter = f"{field_name} >= 'test_3'" + elif field_name in ["array_bool_field"]: + query_filter = f"{field_name} contain_any (false)" + elif field_name in ["array_double_field", "array_float_field"]: + query_filter = f"{field_name} contain_any (3.0, 4.0)" + elif field_name in [ + "array_int64_field", + "array_int32_field", + "array_uint64_field", + "array_uint32_field", + ]: + query_filter = f"{field_name} contain_any (3, 4)" + elif field_name == "array_string_field": + query_filter = f"{field_name} contain_any ('test_3', 'test_4')" + else: + assert False, f"Unsupported field type for index creation: {field_name}" + + query_result_before = full_collection.query(filter=query_filter, topk=5) + + if index_type not in DEFAULT_INDEX_PARAMS: + pytest.fail(f"Unsupported index type for index creation: {index_type}") + index_param = DEFAULT_INDEX_PARAMS[index_type] + + full_collection.create_index( + field_name=field_name, index_param=index_param, option=IndexOption() + ) + + query_result_after = full_collection.query(filter=query_filter, topk=5) + + assert len(query_result_before) == len(query_result_after), ( + f"Query result count mismatch: before={len(query_result_before)}, after={len(query_result_after)}" + ) + + before_ids = set(doc.id for doc in query_result_before) + after_ids = set(doc.id for doc in query_result_after) + assert before_ids == after_ids, ( + f"Query result IDs mismatch: before={before_ids}, after={after_ids}" + ) + + full_collection.create_index( + field_name=field_name, index_param=index_param, option=IndexOption() + ) + + def test_optimize(self, full_collection: Collection): + docs = [generate_doc(i, full_collection.schema) for i in range(10)] + + result = full_collection.insert(docs) + assert bool(result) + for item in result: + assert item.ok() + + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 10 + + full_collection.optimize(option=OptimizeOption()) + + fetched_docs = full_collection.fetch(["1"]) + assert "1" in fetched_docs + assert fetched_docs["1"].id == "1" + + @pytest.mark.parametrize( + "vector_type, index_type", SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP_PARAMS + ) + def test_vector_index_operation( + self, + full_collection: Collection, + vector_type: DataType, + index_type: IndexType, + ): + vector_field_name = DEFAULT_VECTOR_FIELD_NAME[vector_type] + + docs = [generate_doc(i, full_collection.schema) for i in range(5)] + + result = full_collection.insert(docs) + assert len(result) == 5, ( + f"Expected 5 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}" + ) + for i, item in enumerate(result): + assert item.ok(), ( + f"Before create_index,result={result},Insertion result {i} is not OK for vector type {vector_type} and index type {index_type} and result={result}" + ) + + stats = full_collection.stats + assert stats is not None, ( + f"stats is None for vector type {vector_type} and index type {index_type}" + ) + assert stats.doc_count == 5, ( + f"doc_count!=5 for vector type {vector_type} and index type {index_type}" + ) + + if index_type not in DEFAULT_INDEX_PARAMS: + pytest.fail( + f"Unsupported index type {index_type} for vector type {vector_type} in test_vector_all_data_types_index_create_drop_validation" + ) + index_param = DEFAULT_INDEX_PARAMS[index_type] + + full_collection.create_index( + field_name=vector_field_name, + index_param=index_param, + option=IndexOption(), + ) + + stats_after_create = full_collection.stats + assert stats_after_create is not None, ( + f"stats_after_create_index is None for vector type {vector_type} and index type {index_type}" + ) + + new_docs = [generate_doc(i, full_collection.schema) for i in range(5, 8)] + + result = full_collection.insert(new_docs) + assert len(result) == 3, ( + f"Expected 3 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}" + ) + for i, item in enumerate(result): + assert item.ok(), ( + f"Before drop_index,result={result},BInsertion result {i} is not OK for vector type {vector_type} and index type {index_type} and " + ) + + stats_after_insert1 = full_collection.stats + assert stats_after_insert1 is not None, ( + f"stats_after_insert1 is None for vector type {vector_type} and index type {index_type}" + ) + assert stats_after_insert1.doc_count == 8, ( + f"Expected 8 documents, got {stats_after_insert1.doc_count} for vector type {vector_type} and index type {index_type}" + ) + + fetched_docs = full_collection.fetch([f"{i}" for i in range(5, 8)]) + assert len(fetched_docs) == 3, ( + f"Expected 3 fetched documents, got {len(fetched_docs)} for vector type {vector_type} and index type {index_type}" + ) + + for i in range(5, 8): + doc_id = f"{i}" + assert doc_id in fetched_docs, ( + f"Document ID {doc_id} not found in fetched results for vector type {vector_type} and index type {index_type}" + ) + assert fetched_docs[doc_id].id == doc_id, ( + f"Document {doc_id} has incorrect ID field value for vector type {vector_type} and index type {index_type}" + ) + + full_collection.drop_index(field_name=vector_field_name) + + more_docs = [generate_doc(i, full_collection.schema) for i in range(8, 10)] + result = full_collection.insert(more_docs) + assert len(result) == 2, ( + f"Expected 2 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}" + ) + for i, item in enumerate(result): + assert item.ok(), ( + f"After drop_index,Insertion result {i} is not OK for vector type {vector_type} and index type {index_type} and result={result}" + ) + + # Verify document count after second insertion + stats_after_insert2 = full_collection.stats + assert stats_after_insert2 is not None, ( + f"stats_after_insert2 is None for vector type {vector_type} and index type {index_type}" + ) + assert stats_after_insert2.doc_count == 10, ( + f"Expected 10 documents, got {stats_after_insert2.doc_count} for vector type {vector_type} and index type {index_type}" + ) + + # Fetch data + fetched_docs = full_collection.fetch([f"{i}" for i in range(8, 10)]) + assert len(fetched_docs) == 2, ( + f"Expected 2 fetched documents, got {len(fetched_docs)} for vector type {vector_type} and index type {index_type}" + ) + + # Verify fetched documents have correct data + for i in range(8, 10): + doc_id = f"{i}" + assert doc_id in fetched_docs, ( + f"Document ID {doc_id} not found in fetched results for vector type {vector_type} and index type {index_type}" + ) + assert fetched_docs[doc_id].id == doc_id, ( + f"Document {doc_id} has incorrect ID field value for vector type {vector_type} and index type {index_type}" + ) + + # Final verification + final_stats = full_collection.stats + assert final_stats is not None, ( + f"final_stats is None for vector type {vector_type} and index type {index_type}" + ) + assert final_stats.doc_count == 10, ( + f"Expected 10 documents, got {final_stats.doc_count} for vector type {vector_type} and index type {index_type}" + ) + full_collection.destroy() + + @staticmethod + def create_collection( + collection_path, collection_option: CollectionOption + ) -> Collection: + schema = CollectionSchema( + name="test_collection_invalid_vector_index", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=True, + index_param=InvertIndexParam(), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + ], + ) + coll = zvec.create_and_open( + path=collection_path, schema=schema, option=collection_option + ) + assert coll is not None, "Failed to create and open collection" + return coll + + @staticmethod + def check_error_message(exc_info, invalid_name): + if type(invalid_name) is str: + assert INDEX_NON_EXISTENT_COLUMN_ERROR_MSG in str(exc_info.value), ( + "Error message is unreasonable: e=" + str(exc_info.value) + ) + else: + assert INCOMPATIBLE_FUNCTION_ERROR_MSG in str(exc_info.value), ( + "Error message is unreasonable: e=" + str(exc_info.value) + ) + + @pytest.mark.parametrize( + "invalid_field_name,invalid_vector_name", + [ + ("", ""), # Empty string + (" ", " "), # Space only + ("v" * 33, "v" * 33), # Too long (33 characters, exceeds 32) + ("vector name", "vector_name"), # Contains space + ("vector@name", "vector@name"), # Contains special character + ("vector/name", "vector/name"), # Contains slash + ("vector\\name", "vector\\name"), # Contains backslash + ("vector.name", "vector.name"), # Contains dot + ("vector$data", "vector$data"), # Contains dollar sign + ("vector+name", "vector+name"), # Contains plus sign + ("vector=name", "vector=name"), # Contains equals sign + (None, None), # None value, + (1, 1), + (1.1, 1.1), + ], + ) + def test_invalid_field_and_vector_name( + self, + collection_temp_dir, + collection_option: CollectionOption, + invalid_field_name: Any, + invalid_vector_name: Any, + ): + coll = self.create_collection(collection_temp_dir, collection_option) + with pytest.raises(Exception) as exc_info: + coll.create_index( + field_name=invalid_vector_name, + index_param=HnswIndexParam(), + option=IndexOption(), + ) + self.check_error_message(exc_info, invalid_vector_name) + with pytest.raises(Exception) as exc_info: + coll.create_index( + field_name=invalid_field_name, + index_param=InvertIndexParam(), + option=IndexOption(), + ) + self.check_error_message(exc_info, invalid_field_name) + coll.destroy() + coll = self.create_collection(collection_temp_dir, collection_option) + with pytest.raises(Exception) as exc_info: + coll.drop_index(field_name=invalid_vector_name) + self.check_error_message(exc_info, invalid_vector_name) + with pytest.raises(Exception) as exc_info: + coll.drop_index(field_name=invalid_field_name) + self.check_error_message(exc_info, invalid_field_name) + coll.destroy() + + @pytest.mark.parametrize( + "field_name,vector_name", + [ + ("2", "3"), + ("col", "co1"), + ("ID", "IM"), + ("name-1", "name2"), + ("Weigt_12", "Weigt_13"), + ("123age", "123agl"), + ], + ) + def test_valid_field_and_vector_name( + self, + collection_temp_dir, + collection_option: CollectionOption, + field_name: str, + vector_name: str, + ): + schema = zvec.CollectionSchema( + name="test_index_names", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema(field_name, DataType.STRING, nullable=True), + ], + vectors=[ + VectorSchema( + vector_name, + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ], + ) + + coll = zvec.create_and_open( + path=collection_temp_dir, schema=schema, option=collection_option + ) + + assert coll is not None, ( + f"Failed to create and open collection with field_name={field_name}, vector_name={vector_name}" + ) + + # Insert some data + docs = [ + Doc( + id=f"{i}", + fields={"id": i, field_name: f"value_{i}"}, + vectors={vector_name: [float(j % 10) for j in range(128)]}, + ) + for i in range(5) + ] + + result = coll.insert(docs) + assert len(result) == 5, ( + f"Expected 5 insertion results, got {len(result)} for field_name={field_name}, vector_name={vector_name}" + ) + for item in result: + assert item.ok(), ( + f"Insertion failed for field_name={field_name}, vector_name={vector_name}: {item}" + ) + + # Create index on field + coll.create_index( + field_name=field_name, + index_param=InvertIndexParam(), + option=IndexOption(), + ) + + # Create index on vector + coll.create_index( + field_name=vector_name, + index_param=HnswIndexParam(), + option=IndexOption(), + ) + + # Verify indexes were created successfully + stats = coll.stats + assert stats is not None, ( + f"Stats is None for field_name={field_name}, vector_name={vector_name}" + ) + + coll.destroy() + + def test_compicated_workflow( + self, + collection_temp_dir, + basic_schema: CollectionSchema, + collection_option: CollectionOption, + ): + """ + Test the complete workflow: + 1. Create collection + 2. Create index + 3. Insert doc + 4. Upsert + 5. Update doc + 6. Fetch doc + 7. Query doc + 8. Drop index + 9. Insert doc + 10. Update doc + 11. Upsert doc + 12. Fetch doc + 13. Query doc + 14. Flush + 15. Destroy + """ + # Step 1: Create collection + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=basic_schema, + option=collection_option, + ) + + assert coll is not None, "Failed to create and open collection" + assert coll.path == collection_temp_dir + assert coll.schema.name == basic_schema.name + assert coll.stats.doc_count == 0 + + # Step 2: Create index + coll.create_index( + field_name="name", index_param=InvertIndexParam(), option=IndexOption() + ) + # Verify index was created + stats = coll.stats + assert stats is not None, "coll.stats is None!" + + # Step 3: Insert doc + doc1 = Doc( + id="1", + fields={"id": 1, "name": "test1", "weight": 80.5}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + + result = coll.insert(doc1) + assert bool(result) + assert result.ok() + assert coll.stats.doc_count == 1 + + # Step 4: Upsert (existing doc) + doc1_updated = Doc( + id="1", + fields={"id": 1, "name": "test1_updated", "weight": 85.0}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.5, 2: 2.5}, + }, + ) + + result = coll.upsert(doc1_updated) + assert bool(result) + assert result.ok() + assert coll.stats.doc_count == 1 + + # Step 5: Update doc + doc2 = Doc( + id="2", + fields={"id": 2, "name": "test2", "weight": 90.0}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 3.0, 2: 4.0}, + }, + ) + + # First insert doc2 + result = coll.insert(doc2) + assert bool(result) + assert result.ok() + assert coll.stats.doc_count == 2 + + # Then update it + doc2_updated = Doc( + id="2", + fields={"id": 2, "name": "test2_updated", "weight": 95.0}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 3.5, 2: 4.5}, + }, + ) + + result = coll.update(doc2_updated) + assert bool(result) + assert result.ok() + assert coll.stats.doc_count == 2 + + # Step 6: Fetch doc + fetched_docs = coll.fetch(["1", "2"]) + assert len(fetched_docs) == 2 + assert "1" in fetched_docs + assert "2" in fetched_docs + assert fetched_docs["1"].field("name") == "test1_updated" + assert fetched_docs["2"].field("name") == "test2_updated" + + # Step 7: Query doc + query_result = coll.query(filter="id >= 1", topk=10) + assert len(query_result) == 2 + + # Step 8: Drop index + coll.drop_index(field_name="name") + + # Step 9: Insert doc + doc3 = Doc( + id="3", + fields={"id": 3, "name": "test3", "weight": 100.0}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 5.0, 2: 6.0}, + }, + ) + + result = coll.insert(doc3) + assert bool(result) + assert result.ok() + assert coll.stats.doc_count == 3 + + # Step 10: Update doc + doc3_updated = Doc( + id="3", + fields={"id": 3, "name": "test3_updated", "weight": 105.0}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 5.5, 2: 6.5}, + }, + ) + + result = coll.update(doc3_updated) + assert bool(result) + assert result.ok() + assert coll.stats.doc_count == 3 + + # Step 11: Upsert doc + doc4 = Doc( + id="4", + fields={"id": 4, "name": "test4", "weight": 110.0}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 7.0, 2: 8.0}, + }, + ) + + result = coll.upsert(doc4) + assert bool(result) + assert result.ok() + assert coll.stats.doc_count == 4 + + # Step 12: Fetch doc + fetched_docs = coll.fetch(["3", "4"]) + assert len(fetched_docs) == 2 + assert "3" in fetched_docs + assert "4" in fetched_docs + assert fetched_docs["3"].field("name") == "test3_updated" + assert fetched_docs["4"].field("name") == "test4" + + # Step 13: Query doc + query_result = coll.query(filter="id >= 3", topk=10) + assert len(query_result) == 2 + + # Step 14: Flush + coll.flush() + + # Verify data is still accessible after flush + fetched_docs = coll.fetch(["1", "2", "3", "4"]) + assert len(fetched_docs) == 4 + + # Step 15: Destroy + coll.destroy() + + @pytest.mark.parametrize( + "data_type, index_param", VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS + ) + def test_vector_index_params( + self, + collection_temp_dir, + collection_option: CollectionOption, + data_type: DataType, + index_param, + single_vector_schema, + ): + vector_name = DEFAULT_VECTOR_FIELD_NAME[data_type] + dimension = DEFAULT_VECTOR_DIMENSION + + coll = zvec.create_and_open( + path=collection_temp_dir, + schema=single_vector_schema, + option=collection_option, + ) + + assert coll is not None, ( + f"Failed to create and open collection, {data_type}, {index_param}" + ) + + docs = {str(i): generate_doc(i, single_vector_schema) for i in range(5)} + result = coll.insert(docs.values()) + assert len(result) == len(docs), ( + f"Expected 5 results, got {len(result)}, {data_type}, {index_param}" + ) + for item in result: + assert item.ok(), f"Insertion failed for, {data_type}, {index_param}" + + def check_result( + label: str, metric_type: MetricType, quantize_type: QuantizeType + ): + query_vector = [1] * dimension + if data_type in [DataType.SPARSE_VECTOR_FP16, DataType.SPARSE_VECTOR_FP32]: + query_vector = {1: 1} + + fetch_result = coll.fetch([str(i) for i in range(len(docs))]) + assert len(fetch_result) == len(docs), ( + f"{label}, Expected 5 fetched docs, got {len(fetch_result)}, {data_type}, {index_param}" + ) + for i in range(len(docs)): + doc_id = str(i) + assert doc_id in fetch_result, ( + f"{label}, Document ID '{doc_id}' not found, {data_type}, {index_param}" + ) + fetched_doc = fetch_result[doc_id] + # Verify doc equal + assert is_doc_equal(fetched_doc, docs[doc_id], single_vector_schema), ( + f"{label}, doc not equal, insert: {docs[doc_id]}, fetched: {fetched_doc}, {data_type}, {index_param}" + ) + + query_result: list[Doc] = coll.query( + Query(field_name=vector_name, vector=query_vector), + include_vector=False, + topk=len(docs), + ) + assert len(query_result) == len(docs), ( + f"{label}, Expected {len(docs)} result, got {len(query_result)}, {data_type}, {index_param}" + ) + inserted_ids = [str(i) for i in range(len(docs))] + queried_ids = [doc.id for doc in query_result] + assert set(inserted_ids) == set(queried_ids), ( + f"{label}, inserted_ids != queried_ids, insert: {inserted_ids}, query: {queried_ids}, {data_type}, {index_param}" + ) + + last_score = None + for i, doc in enumerate(query_result): + # Get the document's vector for comparison + expect_doc = generate_doc(int(doc.id), single_vector_schema) + doc_vector = expect_doc.vector(vector_name) + expected_score = distance( + doc_vector, + query_vector, + metric_type, + data_type, + quantize_type, + ) + print(f"query: {doc}, expect_core: {expected_score}") + if quantize_type is QuantizeType.UNDEFINED: + assert is_float_equal(doc.score, expected_score), ( + f"{label} top{i} pk{doc.id} score {doc.score:6f} expected:{expected_score:6f}, {data_type}, {index_param}" + ) + if last_score is not None: + if metric_type == MetricType.IP: + assert last_score >= doc.score, ( + f"{label}, score not sorted, last_score: {last_score}, current_score: {doc.score}, {data_type}, {index_param}" + ) + else: + assert last_score <= doc.score, ( + f"{label}, score not sorted, last_score: {last_score}, current_score: {doc.score}, {data_type}, {index_param}" + ) + last_score = doc.score + + # default metric_type=IP, quantize_type=None + check_result("pre_create_index", MetricType.IP, QuantizeType.UNDEFINED) + + # create index + coll.create_index( + field_name=vector_name, + index_param=index_param, + option=IndexOption(), + ) + check_result( + "post_create_index", index_param.metric_type, index_param.quantize_type + ) + + coll.drop_index(field_name=vector_name) + check_result("post_drop_index", MetricType.IP, QuantizeType.UNDEFINED) + + new_docs = {str(i): generate_doc(i, single_vector_schema) for i in range(5, 8)} + new_result = coll.insert(new_docs.values()) + assert len(new_result) == len(new_docs), ( + f"Expected {len(new_docs)} insertion results for new docs, got {len(new_result)} for vector {vector_name}" + ) + for item in new_result: + assert item.ok(), ( + f"New document insertion failed for vector {vector_name}: {item}" + ) + docs |= new_docs + coll.create_index( + field_name=vector_name, + index_param=index_param, + option=IndexOption(), + ) + + check_result( + "post_create_index2", index_param.metric_type, index_param.quantize_type + ) + coll.destroy() + + +class TestColumnDDL: + def test_add_column(self, basic_collection: Collection): + basic_collection.add_column( + field_schema=FieldSchema("income", DataType.INT32), + expression="'weight' * 2", # Simple expression + ) + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, "income": 1}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_add_column_with_default_option(self, basic_collection: Collection): + # Add a new column with default option + basic_collection.add_column( + field_schema=FieldSchema("test_column_default", DataType.INT32), + expression="100", + option=AddColumnOption(), # Default option + ) + # Verify column was added by inserting data + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, "test_column_default": 1}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + # Verify document was inserted + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + @pytest.mark.parametrize("concurrency", [0, 1, 4, 8]) + def test_add_column_with_various_concurrency_options( + self, basic_collection: Collection, concurrency + ): + field_name = f"test_column_concurrent_{concurrency}" + basic_collection.add_column( + field_schema=FieldSchema(field_name, DataType.INT32), + expression="100", + option=AddColumnOption(concurrency=concurrency), + ) + + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, field_name: 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + @pytest.mark.parametrize("data_type", SUPPORT_ADD_COLUMN_DATA_TYPE) + def test_add_column_valid_data_types(self, basic_collection: Collection, data_type): + field_name = f"test_field_{data_type.name.lower()}" + + # Add a new column with specific data type + basic_collection.add_column( + field_schema=FieldSchema(field_name, data_type), + expression="1" if data_type != DataType.STRING else "'test'", + ) + + # Verify column was added by inserting data + if data_type == DataType.STRING: + field_value = "test_value" + elif data_type in [DataType.ARRAY_STRING]: + field_value = ["test_value"] + elif data_type in [DataType.ARRAY_INT32, DataType.ARRAY_INT64]: + field_value = [1, 2, 3] + elif data_type in [DataType.ARRAY_FLOAT, DataType.ARRAY_DOUBLE]: + field_value = [1.1, 2.2, 3.3] + elif data_type == DataType.ARRAY_BOOL: + field_value = [True, False] + elif data_type in [DataType.FLOAT, DataType.DOUBLE]: + field_value = 1.5 + elif data_type in [DataType.INT32, DataType.INT64]: + field_value = 100 + elif data_type == DataType.BOOL: + field_value = True + else: + field_value = 1 + + doc = Doc( + id="1", + fields={ + "id": 1, + "name": "test", + "weight": 80.5, + field_name: field_value, + }, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + # Verify document was inserted + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + @pytest.mark.parametrize("data_type", NOT_SUPPORT_ADD_COLUMN_DATA_TYPE) + def test_add_column_invalid_data_types( + self, basic_collection: Collection, data_type + ): + with pytest.raises(Exception) as exc_info: + field_name = f"test_field_{data_type.name.lower()}" + + # Add a new column with specific data type + basic_collection.add_column( + field_schema=FieldSchema(field_name, data_type), + expression="1" if data_type != DataType.STRING else "'test'", + ) + + assert NOT_SUPPORT_ADD_COLUMN_ERROR_MSG in str(exc_info.value) + + @pytest.mark.parametrize("nullable", [True, False]) + def test_add_column_with_nullable_options( + self, basic_collection: Collection, nullable + ): + field_name = f"test_field_nullable_{str(nullable).lower()}" + + # Add a new column with specific nullable option + basic_collection.add_column( + field_schema=FieldSchema(field_name, DataType.INT32, nullable=nullable), + expression="100", + ) + + # Verify column was added by inserting data + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, field_name: 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + # Verify document was inserted + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + # Verify column was added by inserting data + doc = Doc( + id="2", + fields={"id": 2, "name": "test", "weight": 80.5, field_name: None}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + if nullable: + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + else: + with pytest.raises(ValueError) as e: + basic_collection.insert(doc) + assert ( + "Field 'test_field_nullable_false': expected non-nullable type" + in str(e.value) + ) + + # Verify document was inserted + stats = basic_collection.stats + assert stats is not None + if nullable: + assert stats.doc_count == 2 + else: + assert stats.doc_count == 1 + + @pytest.mark.parametrize( + "expression", + [ + "1", # Constant integer + "1.5", # Constant float + "'test'", # Constant string + "id", # Reference to existing field + "weight * 2", # Simple arithmetic + "weight + id", # Complex arithmetic + "CASE WHEN weight > 50 THEN 1 ELSE 0 END", # Conditional expression + ], + ) + def test_add_column_with_different_expressions( + self, basic_collection: Collection, expression + ): + field_name = f"test_field_expr_{abs(hash(expression)) % 1000}" + + # Add a new column with specific expression + basic_collection.add_column( + field_schema=FieldSchema(field_name, DataType.INT32), + expression=expression, + ) + + # Verify column was added by inserting data + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, field_name: 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + # Verify document was inserted + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_add_column_with_index_param(self, basic_collection: Collection): + basic_collection.add_column( + field_schema=FieldSchema( + "indexed_field", + DataType.INT32, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + expression="id * 2", + ) + + # Verify column was added by inserting data + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, "indexed_field": 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + # Verify document was inserted + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + @pytest.mark.parametrize( + "field_name", + [ + "a", # Minimum length + "a" * 32, # Maximum length (32 characters) + "valid_field_name_123", # Alphanumeric with underscore + "Valid-Field-Name", # With hyphens + "_underscore_start", # Starting with underscore + "field_name_with_123_numbers", # Numbers in middle + "FIELD_NAME_UPPERCASE", # Uppercase + # "field_with_nums_123_and_hyphens-456", # Complex valid name within limit + ], + ) + def test_add_column_with_valid_field_names( + self, basic_collection: Collection, field_name + ): + basic_collection.add_column( + field_schema=FieldSchema(field_name, DataType.INT32), expression="200" + ) + + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, field_name: 300}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + @pytest.mark.parametrize( + "invalid_field_name", + [ + "", # Empty string + " ", # Space only + "a" * 33, # Too long (33 characters, exceeds 32) + "field name", # Contains space + "field.name", # Contains dot + "field@name", # Contains special character + "field/name", # Contains slash + "field\\name", # Contains backslash + "field$name", # Contains dollar sign + "field+name", # Contains plus sign + "field=name", # Contains equals sign + None, # None value + ], + ) + def test_add_column_with_invalid_field_names( + self, basic_collection: Collection, invalid_field_name + ): + with pytest.raises(Exception) as exc_info: + basic_collection.add_column( + field_schema=FieldSchema(invalid_field_name, DataType.INT32), + expression="100", + ) + + if invalid_field_name is None: + assert "validate failed" in str(exc_info.value), ( + "Error message is unreasonable: e=" + str(exc_info.value) + ) + else: + assert ( + "invalid" in str(exc_info.value).lower() + or "name" in str(exc_info.value).lower() + ) + + def test_alter_column_rename(self, basic_collection: Collection): + basic_collection.alter_column( + old_name="weight", + new_name="mass", + option=AlterColumnOption(), + ) + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "mass": 80.5}, # Use new name + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_alter_column_non_exist(self, basic_collection: Collection): + with pytest.raises(Exception) as exc_info: + basic_collection.alter_column( + old_name="non_existing", + new_name="new_name", + field_schema=FieldSchema("new_name", DataType.STRING), + ) + assert "column non_existing not found" in str(exc_info.value), ( + "Error message is unreasonable: e=" + str(exc_info.value) + ) + + def test_alter_column_with_default_option(self, basic_collection: Collection): + basic_collection.add_column( + field_schema=FieldSchema("original_field", DataType.INT32), expression="100" + ) + + basic_collection.alter_column( + old_name="original_field", + new_name="renamed_field", + option=AlterColumnOption(), + ) + + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, "renamed_field": 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + @pytest.mark.parametrize("concurrency", [0, 1, 4, 8]) + def test_alter_column_with_various_concurrency_options( + self, basic_collection: Collection, concurrency + ): + old_field_name = f"orig_field_{concurrency}" + new_field_name = f"modified_field_{concurrency}" + + basic_collection.add_column( + field_schema=FieldSchema(old_field_name, DataType.INT32), + expression="100", + ) + + basic_collection.alter_column( + old_name=old_field_name, + new_name=new_field_name, + option=AlterColumnOption(concurrency=concurrency), + ) + + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, new_field_name: 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + @pytest.mark.parametrize( + "old_field_name,new_field_name", + [ + ("a", "new_a"), # Minimum length + ( + "abcdefghijklmnopqrstuvwxyz123456", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + ), # Maximum length (32 characters) + ("valid_field_name_123", "new_valid_field"), # Alphanumeric with underscore + ("Valid-Field-Name", "New-Field-Name"), # With hyphens + ("_underscore_start", "new_underscore"), # Starting with underscore + ("field_name_with_123_numbers", "new_with_nums"), # Numbers in middle + ("FIELD_NAME_UPPERCASE", "new_uppercase"), # Uppercase + ( + "field_with_nums_3_and_hyphens-6", + "new_field_hyphens", + ), # Complex valid name + ], + ) + def test_alter_column_field_name_valid( + self, basic_collection: Collection, old_field_name, new_field_name + ): + basic_collection.add_column( + field_schema=FieldSchema(old_field_name, DataType.INT32), + expression="100", + ) + basic_collection.alter_column( + old_name=old_field_name, + new_name=new_field_name, + option=AlterColumnOption(), + ) + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, new_field_name: 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + @pytest.mark.parametrize( + "valid_old_name,invalid_new_name", + [ + ("temp_field", ""), # Empty new name + ("temp_field", "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), # Too long new name + ("temp_field", "field name"), # New name with space + ("temp_field", "field.name"), # New name with dot + ("temp_field", "field@name"), # New name with special character + ("temp_field", "field/name"), # New name with slash + ("temp_field", "field\\name"), # New name with backslash + ("temp_field", "field$name"), # New name with dollar sign + ("temp_field", "field+name"), # New name with plus sign + ("temp_field", "field=name"), # New name with equals sign + ("temp_field", None), # None new name + ], + ) + def test_alter_column_with_invalid_field_names( + self, basic_collection: Collection, valid_old_name, invalid_new_name + ): + basic_collection.add_column( + field_schema=FieldSchema("temp_field", DataType.INT32), expression="100" + ) + with pytest.raises(Exception) as exc_info: + basic_collection.alter_column( + old_name=valid_old_name, + new_name=invalid_new_name if invalid_new_name is not None else "", + field_schema=FieldSchema( + invalid_new_name if invalid_new_name is not None else "", + DataType.INT32, + ), + ) + + assert ( + "invalid" in str(exc_info.value).lower() + or "name" in str(exc_info.value).lower() + or "incompatible" in str(exc_info.value).lower() + ) + + def test_drop_column_exist(self, basic_collection: Collection): + basic_collection.add_column( + field_schema=FieldSchema("temp_field", DataType.INT32), expression="100" + ) + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, "temp_field": 1}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + basic_collection.drop_column("temp_field") + doc = Doc( + id="2", + fields={"id": 2, "name": "test", "weight": 80.5, "temp_field": 1}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + with pytest.raises(Exception) as exc_info: + result = basic_collection.insert(doc) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value) + + def test_drop_column_non_exist(self, basic_collection: Collection): + with pytest.raises(Exception) as exc_info: + basic_collection.drop_column("non_existing_column") + assert NOT_EXIST_COLUMN_TO_DROP_ERROR_MSG in str(exc_info.value) + + @pytest.mark.parametrize( + "field_name", + [ + "a", # Minimum length + "a" * 32, # Maximum length (32 characters) + "valid_field_name_123", # Alphanumeric with underscore + "Valid-Field-Name", # With hyphens + "_underscore_start", # Starting with underscore + "field_name_with_123_numbers", # Numbers in middle + "FIELD_NAME_UPPERCASE", # Uppercase + "field_with_nums_3_and_hyphens-6", # Complex valid name within limit + ], + ) + def test_drop_column_field_name_valid( + self, basic_collection: Collection, field_name + ): + basic_collection.add_column( + field_schema=FieldSchema(field_name, DataType.INT32), expression="100" + ) + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5, field_name: 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + + result = basic_collection.insert(doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.ok(), ( + f"result={result},Insert operation failed with code = {result.code()}" + ) + + stats = basic_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + basic_collection.drop_column(field_name) + + doc = Doc( + id="2", + fields={"id": 2, "name": "test", "weight": 80.5, field_name: 200}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + with pytest.raises(Exception) as exc_info: + result = basic_collection.insert(doc) + + assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value) + + def test_add_column_then_query_returns_new_field( + self, basic_collection: Collection + ): + """Regression test for issue #426: query() should return fields added via add_column().""" + basic_collection.add_column( + field_schema=FieldSchema("score", DataType.INT64, nullable=True), + ) + + docs = [ + Doc( + id="1", + fields={"id": 1, "name": "alice", "weight": 60.0, "score": 100}, + vectors={ + "dense": generate_constant_vector(1, 128), + "sparse": generate_sparse_vector(1), + }, + ), + Doc( + id="2", + fields={"id": 2, "name": "bob", "weight": 70.0, "score": 200}, + vectors={ + "dense": generate_constant_vector(2, 128), + "sparse": generate_sparse_vector(2), + }, + ), + ] + result = basic_collection.insert(docs) + assert all(r.ok() for r in result) + + # Query with explicit output_fields + query_result = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + output_fields=["score"], + ) + assert len(query_result) == 2 + for doc in query_result: + assert "score" in doc.fields, ( + f"Doc {doc.id} missing 'score' field after add_column (output_fields explicit)" + ) + assert doc.fields["score"] in (100, 200) + + # Query with select-all (no output_fields) + query_result_all = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + ) + assert len(query_result_all) == 2 + for doc in query_result_all: + assert "score" in doc.fields, ( + f"Doc {doc.id} missing 'score' field after add_column (select all)" + ) + + def test_alter_column_rename_then_query_returns_new_name( + self, basic_collection: Collection + ): + """Regression test: query() should use the new field name after alter_column rename.""" + docs = [ + Doc( + id="1", + fields={"id": 1, "name": "alice", "weight": 60.0}, + vectors={ + "dense": generate_constant_vector(1, 128), + "sparse": generate_sparse_vector(1), + }, + ), + Doc( + id="2", + fields={"id": 2, "name": "bob", "weight": 70.0}, + vectors={ + "dense": generate_constant_vector(2, 128), + "sparse": generate_sparse_vector(2), + }, + ), + ] + result = basic_collection.insert(docs) + assert all(r.ok() for r in result) + + # Rename 'weight' -> 'mass' + basic_collection.alter_column("weight", new_name="mass") + + # Query with explicit output_fields using new name + query_result = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + output_fields=["mass"], + ) + assert len(query_result) == 2 + for doc in query_result: + assert "mass" in doc.fields, ( + f"Doc {doc.id} missing 'mass' field after alter_column rename" + ) + assert "weight" not in doc.fields, ( + f"Doc {doc.id} still has old name 'weight' after rename" + ) + + # Query with select-all + query_result_all = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + ) + assert len(query_result_all) == 2 + for doc in query_result_all: + assert "mass" in doc.fields, ( + f"Doc {doc.id} missing 'mass' in select-all after alter_column rename" + ) + assert "weight" not in doc.fields, ( + f"Doc {doc.id} still has old name 'weight' in select-all after rename" + ) + + def test_drop_column_then_query_excludes_dropped_field( + self, basic_collection: Collection + ): + """Regression test: query() should not return fields removed via drop_column().""" + basic_collection.add_column( + field_schema=FieldSchema("score", DataType.INT64, nullable=True), + ) + + docs = [ + Doc( + id="1", + fields={"id": 1, "name": "alice", "weight": 60.0, "score": 100}, + vectors={ + "dense": generate_constant_vector(1, 128), + "sparse": generate_sparse_vector(1), + }, + ), + Doc( + id="2", + fields={"id": 2, "name": "bob", "weight": 70.0, "score": 200}, + vectors={ + "dense": generate_constant_vector(2, 128), + "sparse": generate_sparse_vector(2), + }, + ), + ] + result = basic_collection.insert(docs) + assert all(r.ok() for r in result) + + # Verify field exists before drop + query_before = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + ) + assert all("score" in doc.fields for doc in query_before) + + # Drop the column + basic_collection.drop_column("score") + + # Query after drop - 'score' should not appear + query_after = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + ) + assert len(query_after) == 2 + for doc in query_after: + assert "score" not in doc.fields, ( + f"Doc {doc.id} still has 'score' after drop_column" + ) + assert "name" in doc.fields, ( + f"Doc {doc.id} missing 'name' - other fields should still be present" + ) diff --git a/python/tests/detail/test_collection_dml.py b/python/tests/detail/test_collection_dml.py new file mode 100644 index 0000000..e1b0c97 --- /dev/null +++ b/python/tests/detail/test_collection_dml.py @@ -0,0 +1,1111 @@ +import logging +import pytest + + +from zvec import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + FieldSchema, + VectorSchema, + CollectionSchema, + Collection, + Doc, + Query, + StatusCode, +) +from distance_helper import * +from fixture_helper import * +from doc_helper import * + +Maximum = 1024 + +DOCID_VALID_LIST = [ + "1valid_Id", + "123.45", + "123abc", + "-!@#$%+=.123abc_+", + "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ123456789012", +] +DOCID_INVALID_LIST = [ + None, + "", + "()qsd123", + " ", + "/&AS12", + "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890121", +] + +FIELD_VALUE_VALID_LIST = [ + ( + "bool_field", + [ + None, + True, + False, + ], + ), + ( + "float_field", + [ + None, + 0.0, + -1.0, + 1.0, + 3.4028235e38, + -3.4028235e38, + 1.17549435e-38, + -1.17549435e-38, + float("inf"), + float("-inf"), + ], + ), + ( + "double_field", + [ + None, + 0.0, + -1.0, + 1.0, + 1.7976931348623157e308, + -1.7976931348623157e308, + 2.2250738585072014e-308, + -2.2250738585072014e-308, + float("inf"), + float("-inf"), + ], + ), + ( + "int32_field", + [ + None, + 0, + 1, + -1, + 2147483647, + -2147483648, + ], + ), + ( + "int64_field", + [ + None, + 0, + 1, + -1, + 9223372036854775807, + -9223372036854775808, + ], + ), + ( + "uint32_field", + [ + None, + 0, + 1, + 4294967295, + ], + ), + ( + "uint64_field", + [ + None, + 0, + 1, + 18446744073709551615, + ], + ), + ( + "string_field", + [ + None, + "", + "a", + "test_name", + "这是一个中文名称测试", + "a" * 1000, + ], + ), + ( + "array_bool_field", + [ + None, + [], + [True], + [False, True], + [True, False, True, False] * 10, + ], + ), + ( + "array_float_field", + [ + None, + [], + [0.0], + [1.0, 2.0, 3.0], + [3.4028235e38, -3.4028235e38], + ], + ), + ( + "array_double_field", + [ + None, + [], + [0.0], + [1.0, 2.0, 3.0], + [1.7976931348623157e308, -1.7976931348623157e308], + ], + ), + ( + "array_int32_field", + [ + None, + [], + [0], + [1, 2, 3], + [2147483647, -2147483648], + ], + ), + ( + "array_int64_field", + [ + None, + [], + [0], + [1, 2, 3], + [9223372036854775807, -9223372036854775808], + ], + ), + ( + "array_uint32_field", + [ + None, + [], + [0], + [1, 2, 3], + [4294967295], + ], + ), + ( + "array_uint64_field", + [ + None, + [], + [0], + [1, 2, 3], + [18446744073709551615], + ], + ), + ( + "array_string_field", + [ + None, + [], + [""], + ["a", "b", "c"], + ["test_string", "测试字符串"], + ["a" * 100] * 5, + ], + ), +] +FIELD_VALUE_INVALID_LIST = [ + ( + "bool_field", + ["True", "False", "", "测试"], + ), + ("float_field", ["invalid", [1.0], {"value": 1.0}, "测试"]), + ("double_field", ["invalid", [1.0], {"value": 1.0}, "测试"]), + ( + "int32_field", + ["invalid", [1], {"value": 1}, 2147483648, -2147483649, "测试"], + ), + ( + "int64_field", + [ + "invalid", + [1], + {"value": 1}, + 9223372036854775808, + -9223372036854775809, + "测试", + ], + ), + ( + "uint32_field", + ["invalid", [1], {"value": 1}, 4294967296, -1, "测试"], + ), + ( + "uint64_field", + ["invalid", [1], {"value": 1}, 18446744073709551616, -1, "测试"], + ), + ( + "string_field", + [ + 123, + 12.34, + True, + ["array"], + {"key": "value"}, + ], + ), + ( + "array_bool_field", + [True, False, [True, "invalid"], {"key": True}, "测试"], + ), + ( + "array_float_field", + [[1.0, "invalid"], [1.0, None], "invalid", [1.0, [2.0]], 1.0, "测试"], + ), + ( + "array_double_field", + [[1.0, "invalid"], [1.0, None], "invalid", [1.0, [2.0]], 1.0, "测试"], + ), + ( + "array_int32_field", + [[1, "invalid"], [1, None], "invalid", [1, [2]], 1, "测试"], + ), + ( + "array_int64_field", + [[1, "invalid"], [1, None], "invalid", [1, [2]], 1, "测试"], + ), + ( + "array_uint32_field", + [[1, "invalid"], [1, None], [1, -1], "invalid", [1, [2]], 1, "测试"], + ), + ( + "array_uint64_field", + [[1, "invalid"], [1, None], [1, -1], "invalid", [1, [2]], 1, "测试"], + ), + ( + "array_string_field", + [["valid", 123], ["valid", None], "invalid", [["nested"]], 123, "测试"], + ), +] + +VECTOR_VALUE_VALID_LIST = [ + ( + "vector_fp32_field", + [ + [0.0] * 128, + [1.0] * 128, + [-1.0] * 128, + [float("inf")] * 128, + [float("-inf")] * 128, + [i / 128.0 for i in range(128)], + [-i / 128.0 for i in range(128)], + ], + ), + ( + "vector_fp16_field", + [ + [0.0] * 128, + [1.0] * 128, + [-1.0] * 128, + [float("inf")] * 128, + [float("-inf")] * 128, + [i / 128.0 for i in range(128)], + [-i / 128.0 for i in range(128)], + ], + ), + ("vector_int8_field", [[100] * 128, [0] * 128, [-100] * 128]), + ( + "sparse_vector_fp32_field", + [ + {0: 1.0}, + {0: 0.0, 1: 1.0, 2: -1.0}, + {0: float("inf"), 1: float("-inf")}, + {i: float(i) for i in range(10)}, + {128: 1.0, 256: -1.0, 512: 0.5}, + ], + ), + ( + "sparse_vector_fp16_field", + [ + {0: 1.0}, + {0: 0.0, 1: 1.0, 2: -1.0}, + {0: float("inf"), 1: float("-inf")}, + {i: float(i) for i in range(10)}, + {128: 1.0, 256: -1.0, 512: 0.5}, + ], + ), +] +VECTOR_VALUE_INVALID_LIST = [ + ( + "vector_fp32_field", + [ + None, + [], + [0.0] * 127, + [0.0] * 129, + [0.0] * 1000, + ["invalid"], + [0, 1, 2], + [None] * 128, + ], + ), + ( + "vector_fp16_field", + [ + None, + [], + [0.0] * 127, + [0.0] * 129, + [0.0] * 1000, + ["invalid"], + [0, 1, 2], + [None] * 128, + ], + ), + ( + "vector_int8_field", + [ + None, + [], + [1] * 127, + [10] * 129, + [0] * 1000, + ["invalid"], + [0, 1, 2], + [None] * 128, + ], + ), + ( + "sparse_vector_fp32_field", + [ + None, + "invalid", + {None: 1.0}, + {"0": 1.0}, + {0: "invalid"}, + {0: None}, + {-1: 1.0}, + ], + ), + ( + "sparse_vector_fp16_field", + [ + None, + "invalid", + {None: 1.0}, + {"0": 1.0}, + {0: "invalid"}, + {0: None}, + {-1: 1.0}, + ], + ), +] + +UPDATE_PARTIAL_VALUE = [ + ( + "partial_fields", + {"string_field": "partially_updated_test", "float_field": 95.5}, + {}, + ), + ("dense_vector_only", {}, {"vector_fp32_field": [0.3] * 128}), + ("dense_vector_only", {}, {"vector_fp16_field": [0.6] * 128}), + ("dense_vector_only", {}, {"vector_int8_field": [3] * 128}), + ("sparse_vector_only", {}, {"sparse_vector_fp32_field": {1: 2.0, 2: 3.0, 4: 4.0}}), + ( + "sparse_vector_only", + {}, + {"sparse_vector_fp16_field": {10: 2.1, 20: 3.1, 40: 4.1}}, + ), + ( + "fields_and_vectors", + {"string_field": "fully_updated_test", "bool_field": False}, + { + "vector_fp32_field": [0.4] * 128, + "sparse_vector_fp32_field": {1: 3.0, 3: 5.0}, + }, + ), +] + + +# ==================== helper ==================== +def singledoc_and_check( + collection: Collection, insert_doc, operator="insert", is_delete=1 +): + if operator == "insert": + result = collection.insert(insert_doc) + elif operator == "upsert": + result = collection.upsert(insert_doc) + elif operator == "update": + result = collection.update(insert_doc) + else: + logging.error("operator value is error!") + + assert bool(result) + assert result.ok() + + stats = collection.stats + assert stats is not None + assert stats.doc_count == 1 + + fetched_docs = collection.fetch([insert_doc.id]) + assert len(fetched_docs) == 1 + assert insert_doc.id in fetched_docs + + fetched_doc = fetched_docs[insert_doc.id] + + assert is_doc_equal(fetched_doc, insert_doc, collection.schema) + assert hasattr(fetched_doc, "score"), "Document should have a score attribute" + assert fetched_doc.score == 0.0, ( + "Fetch operation should return default score of 0.0" + ) + + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v != {}: + query_result = collection.query( + Query(field_name=v, vector=insert_doc.vectors[v]), + topk=10, + ) + assert len(query_result) > 0, ( + f"Expected at least 1 query result, but got {len(query_result)}" + ) + + found_doc = None + for doc in query_result: + if doc.id == insert_doc.id: + found_doc = doc + break + assert found_doc is not None, ( + f"Inserted document {insert_doc.id} not found in query results" + ) + assert is_doc_equal(found_doc, insert_doc, collection.schema, True, False) + if is_delete == 1: + collection.delete(insert_doc.id) + assert collection.stats.doc_count == 0, "Document should be deleted" + + +def updatedoc_partial_check( + collection, update_doc_partial, update_doc_full, operator="update", is_delete=1 +): + if operator == "upsert": + result = collection.upsert(update_doc_partial) + elif operator == "update": + result = collection.update(update_doc_partial) + else: + logging.error("operator value is error!") + + assert bool(result) + assert result.ok() + + stats = collection.stats + assert stats is not None + assert stats.doc_count == 1 + + fetched_docs = collection.fetch([update_doc_partial.id]) + assert len(fetched_docs) == 1, ( + f"fetched_docs={fetched_docs},Expected 1 fetched document, but got {len(fetched_docs)}" + ) + assert update_doc_partial.id in fetched_docs, ( + f"Expected document ID {update_doc_partial.id} in fetched documents" + ) + + fetched_doc = fetched_docs[update_doc_partial.id] + assert is_doc_equal(fetched_doc, update_doc_full, collection.schema) + assert hasattr(fetched_doc, "score"), "Document should have a score attribute" + assert fetched_doc.score == 0.0, ( + "Fetch operation should return default score of 0.0" + ) + + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v != {}: + query_result = collection.query( + Query(field_name=v, vector=update_doc_full.vectors[v]), + topk=10, + ) + assert len(query_result) > 0, ( + f"Expected at least 1 query result, but got {len(query_result)}" + ) + + found_doc = None + for doc in query_result: + if doc.id == update_doc_partial.id: + found_doc = doc + break + assert found_doc is not None, ( + f"Inserted document {update_doc_partial.id} not found in query results" + ) + assert is_doc_equal( + found_doc, update_doc_full, collection.schema, True, False + ) + if is_delete == 1: + collection.delete(update_doc_partial.id) + assert collection.stats.doc_count == 0, "Document should be deleted" + + +def batchdoc_and_check(collection, multiple_docs, doc_num, operator="insert"): + if operator == "insert": + result = collection.insert(multiple_docs) + elif operator == "upsert": + result = collection.upsert(multiple_docs) + + elif operator == "update": + result = collection.update(multiple_docs) + else: + logging.error("operator value is error!") + + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok(), ( + f"result={result},Insert operation failed with code {item.code()}" + ) + + stats = collection.stats + assert stats is not None, "Collection stats should not be None" + assert stats.doc_count == len(multiple_docs), ( + f"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}" + ) + + doc_ids = [doc.id for doc in multiple_docs] + fetched_docs = collection.fetch(doc_ids) + assert len(fetched_docs) == len(multiple_docs), ( + f"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}" + ) + + for original_doc in multiple_docs: + assert original_doc.id in fetched_docs, ( + f"Expected document ID {original_doc.id} in fetched documents" + ) + fetched_doc = fetched_docs[original_doc.id] + + assert is_doc_equal(fetched_doc, original_doc, collection.schema) + + assert hasattr(fetched_doc, "score"), "Document should have a score attribute" + assert fetched_doc.score == 0.0, ( + "Fetch operation should return default score of 0.0" + ) + + first_doc = multiple_docs[doc_num - 1] + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + query_result = collection.query( + Query(field_name=v, vector=first_doc.vectors[v]), + topk=1024, + ) + assert len(query_result) > 0, ( + f"Expected at least 1 query result, but got {len(query_result)}" + ) + + found_doc = None + + for doc in query_result: + if doc.id == first_doc.id: + found_doc = doc + break + assert found_doc is not None, ( + f"Inserted document {first_doc.id} not found in query results" + ) + + assert is_doc_equal(found_doc, first_doc, collection.schema, True, False) + + +# ==================== Tests ==================== +# ---------------------------- +# Collection Insert Test Case +# ---------------------------- + + +class TestCollectionInsert: + def test_insert(self, full_collection: Collection): + single_doc = generate_doc(1, full_collection.schema) + singledoc_and_check(full_collection, single_doc) + + @pytest.mark.parametrize("doc_num", [1, 5, Maximum]) + def test_insert_batch(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num) + + def test_insert_duplicate(self, full_collection: Collection): + insert_doc = generate_doc(1, full_collection.schema) + + result = full_collection.insert(insert_doc) + assert result.code().value == 0 + assert result.ok() + + # Verify documents were inserted + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + insert_doc_duplicate = full_collection.insert(insert_doc) + assert bool(insert_doc_duplicate) + assert insert_doc_duplicate.code() == StatusCode.ALREADY_EXISTS, ( + f"Second insert operation should fail with ALREADY_EXISTS, but got code {insert_doc_duplicate.code()}" + ) + + stats = full_collection.stats + assert stats is not None, "Collection stats should not be None" + assert stats.doc_count == 1, ( + f"Document count should still be 1 after failed insert, but got {stats.doc_count}" + ) + + @pytest.mark.parametrize("doc_id", DOCID_VALID_LIST) + def test_insert_docid_valid(self, full_collection: Collection, doc_id): + insert_doc = generate_doc_random(doc_id, full_collection.schema) + singledoc_and_check(full_collection, insert_doc) + + @pytest.mark.parametrize("doc_id", DOCID_INVALID_LIST) + def test_insert_docid_invalid(self, full_collection: Collection, doc_id): + insert_doc = generate_doc_random(doc_id, full_collection.schema) + + with pytest.raises(Exception) as exc_info: + full_collection.insert(insert_doc) + + assert exc_info.value is not None + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + @pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_VALID_LIST) + @pytest.mark.parametrize( + "full_schema_new", + [(True, True, HnswIndexParam()), (False, True, HnswIndexParam())], + indirect=True, + ) + def test_insert_fields_valid( + self, full_collection_new: Collection, field_name: str, field_values, request + ): + for i, field_value in enumerate(field_values): + doc_id = str(field_value) if field_name == "id" else str(i) + doc_fields, doc_vectors = generate_vectordict_random( + full_collection_new.schema + ) + full_schema_params = request.getfixturevalue("full_schema_new") + target_field = None + for field in full_schema_params.fields: + if field.name == field_name: + target_field = field + break + doc_fields[field_name] = field_value + insert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors) + if target_field and not target_field.nullable and field_value is None: + with pytest.raises(Exception) as exc_info: + full_collection_new.insert(insert_doc) + assert exc_info.value is not None + else: + singledoc_and_check(full_collection_new, insert_doc) + + @pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_INVALID_LIST) + def test_insert_fields_invalid( + self, full_collection: Collection, field_name: str, field_values + ): + for i, field_value in enumerate(field_values): + doc_id = str(field_value) if field_name == "id" else str(i) + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + doc_fields[field_name] = field_value + insert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors) + + with pytest.raises(Exception) as exc_info: + full_collection.insert(insert_doc) + assert exc_info.value is not None + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + @pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_VALID_LIST) + def test_insert_vector_valid( + self, full_collection: Collection, vector_field: str, vector_values + ): + for i, vector_value in enumerate(vector_values): + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + + doc_vectors[vector_field] = vector_value + + insert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + + singledoc_and_check(full_collection, insert_doc) + + @pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_INVALID_LIST) + def test_insert_vector_invalid( + self, full_collection: Collection, vector_field: str, vector_values + ): + for i, vector_value in enumerate(vector_values): + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + doc_vectors[vector_field] = vector_value + insert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + with pytest.raises(Exception) as exc_info: + full_collection.insert(insert_doc) + + assert exc_info.value is not None + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + +class TestCollectionUpdate: + def test_update(self, full_collection: Collection): + insert_doc = generate_doc(1, full_collection.schema) + singledoc_and_check(full_collection, insert_doc, is_delete=0) + updated_doc = generate_update_doc(1, full_collection.schema) + singledoc_and_check(full_collection, updated_doc, operator="update") + + @pytest.mark.parametrize("doc_num", [1, 5, Maximum]) + def test_update_batch(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num) + multiple_update_docs = [ + generate_update_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check( + full_collection, multiple_update_docs, doc_num, operator="update" + ) + + def test_empty_collection_update(self, full_collection: Collection): + updated_doc = generate_update_doc(1, full_collection.schema) + result = full_collection.update(updated_doc) + assert bool(result), f"Expected 1 result, but got {len(result)}" + assert result.code() == StatusCode.NOT_FOUND, ( + f"Update operation should fail with NOT_FOUND, but got code {result.code()}" + ) + fetched_docs = full_collection.fetch([updated_doc.id]) + assert len(fetched_docs) == 0 + + stats = full_collection.stats + assert stats is not None, "Collection stats should not be None" + assert stats.doc_count == 0, ( + f"Document count should be 0, but got {stats.doc_count}" + ) + + @pytest.mark.parametrize("doc_num", [1, 5, Maximum]) + def test_empty_collection_update_batch(self, full_collection: Collection, doc_num): + multiple_update_docs = [ + generate_update_doc(i, full_collection.schema) for i in range(doc_num) + ] + result = full_collection.update(multiple_update_docs) + assert len(result) == len(multiple_update_docs), ( + f"Expected {len(multiple_update_docs)} results, but got {len(result)}" + ) + for item in result: + assert item.code() == StatusCode.NOT_FOUND, ( + f"Update operation should fail with NOT_FOUND, but got code {item.code()}" + ) + + stats = full_collection.stats + assert stats is not None, "Collection stats should not be None" + assert stats.doc_count == 0, ( + f"Document count should be 0, but got {stats.doc_count}" + ) + + doc_ids = [doc.id for doc in multiple_update_docs] + fetched_docs = full_collection.fetch(doc_ids) + assert len(fetched_docs) == 0 + + @pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_VALID_LIST) + @pytest.mark.parametrize( + "full_schema_new", + [(True, True, HnswIndexParam()), (False, True, HnswIndexParam())], + indirect=True, + ) + def test_update_fields_valid( + self, full_collection_new: Collection, field_name: str, field_values, request + ): + for i, field_value in enumerate(field_values): + insert_doc = generate_doc(i, full_collection_new.schema) + singledoc_and_check(full_collection_new, insert_doc, is_delete=0) + update_doc_fields, update_doc_vectors = generate_vectordict_random( + full_collection_new.schema + ) + full_schema_params = request.getfixturevalue("full_schema_new") + target_field = None + for field in full_schema_params.fields: + if field.name == field_name: + target_field = field + break + update_doc_fields[field_name] = field_value + update_doc = Doc( + id=str(i), fields=update_doc_fields, vectors=update_doc_vectors + ) + if target_field and not target_field.nullable and field_value is None: + with pytest.raises(Exception) as exc_info: + update_doc_fields[field_name] = field_value + full_collection_new.update(update_doc) + assert exc_info.value is not None + full_collection_new.delete(insert_doc.id) + else: + singledoc_and_check( + full_collection_new, update_doc, operator="update", is_delete=1 + ) + + @pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_INVALID_LIST) + def test_update_fields_invalid( + self, full_collection: Collection, field_name: str, field_values + ): + for i, field_value in enumerate(field_values): + insert_doc = generate_doc(i, full_collection.schema) + singledoc_and_check(full_collection, insert_doc, is_delete=0) + update_doc_fields, update_doc_vectors = generate_vectordict_random( + full_collection.schema + ) + update_doc_fields[field_name] = field_value + update_doc = Doc( + id=str(i), fields=update_doc_fields, vectors=update_doc_vectors + ) + + with pytest.raises(Exception) as exc_info: + full_collection.update(update_doc) + + assert exc_info.value is not None + full_collection.delete(insert_doc.id) + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + @pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_VALID_LIST) + def test_update_doc_vector_valid( + self, + full_collection: Collection, + collection_temp_dir, + collection_option, + vector_field: str, + vector_values, + ): + for i, vector_value in enumerate(vector_values): + insert_doc = generate_doc(i, full_collection.schema) + singledoc_and_check(full_collection, insert_doc, is_delete=0) + update_doc_fields, update_doc_vectors = generate_vectordict_random( + full_collection.schema + ) + update_doc_vectors[vector_field] = vector_value + update_doc = Doc( + id=str(i), fields=update_doc_fields, vectors=update_doc_vectors + ) + singledoc_and_check(full_collection, update_doc, operator="update") + + @pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_INVALID_LIST) + def test_update_doc_vector_invalid( + self, + full_collection: Collection, + collection_temp_dir, + collection_option, + vector_field: str, + vector_values, + ): + for i, vector_value in enumerate(vector_values): + insert_doc = generate_doc(i, full_collection.schema) + singledoc_and_check(full_collection, insert_doc, is_delete=0) + update_doc_fields, update_doc_vectors = generate_vectordict_random( + full_collection.schema + ) + update_doc_vectors[vector_field] = vector_value + update_doc = Doc( + id=str(i), fields=update_doc_fields, vectors=update_doc_vectors + ) + with pytest.raises(Exception) as exc_info: + full_collection.update(update_doc) + assert exc_info.value is not None + full_collection.delete(insert_doc.id) + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + @pytest.mark.parametrize( + "update_type, fields_to_update, vectors_to_update", UPDATE_PARTIAL_VALUE + ) + def test_update_partial_fields( + self, + full_collection: Collection, + collection_temp_dir, + collection_option, + update_type: str, + fields_to_update: dict, + vectors_to_update: dict, + doc_id=1, + ): + insert_doc = generate_doc(doc_id, full_collection.schema) + singledoc_and_check(full_collection, insert_doc, is_delete=0) + + update_doc_fields, update_doc_vectors = insert_doc.fields, insert_doc.vectors + for k, v in fields_to_update.items(): + update_doc_fields[k] = v + for k, v in vectors_to_update.items(): + update_doc_vectors[k] = v + + update_doc_full = Doc( + id=str(doc_id), fields=update_doc_fields, vectors=update_doc_vectors + ) + + update_doc_partial = Doc( + id=str(doc_id), fields=fields_to_update, vectors=vectors_to_update + ) + + updatedoc_partial_check( + full_collection, + update_doc_partial, + update_doc_full, + operator="update", + is_delete=1, + ) + + +class TestCollectionUpsert: + def test_new_doc_upsert(self, full_collection: Collection): + single_doc = generate_doc(1, full_collection.schema) + singledoc_and_check(full_collection, single_doc, operator="upsert", is_delete=1) + + @pytest.mark.parametrize("doc_num", [1, 5, Maximum]) + def test_new_doc_upsert_batch(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="upsert") + + def test_existing_doc_upsert(self, full_collection: Collection): + insert_doc = generate_doc(1, full_collection.schema) + singledoc_and_check(full_collection, insert_doc, is_delete=0) + updated_doc = generate_update_doc(1, full_collection.schema) + singledoc_and_check(full_collection, updated_doc, operator="upsert") + + @pytest.mark.parametrize("doc_id", DOCID_VALID_LIST) + def test_upsert_docid_valid(self, full_collection: Collection, doc_id): + upsert_doc = generate_doc_random(doc_id, full_collection.schema) + singledoc_and_check(full_collection, upsert_doc, operator="upsert", is_delete=1) + + @pytest.mark.parametrize("doc_id", DOCID_INVALID_LIST) + def test_upsert_docid_invalid(self, full_collection: Collection, doc_id): + upsert_doc = generate_doc_random(doc_id, full_collection.schema) + + with pytest.raises(Exception) as exc_info: + full_collection.upsert(upsert_doc) + + assert exc_info.value is not None + + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + @pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_VALID_LIST) + @pytest.mark.parametrize( + "full_schema_new", + [(True, True, HnswIndexParam()), (False, True, HnswIndexParam())], + indirect=True, + ) + def test_upsert_fields_valid( + self, full_collection_new: Collection, field_name: str, field_values, request + ): + for i, field_value in enumerate(field_values): + doc_id = str(field_value) if field_name == "id" else str(i) + doc_fields, doc_vectors = generate_vectordict_random( + full_collection_new.schema + ) + + full_schema_params = request.getfixturevalue("full_schema_new") + target_field = None + for field in full_schema_params.fields: + if field.name == field_name: + target_field = field + break + doc_fields[field_name] = field_value + upsert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors) + if target_field and not target_field.nullable and field_value is None: + with pytest.raises(Exception) as exc_info: + full_collection_new.upsert(upsert_doc) + assert exc_info.value is not None + else: + singledoc_and_check( + full_collection_new, upsert_doc, operator="upsert", is_delete=1 + ) + + @pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_INVALID_LIST) + def test_upsert_fields_invalid( + self, full_collection: Collection, field_name: str, field_values + ): + for i, field_value in enumerate(field_values): + doc_id = str(field_value) if field_name == "id" else str(i) + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + doc_fields[field_name] = field_value + upsert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors) + + with pytest.raises(Exception) as exc_info: + full_collection.upsert(upsert_doc) + assert exc_info.value is not None + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + @pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_VALID_LIST) + def test_upsert_vector_valid( + self, full_collection: Collection, vector_field: str, vector_values + ): + for i, vector_value in enumerate(vector_values): + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + + doc_vectors[vector_field] = vector_value + + upsert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + + singledoc_and_check( + full_collection, upsert_doc, operator="upsert", is_delete=1 + ) + + @pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_INVALID_LIST) + def test_upsert_vector_invalid( + self, full_collection: Collection, vector_field: str, vector_values + ): + for i, vector_value in enumerate(vector_values): + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + doc_vectors[vector_field] = vector_value + upsert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + with pytest.raises(Exception) as exc_info: + full_collection.upsert(upsert_doc) + + assert exc_info.value is not None + stats = full_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + +class TestCollectionDelete: + @pytest.mark.parametrize("doc_num", [1, 5, Maximum]) + def test_delete_batch(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + + doc_ids = [doc.id for doc in multiple_docs] + result = full_collection.delete(doc_ids) + assert len(result) == len(doc_ids) + for item in result: + assert item.ok() + + def test_delete_non_exist(self, full_collection: Collection): + result = full_collection.delete("non_existing_id") + assert result.code().value == 1 + assert result.code() == StatusCode.NOT_FOUND + + @pytest.mark.parametrize("doc_num", [5]) + def test_delete_batch_part_non_exist(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + doc_ids = [doc.id for doc in multiple_docs] + doc_ids.extend([str(doc_num), str(doc_num + 1)]) + result = full_collection.delete(doc_ids) + + assert len(result) == len(doc_ids) + for i in range(len(result)): + if i < doc_num: + assert result[i].ok() + else: + assert result[i].code().value == 1 + assert result[i].code() == StatusCode.NOT_FOUND + + @pytest.mark.parametrize("doc_num", [5]) + def test_delete_by_filter(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + + result = full_collection.delete_by_filter("int32_field > 0") + assert result is None + + def test_delete_empty_ids(self, full_collection: Collection): + result = full_collection.delete([]) + assert len(result) == 0 diff --git a/python/tests/detail/test_collection_dql.py b/python/tests/detail/test_collection_dql.py new file mode 100644 index 0000000..a49a446 --- /dev/null +++ b/python/tests/detail/test_collection_dql.py @@ -0,0 +1,1434 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from distance_helper import * +from doc_helper import * +from fixture_helper import * +from params_helper import * +from zvec import StatusCode +from zvec.extension import QwenReRanker, RrfReRanker, WeightedReRanker +from zvec.model import Collection, Doc +from zvec.model.param import ( + CollectionOption, + FlatIndexParam, + HnswIndexParam, + HnswQueryParam, + InvertIndexParam, + IVFIndexParam, + IVFQueryParam, +) +from zvec.model.schema import FieldSchema, VectorSchema +from zvec.typing import DataType, MetricType, QuantizeType, StatusCode + + +# ==================== helper ==================== +def batchdoc_and_check( + collection: Collection, multiple_docs, doc_num, operator="insert" +): + if operator == "insert": + result = collection.insert(multiple_docs) + elif operator == "upsert": + result = collection.upsert(multiple_docs) + + elif operator == "update": + result = collection.update(multiple_docs) + else: + logging.error("operator value is error!") + + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok(), ( + f"result={result},Insert operation failed with code {item.code()}" + ) + + stats = collection.stats + assert stats is not None, "Collection stats should not be None" + assert stats.doc_count == len(multiple_docs), ( + f"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}" + ) + + doc_ids = [doc.id for doc in multiple_docs] + fetched_docs = collection.fetch(doc_ids) + assert len(fetched_docs) == len(multiple_docs), ( + f"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}" + ) + + for original_doc in multiple_docs: + assert original_doc.id in fetched_docs, ( + f"Expected document ID {original_doc.id} in fetched documents" + ) + fetched_doc = fetched_docs[original_doc.id] + + assert is_doc_equal(fetched_doc, original_doc, collection.schema) + + assert hasattr(fetched_doc, "score"), "Document should have a score attribute" + assert fetched_doc.score == 0.0, ( + "Fetch operation should return default score of 0.0" + ) + + first_doc = multiple_docs[doc_num - 1] + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + query_result = collection.query( + Query(field_name=v, vector=first_doc.vectors[v]), + topk=1024, + include_vector=True, + ) + assert len(query_result) > 0, ( + f"Expected at least 1 query result, but got {len(query_result)}" + ) + + found_doc = None + + for doc in query_result: + if doc.id == first_doc.id: + found_doc = doc + break + assert found_doc is not None, ( + f"Inserted document {first_doc.id} not found in query results" + ) + + assert is_doc_equal(found_doc, first_doc, collection.schema) + assert hasattr(found_doc, "score") + assert isinstance(found_doc.score, (int, float)) + + +def batchdoc_and_check_ivf( + collection: Collection, multiple_docs, doc_num, operator="insert" +): + if operator == "insert": + result = collection.insert(multiple_docs) + elif operator == "upsert": + result = collection.upsert(multiple_docs) + + elif operator == "update": + result = collection.update(multiple_docs) + else: + logging.error("operator value is error!") + + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok(), ( + f"result={result},Insert operation failed with code {item.code()}" + ) + + stats = collection.stats + assert stats is not None, "Collection stats should not be None" + assert stats.doc_count == len(multiple_docs), ( + f"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}" + ) + + doc_ids = [doc.id for doc in multiple_docs] + fetched_docs = collection.fetch(doc_ids) + assert len(fetched_docs) == len(multiple_docs), ( + f"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}" + ) + + for original_doc in multiple_docs: + assert original_doc.id in fetched_docs, ( + f"Expected document ID {original_doc.id} in fetched documents" + ) + fetched_doc = fetched_docs[original_doc.id] + + assert is_doc_equal(fetched_doc, original_doc, collection.schema) + + assert hasattr(fetched_doc, "score"), "Document should have a score attribute" + assert fetched_doc.score == 0.0, ( + "Fetch operation should return default score of 0.0" + ) + + first_doc = multiple_docs[doc_num - 1] + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + query_result = collection.query( + Query(field_name=v, vector=first_doc.vectors[v]), + topk=1024, + include_vector=True, + ) + assert len(query_result) > 0, ( + f"Expected at least 1 query result, but got {len(query_result)}" + ) + + found_doc = None + + for doc in query_result: + if doc.id == first_doc.id: + found_doc = doc + break + assert found_doc is not None, ( + f"Inserted document {first_doc.id} not found in query results" + ) + + assert is_doc_equal(found_doc, first_doc, collection.schema) + assert hasattr(found_doc, "score") + assert isinstance(found_doc.score, (int, float)) + + +def single_querydoc_check( + multiple_docs, + query_result, + full_collection: Collection, + is_by_vector=0, + query_vector=None, + data_type=None, + vector_name=None, + metric_type=MetricType.IP, + id_include_vector: bool = False, + is_output_fields=0, +): + for original_doc in multiple_docs: + for doc in query_result: + if doc.id == original_doc.id: + found_doc = doc + if is_output_fields == 0: + assert is_doc_equal( + found_doc, + original_doc, + full_collection.schema, + True, + id_include_vector, + ) + assert hasattr(found_doc, "score") + # assert found_doc.score >= 0.0 + if not id_include_vector: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + assert found_doc.vector(v) == {} + else: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + assert found_doc.vector(v) != {} + if is_by_vector: + prev_score = float("inf") + for i, doc in enumerate(query_result): + doc_vector = full_collection.fetch(doc.id)[doc.id].vector( + vector_name + ) + expected_score = distance( + query_vector, doc_vector, metric_type, data_type, k + ) + if ( + full_collection.schema.vector(vector_name).data_type + != DataType.VECTOR_FP16 + ): + assert abs(doc.score - expected_score) < 0.001, ( + f"{data_type} {vector_name} :Expected score {expected_score:.6f}, but got {doc.score:.6f} for document {doc.id}" + ) + assert doc.score <= prev_score, ( + f"{data_type} {vector_name} :Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}" + ) + prev_score = doc.score + + +def multi_querydoc_check(multiple_docs, query_result, full_collection): + for original_doc in multiple_docs: + for doc in query_result: + if doc.id == original_doc.id: + found_doc = doc + assert is_doc_equal( + found_doc, original_doc, full_collection.schema, False, False + ) + assert hasattr(found_doc, "score"), ( + "Document should have a score attribute" + ) + assert found_doc.score >= 0.0, ( + "Fetch operation should return default score of 0.0" + ) + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + assert found_doc.vector(v) == {} + + +# ==================== Tests ==================== +class TestCollectionFetch: + def test_fetch_non_existing(self, full_collection: Collection): + result = full_collection.fetch(ids=["non_existing_id1", "non_existing_id2"]) + assert len(result) == 0 + + @pytest.mark.parametrize("doc_num", [3]) + def test_fetch_partial_non_existing(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + + fetch_id_list = [doc.id for doc in multiple_docs] + fetch_id_list.append("non_existing_id") + result = full_collection.fetch(ids=fetch_id_list) + + assert len(result) == doc_num + assert "non_existing_id" not in result.keys() + + def test_fetch_empty_ids(self, full_collection: Collection): + result = full_collection.fetch(ids=[]) + assert len(result) == 0, ( + f"Expected 0 results for empty ID list, but got {len(result)}" + ) + + @pytest.mark.parametrize("doc_num", [3]) + def test_fetch_with_output_fields(self, full_collection: Collection, doc_num): + """Test that fetch respects output_fields parameter.""" + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + result = full_collection.insert(multiple_docs) + for item in result: + assert item.ok(), f"Insert failed: {item.code()}" + + doc_id = multiple_docs[0].id + + # Case 1: output_fields=None -> all scalar fields returned + fetched_all = full_collection.fetch(ids=[doc_id], output_fields=None) + assert doc_id in fetched_all + doc_all = fetched_all[doc_id] + assert doc_all is not None + assert doc_all.has_field("int32_field"), ( + "int32_field should be present when output_fields=None" + ) + assert doc_all.has_field("string_field"), ( + "string_field should be present when output_fields=None" + ) + + # Case 2: output_fields=["int32_field"] -> only int32_field returned + fetched_partial = full_collection.fetch( + ids=[doc_id], output_fields=["int32_field"] + ) + assert doc_id in fetched_partial + doc_partial = fetched_partial[doc_id] + assert doc_partial is not None + assert doc_partial.has_field("int32_field"), "int32_field should be present" + assert not doc_partial.has_field("string_field"), ( + 'string_field should not be present when output_fields=["int32_field"]' + ) + assert not doc_partial.has_field("float_field"), ( + 'float_field should not be present when output_fields=["int32_field"]' + ) + + # Case 3: output_fields=[] (empty) -> no scalar fields returned + fetched_empty = full_collection.fetch(ids=[doc_id], output_fields=[]) + assert doc_id in fetched_empty + doc_empty = fetched_empty[doc_id] + assert doc_empty is not None + assert doc_empty.id == doc_id, "pk should still be set" + assert not doc_empty.has_field("int32_field"), ( + "int32_field should not be present when output_fields=[]" + ) + assert not doc_empty.has_field("string_field"), ( + "string_field should not be present when output_fields=[]" + ) + + # Case 4: multiple output_fields + fetched_multi = full_collection.fetch( + ids=[doc_id], output_fields=["int32_field", "float_field"] + ) + assert doc_id in fetched_multi + doc_multi = fetched_multi[doc_id] + assert doc_multi is not None + assert doc_multi.has_field("int32_field") + assert doc_multi.has_field("float_field") + assert not doc_multi.has_field("string_field") + + @pytest.mark.parametrize("doc_num", [3]) + def test_fetch_with_include_vector(self, full_collection: Collection, doc_num): + """Test that fetch respects include_vector parameter.""" + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + result = full_collection.insert(multiple_docs) + for item in result: + assert item.ok(), f"Insert failed: {item.code()}" + + doc_id = multiple_docs[0].id + + # Case 1: include_vector=True (default) -> vector data returned + fetched_with_vec = full_collection.fetch(ids=[doc_id]) + assert doc_id in fetched_with_vec + doc_with_vec = fetched_with_vec[doc_id] + assert doc_with_vec is not None + assert doc_with_vec.has_field("int32_field"), ( + "scalar fields should still be present" + ) + assert doc_with_vec.vector("vector_fp32_field"), ( + "vector should be present when include_vector=True (default)" + ) + + # Case 2: include_vector=False -> no vector data returned + fetched_no_vec = full_collection.fetch(ids=[doc_id], include_vector=False) + assert doc_id in fetched_no_vec + doc_no_vec = fetched_no_vec[doc_id] + assert doc_no_vec is not None + assert doc_no_vec.has_field("int32_field"), ( + "scalar fields should still be present" + ) + assert not doc_no_vec.vector("vector_fp32_field"), ( + "vector should not be present when include_vector=False" + ) + + # Case 3: include_vector=False with output_fields + fetched_combo = full_collection.fetch( + ids=[doc_id], output_fields=["int32_field"], include_vector=False + ) + assert doc_id in fetched_combo + doc_combo = fetched_combo[doc_id] + assert doc_combo is not None + assert doc_combo.has_field("int32_field") + assert not doc_combo.has_field("string_field") + assert not doc_combo.vector("vector_fp32_field"), ( + "vector should not be present when include_vector=False" + ) + + +class TestCollectionQuery: + @pytest.mark.parametrize("doc_num", [5]) + def test_query_with_no_condition(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + query_result = full_collection.query() + assert len(query_result) == doc_num + single_querydoc_check(multiple_docs, query_result, full_collection) + + @pytest.mark.parametrize("doc_num", [10]) + def test_query_with_filter_empty(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + result1 = full_collection.query(filter="") + assert len(result1) == doc_num + single_querydoc_check(multiple_docs, result1, full_collection) + result2 = full_collection.query(filter=None) + assert len(result2) == doc_num + single_querydoc_check(multiple_docs, result2, full_collection) + ids1 = set(doc.id for doc in result1) + ids2 = set(doc.id for doc in result2) + assert ids1 == ids2 + + @pytest.mark.parametrize("field_name", ["int32_field"]) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_with_filter_single_condition( + self, full_collection: Collection, doc_num, field_name + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + filter = field_name + " > 5" + query_result = full_collection.query(filter=filter) + assert len(query_result) == doc_num - 6 + + returned_doc_ids = set() + for doc in query_result: + returned_doc_ids.add(doc.id) + + expected_doc_ids = set(str(i) for i in range(6, doc_num)) + + for doc in query_result: + assert doc.id in expected_doc_ids + assert int(doc.field(field_name)) > 5 + + single_querydoc_check(multiple_docs, query_result, full_collection) + + @pytest.mark.parametrize("field_name", ["int32_field"]) + @pytest.mark.parametrize( + "filter", + [ + "int32_field > 3 and int32_field < 9", + "int32_field >= 5 and int32_field <= 7", + ], + ) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_with_filter_and( + self, full_collection: Collection, doc_num, field_name, filter + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + filter = field_name + " > 3 and " + field_name + " < 9" + query_result = full_collection.query(filter=filter) + if filter == "int32_field > 3 and int32_field < 9": + assert len(query_result) == doc_num - 4 - 1 + expected_doc_ids = set(str(i) for i in range(4, 9)) + + for doc in query_result: + assert doc.id in expected_doc_ids + field_value = int(doc.field(field_name)) + assert field_value > 3 and field_value < 9 + else: + assert len(query_result) == 3 + expected_doc_ids = set(str(i) for i in range(5, 8)) + + for doc in query_result: + assert doc.id in expected_doc_ids + field_value = int(doc.field(field_name)) + assert field_value >= 5 and field_value <= 7 + + single_querydoc_check(multiple_docs, query_result, full_collection) + + @pytest.mark.parametrize("field_name", ["int32_field"]) + @pytest.mark.parametrize( + "filter", + [ + "int32_field < 3 or int32_field > 8", + "int32_field = 3 or int32_field = 7", + "int32_field <= 3 or int32_field >= 8", + ], + ) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_with_filter_or( + self, full_collection: Collection, doc_num, field_name, filter + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + query_result = full_collection.query(filter=filter) + if filter == "int32_field < 3 or int32_field > 8": + assert len(query_result) == 4 + expected_doc_ids = set([str(0), str(1), str(2), str(9)]) + for doc in query_result: + assert doc.id in expected_doc_ids + field_value = int(doc.field(field_name)) + assert field_value < 3 or field_value > 8 + elif filter == "int32_field = 3 or int32_field = 7": + assert len(query_result) == 2 + expected_doc_ids = set([str(3), str(7)]) + for doc in query_result: + assert doc.id in expected_doc_ids + field_value = int(doc.field(field_name)) + assert field_value == 3 or field_value == 7 + else: + assert len(query_result) == 6 + expected_doc_ids = set([str(0), str(1), str(2), str(3), str(8), str(9)]) + for doc in query_result: + assert doc.id in expected_doc_ids + field_value = int(doc.field(field_name)) + assert field_value <= 3 or field_value >= 8 + + single_querydoc_check(multiple_docs, query_result, full_collection) + + @pytest.mark.parametrize("field_names", [("int32_field", "bool_field")]) + @pytest.mark.parametrize( + "filter", + [ + "(int32_field < 3 or int32_field > 8) and bool_field = false", + "(int32_field > 2 and int32_field < 5) or (int32_field > 7 and bool_field = true)", + ], + ) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_with_filter_parentheses( + self, full_collection: Collection, doc_num, field_names, filter + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + query_result = full_collection.query(filter=filter) + if filter == "(int32_field < 3 or int32_field > 8) and bool_field = false": + assert len(query_result) == 2 + expected_doc_ids = set([str(1), str(9)]) + for doc in query_result: + assert doc.id in expected_doc_ids + assert ( + int(doc.field(field_names[0])) < 3 + or int(doc.field(field_names[0])) > 8 + ) and doc.field(field_names[1]) == False + else: + assert len(query_result) == 3 + expected_doc_ids = set([str(3), str(4), str(8)]) + for doc in query_result: + assert doc.id in expected_doc_ids + assert ( + ( + int(doc.field(field_names[0])) > 2 + and int(doc.field(field_names[0])) < 5 + ) + or (doc.field(field_names[0])) > 7 + and doc.field(field_names[1]) == True + ) + single_querydoc_check(multiple_docs, query_result, full_collection) + + @pytest.mark.parametrize( + "filter", + [ + "int32_field >", + "int32_field = 'string'", + "nonexistent_field = 5", + "int32_field > 5 and", + "int32_field > > 5", + ], + ) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_filter_invalid(self, full_collection: Collection, doc_num, filter): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + with pytest.raises(Exception) as exc_info: + full_collection.query(filter=filter) + if filter in ["int32_field = 'string'", "nonexistent_field = 5"]: + assert "Analyze SQL info failed" in str(exc_info.value) + else: + assert "Invalid filter" in str(exc_info.value) + + @pytest.mark.parametrize("field_name", ["int32_field"]) + @pytest.mark.parametrize("topk_value", [1, 5, 10, 50, 100, 500, 1000, 1024]) + def test_query_with_filter_topk_valid( + self, full_collection: Collection, topk_value: int, field_name + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(topk_value) + ] + batchdoc_and_check( + full_collection, multiple_docs, topk_value, operator="insert" + ) + filter = ( + field_name + f" >={topk_value - 1} and " + field_name + f" <={topk_value}" + ) + print("filter:\n") + print(filter) + query_result = full_collection.query(filter=filter, topk=topk_value) + assert len(query_result) == 1 + expected_doc_ids = [str(topk_value - 1)] + + for doc in query_result: + assert doc.id in expected_doc_ids + field_value = int(doc.field(field_name)) + assert field_value >= topk_value - 1 and field_value <= topk_value + single_querydoc_check(multiple_docs, query_result, full_collection) + + @pytest.mark.parametrize("field_name", ["int32_field"]) + @pytest.mark.parametrize("topk_value", [1, 5, 10, 50, 100, 500, 1000, 1024]) + def test_query_without_filter_topk_valid( + self, full_collection: Collection, topk_value: int, field_name + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(topk_value) + ] + batchdoc_and_check( + full_collection, multiple_docs, topk_value, operator="insert" + ) + + query_result = full_collection.query(topk=topk_value) + assert len(query_result) == topk_value + single_querydoc_check(multiple_docs, query_result, full_collection) + + @pytest.mark.parametrize("doc_num", [10]) + def test_query_with_include_vector(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + + query_result = full_collection.query(include_vector=True) + assert len(query_result) > 0 + single_querydoc_check( + multiple_docs, query_result, full_collection, id_include_vector=1 + ) + + @pytest.mark.parametrize("output_fields", [["int32_field", "int64_field"]]) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_with_output_fields( + self, full_collection: Collection, doc_num, output_fields + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + query_result = full_collection.query(output_fields=output_fields) + assert len(query_result) > 0 + for doc in query_result: + field_names = doc.field_names() + assert field_names == output_fields + + @pytest.mark.parametrize( + "filter", + [ + "int32_field >= 10 and int32_field <= 20", + "int32_field = 3 and int32_field = 8", + ], + ) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_empty_result(self, full_collection: Collection, doc_num, filter): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + result = full_collection.query(filter=filter) + assert len(result) == 0 + + @pytest.mark.parametrize( + "full_schema_new", + [(True, True, HnswIndexParam()), (False, True, FlatIndexParam())], + indirect=True, + ) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_by_id( + self, full_collection_new: Collection, doc_num, full_schema_new + ): + multiple_docs = [ + generate_doc(i, full_collection_new.schema) for i in range(doc_num) + ] + batchdoc_and_check( + full_collection_new, multiple_docs, doc_num, operator="insert" + ) + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + query_result = full_collection_new.query(Query(field_name=v, id="1")) + assert len(query_result) > 0 + query_doc = full_collection_new.fetch(ids=["1"]) + query_vector = query_doc["1"].vector(v) + single_querydoc_check( + multiple_docs, + query_result, + full_collection_new, + is_by_vector=1, + query_vector=query_vector, + data_type=k, + vector_name=v, + ) + + @pytest.mark.parametrize("doc_num", [10]) + def test_query_by_id_ivf(self, full_collection_ivf: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection_ivf.schema) for i in range(doc_num) + ] + batchdoc_and_check_ivf( + full_collection_ivf, multiple_docs, doc_num, operator="insert" + ) + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + query_result = full_collection_ivf.query(Query(field_name=v, id="1")) + assert len(query_result) > 0 + query_doc = full_collection_ivf.fetch(ids=["1"]) + query_vector = query_doc["1"].vector(v) + single_querydoc_check( + multiple_docs, + query_result, + full_collection_ivf, + is_by_vector=1, + query_vector=query_vector, + data_type=k, + vector_name=v, + ) + + @pytest.mark.parametrize( + "full_schema_new", + [(True, True, HnswIndexParam()), (False, True, FlatIndexParam())], + indirect=True, + ) + @pytest.mark.parametrize("doc_num", [10]) + @pytest.mark.parametrize("topk", [None, 1024]) + @pytest.mark.parametrize("filter", [None, "int32_field >= 3 and int32_field <= 7"]) + def test_query_by_vector( + self, full_collection_new: Collection, doc_num, full_schema_new, topk, filter + ): + multiple_docs = [ + generate_doc(i, full_collection_new.schema) for i in range(doc_num) + ] + batchdoc_and_check( + full_collection_new, multiple_docs, doc_num, operator="insert" + ) + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + doc_fields, doc_vectors = generate_vectordict_random( + full_collection_new.schema + ) + query_vector = doc_vectors[v] + if topk and filter: + query_result = full_collection_new.query( + Query(field_name=v, vector=query_vector), + filter=filter, + topk=topk, + ) + elif topk and not filter: + query_result = full_collection_new.query( + Query(field_name=v, vector=query_vector), topk=topk + ) + elif not topk and filter: + query_result = full_collection_new.query( + Query(field_name=v, vector=query_vector), + filter=filter, + ) + else: + query_result = full_collection_new.query( + Query(field_name=v, vector=query_vector) + ) + assert len(query_result) > 0, ( + f"Expected at least 1 query result, but got {len(query_result)}" + ) + single_querydoc_check( + multiple_docs, + query_result, + full_collection_new, + is_by_vector=1, + query_vector=query_vector, + data_type=k, + vector_name=v, + ) + + @pytest.mark.parametrize("doc_num", [10]) + def test_query_by_vector_ivf(self, full_collection_ivf: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection_ivf.schema) for i in range(doc_num) + ] + batchdoc_and_check_ivf( + full_collection_ivf, multiple_docs, doc_num, operator="insert" + ) + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + doc_fields, doc_vectors = generate_vectordict_random( + full_collection_ivf.schema + ) + query_vector = doc_vectors[v] + query_result = full_collection_ivf.query( + Query(field_name=v, vector=query_vector), + topk=1024, + ) + assert len(query_result) > 0, ( + f"Expected at least 1 query result, but got {len(query_result)}" + ) + single_querydoc_check( + multiple_docs, + query_result, + full_collection_ivf, + is_by_vector=1, + query_vector=query_vector, + data_type=k, + vector_name=v, + ) + + @pytest.mark.parametrize("doc_num", [10]) + def test_query_multivector_rrf(self, full_collection: Collection, doc_num): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + single_query_results = {} + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + single_query_results[v] = full_collection.query( + Query(field_name=v, vector=doc_vectors[v]) + ) + expected_rrf_scores = calculate_multi_vector_rrf_scores(single_query_results) + multi_query_vectors = [] + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + multi_query_vectors.append(Query(field_name=v, vector=doc_vectors[v])) + + rrf_reranker = RrfReRanker() + multi_query_result = full_collection.query( + multi_query_vectors, + topk=3, + reranker=rrf_reranker, + ) + assert len(multi_query_result) > 0, ( + f"Expected at least 1 result, but got {len(multi_query_result)}" + ) + + multi_querydoc_check(multiple_docs, multi_query_result, full_collection) + + prev_score = float("inf") + for i, doc in enumerate(multi_query_result): + doc_id = doc.id + assert doc_id in expected_rrf_scores, ( + f"Document {doc_id} should be in expected RRF scores" + ) + expected_score = expected_rrf_scores[doc_id] + actual_score = doc.score + assert abs(actual_score - expected_score) < 1e-6, ( + f"RRF score mismatch for document {doc_id}: expected {expected_score}, got {actual_score}" + ) + assert doc.score <= prev_score, ( + f"Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}" + ) + prev_score = doc.score + + @pytest.mark.parametrize( + "weights", + [ + { + "vector_fp32_field": 0.3, + "vector_fp16_field": 0.2, + "vector_int8_field": 0.3, + "sparse_vector_fp32_field": 0.1, + "sparse_vector_fp16_field": 0.1, + } + ], + ) + @pytest.mark.parametrize( + "metric_type", [MetricType.L2, MetricType.IP, MetricType.COSINE] + ) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_multivector_weighted( + self, full_collection: Collection, doc_num, weights, metric_type + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + + # Weights are positional, aligned with the multi_query_vectors order + # (DEFAULT_VECTOR_FIELD_NAME insertion order). Metric normalization is + # automatic from each field's schema. + weights_list = [weights[v] for v in DEFAULT_VECTOR_FIELD_NAME.values()] + weighted_reranker = WeightedReRanker(weights_list) + + single_query_results = {} + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + single_query_results[v] = full_collection.query( + Query(field_name=v, vector=doc_vectors[v]) + ) + expected_weighted_scores = calculate_multi_vector_weighted_scores( + single_query_results, weights, MetricType.IP + ) + + multi_query_vectors = [] + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + multi_query_vectors.append(Query(field_name=v, vector=doc_vectors[v])) + + multi_query_result = full_collection.query( + multi_query_vectors, + topk=3, + reranker=weighted_reranker, + ) + assert len(multi_query_result) > 0, ( + f"Expected at least 1 result, but got {len(multi_query_result)}" + ) + + multi_querydoc_check(multiple_docs, multi_query_result, full_collection) + + prev_score = float("inf") + for i, doc in enumerate(multi_query_result): + doc_id = doc.id + assert doc_id in expected_weighted_scores, ( + f"Document {doc_id} should be in expected scores" + ) + expected_score = expected_weighted_scores[doc_id] + actual_score = doc.score + assert abs(actual_score - expected_score) < 1e-6, ( + f"score mismatch for document {doc_id}: expected {expected_score}, got {actual_score}" + ) + assert doc.score <= prev_score, ( + f"Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}" + ) + prev_score = doc.score + + @pytest.mark.parametrize("topk", [5]) + @pytest.mark.parametrize("doc_num", [10]) + @pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"]) + def test_query_consistency( + self, full_collection: Collection, filter, doc_num, topk + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + results = [] + for i in range(5): + query_result = full_collection.query(filter=filter, topk=topk) + single_querydoc_check(multiple_docs, query_result, full_collection) + + results.append(query_result) + assert len(results) == 5 + expected_count = len(results[0]) + for i, result in enumerate(results): + assert len(result) == expected_count + + expected_ids = set(doc.id for doc in results[0]) + for i, result in enumerate(results): + result_ids = set(doc.id for doc in result) + assert result_ids == expected_ids + + for i, result in enumerate(results): + result_ids = [doc.id for doc in result] + expected_sorted_ids = sorted(result_ids, key=lambda x: int(x)) + assert result_ids == expected_sorted_ids + + @pytest.mark.parametrize("ef", [0, 100, 1024, 2048]) + @pytest.mark.parametrize("doc_num", [10]) + @pytest.mark.parametrize("topk", [1024]) + @pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"]) + @pytest.mark.parametrize( + "full_schema_new", [(True, True, HnswIndexParam())], indirect=True + ) + def test_query_vector_with_HnswQueryParam_valid( + self, + full_collection_new: Collection, + doc_num, + full_schema_new, + topk, + filter, + ef, + ): + multiple_docs = [ + generate_doc(i, full_collection_new.schema) for i in range(doc_num) + ] + batchdoc_and_check( + full_collection_new, multiple_docs, doc_num, operator="insert" + ) + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + doc_fields, doc_vectors = generate_vectordict_random( + full_collection_new.schema + ) + query_vector = doc_vectors[v] + query_result = full_collection_new.query( + Query(field_name=v, vector=query_vector, param=HnswQueryParam(ef=ef)), + filter=filter, + topk=topk, + ) + assert len(query_result) > 0, ( + f"Expected at least 1 query result, but got {len(query_result)}" + ) + single_querydoc_check( + multiple_docs, + query_result, + full_collection_new, + is_by_vector=1, + query_vector=query_vector, + data_type=k, + vector_name=v, + ) + + @pytest.mark.parametrize("ef", [None, "invalid", 10.5]) + @pytest.mark.parametrize("doc_num", [10]) + @pytest.mark.parametrize("topk", [10]) + @pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"]) + def test_query_vector_with_HnswQueryParam_invalid( + self, full_collection: Collection, doc_num, topk, ef, filter + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema) + query_vector = doc_vectors[v] + with pytest.raises(Exception) as exc_info: + full_collection.query( + Query( + field_name=v, vector=query_vector, param=HnswQueryParam(ef=ef) + ), + filter=filter, + topk=topk, + ) + assert INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG in str(exc_info.value) + + @pytest.mark.parametrize("nprobe", [1, 10, 100, 2048]) + @pytest.mark.parametrize("doc_num", [10]) + @pytest.mark.parametrize("topk", [10]) + @pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"]) + @pytest.mark.parametrize( + "full_schema_ivf", [(True, True, IVFIndexParam())], indirect=True + ) + def test_query_vector_with_IVFQueryParam_valid( + self, full_collection_ivf: Collection, nprobe, doc_num, topk, filter + ): + multiple_docs = [ + generate_doc(i, full_collection_ivf.schema) for i in range(doc_num) + ] + batchdoc_and_check_ivf( + full_collection_ivf, multiple_docs, doc_num, operator="insert" + ) + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + doc_fields, doc_vectors = generate_vectordict_random( + full_collection_ivf.schema + ) + if v in ["vector_fp32_field"]: + query_vector = doc_vectors[v] + + query_result = full_collection_ivf.query( + Query( + field_name=v, + vector=query_vector, + param=IVFQueryParam(nprobe=nprobe), + ), + filter=filter, + topk=topk, + ) + assert len(query_result) > 0 + single_querydoc_check( + multiple_docs, + query_result, + full_collection_ivf, + is_by_vector=1, + query_vector=query_vector, + data_type=k, + vector_name=v, + ) + + @pytest.mark.parametrize("nprobe", [None, 10.5]) + @pytest.mark.parametrize("doc_num", [10]) + @pytest.mark.parametrize("topk", [10]) + @pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"]) + def test_query_vector_with_IVFQueryParam_invalid( + self, full_collection_ivf: Collection, nprobe, doc_num, topk, filter + ): + multiple_docs = [ + generate_doc(i, full_collection_ivf.schema) for i in range(doc_num) + ] + batchdoc_and_check_ivf( + full_collection_ivf, multiple_docs, doc_num, operator="insert" + ) + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + doc_fields, doc_vectors = generate_vectordict_random( + full_collection_ivf.schema + ) + if v in ["vector_fp32_field"]: + print("v:\n") + print(v) + query_vector = doc_vectors[v] + with pytest.raises(Exception) as exc_info: + full_collection_ivf.query( + Query( + field_name=v, + vector=query_vector, + param=IVFQueryParam(nprobe=nprobe), + ), + # filter=filter, + topk=topk, + ) + assert INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG in str(exc_info.value) + + @pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"]) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_vector_with_param_invalid( + self, full_collection: Collection, doc_num, filter + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + + with pytest.raises(Exception) as exc_info: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + doc_fields, doc_vectors = generate_vectordict_random( + full_collection.schema + ) + query_vector = doc_vectors[v] + if v in ["vector_fp16_field", "vector_fp32_field"]: + full_collection.query( + Query( + field_name=v, vector=query_vector, param=HnswIndexParam() + ), + filter=filter, + ) + assert INCOMPATIBLE_FUNCTION_ERROR_MSG in str(exc_info.value) + + @pytest.mark.parametrize("doc_num", [10]) + @pytest.mark.parametrize( + "test_case_name,vector_query,expected_error_msg", + [ + ( + "Non-existent vector field name", + lambda ref_dense_vector: Query( + field_name="nonexistent_vector", vector=ref_dense_vector + ), + "Expected exception for non-existent vector field name", + ), + ( + "Invalid vector data type for dense vector (string instead of list)", + lambda ref_dense_vector: Query( + field_name="vector_fp32_field", vector="invalid_vector_data" + ), + "Expected exception for invalid dense vector data type", + ), + ( + "Invalid vector data type for sparse vector (list instead of dict)", + lambda ref_dense_vector: Query( + field_name="sparse_fp32", vector=[1.0, 2.0, 3.0] + ), + "Expected exception for invalid sparse vector data type", + ), + ( + "Empty vector data for dense vector", + lambda ref_dense_vector: Query( + field_name="vector_fp32_field", vector=[] + ), + "Expected exception for empty dense vector data", + ), + ( + "Invalid dimension for dense vector", + lambda ref_dense_vector: Query( + field_name="vector_fp32_field", vector=[1.0, 2.0] + ), # Only 2 dimensions instead of 128 + "Expected exception for invalid dense vector dimension", + ), + ( + "Non-existent document ID for by_id query", + lambda ref_dense_vector: Query( + field_name="vector_fp32_field", id="999" + ), # Non-existent ID + "Expected exception for non-existent document ID", + ), + ( + "Neither vector nor id specified", + lambda ref_dense_vector: Query( + field_name="vector_fp32_field" + ), # Neither vector nor id + "Expected exception for specifying neither vector nor id", + ), + ], + ) + def test_query_vector_with_vectors_invalid( + self, + full_collection: Collection, + doc_num, + test_case_name, + vector_query, + expected_error_msg, + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + ref_doc_result = full_collection.fetch(ids=["5"]) + assert "5" in ref_doc_result + ref_doc = ref_doc_result["5"] + ref_dense_vector = ref_doc.vector("vector_fp32_field") + + with pytest.raises(Exception) as exc_info: + full_collection.query([vector_query(ref_dense_vector)]) + assert exc_info.value is not None, expected_error_msg + + @pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"]) + @pytest.mark.parametrize("doc_num", [10]) + def test_query_invalid_param_incompatible_type( + self, full_collection: Collection, doc_num, filter + ): + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(doc_num) + ] + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + + with pytest.raises(Exception) as exc_info: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + doc_fields, doc_vectors = generate_vectordict_random( + full_collection.schema + ) + query_vector = doc_vectors[v] + full_collection.query( + Query(field_name=v, vector=query_vector), + filter=filter, + param=HnswIndexParam(), + topk=3, + ) + + assert "query() got an unexpected keyword argument 'param'" in str( + exc_info.value + ) + + +class TestRRFScoreCalculation: + class MockDoc: + def __init__(self, id, score=0.0): + self._id = id + self._score = score + + @property + def id(self): + return self._id + + @property + def score(self): + return self._score + + @score.setter + def score(self, score): + self._score = score + + def test_rrf_score_calculation_formula(self): + k = 60 + + assert abs(calculate_rrf_score(0, k) - 1.0 / 61) < 1e-10, ( + "RRF score for rank 0 should be 1/61" + ) + assert abs(calculate_rrf_score(1, k) - 1.0 / 62) < 1e-10, ( + "RRF score for rank 1 should be 1/62" + ) + assert abs(calculate_rrf_score(2, k) - 1.0 / 63) < 1e-10, ( + "RRF score for rank 2 should be 1/63" + ) + assert abs(calculate_rrf_score(10, k) - 1.0 / 71) < 1e-10, ( + "RRF score for rank 10 should be 1/71" + ) + + k = 10 + assert abs(calculate_rrf_score(0, k) - 1.0 / 11) < 1e-10, ( + "RRF score for rank 0 with k=10 should be 1/11" + ) + assert abs(calculate_rrf_score(1, k) - 1.0 / 12) < 1e-10, ( + "RRF score for rank 1 with k=10 should be 1/12" + ) + + def test_multi_vector_rrf_scores(self): + query1_results = [self.MockDoc("1"), self.MockDoc("2"), self.MockDoc("3")] + query2_results = [self.MockDoc("3"), self.MockDoc("1"), self.MockDoc("4")] + query3_results = [self.MockDoc("2"), self.MockDoc("4"), self.MockDoc("5")] + query_results = { + "vector1": query1_results, + "vector2": query2_results, + "vector3": query3_results, + } + rrf_scores = calculate_multi_vector_rrf_scores(query_results, k=60) + + expected_doc1_score = 1.0 / 61 + 1.0 / 62 + assert abs(rrf_scores["1"] - expected_doc1_score) < 1e-10, ( + f"RRF score for doc1 mismatch: expected {expected_doc1_score}, got {rrf_scores['1']}" + ) + expected_doc2_score = 1.0 / 62 + 1.0 / 61 + assert abs(rrf_scores["2"] - expected_doc2_score) < 1e-10, ( + f"RRF score for doc2 mismatch: expected {expected_doc2_score}, got {rrf_scores['2']}" + ) + expected_doc3_score = 1.0 / 63 + 1.0 / 61 + assert abs(rrf_scores["3"] - expected_doc3_score) < 1e-10, ( + f"RRF score for doc3 mismatch: expected {expected_doc3_score}, got {rrf_scores['3']}" + ) + expected_doc4_score = 1.0 / 63 + 1.0 / 62 + assert abs(rrf_scores["4"] - expected_doc4_score) < 1e-10, ( + f"RRF score for doc4 mismatch: expected {expected_doc4_score}, got {rrf_scores['4']}" + ) + + expected_doc5_score = 1.0 / 63 + assert abs(rrf_scores["5"] - expected_doc5_score) < 1e-10, ( + f"RRF score for doc5 mismatch: expected {expected_doc5_score}, got {rrf_scores['5']}" + ) + sorted_scores = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True) + expected_order = ["1", "2", "3", "4", "5"] + actual_order = [item[0] for item in sorted_scores] + assert actual_order == expected_order, ( + f"RRF score ranking mismatch: expected {expected_order}, got {actual_order}" + ) + + +class TestCollectionConcurrencyOperations: + @pytest.mark.parametrize("doc_num", [10]) + def test_concurrent_insert_update_upsert_query( + self, full_collection: Collection, doc_num + ): + import threading + + results = [] + errors = [] + multiple_docs = [ + generate_doc(i, full_collection.schema) for i in range(1000, 1010) + ] + + batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert") + + def insert_operation(thread_id): + try: + multiple_docs = [ + generate_doc(i, full_collection.schema) + for i in range(thread_id, thread_id + 5) + ] + result = full_collection.insert(multiple_docs) + results.append(("insert", thread_id, len(result))) + except Exception as e: + errors.append(("insert", thread_id, str(e))) + + def update_operation(thread_id): + try: + multiple_docs = [ + generate_doc_random(i, full_collection.schema) + for i in range(1000, 1001) + ] + result = full_collection.update(multiple_docs) + results.append(("update", thread_id, len(result))) + except Exception as e: + errors.append(("update", thread_id, str(e))) + + def upsert_operation(thread_id): + try: + multiple_docs = [ + generate_doc(i, full_collection.schema) + for i in range(thread_id, thread_id + 5) + ] + result = full_collection.upsert(multiple_docs) + results.append(("upsert", thread_id, len(result))) + except Exception as e: + errors.append(("upsert", thread_id, str(e))) + + def query_operation(thread_id): + try: + if thread_id % 3 == 0: + result = full_collection.query(filter="int32_field > 1", topk=5) + elif thread_id % 3 == 1: + result = full_collection.query(filter="bool_field = true", topk=3) + else: + query_vector = [0.1] * 128 + result = full_collection.query( + Query(field_name="vector_fp32_field", vector=query_vector), + topk=3, + ) + + results.append(("query", thread_id, len(result))) + except Exception as e: + errors.append(("query", thread_id, str(e))) + + def delete_operation(thread_id): + try: + # Delete some existing documents + delete_ids = ( + [f"{thread_id + 1}", f"{thread_id + 2}"] + if thread_id < 5 + else [f"{thread_id % 5 + 1}"] + ) + result = full_collection.delete(delete_ids) + results.append(("delete", thread_id, len(result))) + except Exception as e: + errors.append(("delete", thread_id, str(e))) + + threads = [] + for i in range(1): + thread = threading.Thread(target=insert_operation, args=(i,)) + threads.append(thread) + thread.start() + for i in range(1): + thread = threading.Thread(target=update_operation, args=(i,)) + threads.append(thread) + thread.start() + for i in range(1): + thread = threading.Thread(target=upsert_operation, args=(i,)) + threads.append(thread) + thread.start() + for i in range(1): + thread = threading.Thread(target=query_operation, args=(i,)) + threads.append(thread) + thread.start() + for i in range(1): + thread = threading.Thread(target=delete_operation, args=(i,)) + threads.append(thread) + thread.start() + + for thread in threads: + thread.join() + + insert_results = [r for r in results if r[0] == "insert"] + update_results = [r for r in results if r[0] == "update"] + upsert_results = [r for r in results if r[0] == "upsert"] + query_results = [r for r in results if r[0] == "query"] + delete_results = [r for r in results if r[0] == "delete"] + + assert ( + len(insert_results) + + len(update_results) + + len(upsert_results) + + len(query_results) + + len(delete_results) + > 0 + ), f"No operations succeeded. Errors: {errors}" + + critical_errors = [ + e for e in errors if "critical" in e[2].lower() or "fatal" in e[2].lower() + ] + assert len(critical_errors) == 0, f"Critical errors occurred: {critical_errors}" diff --git a/python/tests/detail/test_collection_exception.py b/python/tests/detail/test_collection_exception.py new file mode 100644 index 0000000..aab0bca --- /dev/null +++ b/python/tests/detail/test_collection_exception.py @@ -0,0 +1,328 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import logging +import pytest +import numpy as np +import zvec + +from zvec import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + DataType, + Collection, + Doc, + FieldSchema, + Query, + VectorSchema, +) + + +class TestCollectionExceptionHandling: + @pytest.fixture(scope="function") + def test_collection(self, tmp_path_factory): + """Fixture to create a test collection""" + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(), + ), + FieldSchema("weight", DataType.FLOAT, nullable=True), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + VectorSchema( + "sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() + ), + ], + ) + + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + coll = zvec.create_and_open( + path=str(collection_path), + schema=collection_schema, + option=collection_option, + ) + + assert coll is not None, "Failed to create and open collection" + + yield coll + + # Clean up + if hasattr(coll, "destroy") and coll is not None: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + def test_create_and_open_missing_path(self, tmp_path_factory): + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ], + ) + + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + with pytest.raises(Exception) as exc_info: + coll = zvec.create_and_open( + schema=collection_schema, option=collection_option + ) + assert exc_info.value is not None, ( + "Expected exception for missing path parameter" + ) + + def test_create_and_open_missing_schema(self, tmp_path_factory): + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + with pytest.raises(Exception) as exc_info: + coll = zvec.create_and_open( + path=str(collection_path), option=collection_option + ) + assert exc_info.value is not None, ( + "Expected exception for missing schema parameter" + ) + + def test_open_missing_path(self): + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + with pytest.raises(Exception) as exc_info: + coll = zvec.open(option=collection_option) + assert exc_info.value is not None, ( + "Expected exception for missing path parameter" + ) + + def test_insert_missing_docs(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + result = test_collection.insert() + assert exc_info.value is not None, ( + "Expected exception for missing docs parameter" + ) + + def test_update_missing_docs(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + result = test_collection.update() + assert exc_info.value is not None, ( + "Expected exception for missing docs parameter" + ) + + def test_upsert_missing_docs(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + result = test_collection.upsert() + assert exc_info.value is not None, ( + "Expected exception for missing docs parameter" + ) + + def test_delete_missing_ids(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + result = test_collection.delete() + assert exc_info.value is not None, ( + "Expected exception for missing ids parameter" + ) + + def test_fetch_missing_ids(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + result = test_collection.fetch() + assert exc_info.value is not None, ( + "Expected exception for missing ids parameter" + ) + + def test_query_missing_query_field_name(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + result = test_collection.query([Query()]) + assert exc_info.value is not None, ( + "Expected exception for missing Query field_name parameter" + ) + + def test_add_column_missing_field_schema(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + test_collection.add_column() + assert exc_info.value is not None, ( + "Expected exception for missing field_schema parameter" + ) + + def test_alter_column_missing_old_name(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + test_collection.alter_column(new_name="new_name") + assert exc_info.value is not None, ( + "Expected exception for missing old_name parameter" + ) + + def test_alter_column_missing_new_name(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + test_collection.alter_column(old_name="old_name") + assert exc_info.value is not None, ( + "Expected exception for missing new_name parameter" + ) + + def test_drop_column_missing_field_name(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + test_collection.drop_column() + assert exc_info.value is not None, ( + "Expected exception for missing field_name parameter" + ) + + def test_invalid_parameter_types(self, test_collection: Collection): + # This test depends on specific implementation details + # Generally, we would expect TypeErrors or similar exceptions + pass + + def test_missing_required_parameters(self, test_collection: Collection): + # This test depends on specific implementation details + # Generally, we would expect TypeErrors or similar exceptions + pass + + def test_empty_collection_operations(self, tmp_path_factory): + collection_schema = zvec.CollectionSchema( + name="empty_test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ], + ) + + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "empty_test_collection" + + coll = zvec.create_and_open( + path=str(collection_path), + schema=collection_schema, + option=collection_option, + ) + + assert coll is not None, "Failed to create and open collection" + + # Test fetch on empty collection + result = coll.fetch(["1"]) + assert len(result) >= 0 # May be empty or have special handling + + # Test query on empty collection + result = coll.query() + assert len(result) == 0 + + # Test update on empty collection + doc = Doc( + id="1", + fields={"id": 1, "name": "test"}, + vectors={"dense": np.random.random(128).tolist()}, + ) + + result = coll.update(doc) + # Should handle gracefully, possibly with NOT_FOUND status + + # Clean up + if hasattr(coll, "destroy") and coll is not None: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + def test_resource_management(self, test_collection: Collection): + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "weight": 80.5}, + vectors={ + "dense": np.random.random(128).tolist(), + "sparse": {1: 1.0, 2: 2.0}, + }, + ) + + # Insert + result = test_collection.insert(doc) + assert result.ok() + + # Fetch + result = test_collection.fetch(["1"]) + assert len(result) == 1 + + # Query + result = test_collection.query() + assert len(result) >= 0 + + # Update + result = test_collection.update(doc) + assert result.ok() + + # Delete + result = test_collection.delete("1") + assert result.ok() + + def test_exception_resource_cleanup(self, test_collection: Collection): + # This test would need to simulate exception conditions + # which is difficult without specific failure injection points + pass diff --git a/python/tests/detail/test_collection_open.py b/python/tests/detail/test_collection_open.py new file mode 100644 index 0000000..2d6c38a --- /dev/null +++ b/python/tests/detail/test_collection_open.py @@ -0,0 +1,967 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import threading +import numpy as np + +from fixture_helper import * + +COLLECTION_OPTION_TEST_CASES_VALID = [ + # (read_only, enable_mmap, description) + (False, True, "Read-write with mmap enabled"), + (False, False, "Read-write with mmap disabled"), + (True, True, "Read-only with mmap enabled"), + (True, False, "Read-only with mmap disabled"), +] + +# Test data for invalid paths +INVALID_PATH_LIST = [ + "/nonexistent/directory/test_collection", + "invalid:path", + "", # Empty path +] + + +@pytest.fixture(scope="session") +def collection_schema(): + return zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", DataType.STRING, nullable=False, index_param=InvertIndexParam() + ), + FieldSchema( + "weight", DataType.FLOAT, nullable=False, index_param=InvertIndexParam() + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + VectorSchema( + "sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() + ), + ], + ) + + +@pytest.fixture +def single_doc(): + id = 0 + return Doc( + id=f"{id}", + fields={"id": id, "name": "test"}, + vectors={ + "dense": [id + 0.1] * 128, + }, + ) + + +@pytest.fixture(scope="function") +def test_collection( + tmp_path_factory, collection_schema, collection_option +) -> Generator[Any, Any, Collection]: + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + coll = zvec.create_and_open( + path=str(collection_path), schema=collection_schema, option=collection_option + ) + + assert coll is not None, "Failed to create and open collection" + assert coll.path == str(collection_path) + assert coll.schema.name == collection_schema.name + assert list(coll.schema.fields) == list(collection_schema.fields) + assert list(coll.schema.vectors) == list(collection_schema.vectors) + assert coll.option.read_only == collection_option.read_only + assert coll.option.enable_mmap == collection_option.enable_mmap + + try: + yield coll + finally: + if hasattr(coll, "destroy") and coll is not None: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + +class TestCollectionOpen: + def test_open_basic_functionality( + self, tmp_path_factory, collection_schema, collection_option + ): + import sys + import time + import os + + # Create unique temp directory + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + # Ensure the path exists + collection_path_str = str(collection_path) + print(f"DEBUG: Collection path: {collection_path_str}") + print(f"DEBUG: Temp directory exists: {temp_dir.exists()}") + + # Create and open collection first + created_coll = zvec.create_and_open( + path=collection_path_str, schema=collection_schema, option=collection_option + ) + + assert created_coll is not None, ( + f"Failed to create collection, returned None instead of valid Collection object. Path: {collection_path_str}" + ) + assert created_coll.path == collection_path_str, ( + f"Collection path mismatch. Expected: {collection_path_str}, Actual: {created_coll.path}" + ) + assert created_coll.schema.name == "test_collection", ( + f"Collection schema name mismatch. Expected: test_collection, Actual: {created_coll.schema.name}" + ) + + # Insert multiple documents to verify persistence + docs = [] + for i in range(3): + doc = Doc( + id=f"{i}", + fields={"id": i, "name": f"test_{i}", "weight": float(i * 10)}, + vectors={ + "dense": [float(j + i) for j in range(128)], + "sparse": {j: float(j + i) for j in range(5)}, + }, + ) + docs.append(doc) + + result = created_coll.insert(docs) + assert len(result) == 3, f"Expected 3 insertion results, but got {len(result)}" + for i, res in enumerate(result): + assert res.ok(), ( + f"Insertion result {i} is not OK. Status code: {res.code()}, Message: {res.message()}" + ) + + # Verify documents were inserted using fetch interface + fetched_docs_after_insert = created_coll.fetch(["0", "1", "2"]) + assert len(fetched_docs_after_insert) == 3, ( + f"Expected 3 fetched documents after insertion, but got {len(fetched_docs_after_insert)}" + ) + assert "0" in fetched_docs_after_insert, ( + "Document with ID '0' not found in fetched results after insertion" + ) + assert "1" in fetched_docs_after_insert, ( + "Document with ID '1' not found in fetched results after insertion" + ) + assert "2" in fetched_docs_after_insert, ( + "Document with ID '2' not found in fetched results after insertion" + ) + + # Verify fetched document content after insertion + for i in range(3): + doc = fetched_docs_after_insert[f"{i}"] + assert doc is not None, ( + f"Fetched document with ID '{i}' is None after insertion" + ) + assert doc.id == f"{i}", ( + f"Document ID mismatch for document '{i}' after insertion. Expected: {i}, Actual: {doc.id}" + ) + assert doc.field("id") == i, ( + f"Document id field mismatch for document '{i}' after insertion. Expected: {i}, Actual: {doc.field('id')}" + ) + assert doc.field("name") == f"test_{i}", ( + f"Document name field mismatch for document '{i}' after insertion. Expected: test_{i}, Actual: {doc.field('name')}" + ) + assert doc.field("weight") == float(i * 10), ( + f"Document weight field mismatch for document '{i}' after insertion. Expected: {float(i * 10)}, Actual: {doc.field('weight')}" + ) + + # Verify vector access after insertion + assert doc.vector("dense") is not None, ( + f"Document {i} should have dense vector after insertion" + ) + assert doc.vector("sparse") is not None, ( + f"Document {i} should have sparse vector after insertion" + ) + + # Verify vector types after insertion + assert isinstance(doc.vector("dense"), list), ( + f"Document {i} dense vector should be dict after insertion, got {type(doc.vector('dense'))}" + ) + assert isinstance(doc.vector("sparse"), dict), ( + f"Document {i} sparse vector should be dict after insertion, got {type(doc.vector('sparse'))}" + ) + + # Verify documents were inserted using stats + stats = created_coll.stats + assert stats is not None, "Collection stats should not be None" + assert stats.doc_count == 3, ( + f"Document count mismatch after insertion. Expected: 3, Actual: {stats.doc_count}" + ) + + # Store the collection path before cleanup + collection_path = created_coll.path + + # Clean up the created collection reference + del created_coll + + # Wait and verify the path still exists + print(f"DEBUG: Collection path after destroy: {collection_path}") + print(f"DEBUG: Path exists after destroy: {os.path.exists(collection_path)}") + + # Now open the existing collection + try: + print(f"DEBUG: Path exists before open: {os.path.exists(collection_path)}") + + # List contents of parent directory for debugging + parent_dir = os.path.dirname(collection_path) + if os.path.exists(parent_dir): + print(f"DEBUG: Parent directory contents: {os.listdir(parent_dir)}") + + opened_coll = zvec.open(path=collection_path, option=collection_option) + + assert opened_coll is not None, ( + f"Failed to open existing collection at path: {collection_path}. Returned None instead of valid Collection object" + ) + assert opened_coll.path == collection_path, ( + f"Opened collection path mismatch. Expected: {collection_path}, Actual: {opened_coll.path}" + ) + assert opened_coll.schema.name == "test_collection", ( + f"Opened collection schema name mismatch. Expected: test_collection, Actual: {opened_coll.schema.name}" + ) + + # Check reference count of opened collection + opened_ref_count = sys.getrefcount(opened_coll) + print(f"DEBUG: Reference count of opened collection: {opened_ref_count}") + + # Verify data persistence + # Verify data persistence using fetch interface + fetched_docs = opened_coll.fetch(["0", "1", "2"]) + assert len(fetched_docs) == 3, ( + f"Expected 3 fetched documents after reopening, but got {len(fetched_docs)}" + ) + assert "0" in fetched_docs, ( + "Document with ID '0' not found in fetched results after reopening" + ) + assert "1" in fetched_docs, ( + "Document with ID '1' not found in fetched results after reopening" + ) + assert "2" in fetched_docs, ( + "Document with ID '2' not found in fetched results after reopening" + ) + + # Verify fetched document content after reopening collection + for i in range(3): + doc = fetched_docs[f"{i}"] + assert doc is not None, ( + f"Fetched document with ID '{i}' is None after reopening collection" + ) + assert doc.id == f"{i}", ( + f"Document ID mismatch for document '{i}' after reopening. Expected: {i}, Actual: {doc.id}" + ) + assert doc.field("id") == i, ( + f"Document id field mismatch for document '{i}' after reopening. Expected: {i}, Actual: {doc.field('id')}" + ) + assert doc.field("name") == f"test_{i}", ( + f"Document name field mismatch for document '{i}' after reopening. Expected: test_{i}, Actual: {doc.field('name')}" + ) + assert doc.field("weight") == float(i * 10), ( + f"Document weight field mismatch for document '{i}' after reopening. Expected: {float(i * 10)}, Actual: {doc.field('weight')}" + ) + + # Verify vector access after reopening + assert doc.vector("dense") is not None, ( + f"Document {i} should have dense vector after reopening" + ) + assert doc.vector("sparse") is not None, ( + f"Document {i} should have sparse vector after reopening" + ) + + # Verify vector types after reopening + assert isinstance(doc.vector("dense"), list), ( + f"Document {i} dense vector should be dict after reopening, got {type(doc.vector('dense'))}" + ) + assert isinstance(doc.vector("sparse"), dict), ( + f"Document {i} sparse vector should be dict after reopening, got {type(doc.vector('sparse'))}" + ) + + # Verify score attribute exists + assert hasattr(doc, "score"), ( + f"Document {i} should have a score attribute after reopening" + ) + assert isinstance(doc.score, (int, float)), ( + f"Document {i} score should be numeric after reopening, got {type(doc.score)}" + ) + # For fetch operations, score is typically 0.0 + assert doc.score == 0.0, ( + f"Document {i} score should be 0.0 for fetch operation after reopening, but got {doc.score}" + ) + + # Test query functionality + query_result = opened_coll.query(include_vector=True) + assert len(query_result) == 3, ( + f"Expected 3 query results, but got {len(query_result)}" + ) + + # Verify query results have proper structure and content with detailed validation + returned_doc_ids = set() + for doc in query_result: + # Verify basic document structure + assert doc.id is not None, f"Query result document should have an ID" + assert doc.id in ["0", "1", "2"], ( + f"Query result document ID should be one of ['0', '1', '2'], but got {doc.id}" + ) + returned_doc_ids.add(doc.id) + + # Verify field access + assert doc.field("id") is not None, ( + f"Document {doc.id} should have id field" + ) + assert doc.field("name") is not None, ( + f"Document {doc.id} should have name field" + ) + assert doc.field("weight") is not None, ( + f"Document {doc.id} should have weight field" + ) + + # Verify field values + expected_id = int(doc.id) + assert doc.field("id") == expected_id, ( + f"Document {doc.id} id field mismatch. Expected: {expected_id}, Actual: {doc.field('id')}" + ) + assert doc.field("name") == f"test_{expected_id}", ( + f"Document {doc.id} name field mismatch. Expected: test_{expected_id}, Actual: {doc.field('name')}" + ) + assert doc.field("weight") == float(expected_id * 10), ( + f"Document {doc.id} weight field mismatch. Expected: {float(expected_id * 10)}, Actual: {doc.field('weight')}" + ) + + # Verify vector access + assert doc.vector("dense") is not None, ( + f"Document {doc.id} should have dense vector" + ) + assert doc.vector("sparse") is not None, ( + f"Document {doc.id} should have sparse vector" + ) + + # Verify vector types + assert isinstance(doc.vector("dense"), list), ( + f"Document {doc.id} dense vector should be list, got {type(doc.vector('dense'))}" + ) + assert isinstance(doc.vector("sparse"), dict), ( + f"Document {doc.id} sparse vector should be dict, got {type(doc.vector('sparse'))}" + ) + + # Verify score attribute exists + assert hasattr(doc, "score"), ( + f"Document {doc.id} should have a score attribute" + ) + assert isinstance(doc.score, (int, float)), ( + f"Document {doc.id} score should be numeric, got {type(doc.score)}" + ) + + # Verify all expected documents are returned + expected_doc_ids = {"0", "1", "2"} + assert returned_doc_ids == expected_doc_ids, ( + f"Query should return all expected documents. Expected: {expected_doc_ids}, Actual: {returned_doc_ids}" + ) + + # === Enhanced validation based on test_collection_dql_operations.py === + + # Verify vector field names accessibility for all documents + for doc in query_result: + vector_names = doc.vector_names() + expected_vector_names = {"dense", "sparse"} + assert set(vector_names) == expected_vector_names, ( + f"Document {doc.id} vector names mismatch. Expected: {expected_vector_names}, Actual: {set(vector_names)}" + ) + + # Verify all vector fields can be accessed + for vector_name in expected_vector_names: + vector_data = doc.vector(vector_name) + assert vector_data is not None, ( + f"Document {doc.id} should have accessible vector '{vector_name}'" + ) + if vector_name == "dense": + assert isinstance(vector_data, list), ( + f"Document {doc.id} vector '{vector_name}' should be list, got {type(vector_data)}" + ) + else: + assert isinstance(vector_data, dict), ( + f"Document {doc.id} vector '{vector_name}' should be dict, got {type(vector_data)}" + ) + + # Test query with filter + filtered_result = opened_coll.query(filter="id >= 1", include_vector=True) + assert len(filtered_result) == 2, ( + f"Expected 2 filtered query results (id >= 1), but got {len(filtered_result)}" + ) + + # Verify filtered query results + filtered_doc_ids = set() + for doc in filtered_result: + assert doc.id is not None, ( + f"Filtered query result document should have an ID" + ) + assert doc.id in ["1", "2"], ( + f"Filtered query result document ID should be one of ['1', '2'], but got {doc.id}" + ) + filtered_doc_ids.add(doc.id) + + # Verify filter condition is satisfied + doc_id = int(doc.id) + assert doc_id >= 1, ( + f"Document {doc.id} should satisfy filter condition id >= 1" + ) + + # Verify document structure + assert doc.field("id") is not None, ( + f"Document {doc.id} should have id field" + ) + assert doc.field("name") is not None, ( + f"Document {doc.id} should have name field" + ) + assert doc.field("weight") is not None, ( + f"Document {doc.id} should have weight field" + ) + + # Verify field values + assert doc.field("id") == doc_id, ( + f"Document {doc.id} id field mismatch. Expected: {doc_id}, Actual: {doc.field('id')}" + ) + assert doc.field("name") == f"test_{doc_id}", ( + f"Document {doc.id} name field mismatch. Expected: test_{doc_id}, Actual: {doc.field('name')}" + ) + assert doc.field("weight") == float(doc_id * 10), ( + f"Document {doc.id} weight field mismatch. Expected: {float(doc_id * 10)}, Actual: {doc.field('weight')}" + ) + + # Verify vector access + assert doc.vector("dense") is not None, ( + f"Document {doc.id} should have dense vector" + ) + assert doc.vector("sparse") is not None, ( + f"Document {doc.id} should have sparse vector" + ) + + # Verify score attribute exists + assert hasattr(doc, "score"), ( + f"Document {doc.id} should have a score attribute" + ) + assert isinstance(doc.score, (int, float)), ( + f"Document {doc.id} score should be numeric, got {type(doc.score)}" + ) + + # Verify filtered documents + expected_filtered_ids = {"1", "2"} + assert filtered_doc_ids == expected_filtered_ids, ( + f"Filtered query should return expected documents. Expected: {expected_filtered_ids}, Actual: {filtered_doc_ids}" + ) + + # Test vector query functionality for dense vectors + query_vector_dense = [0.1] * 128 + vector_query_result = opened_coll.query( + Query(field_name="dense", vector=query_vector_dense) + ) + assert len(vector_query_result) > 0, ( + f"Expected at least 1 vector query result, but got {len(vector_query_result)}" + ) + + # Verify vector query results structure + for doc in vector_query_result[:3]: # Check first 3 results + assert doc.id is not None, ( + f"Vector query result document should have an ID" + ) + assert doc.id in ["0", "1", "2"], ( + f"Vector query result document ID should be one of ['0', '1', '2'], but got {doc.id}" + ) + + # Verify document structure + assert doc.field("id") is not None, ( + f"Document {doc.id} should have id field" + ) + assert doc.field("name") is not None, ( + f"Document {doc.id} should have name field" + ) + assert doc.field("weight") is not None, ( + f"Document {doc.id} should have weight field" + ) + + # Verify vector access + assert doc.vector("dense") is not None, ( + f"Document {doc.id} should have dense vector" + ) + assert doc.vector("sparse") is not None, ( + f"Document {doc.id} should have sparse vector" + ) + + # Verify score attribute exists and is numeric + assert hasattr(doc, "score"), ( + f"Document {doc.id} should have a score attribute" + ) + assert isinstance(doc.score, (int, float)), ( + f"Document {doc.id} score should be numeric, got {type(doc.score)}" + ) + + # For dense vector queries, score should typically be non-negative (depending on metric) + # Note: This may vary based on the metric type used + assert doc.score >= 0 or doc.score < 0, ( + f"Document {doc.id} score should be a valid number" + ) + + # Test vector query functionality for sparse vectors + query_vector_sparse = {1: 1.0, 2: 2.0, 3: 3.0} + sparse_vector_query_result = opened_coll.query( + Query(field_name="sparse", vector=query_vector_sparse) + ) + assert len(sparse_vector_query_result) > 0, ( + f"Expected at least 1 sparse vector query result, but got {len(sparse_vector_query_result)}" + ) + + # Verify sparse vector query results structure + for doc in sparse_vector_query_result[:3]: # Check first 3 results + assert doc.id is not None, ( + f"Sparse vector query result document should have an ID" + ) + assert doc.id in ["0", "1", "2"], ( + f"Sparse vector query result document ID should be one of ['0', '1', '2'], but got {doc.id}" + ) + + # Verify document structure + assert doc.field("id") is not None, ( + f"Document {doc.id} should have id field" + ) + assert doc.field("name") is not None, ( + f"Document {doc.id} should have name field" + ) + assert doc.field("weight") is not None, ( + f"Document {doc.id} should have weight field" + ) + + # Verify vector access + assert doc.vector("dense") is not None, ( + f"Document {doc.id} should have dense vector" + ) + assert doc.vector("sparse") is not None, ( + f"Document {doc.id} should have sparse vector" + ) + + # Verify score attribute exists and is numeric + assert hasattr(doc, "score"), ( + f"Document {doc.id} should have a score attribute" + ) + assert isinstance(doc.score, (int, float)), ( + f"Document {doc.id} score should be numeric, got {type(doc.score)}" + ) + + # Clean up + if hasattr(opened_coll, "destroy") and opened_coll is not None: + opened_coll.destroy() + print("DEBUG: Opened collection destroyed successfully") + + except Exception as e: + logging.error("Exception occurred: [{}]".format(e)) + raise e + + @pytest.mark.parametrize( + "read_only,enable_mmap,description", COLLECTION_OPTION_TEST_CASES_VALID + ) + @pytest.mark.parametrize("createAndopen_enable_mmap", [True, False]) + def test_open_with_different_collection_options_valid( + self, + tmp_path_factory, + createAndopen_enable_mmap, + read_only, + enable_mmap, + description, + collection_schema, + ): + # Create collection with initial option + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + initial_option = CollectionOption( + read_only=False, enable_mmap=createAndopen_enable_mmap + ) + + # Create and open collection first + created_coll = zvec.create_and_open( + path=str(collection_path), schema=collection_schema, option=initial_option + ) + + assert created_coll is not None, "Failed to create collection" + + # Clean up the created collection reference + del created_coll + + # Now open with different options + collection_option = CollectionOption( + read_only=read_only, enable_mmap=enable_mmap + ) + + try: + opened_coll = zvec.open(path=str(collection_path), option=collection_option) + + assert opened_coll is not None, ( + f"Failed to open collection with option: {description}. Returned None instead of valid Collection object. Path: {collection_path}" + ) + assert opened_coll.path == str(collection_path), ( + f"Opened collection path mismatch. Expected: {collection_path}, Actual: {opened_coll.path}" + ) + assert opened_coll.schema.name == collection_schema.name, ( + f"Opened collection schema name mismatch. Expected: {collection_schema.name}, Actual: {opened_coll.schema.name}" + ) + assert opened_coll.option.read_only == read_only, ( + f"Opened collection read_only option mismatch. Expected: {read_only}, Actual: {opened_coll.option.read_only}" + ) + assert opened_coll.option.enable_mmap == createAndopen_enable_mmap, ( + f"Opened collection mmap option mismatch. Expected: {createAndopen_enable_mmap}, Actual: {opened_coll.option.enable_mmap}" + ) + + # Clean up + if ( + hasattr(opened_coll, "destroy") + and opened_coll is not None + and read_only == False + ): + opened_coll.destroy() + + except Exception as e: + logging.error("Exception occurred: [{}]".format(e)) + pytest.fail(f"Failed to open collection with different options: {e}") + + def test_open_with_none_option(self, tmp_path_factory, collection_schema): + # Create collection + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + initial_option = CollectionOption(read_only=False, enable_mmap=True) + + # Create and open collection first + created_coll = zvec.create_and_open( + path=str(collection_path), schema=collection_schema, option=initial_option + ) + + assert created_coll is not None, ( + f"Failed to create collection. Returned None instead of valid Collection object. Path: {collection_path}" + ) + + # Clean up the created collection reference + del created_coll + + # Now open with None option + with pytest.raises(Exception) as exc_info: + zvec.open(path=str(collection_path), option=None) + + assert "incompatible function arguments" in str(exc_info.value), ( + f"Expected 'incompatible function arguments' error, but got: {exc_info.value}" + ) + + def test_reopen_collection(self, tmp_path_factory): + # Prepare schema + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(), + ), + FieldSchema( + "description", + DataType.STRING, + nullable=True, + index_param=InvertIndexParam(), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ], + ) + + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + # Create collection + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + # Create and open collection + coll1 = zvec.create_and_open( + path=str(collection_path), + schema=collection_schema, + option=collection_option, + ) + + assert coll1 is not None, "Failed to create and open collection" + + # Insert some data + doc = Doc( + id="1", + fields={"id": 1, "name": "test", "description": "这是一个中文描述。"}, + vectors={"dense": np.random.random(128).tolist()}, + ) + + result = coll1.insert(doc) + assert result.ok() + + # Close the first collection (delete reference) + del coll1 + + # Reopen the collection + coll2 = zvec.open(path=str(collection_path), option=collection_option) + + assert coll2 is not None, "Failed to reopen collection" + assert coll2.path == str(collection_path) + assert coll2.schema.name == collection_schema.name + + # Verify data is still there + fetched_docs = coll2.fetch(["1"]) + assert "1" in fetched_docs + fetched_doc = fetched_docs["1"] + assert fetched_doc.id == "1" + assert fetched_doc.field("name") == "test" + assert fetched_doc.field("description") == "这是一个中文描述。" + + # Clean up + if hasattr(coll2, "destroy") and coll2 is not None: + try: + coll2.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + def test_open_concurrent_same_path(self, tmp_path_factory): + # First create a collection + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ], + ) + + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + # Create collection path + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + # First create the collection + created_coll = zvec.create_and_open( + path=str(collection_path), + schema=collection_schema, + option=collection_option, + ) + + assert created_coll is not None, "Failed to create collection" + + # Close the collection so we can test concurrent opening + if hasattr(created_coll, "close") and created_coll is not None: + created_coll.close() + + # Shared variables to collect results from threads + results = [] + errors = [] + + # Lock for thread-safe operations + lock = threading.Lock() + # Clean up the created collection reference + del created_coll + + # Function to be executed by each thread + def open_collection_thread(thread_id): + try: + coll = zvec.open(path=str(collection_path), option=collection_option) + with lock: + results.append((thread_id, coll)) + # Close the collection if opened successfully + if hasattr(coll, "close") and coll is not None: + coll.close() + except Exception as e: + with lock: + errors.append((thread_id, str(e))) + + # Create 5 threads to call open concurrently + threads = [] + for i in range(5): + thread = threading.Thread(target=open_collection_thread, args=(i,)) + threads.append(thread) + thread.start() + + # Wait for all threads to complete + for thread in threads: + thread.join() + + # Verify concurrency safety: only one should succeed, others should fail + assert len(results) == 1, ( + f"Expected exactly one successful open, but got {len(results)}" + ) + assert len(errors) == 4, ( + f"Expected exactly four failures, but got {len(errors)}" + ) + + # Additional verification: check that the successful open has a valid collection + successful_thread_id, successful_collection = results[0] + assert successful_collection is not None, ( + "Successful open should return a valid collection" + ) + assert successful_collection.path == str(collection_path), ( + "Collection path mismatch" + ) + + # Clean up the successfully opened collection + if ( + hasattr(successful_collection, "destroy") + and successful_collection is not None + ): + try: + successful_collection.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + def test_open_with_corrupted_files(self, tmp_path_factory): + # First create a collection + collection_schema = zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ) + ], + ) + + collection_option = CollectionOption(read_only=False, enable_mmap=True) + + # Create collection path + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + # First create the collection + created_coll = zvec.create_and_open( + path=str(collection_path), + schema=collection_schema, + option=collection_option, + ) + + assert created_coll is not None, "Failed to create collection" + + # Close the collection so we can manipulate its files + if hasattr(created_coll, "close") and created_coll is not None: + created_coll.close() + + # Test case 1: Delete some files in the collection directory (simulate partial corruption) + import os + import shutil + import random + + # Get the collection directory path + collection_dir = str(collection_path) + + # List all files in the collection directory + files_in_dir = [] + for root, dirs, files in os.walk(collection_dir): + for file in files: + files_in_dir.append(os.path.join(root, file)) + + # Randomly delete approximately half of the files to simulate partial corruption + if files_in_dir: + # Shuffle the list to randomly select files + random.shuffle(files_in_dir) + files_to_delete = files_in_dir[: len(files_in_dir) // 2] + for file_path in files_to_delete: + try: + os.remove(file_path) + except Exception as e: + pass # Ignore errors during deletion + + # Try to open the collection with missing files - should raise an exception + with pytest.raises(Exception): + zvec.open(path=str(collection_path), option=collection_option) + + # Test case 2: Delete all files in the collection directory (simulate complete corruption) + # Recreate the collection + recreated_coll = zvec.create_and_open( + path=str(collection_path) + "_all", + schema=collection_schema, + option=collection_option, + ) + + assert recreated_coll is not None, "Failed to recreate collection" + + # Close the collection so we can manipulate its files + if hasattr(recreated_coll, "close") and recreated_coll is not None: + recreated_coll.close() + + # Delete all files in the collection directory + try: + shutil.rmtree(collection_dir) + os.makedirs(collection_dir) # Recreate empty directory + except Exception as e: + pass # Ignore errors during deletion + + # Try to open the collection with missing files - should raise an exception + with pytest.raises(Exception): + zvec.open(path=str(collection_path), option=collection_option) diff --git a/python/tests/detail/test_collection_recall.py b/python/tests/detail/test_collection_recall.py new file mode 100644 index 0000000..764f79f --- /dev/null +++ b/python/tests/detail/test_collection_recall.py @@ -0,0 +1,740 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from zvec.typing import DataType, StatusCode, MetricType, QuantizeType +from zvec.model import Collection, Doc, Query +from zvec.model.param import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + FlatIndexParam, + IVFIndexParam, + DiskAnnIndexParam, + HnswQueryParam, + IVFQueryParam, + DiskAnnQueryParam, +) + +from zvec.model.schema import FieldSchema, VectorSchema +from zvec.extension import RrfReRanker, WeightedReRanker, QwenReRanker +from distance_helper import * + +from zvec import StatusCode +from distance_helper import * +from fixture_helper import * +from doc_helper import * +from params_helper import * + +import time + + +# ==================== helper ==================== +def batchdoc_and_check(collection: Collection, multiple_docs, operator="insert"): + if operator == "insert": + result = collection.insert(multiple_docs) + elif operator == "upsert": + result = collection.upsert(multiple_docs) + + elif operator == "update": + result = collection.update(multiple_docs) + else: + logging.error("operator value is error!") + + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok(), ( + f"result={result},Insert operation failed with code {item.code()}" + ) + + stats = collection.stats + assert stats is not None, "Collection stats should not be None" + """assert stats.doc_count == len(multiple_docs), ( + f"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}" + )""" + + doc_ids = [doc.id for doc in multiple_docs] + fetched_docs = collection.fetch(doc_ids) + assert len(fetched_docs) == len(multiple_docs), ( + f"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}" + ) + + for original_doc in multiple_docs: + assert original_doc.id in fetched_docs, ( + f"Expected document ID {original_doc.id} in fetched documents" + ) + fetched_doc = fetched_docs[original_doc.id] + + assert is_doc_equal(fetched_doc, original_doc, collection.schema) + + assert hasattr(fetched_doc, "score"), "Document should have a score attribute" + assert fetched_doc.score == 0.0, ( + "Fetch operation should return default score of 0.0" + ) + + +def compute_exact_similarity_scores( + vectors_a, + vectors_b, + metric_type=MetricType.IP, + DataType=DataType.VECTOR_FP32, + QuantizeType=QuantizeType.UNDEFINED, +): + similarities = [] + for i, vec_a in enumerate(vectors_a): + for j, vec_b in enumerate(vectors_b): + similarity = distance_recall(vec_a, vec_b, metric_type, DataType) + similarities.append((j, similarity)) + + # For L2,COSINE metric, smaller distances mean higher similarity, so sort in ascending order + if ( + metric_type in [MetricType.L2] + and DataType + in [DataType.VECTOR_FP32, DataType.VECTOR_FP16, DataType.VECTOR_INT8] + ) or ( + metric_type in [MetricType.COSINE] + and DataType in [DataType.VECTOR_FP32, DataType.VECTOR_FP16] + ): + similarities.sort(key=lambda x: x[1], reverse=False) # Ascending order for L2 + + else: + similarities.sort( + key=lambda x: x[1], reverse=True + ) # Descending order for others + + # Special handling for COSINE in FP16 to address precision issues + if metric_type == MetricType.COSINE and DataType == DataType.VECTOR_FP16: + # Clamp values to valid cosine distance range [0, 2] and handle floating point errors + similarities = [(idx, max(0.0, min(2.0, score))) for idx, score in similarities] + + return similarities + + +def get_ground_truth_for_vector_query( + collection, + query_vector, + field_name, + all_docs, + query_idx, + metric_type, + k, + use_exact_computation=False, +): + if use_exact_computation: + all_vectors = [doc.vectors[field_name] for doc in all_docs] + + for d, f in DEFAULT_VECTOR_FIELD_NAME.items(): + if field_name == f: + DataType = d + break + similarities = compute_exact_similarity_scores( + [query_vector], + all_vectors, + metric_type, + DataType=DataType, + QuantizeType=QuantizeType, + ) + + if metric_type == MetricType.COSINE and DataType == DataType.VECTOR_FP16: + # Filter out tiny non-zero values that may be caused by precision errors + similarities = [ + (idx, max(0.0, min(2.0, score))) for idx, score in similarities + ] + + ground_truth_ids_scores = similarities[:k] + print("Get the most similar k document IDs k:,ground_truth_ids_scores") + print(k, ground_truth_ids_scores) + return ground_truth_ids_scores + + else: + full_result = collection.query( + Query(field_name=field_name, vector=query_vector), + topk=min(len(all_docs), 1024), + include_vector=True, + ) + + ground_truth_ids_scores = [ + (result.id, result.score) for result in full_result[:k] + ] + + if not ground_truth_ids_scores: + ground_truth_ids_scores = [(all_docs[query_idx].id, 0)] + + return ground_truth_ids_scores + + +def get_ground_truth_map(collection, test_docs, query_vectors_map, metric_type, k): + ground_truth_map = {} + + for field_name, query_vectors in query_vectors_map.items(): + ground_truth_map[field_name] = {} + + # Support per-field metric type: metric_type can be a dict mapping + # field_name -> MetricType, or a single MetricType applied to all fields. + if isinstance(metric_type, dict): + field_metric = metric_type.get(field_name, MetricType.IP) + else: + field_metric = metric_type + + for i, query_vector in enumerate(query_vectors): + # Get the ground truth for this query + relevant_doc_ids_scores = get_ground_truth_for_vector_query( + collection, + query_vector, + field_name, + test_docs, + i, + field_metric, + k, + True, + ) + ground_truth_map[field_name][i] = relevant_doc_ids_scores + + print("ground_truth_map:\n") + print(ground_truth_map) + return ground_truth_map + + +def calculate_recall_at_k( + collection: Collection, + test_docs, + query_vectors_map, + schema, + k=1, + expected_doc_ids_scores_map=None, + tolerance=0.01, +): + recall_stats = {} + + for field_name, query_vectors in query_vectors_map.items(): + recall_stats[field_name] = { + "relevant_retrieved_count": 0, + "total_relevant_count": 0, + "retrieved_count": 0, + "recall_at_k": 0.0, + } + + for i, query_vector in enumerate(query_vectors): + print("Starting %dth query" % i) + + query_result_list = collection.query( + Query(field_name=field_name, vector=query_vector), + topk=1024, + include_vector=True, + ) + retrieved_count = len(query_result_list) + + query_result_ids_scores = [] + for word in query_result_list: + query_result_ids_scores.append((word.id, word.score)) + + recall_stats[field_name]["retrieved_count"] += retrieved_count + + print("expected_doc_ids_scores_map:\n") + print(expected_doc_ids_scores_map) + if i in (expected_doc_ids_scores_map[field_name]): + expected_relevant_ids_scores = expected_doc_ids_scores_map[field_name][ + i + ] + print( + "field_name,i,expected_relevant_ids_scores, query_result_ids_scores:\n" + ) + print( + field_name, + i, + "\n", + expected_relevant_ids_scores, + "\n", + len(query_result_ids_scores), + query_result_ids_scores, + ) + + # Update total relevant documents count + recall_stats[field_name]["total_relevant_count"] += len( + expected_relevant_ids_scores + ) + + relevant_found_count = 0 + for ids_scores_except in expected_relevant_ids_scores: + for ids_scores_result in query_result_ids_scores[:k]: + if int(ids_scores_result[0]) == int(ids_scores_except[0]): + relevant_found_count += 1 + break + elif ( + int(ids_scores_result[0]) != int(ids_scores_except[0]) + and abs(ids_scores_result[1] - ids_scores_except[1]) + <= tolerance + ): + print("IDs are not equal, but the error is small, tolerance") + print( + ids_scores_result[0], + ids_scores_except[0], + ids_scores_result[1], + ids_scores_except[1], + tolerance, + ) + relevant_found_count += 1 + break + else: + continue + + recall_stats[field_name]["relevant_retrieved_count"] += relevant_found_count + + # Calculate Recall@K + if recall_stats[field_name]["total_relevant_count"] > 0: + recall_stats[field_name]["recall_at_k"] = ( + recall_stats[field_name]["relevant_retrieved_count"] + / recall_stats[field_name]["total_relevant_count"] + ) + + return recall_stats + + +class TestRecall: + @pytest.mark.parametrize( + "full_schema_new", + [ + (True, True, HnswIndexParam()), + (False, True, IVFIndexParam()), + (False, True, DiskAnnIndexParam()), + (False, True, FlatIndexParam()), # ——ok + ( + True, + True, + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + ), + ), + ( + True, + True, + HnswIndexParam( + metric_type=MetricType.COSINE, + m=24, + ef_construction=150, + ), + ), + ( + True, + True, + HnswIndexParam( + metric_type=MetricType.L2, + m=32, + ef_construction=200, + ), + ), + ( + False, + True, + FlatIndexParam( + metric_type=MetricType.IP, + ), + ), + ( + True, + True, + FlatIndexParam( + metric_type=MetricType.COSINE, + ), + ), + ( + True, + True, + FlatIndexParam( + metric_type=MetricType.L2, + ), + ), + ( + True, + True, + IVFIndexParam( + metric_type=MetricType.IP, + n_list=100, + n_iters=10, + use_soar=False, + ), + ), + ( + True, + True, + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + ), + ( + True, + True, + IVFIndexParam( + metric_type=MetricType.COSINE, + n_list=150, + n_iters=15, + use_soar=False, + ), + ), + ( + True, + True, + DiskAnnIndexParam( + metric_type=MetricType.IP, + max_degree=32, + ), + ), + ( + True, + True, + DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32), + ), + ], + indirect=True, + ) + @pytest.mark.parametrize("doc_num", [500]) + @pytest.mark.parametrize("query_num", [10]) + @pytest.mark.parametrize("top_k", [1]) + def test_recall_with_single_vector_valid_500( + self, + full_collection_new: Collection, + doc_num, + query_num, + top_k, + full_schema_new, + request, + ): + full_schema_params = request.getfixturevalue("full_schema_new") + + # Build per-field metric type map so ground truth uses each field's + # actual index metric (fields may fall back to HnswIndexParam/IP). + field_metric_map = {} + for vector_para in full_schema_params.vectors: + if vector_para.index_param is not None: + field_metric_map[vector_para.name] = vector_para.index_param.metric_type + else: + field_metric_map[vector_para.name] = MetricType.IP + + metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP) + + multiple_docs = [ + generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num) + ] + print("len(multiple_docs):\n") + print(len(multiple_docs)) + # print(multiple_docs) + + for i in range(10): + if i != 0: + pass + # print(multiple_docs[i * 1000:1000 * (i + 1)]) + batchdoc_and_check( + full_collection_new, + multiple_docs[i * 1000 : 1000 * (i + 1)], + operator="insert", + ) + + stats = full_collection_new.stats + assert stats.doc_count == len(multiple_docs) + + doc_ids = ["0", "1"] + fetched_docs = full_collection_new.fetch(doc_ids) + print("fetched_docs,multiple_docs") + print( + fetched_docs[doc_ids[0]].vectors["sparse_vector_fp32_field"], + fetched_docs[doc_ids[0]].vectors["sparse_vector_fp16_field"], + fetched_docs[doc_ids[1]].vectors["sparse_vector_fp32_field"], + fetched_docs[doc_ids[1]].vectors["sparse_vector_fp16_field"], + "\n", + multiple_docs[0].vectors["sparse_vector_fp32_field"], + multiple_docs[0].vectors["sparse_vector_fp32_field"], + multiple_docs[1].vectors["sparse_vector_fp32_field"], + multiple_docs[1].vectors["sparse_vector_fp16_field"], + ) + + full_collection_new.optimize(option=OptimizeOption()) + + time.sleep(2) + + query_vectors_map = {} + for field_name in DEFAULT_VECTOR_FIELD_NAME.values(): + query_vectors_map[field_name] = [ + multiple_docs[i].vectors[field_name] for i in range(query_num) + ] + + # Get ground truth mapping (pass per-field metric map) + ground_truth_map = get_ground_truth_map( + full_collection_new, + multiple_docs, + query_vectors_map, + field_metric_map, + top_k, + ) + + # Validate ground truth mapping structure + for field_name in DEFAULT_VECTOR_FIELD_NAME.values(): + assert field_name in ground_truth_map + field_gt = ground_truth_map[field_name] + assert len(field_gt) == query_num + + for query_idx in range(query_num): + assert query_idx in field_gt + relevant_ids = field_gt[query_idx] + assert isinstance(relevant_ids, list) + assert len(relevant_ids) <= top_k + + # Print ground truth statistics + print(f"Ground Truth for Top-{top_k} Retrieval:") + for field_name, field_gt in ground_truth_map.items(): + print(f" {field_name}:") + for query_idx, relevant_ids in field_gt.items(): + print( + f" Query {query_idx}: {len(relevant_ids)} relevant docs - {relevant_ids[:5]}{'...' if len(relevant_ids) > 5 else ''}" + ) + + # Calculate Recall@K using ground truth + recall_at_k_stats = calculate_recall_at_k( + full_collection_new, + multiple_docs, + query_vectors_map, + full_schema_new, + k=top_k, + expected_doc_ids_scores_map=ground_truth_map, + tolerance=0.01, + ) + print("ground_truth_map:\n") + print(ground_truth_map) + + print("(recall_at_k_stats:\n") + print(recall_at_k_stats) + print("field_metric_map:") + print(field_metric_map) + # Print Recall@K statistics + print(f"Recall@{top_k} using Ground Truth:") + for field_name, stats in recall_at_k_stats.items(): + print(f" {field_name}:") + print( + f" Relevant Retrieved: {stats['relevant_retrieved_count']}/{stats['total_relevant_count']}" + ) + print(f" Recall@{top_k}: {stats['recall_at_k']:.4f}") + for k, v in recall_at_k_stats.items(): + assert v["recall_at_k"] == 1.0 + + @pytest.mark.parametrize( + "full_schema_new", + [ + (True, True, HnswIndexParam()), + (False, True, IVFIndexParam()), + (False, True, FlatIndexParam()), # ——ok + ( + True, + True, + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + ), + ), + ( + True, + True, + HnswIndexParam( + metric_type=MetricType.COSINE, + m=24, + ef_construction=150, + ), + ), + # (True, True, HnswIndexParam(metric_type=MetricType.L2, m=32, ef_construction=200, )), + ( + False, + True, + FlatIndexParam( + metric_type=MetricType.IP, + ), + ), + ( + True, + True, + FlatIndexParam( + metric_type=MetricType.COSINE, + ), + ), + # (True, True, FlatIndexParam(metric_type=MetricType.L2, )), + ( + True, + True, + IVFIndexParam( + metric_type=MetricType.IP, + n_list=100, + n_iters=10, + use_soar=False, + ), + ), + ( + True, + True, + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + ), + ( + True, + True, + DiskAnnIndexParam(metric_type=MetricType.IP, max_degree=32), + ), + ( + True, + True, + DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32), + ), + ( + True, + True, + DiskAnnIndexParam(metric_type=MetricType.COSINE, max_degree=32), + ), + ], + indirect=True, + ) + @pytest.mark.parametrize("doc_num", [2000]) + @pytest.mark.parametrize("query_num", [2]) + @pytest.mark.parametrize("top_k", [1]) + @pytest.mark.skip(reason="known bug") + def test_recall_with_single_vector_valid_2000( + self, + full_collection_new: Collection, + doc_num, + query_num, + top_k, + full_schema_new, + request, + ): + full_schema_params = request.getfixturevalue("full_schema_new") + + # Build per-field metric type map so ground truth uses each field's + # actual index metric (fields may fall back to HnswIndexParam/IP). + field_metric_map = {} + for vector_para in full_schema_params.vectors: + if vector_para.index_param is not None: + field_metric_map[vector_para.name] = vector_para.index_param.metric_type + else: + field_metric_map[vector_para.name] = MetricType.IP + + metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP) + + multiple_docs = [ + generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num) + ] + print("len(multiple_docs):\n") + print(len(multiple_docs)) + # print(multiple_docs) + + for i in range(10): + if i != 0: + pass + # print(multiple_docs[i * 1000:1000 * (i + 1)]) + batchdoc_and_check( + full_collection_new, + multiple_docs[i * 1000 : 1000 * (i + 1)], + operator="insert", + ) + + stats = full_collection_new.stats + assert stats.doc_count == len(multiple_docs) + + doc_ids = ["0", "1"] + fetched_docs = full_collection_new.fetch(doc_ids) + print("fetched_docs,multiple_docs") + print( + fetched_docs[doc_ids[0]].vectors["sparse_vector_fp32_field"], + fetched_docs[doc_ids[0]].vectors["sparse_vector_fp16_field"], + fetched_docs[doc_ids[1]].vectors["sparse_vector_fp32_field"], + fetched_docs[doc_ids[1]].vectors["sparse_vector_fp16_field"], + "\n", + multiple_docs[0].vectors["sparse_vector_fp32_field"], + multiple_docs[0].vectors["sparse_vector_fp32_field"], + multiple_docs[1].vectors["sparse_vector_fp32_field"], + multiple_docs[1].vectors["sparse_vector_fp16_field"], + ) + + full_collection_new.optimize(option=OptimizeOption()) + + time.sleep(2) + + query_vectors_map = {} + for field_name in DEFAULT_VECTOR_FIELD_NAME.values(): + query_vectors_map[field_name] = [ + multiple_docs[i].vectors[field_name] for i in range(query_num) + ] + + # Get ground truth mapping (pass per-field metric map) + ground_truth_map = get_ground_truth_map( + full_collection_new, + multiple_docs, + query_vectors_map, + field_metric_map, + top_k, + ) + + # Validate ground truth mapping structure + for field_name in DEFAULT_VECTOR_FIELD_NAME.values(): + assert field_name in ground_truth_map + field_gt = ground_truth_map[field_name] + assert len(field_gt) == query_num + + for query_idx in range(query_num): + assert query_idx in field_gt + relevant_ids = field_gt[query_idx] + assert isinstance(relevant_ids, list) + assert len(relevant_ids) <= top_k + + # Print ground truth statistics + print(f"Ground Truth for Top-{top_k} Retrieval:") + for field_name, field_gt in ground_truth_map.items(): + print(f" {field_name}:") + for query_idx, relevant_ids in field_gt.items(): + print( + f" Query {query_idx}: {len(relevant_ids)} relevant docs - {relevant_ids[:5]}{'...' if len(relevant_ids) > 5 else ''}" + ) + + # Calculate Recall@K using ground truth + recall_at_k_stats = calculate_recall_at_k( + full_collection_new, + multiple_docs, + query_vectors_map, + full_schema_new, + k=top_k, + expected_doc_ids_scores_map=ground_truth_map, + tolerance=0.01, + ) + print("ground_truth_map:\n") + print(ground_truth_map) + + print("(recall_at_k_stats:\n") + print(recall_at_k_stats) + print("field_metric_map:") + print(field_metric_map) + # Print Recall@K statistics + print(f"Recall@{top_k} using Ground Truth:") + for field_name, stats in recall_at_k_stats.items(): + print(f" {field_name}:") + print( + f" Relevant Retrieved: {stats['relevant_retrieved_count']}/{stats['total_relevant_count']}" + ) + print(f" Recall@{top_k}: {stats['recall_at_k']:.4f}") + for k, v in recall_at_k_stats.items(): + assert v["recall_at_k"] == 1.0 diff --git a/python/tests/detail/test_db_config.py b/python/tests/detail/test_db_config.py new file mode 100644 index 0000000..16203a8 --- /dev/null +++ b/python/tests/detail/test_db_config.py @@ -0,0 +1,307 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import pytest +import tempfile +import os +import sys +import subprocess + +import zvec +import zvec +from zvec import LogType, LogLevel + +# Error messages +INITIALIZATION_ERROR_MSG = "initialization failed" +RUNTIME_ERROR_MSG = "RuntimeError" +VALUE_ERROR_MSG = "ValueError" +TYPE_ERROR_MSG = "TypeError" + + +# ==================== helper ==================== +def run_in_subprocess(func): + def wrapper(*args, **kwargs): + if os.getenv("RUNNING_IN_SUBPROCESS"): + return func(*args, **kwargs) + + env = os.environ.copy() + env["RUNNING_IN_SUBPROCESS"] = "1" + env["PYTEST_CURRENT_TEST"] = func.__name__ + + import inspect + + filepath = inspect.getfile(func) + qualname = func.__qualname__.replace(".", "::") + test_id = f"{filepath}::{qualname}" + + project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + env["PYTHONPATH"] = project_root + ":" + env.get("PYTHONPATH", "") + + cmd = [sys.executable, "-m", "pytest", "-v", "-s", test_id] + + result = subprocess.run(cmd, env=env, capture_output=True, text=True) + if result.returncode != 0: + pytest.fail( + f"Subprocess test {func.__name__} failed with code {result.returncode}\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + + return wrapper + + +# ==================== Fixtures ==================== +@pytest.fixture(scope="function") +def temp_log_dir(tmp_path_factory): + return tmp_path_factory.mktemp("logs") + + +# ==================== Tests ==================== +class TestDbConfigInitialization: + @run_in_subprocess + def test_init_default(self): + # default config + # log_type: Optional[LogType] = LogType.CONSOLE, + # log_level: Optional[LogLevel] = LogLevel.WARN, + # log_dir: Optional[str] = "./logs", + # log_basename: Optional[str] = "zvec.log", + # log_file_size: Optional[int] = 2048, + # log_overdue_days: Optional[int] = 7, + zvec.init() + + @run_in_subprocess + def test_init_file_logger(self): + from pathlib import Path + import shutil + + zvec.init( + log_level=LogLevel.DEBUG, + log_type=LogType.FILE, + ) + # assert logdir exist + log_dir = Path("./logs") + assert log_dir.exists() + + # validate write log + col = zvec.create_and_open( + "/tmp/test/1", + zvec.CollectionSchema( + name="test", + vectors=zvec.VectorSchema( + dimension=4, + data_type=zvec.DataType.VECTOR_FP32, + name="image", + ), + ), + ) + col.insert(docs=[zvec.Doc(id="1", vectors={"image": [1.0, 2.0, 3.0, 4.0]})]) + assert any(log_dir.glob("zvec.log.*")) + + # clear + col.destroy() + shutil.rmtree(log_dir, ignore_errors=True) + + @run_in_subprocess + def test_init_with_mixed_config(self): + zvec.init( + memory_limit_mb=128, + log_type=LogType.FILE, + query_threads=1, + log_level=LogLevel.WARN, + ) + + @run_in_subprocess + def test_repeated_initialization(self): + # Calling init() repeatedly is allowed: + # it succeeds but becomes a no-op after the first successful init() + zvec.init() + + +class TestDbConfigMemoryLimitValidation: + @run_in_subprocess + def test_memory_limit_min_valid(self): + # MIN_MEMORY_LIMIT_BYTES is 100M + with pytest.raises(RuntimeError): + zvec.init(memory_limit_mb=99) + + @run_in_subprocess + def test_memory_limit_invalid_value(self): + # memory_limit_mb must >= 0 and must be int and if None, set default value + with pytest.raises(ValueError): + zvec.init(memory_limit_mb=0) + with pytest.raises(ValueError): + zvec.init(memory_limit_mb=-1) + with pytest.raises(TypeError): + zvec.init(memory_limit_mb="512") + with pytest.raises(TypeError): + zvec.init(memory_limit_mb=512.5) + + +class TestDbConfigThreadValidation: + @run_in_subprocess + def test_query_threads(self): + zvec.init(query_threads=1) + + @run_in_subprocess + def test_query_threads_invalid(self): + # query_threads must >= 0 and must be int and if None, set default value + with pytest.raises(ValueError): + zvec.init(query_threads=0) + with pytest.raises(ValueError): + zvec.init(query_threads=-1) + with pytest.raises(TypeError): + zvec.init(query_threads="value") + with pytest.raises(TypeError): + zvec.init(query_threads=512.5) + with pytest.raises(TypeError): + zvec.init(query_threads="512") + + @run_in_subprocess + def test_optimize_threads(self): + zvec.init(optimize_threads=1) + + @run_in_subprocess + def test_optimize_threads_invalid(self): + # optimize_threads must >= 0 and must be int and if None, set default value + with pytest.raises(ValueError): + zvec.init(optimize_threads=0) + with pytest.raises(ValueError): + zvec.init(optimize_threads=-1) + with pytest.raises(TypeError): + zvec.init(optimize_threads="value") + with pytest.raises(TypeError): + zvec.init(optimize_threads=512.5) + with pytest.raises(TypeError): + zvec.init(optimize_threads="512") + + +class TestDbConfigRatioValidation: + @run_in_subprocess + def test_init_invert_to_forward_scan_ratio(self): + # must be in [0,1] + zvec.init(invert_to_forward_scan_ratio=0.8) + + @run_in_subprocess + def test_init_invert_to_forward_scan_ratio_invalid(self): + with pytest.raises(ValueError): + zvec.init(invert_to_forward_scan_ratio=1.1) + with pytest.raises(ValueError): + zvec.init(invert_to_forward_scan_ratio=-0.1) + with pytest.raises(TypeError): + zvec.init(invert_to_forward_scan_ratio="0.8") + + @run_in_subprocess + def test_init_brute_force_by_keys_ratio(self): + zvec.init(brute_force_by_keys_ratio=0.8) + + @run_in_subprocess + def test_init_brute_force_by_keys_ratio_invalid(self): + with pytest.raises(ValueError): + zvec.init(brute_force_by_keys_ratio=1.1) + with pytest.raises(ValueError): + zvec.init(brute_force_by_keys_ratio=-0.1) + with pytest.raises(TypeError): + zvec.init(brute_force_by_keys_ratio="0.8") + + +class TestDbConfigLogValidation: + @run_in_subprocess + def test_log_type_valid(self): + zvec.init(log_type=LogType.CONSOLE) + + @run_in_subprocess + def test_log_type_invalid(self): + with pytest.raises(TypeError): + zvec.init(log_type="FILE") + with pytest.raises(TypeError): + zvec.init(log_type="") + with pytest.raises(TypeError): + zvec.init(log_type="invalid") + with pytest.raises(TypeError): + zvec.init(log_type=123) + + @run_in_subprocess + def test_log_level_valid(self): + zvec.init(log_level=LogLevel.ERROR) + + @run_in_subprocess + def test_log_level_invalid(self): + with pytest.raises(TypeError): + zvec.init(log_level="WARN") + with pytest.raises(TypeError): + zvec.init(log_level="") + with pytest.raises(TypeError): + zvec.init(log_level="invalid") + with pytest.raises(TypeError): + zvec.init(log_level=123) + + @run_in_subprocess + def test_init_file_logger(self): + from pathlib import Path + import shutil + + temp_dir = tempfile.mkdtemp(prefix="log_test_") + abs_temp_dir = os.path.abspath(temp_dir) + + zvec.init( + log_level=LogLevel.DEBUG, + log_type=LogType.FILE, + log_dir=abs_temp_dir, + log_basename="test", + ) + + # assert logdir exist + log_dir = Path(abs_temp_dir) + assert log_dir.exists() + + # validate write log + col = zvec.create_and_open( + "/tmp/test/1", + zvec.CollectionSchema( + name="test", + vectors=zvec.VectorSchema( + dimension=4, + data_type=zvec.DataType.VECTOR_FP32, + name="image", + ), + ), + ) + col.insert(docs=[zvec.Doc(id="1", vectors={"image": [1.0, 2.0, 3.0, 4.0]})]) + assert any(log_dir.glob("test.*")) + + # clear + col.destroy() + shutil.rmtree(log_dir, ignore_errors=True) + + @run_in_subprocess + def test_log_file_size_invalid(self): + with pytest.raises(TypeError): + zvec.init(log_type=LogType.FILE, log_file_size="df") + + with pytest.raises(ValueError): + zvec.init(log_type=LogType.FILE, log_file_size=0) + + with pytest.raises(ValueError): + zvec.init(log_type=LogType.FILE, log_file_size=-1) + + @run_in_subprocess + def test_log_overdue_days_invalid(self): + with pytest.raises(TypeError): + zvec.init(log_type=LogType.FILE, log_overdue_days="df") + + with pytest.raises(ValueError): + zvec.init(log_type=LogType.FILE, log_overdue_days=0) + + with pytest.raises(ValueError): + zvec.init(log_type=LogType.FILE, log_overdue_days=-1) diff --git a/python/tests/test_collection.py b/python/tests/test_collection.py new file mode 100644 index 0000000..7eba2e2 --- /dev/null +++ b/python/tests/test_collection.py @@ -0,0 +1,1218 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import pytest +import zvec +from zvec import ( + Collection, + CollectionOption, + DataType, + Doc, + FieldSchema, + HnswIndexParam, + IndexOption, + IndexType, + InvertIndexParam, + LogLevel, + LogType, + OptimizeOption, + StatusCode, + Query, + VectorSchema, +) +from zvec.extension.multi_vector_reranker import ( + CallbackReRanker, + RrfReRanker, + WeightedReRanker, +) + +# ==================== Common ==================== + + +@pytest.fixture(scope="session") +def collection_schema(): + return zvec.CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", DataType.STRING, nullable=False, index_param=InvertIndexParam() + ), + FieldSchema("weight", DataType.FLOAT, nullable=True), + FieldSchema("height", DataType.INT32, nullable=True), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + VectorSchema( + "dense2", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + VectorSchema( + "sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() + ), + VectorSchema( + "sparse2", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() + ), + ], + ) + + +@pytest.fixture(scope="session") +def collection_option(): + return CollectionOption(read_only=False, enable_mmap=True) + + +@pytest.fixture +def single_doc(): + id = 0 + return Doc( + id=f"{id}", + fields={"id": id, "name": "test", "weight": 80.0, "height": id + 140}, + vectors={ + "dense": [id + 0.1] * 128, + "dense2": [id + 0.2] * 128, + "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, + "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, + }, + ) + + +@pytest.fixture +def multiple_docs(): + return [ + Doc( + id=f"{id}", + fields={"id": id, "name": "test", "weight": 80.0, "height": 210}, + vectors={ + "dense": [id + 0.1] * 128, + "dense2": [id + 0.2] * 128, + "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, + "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, + }, + ) + for id in range(1, 101) + ] + + +@pytest.fixture(scope="function") +def test_collection( + tmp_path_factory, collection_schema, collection_option +) -> Collection: + """ + Function-scoped fixture: creates and opens a collection. + Uses tmp_path_factory to ensure shared temp dir per class. + """ + # Create unique temp directory for this test class + temp_dir = tmp_path_factory.mktemp("zvec") + collection_path = temp_dir / "test_collection" + + coll = zvec.create_and_open( + path=str(collection_path), schema=collection_schema, option=collection_option + ) + + assert coll is not None, "Failed to create and open collection" + assert coll.path == str(collection_path) + assert coll.schema.name == collection_schema.name + assert list(coll.schema.fields) == list(collection_schema.fields) + assert list(coll.schema.vectors) == list(collection_schema.vectors) + assert coll.option.read_only == collection_option.read_only + assert coll.option.enable_mmap == collection_option.enable_mmap + + try: + yield coll + finally: + if hasattr(coll, "destroy") and coll is not None: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + +@pytest.fixture +def collection_with_single_doc(test_collection: Collection, single_doc) -> Collection: + # Setup: insert single doc + assert test_collection.stats.doc_count == 0 + result = test_collection.insert(single_doc) + assert bool(result) + assert result.ok() + assert test_collection.stats.doc_count == 1 + + yield test_collection + + # Teardown: delete single doc + test_collection.delete(single_doc.id) + assert test_collection.stats.doc_count == 0 + + +@pytest.fixture +def collection_with_multiple_docs( + test_collection: Collection, multiple_docs +) -> Collection: + # Setup: insert multiple docs + assert test_collection.stats.doc_count == 0 + result = test_collection.insert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + assert test_collection.stats.doc_count == len(multiple_docs) + + yield test_collection + + # Teardown: delete multiple docs + test_collection.delete([doc.id for doc in multiple_docs]) + + +# ==================== Tests ==================== + + +# ---------------------------- +# Config Test Case +# ---------------------------- +class TestConfig: + def test_config(self): + zvec.init(log_type=LogType.CONSOLE, log_level=LogLevel.ERROR, log_dir="./log") + + +# ---------------------------- +# Collection DDL Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionDDL: + def test_collection_stats(self, test_collection: Collection): + assert test_collection.stats is not None + stats = test_collection.stats + assert stats.doc_count == 0 + assert len(stats.index_completeness) == 4 + assert stats.index_completeness["dense"] == 1 + assert stats.index_completeness["dense2"] == 1 + assert stats.index_completeness["sparse"] == 1 + assert stats.index_completeness["sparse2"] == 1 + + +# ---------------------------- +# Collection Index DDL Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionIndexDDL: + def test_create_index(self, test_collection: Collection): + # before create + field_schema = test_collection.schema.field("weight") + assert field_schema is not None + assert field_schema.data_type == DataType.FLOAT + assert field_schema.name == "weight" + index_param = field_schema.index_param + assert index_param is None + + # create + test_collection.create_index( + field_name="weight", index_param=InvertIndexParam(), option=IndexOption() + ) + assert test_collection.schema is not None + field_schema = test_collection.schema.field("weight") + assert field_schema is not None + assert field_schema.data_type == DataType.FLOAT + assert field_schema.name == "weight" + + index_param = field_schema.index_param + assert index_param.type == IndexType.INVERT + assert index_param.enable_range_optimization is False + assert index_param.enable_extended_wildcard is False + + def test_drop_index(self, test_collection: Collection): + # before drop + field_schema = test_collection.schema.field("name") + assert field_schema is not None + assert field_schema.data_type == DataType.STRING + assert field_schema.name == "name" + index_param = field_schema.index_param + assert index_param.type == IndexType.INVERT + assert index_param.enable_range_optimization is False + assert index_param.enable_extended_wildcard is False + + # drop + test_collection.drop_index("name") + field_schema = test_collection.schema.field("name") + assert field_schema is not None + assert field_schema.data_type == DataType.STRING + assert field_schema.name == "name" + + # without index + index_param = field_schema.index_param + assert index_param is None + + def test_create_index_field_is_not_exist(self, test_collection: Collection): + with pytest.raises(Exception) as e: + test_collection.create_index( + field_name="not_exist", + index_param=InvertIndexParam(), + ) + + index_param = field_schema.index_param + assert index_param.type == IndexType.INVERT + assert index_param.enable_range_optimization is False + assert index_param.enable_extended_wildcard is False + + def test_drop_index(self, test_collection: Collection): + # before drop + field_schema = test_collection.schema.field("name") + assert field_schema is not None + assert field_schema.data_type == DataType.STRING + assert field_schema.name == "name" + index_param = field_schema.index_param + assert index_param.type == IndexType.INVERT + assert index_param.enable_range_optimization is False + assert index_param.enable_extended_wildcard is False + + # drop + test_collection.drop_index("name") + field_schema = test_collection.schema.field("name") + assert field_schema is not None + assert field_schema.data_type == DataType.STRING + assert field_schema.name == "name" + + # without index + index_param = field_schema.index_param + assert index_param is None + + def test_create_index_field_is_not_exist(self, test_collection: Collection): + with pytest.raises(Exception) as e: + test_collection.create_index( + field_name="not_exist", + index_param=InvertIndexParam(), + ) + + +# ---------------------------- +# Collection Column DDL Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionColumnDDL: + def test_create_column(self, test_collection: Collection): + # before create column + field_schema = test_collection.schema.field("age") + assert field_schema is None + + # create + test_collection.add_column(FieldSchema("age", DataType.INT32, nullable=True)) + + field_schema = test_collection.schema.field("age") + assert field_schema is not None + assert field_schema.data_type == DataType.INT32 + assert field_schema.name == "age" + assert field_schema.index_param is None + + def test_create_column_is_nullable(self, test_collection: Collection): + with pytest.raises(ValueError): + test_collection.add_column( + FieldSchema("age", DataType.INT32, nullable=False) + ) + + def test_drop_column(self, test_collection: Collection): + # before drop column + field_schema = test_collection.schema.field("id") + assert field_schema is not None + assert field_schema.data_type == DataType.INT64 + assert field_schema.name == "id" + index_param = field_schema.index_param + assert index_param is not None + assert index_param.type == IndexType.INVERT + + # drop + test_collection.drop_column("id") + field_schema = test_collection.schema.field("id") + assert field_schema is None + + def test_alert_column_to_rename(self, test_collection: Collection): + # before alert column + field_schema = test_collection.schema.field("id") + assert field_schema is not None + assert field_schema.data_type == DataType.INT64 + assert field_schema.name == "id" + index_param = field_schema.index_param + assert index_param is not None + assert index_param.type == IndexType.INVERT + assert index_param.enable_range_optimization is True + assert index_param.enable_extended_wildcard is False + + # alert rename + test_collection.alter_column("id", "doc_id") + + # validate old column + field_schema = test_collection.schema.field("id") + assert field_schema is None + # validate rename column + field_schema = test_collection.schema.field("doc_id") + assert field_schema is not None + assert field_schema.data_type == DataType.INT64 + assert field_schema.name == "doc_id" + assert field_schema.nullable is False + index_param = field_schema.index_param + assert index_param is not None + assert index_param.type == IndexType.INVERT + assert index_param.enable_range_optimization is True + assert index_param.enable_extended_wildcard is False + + def test_alert_column_to_modify_schema(self, test_collection: Collection): + # before alert column + field_schema = test_collection.schema.field("id") + assert field_schema is not None + assert field_schema.data_type == DataType.INT64 + assert field_schema.name == "id" + index_param = field_schema.index_param + assert index_param.type == IndexType.INVERT + + test_collection.alter_column( + old_name="id", + field_schema=FieldSchema("doc_id", DataType.UINT64, nullable=True), + ) + field_schema = test_collection.schema.field("doc_id") + assert field_schema is not None + assert field_schema.data_type == DataType.UINT64 + assert field_schema.name == "doc_id" + + def test_column_with_other_dtype(self, test_collection: Collection): + # only allow number type + test_collection.add_column(FieldSchema("age", DataType.INT32, nullable=True)) + + with pytest.raises(ValueError): + test_collection.add_column(FieldSchema("full_name", DataType.STRING)) + with pytest.raises(ValueError): + test_collection.drop_column("name") + with pytest.raises(ValueError): + test_collection.alter_column(old_name="name", new_name="full_name") + with pytest.raises(ValueError): + test_collection.alter_column( + old_name="name", field_schema=FieldSchema("full_name", DataType.STRING) + ) + + +# ---------------------------- +# Collection Optimize Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionOptimize: + def test_collection_optimize(self, test_collection: Collection): + test_collection.optimize(option=OptimizeOption()) + + +# ---------------------------- +# Collection Fetch Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionFetch: + def test_collection_fetch( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + result = collection_with_single_doc.fetch(ids=[single_doc.id]) + assert bool(result) + assert single_doc.id in result.keys() + + doc = result[single_doc.id] + assert doc is not None + assert doc.id == single_doc.id + assert set(doc.field_names()) == set(single_doc.field_names()) + for field_name in doc.field_names(): + if field_name in ["dense", "sparse"]: + continue + assert doc.field(field_name) == single_doc.field(field_name) + + def test_collection_fetch_contains_nodata_ids( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + ids = [doc.id for doc in multiple_docs] + no_data_key = "x" + ids_with_no_data = [no_data_key] + ids + result = collection_with_multiple_docs.fetch(ids=ids_with_no_data) + assert bool(result) + assert len(result) == len(ids) + assert no_data_key not in result + + +# ---------------------------- +# Collection Insert Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionInsert: + def test_collection_insert(self, test_collection, single_doc): + result = test_collection.insert(single_doc) + assert bool(result) + assert result.ok() + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_collection_insert_with_nullable_false_field(self, test_collection): + # id, name's nullable == False + # weight, height's nullable == True + + doc = Doc( + id="0", + fields={ + "id": 1, + "name": "test", + }, + vectors={ + "dense": [1 + 0.1] * 128, + "dense2": [1 + 0.2] * 128, + "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, + "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, + }, + ) + result = test_collection.insert(doc) + assert bool(result) + assert result.ok() + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_collection_insert_without_nullable_false_field(self, test_collection): + # id, name's nullable == False + # weight, height's nullable == True + + # without id, name + doc = Doc( + id="0", + vectors={ + "dense": [1 + 0.1] * 128, + "dense2": [1 + 0.2] * 128, + "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, + "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, + }, + ) + with pytest.raises(ValueError) as e: + # ValueError: Invalid doc: field[id] is required but not provided + test_collection.insert(doc) + assert "field[id] is required but not provided" in str(e.value) + + # without name + doc = Doc( + id="0", + fields={ + "id": 1, + }, + vectors={ + "dense": [1 + 0.1] * 128, + "dense2": [1 + 0.2] * 128, + "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, + "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, + }, + ) + with pytest.raises(ValueError) as e: + test_collection.insert(doc) + assert "field[name] is required but not provided" in str(e.value) + + def test_collection_insert_with_nullable_true_field(self, test_collection): + # id, name's nullable == False + # weight, height's nullable == True + + doc = Doc( + id="0", + fields={ + "id": 1, + "name": "test", + }, + vectors={ + "dense": [1 + 0.1] * 128, + "dense2": [1 + 0.2] * 128, + "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, + "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, + }, + ) + result = test_collection.insert(doc) + assert bool(result) + assert result.ok() + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + result = test_collection.fetch(ids=[doc.id]) + assert doc.id in result + ret = result[doc.id] + assert ret.field("id") == 1 + assert ret.field("name") == "test" + assert ret.field("weight") is None + assert ret.field("height") is None + + def test_collection_insert_batch(self, test_collection, multiple_docs): + result = test_collection.insert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == len(multiple_docs) + + def test_collection_insert_duplicate( + self, test_collection, single_doc, multiple_docs + ): + test_collection.insert(single_doc) + result = test_collection.insert(single_doc) + assert bool(result) + assert result.code() == StatusCode.ALREADY_EXISTS + + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + +# ---------------------------- +# Collection Update Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionUpdate: + def test_empty_collection_update( + self, test_collection: Collection, single_doc: Doc + ): + result = test_collection.update(single_doc) + assert bool(result) + assert result.code() == StatusCode.NOT_FOUND + + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + def test_collection_update_with_nullable_false_field( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + # id, name's nullable == False + # weight, height's nullable == True + + # update doc field id + doc = Doc( + id=single_doc.id, + fields={"id": single_doc.field("id") + 1}, + ) + result = collection_with_single_doc.update(doc) + assert bool(result) + assert result.ok() + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 1 + + # fetch + result = collection_with_single_doc.fetch(ids=[doc.id]) + assert doc.id in result + ret = result[doc.id] + assert ret.field("id") == doc.field("id") + assert ret.field("name") == single_doc.field("name") + assert ret.field("weight") == single_doc.field("weight") + assert ret.field("height") == single_doc.field("height") + + def test_collection_update_with_nullable_false_field_is_none( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + # id, name's nullable == False + # weight, height's nullable == True + + # update doc field id + doc = Doc( + id=single_doc.id, + fields={"id": None}, + ) + with pytest.raises(ValueError) as e: + # ValueError: Invalid doc: field[id] is required but its value is null + collection_with_single_doc.update(doc) + + doc = Doc( + id=single_doc.id, + fields={"id": single_doc.field("id") + 1, "weight": None}, + ) + + result = collection_with_single_doc.update(doc) + assert bool(result) + assert result.ok() + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 1 + + ret = collection_with_single_doc.fetch(ids=[doc.id]) + assert doc.id in ret + ret = ret[doc.id] + assert ret.field("id") == doc.field("id") + assert ret.field("name") == single_doc.field("name") + assert ret.field("weight") is None + assert ret.field("height") == single_doc.field("height") + + def test_collection_update_without_nullable_false_field( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + # id, name's nullable == False + # weight, height's nullable == True + + # update doc field weight + doc = Doc( + id=single_doc.id, + fields={"weight": single_doc.field("weight") + 1}, + ) + result = collection_with_single_doc.update(doc) + assert bool(result) + assert result.ok() + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 1 + + # fetch + ret = collection_with_single_doc.fetch(ids=[doc.id]) + assert doc.id in ret + ret = ret[doc.id] + assert ret.field("id") == single_doc.field("id") + assert ret.field("name") == single_doc.field("name") + assert ret.field("weight") == doc.field("weight") + assert ret.field("height") == single_doc.field("height") + + def test_collection_update_without_nullable_false_field_set_null( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + # id, name's nullable == False + # weight, height's nullable == True + + # update doc field weight is None + doc = Doc( + id=single_doc.id, + fields={"weight": None}, + ) + result = collection_with_single_doc.update(doc) + assert bool(result) + assert result.ok() + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 1 + + # fetch + ret = collection_with_single_doc.fetch(ids=[doc.id]) + assert doc.id in ret + ret = ret[doc.id] + assert ret.field("id") == single_doc.field("id") + assert ret.field("name") == single_doc.field("name") + assert ret.field("weight") is None + assert ret.field("height") == single_doc.field("height") + + def test_empty_collection_update_batch( + self, test_collection: Collection, multiple_docs + ): + result = test_collection.update(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.code() == StatusCode.NOT_FOUND + + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == 0 + + def test_collection_update( + self, collection_with_single_doc: Collection, single_doc + ): + result = collection_with_single_doc.update(single_doc) + assert bool(result) == 1 + assert result.ok() + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_collection_update_batch( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.update(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + + stats = collection_with_multiple_docs.stats + assert stats is not None + assert stats.doc_count == len(multiple_docs) + + +# ---------------------------- +# Collection Upsert Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionUpsert: + def test_empty_collection_upsert(self, test_collection: Collection, single_doc): + result = test_collection.upsert(single_doc) + assert bool(result) + assert result.ok() + + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_empty_collection_upsert_batch( + self, test_collection: Collection, multiple_docs + ): + result = test_collection.upsert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + + stats = test_collection.stats + assert stats is not None + assert stats.doc_count == len(multiple_docs) + + def test_collection_upsert( + self, collection_with_single_doc: Collection, single_doc, multiple_docs + ): + # doc is existing + # upsert => update + result = collection_with_single_doc.upsert(single_doc) + assert bool(result) + assert result.ok() + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_collection_upsert_batch( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + # doc is existing + # upsert => update + result = collection_with_multiple_docs.upsert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + + stats = collection_with_multiple_docs.stats + assert stats is not None + assert stats.doc_count == len(multiple_docs) + + +# ---------------------------- +# Collection Upsert Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionDelete: + def test_empty_collection_delete(self, test_collection: Collection, single_doc): + result = test_collection.delete(single_doc.id) + assert bool(result) + assert result.code() == StatusCode.NOT_FOUND + + def test_empty_collection_delete_batch( + self, test_collection: Collection, multiple_docs + ): + result = test_collection.delete([doc.id for doc in multiple_docs]) + assert len(result) == len(multiple_docs) + for item in result: + assert item.code() == StatusCode.NOT_FOUND + + def test_collection_delete( + self, collection_with_single_doc: Collection, single_doc + ): + result = collection_with_single_doc.delete(single_doc.id) + assert bool(result) + assert result.ok() + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 0 + + result = collection_with_single_doc.insert(single_doc) + assert bool(result) + assert result.ok() + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_collection_delete_batch( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.delete([doc.id for doc in multiple_docs]) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + stats = collection_with_multiple_docs.stats + assert stats is not None + assert stats.doc_count == 0 + + def test_collection_delete_by_filter( + self, collection_with_single_doc: Collection, single_doc + ): + collection_with_single_doc.delete_by_filter( + filter=f"height={single_doc.field('height')}" + ) + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 0 + + def test_collection_delete_by_filter_invert_field( + self, collection_with_single_doc: Collection, single_doc + ): + collection_with_single_doc.delete_by_filter( + filter=f"id={single_doc.field('id')}" + ) + stats = collection_with_single_doc.stats + assert stats is not None + assert stats.doc_count == 0 + + +# ---------------------------- +# Collection Upsert Test Case +# ---------------------------- +@pytest.mark.usefixtures("test_collection") +class TestCollectionQuery: + def test_empty_collection_query(self, test_collection: Collection): + result = test_collection.query() + assert len(result) == 0 + + def test_collection_query(self, collection_with_single_doc: Collection, single_doc): + result = collection_with_single_doc.query() + assert len(result) == 1 + doc = result[0] + assert doc.id == single_doc.id + assert "dense" not in doc.field_names() + assert "sparse" not in doc.field_names() + field_without_vector = single_doc.field_names() + assert set(doc.field_names()) == set(field_without_vector) + for name in field_without_vector: + assert doc.field(name) == single_doc.field(name) + + def test_collection_query_with_include_vector( + self, collection_with_single_doc: Collection, single_doc + ): + result = collection_with_single_doc.query(include_vector=True) + assert len(result) == 1 + doc = result[0] + assert doc.vector("dense") is not None + assert doc.vector("sparse") is not None + + def test_collection_query_with_output_fields( + self, collection_with_single_doc: Collection, single_doc + ): + result = collection_with_single_doc.query(output_fields=["id", "name"]) + assert len(result) == 1 + doc = result[0] + assert doc.id == single_doc.id + assert len(doc.field_names()) == 2 + assert set(doc.field_names()) == {"id", "name"} + + def test_collection_query_with_topk( + self, collection_with_multiple_docs: Collection + ): + result = collection_with_multiple_docs.query() + assert len(result) == 10 + + result = collection_with_multiple_docs.query(topk=5) + assert len(result) == 5 + + def test_collection_query_with_range_filter_int_field( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + index = 10 + idx = multiple_docs[index].id + + result = collection_with_multiple_docs.query(filter=f"id>{idx}", topk=100) + assert len(result) == len(multiple_docs) - index - 1 + + result = collection_with_multiple_docs.query(filter=f"id>={idx}", topk=100) + assert len(result) == len(multiple_docs) - index + + result = collection_with_multiple_docs.query(filter=f"id<{idx}", topk=100) + assert len(result) == index + + result = collection_with_multiple_docs.query(filter=f"id<={idx}", topk=100) + assert len(result) == index + 1 + + result = collection_with_multiple_docs.query(filter=f"id={idx}", topk=100) + assert len(result) == 1 + + result = collection_with_multiple_docs.query(filter=f"id!={idx}", topk=100) + assert len(result) == len(multiple_docs) - 1 + + left, right = 10, 90 + l_id, r_id = multiple_docs[left].id, multiple_docs[right].id + result = collection_with_multiple_docs.query( + filter=f"id>{l_id} and id<{r_id}", topk=100 + ) + assert len(result) == right - left - 1 + + result = collection_with_multiple_docs.query( + filter=f"id>={l_id} and id<{r_id}", topk=100 + ) + assert len(result) == right - left + + result = collection_with_multiple_docs.query( + filter=f"id>={l_id} and id<={r_id}", topk=100 + ) + assert len(result) == right - left + 1 + + result = collection_with_multiple_docs.query( + filter=f"id<{l_id} or id>{r_id}", topk=100 + ) + assert len(result) == len(multiple_docs) - (right - left) - 1 + + result = collection_with_multiple_docs.query( + filter=f"id<={l_id} or id>{r_id}", topk=100 + ) + assert len(result) == len(multiple_docs) - (right - left) + + result = collection_with_multiple_docs.query( + filter=f"id<={l_id} or id>={r_id}", topk=100 + ) + assert len(result) == len(multiple_docs) - (right - left) + 1 + + result = collection_with_multiple_docs.query(filter="id in (1)", topk=100) + assert len(result) == 1 + + def test_collection_query_with_filter_not_in( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.query(filter="id not in (1)", topk=100) + assert len(result) == len(multiple_docs) - 1 + + def test_collection_with_error_query_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + query = Query( + field_name="dense", vector=multiple_docs[0].vector("dense"), param=[1, 2, 3] + ) + with pytest.raises(TypeError): + result = collection_with_multiple_docs.query( + query, filter="id in (1)", topk=100 + ) + + def test_collection_query_by_id( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.query( + Query(field_name="dense", id=multiple_docs[0].id) + ) + assert len(result) == 10 + + def test_collection_query_by_dense_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.query( + Query(field_name="dense", vector=multiple_docs[0].vector("dense")), + topk=10, + ) + assert len(result) > 0 + assert len(result) <= 10 + + def test_collection_query_by_sparse_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.query( + Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), + topk=10, + ) + assert len(result) > 0 + assert len(result) <= 10 + + def test_collection_query_by_dense_vector_with_filter( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.query( + Query(field_name="dense", vector=multiple_docs[0].vector("dense")), + topk=10, + filter="id > 50", + ) + assert len(result) > 0 + assert len(result) <= 10 + for doc in result: + assert int(doc.id) > 50 + + def test_collection_query_by_sparse_vector_with_filter( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.query( + Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), + topk=10, + filter="id > 50", + ) + assert len(result) > 0 + assert len(result) <= 10 + for doc in result: + assert int(doc.id) > 50 + + def test_collection_query_with_rrf_reranker_by_multi_dense_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """Test multi-vector query with RRF reranker on multiple dense vectors.""" + reranker = RrfReRanker(rank_constant=60) + result = collection_with_multiple_docs.query( + [ + Query(field_name="dense", vector=multiple_docs[0].vector("dense")), + Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")), + ], + topk=10, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 10 + # Results should have RRF-fused scores + for doc in result: + assert hasattr(doc, "score") + + def test_collection_query_with_rrf_reranker_by_multi_sparse_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """Test multi-vector query with RRF reranker on multiple sparse vectors.""" + reranker = RrfReRanker(rank_constant=60) + result = collection_with_multiple_docs.query( + [ + Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), + Query( + field_name="sparse2", + vector=multiple_docs[0].vector("sparse2"), + ), + ], + topk=10, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 10 + + def test_collection_query_with_rrf_reranker_by_hybrid_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """Test multi-vector query with RRF reranker combining dense + sparse.""" + reranker = RrfReRanker(rank_constant=60) + result = collection_with_multiple_docs.query( + [ + Query(field_name="dense", vector=multiple_docs[0].vector("dense")), + Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), + ], + topk=10, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 10 + + def test_collection_query_with_weighted_reranker_by_multi_dense_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """Test multi-vector query with Weighted reranker on multiple dense vectors.""" + weights = [0.6, 0.4] + reranker = WeightedReRanker(weights=weights) + result = collection_with_multiple_docs.query( + [ + Query(field_name="dense", vector=multiple_docs[0].vector("dense")), + Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")), + ], + topk=10, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 10 + + def test_collection_query_with_weighted_reranker_by_multi_sparse_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """Test multi-vector query with Weighted reranker on multiple sparse vectors.""" + weights = [0.6, 0.4] + reranker = WeightedReRanker(weights=weights) + result = collection_with_multiple_docs.query( + [ + Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), + Query( + field_name="sparse2", + vector=multiple_docs[0].vector("sparse2"), + ), + ], + topk=10, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 10 + + def test_collection_query_with_weighted_reranker_by_hybrid_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """Test multi-vector query with Weighted reranker combining dense + sparse.""" + weights = [0.7, 0.3] + reranker = WeightedReRanker(weights=weights) + result = collection_with_multiple_docs.query( + [ + Query(field_name="dense", vector=multiple_docs[0].vector("dense")), + Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), + ], + topk=10, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 10 + + def test_collection_query_with_callback_reranker_by_multi_dense_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """Test multi-vector query with CallbackReRanker (Python callback via C++).""" + callback_invoked = [] + + def my_rerank_callback(query_results, fields, topn): + callback_invoked.append(True) + all_docs = [] + for docs in query_results: + all_docs.extend(docs) + seen = set() + unique_docs = [] + for doc in all_docs: + if doc.pk() not in seen: + seen.add(doc.pk()) + unique_docs.append(doc) + unique_docs.sort(key=lambda d: d.score(), reverse=True) + return unique_docs[:topn] + + reranker = CallbackReRanker(callback=my_rerank_callback) + result = collection_with_multiple_docs.query( + [ + Query(field_name="dense", vector=multiple_docs[0].vector("dense")), + Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")), + ], + topk=10, + reranker=reranker, + ) + assert len(callback_invoked) == 1 + assert len(result) > 0 + assert len(result) <= 10 + + def test_collection_query_with_callback_reranker_by_hybrid_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """Test multi-vector query with CallbackReRanker combining dense + sparse.""" + + def my_rerank_callback(query_results, fields, topn): + all_docs = [] + for docs in query_results: + all_docs.extend(docs) + seen = set() + unique_docs = [] + for doc in all_docs: + if doc.pk() not in seen: + seen.add(doc.pk()) + unique_docs.append(doc) + unique_docs.sort(key=lambda d: d.score(), reverse=True) + return unique_docs[:topn] + + reranker = CallbackReRanker(callback=my_rerank_callback) + result = collection_with_multiple_docs.query( + [ + Query(field_name="dense", vector=multiple_docs[0].vector("dense")), + Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), + ], + topk=5, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 5 diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py new file mode 100644 index 0000000..b0e12ce --- /dev/null +++ b/python/tests/test_collection_diskann.py @@ -0,0 +1,600 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end collection tests for the DiskAnn index. + +Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn plugin. + +Two platform-level prerequisites are enforced at module import time: + +1. DiskAnn is currently built only for Linux x86_64 — other platforms are + skipped wholesale. +2. The DiskAnn backend lives in a *runtime-loaded* plugin + (``libzvec_diskann_plugin.so``). It must be loaded with ``RTLD_GLOBAL | + RTLD_NOW`` BEFORE ``import zvec`` so that the plugin's ``IndexFactory`` + singleton is unified with the one inside ``_zvec.so``. After ``import + zvec`` we must also call ``zvec.load_diskann_plugin()`` exactly once. + +If either prerequisite fails the whole module is skipped so the rest of the +test-suite is not affected. +""" + +from __future__ import annotations + +import math +import os +import platform +import sys + +import pytest + +# --------------------------------------------------------------------------- # +# Platform gating (must happen BEFORE we touch zvec). +# --------------------------------------------------------------------------- # +pytestmark = pytest.mark.skipif( + not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")), + reason="DiskAnn plugin is only supported on Linux x86_64", +) + +# Promote all symbols in subsequently-loaded DSOs to the global namespace and +# resolve relocations eagerly. This is REQUIRED so the DiskAnn plugin can see +# the ``IndexFactory`` singleton that lives in ``_zvec.so`` and vice versa. +# See: DiskAnn RTLD_GLOBAL + RTLD_NOW Requirement. +if sys.platform == "linux": + sys.setdlopenflags(sys.getdlopenflags() | os.RTLD_GLOBAL | os.RTLD_NOW) + +import zvec # noqa: E402 + +from zvec import ( # noqa: E402 + Collection, + CollectionOption, + DataType, + DiskAnnIndexParam, + DiskAnnQueryParam, + Doc, + FieldSchema, + MetricType, + Query, + VectorSchema, +) +from zvec.typing import QuantizeType # noqa: E402 + + +@pytest.fixture(scope="session") +def diskann_collection_schema(): + """Create a collection schema with a DiskAnn index.""" + return zvec.CollectionSchema( + name="test_diskann_collection", + fields=[ + FieldSchema("id", DataType.INT64, nullable=False), + FieldSchema("name", DataType.STRING, nullable=False), + ], + vectors=[ + VectorSchema( + "embedding", + DataType.VECTOR_FP32, + dimension=128, + index_param=DiskAnnIndexParam( + metric_type=MetricType.L2, + max_degree=64, + list_size=100, + pq_chunk_num=0, + quantize_type=QuantizeType.UNDEFINED, + ), + ), + ], + ) + + +@pytest.fixture(scope="session") +def collection_option(): + """Create collection options.""" + return CollectionOption(read_only=False, enable_mmap=True) + + +@pytest.fixture +def single_doc(): + """Create a single document for testing.""" + return Doc( + id="0", + fields={"id": 0, "name": "test_doc_0"}, + vectors={"embedding": [0.1 + i * 0.01 for i in range(128)]}, + ) + + +@pytest.fixture +def multiple_docs(): + """Create multiple documents for testing.""" + return [ + Doc( + id=f"{i}", + fields={"id": i, "name": f"test_doc_{i}"}, + vectors={"embedding": [i * 0.1 + j * 0.01 for j in range(128)]}, + ) + for i in range(1, 101) + ] + + +@pytest.fixture(scope="function") +def diskann_collection( + tmp_path_factory, diskann_collection_schema, collection_option +) -> Collection: + """ + Function-scoped fixture: creates and opens a collection with DiskAnn index. + """ + temp_dir = tmp_path_factory.mktemp("zvec_diskann") + collection_path = temp_dir / "test_diskann_collection" + + coll = zvec.create_and_open( + path=str(collection_path), + schema=diskann_collection_schema, + option=collection_option, + ) + + assert coll is not None, "Failed to create and open DiskAnn collection" + assert coll.path == str(collection_path) + assert coll.schema.name == diskann_collection_schema.name + + try: + yield coll + finally: + if hasattr(coll, "destroy") and coll is not None: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + +@pytest.fixture +def collection_with_single_doc( + diskann_collection: Collection, single_doc: Doc +) -> Collection: + """Setup: insert single doc into collection.""" + assert diskann_collection.stats.doc_count == 0 + result = diskann_collection.insert(single_doc) + assert bool(result) + assert result.ok() + assert diskann_collection.stats.doc_count == 1 + + yield diskann_collection + + # Teardown: delete single doc + diskann_collection.delete(single_doc.id) + assert diskann_collection.stats.doc_count == 0 + + +@pytest.fixture +def collection_with_multiple_docs( + diskann_collection: Collection, multiple_docs: list[Doc] +) -> Collection: + """Setup: insert multiple docs into collection.""" + assert diskann_collection.stats.doc_count == 0 + result = diskann_collection.insert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + assert diskann_collection.stats.doc_count == len(multiple_docs) + + yield diskann_collection + + # Teardown: delete multiple docs + diskann_collection.delete([doc.id for doc in multiple_docs]) + + +# ==================== Tests ==================== + + +@pytest.mark.usefixtures("diskann_collection") +class TestDiskAnnCollectionCreation: + """Test DiskAnn collection creation and schema validation.""" + + def test_collection_creation( + self, diskann_collection: Collection, diskann_collection_schema + ): + """Test that collection is created with correct schema.""" + assert diskann_collection is not None + assert diskann_collection.schema.name == diskann_collection_schema.name + assert len(diskann_collection.schema.fields) == len( + diskann_collection_schema.fields + ) + assert len(diskann_collection.schema.vectors) == len( + diskann_collection_schema.vectors + ) + + def test_vector_schema_validation(self, diskann_collection: Collection): + """Test that vector schema has correct DiskAnn configuration.""" + vector_schema = diskann_collection.schema.vector("embedding") + assert vector_schema is not None + assert vector_schema.name == "embedding" + assert vector_schema.data_type == DataType.VECTOR_FP32 + assert vector_schema.dimension == 128 + + index_param = vector_schema.index_param + assert index_param is not None + assert index_param.metric_type == MetricType.L2 + assert index_param.max_degree == 64 + assert index_param.list_size == 100 + assert index_param.pq_chunk_num == 0 + + def test_collection_stats(self, diskann_collection: Collection): + """Test initial collection statistics.""" + stats = diskann_collection.stats + assert stats is not None + assert stats.doc_count == 0 + assert len(stats.index_completeness) == 1 + assert stats.index_completeness["embedding"] == 1 + + +@pytest.mark.usefixtures("diskann_collection") +class TestDiskAnnCollectionInsert: + """Test document insertion into DiskAnn collection.""" + + def test_insert_single_doc(self, diskann_collection: Collection, single_doc: Doc): + """Test inserting a single document.""" + result = diskann_collection.insert(single_doc) + assert bool(result) + assert result.ok() + + stats = diskann_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_insert_multiple_docs( + self, diskann_collection: Collection, multiple_docs: list[Doc] + ): + """Test inserting multiple documents.""" + result = diskann_collection.insert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + + stats = diskann_collection.stats + assert stats is not None + assert stats.doc_count == len(multiple_docs) + + +@pytest.mark.usefixtures("diskann_collection") +class TestDiskAnnCollectionFetch: + """Test document fetching from DiskAnn collection.""" + + def test_fetch_single_doc( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + """Test fetching a single document by ID.""" + result = collection_with_single_doc.fetch(ids=[single_doc.id]) + assert bool(result) + assert single_doc.id in result.keys() + + doc = result[single_doc.id] + assert doc is not None + assert doc.id == single_doc.id + assert doc.field("id") == single_doc.field("id") + assert doc.field("name") == single_doc.field("name") + + def test_fetch_multiple_docs( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test fetching multiple documents by IDs.""" + ids = [doc.id for doc in multiple_docs[:10]] + result = collection_with_multiple_docs.fetch(ids=ids) + assert bool(result) + assert len(result) == len(ids) + + for doc_id in ids: + assert doc_id in result + doc = result[doc_id] + assert doc is not None + assert doc.id == doc_id + + def test_fetch_nonexistent_doc(self, collection_with_single_doc: Collection): + """Test fetching a non-existent document.""" + result = collection_with_single_doc.fetch(ids=["nonexistent_id"]) + assert len(result) == 0 + + +@pytest.mark.usefixtures("diskann_collection") +class TestDiskAnnCollectionQuery: + """Test vector search queries on DiskAnn collection.""" + + def test_query_by_vector( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying by vector with DiskAnn index.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=100), + ) + + result = collection_with_multiple_docs.query(queries=query, topk=10) + assert len(result) > 0 + assert len(result) <= 10 + + # First result should be the query document itself (or very close) + first_doc = result[0] + assert first_doc is not None + assert first_doc.id is not None + + def test_query_by_id( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying by document ID with DiskAnn index.""" + query = Query( + field_name="embedding", + id=multiple_docs[0].id, + param=DiskAnnQueryParam(list_size=100), + ) + + result = collection_with_multiple_docs.query(queries=query, topk=10) + assert len(result) > 0 + assert len(result) <= 10 + + def test_query_with_different_list_size( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with different list_size parameter values.""" + query_vector = multiple_docs[0].vector("embedding") + + # Test with list_size=50 + query_small = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=50), + ) + result_small = collection_with_multiple_docs.query(queries=query_small, topk=10) + assert len(result_small) > 0 + + # Test with list_size=200 + query_large = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=200), + ) + result_large = collection_with_multiple_docs.query(queries=query_large, topk=10) + assert len(result_large) > 0 + + def test_query_with_topk( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with different topk values.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=100), + ) + + # Test topk=5 + result_5 = collection_with_multiple_docs.query(queries=query, topk=5) + assert len(result_5) <= 5 + + # Test topk=20 + result_20 = collection_with_multiple_docs.query(queries=query, topk=20) + assert len(result_20) <= 20 + + def test_query_with_filter( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with filter conditions.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=100), + ) + + # Query with id filter + result = collection_with_multiple_docs.query( + queries=query, topk=10, filter="id < 50" + ) + assert len(result) > 0 + for doc in result: + assert doc.field("id") < 50 + + def test_query_with_output_fields( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with specific output fields.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=100), + ) + + result = collection_with_multiple_docs.query( + queries=query, topk=10, output_fields=["id", "name"] + ) + assert len(result) > 0 + + first_doc = result[0] + assert "id" in first_doc.field_names() + assert "name" in first_doc.field_names() + + def test_query_with_include_vector( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with vector data included in results.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=100), + ) + + result = collection_with_multiple_docs.query( + queries=query, topk=10, include_vector=True + ) + assert len(result) > 0 + + first_doc = result[0] + assert first_doc.vector("embedding") is not None + assert len(first_doc.vector("embedding")) == 128 + + +@pytest.mark.usefixtures("diskann_collection") +class TestDiskAnnCollectionUpdate: + """Test document update in DiskAnn collection.""" + + def test_update_doc_fields( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + """Test updating document fields.""" + updated_doc = Doc( + id=single_doc.id, + fields={"id": single_doc.field("id"), "name": "updated_name"}, + ) + + result = collection_with_single_doc.update(updated_doc) + assert bool(result) + assert result.ok() + + # Verify update + fetched = collection_with_single_doc.fetch(ids=[single_doc.id]) + assert single_doc.id in fetched + doc = fetched[single_doc.id] + assert doc.field("name") == "updated_name" + + def test_update_doc_vector( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + """Test updating document vector.""" + new_vector = [0.5 + i * 0.01 for i in range(128)] + updated_doc = Doc( + id=single_doc.id, + vectors={"embedding": new_vector}, + ) + + result = collection_with_single_doc.update(updated_doc) + assert bool(result) + assert result.ok() + + # Verify update + fetched = collection_with_single_doc.fetch( + ids=[single_doc.id], + ) + assert single_doc.id in fetched + doc = fetched[single_doc.id] + assert doc.vector("embedding") is not None + embedding = doc.vector("embedding") + assert len(embedding) == 128 + # Verify vector values are approximately equal (float comparison) + for i in range(128): + assert math.isclose(embedding[i], new_vector[i], rel_tol=1e-5) + + +@pytest.mark.usefixtures("diskann_collection") +class TestDiskAnnCollectionDelete: + """Test document deletion from DiskAnn collection.""" + + def test_delete_single_doc( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + """Test deleting a single document.""" + result = collection_with_single_doc.delete(single_doc.id) + assert bool(result) + assert result.ok() + + stats = collection_with_single_doc.stats + assert stats.doc_count == 0 + + def test_delete_multiple_docs( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test deleting multiple documents.""" + ids_to_delete = [doc.id for doc in multiple_docs[:10]] + result = collection_with_multiple_docs.delete(ids_to_delete) + assert len(result) == len(ids_to_delete) + for item in result: + assert item.ok() + + stats = collection_with_multiple_docs.stats + assert stats.doc_count == len(multiple_docs) - len(ids_to_delete) + + +@pytest.mark.usefixtures("diskann_collection") +class TestDiskAnnCollectionOptimizeAndReopen: + """Test collection optimize and reopen functionality.""" + + def test_optimize_close_reopen_and_query( + self, + tmp_path_factory, + diskann_collection_schema, + collection_option, + multiple_docs: list[Doc], + ): + """Test inserting 100 docs, optimize, close, reopen and query.""" + # Create collection and insert 100 documents + temp_dir = tmp_path_factory.mktemp("zvec_diskann_optimize") + collection_path = temp_dir / "test_optimize_collection" + + coll = zvec.create_and_open( + path=str(collection_path), + schema=diskann_collection_schema, + option=collection_option, + ) + + assert coll is not None + assert coll.stats.doc_count == 0 + + # Insert 100 documents + result = coll.insert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + assert coll.stats.doc_count == len(multiple_docs) + + # Call optimize + from zvec import OptimizeOption + + coll.optimize(option=OptimizeOption()) + + # Verify data is still accessible after optimize + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=100), + ) + result_before_close = coll.query(queries=query, topk=10) + assert len(result_before_close) > 0 + + # Close collection (destroy will close it) + collection_path_str = str(collection_path) + del coll + + # Reopen collection + reopened_coll = zvec.open(path=collection_path_str, option=collection_option) + assert reopened_coll is not None + assert reopened_coll.stats.doc_count == len(multiple_docs) + + # Execute query on reopened collection + query_after_reopen = Query( + field_name="embedding", + vector=query_vector, + param=DiskAnnQueryParam(list_size=100), + ) + result_after_reopen = reopened_coll.query(queries=query_after_reopen, topk=10) + assert len(result_after_reopen) > 0 + assert len(result_after_reopen) <= 10 + + # Verify query results are valid + first_doc = result_after_reopen[0] + assert first_doc is not None + assert first_doc.id is not None + assert first_doc.field("id") is not None + assert first_doc.field("name") is not None + + # Cleanup + reopened_coll.destroy() diff --git a/python/tests/test_collection_fts.py b/python/tests/test_collection_fts.py new file mode 100644 index 0000000..6cc884b --- /dev/null +++ b/python/tests/test_collection_fts.py @@ -0,0 +1,188 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end tests for FTS-only collections (no vector field). + +The schema validation rule "must have at least one vector field" has been +lifted; these tests pin the new behavior so insert / query / delete / +optimize all work on a vector-less collection. +""" + +from __future__ import annotations + +import pytest +import zvec +from zvec import ( + Collection, + CollectionOption, + DataType, + Doc, + FieldSchema, + FtsIndexParam, + OptimizeOption, +) +from zvec.model.param.query import Fts, Query + + +# ==================== Fixtures ==================== + + +@pytest.fixture(scope="function") +def fts_collection(tmp_path_factory) -> Collection: + """FTS-only collection: a STRING field for forward + an FTS-indexed STRING.""" + temp_dir = tmp_path_factory.mktemp("zvec_fts_only") + collection_path = temp_dir / "fts_collection" + + schema = zvec.CollectionSchema( + name="fts_only", + fields=[ + FieldSchema("title", DataType.STRING, nullable=False), + FieldSchema( + "content", + DataType.STRING, + nullable=False, + index_param=FtsIndexParam( + tokenizer_name="standard", + filters=["lowercase"], + ), + ), + ], + # vectors omitted on purpose — schema validation must accept this. + ) + + coll = zvec.create_and_open( + path=str(collection_path), + schema=schema, + option=CollectionOption(read_only=False, enable_mmap=True), + ) + assert coll is not None + + try: + yield coll + finally: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + +def _make_docs() -> list[Doc]: + """5-doc corpus where 4 contain 'hello' and doc 4 is the only outlier.""" + return [ + Doc(id="pk_0", fields={"title": "intro", "content": "hello world"}), + Doc(id="pk_1", fields={"title": "guide", "content": "hello foo bar"}), + Doc(id="pk_2", fields={"title": "tips", "content": "hello baz"}), + Doc(id="pk_3", fields={"title": "more", "content": "hello hello"}), + Doc(id="pk_4", fields={"title": "other", "content": "nothing relevant"}), + ] + + +def _fts_query(coll: Collection, term: str) -> list[Doc]: + """Run a single-term FTS match query against the `content` field.""" + return coll.query( + queries=Query(field_name="content", fts=Fts(match_string=term)), + topk=10, + ) + + +# ==================== Tests ==================== + + +class TestFtsOnlyCollectionSchema: + def test_create_and_open_without_vectors(self, fts_collection: Collection): + """Schema with zero vector fields must be accepted by validate().""" + assert fts_collection.schema.name == "fts_only" + assert {f.name for f in fts_collection.schema.fields} == {"title", "content"} + # Empty vectors is the whole point of the test. + assert list(fts_collection.schema.vectors) == [] + assert fts_collection.stats.doc_count == 0 + + def test_create_schema_omitting_vectors_kwarg(self): + """Constructing CollectionSchema without `vectors=` argument is valid.""" + schema = zvec.CollectionSchema( + name="bare_fts", + fields=[ + FieldSchema( + "content", + DataType.STRING, + nullable=False, + index_param=FtsIndexParam(), + ), + ], + ) + assert list(schema.vectors) == [] + assert {f.name for f in schema.fields} == {"content"} + + +class TestFtsOnlyCollectionLifecycle: + def test_insert_and_fts_query(self, fts_collection: Collection): + """FTS-only collection supports insert + FTS query end-to-end.""" + results = fts_collection.insert(_make_docs()) + assert all(r.ok() for r in results) + assert fts_collection.stats.doc_count == 5 + + hits = _fts_query(fts_collection, "hello") + assert len(hits) == 4 + assert {doc.id for doc in hits} == {"pk_0", "pk_1", "pk_2", "pk_3"} + + # Term that nothing in the surviving corpus contains. + assert _fts_query(fts_collection, "missing_term_xyz") == [] + + def test_delete_then_query(self, fts_collection: Collection): + """Tombstone filter must drop deleted docs from FTS results.""" + fts_collection.insert(_make_docs()) + statuses = fts_collection.delete(["pk_0", "pk_4"]) + assert all(s.ok() for s in statuses) + assert fts_collection.stats.doc_count == 3 + + hits = _fts_query(fts_collection, "hello") + assert len(hits) == 3 + assert {doc.id for doc in hits} == {"pk_1", "pk_2", "pk_3"} + # pk_4's unique term is filtered out post-delete. + assert _fts_query(fts_collection, "nothing") == [] + + def test_optimize_rebuilds_fts(self, fts_collection: Collection): + """Optimize with >30% deletes triggers ReduceFts; recall unchanged.""" + fts_collection.insert(_make_docs()) + # 40% delete ratio — above COMPACT_DELETE_RATIO_THRESHOLD=0.3, so + # build_compact_task picks the rebuild path and ReduceFts runs. + fts_collection.delete(["pk_0", "pk_4"]) + + before = {doc.id for doc in _fts_query(fts_collection, "hello")} + assert before == {"pk_1", "pk_2", "pk_3"} + + fts_collection.optimize(option=OptimizeOption()) + assert fts_collection.stats.doc_count == 3 + + after = {doc.id for doc in _fts_query(fts_collection, "hello")} + assert after == before + assert _fts_query(fts_collection, "nothing") == [] + + +class TestFtsOnlyCollectionQueryValidation: + def test_vector_query_rejected(self, fts_collection: Collection): + """Vector query on a no-vector collection must raise.""" + with pytest.raises(ValueError, match="No vector field found"): + fts_collection.query( + queries=Query(field_name="content", vector=[0.1, 0.2, 0.3]), + topk=5, + ) + + def test_id_query_rejected(self, fts_collection: Collection): + """ID-based query on a no-vector collection must raise.""" + fts_collection.insert(_make_docs()[:1]) + with pytest.raises(ValueError, match="No vector field found"): + fts_collection.query( + queries=Query(field_name="content", id="pk_0"), + topk=5, + ) diff --git a/python/tests/test_collection_fts_vector_hybrid.py b/python/tests/test_collection_fts_vector_hybrid.py new file mode 100644 index 0000000..0478b89 --- /dev/null +++ b/python/tests/test_collection_fts_vector_hybrid.py @@ -0,0 +1,391 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for FTS + vector hybrid retrieval via multi-query with reranker.""" + +from __future__ import annotations + +import pytest +import zvec +from zvec import ( + Collection, + CollectionOption, + DataType, + Doc, + FieldSchema, + FtsIndexParam, + HnswIndexParam, + VectorSchema, +) +from zvec.extension.multi_vector_reranker import RrfReRanker, WeightedReRanker +from zvec.model.param.query import Fts, Query + + +DIM = 16 + + +# ==================== Fixtures ==================== + + +@pytest.fixture(scope="function") +def hybrid_collection(tmp_path_factory) -> Collection: + """Collection with one vector field + one FTS field.""" + temp_dir = tmp_path_factory.mktemp("zvec_hybrid") + collection_path = temp_dir / "hybrid_collection" + + schema = zvec.CollectionSchema( + name="hybrid_test", + fields=[ + FieldSchema("title", DataType.STRING, nullable=False), + FieldSchema( + "content", + DataType.STRING, + nullable=False, + index_param=FtsIndexParam( + tokenizer_name="standard", + filters=["lowercase"], + ), + ), + ], + vectors=[ + VectorSchema( + "embedding", + DataType.VECTOR_FP32, + dimension=DIM, + index_param=HnswIndexParam(), + ), + ], + ) + + coll = zvec.create_and_open( + path=str(collection_path), + schema=schema, + option=CollectionOption(read_only=False, enable_mmap=True), + ) + assert coll is not None + + try: + yield coll + finally: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + +def _make_docs() -> list[Doc]: + """Corpus with both text content and vectors. + + Docs 0-2: AI/ML topic, vectors clustered in one region. + Docs 3-4: retrieval topic, vectors clustered in another region. + Doc 5: unrelated topic. + """ + # AI cluster vectors + ai_vec = [1.0] * 8 + [0.0] * 8 + # Retrieval cluster vectors + ret_vec = [0.0] * 8 + [1.0] * 8 + # Unrelated vector + other_vec = [0.5] * 16 + + return [ + Doc( + id="pk_0", + fields={ + "title": "ML Intro", + "content": "machine learning is a branch of artificial intelligence", + }, + vectors={"embedding": ai_vec}, + ), + Doc( + id="pk_1", + fields={ + "title": "Deep Learning", + "content": "deep learning uses neural networks for pattern recognition", + }, + vectors={"embedding": [0.9] * 8 + [0.1] * 8}, + ), + Doc( + id="pk_2", + fields={ + "title": "NLP", + "content": "natural language processing handles text with artificial intelligence", + }, + vectors={"embedding": [0.8] * 8 + [0.2] * 8}, + ), + Doc( + id="pk_3", + fields={ + "title": "Search Engine", + "content": "search engine uses inverted index for text retrieval", + }, + vectors={"embedding": ret_vec}, + ), + Doc( + id="pk_4", + fields={ + "title": "Vector DB", + "content": "vector database enables similarity retrieval and search", + }, + vectors={"embedding": [0.1] * 8 + [0.9] * 8}, + ), + Doc( + id="pk_5", + fields={ + "title": "Cooking", + "content": "baking bread requires flour water yeast and salt", + }, + vectors={"embedding": other_vec}, + ), + ] + + +@pytest.fixture(scope="function") +def hybrid_collection_with_docs(hybrid_collection: Collection) -> Collection: + """Hybrid collection pre-populated with test documents.""" + results = hybrid_collection.insert(_make_docs()) + assert all(r.ok() for r in results) + return hybrid_collection + + +# ==================== Tests ==================== + + +class TestFtsVectorHybridQuery: + """Test FTS + vector hybrid retrieval using multi-query with RRF reranker.""" + + def test_hybrid_fts_and_vector_basic(self, hybrid_collection_with_docs: Collection): + """FTS + vector multi-query with RRF reranker returns results.""" + reranker = RrfReRanker(rank_constant=60) + result = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="retrieval")), + Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8), + ], + topk=5, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 5 + # Results should have scores + for doc in result: + assert doc.score > 0 + + def test_hybrid_fts_and_vector_ranking( + self, hybrid_collection_with_docs: Collection + ): + """Docs relevant in both FTS and vector should rank higher.""" + reranker = RrfReRanker(rank_constant=60) + # FTS: "retrieval search" matches pk_3, pk_4 + # Vector: ret_vec cluster matches pk_3, pk_4 + # Both signals agree: pk_3 and pk_4 should rank top + result = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="retrieval search")), + Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8), + ], + topk=5, + reranker=reranker, + ) + top_ids = {doc.id for doc in result[:3]} + assert "pk_3" in top_ids or "pk_4" in top_ids + + def test_hybrid_scores_descending(self, hybrid_collection_with_docs: Collection): + """Hybrid query results must be sorted by score descending.""" + reranker = RrfReRanker(rank_constant=60) + result = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="intelligence")), + Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8), + ], + topk=6, + reranker=reranker, + ) + assert len(result) >= 2 + scores = [doc.score for doc in result] + assert scores == sorted(scores, reverse=True) + + def test_hybrid_with_filter(self, hybrid_collection_with_docs: Collection): + """Hybrid query respects SQL filter.""" + reranker = RrfReRanker(rank_constant=60) + result = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="learning")), + Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8), + ], + topk=10, + reranker=reranker, + filter="title like '%Learning%'", + ) + for doc in result: + assert "Learning" in doc.fields["title"] + + def test_hybrid_fts_no_match_still_returns_vector_results( + self, hybrid_collection_with_docs: Collection + ): + """When FTS matches nothing, vector results still appear.""" + reranker = RrfReRanker(rank_constant=60) + result = hybrid_collection_with_docs.query( + queries=[ + Query( + field_name="content", + fts=Fts(match_string="nonexistent_term_xyz"), + ), + Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8), + ], + topk=5, + reranker=reranker, + ) + # Vector query alone should still produce results + assert len(result) > 0 + + def test_hybrid_query_string_syntax(self, hybrid_collection_with_docs: Collection): + """Hybrid query works with FTS query_string (advanced syntax).""" + reranker = RrfReRanker(rank_constant=60) + result = hybrid_collection_with_docs.query( + queries=[ + Query( + field_name="content", + fts=Fts(query_string="artificial AND intelligence"), + ), + Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8), + ], + topk=5, + reranker=reranker, + ) + assert len(result) > 0 + # pk_0 and pk_2 contain "artificial intelligence" + hit_ids = {doc.id for doc in result} + assert "pk_0" in hit_ids or "pk_2" in hit_ids + + +class TestFtsVectorHybridValidation: + """Test validation rules for FTS + vector hybrid queries.""" + + def test_hybrid_requires_reranker(self, hybrid_collection_with_docs: Collection): + """Multi-query with FTS + vector without reranker should raise.""" + with pytest.raises(ValueError, match="[Rr]eranker"): + hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="learning")), + Query(field_name="embedding", vector=[1.0] * DIM), + ], + topk=5, + ) + + def test_duplicate_field_name_allowed( + self, hybrid_collection_with_docs: Collection + ): + """Multi-query with duplicate field names is allowed and returns results.""" + reranker = RrfReRanker(rank_constant=60) + result = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="learning")), + Query(field_name="content", fts=Fts(match_string="intelligence")), + ], + topk=5, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 5 + + def test_multiple_vectors_allowed(self, hybrid_collection_with_docs: Collection): + """Two vector queries on the same field are allowed with a reranker.""" + reranker = RrfReRanker(rank_constant=60) + result = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="embedding", vector=[1.0] * DIM), + Query(field_name="embedding", vector=[0.5] * DIM), + ], + topk=5, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 5 + + +class TestFtsVectorHybridWeightedReranker: + """Test FTS + vector hybrid retrieval using WeightedReranker.""" + + def test_weighted_reranker_fts_and_vector( + self, hybrid_collection_with_docs: Collection + ): + """WeightedReranker correctly normalizes FTS scores alongside vector scores.""" + weights = [0.5, 0.5] + reranker = WeightedReRanker(weights=weights) + result = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="retrieval search")), + Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8), + ], + topk=5, + reranker=reranker, + ) + assert len(result) > 0 + assert len(result) <= 5 + for doc in result: + assert doc.score > 0 + + def test_weighted_reranker_scores_descending( + self, hybrid_collection_with_docs: Collection + ): + """WeightedReranker hybrid results are sorted by score descending.""" + weights = [0.4, 0.6] + reranker = WeightedReRanker(weights=weights) + result = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="intelligence")), + Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8), + ], + topk=6, + reranker=reranker, + ) + assert len(result) >= 2 + scores = [doc.score for doc in result] + assert scores == sorted(scores, reverse=True) + + def test_weighted_reranker_fts_weight_influence( + self, hybrid_collection_with_docs: Collection + ): + """Higher FTS weight should boost FTS-relevant docs in ranking.""" + # High FTS weight: FTS signal dominates + weights_fts_heavy = [0.9, 0.1] + reranker_fts = WeightedReRanker(weights=weights_fts_heavy) + result_fts = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="retrieval")), + Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8), + ], + topk=5, + reranker=reranker_fts, + ) + + # High vector weight: vector signal dominates + weights_vec_heavy = [0.1, 0.9] + reranker_vec = WeightedReRanker(weights=weights_vec_heavy) + result_vec = hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="retrieval")), + Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8), + ], + topk=5, + reranker=reranker_vec, + ) + + # Both should return results + assert len(result_fts) > 0 + assert len(result_vec) > 0 + # With FTS-heavy weight, FTS-relevant docs (pk_3, pk_4) should rank higher + fts_top = [doc.id for doc in result_fts[:2]] + vec_top = [doc.id for doc in result_vec[:2]] + # The rankings should differ due to weight difference + assert fts_top != vec_top or len(result_fts) == len(result_vec) == 1 diff --git a/python/tests/test_collection_hnsw_rabitq.py b/python/tests/test_collection_hnsw_rabitq.py new file mode 100644 index 0000000..4051e29 --- /dev/null +++ b/python/tests/test_collection_hnsw_rabitq.py @@ -0,0 +1,574 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import platform +import sys + +import pytest +import math +import zvec + +pytestmark = pytest.mark.skipif( + not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")), + reason="HNSW RaBitQ only supported on Linux x86_64", +) +from zvec import ( + Collection, + CollectionOption, + DataType, + Doc, + FieldSchema, + HnswRabitqIndexParam, + HnswRabitqQueryParam, + MetricType, + VectorSchema, + Query, +) + + +# ==================== Fixtures ==================== + + +@pytest.fixture(scope="session") +def hnsw_rabitq_collection_schema(): + """Create a collection schema with HNSW RaBitQ index.""" + return zvec.CollectionSchema( + name="test_hnsw_rabitq_collection", + fields=[ + FieldSchema("id", DataType.INT64, nullable=False), + FieldSchema("name", DataType.STRING, nullable=False), + ], + vectors=[ + VectorSchema( + "embedding", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswRabitqIndexParam( + metric_type=MetricType.L2, + m=16, + ef_construction=200, + total_bits=7, + num_clusters=64, + ), + ), + ], + ) + + +@pytest.fixture(scope="session") +def collection_option(): + """Create collection options.""" + return CollectionOption(read_only=False, enable_mmap=True) + + +@pytest.fixture +def single_doc(): + """Create a single document for testing.""" + return Doc( + id="0", + fields={"id": 0, "name": "test_doc_0"}, + vectors={"embedding": [0.1 + i * 0.01 for i in range(128)]}, + ) + + +@pytest.fixture +def multiple_docs(): + """Create multiple documents for testing.""" + return [ + Doc( + id=f"{i}", + fields={"id": i, "name": f"test_doc_{i}"}, + vectors={"embedding": [i * 0.1 + j * 0.01 for j in range(128)]}, + ) + for i in range(1, 101) + ] + + +@pytest.fixture(scope="function") +def hnsw_rabitq_collection( + tmp_path_factory, hnsw_rabitq_collection_schema, collection_option +) -> Collection: + """ + Function-scoped fixture: creates and opens a collection with HNSW RaBitQ index. + """ + temp_dir = tmp_path_factory.mktemp("zvec_hnsw_rabitq") + collection_path = temp_dir / "test_hnsw_rabitq_collection" + + coll = zvec.create_and_open( + path=str(collection_path), + schema=hnsw_rabitq_collection_schema, + option=collection_option, + ) + + assert coll is not None, "Failed to create and open HNSW RaBitQ collection" + assert coll.path == str(collection_path) + assert coll.schema.name == hnsw_rabitq_collection_schema.name + + try: + yield coll + finally: + if hasattr(coll, "destroy") and coll is not None: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + +@pytest.fixture +def collection_with_single_doc( + hnsw_rabitq_collection: Collection, single_doc: Doc +) -> Collection: + """Setup: insert single doc into collection.""" + assert hnsw_rabitq_collection.stats.doc_count == 0 + result = hnsw_rabitq_collection.insert(single_doc) + assert bool(result) + assert result.ok() + assert hnsw_rabitq_collection.stats.doc_count == 1 + + yield hnsw_rabitq_collection + + # Teardown: delete single doc + hnsw_rabitq_collection.delete(single_doc.id) + assert hnsw_rabitq_collection.stats.doc_count == 0 + + +@pytest.fixture +def collection_with_multiple_docs( + hnsw_rabitq_collection: Collection, multiple_docs: list[Doc] +) -> Collection: + """Setup: insert multiple docs into collection.""" + assert hnsw_rabitq_collection.stats.doc_count == 0 + result = hnsw_rabitq_collection.insert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + assert hnsw_rabitq_collection.stats.doc_count == len(multiple_docs) + + yield hnsw_rabitq_collection + + # Teardown: delete multiple docs + hnsw_rabitq_collection.delete([doc.id for doc in multiple_docs]) + + +# ==================== Tests ==================== + + +@pytest.mark.usefixtures("hnsw_rabitq_collection") +class TestHnswRabitqCollectionCreation: + """Test HNSW RaBitQ collection creation and schema validation.""" + + def test_collection_creation( + self, hnsw_rabitq_collection: Collection, hnsw_rabitq_collection_schema + ): + """Test that collection is created with correct schema.""" + assert hnsw_rabitq_collection is not None + assert hnsw_rabitq_collection.schema.name == hnsw_rabitq_collection_schema.name + assert len(hnsw_rabitq_collection.schema.fields) == len( + hnsw_rabitq_collection_schema.fields + ) + assert len(hnsw_rabitq_collection.schema.vectors) == len( + hnsw_rabitq_collection_schema.vectors + ) + + def test_vector_schema_validation(self, hnsw_rabitq_collection: Collection): + """Test that vector schema has correct HNSW RaBitQ configuration.""" + vector_schema = hnsw_rabitq_collection.schema.vector("embedding") + assert vector_schema is not None + assert vector_schema.name == "embedding" + assert vector_schema.data_type == DataType.VECTOR_FP32 + assert vector_schema.dimension == 128 + + index_param = vector_schema.index_param + assert index_param is not None + assert index_param.metric_type == MetricType.L2 + assert index_param.m == 16 + assert index_param.ef_construction == 200 + assert index_param.total_bits == 7 + assert index_param.num_clusters == 64 + + def test_collection_stats(self, hnsw_rabitq_collection: Collection): + """Test initial collection statistics.""" + stats = hnsw_rabitq_collection.stats + assert stats is not None + assert stats.doc_count == 0 + assert len(stats.index_completeness) == 1 + assert stats.index_completeness["embedding"] == 1 + + +@pytest.mark.usefixtures("hnsw_rabitq_collection") +class TestHnswRabitqCollectionInsert: + """Test document insertion into HNSW RaBitQ collection.""" + + def test_insert_single_doc( + self, hnsw_rabitq_collection: Collection, single_doc: Doc + ): + """Test inserting a single document.""" + result = hnsw_rabitq_collection.insert(single_doc) + assert bool(result) + assert result.ok() + + stats = hnsw_rabitq_collection.stats + assert stats is not None + assert stats.doc_count == 1 + + def test_insert_multiple_docs( + self, hnsw_rabitq_collection: Collection, multiple_docs: list[Doc] + ): + """Test inserting multiple documents.""" + result = hnsw_rabitq_collection.insert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + + stats = hnsw_rabitq_collection.stats + assert stats is not None + assert stats.doc_count == len(multiple_docs) + + +@pytest.mark.usefixtures("hnsw_rabitq_collection") +class TestHnswRabitqCollectionFetch: + """Test document fetching from HNSW RaBitQ collection.""" + + def test_fetch_single_doc( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + """Test fetching a single document by ID.""" + result = collection_with_single_doc.fetch(ids=[single_doc.id]) + assert bool(result) + assert single_doc.id in result.keys() + + doc = result[single_doc.id] + assert doc is not None + assert doc.id == single_doc.id + assert doc.field("id") == single_doc.field("id") + assert doc.field("name") == single_doc.field("name") + + def test_fetch_multiple_docs( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test fetching multiple documents by IDs.""" + ids = [doc.id for doc in multiple_docs[:10]] + result = collection_with_multiple_docs.fetch(ids=ids) + assert bool(result) + assert len(result) == len(ids) + + for doc_id in ids: + assert doc_id in result + doc = result[doc_id] + assert doc is not None + assert doc.id == doc_id + + def test_fetch_nonexistent_doc(self, collection_with_single_doc: Collection): + """Test fetching a non-existent document.""" + result = collection_with_single_doc.fetch(ids=["nonexistent_id"]) + assert len(result) == 0 + + +@pytest.mark.usefixtures("hnsw_rabitq_collection") +class TestHnswRabitqCollectionQuery: + """Test vector search queries on HNSW RaBitQ collection.""" + + def test_query_by_vector( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying by vector with HNSW RaBitQ index.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=300), + ) + + result = collection_with_multiple_docs.query(queries=query, topk=10) + assert len(result) > 0 + assert len(result) <= 10 + + # First result should be the query document itself (or very close) + first_doc = result[0] + assert first_doc is not None + assert first_doc.id is not None + + def test_query_by_id( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying by document ID with HNSW RaBitQ index.""" + query = Query( + field_name="embedding", + id=multiple_docs[0].id, + param=HnswRabitqQueryParam(ef=300), + ) + + result = collection_with_multiple_docs.query(queries=query, topk=10) + assert len(result) > 0 + assert len(result) <= 10 + + def test_query_with_different_ef_values( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with different ef parameter values.""" + query_vector = multiple_docs[0].vector("embedding") + + # Test with ef=100 + query_100 = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=100), + ) + result_100 = collection_with_multiple_docs.query(queries=query_100, topk=10) + assert len(result_100) > 0 + + # Test with ef=500 + query_500 = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=500), + ) + result_500 = collection_with_multiple_docs.query(queries=query_500, topk=10) + assert len(result_500) > 0 + + def test_query_with_topk( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with different topk values.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=300), + ) + + # Test topk=5 + result_5 = collection_with_multiple_docs.query(queries=query, topk=5) + assert len(result_5) <= 5 + + # Test topk=20 + result_20 = collection_with_multiple_docs.query(queries=query, topk=20) + assert len(result_20) <= 20 + + def test_query_with_filter( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with filter conditions.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=300), + ) + + # Query with id filter + result = collection_with_multiple_docs.query( + queries=query, topk=10, filter="id < 50" + ) + assert len(result) > 0 + for doc in result: + assert doc.field("id") < 50 + + def test_query_with_output_fields( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with specific output fields.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=300), + ) + + result = collection_with_multiple_docs.query( + queries=query, topk=10, output_fields=["id", "name"] + ) + assert len(result) > 0 + + first_doc = result[0] + assert "id" in first_doc.field_names() + assert "name" in first_doc.field_names() + + def test_query_with_include_vector( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test querying with vector data included in results.""" + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=300), + ) + + result = collection_with_multiple_docs.query( + queries=query, topk=10, include_vector=True + ) + assert len(result) > 0 + + first_doc = result[0] + assert first_doc.vector("embedding") is not None + assert len(first_doc.vector("embedding")) == 128 + + +@pytest.mark.usefixtures("hnsw_rabitq_collection") +class TestHnswRabitqCollectionUpdate: + """Test document update in HNSW RaBitQ collection.""" + + def test_update_doc_fields( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + """Test updating document fields.""" + updated_doc = Doc( + id=single_doc.id, + fields={"id": single_doc.field("id"), "name": "updated_name"}, + ) + + result = collection_with_single_doc.update(updated_doc) + assert bool(result) + assert result.ok() + + # Verify update + fetched = collection_with_single_doc.fetch(ids=[single_doc.id]) + assert single_doc.id in fetched + doc = fetched[single_doc.id] + assert doc.field("name") == "updated_name" + + def test_update_doc_vector( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + """Test updating document vector.""" + new_vector = [0.5 + i * 0.01 for i in range(128)] + updated_doc = Doc( + id=single_doc.id, + vectors={"embedding": new_vector}, + ) + + result = collection_with_single_doc.update(updated_doc) + assert bool(result) + assert result.ok() + + # Verify update + fetched = collection_with_single_doc.fetch( + ids=[single_doc.id], + ) + assert single_doc.id in fetched + doc = fetched[single_doc.id] + assert doc.vector("embedding") is not None + embedding = doc.vector("embedding") + assert len(embedding) == 128 + # Verify vector values are approximately equal (float comparison) + for i in range(128): + assert math.isclose(embedding[i], new_vector[i], rel_tol=1e-5) + + +@pytest.mark.usefixtures("hnsw_rabitq_collection") +class TestHnswRabitqCollectionDelete: + """Test document deletion from HNSW RaBitQ collection.""" + + def test_delete_single_doc( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + """Test deleting a single document.""" + result = collection_with_single_doc.delete(single_doc.id) + assert bool(result) + assert result.ok() + + stats = collection_with_single_doc.stats + assert stats.doc_count == 0 + + def test_delete_multiple_docs( + self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] + ): + """Test deleting multiple documents.""" + ids_to_delete = [doc.id for doc in multiple_docs[:10]] + result = collection_with_multiple_docs.delete(ids_to_delete) + assert len(result) == len(ids_to_delete) + for item in result: + assert item.ok() + + stats = collection_with_multiple_docs.stats + assert stats.doc_count == len(multiple_docs) - len(ids_to_delete) + + +@pytest.mark.usefixtures("hnsw_rabitq_collection") +class TestHnswRabitqCollectionOptimizeAndReopen: + """Test collection optimize and reopen functionality.""" + + def test_optimize_close_reopen_and_query( + self, + tmp_path_factory, + hnsw_rabitq_collection_schema, + collection_option, + multiple_docs: list[Doc], + ): + """Test inserting 100 docs, optimize, close, reopen and query.""" + # Create collection and insert 100 documents + temp_dir = tmp_path_factory.mktemp("zvec_hnsw_rabitq_optimize") + collection_path = temp_dir / "test_optimize_collection" + + coll = zvec.create_and_open( + path=str(collection_path), + schema=hnsw_rabitq_collection_schema, + option=collection_option, + ) + + assert coll is not None + assert coll.stats.doc_count == 0 + + # Insert 100 documents + result = coll.insert(multiple_docs) + assert len(result) == len(multiple_docs) + for item in result: + assert item.ok() + assert coll.stats.doc_count == len(multiple_docs) + + # Call optimize + from zvec import OptimizeOption + + coll.optimize(option=OptimizeOption()) + + # Verify data is still accessible after optimize + query_vector = multiple_docs[0].vector("embedding") + query = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=300), + ) + result_before_close = coll.query(query, topk=10) + assert len(result_before_close) > 0 + + # Close collection (destroy will close it) + collection_path_str = str(collection_path) + del coll + + # Reopen collection + reopened_coll = zvec.open(path=collection_path_str, option=collection_option) + assert reopened_coll is not None + assert reopened_coll.stats.doc_count == len(multiple_docs) + + # Execute query on reopened collection + query_after_reopen = Query( + field_name="embedding", + vector=query_vector, + param=HnswRabitqQueryParam(ef=300), + ) + result_after_reopen = reopened_coll.query(query_after_reopen, topk=10) + assert len(result_after_reopen) > 0 + assert len(result_after_reopen) <= 10 + + # Verify query results are valid + first_doc = result_after_reopen[0] + assert first_doc is not None + assert first_doc.id is not None + assert first_doc.field("id") is not None + assert first_doc.field("name") is not None + + # Cleanup + reopened_coll.destroy() diff --git a/python/tests/test_convert.py b/python/tests/test_convert.py new file mode 100644 index 0000000..2fac2f5 --- /dev/null +++ b/python/tests/test_convert.py @@ -0,0 +1,584 @@ +from __future__ import annotations + +import math + +import pytest +from zvec._zvec import _Doc +from zvec.model.convert import convert_to_py_doc, convert_to_cpp_doc +from zvec import Doc, CollectionSchema, DataType, FieldSchema, VectorSchema + + +# ---------------------------- +# Convert Cpp Doc Test Case +# ---------------------------- +class TestConvertCppDoc: + def test_default(self): + doc = Doc(id="1") + schema = CollectionSchema( + name="test_collection", + fields=FieldSchema("name", DataType.STRING), + ) + + cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema) + + assert cpp_doc is not None + assert cpp_doc.pk() == doc.id + + def test_with_field_notin_schema(self): + doc = Doc(id="1", fields={"name": "Tom"}) + schema = CollectionSchema( + name="test_collection", + fields=[ + FieldSchema("id", DataType.UINT64), + FieldSchema("salary", DataType.UINT32), + FieldSchema("age", DataType.INT32), + FieldSchema("create_at", DataType.INT64), + FieldSchema("author", DataType.STRING), + FieldSchema("weight", DataType.FLOAT), + ], + ) + with pytest.raises(ValueError): + convert_to_cpp_doc(doc, collection_schema=schema) + + def test_with_scalar_fields(self): + schema = CollectionSchema( + name="test_collection", + fields=[ + FieldSchema("id", DataType.UINT64), + FieldSchema("salary", DataType.UINT32), + FieldSchema("age", DataType.INT32), + FieldSchema("create_at", DataType.INT64), + FieldSchema("author", DataType.STRING), + FieldSchema("weight", DataType.FLOAT), + FieldSchema("bmi", DataType.DOUBLE), + FieldSchema("is_male", DataType.BOOL), + ], + ) + doc = Doc( + id="1", + fields={ + "id": 1, + "salary": 1000, + "age": 18, + "create_at": 1640995200, + "bmi": 80.0 / 200.0, + "author": "Tom", + "weight": 80.0, + "is_male": True, + }, + ) + cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema) + assert cpp_doc is not None + assert cpp_doc.pk() == doc.id + assert cpp_doc.get_any("id", DataType.UINT64) == 1 + assert cpp_doc.get_any("salary", DataType.UINT32) == 1000 + assert cpp_doc.get_any("age", DataType.INT32) == 18 + assert cpp_doc.get_any("create_at", DataType.INT64) == 1640995200 + assert cpp_doc.get_any("author", DataType.STRING) == "Tom" + assert math.isclose( + cpp_doc.get_any("weight", DataType.FLOAT), 80.0, rel_tol=1e-6 + ) + assert math.isclose( + cpp_doc.get_any("bmi", DataType.DOUBLE), 80.0 / 200.0, rel_tol=1e-6 + ) + assert cpp_doc.get_any("is_male", DataType.BOOL) == True + + def test_with_array_fields(self): + schema = CollectionSchema( + name="test_collection", + fields=[ + FieldSchema("tags", DataType.ARRAY_STRING), + FieldSchema("ids", DataType.ARRAY_UINT64), + FieldSchema("marks", DataType.ARRAY_UINT32), + FieldSchema("x", DataType.ARRAY_INT32), + FieldSchema("y", DataType.ARRAY_INT64), + FieldSchema("scores", DataType.ARRAY_FLOAT), + FieldSchema("ratios", DataType.ARRAY_DOUBLE), + FieldSchema("results", DataType.ARRAY_BOOL), + ], + ) + + doc = Doc( + id="1", + fields={ + "tags": ["tag1", "tag2", "tag3"], + "ids": [111111111111, 222222222222, 333333333333], + "marks": [100, 200, 300], + "x": [1, 2, 3], + "y": [100, 200, 300], + "scores": [1.1, 2.2, 3.3], + "ratios": [0.1, 0.2, 0.3], + "results": [True, False, True], + }, + ) + cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema) + + assert cpp_doc is not None + assert cpp_doc.pk() == doc.id + assert cpp_doc.get_any("tags", DataType.ARRAY_STRING) == doc.field("tags") + assert cpp_doc.get_any("ids", DataType.ARRAY_UINT64) == doc.field("ids") + assert cpp_doc.get_any("marks", DataType.ARRAY_UINT32) == doc.field("marks") + assert cpp_doc.get_any("x", DataType.ARRAY_INT32) == doc.field("x") + assert cpp_doc.get_any("y", DataType.ARRAY_INT64) == doc.field("y") + scores = cpp_doc.get_any("scores", DataType.ARRAY_FLOAT) + for i in range(len(doc.field("scores"))): + assert math.isclose(scores[i], doc.field("scores")[i], rel_tol=1e-1) + ratios = cpp_doc.get_any("ratios", DataType.ARRAY_DOUBLE) + for i in range(len(doc.field("ratios"))): + assert math.isclose(ratios[i], doc.field("ratios")[i], rel_tol=1e-1) + results = cpp_doc.get_any("results", DataType.ARRAY_BOOL) + for i in range(len(doc.field("results"))): + assert results[i] == doc.field("results")[i] + + def test_with_dense_vector_fields(self): + schema = CollectionSchema( + name="test_collection", + vectors=[ + VectorSchema( + name="embedding", + data_type=DataType.VECTOR_FP16, + dimension=4, + ), + VectorSchema( + name="image", + data_type=DataType.VECTOR_FP32, + dimension=8, + ), + VectorSchema( + name="text", + data_type=DataType.VECTOR_INT8, + dimension=32, + ), + ], + ) + + doc = Doc( + id="1", + vectors={ + "embedding": [1.1] * 4, + "image": [2.2] * 8, + "text": [4] * 32, + }, + ) + cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema) + assert cpp_doc is not None + assert cpp_doc.pk() == doc.id + + embedding_vector = cpp_doc.get_any("embedding", DataType.VECTOR_FP16) + assert len(embedding_vector) == 4 + for i in range(4): + assert math.isclose( + embedding_vector[i], doc.vector("embedding")[i], rel_tol=1e-1 + ) + + image_vector = cpp_doc.get_any("image", DataType.VECTOR_FP32) + assert len(image_vector) == 8 + for i in range(8): + assert math.isclose(image_vector[i], doc.vector("image")[i], rel_tol=1e-1) + + text_vector = cpp_doc.get_any("text", DataType.VECTOR_INT8) + assert len(text_vector) == 32 + for i in range(32): + assert text_vector[i] == doc.vectors["text"][i] + + def test_with_sparse_vector_fields(self): + schema = CollectionSchema( + name="test_collection", + vectors=[ + VectorSchema( + name="author", + data_type=DataType.SPARSE_VECTOR_FP32, + ), + VectorSchema( + name="content", + data_type=DataType.SPARSE_VECTOR_FP16, + ), + ], + ) + doc = Doc( + id="1", + vectors={ + "author": {1: 1.1, 2: 2.2, 3: 3.3}, + "content": {4: 4.4, 5: 5.5, 6: 6.6}, + }, + ) + + cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema) + assert cpp_doc is not None + assert cpp_doc.pk() == doc.id + + author_vector = cpp_doc.get_any("author", DataType.SPARSE_VECTOR_FP32) + assert isinstance(author_vector, dict) + for key, value in doc.vector("author").items(): + assert math.isclose(author_vector[key], value, rel_tol=1e-1) + + content_vector = cpp_doc.get_any("content", DataType.SPARSE_VECTOR_FP16) + assert isinstance(content_vector, dict) + for key, value in doc.vector("content").items(): + assert math.isclose(content_vector[key], value, rel_tol=1e-1) + + def test_with_scalar_fields_error_datatype(self): + schema = CollectionSchema( + name="test_collection", + fields=[ + FieldSchema("id", DataType.UINT64), + FieldSchema("salary", DataType.UINT32), + FieldSchema("age", DataType.INT32), + FieldSchema("create_at", DataType.INT64), + FieldSchema("author", DataType.STRING), + FieldSchema("weight", DataType.FLOAT), + FieldSchema("bmi", DataType.DOUBLE), + FieldSchema("is_male", DataType.BOOL), + ], + ) + doc = Doc( + id="1", + fields={ + "id": "1", + }, + ) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"salary": "1000"}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"age": "18"}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"create_at": "2021-01-01"}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"author": 1}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"weight": "80.5"}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"bmi": "25.0"}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"is_male": "true"}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + def test_with_array_fields_error_datatype(self): + schema = CollectionSchema( + name="test_collection", + fields=[ + FieldSchema("tags", DataType.ARRAY_STRING), + FieldSchema("ids", DataType.ARRAY_UINT64), + FieldSchema("marks", DataType.ARRAY_UINT32), + FieldSchema("x", DataType.ARRAY_INT32), + FieldSchema("y", DataType.ARRAY_INT64), + FieldSchema("scores", DataType.ARRAY_FLOAT), + FieldSchema("ratios", DataType.ARRAY_DOUBLE), + FieldSchema("results", DataType.ARRAY_BOOL), + ], + ) + + doc = Doc(id="1", fields={"tags": [1, 2, 3]}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"ids": ["1", "2", "3"]}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"marks": [1.1, 2.2, 3.3]}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"x": [1.1, 2.2, 3.3]}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"y": [1.1, 2.2, 3.3]}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"scores": ["1", "2", "3"]}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"ratios": ["1", "2", "3"]}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", fields={"results": ["1", "2", "3"]}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + def test_with_vector_fields_error_datatype(self): + schema = CollectionSchema( + name="test_collection", + vectors=[ + VectorSchema( + name="embedding", + data_type=DataType.VECTOR_FP16, + dimension=4, + ), + VectorSchema( + name="image", + data_type=DataType.VECTOR_FP32, + dimension=8, + ), + VectorSchema( + name="text", + data_type=DataType.VECTOR_INT8, + dimension=32, + ), + ], + ) + + doc = Doc(id="1", vectors={"image": ["1.1"] * 4}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", vectors={"text": ["1"] * 4}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc(id="1", vectors={"embedding": ["1"] * 4}) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + def test_with_sparse_vector_error_datatype(self): + schema = CollectionSchema( + name="test_collection", + vectors=[ + VectorSchema( + name="author", + data_type=DataType.SPARSE_VECTOR_FP32, + ), + VectorSchema( + name="content", + data_type=DataType.SPARSE_VECTOR_FP16, + ), + ], + ) + doc = Doc( + id="1", + vectors={ + "author": {"1": 1.1, "2": 2.2, "3": 3.3}, + }, + ) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc( + id="1", + vectors={ + "content": {"1": 1.1, "2": 2.2, "3": 3.3}, + }, + ) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + doc = Doc( + id="1", + vectors={ + "author": {1: "1", 2: "2", 3: "3"}, + }, + ) + with pytest.raises(TypeError): + convert_to_cpp_doc(doc, collection_schema=schema) + + +# ---------------------------- +# Convert Py Doc Test Case +# ---------------------------- +class TestConvertPyDoc: + def test_default(self): + doc = _Doc() + doc.set_pk("1") + doc.set_score(1.0) + + schema = CollectionSchema( + name="test_collection", + fields=FieldSchema("name", DataType.STRING), + ) + + py_doc = convert_to_py_doc(doc, schema) + assert py_doc.id == "1" + assert py_doc.score == 1.0 + + def test_with_scalar_fields(self): + schema = CollectionSchema( + name="test_collection", + fields=[ + FieldSchema("id", DataType.UINT64), + FieldSchema("salary", DataType.UINT32), + FieldSchema("age", DataType.INT32), + FieldSchema("create_at", DataType.INT64), + FieldSchema("author", DataType.STRING), + FieldSchema("weight", DataType.FLOAT), + FieldSchema("bmi", DataType.DOUBLE), + FieldSchema("is_male", DataType.BOOL), + ], + ) + doc = _Doc() + doc.set_pk("1") + doc.set_any("id", schema.field("id")._get_object(), 1111111111111111) + doc.set_any("salary", schema.field("salary")._get_object(), 1000) + doc.set_any("age", schema.field("age")._get_object(), 18) + doc.set_any("create_at", schema.field("create_at")._get_object(), 1640995200) + doc.set_any("author", schema.field("author")._get_object(), "Tom") + doc.set_any("weight", schema.field("weight")._get_object(), 80.0) + doc.set_any("bmi", schema.field("bmi")._get_object(), 80.0 / 200.0) + doc.set_any("is_male", schema.field("is_male")._get_object(), True) + + py_doc = convert_to_py_doc(doc, schema) + assert py_doc.id == "1" + assert py_doc.field("id") == 1111111111111111 + assert py_doc.field("salary") == 1000 + assert py_doc.field("age") == 18 + assert py_doc.field("create_at") == 1640995200 + assert py_doc.field("author") == "Tom" + assert py_doc.field("weight") == 80.0 + assert py_doc.field("bmi") == 80.0 / 200.0 + assert py_doc.field("is_male") == True + + def test_with_array_fields(self): + schema = CollectionSchema( + name="test_collection", + fields=[ + FieldSchema("tags", DataType.ARRAY_STRING), + FieldSchema("ids", DataType.ARRAY_UINT64), + FieldSchema("marks", DataType.ARRAY_UINT32), + FieldSchema("x", DataType.ARRAY_INT32), + FieldSchema("y", DataType.ARRAY_INT64), + FieldSchema("scores", DataType.ARRAY_FLOAT), + FieldSchema("ratios", DataType.ARRAY_DOUBLE), + FieldSchema("results", DataType.ARRAY_BOOL), + ], + ) + + doc = _Doc() + doc.set_pk("1") + doc.set_any( + "tags", schema.field("tags")._get_object(), ["tag1", "tag2", "tag3"] + ) + doc.set_any( + "ids", + schema.field("ids")._get_object(), + [111111111111, 222222222222, 3333333333333], + ) + doc.set_any("marks", schema.field("marks")._get_object(), [1000, 2000, 3000]) + doc.set_any("x", schema.field("x")._get_object(), [1, 2, 3]) + doc.set_any("y", schema.field("y")._get_object(), [100, 200, 300]) + doc.set_any("scores", schema.field("scores")._get_object(), [0.1, 0.2, 0.3]) + doc.set_any("ratios", schema.field("ratios")._get_object(), [0.1, 0.2, 0.3]) + doc.set_any( + "results", schema.field("results")._get_object(), [True, False, True] + ) + + py_doc = convert_to_py_doc(doc, schema) + assert py_doc.field("tags") == ["tag1", "tag2", "tag3"] + assert py_doc.field("ids") == [111111111111, 222222222222, 3333333333333] + assert py_doc.field("marks") == [1000, 2000, 3000] + assert py_doc.field("x") == [1, 2, 3] + assert py_doc.field("y") == [100, 200, 300] + + scores = doc.get_any("scores", DataType.ARRAY_FLOAT) + for i in range(len(scores)): + assert math.isclose(scores[i], py_doc.field("scores")[i], rel_tol=1e-1) + ratios = doc.get_any("ratios", DataType.ARRAY_DOUBLE) + for i in range(len(ratios)): + assert math.isclose(ratios[i], py_doc.field("ratios")[i], rel_tol=1e-1) + results = doc.get_any("results", DataType.ARRAY_BOOL) + for i in range(len(results)): + assert results[i] == py_doc.field("results")[i] + + def test_with_dense_vector_fields(self): + schema = CollectionSchema( + name="test_collection", + vectors=[ + VectorSchema( + name="embedding", + data_type=DataType.VECTOR_FP16, + dimension=4, + ), + VectorSchema( + name="image", + data_type=DataType.VECTOR_FP32, + dimension=8, + ), + VectorSchema( + name="text", + data_type=DataType.VECTOR_INT8, + dimension=32, + ), + ], + ) + + doc = _Doc() + doc.set_pk("1") + doc.set_any("embedding", schema.vector("embedding")._get_object(), [1.1] * 4) + doc.set_any("image", schema.vector("image")._get_object(), [2.2] * 8) + doc.set_any("text", schema.vector("text")._get_object(), [4] * 32) + + py_doc = convert_to_py_doc(doc, schema) + assert py_doc.id == "1" + + embedding_vector = py_doc.vector("embedding") + assert len(embedding_vector) == 4 + for i in range(4): + assert math.isclose( + py_doc.vector("embedding")[i], embedding_vector[i], rel_tol=1e-1 + ) + + image_vector = py_doc.vector("image") + assert len(image_vector) == 8 + for i in range(8): + assert math.isclose( + py_doc.vector("image")[i], image_vector[i], rel_tol=1e-1 + ) + + text_vector = py_doc.vector("text") + assert len(text_vector) == 32 + for i in range(32): + assert py_doc.vector("text")[i] == text_vector[i] + + def test_with_sparse_vector_fields(self): + schema = CollectionSchema( + name="test_collection", + vectors=[ + VectorSchema( + name="author", + data_type=DataType.SPARSE_VECTOR_FP32, + ), + VectorSchema( + name="content", + data_type=DataType.SPARSE_VECTOR_FP16, + ), + ], + ) + + doc = _Doc() + doc.set_pk("1") + doc.set_any( + "author", schema.vector("author")._get_object(), {1: 1.1, 2: 2.2, 3: 3.3} + ) + doc.set_any( + "content", schema.vector("content")._get_object(), {4: 4.4, 5: 5.5, 6: 6.6} + ) + + py_doc = convert_to_py_doc(doc, schema) + assert py_doc.id == "1" + + author_vector = py_doc.vector("author") + assert isinstance(author_vector, dict) + for key, value in doc.get_any("author", DataType.SPARSE_VECTOR_FP32).items(): + assert math.isclose(author_vector[key], value, rel_tol=1e-1) + + content_vector = py_doc.vector("content") + assert isinstance(content_vector, dict) + for key, value in doc.get_any("content", DataType.SPARSE_VECTOR_FP16).items(): + assert math.isclose(content_vector[key], value, rel_tol=1e-1) diff --git a/python/tests/test_doc.py b/python/tests/test_doc.py new file mode 100644 index 0000000..ecfd31e --- /dev/null +++ b/python/tests/test_doc.py @@ -0,0 +1,269 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import math +import pytest + + +from zvec._zvec import _Doc +from zvec import FieldSchema, VectorSchema, Doc, DataType + + +# ---------------------------- +# PyDoc Test Case +# ---------------------------- +class TestPyDoc: + def test_default(self): + Doc(id="1") + + def test_with_single_vector(self): + doc = Doc(id="1", vectors={"dense": [1, 2, 3]}) + assert doc is not None + assert doc.id == "1" + assert doc.vector("dense") == [1, 2, 3] + + def test_with_hybrid_vectors(self): + doc = Doc( + id="1", vectors={"dense": [1, 2, 3], "sparse": {1: 1.0, 2: 2.0, 3: 3.0}} + ) + assert doc is not None + assert doc.id == "1" + assert doc.vector("dense") == [1, 2, 3] + assert doc.vector("sparse") == {1: 1.0, 2: 2.0, 3: 3.0} + + def test_with_multi_vectors(self): + doc = Doc( + id="1", + vectors={ + "image": [1, 2, 3], + "description": [4, 5, 6], + "keys": {1: 1.0, 2: 2.0, 3: 3.0}, + }, + fields={"author": "Tom", "age": 19, "is_male": True, "weight": 60.5}, + ) + assert doc is not None + assert doc.id == "1" + assert doc.vector("image") == [1, 2, 3] + assert doc.vector("description") == [4, 5, 6] + assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0} + assert doc.field("author") == "Tom" + assert doc.field("age") == 19 + assert doc.field("is_male") == True + assert doc.field("weight") == 60.5 + + def test_with_numpy_array(self): + import numpy as np + + doc = Doc._from_tuple( + ( + "1", + 0.0, + None, + { + "image": np.array([1, 2, 3]), + "description": np.random.random(512), + "keys": {1: 1.0, 2: 2.0, 3: 3.0}, + }, + ) + ) + assert doc is not None + assert doc.id == "1" + assert doc.vector("image") == [1, 2, 3] + assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0} + + +# ---------------------------- +# CppDoc Test Case +# ---------------------------- +class TestCppDoc: + def test_default(self): + doc = _Doc() + assert doc is not None + + def test_doc_set_pk(self): + doc = _Doc() + doc.set_pk("1") + assert doc.pk() == "1" + + def test_doc_set_score(self): + doc = _Doc() + doc.set_score(0.9) + assert math.isclose(doc.score(), 0.9, rel_tol=1e-6) + + def test_doc_get_null_field(self): + doc = _Doc() + schema = FieldSchema("author", DataType.STRING, nullable=True) + doc.set_any("author", schema._get_object(), None) + assert doc.has_field("author") + assert doc.get_any("author", schema.data_type) is None + + def test_doc_get_set_has_null_field(self): + doc = _Doc() + schema = FieldSchema("author", DataType.STRING, nullable=False) + with pytest.raises(ValueError): + doc.set_any("author", schema._get_object(), None) + + def test_doc_get_set_has_string_field(self): + doc = _Doc() + schema = FieldSchema("author", DataType.STRING) + doc.set_any("author", schema._get_object(), "Tom") + assert doc.has_field("author") + assert doc.get_any("author", DataType.STRING) == "Tom" + + def test_doc_get_set_has_bool_field(self): + doc = _Doc() + schema = FieldSchema("is_male", DataType.BOOL) + doc.set_any("is_male", schema._get_object(), True) + assert doc.has_field("is_male") + assert doc.get_any("is_male", DataType.BOOL) == True + + def test_doc_get_set_has_int32_field(self): + doc = _Doc() + schema = FieldSchema("age", DataType.INT32) + doc.set_any("age", schema._get_object(), 19) + assert doc.has_field("age") + assert doc.get_any("age", DataType.INT32) == 19 + + def test_doc_get_set_has_int64_field(self): + doc = _Doc() + schema = FieldSchema("id", DataType.INT64) + doc.set_any("id", schema._get_object(), 1111111111111111111) + assert doc.has_field("id") + assert doc.get_any("id", DataType.INT64) == 1111111111111111111 + + def test_doc_get_set_has_float_field(self): + doc = _Doc() + schema = FieldSchema("weight", DataType.FLOAT) + doc.set_any("weight", schema._get_object(), 60.5) + assert doc.has_field("weight") + assert math.isclose(doc.get_any("weight", DataType.FLOAT), 60.5, rel_tol=1e-6) + + def test_doc_get_set_has_double_field(self): + doc = _Doc() + schema = FieldSchema("height", DataType.DOUBLE) + doc.set_any("height", schema._get_object(), 1.77777777777) + assert doc.has_field("height") + assert math.isclose( + doc.get_any("height", DataType.DOUBLE), 1.7777777777, rel_tol=1e-9 + ) + + def test_doc_get_set_has_uint32_field(self): + doc = _Doc() + schema = FieldSchema("id", DataType.UINT32) + doc.set_any("id", schema._get_object(), 4294967295) + assert doc.has_field("id") + assert doc.get_any("id", DataType.UINT32) == 4294967295 + + def test_doc_get_set_has_uint64_field(self): + doc = _Doc() + schema = FieldSchema("id", DataType.UINT64) + doc.set_any("id", schema._get_object(), 18446744073709551615) + assert doc.has_field("id") + assert doc.get_any("id", DataType.UINT64) == 18446744073709551615 + + def test_doc_get_set_has_array_string_field(self): + doc = _Doc() + schema = FieldSchema("tags", DataType.ARRAY_STRING) + doc.set_any("tags", schema._get_object(), ["tag1", "tag2", "tag3"]) + assert doc.has_field("tags") + assert doc.get_any("tags", DataType.ARRAY_STRING) == ["tag1", "tag2", "tag3"] + + def test_doc_get_set_has_array_int32_field(self): + doc = _Doc() + schema = FieldSchema("ids", DataType.ARRAY_INT32) + doc.set_any("ids", schema._get_object(), [1, 2, 3]) + assert doc.has_field("ids") + assert doc.get_any("ids", DataType.ARRAY_INT32) == [1, 2, 3] + + def test_doc_get_set_has_array_int64_field(self): + doc = _Doc() + schema = FieldSchema("ids", DataType.ARRAY_INT64) + doc.set_any("ids", schema._get_object(), [1, 2, 3]) + assert doc.has_field("ids") + assert doc.get_any("ids", DataType.ARRAY_INT64) == [1, 2, 3] + + def test_doc_get_set_has_array_float_field(self): + doc = _Doc() + schema = FieldSchema("weights", DataType.ARRAY_FLOAT) + doc.set_any("weights", schema._get_object(), [1.0, 2.0, 3.0]) + assert doc.has_field("weights") + assert doc.get_any("weights", DataType.ARRAY_FLOAT) == [1.0, 2.0, 3.0] + + def test_doc_get_set_has_array_double_field(self): + doc = _Doc() + schema = FieldSchema("heights", DataType.ARRAY_DOUBLE) + doc.set_any("heights", schema._get_object(), [1.0, 2.0, 3.0]) + assert doc.has_field("heights") + assert doc.get_any("heights", DataType.ARRAY_DOUBLE) == [1.0, 2.0, 3.0] + + def test_doc_get_set_has_array_bool_field(self): + doc = _Doc() + schema = FieldSchema("bools", DataType.ARRAY_BOOL) + doc.set_any("bools", schema._get_object(), [True, False, True]) + assert doc.has_field("bools") + assert doc.get_any("bools", DataType.ARRAY_BOOL) == [True, False, True] + + def test_doc_get_set_has_vector_fp16(self): + doc = _Doc() + schema = VectorSchema("image", DataType.VECTOR_FP16) + doc.set_any("image", schema._get_object(), [1.0, 2.0, 3.0]) + assert doc.has_field("image") + image_vector = doc.get_any("image", DataType.VECTOR_FP16) + assert image_vector is not None + for i in range(len(image_vector)): + assert math.isclose(image_vector[i], [1.0, 2.0, 3.0][i], rel_tol=1e-6) + + def test_doc_get_set_has_vector_fp32(self): + doc = _Doc() + schema = VectorSchema("image", DataType.VECTOR_FP32) + doc.set_any("image", schema._get_object(), [1.111111, 2.222222, 3.333333]) + assert doc.has_field("image") + vector = doc.get_any("image", DataType.VECTOR_FP32) + assert vector is not None + for i in range(len(vector)): + assert math.isclose( + vector[i], [1.111111, 2.222222, 3.333333][i], rel_tol=1e-6 + ) + + def test_doc_get_set_has_vector_int8(self): + doc = _Doc() + schema = VectorSchema("image", DataType.VECTOR_INT8) + doc.set_any("image", schema._get_object(), [1, 2, 3]) + assert doc.has_field("image") + assert doc.get_any("image", DataType.VECTOR_INT8) == [1, 2, 3] + + def test_doc_get_set_has_sparse_vector_fp32(self): + doc = _Doc() + sparse = {1: 1.111111, 2: 2.222222, 3: 3.333333} + schema = VectorSchema("key", DataType.SPARSE_VECTOR_FP32) + doc.set_any("key", schema._get_object(), sparse) + assert doc.has_field("key") + vector = doc.get_any("key", DataType.SPARSE_VECTOR_FP32) + assert vector is not None + assert isinstance(vector, dict) + for key, value in sparse.items(): + assert math.isclose(vector[key], value, rel_tol=1e-6) + + def test_doc_get_set_has_sparse_vector_fp16(self): + doc = _Doc() + sparse = {1: 1.1, 2: 2.2, 3: 3.3} + schema = VectorSchema("key", DataType.SPARSE_VECTOR_FP16) + doc.set_any("key", schema._get_object(), sparse) + assert doc.has_field("key") + vector = doc.get_any("key", DataType.SPARSE_VECTOR_FP16) + assert vector is not None + assert isinstance(vector, dict) + for key, value in sparse.items(): + assert math.isclose(vector[key], value, rel_tol=1e-1) diff --git a/python/tests/test_embedding.py b/python/tests/test_embedding.py new file mode 100644 index 0000000..1b0622b --- /dev/null +++ b/python/tests/test_embedding.py @@ -0,0 +1,2099 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import os +from http import HTTPStatus +from unittest.mock import MagicMock, patch, Mock + +import numpy as np +import pytest +from zvec.extension import ( + BM25EmbeddingFunction, + DefaultLocalDenseEmbedding, + DefaultLocalSparseEmbedding, + OpenAIDenseEmbedding, + QwenDenseEmbedding, + QwenSparseEmbedding, +) + +# Environment variable to control integration tests +# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API/model tests +RUN_INTEGRATION_TESTS = os.environ.get("ZVEC_RUN_INTEGRATION_TESTS", "0") == "1" + + +# ---------------------------- +# QwenDenseEmbedding Test Case +# ---------------------------- +class TestQwenDenseEmbedding: + def test_init_with_api_key(self): + # Test initialization with explicit API key + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") + assert embedding_func.dimension == 128 + assert embedding_func.model == "text-embedding-v4" + assert embedding_func._api_key == "test_key" + + @patch.dict(os.environ, {"DASHSCOPE_API_KEY": "env_key"}) + def test_init_with_env_api_key(self): + # Test initialization with API key from environment + embedding_func = QwenDenseEmbedding(dimension=128) + assert embedding_func._api_key == "env_key" + + @patch.dict(os.environ, {"DASHSCOPE_API_KEY": ""}) + def test_init_with_empty_env_api_key(self): + # Test initialization with empty API key from environment + with pytest.raises(ValueError, match="DashScope API key is required"): + QwenDenseEmbedding(dimension=128) + + def test_model_property(self): + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") + assert embedding_func.model == "text-embedding-v4" + + embedding_func = QwenDenseEmbedding( + dimension=128, model="custom-model", api_key="test_key" + ) + assert embedding_func.model == "custom-model" + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_with_empty_text(self, mock_require_module): + # Test embed method with empty text raises ValueError + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed("") + + with pytest.raises(TypeError): + embedding_func.embed(None) + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_success(self, mock_require_module): + # Test successful embedding + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + mock_response.output = {"embeddings": [{"embedding": [0.1, 0.2, 0.3]}]} + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenDenseEmbedding(dimension=3, api_key="test_key") + # Clear cache to avoid interference + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + assert result == [0.1, 0.2, 0.3] + mock_dashscope.TextEmbedding.call.assert_called_once_with( + model="text-embedding-v4", + input="test text", + dimension=3, + output_type="dense", + ) + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_http_error(self, mock_require_module): + # Test embedding with HTTP error + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.BAD_REQUEST + mock_response.message = "Bad Request" + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError): + embedding_func.embed("test text") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_invalid_response(self, mock_require_module): + # Test embedding with invalid response (wrong number of embeddings) + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + mock_response.output = {"embeddings": []} + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError): + embedding_func.embed("test text") + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_embed_success(self): + """Integration test with real DashScope API. + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + export DASHSCOPE_API_KEY=your-api-key + """ + embedding_func = QwenDenseEmbedding(dimension=128) + dense = embedding_func("test text") + assert len(dense) == 128 + + +# ---------------------------- +# QwenSparseEmbedding Test Case +# ---------------------------- +class TestQwenSparseEmbedding: + """Test suite for QwenSparseEmbedding (Qwen sparse embedding via DashScope API).""" + + def test_init_with_api_key(self): + """Test initialization with explicit API key.""" + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + assert embedding_func._dimension == 1024 + assert embedding_func.model == "text-embedding-v4" + assert embedding_func._api_key == "test_key" + # encoding_type defaults to "query" via extra_params + assert embedding_func.extra_params.get("encoding_type", "query") == "query" + + def test_init_with_custom_encoding_type(self): + """Test initialization with custom encoding type.""" + embedding_func = QwenSparseEmbedding( + dimension=1024, encoding_type="document", api_key="test_key" + ) + assert embedding_func.extra_params.get("encoding_type") == "document" + + @patch.dict(os.environ, {"DASHSCOPE_API_KEY": "env_key"}) + def test_init_with_env_api_key(self): + """Test initialization with API key from environment.""" + embedding_func = QwenSparseEmbedding(dimension=1024) + assert embedding_func._api_key == "env_key" + + @patch.dict(os.environ, {"DASHSCOPE_API_KEY": ""}) + def test_init_without_api_key(self): + """Test initialization fails without API key.""" + with pytest.raises(ValueError, match="DashScope API key is required"): + QwenSparseEmbedding(dimension=1024) + + def test_model_property(self): + """Test model property.""" + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + assert embedding_func.model == "text-embedding-v4" + + embedding_func = QwenSparseEmbedding( + dimension=1024, model="text-embedding-v3", api_key="test_key" + ) + assert embedding_func.model == "text-embedding-v3" + + def test_encoding_type_property(self): + """Test encoding_type via extra_params.""" + query_emb = QwenSparseEmbedding( + dimension=1024, encoding_type="query", api_key="test_key" + ) + assert query_emb.extra_params.get("encoding_type") == "query" + + doc_emb = QwenSparseEmbedding( + dimension=1024, encoding_type="document", api_key="test_key" + ) + assert doc_emb.extra_params.get("encoding_type") == "document" + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_with_empty_text(self, mock_require_module): + """Test embed method with empty text raises ValueError.""" + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed("") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed(" ") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_with_non_string_input(self, mock_require_module): + """Test embed method with non-string input raises TypeError.""" + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + embedding_func.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + embedding_func.embed(None) + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_success(self, mock_require_module): + """Test successful sparse embedding generation.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + # Sparse embedding returns array of {index, value, token} objects + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 10, "value": 0.5, "token": "机器"}, + {"index": 245, "value": 0.8, "token": "学习"}, + {"index": 1023, "value": 1.2, "token": "算法"}, + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + # Clear cache to avoid interference + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + # Verify result is a dict + assert isinstance(result, dict) + # Verify keys are integers + assert all(isinstance(k, int) for k in result.keys()) + # Verify values are floats + assert all(isinstance(v, float) for v in result.values()) + # Verify all values are positive + assert all(v > 0 for v in result.values()) + # Verify sorted by indices + keys = list(result.keys()) + assert keys == sorted(keys) + # Verify specific keys + assert keys == [10, 245, 1023] + + mock_dashscope.TextEmbedding.call.assert_called_once_with( + model="text-embedding-v4", + input="test text", + dimension=1024, + output_type="sparse", + text_type="query", + ) + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_with_document_encoding_type(self, mock_require_module): + """Test embedding with document encoding type.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 5, "value": 0.3, "token": "文档"}, + {"index": 100, "value": 0.7, "token": "内容"}, + {"index": 500, "value": 0.9, "token": "检索"}, + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding( + dimension=1024, encoding_type="document", api_key="test_key" + ) + embedding_func.embed.cache_clear() + result = embedding_func.embed("test document") + + assert isinstance(result, dict) + assert list(result.keys()) == [5, 100, 500] + + # Verify text_type parameter is "document" + call_args = mock_dashscope.TextEmbedding.call.call_args + assert call_args[1]["text_type"] == "document" + assert call_args[1]["output_type"] == "sparse" + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_output_sorted_by_indices(self, mock_require_module): + """Test that output is always sorted by indices in ascending order.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + # Return unsorted indices + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 9999, "value": 1.5, "token": "A"}, + {"index": 5, "value": 2.0, "token": "B"}, + {"index": 1234, "value": 0.8, "token": "C"}, + {"index": 77, "value": 3.2, "token": "D"}, + {"index": 500, "value": 1.1, "token": "E"}, + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + result = embedding_func.embed("test sorting") + + # Verify keys are sorted + result_keys = list(result.keys()) + assert result_keys == sorted(result_keys) + # Verify expected sorted order + assert result_keys == [5, 77, 500, 1234, 9999] + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_filters_zero_values(self, mock_require_module): + """Test that zero and negative values are filtered out.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + # Include zero and negative values + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 10, "value": 0.5, "token": "正"}, + { + "index": 20, + "value": 0.0, + "token": "零", + }, # Should be filtered + { + "index": 30, + "value": -0.3, + "token": "负", + }, # Should be filtered + {"index": 40, "value": 0.8, "token": "正"}, + { + "index": 50, + "value": 0.0, + "token": "零", + }, # Should be filtered + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + result = embedding_func.embed("test filtering") + + # Only positive values should remain + assert list(result.keys()) == [10, 40] + assert all(v > 0 for v in result.values()) + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_http_error(self, mock_require_module): + """Test embedding with HTTP error.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.BAD_REQUEST + mock_response.message = "Bad Request" + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError, match="DashScope API error"): + embedding_func.embed("test text") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_invalid_response_no_embeddings(self, mock_require_module): + """Test embedding with invalid response (no embeddings).""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + mock_response.output = {"embeddings": []} + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError, match="Expected exactly 1 embedding"): + embedding_func.embed("test text") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_invalid_response_not_dict(self, mock_require_module): + """Test embedding with invalid response (sparse_embedding not list).""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + # sparse_embedding should be list, not dict + mock_response.output = { + "embeddings": [{"sparse_embedding": {"index": 10, "value": 0.5}}] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises( + ValueError, match="'sparse_embedding' field is missing or not a list" + ): + embedding_func.embed("test text") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_callable_interface(self, mock_require_module): + """Test that embedding function is callable.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 100, "value": 1.0, "token": "测试"}, + {"index": 200, "value": 0.5, "token": "调用"}, + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + # Test calling the function directly + result = embedding_func("test text") + assert isinstance(result, dict) + assert list(result.keys()) == [100, 200] + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_api_connection_error(self, mock_require_module): + """Test handling of API connection errors.""" + mock_dashscope = MagicMock() + mock_dashscope.TextEmbedding.call.side_effect = Exception("Connection timeout") + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises(RuntimeError, match="Failed to call DashScope API"): + embedding_func.embed("test text") + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_embed_success(self): + """Integration test with real DashScope API. + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + export DASHSCOPE_API_KEY=your-api-key + """ + # Test query embedding + query_emb = QwenSparseEmbedding(dimension=1024, encoding_type="query") + query_vec = query_emb.embed("machine learning") + + assert isinstance(query_vec, dict) + assert len(query_vec) > 0 + assert all(isinstance(k, int) for k in query_vec.keys()) + assert all(isinstance(v, float) and v > 0 for v in query_vec.values()) + + # Verify sorted output + keys = list(query_vec.keys()) + assert keys == sorted(keys) + + # Test document embedding + doc_emb = QwenSparseEmbedding(dimension=1024, encoding_type="document") + doc_vec = doc_emb.embed("Machine learning is a subset of AI") + + assert isinstance(doc_vec, dict) + assert len(doc_vec) > 0 + + # Verify sorted output + doc_keys = list(doc_vec.keys()) + assert doc_keys == sorted(doc_keys) + + +# ---------------------------- +# OpenAIDenseEmbedding Test Case +# ---------------------------- +class TestOpenAIDenseEmbedding: + def test_init_with_api_key(self): + """Test initialization with explicit API key.""" + embedding_func = OpenAIDenseEmbedding(api_key="sk-test-key") + assert embedding_func.dimension == 1536 # Default for text-embedding-3-small + assert embedding_func.model == "text-embedding-3-small" + assert embedding_func._api_key == "sk-test-key" + + @patch.dict(os.environ, {"OPENAI_API_KEY": "sk-env-key"}) + def test_init_with_env_api_key(self): + """Test initialization with API key from environment.""" + embedding_func = OpenAIDenseEmbedding() + assert embedding_func._api_key == "sk-env-key" + + @patch.dict(os.environ, {"OPENAI_API_KEY": ""}) + def test_init_without_api_key(self): + """Test initialization fails without API key.""" + with pytest.raises(ValueError, match="OpenAI API key is required"): + OpenAIDenseEmbedding() + + def test_init_with_custom_dimension(self): + """Test initialization with custom dimension.""" + embedding_func = OpenAIDenseEmbedding( + model="text-embedding-3-large", dimension=1024, api_key="sk-test" + ) + assert embedding_func.dimension == 1024 + assert embedding_func.model == "text-embedding-3-large" + + def test_init_with_base_url(self): + """Test initialization with custom base URL.""" + embedding_func = OpenAIDenseEmbedding( + api_key="sk-test", base_url="https://custom.openai.com/" + ) + assert embedding_func._base_url == "https://custom.openai.com/" + + def test_model_property(self): + """Test model property.""" + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + assert embedding_func.model == "text-embedding-3-small" + + embedding_func = OpenAIDenseEmbedding( + model="text-embedding-ada-002", api_key="sk-test" + ) + assert embedding_func.model == "text-embedding-ada-002" + + def test_extra_params(self): + """Test extra_params property.""" + # Test without extra params + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + assert embedding_func.extra_params == {} + + # Test with extra params + embedding_func = OpenAIDenseEmbedding( + api_key="sk-test", + encoding_format="float", + user="test-user", + ) + assert embedding_func.extra_params == { + "encoding_format": "float", + "user": "test-user", + } + + @patch("zvec.extension.openai_function.require_module") + def test_embed_with_empty_text(self, mock_require_module): + """Test embed method with empty text raises ValueError.""" + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed("") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed(" ") + + @patch("zvec.extension.openai_function.require_module") + def test_embed_with_non_string_input(self, mock_require_module): + """Test embed method with non-string input raises TypeError.""" + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + embedding_func.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + embedding_func.embed(None) + + @patch("zvec.extension.openai_function.require_module") + def test_embed_success(self, mock_require_module): + """Test successful embedding generation.""" + # Mock OpenAI client + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + # Create mock embedding data + fake_embedding = [0.1, 0.2, 0.3] + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(dimension=3, api_key="sk-test") + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + assert result == [0.1, 0.2, 0.3] + mock_client.embeddings.create.assert_called_once_with( + model="text-embedding-3-small", input="test text", dimensions=3 + ) + + @patch("zvec.extension.openai_function.require_module") + def test_embed_with_custom_model(self, mock_require_module): + """Test embedding with custom model.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + fake_embedding = [0.1] * 1536 + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding( + model="text-embedding-ada-002", api_key="sk-test" + ) + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + assert len(result) == 1536 + mock_client.embeddings.create.assert_called_once_with( + model="text-embedding-ada-002", input="test text" + ) + + @patch("zvec.extension.openai_function.require_module") + def test_embed_api_error(self, mock_require_module): + """Test handling of API errors.""" + mock_openai = Mock() + mock_client = Mock() + + # Simulate API error + api_error = Mock() + api_error.__class__.__name__ = "APIError" + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + + mock_client.embeddings.create.side_effect = mock_openai.APIError( + "Rate limit exceeded" + ) + mock_openai.OpenAI.return_value = mock_client + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + embedding_func.embed.cache_clear() + + with pytest.raises(RuntimeError, match="Failed to call OpenAI API"): + embedding_func.embed("test text") + + @patch("zvec.extension.openai_function.require_module") + def test_embed_invalid_response(self, mock_require_module): + """Test handling of invalid API response.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + # Empty response data + mock_response.data = [] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError, match="no embedding data returned"): + embedding_func.embed("test text") + + @patch("zvec.extension.openai_function.require_module") + def test_embed_dimension_mismatch(self, mock_require_module): + """Test handling of dimension mismatch.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + # Return embedding with wrong dimension + fake_embedding = [0.1] * 512 + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(dimension=1536, api_key="sk-test") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError, match="Dimension mismatch"): + embedding_func.embed("test text") + + @patch("zvec.extension.openai_function.require_module") + def test_embed_callable(self, mock_require_module): + """Test that embedding function is callable.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + fake_embedding = [0.1] * 1536 + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + embedding_func.embed.cache_clear() + + # Test calling the function directly + result = embedding_func("test text") + assert isinstance(result, list) + assert len(result) == 1536 + + @patch("zvec.extension.openai_function.require_module") + def test_embed_with_base_url(self, mock_require_module): + """Test embedding with custom base URL.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + fake_embedding = [0.1] * 1536 + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding( + api_key="sk-test", base_url="https://custom.openai.com/" + ) + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + # Verify client was created with custom base URL + mock_openai.OpenAI.assert_called_once_with( + api_key="sk-test", base_url="https://custom.openai.com/" + ) + assert len(result) == 1536 + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_embed_success(self): + """Integration test with real OpenAI API. + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + export OPENAI_API_KEY=sk-... + """ + embedding_func = OpenAIDenseEmbedding( + model="text-embedding-v4", + dimension=256, + base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", + ) + vector = embedding_func.embed("Hello, world!") + assert len(vector) == 256 + assert isinstance(vector, list) + assert all(isinstance(x, float) for x in vector) + + +# ---------------------------- +# DefaultLocalDenseEmbedding Test Case +# ---------------------------- +class TestDefaultLocalDenseEmbedding: + """Test cases for DefaultLocalDenseEmbedding.""" + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_success(self, mock_require_module): + """Test successful initialization with mocked model.""" + # Mock sentence_transformers module + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + # Initialize embedding function + emb_func = DefaultLocalDenseEmbedding() + + # Assertions + assert emb_func.dimension == 384 + assert emb_func.model_name == "all-MiniLM-L6-v2" + assert emb_func.model_source == "huggingface" + assert emb_func.device == "cpu" + mock_st.SentenceTransformer.assert_called_once_with( + "all-MiniLM-L6-v2", device=None, trust_remote_code=True + ) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_with_custom_device(self, mock_require_module): + """Test initialization with custom device.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_model.device = "cuda" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding(device="cuda") + + assert emb_func.device == "cuda" + mock_st.SentenceTransformer.assert_called_once_with( + "all-MiniLM-L6-v2", device="cuda", trust_remote_code=True + ) + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_with_modelscope(self, mock_require_module): + """Test initialization with ModelScope as model source.""" + mock_st = Mock() + mock_ms = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + + def require_module_side_effect(module_name): + if module_name == "sentence_transformers": + return mock_st + elif module_name == "modelscope": + return mock_ms + raise ImportError(f"No module named '{module_name}'") + + mock_require_module.side_effect = require_module_side_effect + + # Mock snapshot_download at the correct import location + with patch( + "modelscope.hub.snapshot_download.snapshot_download", + return_value="/path/to/cached/model", + ): + emb_func = DefaultLocalDenseEmbedding(model_source="modelscope") + + # Assertions + assert emb_func.dimension == 384 + assert emb_func.model_name == "iic/nlp_gte_sentence-embedding_chinese-small" + assert emb_func.model_source == "modelscope" + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_with_invalid_model_source(self, mock_require_module): + """Test initialization with invalid model_source raises ValueError.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + with pytest.raises(ValueError, match="Invalid model_source"): + DefaultLocalDenseEmbedding(model_source="invalid_source") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_success(self, mock_require_module): + """Test successful embedding generation.""" + # Mock embedding output + fake_embedding = np.random.rand(384).astype(np.float32) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method + mock_model.encode = Mock(return_value=fake_embedding) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + result = emb_func.embed("Hello, world!") + + # Assertions + assert isinstance(result, list) + assert len(result) == 384 + assert all(isinstance(x, float) for x in result) + mock_model.encode.assert_called_once_with( + "Hello, world!", + convert_to_numpy=True, + normalize_embeddings=True, + batch_size=32, + ) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_with_normalization(self, mock_require_module): + """Test embedding with L2 normalization.""" + # Create a normalized vector + fake_embedding = np.random.rand(384).astype(np.float32) + fake_embedding = fake_embedding / np.linalg.norm(fake_embedding) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method + mock_model.encode = Mock(return_value=fake_embedding) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding(normalize_embeddings=True) + result = emb_func.embed("Test sentence") + + # Check if vector is normalized (L2 norm should be close to 1.0) + result_array = np.array(result) + norm = np.linalg.norm(result_array) + assert abs(norm - 1.0) < 1e-5 + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_empty_string(self, mock_require_module): + """Test embedding with empty string raises ValueError.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + with pytest.raises(ValueError, match="Input text cannot be empty"): + emb_func.embed("") + + with pytest.raises(ValueError, match="Input text cannot be empty"): + emb_func.embed(" ") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_non_string_input(self, mock_require_module): + """Test embedding with non-string input raises TypeError.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + emb_func.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + emb_func.embed(None) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_callable(self, mock_require_module): + """Test that embedding function is callable.""" + fake_embedding = np.random.rand(384).astype(np.float32) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method + mock_model.encode = Mock(return_value=fake_embedding) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + # Test calling the function directly + result = emb_func("Test text") + assert isinstance(result, list) + assert len(result) == 384 + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_semantic_similarity(self, mock_require_module): + """Test semantic similarity between similar and different texts.""" + # Create mock embeddings for similar and different texts + similar_emb_1 = np.array([1.0, 0.0, 0.0] + [0.0] * 381, dtype=np.float32) + similar_emb_2 = np.array([0.9, 0.1, 0.0] + [0.0] * 381, dtype=np.float32) + different_emb = np.array([0.0, 0.0, 1.0] + [0.0] * 381, dtype=np.float32) + + # Normalize + similar_emb_1 = similar_emb_1 / np.linalg.norm(similar_emb_1) + similar_emb_2 = similar_emb_2 / np.linalg.norm(similar_emb_2) + different_emb = different_emb / np.linalg.norm(different_emb) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method with side_effect for multiple calls + mock_model.encode = Mock( + side_effect=[similar_emb_1, similar_emb_2, different_emb] + ) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + v1 = emb_func.embed("The cat sits on the mat") + v2 = emb_func.embed("A feline rests on a rug") + v3 = emb_func.embed("Python programming") + + # Calculate similarities + similarity_high = np.dot(v1, v2) + similarity_low = np.dot(v1, v3) + + assert similarity_high > similarity_low + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_model_loading_error(self, mock_require_module): + """Test handling of model loading failure.""" + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + mock_st = Mock() + mock_st.SentenceTransformer.side_effect = Exception("Model not found") + mock_require_module.return_value = mock_st + + with pytest.raises( + ValueError, match="Failed to load Sentence Transformer model" + ): + DefaultLocalDenseEmbedding() + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_modelscope_import_error(self, mock_require_module): + """Test handling of ModelScope import error.""" + mock_st = Mock() + + def require_module_side_effect(module_name): + if module_name == "sentence_transformers": + return mock_st + elif module_name == "modelscope": + raise ImportError("No module named 'modelscope'") + + mock_require_module.side_effect = require_module_side_effect + + with pytest.raises( + ImportError, match="ModelScope support requires the 'modelscope' package" + ): + DefaultLocalDenseEmbedding(model_source="modelscope") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_dimension_mismatch(self, mock_require_module): + """Test handling of dimension mismatch in embedding output.""" + # Return embedding with wrong dimension + fake_embedding = np.random.rand(256).astype(np.float32) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method + mock_model.encode = Mock(return_value=fake_embedding) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + with pytest.raises(ValueError, match="Dimension mismatch"): + emb_func.embed("Test text") + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_embedding_generation(self): + """Integration test with real model (requires sentence-transformers). + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + + Note: First run will download the model (~80MB). + """ + emb_func = DefaultLocalDenseEmbedding() + + # Test basic embedding + vector = emb_func.embed("Hello, world!") + assert len(vector) == 384 + assert isinstance(vector, list) + assert all(isinstance(x, float) for x in vector) + + # Test normalization + norm = np.linalg.norm(vector) + assert abs(norm - 1.0) < 1e-5 + + # Test semantic similarity + v1 = emb_func.embed("The cat sits on the mat") + v2 = emb_func.embed("A feline rests on a rug") + v3 = emb_func.embed("Python programming language") + + similarity_high = np.dot(v1, v2) + similarity_low = np.dot(v1, v3) + assert similarity_high > similarity_low + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_model_properties(self, mock_require_module): + """Test model_name and model_source properties.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + # Test Hugging Face + emb_func_hf = DefaultLocalDenseEmbedding(model_source="huggingface") + assert emb_func_hf.model_name == "all-MiniLM-L6-v2" + assert emb_func_hf.model_source == "huggingface" + + # Test ModelScope + with patch( + "modelscope.hub.snapshot_download.snapshot_download", + return_value="/path/to/model", + ): + mock_ms = Mock() + mock_require_module.side_effect = lambda m: ( + mock_st if m == "sentence_transformers" else mock_ms + ) + emb_func_ms = DefaultLocalDenseEmbedding(model_source="modelscope") + assert ( + emb_func_ms.model_name == "iic/nlp_gte_sentence-embedding_chinese-small" + ) + assert emb_func_ms.model_source == "modelscope" + + +# ----------------------------------- +# DefaultLocalSparseEmbedding Test Case +# ----------------------------------- +class TestDefaultLocalSparseEmbedding: + """Test suite for DefaultLocalSparseEmbedding (SPLADE sparse embedding). + + Note: + DefaultLocalSparseEmbedding uses naver/splade-cocondenser-ensembledistil + instead of naver/splade-v3 because: + + - splade-v3 is a gated model requiring Hugging Face authentication + - cocondenser-ensembledistil is publicly accessible + - Performance difference is minimal (~2%) + - Avoids "Access to model is restricted" errors + + This allows all users to run tests without authentication setup. + """ + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_success(self, mock_require_module): + """Test successful initialization. + + Verifies that DefaultLocalSparseEmbedding initializes with the publicly + accessible naver/splade-cocondenser-ensembledistil model instead of + the gated naver/splade-v3 model. + """ + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + assert sparse_emb.model_name == "naver/splade-cocondenser-ensembledistil" + assert sparse_emb.model_source == "huggingface" + assert sparse_emb.device == "cpu" + mock_st.SentenceTransformer.assert_called_once_with( + "naver/splade-cocondenser-ensembledistil", + device=None, + trust_remote_code=True, + ) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_with_custom_device(self, mock_require_module): + """Test initialization with custom device.""" + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cuda" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding(device="cuda") + + assert sparse_emb.device == "cuda" + mock_st.SentenceTransformer.assert_called_once_with( + "naver/splade-cocondenser-ensembledistil", + device="cuda", + trust_remote_code=True, + ) + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_success(self, mock_require_module): + """Test successful sparse embedding generation with official API.""" + import numpy as np + + # Clear model cache to ensure fresh mock + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + # Create a mock sparse matrix that simulates scipy.sparse behavior + # The code will call: sparse_matrix[0].toarray().flatten() + mock_sparse_matrix = Mock() + + # Create a dense array representation with vocab_size=30522 + vocab_size = 30522 + dense_array = np.zeros(vocab_size) + # Set specific non-zero values at indices [10, 245, 1023, 5678] + dense_array[10] = 0.5 + dense_array[245] = 0.8 + dense_array[1023] = 1.2 + dense_array[5678] = 0.3 + + # Mock the method chain: sparse_matrix[0].toarray().flatten() + mock_row = Mock() + mock_dense = Mock() + mock_row.toarray.return_value = mock_dense + mock_dense.flatten.return_value = dense_array + mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row) + + # Also mock hasattr check for 'toarray' + mock_sparse_matrix.toarray = Mock() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods to return sparse matrix + # Must set return_value BEFORE hasattr() check in the code + mock_model.encode_query = Mock(return_value=mock_sparse_matrix) + mock_model.encode_document = Mock(return_value=mock_sparse_matrix) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + result = sparse_emb.embed("machine learning") + + # Verify result is a dictionary + assert isinstance(result, dict) + # Verify keys are integers and values are floats + assert all(isinstance(k, int) for k in result.keys()) + assert all(isinstance(v, float) for v in result.values()) + # Verify all values are positive + assert all(v > 0 for v in result.values()) + # Sparse vectors should have specific dimensions + assert len(result) == 4 + + # Verify output is sorted by indices (keys) + keys = list(result.keys()) + assert keys == sorted(keys), ( + "Sparse vector keys must be sorted in ascending order" + ) + + # Verify expected keys + assert keys == [10, 245, 1023, 5678] + + # Verify encode_query was called with a list + mock_model.encode_query.assert_called_once() + call_args = mock_model.encode_query.call_args[0][0] + assert isinstance(call_args, list) + assert call_args == ["machine learning"] + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_empty_input(self, mock_require_module): + """Test embedding with empty input.""" + mock_st = Mock() + mock_model = Mock() + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + with pytest.raises(ValueError, match="Input text cannot be empty"): + sparse_emb.embed("") + + with pytest.raises(ValueError, match="Input text cannot be empty"): + sparse_emb.embed(" ") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_non_string_input(self, mock_require_module): + """Test embedding with non-string input.""" + mock_st = Mock() + mock_model = Mock() + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + sparse_emb.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + sparse_emb.embed(["text"]) + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_callable_interface(self, mock_require_module): + """Test that DefaultSparseEmbedding is callable.""" + import numpy as np + + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + # Create a mock sparse matrix + mock_sparse_matrix = Mock() + + # Create a dense array representation with vocab_size=30522 + vocab_size = 30522 + dense_array = np.zeros(vocab_size) + # Set specific non-zero values at indices [100, 200, 300] + dense_array[100] = 1.0 + dense_array[200] = 0.5 + dense_array[300] = 0.8 + + # Mock the method chain: sparse_matrix[0].toarray().flatten() + mock_row = Mock() + mock_dense = Mock() + mock_row.toarray.return_value = mock_dense + mock_dense.flatten.return_value = dense_array + mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row) + + # Also mock hasattr check for 'toarray' + mock_sparse_matrix.toarray = Mock() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods + mock_model.encode_query = Mock(return_value=mock_sparse_matrix) + mock_model.encode_document = Mock(return_value=mock_sparse_matrix) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + # Test callable interface + result = sparse_emb("test input") + assert isinstance(result, dict) + assert all(isinstance(k, int) for k in result.keys()) + + # Verify sorted output + keys = list(result.keys()) + assert keys == sorted(keys), "Callable interface must also return sorted keys" + assert keys == [100, 200, 300] + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_model_loading_failure(self, mock_require_module): + """Test handling of model loading failure.""" + # Clear model cache to ensure the test actually tries to load the model + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + mock_st = Mock() + mock_st.SentenceTransformer.side_effect = Exception("Model not found") + mock_require_module.return_value = mock_st + + with pytest.raises( + ValueError, match="Failed to load Sentence Transformer model" + ): + DefaultLocalSparseEmbedding() + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_inference_failure(self, mock_require_module): + """Test handling of inference failure.""" + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods to raise RuntimeError + mock_model.encode_query = Mock(side_effect=RuntimeError("CUDA out of memory")) + mock_model.encode_document = Mock( + side_effect=RuntimeError("CUDA out of memory") + ) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + with pytest.raises(RuntimeError, match="Failed to generate sparse embedding"): + sparse_emb.embed("test input") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_sparse_vector_properties(self, mock_require_module): + """Test properties of sparse vectors (sparsity, non-zero values, sorted order).""" + import numpy as np + + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + # Create a mock sparse matrix that simulates scipy.sparse behavior + # The code will call: sparse_matrix[0].toarray().flatten() + mock_sparse_matrix = Mock() + + # Create a dense array representation with vocab_size=30522 + vocab_size = 30522 + dense_array = np.zeros(vocab_size) + # Set specific non-zero values at indices [50, 100, 200, 400, 500] + dense_array[50] = 3.0 + dense_array[100] = 2.0 + dense_array[200] = 1.5 + dense_array[400] = 2.5 + dense_array[500] = 1.8 + + # Mock the method chain: sparse_matrix[0].toarray().flatten() + mock_row = Mock() + mock_dense = Mock() + mock_row.toarray.return_value = mock_dense + mock_dense.flatten.return_value = dense_array + mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row) + + # Also mock hasattr check for 'toarray' + mock_sparse_matrix.toarray = Mock() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods + mock_model.encode_query = Mock(return_value=mock_sparse_matrix) + mock_model.encode_document = Mock(return_value=mock_sparse_matrix) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + result = sparse_emb.embed("test") + + # Verify sparsity: result should have much fewer dimensions than vocab_size + assert len(result) < vocab_size + # All values should be positive + assert all(v > 0 for v in result.values()) + + # Verify keys are sorted in ascending order + keys = list(result.keys()) + assert keys == sorted(keys), "Sparse vector keys must be sorted" + + # Verify the specific non-zero indices are present and sorted + # Expected order: [50, 100, 200, 400, 500] (sorted) + expected_keys = [50, 100, 200, 400, 500] + assert keys == expected_keys, f"Expected {expected_keys}, got {keys}" + + # First key should be smallest + if len(result) > 0: + first_key = next(iter(result.keys())) + assert first_key == min(result.keys()), "First key must be the smallest" + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_output_sorted_by_indices(self, mock_require_module): + """Test that output dictionary is always sorted by indices (keys) in ascending order.""" + import numpy as np + + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + # Create sparse output with deliberately out-of-order indices + # Non-sequential indices: 9999, 5, 1234, 77, 500 + mock_sparse_matrix = Mock() + + # Create a dense array representation with vocab_size=30522 + vocab_size = 30522 + dense_array = np.zeros(vocab_size) + # Set specific non-zero values at out-of-order indices + dense_array[9999] = 1.5 + dense_array[5] = 2.0 + dense_array[1234] = 0.8 + dense_array[77] = 3.2 + dense_array[500] = 1.1 + + # Mock the method chain: sparse_matrix[0].toarray().flatten() + mock_row = Mock() + mock_dense = Mock() + mock_row.toarray.return_value = mock_dense + mock_dense.flatten.return_value = dense_array + mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row) + + # Also mock hasattr check for 'toarray' + mock_sparse_matrix.toarray = Mock() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods + mock_model.encode_query = Mock(return_value=mock_sparse_matrix) + mock_model.encode_document = Mock(return_value=mock_sparse_matrix) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + result = sparse_emb.embed("test sorting") + + # Extract keys from result + result_keys = list(result.keys()) + + # Verify keys are sorted + assert result_keys == sorted(result_keys), ( + f"Keys must be sorted in ascending order. " + f"Got: {result_keys}, Expected: {sorted(result_keys)}" + ) + + # Verify expected keys are present and in correct order + # Expected sorted order: [5, 77, 500, 1234, 9999] + expected_sorted_keys = [5, 77, 500, 1234, 9999] + assert result_keys == expected_sorted_keys, ( + f"All expected keys should be present in sorted order. " + f"Expected: {expected_sorted_keys}, Got: {result_keys}" + ) + + # Verify first and last keys + assert result_keys[0] == 5, "First key must be minimum" + assert result_keys[-1] == 9999, "Last key must be maximum" + + # Verify iteration order matches sorted order + for i, (key, value) in enumerate(result.items()): + if i > 0: + prev_key = list(result.keys())[i - 1] + assert key > prev_key, ( + f"Key at position {i} must be greater than previous key" + ) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_device_property(self, mock_require_module): + """Test device property returns correct device.""" + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cuda" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding(device="cuda") + assert sparse_emb.device == "cuda" + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1 and model download", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_modelscope_source(self, mock_require_module): + """Test initialization with ModelScope source.""" + mock_st = Mock() + mock_ms = Mock() + mock_model = Mock() + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + + # Mock ModelScope snapshot_download + with patch( + "modelscope.hub.snapshot_download.snapshot_download", + return_value="/cache/splade-cocondenser", + ): + mock_require_module.side_effect = lambda m: ( + mock_st if m == "sentence_transformers" else mock_ms + ) + + sparse_emb = DefaultLocalSparseEmbedding(model_source="modelscope") + + assert sparse_emb.model_name == "naver/splade-cocondenser-ensembledistil" + assert sparse_emb.model_source == "modelscope" + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1 and model download", + ) + def test_integration_real_model(self): + """Integration test with real SPLADE model (requires model download). + + This test uses naver/splade-cocondenser-ensembledistil instead of + naver/splade-v3 because splade-v3 requires Hugging Face authentication. + The cocondenser-ensembledistil model is publicly accessible and provides + comparable performance. + + To run this test: + export ZVEC_RUN_INTEGRATION_TESTS=1 + pytest tests/test_embedding.py::TestDefaultSparseEmbedding::test_integration_real_model -v + + Note: First run will download ~100MB model from Hugging Face. + + Alternative models: + If you have access to splade-v3, you can create a custom embedding + class following the example in DefaultSparseEmbedding docstring. + """ + # Clear model cache to ensure fresh load + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + sparse_emb = DefaultLocalSparseEmbedding() + + # Test with real input + text = "machine learning and artificial intelligence" + result = sparse_emb.embed(text) + + # Verify result structure + assert isinstance(result, dict) + assert len(result) > 0 + assert all(isinstance(k, int) and k >= 0 for k in result.keys()) + assert all(isinstance(v, float) and v > 0 for v in result.values()) + + # SPLADE typically produces 100-300 non-zero dimensions + assert 50 < len(result) < 500 + + # Verify keys are sorted in ascending order + keys = list(result.keys()) + assert keys == sorted(keys), "Real model output must be sorted by indices" + + # Test callable interface + result2 = sparse_emb(text) + assert result == result2 + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1", + ) + def test_integration_multiple_inputs(self): + """Integration test with multiple different inputs.""" + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + sparse_emb = DefaultLocalSparseEmbedding() + + texts = [ + "Hello, world!", + "Machine learning is fascinating", + "Python programming language", + ] + + results = [sparse_emb.embed(text) for text in texts] + + # All results should be different + assert len(results) == 3 + assert all(isinstance(r, dict) for r in results) + + # Different inputs should produce different sparse vectors + assert results[0] != results[1] + assert results[1] != results[2] + + # All results must be sorted by indices + for i, result in enumerate(results): + keys = list(result.keys()) + assert keys == sorted(keys), f"Result {i} must have sorted keys" + + +# ---------------------------- +# BM25EmbeddingFunction Test Case +# ---------------------------- +class TestBM25EmbeddingFunction: + """Test suite for BM25EmbeddingFunction (BM25-based sparse embedding using DashText SDK).""" + + def test_init_with_built_in_encoder(self): + """Test successful initialization with built-in encoder (no corpus).""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + # Test with default language (Chinese) + bm25 = BM25EmbeddingFunction() + + assert bm25.corpus_size == 0 + assert bm25.encoding_type == "query" + assert bm25.language == "zh" + mock_dashtext.SparseVectorEncoder.default.assert_called_once_with(name="zh") + + def test_init_with_custom_encoder(self): + """Test successful initialization with custom encoder (with corpus).""" + corpus = [ + "a cat is a feline and likes to purr", + "a dog is the human's best friend", + "a bird is a beautiful animal that can fly", + ] + + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction(corpus=corpus, b=0.75, k1=1.2) + + assert bm25.corpus_size == 3 + assert bm25.encoding_type == "query" + mock_dashtext.SparseVectorEncoder.assert_called_once_with(b=0.75, k1=1.2) + mock_encoder.train.assert_called_once_with(corpus) + + def test_init_with_empty_corpus(self): + """Test initialization with empty corpus raises ValueError.""" + with pytest.raises(ValueError, match="Corpus must be a non-empty list"): + BM25EmbeddingFunction(corpus=[]) + + def test_init_with_invalid_corpus(self): + """Test initialization with invalid corpus elements.""" + with pytest.raises(ValueError, match="All corpus documents must be strings"): + BM25EmbeddingFunction(corpus=["text", 123, "another"]) + + with pytest.raises(ValueError, match="All corpus documents must be strings"): + BM25EmbeddingFunction(corpus=[None, "text"]) + + def test_init_with_language_parameter(self): + """Test initialization with different language settings.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + # Test English language + bm25_en = BM25EmbeddingFunction(language="en") + assert bm25_en.language == "en" + mock_dashtext.SparseVectorEncoder.default.assert_called_with(name="en") + + def test_init_with_encoding_type(self): + """Test initialization with different encoding types.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + # Test document encoding type + bm25_doc = BM25EmbeddingFunction(encoding_type="document") + assert bm25_doc.encoding_type == "document" + + def test_init_with_missing_dashtext_library(self): + """Test initialization fails when dashtext library is not installed.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_require.side_effect = ImportError("dashtext package is required") + + with pytest.raises(ImportError, match="dashtext package is required"): + BM25EmbeddingFunction() + + def test_embed_with_query_encoding(self): + """Test successful sparse embedding generation with query encoding.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + + # Mock encode_queries to return sparse vector + mock_encoder.encode_queries.return_value = { + 5: 0.89, + 12: 1.45, + 23: 0.67, + 45: 1.12, + } + + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction(encoding_type="query") + # Clear LRU cache to ensure fresh call + bm25.embed.cache_clear() + result = bm25.embed("cat purr loud") + + # Verify result structure + assert isinstance(result, dict) + assert all(isinstance(k, int) for k in result.keys()) + assert all(isinstance(v, float) for v in result.values()) + + # Verify all values are positive + assert all(v > 0 for v in result.values()) + + # Verify output is sorted by indices + keys = list(result.keys()) + assert keys == sorted(keys), "Output must be sorted by indices" + + # Verify expected keys from mock response + assert result == {5: 0.89, 12: 1.45, 23: 0.67, 45: 1.12} + + # Verify encode_queries was called + mock_encoder.encode_queries.assert_called_once_with("cat purr loud") + + def test_embed_with_document_encoding(self): + """Test successful sparse embedding generation with document encoding.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + + # Mock encode_documents to return sparse vector + mock_encoder.encode_documents.return_value = {10: 1.5, 20: 2.3} + + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction(encoding_type="document") + bm25.embed.cache_clear() + result = bm25.embed("document text") + + assert result == {10: 1.5, 20: 2.3} + mock_encoder.encode_documents.assert_called_once_with("document text") + + def test_embed_with_empty_input(self): + """Test embedding with empty input raises ValueError.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + + with pytest.raises(ValueError, match="Input text cannot be empty"): + bm25.embed("") + + with pytest.raises(ValueError, match="Input text cannot be empty"): + bm25.embed(" ") + + def test_embed_with_non_string_input(self): + """Test embedding with non-string input raises TypeError.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + + # Test with hashable non-string types - should get our custom error message + with pytest.raises(TypeError, match="Expected 'input' to be str"): + bm25.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + bm25.embed(None) + + # Test with unhashable type (list) + # Note: lru_cache raises TypeError("unhashable type: 'list'") before our type check + # This is still a valid type error, just caught at a different layer + with pytest.raises(TypeError, match="unhashable type"): + bm25.embed(["text"]) + + def test_embed_callable_interface(self): + """Test that BM25EmbeddingFunction is callable.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_encoder.encode_queries.return_value = {10: 1.5} + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + bm25.embed.cache_clear() + + # Test callable interface + result = bm25("test query") + assert isinstance(result, dict) + assert 10 in result + + def test_embed_output_sorted_by_indices(self): + """Test that output is always sorted by indices in ascending order.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + + # Mock encode_queries with unsorted indices + mock_encoder.encode_queries.return_value = { + 9999: 1.5, + 5: 2.0, + 1234: 0.8, + 77: 3.2, + 500: 1.1, + } + + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + bm25.embed.cache_clear() + result = bm25.embed("test query") + + # Verify keys are sorted + result_keys = list(result.keys()) + assert result_keys == sorted(result_keys), ( + f"Keys must be sorted. Got: {result_keys}, Expected: {sorted(result_keys)}" + ) + + # Verify expected sorted order: [5, 77, 500, 1234, 9999] + expected_keys = [5, 77, 500, 1234, 9999] + assert result_keys == expected_keys + + def test_embed_filters_zero_values(self): + """Test that zero and negative values are filtered out.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + + # Mock encode_queries with zero and negative values + mock_encoder.encode_queries.return_value = { + 0: 1.5, # Positive - should be included + 1: 0.0, # Zero - should be filtered + 2: -0.5, # Negative - should be filtered + } + + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + bm25.embed.cache_clear() + result = bm25.embed("test") + + # Only positive token should be in result + assert 0 in result + assert 1 not in result # Zero value filtered + assert 2 not in result # Negative value filtered + assert all(v > 0 for v in result.values()) + + def test_properties(self): + """Test property accessors.""" + corpus = ["doc1", "doc2", "doc3"] + + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction( + corpus=corpus, + encoding_type="document", + language="en", + b=0.8, + k1=1.5, + custom_param="test", + ) + + assert bm25.corpus_size == 3 + assert bm25.encoding_type == "document" + assert bm25.language == "en" + assert bm25.extra_params == {"custom_param": "test"} + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_dashtext_bm25_embedding(self): + """Integration test with real DashText library. + + To run this test: + export ZVEC_RUN_INTEGRATION_TESTS=1 + pip install dashtext + + Note: This test requires the dashtext package to be installed. + """ + # Test built-in encoder (Chinese) + bm25_zh = BM25EmbeddingFunction(language="zh", encoding_type="query") + + query_zh = "什么是向量检索服务" + result_zh = bm25_zh.embed(query_zh) + + assert isinstance(result_zh, dict) + assert len(result_zh) > 0 + assert all(isinstance(k, int) for k in result_zh.keys()) + assert all(isinstance(v, float) and v > 0 for v in result_zh.values()) + + # Verify sorted output + keys = list(result_zh.keys()) + assert keys == sorted(keys), "Real DashText BM25 output must be sorted" + + # Test custom corpus + corpus = [ + "The cat sits on the mat", + "The dog plays in the garden", + "Birds fly in the sky", + "Fish swim in the water", + ] + + bm25_custom = BM25EmbeddingFunction(corpus=corpus, encoding_type="query") + + query_en = "cat on mat" + result_en = bm25_custom.embed(query_en) + + assert isinstance(result_en, dict) + assert len(result_en) > 0 + assert all(isinstance(k, int) for k in result_en.keys()) + assert all(isinstance(v, float) and v > 0 for v in result_en.values()) + + # Test callable interface + result2 = bm25_custom(query_en) + assert result_en == result2 + + # Verify properties + assert bm25_custom.corpus_size == 4 diff --git a/python/tests/test_fts_query.py b/python/tests/test_fts_query.py new file mode 100644 index 0000000..e87c7c1 --- /dev/null +++ b/python/tests/test_fts_query.py @@ -0,0 +1,158 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for FTS (Full-Text Search) query support in the Python SDK.""" + +import pickle + +import pytest + +from zvec.model.param.query import Fts, Query + + +class TestFtsQueryValidation: + """Test FTS parameter validation in Query dataclass.""" + + def test_fts_query_string_only(self): + """Query with only query_string in Fts should be valid.""" + q = Query( + field_name="content", fts=Fts(query_string='+hello -world "exact phrase"') + ) + q._validate() + assert q.fts.query_string == '+hello -world "exact phrase"' + assert q.fts.match_string is None + assert q.has_fts() is True + + def test_fts_match_string_only(self): + """Query with only match_string in Fts should be valid.""" + q = Query(field_name="content", fts=Fts(match_string="machine learning")) + q._validate() + assert q.fts.match_string == "machine learning" + assert q.fts.query_string is None + assert q.has_fts() is True + + def test_fts_query_string_and_match_string_mutually_exclusive(self): + """Cannot provide both query_string and match_string in Fts.""" + q = Query( + field_name="content", + fts=Fts(query_string="+hello", match_string="hello world"), + ) + with pytest.raises(ValueError, match="mutually exclusive"): + q._validate() + + def test_no_fts(self): + """Query without FTS fields should have has_fts() == False.""" + q = Query(field_name="embedding", vector=[0.1, 0.2, 0.3]) + assert q.has_fts() is False + + def test_vector_and_fts_mutually_exclusive(self): + """Cannot combine vector search with FTS in a single Query.""" + q = Query( + field_name="embedding", + vector=[0.1, 0.2, 0.3], + fts=Fts(match_string="deep learning"), + ) + with pytest.raises(ValueError, match="Cannot combine fts with vector search"): + q._validate() + + def test_fts_without_vector_or_id(self): + """Query with only FTS (no vector, no id) should be valid.""" + q = Query(field_name="content", fts=Fts(query_string="hello")) + q._validate() + assert q.has_vector() is False + assert q.has_id() is False + assert q.has_fts() is True + + +class TestFtsQueryBinding: + """Test FTS binding layer (_Fts).""" + + def test_import_fts_query(self): + """_Fts should be importable from _zvec.param.""" + from zvec._zvec.param import _Fts + + fts = _Fts() + assert fts.query_string == "" + assert fts.match_string == "" + + def test_fts_query_set_fields(self): + """Setting fields on _Fts should work.""" + from zvec._zvec.param import _Fts + + fts = _Fts() + fts.query_string = "+hello -world" + assert fts.query_string == "+hello -world" + + fts2 = _Fts() + fts2.match_string = "machine learning" + assert fts2.match_string == "machine learning" + + def test_fts_query_pickle(self): + """_Fts should support pickling.""" + from zvec._zvec.param import _Fts + + fts = _Fts() + fts.query_string = "+vector search" + fts.match_string = "" + + data = pickle.dumps(fts) + restored = pickle.loads(data) + assert restored.query_string == "+vector search" + assert restored.match_string == "" + + def test_search_query_fts_field(self): + """_SearchQuery should have fts field.""" + from zvec._zvec.param import _Fts, _SearchQuery + + vq = _SearchQuery() + # fts should be None by default (optional) + assert vq.fts is None + + # set fts + fts = _Fts() + fts.query_string = "hello" + vq.fts = fts + assert vq.fts is not None + assert vq.fts.query_string == "hello" + + def test_search_query_pickle_with_fts(self): + """_SearchQuery with fts should survive pickling.""" + from zvec._zvec.param import _Fts, _SearchQuery + + vq = _SearchQuery() + vq.topk = 10 + vq.field_name = "embedding" + fts = _Fts() + fts.match_string = "test query" + vq.fts = fts + + data = pickle.dumps(vq) + restored = pickle.loads(data) + assert restored.topk == 10 + assert restored.field_name == "embedding" + assert restored.fts is not None + assert restored.fts.match_string == "test query" + + def test_search_query_pickle_without_fts(self): + """_SearchQuery without fts should survive pickling.""" + from zvec._zvec.param import _SearchQuery + + vq = _SearchQuery() + vq.topk = 5 + vq.field_name = "vec" + + data = pickle.dumps(vq) + restored = pickle.loads(data) + assert restored.topk == 5 + assert restored.field_name == "vec" + assert restored.fts is None diff --git a/python/tests/test_gil_release.py b/python/tests/test_gil_release.py new file mode 100644 index 0000000..0d57369 --- /dev/null +++ b/python/tests/test_gil_release.py @@ -0,0 +1,273 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests to verify that the GIL is released during native C++ query calls, +enabling true thread-level concurrency for multi-threaded Python applications.""" + +from __future__ import annotations + +import os +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed + +import pytest +import zvec +from zvec import ( + Collection, + CollectionOption, + DataType, + Doc, + FieldSchema, + HnswIndexParam, + Query, + VectorSchema, +) + + +@pytest.fixture(scope="module") +def gil_test_collection(tmp_path_factory) -> Collection: + """Create a collection with enough data to make queries take measurable time.""" + schema = zvec.CollectionSchema( + name="gil_test", + fields=[ + FieldSchema("id", DataType.INT64, nullable=False), + ], + vectors=[ + VectorSchema( + "vec", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + ], + ) + option = CollectionOption(read_only=False, enable_mmap=True) + temp_dir = tmp_path_factory.mktemp("zvec_gil_test") + collection_path = temp_dir / "gil_test_collection" + + coll = zvec.create_and_open(path=str(collection_path), schema=schema, option=option) + + # Insert enough docs to make queries non-trivial + docs = [ + Doc( + id=str(i), + fields={"id": i}, + vectors={"vec": [float(i % 100) + 0.1 * j for j in range(128)]}, + ) + for i in range(500) + ] + result = coll.insert(docs) + for r in result: + assert r.ok() + + yield coll + + try: + coll.destroy() + except Exception: + pass + + +class TestGILRelease: + """Verify that C++ query calls release the GIL, allowing true thread concurrency.""" + + def test_gil_released_during_query(self, gil_test_collection: Collection): + """Prove the GIL is explicitly released during C++ Query calls. + + Strategy: + - Calibrate per-query latency on the current platform (slow archs like + RISC-V can be 10x slower than x86), then dynamically pick a query count + whose total runtime fits comfortably inside switch_interval. + - Set switch_interval well above the projected total query time so that + CPython's involuntary GIL switching will NOT trigger during the run. + - A background thread (using time.sleep(0) to avoid deadlock) counts how + many times it got to run. + - Since total query time < switch_interval, the bg thread can ONLY run if + the C++ code explicitly releases the GIL. + - Reset counter just before queries; check counter > 0 after queries. + """ + query_vec = [1.0] * 128 + + def run_query(): + gil_test_collection.query( + Query(field_name="vec", vector=query_vec), + topk=100, + ) + + # --- Calibrate: estimate per-query latency on this platform --- + # Warm up to avoid first-call overhead skewing the measurement. + for _ in range(3): + run_query() + + calib_iters = 10 + calib_start = time.monotonic() + for _ in range(calib_iters): + run_query() + per_query = max((time.monotonic() - calib_start) / calib_iters, 1e-6) + + # Target total query window ~200ms, capped to a sane range so the test + # remains meaningful on both fast and slow archs. + target_total = 0.2 + num_iters = max(1, min(500, int(target_total / per_query))) + projected_total = per_query * num_iters + # Pick switch_interval with a large safety margin (>=10x, >=2s) to absorb + # GC pauses, CPU throttling, and noisy-neighbor effects on CI / shared VMs. + switch_interval = max(2.0, projected_total * 10.0) + + old_interval = sys.getswitchinterval() + sys.setswitchinterval(switch_interval) + + try: + counter = {"value": 0} + stop_event = threading.Event() + + def background_counter(): + while not stop_event.is_set(): + counter["value"] += 1 + time.sleep(0) # Yield GIL to prevent deadlock + + bg_thread = threading.Thread(target=background_counter, daemon=True) + bg_thread.start() + + # Let bg thread start (sleep releases GIL) + time.sleep(0.05) + + # --- Critical section: reset counter, run queries, capture counter --- + counter["value"] = 0 + + start = time.monotonic() + for _ in range(num_iters): + run_query() + elapsed = time.monotonic() - start + + count_during_queries = counter["value"] + # --- End critical section --- + + stop_event.set() + time.sleep(0.01) + bg_thread.join(timeout=5) + + print( + f"\nPer-query: {per_query * 1000:.2f}ms, iters: {num_iters}, " + f"elapsed: {elapsed:.4f}s, switch_interval: {switch_interval:.2f}s" + ) + print(f"Counter during queries: {count_during_queries}") + + # Verify queries completed within the switch_interval window. + # If they did NOT, the run was contaminated by external jitter (GC, + # throttling, noisy neighbor) rather than a real GIL-release defect, + # so skip instead of failing to avoid flaky CI noise. + if elapsed >= switch_interval: + pytest.skip( + f"Queries took {elapsed:.3f}s >= switch_interval " + f"({switch_interval:.3f}s); calibration was outpaced by " + "runtime jitter, result is inconclusive." + ) + # If elapsed < switch_interval, the ONLY way bg thread could run is + # via explicit GIL release. + assert count_during_queries > 0, ( + "Background thread could not run during C++ execution despite " + "query time < switch_interval. GIL was NOT released." + ) + finally: + sys.setswitchinterval(old_interval) + + def test_parallel_queries_correctness(self, gil_test_collection: Collection): + """Verify parallel queries return correct results and print timing info. + + NOTE: The definitive proof of GIL release is test_gil_released_during_query + (counter + setswitchinterval). This test focuses on parallel correctness and + logs timing for manual inspection, since CI timing is too noisy for assertions. + """ + num_queries = 1000 + query_vec = [1.0] * 128 + + def do_query(): + return gil_test_collection.query( + Query(field_name="vec", vector=query_vec), + topk=100, + ) + + # Serial execution (baseline) + start_serial = time.monotonic() + for _ in range(num_queries): + do_query() + serial_time = time.monotonic() - start_serial + + # Parallel execution + num_workers = os.cpu_count() or 2 + start_parallel = time.monotonic() + with ThreadPoolExecutor(max_workers=num_workers) as executor: + futures = [executor.submit(do_query) for _ in range(num_queries)] + for future in as_completed(futures): + result = future.result() + assert len(result) > 0 + parallel_time = time.monotonic() - start_parallel + + print(f"\nSerial time: {serial_time:.4f}s, Parallel time: {parallel_time:.4f}s") + print( + f"Speedup ratio: {serial_time / parallel_time:.2f}x (workers={num_workers})" + ) + + def test_thread_safety_concurrent_queries(self, gil_test_collection: Collection): + """Verify no crashes or data corruption under concurrent query load.""" + num_threads = 8 + queries_per_thread = 10 + errors = [] + + def worker(thread_id): + try: + for i in range(queries_per_thread): + vec = [float(thread_id + i) + 0.1 * j for j in range(128)] + result = gil_test_collection.query( + Query(field_name="vec", vector=vec), + topk=10, + ) + assert len(result) > 0 + except Exception as e: + errors.append((thread_id, e)) + + threads = [ + threading.Thread(target=worker, args=(tid,)) for tid in range(num_threads) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + + assert len(errors) == 0, f"Errors in threads: {errors}" + + def test_concurrent_fetch_release_gil(self, gil_test_collection: Collection): + """Verify Fetch operations also release the GIL correctly.""" + num_threads = 4 + errors = [] + + def worker(thread_id): + try: + ids = [str(i) for i in range(thread_id * 10, thread_id * 10 + 10)] + result = gil_test_collection.fetch(ids) + assert len(result) > 0 + except Exception as e: + errors.append((thread_id, e)) + + threads = [ + threading.Thread(target=worker, args=(tid,)) for tid in range(num_threads) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + + assert len(errors) == 0, f"Errors in threads: {errors}" diff --git a/python/tests/test_hnsw_contiguous_memory.py b/python/tests/test_hnsw_contiguous_memory.py new file mode 100644 index 0000000..380a0f9 --- /dev/null +++ b/python/tests/test_hnsw_contiguous_memory.py @@ -0,0 +1,415 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for the ``use_contiguous_memory`` optimization on ``HnswIndexParam``. + +The HNSW streamer supports two allocation strategies for graph nodes: + +* ``use_contiguous_memory=False`` (default): each node allocates its own + linked buffer. Lower peak memory usage, worse cache locality. +* ``use_contiguous_memory=True``: a single contiguous arena holds every + node. Higher peak memory usage, better cache locality and search + throughput. + +These tests exercise the Python surface end-to-end and make sure that +when a collection is created / reopened with ``use_contiguous_memory=True`` +the underlying HNSW streamer entity is constructed correctly and serves +search traffic. +""" + +from __future__ import annotations + +import pickle +import sys + +import numpy as np +import pytest + +import zvec +from zvec import ( + Collection, + CollectionOption, + CollectionSchema, + Doc, + FieldSchema, + HnswIndexParam, + HnswQueryParam, + InvertIndexParam, + Query, + VectorSchema, +) +from zvec.typing import DataType, IndexType, MetricType, QuantizeType + + +DIMENSION = 32 +NUM_DOCS = 128 +TOPK = 5 + +# --------------------------------------------------------------------------- + + +def _debug_hnsw_storage_mode(coll: Collection, column: str = "dense") -> str: + """Return the internal HNSW entity storage mode for ``column``. + + Exposes the debug-only introspection hook on the pybind11 ``_Collection``. + Only meaningful after ``optimize()`` has built a persisted HNSW index; on + a pure writing segment it will raise ``KeyError``. + """ + underlying = coll._obj # type: ignore[attr-defined] + return underlying._debug_hnsw_storage_mode(column) + + +def _build_schema(name: str, *, use_contiguous_memory: bool) -> CollectionSchema: + """Create a simple schema with a single FP32 HNSW vector column.""" + return CollectionSchema( + name=name, + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=DIMENSION, + index_param=HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + use_contiguous_memory=use_contiguous_memory, + ), + ), + ], + ) + + +def _generate_docs(rng: np.random.Generator, num: int = NUM_DOCS) -> list[Doc]: + """Produce deterministic documents for insertion.""" + docs: list[Doc] = [] + for i in range(num): + vec = rng.standard_normal(DIMENSION).astype(np.float32) + docs.append( + Doc( + id=str(i), + fields={"id": i}, + vectors={"dense": vec.tolist()}, + ) + ) + return docs + + +def _assert_query_matches(coll: Collection, query_vec: list[float]) -> list[str]: + """Run a top-k vector query and return the returned ids in order.""" + vector_query = Query( + field_name="dense", + vector=query_vec, + param=HnswQueryParam(ef=128), + ) + hits = coll.query(vector_query, topk=TOPK) + # Expect a single result group for the single vector query. + assert hits is not None, "query returned None" + assert len(hits) >= 1, f"expected at least one hit, got {hits!r}" + return [doc.id for doc in hits] + + +# --------------------------------------------------------------------------- +# 1) Pure Python surface: construction / property / to_dict / repr / pickle +# --------------------------------------------------------------------------- + + +class TestHnswIndexParamContiguousMemorySurface: + """Verify the Python binding exposes ``use_contiguous_memory`` correctly.""" + + def test_default_is_false(self): + param = HnswIndexParam() + assert param.use_contiguous_memory is False + + def test_custom_true(self): + param = HnswIndexParam(use_contiguous_memory=True) + assert param.use_contiguous_memory is True + assert param.type == IndexType.HNSW + # other fields keep their default values + assert param.m == 50 + assert param.ef_construction == 500 + + def test_to_dict_includes_use_contiguous_memory(self): + param = HnswIndexParam( + metric_type=MetricType.L2, + m=16, + ef_construction=100, + quantize_type=QuantizeType.FP16, + use_contiguous_memory=True, + ) + data = param.to_dict() + assert data["use_contiguous_memory"] is True + # Make sure existing fields are still present. + assert data["metric_type"] == "L2" + assert data["m"] == 16 + assert data["ef_construction"] == 100 + assert data["quantize_type"] == "FP16" + + def test_repr_contains_flag(self): + on = repr(HnswIndexParam(use_contiguous_memory=True)) + off = repr(HnswIndexParam(use_contiguous_memory=False)) + assert "use_contiguous_memory" in on + assert "use_contiguous_memory" in off + assert "true" in on + assert "false" in off + + def test_readonly_property(self): + param = HnswIndexParam(use_contiguous_memory=True) + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + param.use_contiguous_memory = False # type: ignore[misc] + + def test_pickle_roundtrip(self): + original = HnswIndexParam( + metric_type=MetricType.COSINE, + m=24, + ef_construction=150, + quantize_type=QuantizeType.INT8, + use_contiguous_memory=True, + ) + restored = pickle.loads(pickle.dumps(original)) + assert restored.use_contiguous_memory is True + assert restored.metric_type == MetricType.COSINE + assert restored.m == 24 + assert restored.ef_construction == 150 + assert restored.quantize_type == QuantizeType.INT8 + + +# --------------------------------------------------------------------------- +# 2) End-to-end: create collection, insert, query with contiguous memory on +# --------------------------------------------------------------------------- + + +@pytest.fixture +def rng() -> np.random.Generator: + return np.random.default_rng(seed=42) + + +# NOTE: the ``enable_mmap=False`` (BufferPool) variant is intentionally +# omitted from this fixture. Building a persisted HNSW index via +# ``optimize()`` / ``create_vector_index`` / ``drop_vector_index`` +# currently requires mmap-backed storage, because the BufferPool backend +# has not implemented the ``create_new`` semantics yet and the guard in +# ``SegmentImpl::merge_vector_indexer`` rejects that combination. Once +# BufferPool gains write support, re-add ``False`` to ``params`` (and +# drop the guard in segment.cc) so these end-to-end tests cover both +# storage modes again. +@pytest.fixture(params=[True], ids=["mmap_on"]) +def collection_option(request) -> CollectionOption: + return CollectionOption(read_only=False, enable_mmap=request.param) + + +# Building a new persisted HNSW index currently requires mmap-backed storage +# because the BufferPool backend has not implemented `create_new` semantics +# yet. Collections opened with ``enable_mmap=False`` therefore cannot run +# optimize()/create_vector_index/drop_vector_index. Tests use this fixture +# to know which behaviour to assert, and once BufferPool gains write support +# the guard in segment.cc (and these branches) can be removed together. +@pytest.fixture +def build_index_supported(collection_option: CollectionOption) -> bool: + return bool(collection_option.enable_mmap) + + +# Error message fragments emitted by the NotSupported guard in +# SegmentImpl::merge_vector_indexer / drop_vector_index. If the C++ message +# changes, update these together. +_BUILD_NOT_SUPPORTED_FRAGMENTS = ("not yet supported", "enable_mmap=false") + + +class TestHnswContiguousMemoryEndToEnd: + """End-to-end: schema -> create_and_open -> insert -> query works.""" + + def test_create_with_contiguous_memory_and_query( + self, + tmp_path_factory, + collection_option, + rng, + ): + """With the flag on, the schema round-trips and search works end-to-end. + + After ``optimize()`` the writing segment is compacted into a persisted + segment backed by the configured HNSW entity. We assert both the + user-observable behaviour (schema + search) and, via the debug hook, + that the entity type actually honours ``use_contiguous_memory``. + """ + schema = _build_schema("hnsw_contig_create", use_contiguous_memory=True) + + path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_create" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + # Schema round-trips with the flag set. + vec_schema = coll.schema.vectors[0] + assert vec_schema.index_param.use_contiguous_memory is True + + docs = _generate_docs(rng) + insert_result = coll.insert(docs=docs) + for r in insert_result: + assert r.ok(), f"insert failed: code={r.code()}" + assert coll.stats.doc_count == NUM_DOCS + + # Build persisted HNSW index; this is where the contiguous entity + # is actually instantiated. + coll.optimize() + assert _debug_hnsw_storage_mode(coll) == "contiguous", ( + "use_contiguous_memory=True should produce a contiguous entity" + ) + + # Pick an existing vector as the query; top-1 must be itself. + query_vec = docs[0].vector("dense") + ids = _assert_query_matches(coll, query_vec) + assert ids[0] == "0", f"expected self-recall, got top-1 id={ids[0]}" + finally: + coll.destroy() + + def test_create_without_contiguous_memory_uses_mmap_entity( + self, + tmp_path_factory, + collection_option, + rng, + ): + """Baseline: when the flag is omitted the default (mmap) entity is used.""" + schema = _build_schema("hnsw_contig_default", use_contiguous_memory=False) + path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_default" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + vec_schema = coll.schema.vectors[0] + assert vec_schema.index_param.use_contiguous_memory is False + + docs = _generate_docs(rng) + for r in coll.insert(docs=docs): + assert r.ok() + assert coll.stats.doc_count == NUM_DOCS + + coll.optimize() + # With the flag off and mmap on, the persisted entity must be the + # default mmap layout — specifically, not the contiguous arena. + assert _debug_hnsw_storage_mode(coll) == "mmap", ( + "use_contiguous_memory=False + enable_mmap=True should " + "produce the mmap entity" + ) + + # Search still functions with the default entity backing. + query_vec = docs[0].vector("dense") + ids = _assert_query_matches(coll, query_vec) + assert ids[0] == "0" + finally: + coll.destroy() + + def test_close_and_reopen_with_contiguous_memory( + self, + tmp_path_factory, + collection_option, + rng, + ): + """Reopening a collection must preserve the ``use_contiguous_memory`` flag. + + The core property: the flag survives the schema persist/reload + round-trip so the HNSW streamer entity — constructed lazily on first + persisted-segment build — honours the user's choice. We run + ``optimize()`` after reopen and confirm the contiguous entity was + materialized. + """ + schema = _build_schema("hnsw_contig_reopen", use_contiguous_memory=True) + path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_reopen" + path_str = str(path) + + created = zvec.create_and_open( + path=path_str, schema=schema, option=collection_option + ) + docs = _generate_docs(rng) + for r in created.insert(docs=docs): + assert r.ok() + assert created.stats.doc_count == NUM_DOCS + # Persist pending writes so that reopen reconstructs state from disk. + created.flush() + del created # close the handle + + reopened = zvec.open(path=path_str, option=collection_option) + try: + assert reopened is not None + assert reopened.stats.doc_count == NUM_DOCS + + # Schema persisted the flag across the reopen boundary. + vec_schema = reopened.schema.vectors[0] + assert vec_schema.index_param.use_contiguous_memory is True + + reopened.optimize() + assert _debug_hnsw_storage_mode(reopened) == "contiguous" + + # Entity actually works: exact self-recall + fetch parity. + query_vec = docs[7].vector("dense") + ids = _assert_query_matches(reopened, query_vec) + assert ids[0] == "7" + + fetched = reopened.fetch([d.id for d in docs[:10]]) + assert len(fetched) == 10 + finally: + reopened.destroy() + + def test_result_parity_with_and_without_contiguous_memory( + self, + tmp_path_factory, + rng, + ): + """ + Two collections built from the same documents must return the same + top-k neighbors regardless of whether contiguous memory is enabled: + the flag is a memory-layout optimization and must not alter recall + for identical graph construction parameters on the same data. + """ + docs = _generate_docs(rng) + query_vec = docs[3].vector("dense") + + def _build_and_query(tag: str, flag: bool) -> list[str]: + schema = _build_schema(f"hnsw_parity_{tag}", use_contiguous_memory=flag) + option = CollectionOption(read_only=False, enable_mmap=True) + path = tmp_path_factory.mktemp("zvec") / f"hnsw_parity_{tag}" + coll = zvec.create_and_open(path=str(path), schema=schema, option=option) + try: + for r in coll.insert(docs=docs): + assert r.ok() + coll.optimize() + expected_mode = "contiguous" if flag else "mmap" + assert _debug_hnsw_storage_mode(coll) == expected_mode, ( + f"{tag}: unexpected entity type" + ) + return _assert_query_matches(coll, query_vec) + finally: + coll.destroy() + + ids_off = _build_and_query("off", flag=False) + ids_on = _build_and_query("on", flag=True) + + # The graph is built with the same (m, ef_construction, data, order), + # so top-k results must match exactly. + assert ids_on == ids_off, ( + f"top-{TOPK} results diverged between use_contiguous_memory modes: " + f"on={ids_on}, off={ids_off}" + ) + # Sanity: self-recall is still perfect. + assert ids_on[0] == "3" diff --git a/python/tests/test_jieba_default_dict.py b/python/tests/test_jieba_default_dict.py new file mode 100644 index 0000000..523da8a --- /dev/null +++ b/python/tests/test_jieba_default_dict.py @@ -0,0 +1,201 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end: jieba FTS works without any user configuration. + +`import zvec` is supposed to register the wheel-bundled jieba dict +directory via `set_default_jieba_dict_dir`. With that in place a user can +declare an FTS field with `tokenizer_name="jieba"`, leave `extra_params` +empty, and Chinese full-text search just works. + +Falls back to GTEST_SKIP-equivalent when running against a build that did +not bundle the dict (e.g., source-tree dev install without the install +step). In that case CI will rely on the C++ unit tests instead. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import zvec +from zvec import ( + Collection, + CollectionOption, + DataType, + Doc, + FieldSchema, + FtsIndexParam, +) +from zvec.model.param.query import Fts, Query + + +def _bundled_dict_dir() -> str: + """Path zvec.__init__ would have registered; empty when not bundled.""" + return zvec.get_default_jieba_dict_dir() + + +def _bundled_dict_files_exist() -> bool: + """Whether the registered default actually contains the dict files. + + `importlib.resources` happily returns a path even when the data dir was + not installed (e.g. source-tree dev runs); only an installed wheel has + the files on disk. + """ + import os + + base = _bundled_dict_dir() + if not base: + return False + return os.path.isfile(os.path.join(base, "jieba.dict.utf8")) and os.path.isfile( + os.path.join(base, "hmm_model.utf8") + ) + + +@pytest.fixture(scope="module", autouse=True) +def _require_bundled_dict(): + if not _bundled_dict_files_exist(): + pytest.skip( + "Bundled jieba dict not found at zvec/data/jieba_dict/ — " + "this test requires an installed wheel (not a source-tree dev " + "build without the install step).", + ) + + +@pytest.fixture(scope="function") +def jieba_collection(tmp_path_factory) -> Collection: + """FTS-only collection using jieba tokenizer and no explicit dict path.""" + # env-var shadows GlobalConfig in the priority chain. + if os.environ.get("ZVEC_JIEBA_DICT_DIR"): + pytest.skip("ZVEC_JIEBA_DICT_DIR shadows the bundled default") + temp_dir = tmp_path_factory.mktemp("zvec_jieba_default") + collection_path = temp_dir / "fts_jieba" + + schema = zvec.CollectionSchema( + name="fts_jieba_default", + fields=[ + FieldSchema("title", DataType.STRING, nullable=False), + FieldSchema( + "content", + DataType.STRING, + nullable=False, + # Deliberately omit extra_params — the bundled default must + # be picked up via GlobalConfig.jieba_dict_dir. + index_param=FtsIndexParam( + tokenizer_name="jieba", + filters=["lowercase"], + ), + ), + ], + ) + + coll = zvec.create_and_open( + path=str(collection_path), + schema=schema, + option=CollectionOption(read_only=False, enable_mmap=True), + ) + assert coll is not None + try: + yield coll + finally: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + +def test_jieba_works_without_explicit_dict_path(jieba_collection: Collection): + """User opens collection, inserts CJK doc, searches — no init() / no + extra_params / no env var / no manual setter call. Just `import zvec`.""" + docs = [ + Doc(id="pk_1", fields={"title": "t1", "content": "中华人民共和国成立"}), + Doc(id="pk_2", fields={"title": "t2", "content": "无关文档"}), + ] + insert_results = jieba_collection.insert(docs) + assert all(r.ok() for r in insert_results) + + hits = jieba_collection.query( + queries=Query(field_name="content", fts=Fts(match_string="中华")), + topk=10, + ) + ids = {doc.id for doc in hits} + assert "pk_1" in ids + assert "pk_2" not in ids + + +def test_default_dict_dir_is_registered_on_import(): + """Sanity check: zvec.__init__ registered a non-empty default.""" + assert _bundled_dict_dir() != "" + + +def test_user_can_override_default_at_runtime(): + """zvec.set_default_jieba_dict_dir can be called any time to override.""" + saved = zvec.get_default_jieba_dict_dir() + try: + zvec.set_default_jieba_dict_dir("/tmp/zvec/jieba-override") + assert zvec.get_default_jieba_dict_dir() == "/tmp/zvec/jieba-override" + finally: + zvec.set_default_jieba_dict_dir(saved) + + +@pytest.mark.skipif( + sys.platform == "win32", + reason="os.environ writes may not propagate across CRT to zvec.pyd", +) +def test_env_var_overrides_global_config(monkeypatch, tmp_path_factory): + """ZVEC_JIEBA_DICT_DIR beats GlobalConfig in jieba's resolution chain.""" + bundled = _bundled_dict_dir() + monkeypatch.setenv("ZVEC_JIEBA_DICT_DIR", bundled) + saved_global = zvec.get_default_jieba_dict_dir() + try: + zvec.set_default_jieba_dict_dir("/zvec/intentionally/missing/global") + + temp_dir = tmp_path_factory.mktemp("zvec_jieba_env") + schema = zvec.CollectionSchema( + name="fts_jieba_env", + fields=[ + FieldSchema("title", DataType.STRING, nullable=False), + FieldSchema( + "content", + DataType.STRING, + nullable=False, + index_param=FtsIndexParam( + tokenizer_name="jieba", + filters=["lowercase"], + ), + ), + ], + ) + coll = zvec.create_and_open( + path=str(temp_dir / "fts_jieba_env"), + schema=schema, + option=CollectionOption(read_only=False, enable_mmap=True), + ) + assert coll is not None + try: + results = coll.insert( + [ + Doc(id="pk_1", fields={"title": "t", "content": "搜索引擎技术"}), + ] + ) + assert all(r.ok() for r in results) + hits = coll.query( + queries=Query(field_name="content", fts=Fts(match_string="搜索")), + topk=10, + ) + assert {d.id for d in hits} == {"pk_1"} + finally: + coll.destroy() + finally: + zvec.set_default_jieba_dict_dir(saved_global) diff --git a/python/tests/test_params.py b/python/tests/test_params.py new file mode 100644 index 0000000..8f61393 --- /dev/null +++ b/python/tests/test_params.py @@ -0,0 +1,540 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import sys +import time + + +import numpy as np +import pytest +from zvec import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + FlatIndexParam, + HnswIndexParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + OptimizeOption, + HnswQueryParam, + IVFQueryParam, + Query, + VectorQuery, + IndexType, + MetricType, + QuantizeType, + QuantizerParam, + DataType, + VectorSchema, +) + +from zvec._zvec.param import _SearchQuery + +# ---------------------------- +# Invert Index Param Test Case +# ---------------------------- + + +class TestInvertIndexParam: + def test_default(self): + param = InvertIndexParam() + assert param.enable_range_optimization is False + assert param.enable_extended_wildcard is False + assert param.type == IndexType.INVERT + + def test_custom(self): + param = InvertIndexParam( + enable_range_optimization=True, enable_extended_wildcard=True + ) + assert param.enable_range_optimization is True + assert param.enable_extended_wildcard is True + + def test_readonly(self): + param = InvertIndexParam() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + param.enable_range_optimization = False + param.enable_extended_wildcard = False + + +# ---------------------------- +# Hnsw Index Param Test Case +# ---------------------------- + + +class TestHnswIndexParam: + def test_default(self): + param = HnswIndexParam() + assert param.metric_type == MetricType.IP + assert param.m == 50 + assert param.ef_construction == 500 + assert param.quantize_type == QuantizeType.UNDEFINED + assert param.type == IndexType.HNSW + + def test_custom(self): + param = HnswIndexParam( + metric_type=MetricType.L2, + m=10, + ef_construction=1000, + quantize_type=QuantizeType.FP16, + ) + assert param.metric_type == MetricType.L2 + assert param.m == 10 + assert param.ef_construction == 1000 + assert param.quantize_type == QuantizeType.FP16 + + @pytest.mark.parametrize( + "attr", ["metric_type", "m", "ef_construction", "quantize_type"] + ) + def test_readonly_attributes(self, attr): + param = HnswIndexParam() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + +# ---------------------------- +# Flat Index Param Test Case +# ---------------------------- +class TestFlatIndexParam: + def test_default(self): + param = FlatIndexParam() + assert param.type == IndexType.FLAT + assert param.quantize_type == QuantizeType.UNDEFINED + assert param.metric_type == MetricType.IP + + def test_custom(self): + param = FlatIndexParam( + metric_type=MetricType.L2, quantize_type=QuantizeType.INT8 + ) + assert param.metric_type == MetricType.L2 + assert param.quantize_type == QuantizeType.INT8 + + @pytest.mark.parametrize("attr", ["metric_type", "quantize_type"]) + def test_readonly_attributes(self, attr): + param = FlatIndexParam() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + +# ---------------------------- +# Ivf Index Param Test Case +# ---------------------------- +class TestIVFIndexParam: + def test_default(self): + param = IVFIndexParam() + assert param.metric_type == MetricType.IP + assert param.n_list == 10 + assert param.quantize_type == QuantizeType.UNDEFINED + assert param.type == IndexType.IVF + + def test_custom(self): + param = IVFIndexParam( + metric_type=MetricType.L2, n_list=1000, quantize_type=QuantizeType.FP16 + ) + assert param.metric_type == MetricType.L2 + assert param.n_list == 1000 + assert param.quantize_type == QuantizeType.FP16 + assert param.type == IndexType.IVF + + @pytest.mark.parametrize("attr", ["metric_type", "n_list", "quantize_type"]) + def test_readonly_attributes(self, attr): + param = IVFIndexParam() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + +# ---------------------------- +# CollectionOption Test Case +# ---------------------------- +class TestCollectionOption: + def test_default(self): + option = CollectionOption() + assert option is not None + assert option.read_only == False + assert option.enable_mmap == True + + def test_custom(self): + option = CollectionOption(read_only=True, enable_mmap=False) + assert option.read_only == True + assert option.enable_mmap == False + + option = CollectionOption(read_only=False, enable_mmap=True) + assert option.read_only == False + assert option.enable_mmap == True + + @pytest.mark.parametrize("attr", ["read_only", "enable_mmap"]) + def test_readonly_attributes(self, attr): + param = CollectionOption() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + +# ---------------------------- +# IndexOption Test Case +# ---------------------------- +class TestIndexOption: + def test_default(self): + option = IndexOption() + assert option is not None + assert option.concurrency == 0 + + def test_custom(self): + option = IndexOption(concurrency=10) + assert option.concurrency == 10 + + @pytest.mark.parametrize("attr", ["concurrency"]) + def test_readonly_attributes(self, attr): + param = IndexOption() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + +# ---------------------------- +# AddColumnOption Test Case +# ---------------------------- +class TestAddColumnOption: + def test_default(self): + option = AddColumnOption() + assert option is not None + assert option.concurrency == 0 + + def test_custom(self): + option = AddColumnOption(concurrency=10) + assert option.concurrency == 10 + + @pytest.mark.parametrize("attr", ["concurrency"]) + def test_readonly_attributes(self, attr): + param = AddColumnOption() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + +# ---------------------------- +# AlterColumnOption Test Case +# ---------------------------- +class TestAlterColumnOption: + def test_default(self): + option = AlterColumnOption() + assert option is not None + assert option.concurrency == 0 + + def test_custom(self): + option = AlterColumnOption(concurrency=10) + assert option.concurrency == 10 + + @pytest.mark.parametrize("attr", ["concurrency"]) + def test_readonly_attributes(self, attr): + param = AlterColumnOption() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + +# ---------------------------- +# OptimizeOption Test Case +# ---------------------------- +class TestOptimizeOption: + def test_default(self): + option = OptimizeOption() + assert option is not None + assert option.concurrency == 0 + + def test_custom(self): + option = OptimizeOption(concurrency=10) + assert option.concurrency == 10 + + @pytest.mark.parametrize("attr", ["concurrency"]) + def test_readonly_attributes(self, attr): + param = OptimizeOption() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + +# ---------------------------- +# HnswQueryParam Test Case +# ---------------------------- +class TestHnswQueryParam: + def test_default(self): + param = HnswQueryParam() + assert param is not None + assert param.ef == 300 + assert param.is_using_refiner == False + assert param.radius == 0 + assert param.is_linear == False + assert param.prefetch_offset == 8 + assert param.prefetch_lines == 0 + + def test_custom(self): + param = HnswQueryParam( + ef=10, + is_using_refiner=True, + radius=30, + is_linear=True, + extra_params={ + "prefetch_offset": 16, + "prefetch_lines": 4, + }, + ) + assert param.ef == 10 + assert param.is_using_refiner == True + assert param.radius == 30 + assert param.is_linear == True + assert param.prefetch_offset == 16 + assert param.prefetch_lines == 4 + + def test_readonly_attributes(self): + param = HnswQueryParam() + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + param.ef = 10 + param.is_using_refiner = True + param.radius = 30 + param.is_linear = True + + +# # ---------------------------- +# # IVFQueryParam Test Case +# # ---------------------------- +# class TestIVFQueryParam: +# def test_default(self): +# param = IVFQueryParam() +# assert param is not None +# assert param.nprobe == 10 +# assert param.is_using_refiner == False +# assert param.radius == 0 +# assert param.is_linear == False +# assert param.scale_factor == 10 +# +# def test_custom(self): +# param = IVFQueryParam( +# nprobe=20, +# is_using_refiner=True, +# radius=30, +# is_linear=True, +# scale_factor=40 +# ) +# assert param.nprobe == 20 +# assert param.is_using_refiner == True +# assert param.radius == 30 +# assert param.is_linear == True +# assert param.scale_factor == 40 + + +class TestQuery: + def test_init_with_valid_id(self): + vq = Query(field_name="embedding", id="doc123") + assert vq.field_name == "embedding" + assert vq.id == "doc123" + assert vq.vector is None + assert vq.param is None + + def test_init_with_valid_vector(self): + vec = [0.1, 0.2, 0.3] + param = HnswQueryParam(ef=300) + vq = Query(field_name="embedding", vector=vec, param=param) + assert vq.field_name == "embedding" + assert vq.vector == vec + assert vq.param == param + + def test_init_both_id_and_vector_raises_error(self): + with pytest.raises(ValueError): + Query(field_name="embedding", id="doc123", vector=[0.1])._validate() + + def test_init_without_field_name_raises_error(self): + with pytest.raises(ValueError): + Query(field_name=None)._validate() + + def test_has_id_returns_true_when_id_set(self): + vq = Query(field_name="embedding", id="doc123") + assert vq.has_id() + + def test_has_id_returns_false_when_no_id(self): + vq = Query(field_name="embedding", vector=[0.1]) + assert not vq.has_id() + + def test_has_vector_returns_true_with_non_empty_vector(self): + vq = Query(field_name="embedding", vector=[0.1]) + assert vq.has_vector() + + def test_validate_fails_on_both_id_and_vector(self): + vq = Query(field_name="test", id="doc123", vector=[0.1]) + with pytest.raises(ValueError): + vq._validate() + + def test_validate_fails_on_both_id_and_numpy_vector(self): + vq = Query(field_name="test", id="doc123", vector=np.array([0.1])) + with pytest.raises(ValueError, match="Cannot provide both id and vector"): + vq._validate() + + +class TestVectorQueryDeprecated: + def test_deprecation_warning(self): + import warnings + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + vq = VectorQuery(field_name="embedding", id="doc123") + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert "Query" in str(w[0].message) + + def test_isinstance_compatibility(self): + import warnings + + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + vq = VectorQuery(field_name="embedding", id="doc123") + assert isinstance(vq, Query) + + +# ---------------------------- +# QuantizerParam Test Case +# ---------------------------- + + +class TestQuantizerParam: + def test_default(self): + qp = QuantizerParam() + assert qp.enable_rotate is False + + def test_enable_rotate_true(self): + qp = QuantizerParam(enable_rotate=True) + assert qp.enable_rotate is True + + def test_enable_rotate_false(self): + qp = QuantizerParam(enable_rotate=False) + assert qp.enable_rotate is False + + def test_equality(self): + qp1 = QuantizerParam(enable_rotate=True) + qp2 = QuantizerParam(enable_rotate=True) + qp3 = QuantizerParam(enable_rotate=False) + assert qp1 == qp2 + assert qp1 != qp3 + + def test_to_dict(self): + qp = QuantizerParam(enable_rotate=True) + d = qp.to_dict() + assert isinstance(d, dict) + assert d.get("enable_rotate") is True + + def test_repr(self): + qp = QuantizerParam(enable_rotate=True) + r = repr(qp) + assert "enable_rotate" in r or "QuantizerParam" in r + + def test_pickle_roundtrip(self): + import pickle + + qp = QuantizerParam(enable_rotate=True) + data = pickle.dumps(qp) + qp2 = pickle.loads(data) + assert qp2.enable_rotate is True + assert qp == qp2 + + +# ---------------------------- +# HnswIndexParam with QuantizerParam +# ---------------------------- + + +class TestHnswIndexParamQuantizer: + def test_default_quantizer_param(self): + param = HnswIndexParam() + assert param.quantizer_param is not None + assert param.quantizer_param.enable_rotate is False + + def test_with_quantizer_param(self): + qp = QuantizerParam(enable_rotate=True) + param = HnswIndexParam( + metric_type=MetricType.L2, + quantize_type=QuantizeType.INT8, + quantizer_param=qp, + ) + assert param.quantizer_param.enable_rotate is True + assert param.quantize_type == QuantizeType.INT8 + + +# ---------------------------- +# FlatIndexParam with QuantizerParam +# ---------------------------- + + +class TestFlatIndexParamQuantizer: + def test_with_quantizer_param(self): + qp = QuantizerParam(enable_rotate=True) + param = FlatIndexParam( + metric_type=MetricType.L2, + quantize_type=QuantizeType.INT8, + quantizer_param=qp, + ) + assert param.quantizer_param.enable_rotate is True + assert param.quantize_type == QuantizeType.INT8 diff --git a/python/tests/test_query_executor.py b/python/tests/test_query_executor.py new file mode 100644 index 0000000..1bf0841 --- /dev/null +++ b/python/tests/test_query_executor.py @@ -0,0 +1,329 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import Dict, Union +from unittest.mock import MagicMock, patch + +import numpy as np +import math +from zvec._zvec.param import _SearchQuery + +import pytest +from zvec.executor.query_executor import ( + QueryContext, + QueryExecutor, +) +from zvec import ( + RrfReRanker, + WeightedReRanker, + HnswQueryParam, + CollectionSchema, + VectorSchema, + DataType, + MetricType, + Query, + VectorQuery, +) +from zvec.extension.multi_vector_reranker import CallbackReRanker + + +# ---------------------------- +# Mock Collection Schema +# ---------------------------- +class MockCollectionSchema(CollectionSchema): + def __init__(self, vectors=Union[VectorSchema, Dict[str, VectorSchema]]): + self._vectors = ( + [vectors] if not isinstance(vectors, Dict) else list(vectors.values()) + ) + + @property + def vectors(self): + return self._vectors + + +# ---------------------------- +# VectorQuery Test Case +# ---------------------------- +class TestQuery: + def test_init(self): + query = Query(field_name="test_field") + assert query.field_name == "test_field" + assert query.id is None + assert query.vector is None + assert query.param is None + + param = HnswQueryParam() + query = Query( + field_name="test_field", id="test_id", vector=[1, 2, 3], param=param + ) + assert query.field_name == "test_field" + assert query.id == "test_id" + assert query.vector == [1, 2, 3] + assert query.param == param + + def test_has_id(self): + query = Query(field_name="test_field") + assert not query.has_id() + + query = Query(field_name="test_field", id="test_id") + assert query.has_id() + + def test_has_vector(self): + query = Query(field_name="test_field") + assert not query.has_vector() + + query = Query(field_name="test_field", vector=[]) + assert not query.has_vector() + + query = Query(field_name="test_field", vector=[1, 2, 3]) + assert query.has_vector() + + def test_validate_dense_fp16_convert(self): + v = _SearchQuery() + schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP16) + vec = np.array([1.1, 2.1, 3.1], dtype=np.float16) + v.set_vector(schema._get_object(), vec) + ret = v.get_vector(schema._get_object()) + assert np.array_equal(vec, ret) + + def test_validate_dense_fp32_convert(self): + v = _SearchQuery() + schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32) + vec = np.array([1.1, 2.1, 3.1], dtype=np.float32) + v.set_vector(schema._get_object(), vec) + ret = v.get_vector(schema._get_object()) + assert np.array_equal(vec, ret) + + def test_validate_dense_fp64_convert(self): + v = _SearchQuery() + schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP64) + vec = np.array([1.1, 2.1, 3.1], dtype=np.float64) + v.set_vector(schema._get_object(), vec) + ret = v.get_vector(schema._get_object()) + assert np.array_equal(vec, ret) + + def test_validate_dense_int8_convert(self): + v = _SearchQuery() + schema = VectorSchema(name="test", data_type=DataType.VECTOR_INT8) + vec = np.array([1, 2, 3], dtype=np.int8) + v.set_vector(schema._get_object(), vec) + ret = v.get_vector(schema._get_object()) + assert np.array_equal(vec, ret) + + def test_validate_sparse_fp32_convert(self): + v = _SearchQuery() + schema = VectorSchema(name="test", data_type=DataType.SPARSE_VECTOR_FP32) + vec = {1: 1.1, 2: 2.2, 3: 3.3} + v.set_vector(schema._get_object(), vec) + ret = v.get_vector(schema._get_object()) + for k in vec.keys(): + assert math.isclose(vec[k], ret[k], abs_tol=1e-6) + + def test_validate_sparse_fp16_convert(self): + v = _SearchQuery() + schema = VectorSchema(name="test", data_type=DataType.SPARSE_VECTOR_FP16) + vec = {1: 1.1, 2: 2.2, 3: 3.3} + v.set_vector(schema._get_object(), vec) + ret = v.get_vector(schema._get_object()) + for k in vec.keys(): + assert math.isclose(np.float16(vec[k]), ret[k], abs_tol=1e-6) + + +class TestVectorQueryDeprecated: + def test_deprecation_warning(self): + import warnings + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + vq = VectorQuery(field_name="test_field") + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert "Query" in str(w[0].message) + + def test_isinstance_compatibility(self): + import warnings + + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + vq = VectorQuery(field_name="test_field") + assert isinstance(vq, Query) + + +class TestQueryContext: + def test_init(self): + ctx = QueryContext(topk=10) + assert ctx.topk == 10 + assert ctx.queries == [] + assert ctx.filter is None + assert ctx.reranker is None + assert ctx.output_fields is None + assert ctx.include_vector is False + + def test_properties(self): + queries = [Query(field_name="test")] + reranker = RrfReRanker() + output_fields = ["field1", "field2"] + + ctx = QueryContext( + topk=5, + filter="test_filter", + include_vector=True, + queries=queries, + output_fields=output_fields, + reranker=reranker, + ) + + assert ctx.topk == 5 + assert ctx.queries == queries + assert ctx.filter == "test_filter" + assert ctx.reranker == reranker + assert ctx.output_fields == output_fields + assert ctx.include_vector is True + + def test_properties_with_weighted_reranker(self): + queries = [Query(field_name="test")] + reranker = WeightedReRanker( + weights=[1.0], + ) + + ctx = QueryContext( + topk=5, + queries=queries, + reranker=reranker, + ) + + assert ctx.reranker == reranker + assert ctx.reranker.weights == [1.0] + + def test_properties_with_callback_reranker(self): + queries = [Query(field_name="test")] + cb = lambda query_results, topn: [] + reranker = CallbackReRanker(callback=cb) + + ctx = QueryContext( + topk=5, + queries=queries, + reranker=reranker, + ) + + assert ctx.reranker == reranker + + +class TestQueryExecutor: + def test_init(self): + schema = MockCollectionSchema() + executor = QueryExecutor(schema) + assert isinstance(executor, QueryExecutor) + + def test_do_build_without_queries(self): + # When no queries are given, build a single vector-less query. + schema = MockCollectionSchema() + executor = QueryExecutor(schema) + ctx = QueryContext(topk=5, filter="test_filter") + + result = executor._build_queries(ctx, MagicMock()) + assert len(result) == 1 + assert result[0].topk == 5 + assert result[0].filter == "test_filter" + + def test_do_build_query_wo_vector(self): + # Vector-less core query should carry the context query params. + schema = MockCollectionSchema() + executor = QueryExecutor(schema) + ctx = QueryContext(topk=7, filter="f", include_vector=True) + + core_vector = executor._build_base_search_query(ctx) + assert core_vector.topk == 7 + assert core_vector.filter == "f" + assert core_vector.include_vector is True + + def test_do_merge_rerank_results_single_without_reranker(self): + # A single result list without a reranker is returned as-is. + schema = MockCollectionSchema() + executor = QueryExecutor(schema) + ctx = QueryContext(topk=5) + docs_list = [["doc1", "doc2"]] + + result = executor._merge_and_rerank(ctx, docs_list) + assert result == ["doc1", "doc2"] + + def test_do_merge_rerank_results_empty(self): + # Empty results should raise an error. + schema = MockCollectionSchema() + executor = QueryExecutor(schema) + ctx = QueryContext(topk=5) + + with pytest.raises(ValueError, match="Query results is empty"): + executor._merge_and_rerank(ctx, []) + + def test_do_merge_rerank_results_with_reranker(self): + # Multiple result lists are merged through the reranker. + schema = MockCollectionSchema() + executor = QueryExecutor(schema) + reranker = MagicMock() + reranker.rerank.return_value = ["merged"] + ctx = QueryContext( + topk=5, + queries=[Query(field_name="test1"), Query(field_name="test2")], + reranker=reranker, + ) + docs_list = [["d1"], ["d2"]] + + result = executor._merge_and_rerank(ctx, docs_list) + assert result == ["merged"] + reranker.rerank.assert_called_once_with(docs_list, ctx.topk) + + def test_execute_python_pipeline(self): + # Each query is executed serially and converted into a result list. + schema = MockCollectionSchema() + executor = QueryExecutor(schema) + collection = MagicMock() + collection.Query.side_effect = [["raw1"], ["raw2"]] + vectors = [MagicMock(), MagicMock()] + + with patch( + "zvec.executor.query_executor.convert_to_py_doc", + side_effect=lambda doc, schema: doc, + ): + results = executor._execute_python_pipeline(vectors, collection) + assert results == [["raw1"], ["raw2"]] + assert collection.Query.call_count == 2 + + def test_build_search_query_by_missing_id_raises_value_error(self): + vector_schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32) + schema = CollectionSchema(name="test_collection", vectors=[vector_schema]) + executor = QueryExecutor(schema) + ctx = QueryContext(topk=5) + collection = MagicMock() + collection.Fetch.return_value = {} + + with pytest.raises(ValueError, match="Document with id 'missing' not found"): + executor._build_search_query( + ctx, Query(field_name="test", id="missing"), collection + ) + + def test_build_search_query_validates_query(self): + vector_schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32) + schema = CollectionSchema(name="test_collection", vectors=[vector_schema]) + executor = QueryExecutor(schema) + ctx = QueryContext(topk=5) + collection = MagicMock() + + with pytest.raises(ValueError, match="Cannot provide both id and vector"): + executor._build_search_query( + ctx, + Query(field_name="test", id="doc1", vector=np.array([0.1])), + collection, + ) diff --git a/python/tests/test_reranker.py b/python/tests/test_reranker.py new file mode 100644 index 0000000..4c0461b --- /dev/null +++ b/python/tests/test_reranker.py @@ -0,0 +1,948 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from unittest.mock import patch, MagicMock +import pytest +import os + +from zvec import Doc, MetricType, VectorSchema, DataType, FlatIndexParam +from zvec.extension.multi_vector_reranker import ( + CallbackReRanker, + RrfReRanker, + WeightedReRanker, +) +from zvec.extension.sentence_transformer_rerank_function import ( + DefaultLocalReRanker, +) +from zvec.extension.qwen_rerank_function import QwenReRanker + +# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API tests +RUN_INTEGRATION_TESTS = os.environ.get("ZVEC_RUN_INTEGRATION_TESTS", "0") == "1" + + +# ---------------------------- +# RrfReRanker Test Case +# ---------------------------- +class TestRrfReRanker: + def test_init(self): + reranker = RrfReRanker(rank_constant=100) + assert reranker.rank_constant == 100 + + def test_default_rank_constant(self): + reranker = RrfReRanker() + assert reranker.rank_constant == 60 + + def test_rerank(self): + reranker = RrfReRanker(rank_constant=60) + + doc1 = Doc(id="1", score=0.8) + doc2 = Doc(id="2", score=0.7) + doc3 = Doc(id="3", score=0.9) + doc4 = Doc(id="4", score=0.6) + + query_results = [[doc1, doc2, doc3], [doc3, doc1, doc4]] + + results = reranker.rerank(query_results, topn=3) + + assert len(results) <= 3 + + for doc in results: + assert hasattr(doc, "score") + + scores = [doc.score for doc in results] + assert scores == sorted(scores, reverse=True) + + +# ---------------------------- +# WeightedReRanker Test Case +# ---------------------------- +class TestWeightedReRanker: + @staticmethod + def _make_fields(metrics): + return [ + VectorSchema( + name=f"vector{i}", + data_type=DataType.VECTOR_FP32, + dimension=4, + index_param=FlatIndexParam(metric_type=metric), + ) + for i, metric in enumerate(metrics) + ] + + def test_init(self): + reranker = WeightedReRanker([0.7, 0.3]) + assert reranker.weights == [0.7, 0.3] + + def test_rerank(self): + reranker = WeightedReRanker([0.7, 0.3]) + + doc1 = Doc(id="1", score=0.8) + doc2 = Doc(id="2", score=0.7) + doc3 = Doc(id="3", score=0.9) + + query_results = [[doc1, doc2], [doc2, doc3]] + fields = self._make_fields([MetricType.L2, MetricType.L2]) + + results = reranker.rerank(query_results, topn=3, fields=fields) + + assert len(results) <= 3 + + for doc in results: + assert hasattr(doc, "score") + + +# ---------------------------- +# CallbackReRanker Test Case +# ---------------------------- +class TestCallbackReRanker: + def test_rerank(self): + def my_callback(query_results, fields, topn): + all_docs = [] + for docs in query_results: + all_docs.extend(docs) + all_docs.sort(key=lambda d: d.score, reverse=True) + return all_docs[:topn] + + reranker = CallbackReRanker(my_callback) + + doc1 = Doc(id="1", score=0.8) + doc2 = Doc(id="2", score=0.9) + doc3 = Doc(id="3", score=0.7) + doc4 = Doc(id="4", score=0.6) + + query_results = [[doc1, doc2], [doc3, doc4]] + + results = reranker.rerank(query_results, topn=3) + + assert len(results) == 3 + scores = [doc.score for doc in results] + assert scores == sorted(scores, reverse=True) + + def test_callback_with_topn(self): + received_topn = [] + + def my_callback(query_results, fields, topn): + received_topn.append(topn) + return [] + + reranker = CallbackReRanker(my_callback) + reranker.rerank([[Doc(id="1", score=0.5)]], topn=7) + + assert received_topn == [7] + + +# ---------------------------- +# QwenReRanker Test Case +# ---------------------------- +class TestQwenReRanker: + def test_init_without_query(self): + with pytest.raises(ValueError, match="Query is required for QwenReRanker"): + QwenReRanker(api_key="test_key") + + def test_init_without_api_key(self): + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(ValueError, match="DashScope API key is required"): + QwenReRanker(query="test") + + @patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}) + def test_init_with_env_api_key(self): + reranker = QwenReRanker(query="test", rerank_field="content") + assert reranker.query == "test" + assert reranker._api_key == "test_key" + assert reranker.rerank_field == "content" + + def test_init_with_explicit_api_key(self): + reranker = QwenReRanker( + query="test", api_key="explicit_key", rerank_field="content" + ) + assert reranker.query == "test" + assert reranker._api_key == "explicit_key" + + def test_model_property(self): + reranker = QwenReRanker( + query="test", api_key="test_key", rerank_field="content" + ) + assert reranker.model == "gte-rerank-v2" + + reranker = QwenReRanker( + query="test", + model="custom-model", + api_key="test_key", + rerank_field="content", + ) + assert reranker.model == "custom-model" + + def test_query_property(self): + reranker = QwenReRanker( + query="test query", api_key="test_key", rerank_field="content" + ) + assert reranker.query == "test query" + + def test_rerank_field_property(self): + reranker = QwenReRanker(query="test", api_key="test_key", rerank_field="title") + assert reranker.rerank_field == "title" + + def test_rerank_empty_results(self): + reranker = QwenReRanker( + query="test", api_key="test_key", rerank_field="content" + ) + results = reranker.rerank({}) + assert results == [] + + def test_rerank_no_valid_documents(self): + reranker = QwenReRanker( + query="test", api_key="test_key", rerank_field="content" + ) + # Document without the rerank_field + query_results = {"vector1": [Doc(id="1")]} + with pytest.raises(ValueError, match="No documents to rerank"): + reranker.rerank(query_results) + + def test_rerank_skip_empty_content(self): + reranker = QwenReRanker( + query="test", api_key="test_key", rerank_field="content" + ) + query_results = { + "vector1": [ + Doc(id="1", fields={"content": ""}), + Doc(id="2", fields={"content": " "}), + ] + } + with pytest.raises(ValueError, match="No documents to rerank"): + reranker.rerank(query_results) + + @patch("zvec.extension.qwen_function.require_module") + def test_rerank_success(self, mock_require_module): + # Mock dashscope module + mock_dashscope = MagicMock() + mock_require_module.return_value = mock_dashscope + + # Mock API response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.output = { + "results": [ + {"index": 0, "relevance_score": 0.95}, + {"index": 1, "relevance_score": 0.85}, + ] + } + mock_dashscope.TextReRank.call.return_value = mock_response + + reranker = QwenReRanker( + query="test query", api_key="test_key", rerank_field="content" + ) + + query_results = { + "vector1": [ + Doc(id="1", fields={"content": "Document 1"}), + Doc(id="2", fields={"content": "Document 2"}), + ] + } + + results = reranker.rerank(query_results, topn=2) + + assert len(results) == 2 + assert results[0].id == "1" + assert results[0].score == 0.95 + assert results[1].id == "2" + assert results[1].score == 0.85 + + # Verify API call + mock_dashscope.TextReRank.call.assert_called_once_with( + model="gte-rerank-v2", + query="test query", + documents=["Document 1", "Document 2"], + top_n=2, + return_documents=False, + ) + + @patch("zvec.extension.qwen_function.require_module") + def test_rerank_deduplicate_documents(self, mock_require_module): + # Mock dashscope module + mock_dashscope = MagicMock() + mock_require_module.return_value = mock_dashscope + + # Mock API response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.output = { + "results": [ + {"index": 0, "relevance_score": 0.9}, + ] + } + mock_dashscope.TextReRank.call.return_value = mock_response + + reranker = QwenReRanker( + query="test", api_key="test_key", rerank_field="content" + ) + + # Same document in multiple vector results + doc1 = Doc(id="1", fields={"content": "Document 1"}) + query_results = {"vector1": [doc1], "vector2": [doc1]} + + results = reranker.rerank(query_results, topn=5) + + # Should only call API with document once + call_args = mock_dashscope.TextReRank.call.call_args + assert len(call_args[1]["documents"]) == 1 + + @patch("zvec.extension.qwen_function.require_module") + def test_rerank_api_error(self, mock_require_module): + # Mock dashscope module + mock_dashscope = MagicMock() + mock_require_module.return_value = mock_dashscope + + # Mock API error response + mock_response = MagicMock() + mock_response.status_code = 400 + mock_response.message = "Invalid request" + mock_response.code = "InvalidParameter" + mock_dashscope.TextReRank.call.return_value = mock_response + + reranker = QwenReRanker( + query="test", api_key="test_key", rerank_field="content" + ) + + query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]} + + with pytest.raises(ValueError, match="DashScope API error"): + reranker.rerank(query_results) + + @patch("zvec.extension.qwen_function.require_module") + def test_rerank_runtime_error(self, mock_require_module): + # Mock dashscope module that raises exception + mock_dashscope = MagicMock() + mock_require_module.return_value = mock_dashscope + mock_dashscope.TextReRank.call.side_effect = Exception("Network error") + + reranker = QwenReRanker( + query="test", api_key="test_key", rerank_field="content" + ) + + query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]} + + with pytest.raises(RuntimeError, match="Failed to call DashScope API"): + reranker.rerank(query_results) + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_qwen_rerank(self): + """Integration test with real DashScope TextReRank API. + + To run this test, set environment variables: + export ZVEC_RUN_INTEGRATION_TESTS=1 + export DASHSCOPE_API_KEY=your-api-key + """ + # Create reranker with real API + reranker = QwenReRanker( + query="What is machine learning?", + rerank_field="content", + model="gte-rerank-v2", + ) + + # Prepare test documents + query_results = { + "vector1": [ + Doc( + id="1", + score=0.8, + fields={ + "content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data." + }, + ), + Doc( + id="2", + score=0.7, + fields={ + "content": "The weather is nice today with clear skies and sunshine." + }, + ), + Doc( + id="3", + score=0.75, + fields={ + "content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers." + }, + ), + ], + "vector2": [ + Doc( + id="4", + score=0.6, + fields={ + "content": "Python is a popular programming language for data science and machine learning applications." + }, + ), + Doc( + id="5", + score=0.65, + fields={ + "content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder." + }, + ), + ], + } + + # Call real API + results = reranker.rerank(query_results, topn=3) + + # Verify results + assert len(results) <= 3, "Should return at most topn documents" + assert len(results) > 0, "Should return at least one document" + + # All results should have valid scores + for doc in results: + assert hasattr(doc, "score"), "Each document should have a score" + assert isinstance(doc.score, (int, float)), "Score should be numeric" + assert doc.score > 0, "Score should be positive" + + # Verify scores are in descending order + scores = [doc.score for doc in results] + assert scores == sorted(scores, reverse=True), ( + "Results should be sorted by score in descending order" + ) + + # Verify relevant documents are ranked higher + # Document 1 and 3 are about machine learning, should rank higher than weather/recipe docs + result_ids = [doc.id for doc in results] + + # At least one of the ML-related documents should be in top results + ml_related_docs = {"1", "3", "4"} + assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), ( + "ML-related documents should rank higher" + ) + + # Print results for manual verification (useful during development) + print("\nReranking results:") + for i, doc in enumerate(results, 1): + print(f"{i}. ID={doc.id}, Score={doc.score:.4f}") + if doc.fields: + content = doc.field("content") + if content: + print(f" Content: {content[:80]}...") + + +# ---------------------------- +# DefaultLocalReRanker Test Case +# ---------------------------- +class TestDefaultLocalReRanker: + """Test cases for DefaultLocalReRanker.""" + + def test_init_without_query(self): + """Test initialization fails without query.""" + with pytest.raises( + ValueError, match="Query is required for DefaultLocalReRanker" + ): + DefaultLocalReRanker(rerank_field="content") + + def test_init_with_empty_query(self): + """Test initialization fails with empty query.""" + with pytest.raises( + ValueError, match="Query is required for DefaultLocalReRanker" + ): + DefaultLocalReRanker(query="", rerank_field="content") + + @patch("zvec.extension.sentence_transformer_rerank_function.require_module") + def test_init_success(self, mock_require_module): + """Test successful initialization with mocked model.""" + # Mock sentence_transformers module + mock_st = MagicMock() + mock_model = MagicMock() + mock_model.predict = MagicMock() # Cross-encoder has predict method + mock_model.device = "cpu" + mock_st.CrossEncoder.return_value = mock_model + mock_require_module.return_value = mock_st + + reranker = DefaultLocalReRanker( + query="test query", + rerank_field="content", + model_name="cross-encoder/ms-marco-MiniLM-L6-v2", + ) + + assert reranker.query == "test query" + assert reranker.rerank_field == "content" + assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L6-v2" + assert reranker.model_source == "huggingface" + assert reranker.batch_size == 32 + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_rerank_function.require_module") + def test_init_with_custom_params(self, mock_require_module): + """Test initialization with custom parameters.""" + mock_st = MagicMock() + mock_model = MagicMock() + mock_model.predict = MagicMock() + mock_model.device = "cuda" + mock_st.CrossEncoder.return_value = mock_model + mock_require_module.return_value = mock_st + + reranker = DefaultLocalReRanker( + query="custom query", + rerank_field="title", + model_name="cross-encoder/ms-marco-MiniLM-L12-v2", + model_source="modelscope", + device="cuda", + batch_size=64, + ) + + assert reranker.query == "custom query" + assert reranker.rerank_field == "title" + assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L12-v2" + assert reranker.model_source == "modelscope" + assert reranker.batch_size == 64 + + @patch("zvec.extension.sentence_transformer_rerank_function.require_module") + def test_init_invalid_model(self, mock_require_module): + """Test initialization fails with non-cross-encoder model.""" + # Mock a model without predict method (not a cross-encoder) + mock_st = MagicMock() + mock_model = MagicMock(spec=[]) # No predict method + mock_st.CrossEncoder.return_value = mock_model + mock_require_module.return_value = mock_st + + with pytest.raises(ValueError, match="does not appear to be a cross-encoder"): + DefaultLocalReRanker(query="test", rerank_field="content") + + def test_query_property(self): + """Test query property.""" + mock_model = MagicMock() + mock_model.predict = MagicMock() + + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test query", rerank_field="content") + assert reranker.query == "test query" + + def test_rerank_field_property(self): + """Test rerank_field property.""" + mock_model = MagicMock() + mock_model.predict = MagicMock() + + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test", rerank_field="title") + assert reranker.rerank_field == "title" + + def test_batch_size_property(self): + """Test batch_size property.""" + mock_model = MagicMock() + mock_model.predict = MagicMock() + + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker( + query="test", rerank_field="content", batch_size=128 + ) + assert reranker.batch_size == 128 + + def test_rerank_empty_results(self): + """Test rerank with empty query_results.""" + mock_model = MagicMock() + mock_model.predict = MagicMock() + + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test", rerank_field="content") + results = reranker.rerank({}) + assert results == [] + + def test_rerank_no_valid_documents(self): + """Test rerank with documents missing rerank_field.""" + mock_model = MagicMock() + mock_model.predict = MagicMock() + + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test", rerank_field="content") + + # Document without the rerank_field + query_results = {"vector1": [Doc(id="1")]} + with pytest.raises(ValueError, match="No documents to rerank"): + reranker.rerank(query_results) + + def test_rerank_skip_empty_content(self): + """Test rerank skips documents with empty content.""" + mock_model = MagicMock() + mock_model.predict = MagicMock() + + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test", rerank_field="content") + + query_results = { + "vector1": [ + Doc(id="1", fields={"content": ""}), + Doc(id="2", fields={"content": " "}), + ] + } + with pytest.raises(ValueError, match="No documents to rerank"): + reranker.rerank(query_results) + + def test_rerank_success(self): + """Test successful rerank with mocked model.""" + # Mock standard cross-encoder model + mock_model = MagicMock() + + # Mock predict method to return scores + import numpy as np + + mock_scores = np.array([0.95, 0.85, 0.75]) + mock_model.predict.return_value = mock_scores + mock_model.device = "cpu" + + # Mock sentence_transformers module + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test query", rerank_field="content") + + query_results = { + "vector1": [ + Doc(id="1", score=0.8, fields={"content": "Document 1"}), + Doc(id="2", score=0.7, fields={"content": "Document 2"}), + Doc(id="3", score=0.6, fields={"content": "Document 3"}), + ] + } + + results = reranker.rerank(query_results, topn=3) + + # Verify results + assert len(results) == 3 + assert results[0].id == "1" + assert results[0].score == 0.95 + assert results[1].id == "2" + assert results[1].score == 0.85 + assert results[2].id == "3" + assert results[2].score == 0.75 + + # Verify model.predict was called correctly + assert mock_model.predict.called + call_args = mock_model.predict.call_args + pairs = call_args[0][0] + assert len(pairs) == 3 + assert pairs[0] == ["test query", "Document 1"] + assert pairs[1] == ["test query", "Document 2"] + assert pairs[2] == ["test query", "Document 3"] + assert call_args[1]["batch_size"] == 32 + assert call_args[1]["show_progress_bar"] is False + + def test_rerank_with_topn_limit(self): + """Test rerank respects topn limit.""" + mock_model = MagicMock() + + import numpy as np + + mock_scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5]) + mock_model.predict.return_value = mock_scores + + # Mock sentence_transformers module + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test", rerank_field="content") + + query_results = { + "vector1": [ + Doc(id="1", fields={"content": "Doc 1"}), + Doc(id="2", fields={"content": "Doc 2"}), + Doc(id="3", fields={"content": "Doc 3"}), + Doc(id="4", fields={"content": "Doc 4"}), + Doc(id="5", fields={"content": "Doc 5"}), + ] + } + + results = reranker.rerank(query_results, topn=2) + + # Should only return top 2 + assert len(results) == 2 + assert results[0].id == "1" + assert results[0].score == 0.9 + assert results[1].id == "2" + assert results[1].score == 0.8 + + def test_rerank_deduplicate_documents(self): + """Test rerank deduplicates documents across multiple vectors.""" + mock_model = MagicMock() + + import numpy as np + + mock_scores = np.array([0.95, 0.85]) + mock_model.predict.return_value = mock_scores + + # Mock sentence_transformers module + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test", rerank_field="content") + + # Same document in multiple vector results + doc1 = Doc(id="1", fields={"content": "Document 1"}) + doc2 = Doc(id="2", fields={"content": "Document 2"}) + + query_results = { + "vector1": [doc1, doc2], + "vector2": [doc1], # doc1 appears in both + } + + results = reranker.rerank(query_results, topn=5) + + # Should only process each document once + assert len(results) == 2 + assert mock_model.predict.call_count == 1 + + call_args = mock_model.predict.call_args + pairs = call_args[0][0] + assert len(pairs) == 2 # Only 2 unique documents + + def test_rerank_sorting(self): + """Test rerank sorts documents by score in descending order.""" + mock_model = MagicMock() + + import numpy as np + + # Return scores in non-sorted order + mock_scores = np.array([0.6, 0.9, 0.7]) + mock_model.predict.return_value = mock_scores + + # Mock sentence_transformers module + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test", rerank_field="content") + + query_results = { + "vector1": [ + Doc(id="1", fields={"content": "Doc 1"}), + Doc(id="2", fields={"content": "Doc 2"}), + Doc(id="3", fields={"content": "Doc 3"}), + ] + } + + results = reranker.rerank(query_results, topn=3) + + # Should be sorted by score (descending) + assert len(results) == 3 + assert results[0].id == "2" # score 0.9 + assert results[0].score == 0.9 + assert results[1].id == "3" # score 0.7 + assert results[1].score == 0.7 + assert results[2].id == "1" # score 0.6 + assert results[2].score == 0.6 + + def test_rerank_model_error(self): + """Test rerank handles model prediction errors.""" + mock_model = MagicMock() + + # Mock predict to raise exception + mock_model.predict.side_effect = Exception("Model inference error") + + # Mock sentence_transformers module + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker(query="test", rerank_field="content") + + query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]} + + with pytest.raises(RuntimeError, match="Failed to compute rerank scores"): + reranker.rerank(query_results) + + def test_rerank_with_custom_batch_size(self): + """Test rerank uses custom batch_size.""" + mock_model = MagicMock() + + import numpy as np + + mock_scores = np.array([0.9, 0.8]) + mock_model.predict.return_value = mock_scores + + # Mock sentence_transformers module + mock_st = MagicMock() + mock_st.CrossEncoder.return_value = mock_model + + with patch( + "zvec.extension.sentence_transformer_rerank_function.require_module", + return_value=mock_st, + ): + reranker = DefaultLocalReRanker( + query="test", rerank_field="content", batch_size=64 + ) + + query_results = { + "vector1": [ + Doc(id="1", fields={"content": "Doc 1"}), + Doc(id="2", fields={"content": "Doc 2"}), + ] + } + + reranker.rerank(query_results) + + # Verify batch_size is passed to predict + call_args = mock_model.predict.call_args + assert call_args[1]["batch_size"] == 64 + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_sentence_transformer_rerank(self): + """Integration test with real SentenceTransformer cross-encoder model. + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + + Note: This test requires sentence-transformers package and will + download the MS MARCO MiniLM model (~80MB) on first run. + """ + # Create reranker with real model (using default lightweight model) + reranker = DefaultLocalReRanker( + query="What is machine learning?", + rerank_field="content", + ) + + # Prepare test documents + query_results = { + "vector1": [ + Doc( + id="1", + score=0.8, + fields={ + "content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data." + }, + ), + Doc( + id="2", + score=0.7, + fields={ + "content": "The weather is nice today with clear skies and sunshine." + }, + ), + Doc( + id="3", + score=0.75, + fields={ + "content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers." + }, + ), + ], + "vector2": [ + Doc( + id="4", + score=0.6, + fields={ + "content": "Python is a popular programming language for data science and machine learning applications." + }, + ), + Doc( + id="5", + score=0.65, + fields={ + "content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder." + }, + ), + ], + } + + # Call real model + results = reranker.rerank(query_results, topn=3) + + # Verify results + assert len(results) <= 3, "Should return at most topn documents" + assert len(results) > 0, "Should return at least one document" + + # All results should have valid scores + for doc in results: + assert hasattr(doc, "score"), "Each document should have a score" + assert isinstance(doc.score, (int, float)), "Score should be numeric" + + # Verify scores are in descending order + scores = [doc.score for doc in results] + assert scores == sorted(scores, reverse=True), ( + "Results should be sorted by score in descending order" + ) + + # Verify relevant documents are ranked higher + # Documents 1, 3, and 4 are about machine learning, should rank higher + result_ids = [doc.id for doc in results] + + # At least one of the ML-related documents should be in top results + ml_related_docs = {"1", "3", "4"} + assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), ( + "ML-related documents should rank higher" + ) + + # Print results for manual verification (useful during development) + print("\nSentenceTransformer Reranking results:") + for i, doc in enumerate(results, 1): + print(f"{i}. ID={doc.id}, Score={doc.score:.4f}") + if doc.fields: + content = doc.field("content") + if content: + print(f" Content: {content[:80]}...") diff --git a/python/tests/test_schema.py b/python/tests/test_schema.py new file mode 100644 index 0000000..4a4e1a5 --- /dev/null +++ b/python/tests/test_schema.py @@ -0,0 +1,246 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import pytest +from zvec import ( + CollectionSchema, + CollectionStats, + FieldSchema, + VectorSchema, + HnswIndexParam, + InvertIndexParam, + DataType, + IndexType, + MetricType, +) + +# ---------------------------- +# FieldSchema Test Case +# ---------------------------- + + +class TestFieldSchema: + def test_default(self): + field = FieldSchema("field", data_type=DataType.FLOAT) + assert field.name == "field" + assert field.data_type == DataType.FLOAT + assert field.nullable is False + assert field.index_param is None + + def test_custom(self): + field_1 = FieldSchema( + name="float", + data_type=DataType.FLOAT, + nullable=True, + index_param=InvertIndexParam(), + ) + assert field_1.name == "float" + assert field_1.data_type == DataType.FLOAT + assert field_1.nullable is True + assert field_1.index_param.enable_range_optimization is False + + field_2 = FieldSchema( + name="str", + data_type=DataType.STRING, + nullable=True, + index_param=InvertIndexParam(enable_range_optimization=True), + ) + assert field_2.name == "str" + assert field_2.data_type == DataType.STRING + assert field_2.nullable is True + assert field_2.index_param.enable_range_optimization is True + + def test_readonly(self): + field = FieldSchema( + name="float", + data_type=DataType.FLOAT, + nullable=True, + index_param=InvertIndexParam(), + ) + + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + field.index_param = InvertIndexParam(enable_range_optimization=True) + + +# ---------------------------- +# VectorSchema Test Case +# ---------------------------- +class TestVectorSchema: + def test_default(self): + field = VectorSchema("vector", data_type=DataType.VECTOR_FP32, dimension=128) + assert field.name == "vector" + assert field.data_type == DataType.VECTOR_FP32 + assert field.dimension == 128 + assert field.index_param is not None + assert field.index_param.type == IndexType.FLAT + assert field.index_param.metric_type == MetricType.IP + + def test_custom(self): + field = VectorSchema( + name="vector", + data_type=DataType.VECTOR_INT8, + dimension=512, + index_param=HnswIndexParam( + metric_type=MetricType.COSINE, m=15, ef_construction=300 + ), + ) + assert field.name == "vector" + assert field.data_type == DataType.VECTOR_INT8 + assert field.index_param.metric_type == MetricType.COSINE + assert field.index_param.m == 15 + assert field.index_param.ef_construction == 300 + + def test_readonly(self): + field = VectorSchema( + name="vector", + dimension=128, + data_type=DataType.VECTOR_INT8, + ) + + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + field.dimension = 4 + + +# ---------------------------- +# CollectionSchema Test Case +# ---------------------------- +class TestCollectionSchema: + def test_collection_schema_with_single_field(self): + collection_schema = CollectionSchema( + name="test_collection", + fields=FieldSchema( + name="id", + data_type=DataType.INT64, + index_param=InvertIndexParam(), + nullable=False, + ), + vectors=VectorSchema( + name="vector", + data_type=DataType.VECTOR_INT8, + dimension=128, + index_param=HnswIndexParam(), + ), + ) + + assert collection_schema is not None + assert collection_schema.name == "test_collection" + assert len(collection_schema.fields) == 1 + assert len(collection_schema.vectors) == 1 + + field = collection_schema.field("id") + assert field is not None + assert field.name == "id" + assert field.data_type == DataType.INT64 + assert not field.nullable + assert field.index_param.type == IndexType.INVERT + assert not field.index_param.enable_range_optimization + + vector = collection_schema.vector("vector") + assert vector is not None + assert vector.name == "vector" + assert vector.data_type == DataType.VECTOR_INT8 + assert vector.dimension == 128 + assert vector.index_param.type == IndexType.HNSW + assert vector.index_param.m == 50 + assert vector.index_param.ef_construction == 500 + assert vector.index_param.metric_type == MetricType.IP + + def test_collection_schema_with_multi_fields(self): + collection_schema = CollectionSchema( + name="test_collection", + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema( + "name", + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(), + ), + FieldSchema( + "weight", + DataType.INT32, + nullable=True, + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=128, + index_param=HnswIndexParam(), + ), + VectorSchema( + "sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() + ), + ], + ) + assert collection_schema is not None + assert collection_schema.name == "test_collection" + assert len(collection_schema.fields) == 3 + assert len(collection_schema.vectors) == 2 + + field_id = collection_schema.field("id") + assert field_id is not None + assert field_id.name == "id" + assert field_id.data_type == DataType.INT64 + assert not field_id.nullable + assert field_id.index_param.type == IndexType.INVERT + + dense = collection_schema.vector("dense") + assert dense is not None + assert dense.name == "dense" + assert dense.data_type == DataType.VECTOR_FP32 + assert dense.dimension == 128 + assert dense.index_param.type == IndexType.HNSW + + sparse = collection_schema.vector("sparse") + assert sparse is not None + assert sparse.name == "sparse" + assert sparse.data_type == DataType.SPARSE_VECTOR_FP32 + assert sparse.dimension == 0 + assert sparse.index_param.type == IndexType.HNSW + + assert str(collection_schema) is not None + + +# ---------------------------- +# CollectionStats Test Case +# ---------------------------- +class TestCollectionStats: + """ + The constructor of CollectionStats is not provided. + It can only be obtained through collection.stats() + """ + + def test_collection_stats(self): + stats = CollectionStats() + assert stats is not None diff --git a/python/tests/test_typing.py b/python/tests/test_typing.py new file mode 100644 index 0000000..14da5c9 --- /dev/null +++ b/python/tests/test_typing.py @@ -0,0 +1,141 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import pytest +from zvec import ( + DataType, + IndexType, + MetricType, + QuantizeType, + Status, + StatusCode, +) + + +# ---------------------------- +# Enum Test Case +# ---------------------------- +@pytest.mark.parametrize( + "member, name", + [ + (DataType.FLOAT, "FLOAT"), + (IndexType.HNSW, "HNSW"), + (MetricType.COSINE, "COSINE"), + (QuantizeType.INT8, "INT8"), + (StatusCode.OK, "OK"), + ], +) +def test_enum_names(member, name): + assert member.name == name + + +@pytest.mark.parametrize( + "member, value", + [ + (DataType.FLOAT, 8), + (IndexType.HNSW, 1), + (MetricType.COSINE, 3), + (QuantizeType.INT8, 2), + (StatusCode.OK, 0), + ], +) +def test_enum_values(member, value): + assert member.value == value + + +@pytest.mark.parametrize("member", ["L2", "IP", "COSINE"]) +def test_metric_type_has_member(member): + assert member in MetricType.__members__ + + +@pytest.mark.parametrize( + "member", + [ + "STRING", + "BOOL", + "INT32", + "INT64", + "FLOAT", + "DOUBLE", + "UINT32", + "UINT64", + "VECTOR_FP16", + "VECTOR_FP32", + "VECTOR_FP64", + "VECTOR_INT8", + "SPARSE_VECTOR_FP32", + "SPARSE_VECTOR_FP16", + "ARRAY_STRING", + "ARRAY_INT32", + "ARRAY_INT64", + "ARRAY_FLOAT", + "ARRAY_DOUBLE", + "ARRAY_BOOL", + "ARRAY_UINT32", + "ARRAY_UINT64", + ], +) +def test_data_type_has_member(member): + assert member in DataType.__members__ + + +@pytest.mark.parametrize("member", ["HNSW", "IVF", "FLAT", "INVERT"]) +def test_index_type_has_member(member): + assert member in IndexType.__members__ + + +@pytest.mark.parametrize("member", ["FP16", "INT8", "INT4", "UNDEFINED"]) +def test_quantize_type_has_member(member): + assert member in QuantizeType.__members__ + + +@pytest.mark.parametrize( + "member", + [ + "OK", + "UNKNOWN", + "NOT_FOUND", + "ALREADY_EXISTS", + "INVALID_ARGUMENT", + "PERMISSION_DENIED", + "FAILED_PRECONDITION", + "RESOURCE_EXHAUSTED", + "UNAVAILABLE", + "INTERNAL_ERROR", + "NOT_SUPPORTED", + ], +) +def test_status_code_has_member(member): + assert member in StatusCode.__members__ + + +# ---------------------------- +# Status Test Case +# ---------------------------- +class TestStatus: + def test_status_code(self): + status = Status(StatusCode.OK) + assert status.code() == StatusCode.OK + + def test_status_message(self): + status = Status(StatusCode.OK, "OK") + assert status.message() == "OK" + + status = Status(StatusCode.NOT_FOUND, "Not Found") + assert status.message() == "Not Found" + + def test_status_ok(self): + status = Status(StatusCode.OK) + assert status.ok() diff --git a/python/tests/test_util.py b/python/tests/test_util.py new file mode 100644 index 0000000..c5a56c1 --- /dev/null +++ b/python/tests/test_util.py @@ -0,0 +1,89 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from zvec import require_module + + +# ---------------------------- +# require_module func Test Case +# ---------------------------- +def test_require_module_success(): + module = require_module("os") + assert module is not None + assert hasattr(module, "path") + + +def test_require_module_with_submodule_success(): + module = require_module("os.path") + assert module is not None + assert hasattr(module, "join") + + +def test_require_module_import_error(): + with pytest.raises(ImportError) as exc_info: + require_module("nonexistent_module") + + exception_msg = str(exc_info.value) + assert "Required package 'nonexistent_module' is not installed." in exception_msg + + +def test_require_module_with_mitigation_import_error(): + with pytest.raises(ImportError) as exc_info: + require_module("nonexistent_module.submodule", mitigation="custom_package") + + exception_msg = str(exc_info.value) + assert "Required package 'custom_package' is not installed." in exception_msg + assert ( + "Module 'nonexistent_module.submodule' is part of 'nonexistent_module'" + in exception_msg + ) + assert "please pip install 'custom_package'." in exception_msg + + +def test_require_module_submodule_import_error(): + with pytest.raises(ImportError) as exc_info: + require_module("os.nonexistent_submodule") + + exception_msg = str(exc_info.value) + assert ( + "Required package 'os.nonexistent_submodule' is not installed." in exception_msg + ) + assert "Module 'os.nonexistent_submodule' is part of 'os'" in exception_msg + assert "please pip install 'os'." in exception_msg + + +@patch("importlib.import_module") +def test_require_module_wraps_original_exception(mock_import_module): + original_exception = ImportError("Original error") + mock_import_module.side_effect = original_exception + + with pytest.raises(ImportError) as exc_info: + require_module("some_module") + + assert exc_info.value.__cause__ is original_exception + + +@patch("importlib.import_module") +def test_require_module_calls_importlib(mock_import_module): + mock_module = MagicMock() + mock_import_module.return_value = mock_module + + result = require_module("test_module") + + mock_import_module.assert_called_once_with("test_module") + assert result is mock_module diff --git a/python/tests/test_vamana.py b/python/tests/test_vamana.py new file mode 100644 index 0000000..d05e0a8 --- /dev/null +++ b/python/tests/test_vamana.py @@ -0,0 +1,501 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for the Python entry point of the Vamana (DiskANN) dense vector index. + +Mirrors the structure of ``test_hnsw_contiguous_memory.py`` (the closest +hnsw dense reference), and is split into two parts: + +1. **Surface tests** — verify that ``VamanaIndexParam`` / ``VamanaQueryParam`` + are correctly bound: construction defaults, readonly properties, + ``to_dict``, ``__repr__``, pickle round-trip, and that they appear in the + public ``zvec`` namespace with the expected ``IndexType.VAMANA`` value. + +2. **End-to-end tests** — build a collection that uses Vamana on a dense + FP32 column, insert deterministic documents, then run a top-k query + through ``VamanaQueryParam`` on both the writer segment and the + persisted (post-``optimize()``) segment. +""" + +from __future__ import annotations + +import pickle +import sys + +import numpy as np +import pytest + +import zvec +from zvec import ( + Collection, + CollectionOption, + CollectionSchema, + Doc, + FieldSchema, + InvertIndexParam, + VamanaIndexParam, + VamanaQueryParam, + Query, + VectorSchema, +) +from zvec.typing import DataType, IndexType, MetricType, QuantizeType + +DIMENSION = 32 +NUM_DOCS = 128 +TOPK = 5 + +# Defaults pulled from src/include/zvec/core/interface/constants.h. Keep +# in sync with kDefaultVamana* if the engine defaults ever change. +DEFAULT_MAX_DEGREE = 64 +DEFAULT_SEARCH_LIST_SIZE = 100 +DEFAULT_ALPHA = 1.2 +DEFAULT_EF_SEARCH = 200 +DEFAULT_SATURATE_GRAPH = False + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _build_schema( + name: str, + *, + metric_type: MetricType = MetricType.IP, + max_degree: int = 32, + search_list_size: int = 64, + alpha: float = 1.2, + use_contiguous_memory: bool = False, +) -> CollectionSchema: + """Create a simple schema with a single FP32 Vamana vector column.""" + return CollectionSchema( + name=name, + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=DIMENSION, + index_param=VamanaIndexParam( + metric_type=metric_type, + max_degree=max_degree, + search_list_size=search_list_size, + alpha=alpha, + use_contiguous_memory=use_contiguous_memory, + ), + ), + ], + ) + + +def _generate_docs(rng: np.random.Generator, num: int = NUM_DOCS) -> list[Doc]: + """Produce deterministic documents for insertion.""" + docs: list[Doc] = [] + for i in range(num): + vec = rng.standard_normal(DIMENSION).astype(np.float32) + docs.append( + Doc( + id=str(i), + fields={"id": i}, + vectors={"dense": vec.tolist()}, + ) + ) + return docs + + +def _query_topk( + coll: Collection, query_vec: list[float], *, ef_search: int = 64 +) -> list[str]: + """Run a top-k vector query and return the returned ids in order.""" + vector_query = Query( + field_name="dense", + vector=query_vec, + param=VamanaQueryParam(ef_search=ef_search), + ) + hits = coll.query(vector_query, topk=TOPK) + assert hits is not None, "query returned None" + assert len(hits) >= 1, f"expected at least one hit, got {hits!r}" + return [doc.id for doc in hits] + + +# --------------------------------------------------------------------------- +# 1) Surface: construction / property / to_dict / repr / pickle / namespace +# --------------------------------------------------------------------------- + + +class TestVamanaIndexParamSurface: + """Verify the Python binding for ``VamanaIndexParam``.""" + + def test_defaults(self): + param = VamanaIndexParam() + assert param.type == IndexType.VAMANA + assert param.metric_type == MetricType.IP + assert param.max_degree == DEFAULT_MAX_DEGREE + assert param.search_list_size == DEFAULT_SEARCH_LIST_SIZE + assert param.alpha == pytest.approx(DEFAULT_ALPHA) + assert param.saturate_graph is DEFAULT_SATURATE_GRAPH + assert param.use_contiguous_memory is False + assert param.use_id_map is False + assert param.quantize_type == QuantizeType.UNDEFINED + + def test_custom_construction(self): + param = VamanaIndexParam( + metric_type=MetricType.COSINE, + max_degree=48, + search_list_size=128, + alpha=1.5, + saturate_graph=True, + use_contiguous_memory=True, + use_id_map=False, + quantize_type=QuantizeType.INT8, + ) + assert param.type == IndexType.VAMANA + assert param.metric_type == MetricType.COSINE + assert param.max_degree == 48 + assert param.search_list_size == 128 + assert param.alpha == pytest.approx(1.5) + assert param.saturate_graph is True + assert param.use_contiguous_memory is True + assert param.use_id_map is False + assert param.quantize_type == QuantizeType.INT8 + + def test_to_dict_includes_all_fields(self): + param = VamanaIndexParam( + metric_type=MetricType.L2, + max_degree=32, + search_list_size=80, + alpha=1.3, + saturate_graph=True, + use_contiguous_memory=True, + use_id_map=False, + quantize_type=QuantizeType.FP16, + ) + data = param.to_dict() + assert data["type"] == "VAMANA" + assert data["metric_type"] == "L2" + assert data["max_degree"] == 32 + assert data["search_list_size"] == 80 + assert data["alpha"] == pytest.approx(1.3) + assert data["saturate_graph"] is True + assert data["use_contiguous_memory"] is True + assert data["use_id_map"] is False + assert data["quantize_type"] == "FP16" + + def test_repr_contains_key_fields(self): + text = repr( + VamanaIndexParam( + metric_type=MetricType.COSINE, + max_degree=24, + search_list_size=72, + alpha=1.4, + saturate_graph=True, + use_contiguous_memory=True, + ) + ) + # Spot-check the most diagnostic fields are rendered. + assert "VAMANA" in text + assert "COSINE" in text + assert "max_degree" in text and "24" in text + assert "search_list_size" in text and "72" in text + assert "alpha" in text + assert "saturate_graph" in text and "true" in text + assert "use_contiguous_memory" in text and "true" in text + + @pytest.mark.parametrize( + "field, kwargs", + [ + ("max_degree", dict(max_degree=99)), + ("search_list_size", dict(search_list_size=99)), + ("alpha", dict(alpha=1.7)), + ("saturate_graph", dict(saturate_graph=True)), + ("use_contiguous_memory", dict(use_contiguous_memory=True)), + ("use_id_map", dict(use_id_map=True)), + ], + ) + def test_readonly_properties(self, field, kwargs): + param = VamanaIndexParam(**kwargs) + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, field, getattr(param, field)) + + def test_pickle_roundtrip(self): + original = VamanaIndexParam( + metric_type=MetricType.COSINE, + max_degree=48, + search_list_size=120, + alpha=1.4, + saturate_graph=True, + use_contiguous_memory=True, + use_id_map=False, + quantize_type=QuantizeType.INT8, + ) + restored = pickle.loads(pickle.dumps(original)) + assert restored.type == IndexType.VAMANA + assert restored.metric_type == MetricType.COSINE + assert restored.max_degree == 48 + assert restored.search_list_size == 120 + assert restored.alpha == pytest.approx(1.4) + assert restored.saturate_graph is True + assert restored.use_contiguous_memory is True + assert restored.use_id_map is False + assert restored.quantize_type == QuantizeType.INT8 + # to_dict equality is the strongest end-to-end equivalence we have. + assert restored.to_dict() == original.to_dict() + + +class TestVamanaQueryParamSurface: + """Verify the Python binding for ``VamanaQueryParam``.""" + + def test_defaults(self): + q = VamanaQueryParam() + assert q.type == IndexType.VAMANA + assert q.ef_search == DEFAULT_EF_SEARCH + assert q.radius == pytest.approx(0.0) + assert q.is_linear is False + assert q.is_using_refiner is False + assert q.prefetch_offset == 8 + assert q.prefetch_lines == 0 + + def test_custom_construction(self): + q = VamanaQueryParam( + ef_search=300, + radius=0.5, + is_linear=True, + is_using_refiner=True, + extra_params={ + "prefetch_offset": 8, + "prefetch_lines": 2, + }, + ) + assert q.type == IndexType.VAMANA + assert q.ef_search == 300 + assert q.radius == pytest.approx(0.5) + assert q.is_linear is True + assert q.is_using_refiner is True + assert q.prefetch_offset == 8 + assert q.prefetch_lines == 2 + + def test_repr_contains_key_fields(self): + text = repr(VamanaQueryParam(ef_search=128, radius=0.25)) + assert "VAMANA" in text + assert "ef_search" in text and "128" in text + assert "radius" in text + + def test_readonly_ef_search(self): + q = VamanaQueryParam(ef_search=100) + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + q.ef_search = 200 # type: ignore[misc] + + def test_pickle_roundtrip(self): + original = VamanaQueryParam( + ef_search=256, + radius=0.3, + is_linear=False, + is_using_refiner=True, + extra_params={ + "prefetch_offset": 4, + "prefetch_lines": 3, + }, + ) + restored = pickle.loads(pickle.dumps(original)) + assert restored.type == IndexType.VAMANA + assert restored.ef_search == 256 + assert restored.radius == pytest.approx(0.3) + assert restored.is_linear is False + assert restored.is_using_refiner is True + assert restored.prefetch_offset == 4 + assert restored.prefetch_lines == 3 + + +class TestVamanaPublicNamespace: + """The Vamana entry points must be importable from the top-level ``zvec``.""" + + def test_top_level_exports(self): + assert zvec.VamanaIndexParam is VamanaIndexParam + assert zvec.VamanaQueryParam is VamanaQueryParam + assert "VamanaIndexParam" in zvec.__all__ + assert "VamanaQueryParam" in zvec.__all__ + + def test_index_type_enum_member(self): + # Sanity: the IndexType enum exposes VAMANA and it is what the + # bound params advertise. + assert IndexType.VAMANA is not None + assert VamanaIndexParam().type == IndexType.VAMANA + assert VamanaQueryParam().type == IndexType.VAMANA + + +# --------------------------------------------------------------------------- +# 2) End-to-end: create collection, insert, query through the writer segment +# --------------------------------------------------------------------------- + + +@pytest.fixture +def rng() -> np.random.Generator: + return np.random.default_rng(seed=42) + + +# Mirror the hnsw dense test fixture: only the mmap-backed variant is +# currently usable for vector index construction. BufferPool (enable_mmap= +# False) is intentionally omitted because the same write-path guard in +# ``SegmentImpl::merge_vector_indexer`` rejects that combination. +@pytest.fixture(params=[True], ids=["mmap_on"]) +def collection_option(request) -> CollectionOption: + return CollectionOption(read_only=False, enable_mmap=request.param) + + +class TestVamanaEndToEnd: + """End-to-end: schema -> create_and_open -> insert -> query works.""" + + def test_schema_round_trip(self, tmp_path_factory, collection_option): + """The Vamana index params survive the schema persist path.""" + schema = _build_schema( + "vamana_schema_rt", + metric_type=MetricType.COSINE, + max_degree=32, + search_list_size=80, + alpha=1.3, + use_contiguous_memory=True, + ) + path = tmp_path_factory.mktemp("zvec") / "vamana_schema_rt" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + vec_schema = coll.schema.vectors[0] + ip = vec_schema.index_param + assert ip.type == IndexType.VAMANA + assert ip.metric_type == MetricType.COSINE + assert ip.max_degree == 32 + assert ip.search_list_size == 80 + assert ip.alpha == pytest.approx(1.3) + assert ip.use_contiguous_memory is True + finally: + coll.destroy() + + def test_insert_and_query_self_recall( + self, tmp_path_factory, collection_option, rng + ): + """Top-1 of a query equal to an inserted vector must be that vector. + + Exercises the writer-segment Vamana streamer end-to-end through the + Python entry point: ``VamanaIndexParam`` for build and + ``VamanaQueryParam`` for search. + """ + schema = _build_schema("vamana_e2e_recall") + path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_recall" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + docs = _generate_docs(rng) + for r in coll.insert(docs=docs): + assert r.ok(), f"insert failed: code={r.code()}" + assert coll.stats.doc_count == NUM_DOCS + + # Self-recall: query with the i-th inserted vector, expect id i + # to be the top result. + for probe in (0, 7, 42, NUM_DOCS - 1): + query_vec = docs[probe].vector("dense") + ids = _query_topk(coll, query_vec) + assert ids[0] == str(probe), ( + f"expected self-recall at probe={probe}, got top-1 id={ids[0]} " + f"(top-{TOPK}={ids})" + ) + finally: + coll.destroy() + + def test_query_param_ef_search_affects_only_quality( + self, tmp_path_factory, collection_option, rng + ): + """``ef_search`` is a search-time knob and must not crash for any + sensible value. Larger ``ef_search`` should be at least as good as + smaller for self-recall.""" + schema = _build_schema("vamana_e2e_ef") + path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_ef" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + docs = _generate_docs(rng) + for r in coll.insert(docs=docs): + assert r.ok() + + query_vec = docs[3].vector("dense") + ids_small = _query_topk(coll, query_vec, ef_search=16) + ids_large = _query_topk(coll, query_vec, ef_search=256) + + # Both should self-recall the probe vector at top-1. + assert ids_small[0] == "3" + assert ids_large[0] == "3" + assert len(ids_small) == TOPK + assert len(ids_large) == TOPK + finally: + coll.destroy() + + def test_optimize_then_query(self, tmp_path_factory, collection_option, rng): + """The persisted Vamana segment built by ``optimize()`` must serve + queries correctly. + + Until the cmake fix to force-load ``core_knn_vamana_static`` into the + ``_zvec`` pybind module, this path failed at ``VamanaStreamer`` + creation because the global factory registration in + ``vamana_streamer.cc`` was never linked in. This test pins down the + regression. + """ + schema = _build_schema("vamana_e2e_optimize") + path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_optimize" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + docs = _generate_docs(rng) + for r in coll.insert(docs=docs): + assert r.ok() + assert coll.stats.doc_count == NUM_DOCS + + # Snapshot the writer-segment top-k for a probe vector. + query_vec = docs[5].vector("dense") + ids_pre = _query_topk(coll, query_vec) + assert ids_pre[0] == "5" + + # Trigger persisted segment build. Pre-fix this raised + # RuntimeError("Failed to create index"). + coll.optimize() + + # Persisted segment must still serve queries with the same + # top-1 self-recall guarantee. We do not assert full top-k + # equality with the writer segment because the persisted + # streamer may visit nodes in a different order; top-1 self- + # recall is the strong invariant. + ids_post = _query_topk(coll, query_vec) + assert ids_post[0] == "5", ( + f"post-optimize top-1 should still be probe id, got {ids_post}" + ) + assert len(ids_post) == TOPK + finally: + coll.destroy() diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py new file mode 100644 index 0000000..d75a391 --- /dev/null +++ b/python/zvec/__init__.py @@ -0,0 +1,238 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from importlib.metadata import PackageNotFoundError + +# zvec ships a native C++ extension that is only built and tested for 64-bit +# CPython. A 32-bit interpreter would fail to load the extension with an +# obscure error, so fail fast here with an actionable message. +if sys.maxsize <= 2**32: + raise ImportError( + "zvec requires a 64-bit Python interpreter; " + "the current interpreter is 32-bit and is not supported." + ) + + +# Register the wheel-bundled jieba dict dir so `import zvec` alone makes +# the jieba FTS tokenizer usable. Users can still override via +# zvec.init(jieba_dict_dir=...), zvec.set_default_jieba_dict_dir(...), +# ZVEC_JIEBA_DICT_DIR, or per-field FtsIndexParam.extra_params. +try: + from importlib.resources import files as _resource_files + + from zvec._zvec import ( + get_default_jieba_dict_dir, + set_default_jieba_dict_dir, + ) + + set_default_jieba_dict_dir(str(_resource_files("zvec").joinpath("data/jieba_dict"))) +except Exception: + # Custom builds without bundled dict; users must configure explicitly. + pass + + +# ============================== +# Public API — grouped by category +# ============================== + +# —— DiskAnn runtime plugin —— +# Re-export the plugin management entry points defined by the C++ extension. +# DiskAnn normally auto-loads on first use; these APIs let tests and +# diagnostic tools preload the plugin and get a clear error if libaio is +# missing or the plugin shared object cannot be located. +from zvec._zvec import ( + DISKANN_PLUGIN_DLOPEN_FAILED, + DISKANN_PLUGIN_LIBAIO_MISSING, + DISKANN_PLUGIN_OK, + DISKANN_PLUGIN_UNSUPPORTED_PLATFORM, + is_diskann_plugin_loaded, + is_libaio_available, + load_diskann_plugin, +) + +from . import model as model + +# —— Extensions —— +from .extension import ( + BM25EmbeddingFunction, + DefaultLocalDenseEmbedding, + DefaultLocalReRanker, + DefaultLocalSparseEmbedding, + DenseEmbeddingFunction, + OpenAIDenseEmbedding, + OpenAIFunctionBase, + QwenDenseEmbedding, + QwenFunctionBase, + QwenReRanker, + QwenSparseEmbedding, + ReRanker, + RrfReRanker, + SentenceTransformerFunctionBase, + SparseEmbeddingFunction, + WeightedReRanker, +) + +# —— Typing —— +from .model import param as param +from .model import schema as schema + +# —— Core data structures —— +from .model.collection import Collection +from .model.doc import Doc, DocList + +# —— Query & index parameters —— +# —— FTS params (C++ binding) —— +from .model.param import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + DiskAnnIndexParam, + DiskAnnQueryParam, + FlatIndexParam, + FtsIndexParam, + FtsQueryParam, + HnswIndexParam, + HnswQueryParam, + HnswRabitqIndexParam, + HnswRabitqQueryParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + IVFQueryParam, + OptimizeOption, + QuantizerParam, + VamanaIndexParam, + VamanaQueryParam, +) +from .model.param.query import Fts, Query, VectorQuery + +# —— Schema & field definitions —— +from .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema + +# —— tools —— +from .tool import require_module +from .typing import ( + DataType, + IndexType, + MetricType, + QuantizeType, + Status, + StatusCode, +) +from .typing.enum import LogLevel, LogType + +# —— lifecycle —— +from .zvec import create_and_open, init, open + +# ============================== +# Public interface declaration +# ============================== +__all__ = [ + # Zvec functions + "create_and_open", + "init", + "open", + "set_default_jieba_dict_dir", + "get_default_jieba_dict_dir", + # Core classes + "Collection", + "Doc", + "DocList", + # Schema + "CollectionSchema", + "FieldSchema", + "VectorSchema", + "CollectionStats", + # Parameters + "Query", + "VectorQuery", + "Fts", + "FtsIndexParam", + "FtsQueryParam", + "InvertIndexParam", + "HnswIndexParam", + "HnswRabitqIndexParam", + "FlatIndexParam", + "IVFIndexParam", + "DiskAnnIndexParam", + "DiskAnnQueryParam", + "CollectionOption", + "IndexOption", + "OptimizeOption", + "AddColumnOption", + "AlterColumnOption", + "HnswQueryParam", + "HnswRabitqQueryParam", + "IVFQueryParam", + "QuantizerParam", + "VamanaIndexParam", + "VamanaQueryParam", + # Extensions + "DenseEmbeddingFunction", + "SparseEmbeddingFunction", + "QwenFunctionBase", + "OpenAIFunctionBase", + "SentenceTransformerFunctionBase", + "ReRanker", + "DefaultLocalDenseEmbedding", + "DefaultLocalSparseEmbedding", + "BM25EmbeddingFunction", + "OpenAIDenseEmbedding", + "QwenDenseEmbedding", + "QwenSparseEmbedding", + "RrfReRanker", + "WeightedReRanker", + "DefaultLocalReRanker", + "QwenReRanker", + # Typing + "DataType", + "MetricType", + "QuantizeType", + "IndexType", + "LogLevel", + "LogType", + "Status", + "StatusCode", + # Tools + "require_module", + # DiskAnn plugin + "load_diskann_plugin", + "is_diskann_plugin_loaded", + "is_libaio_available", + "DISKANN_PLUGIN_OK", + "DISKANN_PLUGIN_UNSUPPORTED_PLATFORM", + "DISKANN_PLUGIN_LIBAIO_MISSING", + "DISKANN_PLUGIN_DLOPEN_FAILED", +] + +# ============================== +# Version handling +# ============================== +__version__: str + +try: + from importlib.metadata import version +except ImportError: + from importlib_metadata import version # Python < 3.8 + +try: + __version__ = version("zvec") +except Exception: + __version__ = "unknown" diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi new file mode 100644 index 0000000..3e75f93 --- /dev/null +++ b/python/zvec/__init__.pyi @@ -0,0 +1,202 @@ +""" +Zvec core module +""" + +from __future__ import annotations + +import collections + +from . import typing +from .extension import ReRanker, RrfReRanker, WeightedReRanker +from .extension.embedding import DenseEmbeddingFunction +from .model import param, schema +from .model.collection import Collection +from .model.doc import Doc, DocList +from .model.param import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + DiskAnnIndexParam, + DiskAnnQueryParam, + FlatIndexParam, + FtsIndexParam, + FtsQueryParam, + HnswIndexParam, + HnswQueryParam, + HnswRabitqIndexParam, + HnswRabitqQueryParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + IVFQueryParam, + OptimizeOption, + QuantizerParam, + VamanaIndexParam, + VamanaQueryParam, +) +from .model.param.query import Fts, Query, VectorQuery +from .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema +from .tool import require_module +from .typing import ( + DataType, + IndexType, + MetricType, + QuantizeType, + Status, + StatusCode, +) +from .typing.enum import LogLevel, LogType +from .zvec import create_and_open, init, open + +__all__: list = [ + "AddColumnOption", + "AlterColumnOption", + "Collection", + "CollectionOption", + "CollectionSchema", + "CollectionStats", + "DataType", + "DenseEmbeddingFunction", + "DiskAnnIndexParam", + "DiskAnnQueryParam", + "Doc", + "DocList", + "FieldSchema", + "FlatIndexParam", + "Fts", + "FtsIndexParam", + "FtsQueryParam", + "HnswIndexParam", + "HnswQueryParam", + "HnswRabitqIndexParam", + "HnswRabitqQueryParam", + "IVFIndexParam", + "IVFQueryParam", + "IndexOption", + "IndexType", + "InvertIndexParam", + "LogLevel", + "LogType", + "MetricType", + "OptimizeOption", + "QuantizeType", + "QuantizerParam", + "Query", + "ReRanker", + "RrfReRanker", + "Status", + "StatusCode", + "VamanaIndexParam", + "VamanaQueryParam", + "VectorQuery", + "VectorSchema", + "WeightedReRanker", + "create_and_open", + "init", + "open", + "require_module", +] + +class _Collection: + @staticmethod + def CreateAndOpen( + arg0: str, arg1: schema._CollectionSchema, arg2: param.CollectionOption + ) -> _Collection: ... + @staticmethod + def Open(arg0: str, arg1: param.CollectionOption) -> _Collection: ... + def AddColumn( + self, + arg0: schema._FieldSchema, + arg1: str, + arg2: param.AddColumnOption, + ) -> None: ... + def AlterColumn( + self, + arg0: str, + arg1: str, + arg2: schema._FieldSchema, + arg3: param.AlterColumnOption, + ) -> None: ... + def CreateIndex( + self, arg0: str, arg1: param.IndexParam, arg2: param.IndexOption + ) -> None: ... + def Delete(self, arg0: collections.abc.Sequence[str]) -> list[typing.Status]: ... + def DeleteByFilter(self, arg0: str) -> None: ... + def Destroy(self) -> None: ... + def DropColumn(self, arg0: str) -> None: ... + def DropIndex(self, arg0: str) -> None: ... + def Fetch( + self, + pks: collections.abc.Sequence[str], + output_fields: list[str] | None = None, + include_vector: bool = True, + ) -> dict[str, _Doc]: ... + def Flush(self) -> None: ... + def GroupByQuery(self, arg0: ...) -> list[...]: ... + def Insert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ... + def Optimize(self, arg0: param.OptimizeOption) -> None: ... + def Options(self) -> param.CollectionOption: ... + def Path(self) -> str: ... + def Query(self, arg0: param._SearchQuery) -> list[_Doc]: ... + def Schema(self) -> schema._CollectionSchema: ... + def Stats(self) -> schema.CollectionStats: ... + def Update(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ... + def Upsert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ... + def _debug_hnsw_storage_mode(self, column_name: str) -> str: + """Debug-only: returns the storage mode of the HNSW entity on the + given vector column. One of 'mmap', 'buffer_pool', 'contiguous'. + Raises KeyError if no HNSW index exists on the column, or + ValueError if the column's index is not an HNSW index. Intended + for introspection and testing only; not part of the stable API.""" + + def __getstate__(self) -> tuple: ... + def __setstate__(self, arg0: tuple) -> None: ... + +class _Doc: + def __getstate__(self) -> bytes: ... + def __init__(self) -> None: ... + def __setstate__(self, arg0: bytes) -> None: ... + def field_names(self) -> list[str]: ... + def get_any(self, arg0: str, arg1: typing.DataType) -> typing.Any: ... + def has_field(self, arg0: str) -> bool: ... + def pk(self) -> str: ... + def score(self) -> float: ... + def set_any(self, arg0: str, arg1: typing.DataType, arg2: typing.Any) -> bool: ... + def set_pk(self, arg0: str) -> None: ... + def set_score(self, arg0: typing.SupportsFloat) -> None: ... + +class _DocOp: + """ + Members: + + INSERT + + UPDATE + + DELETE + + UPSERT + """ + + DELETE: typing.ClassVar[_DocOp] # value = <_DocOp.DELETE: 3> + INSERT: typing.ClassVar[_DocOp] # value = <_DocOp.INSERT: 0> + UPDATE: typing.ClassVar[_DocOp] # value = <_DocOp.UPDATE: 2> + UPSERT: typing.ClassVar[_DocOp] # value = <_DocOp.UPSERT: 1> + __members__: typing.ClassVar[ + dict[str, _DocOp] + ] # value = {'INSERT': <_DocOp.INSERT: 0>, 'UPDATE': <_DocOp.UPDATE: 2>, 'DELETE': <_DocOp.DELETE: 3>, 'UPSERT': <_DocOp.UPSERT: 1>} + + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + def __init__(self, value: typing.SupportsInt) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: typing.SupportsInt) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... diff --git a/python/zvec/common/__init__.py b/python/zvec/common/__init__.py new file mode 100644 index 0000000..1287488 --- /dev/null +++ b/python/zvec/common/__init__.py @@ -0,0 +1,18 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from .constants import DenseVectorType, SparseVectorType, VectorType + +__all__ = ["DenseVectorType", "SparseVectorType", "VectorType"] diff --git a/python/zvec/common/constants.py b/python/zvec/common/constants.py new file mode 100644 index 0000000..6a1654d --- /dev/null +++ b/python/zvec/common/constants.py @@ -0,0 +1,33 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import Optional, TypeVar, Union + +import numpy as np + +# VectorType: DenseVectorType | SparseVectorType +DenseVectorType = Union[list[float], list[int], np.ndarray] +SparseVectorType = dict[int, float] +VectorType = Optional[Union[DenseVectorType, SparseVectorType]] + +# Embeddable: Text | Image | Audio +TEXT = str +IMAGE = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array +AUDIO = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array + +Embeddable = Optional[Union[TEXT, IMAGE, AUDIO]] + +# Multimodal Embeddable +MD = TypeVar("MD", bound=Embeddable, contravariant=True) diff --git a/python/zvec/executor/__init__.py b/python/zvec/executor/__init__.py new file mode 100644 index 0000000..96582e3 --- /dev/null +++ b/python/zvec/executor/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from .query_executor import ( + QueryContext, + QueryExecutor, +) + +__all__ = [ + "QueryContext", + "QueryExecutor", +] diff --git a/python/zvec/executor/query_executor.py b/python/zvec/executor/query_executor.py new file mode 100644 index 0000000..3faf160 --- /dev/null +++ b/python/zvec/executor/query_executor.py @@ -0,0 +1,266 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import Optional, Union + +import numpy as np + +from zvec._zvec import _Collection, _MultiQuery +from zvec._zvec.param import _Fts, _SearchQuery, _SubQuery + +from ..extension import CallbackReRanker, ReRanker, RrfReRanker, WeightedReRanker +from ..model.convert import convert_to_py_doc +from ..model.doc import DocList +from ..model.param.query import Query +from ..model.schema import CollectionSchema +from ..typing import DataType + +__all__ = [ + "QueryContext", + "QueryExecutor", +] + +DTYPE_MAP = { + DataType.VECTOR_FP16.value: np.float16, + DataType.VECTOR_FP32.value: np.float32, + DataType.VECTOR_FP64.value: np.float64, + DataType.VECTOR_INT8.value: np.int8, +} + + +def convert_to_numpy(vec: Union[list, np.ndarray], dtype: np.dtype) -> np.ndarray: + if isinstance(vec, np.ndarray): + if vec.dtype == dtype and vec.ndim == 1: + return vec + return np.asarray(vec, dtype=dtype).flatten() + + try: + arr = np.asarray(vec, dtype=dtype) + if arr.ndim != 1: + arr = arr.flatten() + return arr + except (ValueError, TypeError) as e: + raise TypeError( + f"Cannot convert input to 1D numpy array with dtype={dtype}: {type(vec)}" + ) from e + + +class QueryContext: + def __init__( + self, + topk: int, + filter: Optional[str] = None, + include_vector: bool = False, + queries: Optional[list[Query]] = None, + output_fields: Optional[list[str]] = None, + reranker: Optional[ReRanker] = None, + ): + # query param + self._filter = filter + self._queries = queries or [] + self._topk = topk + self._include_vector = include_vector + self._output_fields = output_fields + + # reranker + self._reranker = reranker + + @property + def topk(self): + return self._topk + + @property + def queries(self): + return self._queries + + @property + def filter(self): + return self._filter + + @property + def reranker(self): + return self._reranker + + @property + def output_fields(self): + return self._output_fields + + @property + def include_vector(self): + return self._include_vector + + +class QueryExecutor: + """Unified query executor that routes based on query count and reranker type.""" + + def __init__(self, schema: CollectionSchema): + self._schema = schema + + def _build_queries( + self, ctx: QueryContext, collection: _Collection + ) -> list[_SearchQuery]: + """Build query vector list (no validation, conversion only).""" + if not ctx.queries: + return [self._build_base_search_query(ctx)] + return [ + self._build_search_query(ctx, query, collection) for query in ctx.queries + ] + + def execute(self, ctx: QueryContext, collection: _Collection) -> DocList: + """Execute a query, routing by query count. + + A single (or vector-less) query is sent to C++ as a ``_SearchQuery``; + multiple queries are assembled into a ``_MultiQuery``. + """ + queries = self._build_queries(ctx, collection) + if not queries: + raise ValueError("No query to execute") + + if len(queries) == 1: + return self._execute_single_query(queries[0], collection) + return self._execute_multi_query(ctx, queries, collection) + + def _execute_single_query( + self, query: _SearchQuery, collection: _Collection + ) -> DocList: + """Single/vector-less query: send a ``_SearchQuery`` to C++.""" + docs = collection.Query(query) + return [convert_to_py_doc(doc, self._schema) for doc in docs] + + def _execute_multi_query( + self, ctx: QueryContext, queries: list[_SearchQuery], collection: _Collection + ) -> DocList: + """Multiple queries: send a ``_MultiQuery`` to C++. + + A Python-only reranker (e.g. a model/API-based one) cannot run inside + the C++ MultiQuery, so each route is executed individually and merged by + the reranker in Python. The built-in RRF/Weighted/Callback rerankers use + the C++ variant-based fast path. + """ + reranker = ctx.reranker + if reranker is None: + raise ValueError( + "A reranker is required to merge results from multiple queries; " + "specify the 'reranker' argument." + ) + if not isinstance(reranker, (RrfReRanker, WeightedReRanker, CallbackReRanker)): + docs_list = self._execute_python_pipeline(queries, collection) + return self._merge_and_rerank(ctx, docs_list) + + multi_query = self._build_multi_query(ctx, queries) + docs = collection.Query(multi_query) + return [convert_to_py_doc(doc, self._schema) for doc in docs] + + def _build_multi_query( + self, ctx: QueryContext, queries: list[_SearchQuery] + ) -> _MultiQuery: + """Assemble a C++ ``_MultiQuery`` from per-route ``_SearchQuery`` objects.""" + multi_query = _MultiQuery() + multi_query.queries = [_SubQuery.from_search_query(query) for query in queries] + # num_candidates controls per-sub-query candidate count for reranking pool. + # It must NOT be limited to the final output topk; use at least the C++ + # SubQuery default of 10 to ensure sufficient candidates for reranking. + _DEFAULT_NUM_CANDIDATES = 10 + for sub in multi_query.queries: + sub.num_candidates = max(ctx.topk, _DEFAULT_NUM_CANDIDATES) + multi_query.topk = ctx.topk + if ctx.filter: + multi_query.filter = ctx.filter + multi_query.include_vector = ctx.include_vector + if ctx.output_fields is not None: + multi_query.output_fields = ctx.output_fields + # Set rerank strategy via the C++ variant-based API. + reranker = ctx.reranker + if isinstance(reranker, RrfReRanker): + multi_query.set_rerank_rrf(reranker.rank_constant) + elif isinstance(reranker, WeightedReRanker): + multi_query.set_rerank_weighted(reranker.weights) + elif isinstance(reranker, CallbackReRanker): + multi_query.set_rerank_callback(reranker._callback) + return multi_query + + def _execute_python_pipeline( + self, vectors: list[_SearchQuery], collection: _Collection + ) -> list[DocList]: + """Execute queries serially for the Python-only reranker path.""" + return [self._execute_single_query(query, collection) for query in vectors] + + def _merge_and_rerank(self, ctx: QueryContext, docs_list: list[DocList]) -> DocList: + """Merge and rerank results from the Python pipeline path.""" + if not docs_list: + raise ValueError("Query results is empty") + if len(docs_list) == 1 and not ctx.reranker: + return docs_list[0] + return ctx.reranker.rerank(docs_list, ctx.topk) + + def _build_base_search_query(self, ctx: QueryContext) -> _SearchQuery: + search_query = _SearchQuery() + search_query.topk = ctx.topk + search_query.include_vector = ctx.include_vector + if ctx.filter: + search_query.filter = ctx.filter + if ctx.output_fields is not None: + search_query.output_fields = ctx.output_fields + return search_query + + def _apply_fts(self, query: Query, search_query: _SearchQuery) -> None: + """Set FTS query on search_query if the query has FTS parameters.""" + if query.has_fts(): + fts = _Fts() + fts.query_string = query.fts.query_string or "" + fts.match_string = query.fts.match_string or "" + search_query.fts = fts + + def _build_search_query( + self, ctx: QueryContext, query: Query, collection: _Collection + ) -> _SearchQuery: + query._validate() + search_query = self._build_base_search_query(ctx) + search_query.field_name = query.field_name + if query.param: + search_query.query_params = query.param + + # set FTS query if provided + self._apply_fts(query, search_query) + + vector_schema = None + if query.has_vector() or query.has_id(): + vector_schema = ( + self._schema.vector(query.field_name) + if query + else self._schema.vectors[0] + ) + + if vector_schema is None: + raise ValueError("No vector field found") + + # set vector + if query.has_vector(): + vec_data = query.vector + elif query.has_id(): + fetched = collection.Fetch([query.id]) + doc = next(iter(fetched.values()), None) + if not doc: + raise ValueError(f"Document with id '{query.id}' not found") + vec_data = doc.get_any(vector_schema.name, vector_schema.data_type) + else: + return search_query + + target_dtype = DTYPE_MAP.get(vector_schema.data_type.value) + search_query.set_vector( + vector_schema._get_object(), + convert_to_numpy(vec_data, target_dtype) if target_dtype else vec_data, + ) + return search_query diff --git a/python/zvec/extension/__init__.py b/python/zvec/extension/__init__.py new file mode 100644 index 0000000..a1f3a8c --- /dev/null +++ b/python/zvec/extension/__init__.py @@ -0,0 +1,58 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from .bm25_embedding_function import BM25EmbeddingFunction +from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction +from .http_embedding_function import HTTPDenseEmbedding +from .jina_embedding_function import JinaDenseEmbedding +from .jina_function import JinaFunctionBase +from .multi_vector_reranker import CallbackReRanker, RrfReRanker, WeightedReRanker +from .openai_embedding_function import OpenAIDenseEmbedding +from .openai_function import OpenAIFunctionBase +from .qwen_embedding_function import QwenDenseEmbedding, QwenSparseEmbedding +from .qwen_function import QwenFunctionBase +from .qwen_rerank_function import QwenReRanker +from .rerank_function import RerankFunction +from .rerank_function import RerankFunction as ReRanker +from .sentence_transformer_embedding_function import ( + DefaultLocalDenseEmbedding, + DefaultLocalSparseEmbedding, +) +from .sentence_transformer_function import SentenceTransformerFunctionBase +from .sentence_transformer_rerank_function import DefaultLocalReRanker + +__all__ = [ + "BM25EmbeddingFunction", + "CallbackReRanker", + "DefaultLocalDenseEmbedding", + "DefaultLocalReRanker", + "DefaultLocalSparseEmbedding", + "DenseEmbeddingFunction", + "HTTPDenseEmbedding", + "JinaDenseEmbedding", + "JinaFunctionBase", + "OpenAIDenseEmbedding", + "OpenAIFunctionBase", + "QwenDenseEmbedding", + "QwenFunctionBase", + "QwenReRanker", + "QwenSparseEmbedding", + "ReRanker", + "RerankFunction", + "RrfReRanker", + "SentenceTransformerFunctionBase", + "SparseEmbeddingFunction", + "WeightedReRanker", +] diff --git a/python/zvec/extension/bm25_embedding_function.py b/python/zvec/extension/bm25_embedding_function.py new file mode 100644 index 0000000..51ab5ac --- /dev/null +++ b/python/zvec/extension/bm25_embedding_function.py @@ -0,0 +1,375 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from functools import lru_cache +from typing import Literal, Optional + +from ..common.constants import TEXT, SparseVectorType +from ..tool import require_module +from .embedding_function import SparseEmbeddingFunction + + +class BM25EmbeddingFunction(SparseEmbeddingFunction[TEXT]): + """BM25-based sparse embedding function using DashText SDK. + + This class provides text-to-sparse-vector embedding capabilities using + the DashText library with BM25 algorithm. BM25 (Best Matching 25) is a + probabilistic retrieval function used for lexical search and document + ranking based on term frequency and inverse document frequency. + + BM25 generates sparse vectors where each dimension corresponds to a term in + the vocabulary, and the value represents the BM25 score for that term. It's + particularly effective for: + + - Lexical search and keyword matching + - Document ranking and information retrieval + - Combining with dense embeddings for hybrid search + - Traditional IR tasks where exact term matching is important + + This implementation uses DashText's SparseVectorEncoder, which provides + efficient BM25 computation for Chinese and English text using either a + built-in encoder or custom corpus training. + + Args: + corpus (Optional[list[str]], optional): List of documents to train the + BM25 encoder. If provided, creates a custom encoder trained on this + corpus for better domain-specific accuracy. If ``None``, uses the + built-in encoder. Defaults to ``None``. + encoding_type (Literal["query", "document"], optional): Encoding mode + for text processing. Use ``"query"`` for search queries (default) and + ``"document"`` for document indexing. This distinction optimizes the + BM25 scoring for asymmetric retrieval tasks. Defaults to ``"query"``. + language (Literal["zh", "en"], optional): Language for built-in encoder. + Only used when corpus is None. ``"zh"`` for Chinese (trained on Chinese + Wikipedia), ``"en"`` for English. Defaults to ``"zh"``. + b (float, optional): Document length normalization parameter for BM25. + Range [0, 1]. 0 means no normalization, 1 means full normalization. + Only used with custom corpus. Defaults to ``0.75``. + k1 (float, optional): Term frequency saturation parameter for BM25. + Higher values give more weight to term frequency. Only used with + custom corpus. Defaults to ``1.2``. + **kwargs: Additional parameters for DashText encoder customization. + + Attributes: + corpus_size (int): Number of documents in the training corpus (0 if using built-in encoder). + encoding_type (str): The encoding type being used ("query" or "document"). + language (str): The language of the built-in encoder ("zh" or "en"). + + Raises: + ValueError: If corpus is provided but empty or contains non-string elements. + TypeError: If input to ``embed()`` is not a string. + RuntimeError: If DashText encoder initialization or training fails. + + Note: + - Requires Python 3.10, 3.11, or 3.12 + - Requires the ``dashtext`` package: ``pip install dashtext`` + - Two encoder options available: + + 1. **Built-in encoder** (no corpus needed): Pre-trained models for + Chinese (zh) and English (en), good generalization, works out-of-the-box + 2. **Custom encoder** (corpus required): Better accuracy for domain-specific + terminology, requires training on your full corpus with BM25 parameters + + - Encoding types: + + * ``encoding_type="query"``: Optimized for search queries (shorter text) + * ``encoding_type="document"``: Optimized for document indexing (longer text) + + - BM25 parameters (b, k1) only apply to custom encoder training + - Output is sorted by indices (vocabulary term IDs) for consistency + - Results are cached (LRU cache, maxsize=10) to reduce computation + - No API key or network connectivity required (local computation) + + Examples: + >>> # Option 1: Using built-in encoder for Chinese (no corpus needed) + >>> from zvec.extension import BM25EmbeddingFunction + >>> + >>> # For query encoding (Chinese) + >>> bm25_query_zh = BM25EmbeddingFunction(language="zh", encoding_type="query") + >>> query_vec = bm25_query_zh.embed("什么是机器学习") + >>> isinstance(query_vec, dict) + True + >>> # query_vec: {1169440797: 0.29, 2045788977: 0.70, ...} + + >>> # For document encoding (Chinese) + >>> bm25_doc_zh = BM25EmbeddingFunction(language="zh", encoding_type="document") + >>> doc_vec = bm25_doc_zh.embed("机器学习是人工智能的一个重要分支...") + >>> isinstance(doc_vec, dict) + True + + >>> # Using built-in encoder for English + >>> bm25_query_en = BM25EmbeddingFunction(language="en", encoding_type="query") + >>> query_vec_en = bm25_query_en.embed("what is vector search service") + >>> isinstance(query_vec_en, dict) + True + + >>> # Option 2: Using custom corpus for domain-specific accuracy + >>> corpus = [ + ... "机器学习是人工智能的一个重要分支", + ... "深度学习使用多层神经网络进行特征提取", + ... "自然语言处理技术用于理解和生成人类语言" + ... ] + >>> bm25_custom = BM25EmbeddingFunction( + ... corpus=corpus, + ... encoding_type="query", + ... b=0.75, + ... k1=1.2 + ... ) + >>> custom_vec = bm25_custom.embed("机器学习算法") + >>> isinstance(custom_vec, dict) + True + + >>> # Hybrid search: combining with dense embeddings + >>> from zvec.extension import DefaultLocalDenseEmbedding + >>> dense_emb = DefaultLocalDenseEmbedding() + >>> bm25_emb = BM25EmbeddingFunction(language="zh", encoding_type="query") + >>> + >>> query = "machine learning algorithms" + >>> dense_vec = dense_emb.embed(query) # Semantic similarity + >>> sparse_vec = bm25_emb.embed(query) # Lexical matching + >>> # Combine scores for hybrid retrieval + + >>> # Callable interface + >>> sparse_vec = bm25_query_zh("information retrieval") + >>> isinstance(sparse_vec, dict) + True + + >>> # Error handling + >>> try: + ... bm25_query_zh.embed("") # Empty query + ... except ValueError as e: + ... print(f"Error: {e}") + Error: Input text cannot be empty or whitespace only + + See Also: + - ``SparseEmbeddingFunction``: Base class for sparse embeddings + - ``DefaultLocalSparseEmbedding``: SPLADE-based sparse embedding + - ``QwenSparseEmbedding``: API-based sparse embedding using Qwen + - ``DefaultLocalDenseEmbedding``: Dense embedding for semantic search + + References: + - DashText Documentation: https://help.aliyun.com/zh/document_detail/2546039.html + - DashText PyPI: https://pypi.org/project/dashtext/ + - BM25 Algorithm: Robertson & Zaragoza (2009) + """ + + def __init__( + self, + corpus: Optional[list[str]] = None, + encoding_type: Literal["query", "document"] = "query", + language: Literal["zh", "en"] = "zh", + b: float = 0.75, + k1: float = 1.2, + **kwargs, + ): + """Initialize the BM25 embedding function. + + Args: + corpus (Optional[list[str]]): Optional corpus for training custom encoder. + If None, uses built-in encoder. Defaults to None. + encoding_type (Literal["query", "document"]): Text encoding mode. + Use "query" for search queries, "document" for indexing. + Defaults to "query". + language (Literal["zh", "en"]): Language for built-in encoder. + "zh" for Chinese, "en" for English. Defaults to "zh". + b (float): Document length normalization for BM25 [0, 1]. + Only used with custom corpus. Defaults to 0.75. + k1 (float): Term frequency saturation for BM25. + Only used with custom corpus. Defaults to 1.2. + **kwargs: Additional DashText encoder parameters. + + Raises: + ValueError: If corpus is provided but empty or invalid. + ImportError: If dashtext package is not installed. + RuntimeError: If encoder initialization or training fails. + """ + # Validate corpus if provided + if corpus is not None: + if not corpus or not isinstance(corpus, list): + raise ValueError("Corpus must be a non-empty list of strings") + + if not all(isinstance(doc, str) for doc in corpus): + raise ValueError("All corpus documents must be strings") + + # Import dashtext + self._dashtext = require_module("dashtext") + + self._corpus = corpus + self._encoding_type = encoding_type + self._language = language + self._b = b + self._k1 = k1 + self._extra_params = kwargs + + # Initialize the BM25 encoder + self._build_encoder() + + def _build_encoder(self): + """Build the BM25 sparse vector encoder. + + Creates either a built-in encoder (pre-trained) or a custom encoder + trained on the provided corpus. + + Raises: + RuntimeError: If encoder initialization or training fails. + ImportError: If dashtext package is not installed. + """ + try: + if self._corpus is None: + # Use built-in encoder (pre-trained on Wikipedia) + # language: 'zh' for Chinese, 'en' for English + self._encoder = self._dashtext.SparseVectorEncoder.default( + name=self._language + ) + else: + # Create custom encoder with BM25 parameters + self._encoder = self._dashtext.SparseVectorEncoder( + b=self._b, k1=self._k1, **self._extra_params + ) + + # Train encoder with the corpus + self._encoder.train(self._corpus) + + except ImportError as e: + raise ImportError( + "dashtext package is required for BM25EmbeddingFunction. " + "Install it with: pip install dashtext" + ) from e + except Exception as e: + if isinstance(e, (ValueError, RuntimeError)): + raise + raise RuntimeError(f"Failed to build BM25 encoder: {e!s}") from e + + @property + def corpus_size(self) -> int: + """int: Number of documents in the training corpus (0 if using built-in encoder).""" + return len(self._corpus) if self._corpus is not None else 0 + + @property + def encoding_type(self) -> str: + """str: The encoding type being used ("query" or "document").""" + return self._encoding_type + + @property + def language(self) -> str: + """str: The language of the built-in encoder ("zh" or "en").""" + return self._language + + @property + def extra_params(self) -> dict: + """dict: Extra parameters for DashText encoder customization.""" + return self._extra_params + + def __call__(self, input: TEXT) -> SparseVectorType: + """Make the embedding function callable. + + Args: + input (TEXT): Input text to embed. + + Returns: + SparseVectorType: Sparse vector as dictionary. + """ + return self.embed(input) + + @lru_cache(maxsize=10) + def embed(self, input: TEXT) -> SparseVectorType: + """Generate BM25 sparse embedding for the input text. + + This method computes BM25 scores for the input text using DashText's + SparseVectorEncoder. The encoding behavior depends on the encoding_type: + + - ``encoding_type="query"``: Uses ``encode_queries()`` for search queries + - ``encoding_type="document"``: Uses ``encode_documents()`` for documents + + The result is a sparse vector where keys are term indices in the + vocabulary and values are BM25 scores. + + Args: + input (TEXT): Input text string to embed. Must be non-empty after + stripping whitespace. + + Returns: + SparseVectorType: A dictionary mapping vocabulary term index to BM25 score. + Only non-zero scores are included. The dictionary is sorted by indices + (keys) in ascending order for consistent output. + Example: ``{1169440797: 0.29, 2045788977: 0.70, ...}`` + + Raises: + TypeError: If ``input`` is not a string. + ValueError: If input is empty or whitespace-only. + RuntimeError: If BM25 encoding fails. + + Examples: + >>> bm25 = BM25EmbeddingFunction(language="zh", encoding_type="query") + >>> sparse_vec = bm25.embed("query text") + >>> isinstance(sparse_vec, dict) + True + >>> all(isinstance(k, int) and isinstance(v, float) for k, v in sparse_vec.items()) + True + + >>> # Verify sorted output + >>> keys = list(sparse_vec.keys()) + >>> keys == sorted(keys) + True + + >>> # Error: empty input + >>> bm25.embed(" ") + ValueError: Input text cannot be empty or whitespace only + + >>> # Error: non-string input + >>> bm25.embed(123) + TypeError: Expected 'input' to be str, got int + + Note: + - BM25 scores are relative to the vocabulary statistics + - Output dictionary is always sorted by indices for consistency + - Terms not in the vocabulary will have zero scores (not included) + - This method is cached (maxsize=10) for performance + - DashText automatically handles Chinese/English text segmentation + """ + if not isinstance(input, str): + raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}") + + input = input.strip() + if not input: + raise ValueError("Input text cannot be empty or whitespace only") + + try: + # Encode based on encoding_type + if self._encoding_type == "query": + sparse_vector = self._encoder.encode_queries(input) + else: # encoding_type == "document" + sparse_vector = self._encoder.encode_documents(input) + + # DashText returns dict with int/long keys and float values + # Convert to standard format: {int: float} + sparse_dict: dict[int, float] = {} + for key, value in sparse_vector.items(): + try: + idx = int(key) + val = float(value) + if val > 0: + sparse_dict[idx] = val + except (ValueError, TypeError): + # Skip invalid entries + continue + + # Sort by indices (keys) to ensure consistent ordering + return dict(sorted(sparse_dict.items())) + + except Exception as e: + if isinstance(e, (TypeError, ValueError)): + raise + raise RuntimeError(f"Failed to generate BM25 embedding: {e!s}") from e diff --git a/python/zvec/extension/embedding_function.py b/python/zvec/extension/embedding_function.py new file mode 100644 index 0000000..a421f1e --- /dev/null +++ b/python/zvec/extension/embedding_function.py @@ -0,0 +1,147 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from abc import abstractmethod +from typing import Protocol, runtime_checkable + +from ..common.constants import MD, DenseVectorType, SparseVectorType + + +@runtime_checkable +class DenseEmbeddingFunction(Protocol[MD]): + """Protocol for dense vector embedding functions. + + Dense embedding functions map multimodal input (text, image, or audio) to + fixed-length real-valued vectors. This is a Protocol class that defines + the interface - implementations should provide their own initialization + and properties. + + Type Parameters: + MD: The type of input data (bound to Embeddable: TEXT, IMAGE, or AUDIO). + + Note: + - This is a Protocol class - it only defines the ``embed()`` interface. + - Implementations are free to define their own ``__init__``, properties, + and additional methods as needed. + - The ``embed()`` method is the only required interface. + + Examples: + >>> # Custom text embedding implementation + >>> class MyTextEmbedding: + ... def __init__(self, dimension: int, model_name: str): + ... self.dimension = dimension + ... self.model = load_model(model_name) + ... + ... def embed(self, input: str) -> list[float]: + ... return self.model.encode(input).tolist() + + >>> # Custom image embedding implementation + >>> class MyImageEmbedding: + ... def __init__(self, dimension: int = 512): + ... self.dimension = dimension + ... self.model = load_image_model() + ... + ... def embed(self, input: Union[str, bytes, np.ndarray]) -> list[float]: + ... if isinstance(input, str): + ... image = load_image_from_path(input) + ... else: + ... image = input + ... return self.model.extract_features(image).tolist() + + >>> # Using built-in implementations + >>> from zvec.extension import QwenDenseEmbedding + >>> text_emb = QwenDenseEmbedding(dimension=768, api_key="sk-xxx") + >>> vector = text_emb.embed("Hello world") + """ + + @abstractmethod + def embed(self, input: MD) -> DenseVectorType: + """Generate a dense embedding vector for the input data. + + Args: + input (MD): Multimodal input data to embed. Can be: + - TEXT (str): Text string + - IMAGE (str | bytes | np.ndarray): Image file path, raw bytes, or array + - AUDIO (str | bytes | np.ndarray): Audio file path, raw bytes, or array + + Returns: + DenseVectorType: A dense vector representing the embedding. + Can be list[float], list[int], or np.ndarray. + Length should match the implementation's dimension. + """ + ... + + +@runtime_checkable +class SparseEmbeddingFunction(Protocol[MD]): + """Abstract base class for sparse vector embedding functions. + + Sparse embedding functions map multimodal input (text, image, or audio) to + a dictionary of {index: weight}, where only non-zero dimensions are stored. + You can inherit this class to create custom sparse embedding functions. + + Type Parameters: + MD: The type of input data (bound to Embeddable: TEXT, IMAGE, or AUDIO). + + Note: + Subclasses must implement the ``embed()`` method. + + Examples: + >>> # Using built-in text sparse embedding (e.g., BM25, TF-IDF) + >>> sparse_emb = SomeSparseEmbedding() + >>> vector = sparse_emb.embed("Hello world") + >>> # Returns: {0: 0.5, 42: 1.2, 100: 0.8} + + >>> # Custom BM25 sparse embedding function + >>> class MyBM25Embedding(SparseEmbeddingFunction): + ... def __init__(self, vocab_size: int = 10000): + ... self.vocab_size = vocab_size + ... self.tokenizer = MyTokenizer() + ... + ... def embed(self, input: str) -> dict[int, float]: + ... tokens = self.tokenizer.tokenize(input) + ... sparse_vector = {} + ... for token_id, weight in self._calculate_bm25(tokens): + ... if weight > 0: + ... sparse_vector[token_id] = weight + ... return sparse_vector + ... + ... def _calculate_bm25(self, tokens): + ... # BM25 calculation logic + ... pass + + >>> # Custom sparse image feature extractor + >>> class MySparseImageEmbedding(SparseEmbeddingFunction): + ... def embed(self, input: Union[str, bytes, np.ndarray]) -> dict[int, float]: + ... image = self._load_image(input) + ... features = self._extract_sparse_features(image) + ... return {idx: val for idx, val in enumerate(features) if val != 0} + """ + + @abstractmethod + def embed(self, input: MD) -> SparseVectorType: + """Generate a sparse embedding for the input data. + + Args: + input (MD): Multimodal input data to embed. Can be: + - TEXT (str): Text string + - IMAGE (str | bytes | np.ndarray): Image file path, raw bytes, or array + - AUDIO (str | bytes | np.ndarray): Audio file path, raw bytes, or array + + Returns: + SparseVectorType: Mapping from dimension index to non-zero weight. + Only dimensions with non-zero values are included. + """ + ... diff --git a/python/zvec/extension/http_embedding_function.py b/python/zvec/extension/http_embedding_function.py new file mode 100644 index 0000000..3386c65 --- /dev/null +++ b/python/zvec/extension/http_embedding_function.py @@ -0,0 +1,162 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import json +import os +import urllib.request +from functools import lru_cache +from typing import Optional + +from ..common.constants import TEXT, DenseVectorType +from .embedding_function import DenseEmbeddingFunction + + +class HTTPDenseEmbedding(DenseEmbeddingFunction[TEXT]): + """Dense text embedding function using any OpenAI-compatible HTTP endpoint. + + This class calls any server that implements the ``/v1/embeddings`` API + (LM Studio, Ollama, vLLM, LocalAI, etc.) using only the Python standard + library — no extra dependencies are required. + + The embedding dimension is detected automatically from the first server + response. + + Args: + base_url (str, optional): Base URL of the embedding server. + Defaults to ``"http://localhost:1234"`` (LM Studio). + Common values: + + - ``"http://localhost:1234"`` — LM Studio + - ``"http://localhost:11434"`` — Ollama + model (str, optional): Model identifier as expected by the server. + Defaults to ``"text-embedding-nomic-embed-text-v1.5@f16"``. + api_key (Optional[str], optional): Bearer token for authenticated + endpoints. Falls back to the ``OPENAI_API_KEY`` environment + variable. Leave as ``None`` for local servers that do not + require authentication. + timeout (int, optional): HTTP request timeout in seconds. + Defaults to 30. + + Attributes: + dimension (int): Embedding vector dimensionality (auto-detected). + + Raises: + TypeError: If ``embed()`` receives a non-string input. + ValueError: If input is empty/whitespace-only or the server returns + an unexpected response format. + RuntimeError: If the HTTP request fails or the server is unreachable. + + Examples: + >>> from zvec.extension import HTTPDenseEmbedding + >>> + >>> # LM Studio (default) + >>> emb = HTTPDenseEmbedding() + >>> vector = emb.embed("Hello, world!") + >>> len(vector) + 768 + >>> + >>> # Ollama + >>> emb = HTTPDenseEmbedding( + ... base_url="http://localhost:11434", + ... model="nomic-embed-text", + ... ) + >>> vector = emb.embed("Semantic search with local models") + + See Also: + - ``DenseEmbeddingFunction``: Protocol for dense embeddings. + - ``OpenAIDenseEmbedding``: Cloud embedding via the OpenAI API. + """ + + ENDPOINT = "/v1/embeddings" + + def __init__( + self, + base_url: str = "http://localhost:1234", + model: str = "text-embedding-nomic-embed-text-v1.5@f16", + api_key: Optional[str] = None, + timeout: int = 30, + ) -> None: + self._base_url = base_url.rstrip("/") + self._model = model + self._api_key = api_key or os.environ.get("OPENAI_API_KEY", "") + self._timeout = timeout + self._dimension: Optional[int] = None + + @property + def dimension(self) -> int: + """int: Embedding vector dimensionality (auto-detected on first call).""" + if self._dimension is None: + self._dimension = len(self.embed("dimension probe")) + return self._dimension + + def __call__(self, input: TEXT) -> DenseVectorType: + """Make the embedding function callable.""" + return self.embed(input) + + @lru_cache(maxsize=256) + def embed(self, input: TEXT) -> DenseVectorType: + """Generate a dense embedding vector for the input text. + + Results are cached (LRU, up to 256 entries) so repeated strings + do not trigger extra HTTP requests. + + Args: + input (TEXT): Input text string to embed. Must be non-empty + after stripping whitespace. + + Returns: + DenseVectorType: A list of floats representing the embedding. + + Raises: + TypeError: If *input* is not a string. + ValueError: If *input* is empty/whitespace-only or the server + returns an unexpected response format. + RuntimeError: If the HTTP request fails. + """ + if not isinstance(input, TEXT): + raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}") + + input = input.strip() + if not input: + raise ValueError("Input text cannot be empty or whitespace only") + + url = self._base_url + self.ENDPOINT + payload = json.dumps({"model": self._model, "input": input}).encode() + + headers: dict[str, str] = {"Content-Type": "application/json"} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + + req = urllib.request.Request(url, data=payload, headers=headers, method="POST") + try: + with urllib.request.urlopen(req, timeout=self._timeout) as resp: + body = json.loads(resp.read()) + except urllib.error.HTTPError as exc: + raise RuntimeError( + f"Embedding server returned HTTP {exc.code}: {exc.read().decode()}" + ) from exc + except OSError as exc: + raise RuntimeError( + f"Could not reach embedding server at {url}: {exc}" + ) from exc + + try: + vector: list[float] = body["data"][0]["embedding"] + except (KeyError, IndexError) as exc: + raise ValueError( + f"Unexpected response format from embedding server: {body}" + ) from exc + + return vector diff --git a/python/zvec/extension/jina_embedding_function.py b/python/zvec/extension/jina_embedding_function.py new file mode 100644 index 0000000..2f8b02a --- /dev/null +++ b/python/zvec/extension/jina_embedding_function.py @@ -0,0 +1,240 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from functools import lru_cache +from typing import Optional + +from ..common.constants import TEXT, DenseVectorType +from .embedding_function import DenseEmbeddingFunction +from .jina_function import JinaFunctionBase + + +class JinaDenseEmbedding(JinaFunctionBase, DenseEmbeddingFunction[TEXT]): + """Dense text embedding function using Jina AI API. + + This class provides text-to-vector embedding capabilities using Jina AI's + embedding models. It inherits from ``DenseEmbeddingFunction`` and implements + dense text embedding via the Jina Embeddings API (OpenAI-compatible). + + Jina Embeddings v5 models support task-specific embedding through the + ``task`` parameter, which optimizes the embedding for different use cases + such as retrieval, text matching, or classification. They also support + Matryoshka Representation Learning, allowing flexible output dimensions. + + Args: + model (str, optional): Jina embedding model identifier. + Defaults to ``"jina-embeddings-v5-text-nano"``. Available models: + - ``"jina-embeddings-v5-text-nano"``: 768 dims, 239M params, 8K context + - ``"jina-embeddings-v5-text-small"``: 1024 dims, 677M params, 32K context + dimension (Optional[int], optional): Desired output embedding dimension. + If ``None``, uses model's default dimension. Supports Matryoshka + dimensions: 32, 64, 128, 256, 512, 768 (nano) / 1024 (small). + Defaults to ``None``. + api_key (Optional[str], optional): Jina API authentication key. + If ``None``, reads from ``JINA_API_KEY`` environment variable. + Obtain your key from: https://jina.ai/api-dashboard + task (Optional[str], optional): Task type to optimize embeddings for. + Defaults to ``None``. Valid values: + - ``"retrieval.query"``: For search queries + - ``"retrieval.passage"``: For documents/passages to be searched + - ``"text-matching"``: For symmetric text similarity + - ``"classification"``: For text classification + - ``"separation"``: For clustering/separation tasks + + Attributes: + dimension (int): The embedding vector dimension. + data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation. + model (str): The Jina model name being used. + task (Optional[str]): The task type for embedding optimization. + + Raises: + ValueError: If API key is not provided and not found in environment, + if task is not a valid task type, or if API returns an error response. + TypeError: If input to ``embed()`` is not a string. + RuntimeError: If network error or Jina service error occurs. + + Note: + - Requires Python 3.10, 3.11, or 3.12 + - Requires the ``openai`` package: ``pip install openai`` + - Jina API is OpenAI-compatible, so it uses the ``openai`` Python client + - Embedding results are cached (LRU cache, maxsize=10) to reduce API calls + - For retrieval tasks, use ``"retrieval.query"`` for queries and + ``"retrieval.passage"`` for documents + - API usage requires a Jina API key from https://jina.ai/api-dashboard + + Examples: + >>> # Basic usage with default model + >>> from zvec.extension import JinaDenseEmbedding + >>> import os + >>> os.environ["JINA_API_KEY"] = "jina_..." + >>> + >>> emb_func = JinaDenseEmbedding() + >>> vector = emb_func.embed("Hello, world!") + >>> len(vector) + 768 + + >>> # Retrieval use case: embed queries and documents differently + >>> query_emb = JinaDenseEmbedding(task="retrieval.query") + >>> doc_emb = JinaDenseEmbedding(task="retrieval.passage") + >>> + >>> query_vector = query_emb.embed("What is machine learning?") + >>> doc_vector = doc_emb.embed("Machine learning is a subset of AI...") + + >>> # Using larger model with custom dimension (Matryoshka) + >>> emb_func = JinaDenseEmbedding( + ... model="jina-embeddings-v5-text-small", + ... dimension=256, + ... api_key="jina_...", + ... task="text-matching", + ... ) + >>> vector = emb_func.embed("Semantic similarity comparison") + >>> len(vector) + 256 + + >>> # Using with zvec collection + >>> import zvec + >>> emb_func = JinaDenseEmbedding(task="retrieval.passage") + >>> schema = zvec.CollectionSchema( + ... name="docs", + ... vectors=zvec.VectorSchema( + ... "embedding", zvec.DataType.VECTOR_FP32, emb_func.dimension + ... ), + ... ) + >>> collection = zvec.create_and_open(path="./my_docs", schema=schema) + + See Also: + - ``DenseEmbeddingFunction``: Base class for dense embeddings + - ``OpenAIDenseEmbedding``: Alternative using OpenAI API + - ``QwenDenseEmbedding``: Alternative using Qwen/DashScope API + - ``DefaultLocalDenseEmbedding``: Local model without API calls + """ + + def __init__( + self, + model: str = "jina-embeddings-v5-text-nano", + dimension: Optional[int] = None, + api_key: Optional[str] = None, + task: Optional[str] = None, + **kwargs, + ): + """Initialize the Jina dense embedding function. + + Args: + model (str): Jina model name. Defaults to "jina-embeddings-v5-text-nano". + dimension (Optional[int]): Target embedding dimension or None for default. + api_key (Optional[str]): API key or None to use environment variable. + task (Optional[str]): Task type for embedding optimization or None. + **kwargs: Additional parameters for API calls. + + Raises: + ValueError: If API key is not provided and not in environment, + or if task is not a valid task type. + """ + # Initialize base class for API connection + JinaFunctionBase.__init__(self, model=model, api_key=api_key, task=task) + + # Store dimension configuration + self._custom_dimension = dimension + + # Determine actual dimension + if dimension is None: + self._dimension = self._MODEL_DIMENSIONS.get(model, 768) + else: + self._dimension = dimension + + # Store extra attributes + self._extra_params = kwargs + + @property + def dimension(self) -> int: + """int: The expected dimensionality of the embedding vector.""" + return self._dimension + + @property + def extra_params(self) -> dict: + """dict: Extra parameters for model-specific customization.""" + return self._extra_params + + def __call__(self, input: TEXT) -> DenseVectorType: + """Make the embedding function callable.""" + return self.embed(input) + + @lru_cache(maxsize=10) + def embed(self, input: TEXT) -> DenseVectorType: + """Generate dense embedding vector for the input text. + + This method calls the Jina Embeddings API to convert input text + into a dense vector representation. Results are cached to improve + performance for repeated inputs. + + Args: + input (TEXT): Input text string to embed. Must be non-empty after + stripping whitespace. Maximum length depends on model: + 8192 tokens for v5-nano, 32768 tokens for v5-small. + + Returns: + DenseVectorType: A list of floats representing the embedding vector. + Length equals ``self.dimension``. Example: + ``[0.123, -0.456, 0.789, ...]`` + + Raises: + TypeError: If ``input`` is not a string. + ValueError: If input is empty/whitespace-only, or if the API returns + an error or malformed response. + RuntimeError: If network connectivity issues or Jina service + errors occur. + + Examples: + >>> emb = JinaDenseEmbedding(task="retrieval.query") + >>> vector = emb.embed("What is deep learning?") + >>> len(vector) + 768 + >>> isinstance(vector[0], float) + True + + >>> # Error: empty input + >>> emb.embed(" ") + ValueError: Input text cannot be empty or whitespace only + + >>> # Error: non-string input + >>> emb.embed(123) + TypeError: Expected 'input' to be str, got int + + Note: + - This method is cached (maxsize=10). Identical inputs return cached results. + - The cache is based on exact string match (case-sensitive). + - Task type affects embedding optimization but not caching behavior. + """ + if not isinstance(input, TEXT): + raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}") + + input = input.strip() + if not input: + raise ValueError("Input text cannot be empty or whitespace only") + + # Call API + embedding_vector = self._call_text_embedding_api( + input=input, + dimension=self._custom_dimension, + ) + + # Verify dimension + if len(embedding_vector) != self.dimension: + raise ValueError( + f"Dimension mismatch: expected {self.dimension}, " + f"got {len(embedding_vector)}" + ) + + return embedding_vector diff --git a/python/zvec/extension/jina_function.py b/python/zvec/extension/jina_function.py new file mode 100644 index 0000000..f20b679 --- /dev/null +++ b/python/zvec/extension/jina_function.py @@ -0,0 +1,182 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import os +from typing import ClassVar, Optional + +from ..common.constants import TEXT +from ..tool import require_module + + +class JinaFunctionBase: + """Base class for Jina AI functions. + + This base class provides common functionality for calling Jina AI APIs + and handling responses. It supports embeddings (dense) operations via + the OpenAI-compatible Jina Embeddings API. + + This class is not meant to be used directly. Use concrete implementations: + - ``JinaDenseEmbedding`` for dense embeddings + + Args: + model (str): Jina embedding model identifier. + api_key (Optional[str]): Jina API authentication key. + task (Optional[str]): Task type for the embedding model. + + Note: + - This is an internal base class for code reuse across Jina features + - Subclasses should inherit from appropriate Protocol + - Provides unified API connection and response handling + - Jina API is OpenAI-compatible, using the ``openai`` Python client + """ + + _BASE_URL: ClassVar[str] = "https://api.jina.ai/v1" + + # Model default dimensions + _MODEL_DIMENSIONS: ClassVar[dict[str, int]] = { + "jina-embeddings-v5-text-nano": 768, + "jina-embeddings-v5-text-small": 1024, + } + + # Model max tokens + _MODEL_MAX_TOKENS: ClassVar[dict[str, int]] = { + "jina-embeddings-v5-text-nano": 8192, + "jina-embeddings-v5-text-small": 32768, + } + + # Valid task types + _VALID_TASKS: ClassVar[tuple[str, ...]] = ( + "retrieval.query", + "retrieval.passage", + "text-matching", + "classification", + "separation", + ) + + def __init__( + self, + model: str, + api_key: Optional[str] = None, + task: Optional[str] = None, + ): + """Initialize the base Jina functionality. + + Args: + model (str): Jina model name. + api_key (Optional[str]): API key or None to use environment variable. + task (Optional[str]): Task type for the embedding model. + Valid values: "retrieval.query", "retrieval.passage", + "text-matching", "classification", "separation". + + Raises: + ValueError: If API key is not provided and not in environment, + or if task is not a valid task type. + """ + self._model = model + self._api_key = api_key or os.environ.get("JINA_API_KEY") + self._task = task + + if not self._api_key: + raise ValueError( + "Jina API key is required. Please provide 'api_key' parameter " + "or set the 'JINA_API_KEY' environment variable. " + "Get your key from: https://jina.ai/api-dashboard" + ) + + if task is not None and task not in self._VALID_TASKS: + raise ValueError( + f"Invalid task '{task}'. Valid tasks: {', '.join(self._VALID_TASKS)}" + ) + + @property + def model(self) -> str: + """str: The Jina model name currently in use.""" + return self._model + + @property + def task(self) -> Optional[str]: + """Optional[str]: The task type for the embedding model.""" + return self._task + + def _get_client(self): + """Get OpenAI-compatible client instance configured for Jina API. + + Returns: + OpenAI: Configured OpenAI client pointing to Jina API. + + Raises: + ImportError: If openai package is not installed. + """ + openai = require_module("openai") + return openai.OpenAI(api_key=self._api_key, base_url=self._BASE_URL) + + def _call_text_embedding_api( + self, + input: TEXT, + dimension: Optional[int] = None, + ) -> list: + """Call Jina Embeddings API. + + Args: + input (TEXT): Input text to embed. + dimension (Optional[int]): Target dimension for Matryoshka embeddings. + + Returns: + list: Embedding vector as list of floats. + + Raises: + RuntimeError: If API call fails. + ValueError: If API returns error response. + """ + try: + client = self._get_client() + + # Prepare embedding parameters + params = {"model": self.model, "input": input} + + # Add dimension parameter for Matryoshka support + if dimension is not None: + params["dimensions"] = dimension + + # Add task parameter via extra_body + if self._task is not None: + params["extra_body"] = {"task": self._task} + + # Call Jina API (OpenAI-compatible) + response = client.embeddings.create(**params) + + except Exception as e: + # Check if it's an OpenAI API error + openai = require_module("openai") + if isinstance(e, (openai.APIError, openai.APIConnectionError)): + raise RuntimeError(f"Failed to call Jina API: {e!s}") from e + raise RuntimeError(f"Unexpected error during API call: {e!s}") from e + + # Extract embedding from response + try: + if not response.data: + raise ValueError("Invalid API response: no embedding data returned") + + embedding_vector = response.data[0].embedding + + if not isinstance(embedding_vector, list): + raise ValueError( + "Invalid API response: embedding is not a list of numbers" + ) + + return embedding_vector + + except (AttributeError, IndexError, TypeError) as e: + raise ValueError(f"Failed to parse API response: {e!s}") from e diff --git a/python/zvec/extension/multi_vector_reranker.py b/python/zvec/extension/multi_vector_reranker.py new file mode 100644 index 0000000..a04c9f8 --- /dev/null +++ b/python/zvec/extension/multi_vector_reranker.py @@ -0,0 +1,197 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from zvec._zvec import ( + _CallbackParams, + _Doc, + _reranker_rerank, + _RrfParams, + _WeightedParams, +) + +from ..model.doc import Doc, DocList +from .rerank_function import RerankFunction + +if TYPE_CHECKING: + from ..model.schema import FieldSchema, VectorSchema + + +def _to_cpp_doc_lists( + query_results: list[list[Doc]], +) -> tuple[list[list], dict[str, Doc]]: + """Convert Python Doc lists to C++ _Doc lists for reranker input.""" + id_to_doc: dict[str, Doc] = {} + cpp_results: list[list] = [] + for query_result in query_results: + cpp_list: list = [] + for doc in query_result: + _doc = _Doc() + _doc.set_pk(doc.id) + _doc.set_score(doc.score if doc.score is not None else 0.0) + cpp_list.append(_doc) + if doc.id not in id_to_doc: + id_to_doc[doc.id] = doc + cpp_results.append(cpp_list) + return cpp_results, id_to_doc + + +def _from_cpp_docs(cpp_docs: list, id_to_doc: dict[str, Doc]) -> DocList: + """Convert C++ rerank result _Doc list back to Python DocList.""" + results: DocList = [] + for _doc in cpp_docs: + doc_id = _doc.pk() + new_score = _doc.score() + original = id_to_doc.get(doc_id) + if original is not None: + results.append(original._replace(score=new_score)) + else: + results.append(Doc(id=doc_id, score=new_score)) + return results + + +class RrfReRanker(RerankFunction): + """Re-ranker using Reciprocal Rank Fusion (RRF) for multi-vector search. + + RRF combines results from multiple vector queries without requiring + relevance scores. The RRF score for a document at rank r is: + score = 1 / (k + r + 1) + where k is the rank constant. + + Args: + rank_constant: RRF smoothing constant (default: 60). + Higher values reduce the influence of rank position. + + Example: + >>> reranker = RrfReRanker(rank_constant=60) + >>> merged = reranker.rerank([results_a, results_b], topn=10) + """ + + def __init__(self, rank_constant: int = 60): + self._rank_constant = rank_constant + + @property + def rank_constant(self) -> int: + """int: RRF rank constant.""" + return self._rank_constant + + def _to_cpp_params(self): + return _RrfParams(self._rank_constant) + + def rerank( + self, + query_results: list[list[Doc]], + topn: int = 10, + *, + fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002 + ) -> DocList: + """Apply RRF to combine multiple query results via C++ reranker.""" + cpp_results, id_to_doc = _to_cpp_doc_lists(query_results) + cpp_docs = _reranker_rerank(self._to_cpp_params(), cpp_results, [], topn) + return _from_cpp_docs(cpp_docs, id_to_doc) + + +class WeightedReRanker(RerankFunction): + """Re-ranker that combines scores using per-sub-query weights. + + Each sub-query's score is normalized by metric type (automatic when used + via collection.multi_query), then multiplied by the corresponding weight. + + Args: + weights: Per-sub-query weights. Length must match the number of + sub-queries. + + Example: + >>> reranker = WeightedReRanker([0.7, 0.3]) + >>> merged = reranker.rerank([results_a, results_b], topn=10, + ... fields=field_schemas) + """ + + def __init__(self, weights: list[float]): + self._weights = list(weights) + + @property + def weights(self) -> list[float]: + """list[float]: Per-sub-query weights.""" + return self._weights + + def _to_cpp_params(self): + return _WeightedParams(self._weights) + + def rerank( + self, + query_results: list[list[Doc]], + topn: int = 10, + *, + fields: list[FieldSchema | VectorSchema] | None = None, + ) -> DocList: + """Combine scores from multiple sub-queries using weighted sum via C++ reranker. + + Args: + query_results: Per-sub-query document lists. + topn: Maximum results to return. + fields: Per-sub-query Python FieldSchema/VectorSchema objects + (required for score normalization by metric type). + + Raises: + ValueError: If fields is None (required for normalization). + """ + if not fields: + raise ValueError( + "WeightedReRanker.rerank() requires 'fields' for score normalization. " + "Pass field schemas via fields= parameter." + ) + cpp_fields = [f._get_object() for f in fields] + cpp_results, id_to_doc = _to_cpp_doc_lists(query_results) + cpp_docs = _reranker_rerank( + self._to_cpp_params(), cpp_results, cpp_fields, topn + ) + return _from_cpp_docs(cpp_docs, id_to_doc) + + +class CallbackReRanker(RerankFunction): + """Re-ranker that delegates to a user-provided callback. + + The callback receives sub-query results, field schemas, and topn. + + Args: + callback: A callable with signature + (results: list[list[Doc]], fields: list, topn: int) -> list[Doc] + + Example: + >>> def my_rerank(results, fields, topn): + ... # custom logic + ... return merged[:topn] + >>> reranker = CallbackReRanker(my_rerank) + >>> merged = reranker.rerank([results_a, results_b], topn=10) + """ + + def __init__(self, callback: Callable): + self._callback = callback + + def _to_cpp_params(self): + return _CallbackParams(self._callback) + + def rerank( + self, + query_results: list[list[Doc]], + topn: int = 10, + *, + fields: list[FieldSchema | VectorSchema] | None = None, + ) -> DocList: + """Invoke the callback to re-rank documents.""" + return self._callback(query_results, fields, topn) diff --git a/python/zvec/extension/openai_embedding_function.py b/python/zvec/extension/openai_embedding_function.py new file mode 100644 index 0000000..03a34ed --- /dev/null +++ b/python/zvec/extension/openai_embedding_function.py @@ -0,0 +1,238 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from functools import lru_cache +from typing import Optional + +from ..common.constants import TEXT, DenseVectorType +from .embedding_function import DenseEmbeddingFunction +from .openai_function import OpenAIFunctionBase + + +class OpenAIDenseEmbedding(OpenAIFunctionBase, DenseEmbeddingFunction[TEXT]): + """Dense text embedding function using OpenAI API. + + This class provides text-to-vector embedding capabilities using OpenAI's + embedding models. It inherits from ``DenseEmbeddingFunction`` and implements + dense text embedding via the OpenAI API. + + The implementation supports various OpenAI embedding models with different + dimensions and includes automatic result caching for improved performance. + + Args: + model (str, optional): OpenAI embedding model identifier. + Defaults to ``"text-embedding-3-small"``. Common options: + - ``"text-embedding-3-small"``: 1536 dims, cost-efficient, good performance + - ``"text-embedding-3-large"``: 3072 dims, highest quality + - ``"text-embedding-ada-002"``: 1536 dims, legacy model + dimension (Optional[int], optional): Desired output embedding dimension. + If ``None``, uses model's default dimension. For text-embedding-3 models, + you can specify custom dimensions (e.g., 256, 512, 1024, 1536). + Defaults to ``None``. + api_key (Optional[str], optional): OpenAI API authentication key. + If ``None``, reads from ``OPENAI_API_KEY`` environment variable. + Obtain your key from: https://platform.openai.com/api-keys + base_url (Optional[str], optional): Custom API base URL for OpenAI-compatible + services. Defaults to ``None`` (uses official OpenAI endpoint). + + Attributes: + dimension (int): The embedding vector dimension. + data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation. + model (str): The OpenAI model name being used. + + Raises: + ValueError: If API key is not provided and not found in environment, + or if API returns an error response. + TypeError: If input to ``embed()`` is not a string. + RuntimeError: If network error or OpenAI service error occurs. + + Note: + - Requires Python 3.10, 3.11, or 3.12 + - Requires the ``openai`` package: ``pip install openai`` + - Embedding results are cached (LRU cache, maxsize=10) to reduce API calls + - Network connectivity to OpenAI API endpoints is required + - API usage incurs costs based on your OpenAI subscription plan + - Rate limits apply based on your OpenAI account tier + + Examples: + >>> # Basic usage with default model + >>> from zvec.extension import OpenAIDenseEmbedding + >>> import os + >>> os.environ["OPENAI_API_KEY"] = "sk-..." + >>> + >>> emb_func = OpenAIDenseEmbedding() + >>> vector = emb_func.embed("Hello, world!") + >>> len(vector) + 1536 + + >>> # Using specific model with custom dimension + >>> emb_func = OpenAIDenseEmbedding( + ... model="text-embedding-3-large", + ... dimension=1024, + ... api_key="sk-..." + ... ) + >>> vector = emb_func.embed("Machine learning is fascinating") + >>> len(vector) + 1024 + + >>> # Using with custom base URL (e.g., Azure OpenAI) + >>> emb_func = OpenAIDenseEmbedding( + ... model="text-embedding-ada-002", + ... api_key="your-azure-key", + ... base_url="https://your-resource.openai.azure.com/" + ... ) + >>> vector = emb_func("Natural language processing") + >>> isinstance(vector, list) + True + + >>> # Batch processing with caching benefit + >>> texts = ["First text", "Second text", "First text"] + >>> vectors = [emb_func.embed(text) for text in texts] + >>> # Third call uses cached result for "First text" + + >>> # Error handling + >>> try: + ... emb_func.embed("") # Empty string + ... except ValueError as e: + ... print(f"Error: {e}") + Error: Input text cannot be empty or whitespace only + + See Also: + - ``DenseEmbeddingFunction``: Base class for dense embeddings + - ``QwenDenseEmbedding``: Alternative using Qwen/DashScope API + - ``DefaultDenseEmbedding``: Local model without API calls + - ``SparseEmbeddingFunction``: Base class for sparse embeddings + """ + + def __init__( + self, + model: str = "text-embedding-3-small", + dimension: Optional[int] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + **kwargs, + ): + """Initialize the OpenAI dense embedding function. + + Args: + model (str): OpenAI model name. Defaults to "text-embedding-3-small". + dimension (Optional[int]): Target embedding dimension or None for default. + api_key (Optional[str]): API key or None to use environment variable. + base_url (Optional[str]): Custom API base URL or None for default. + **kwargs: Additional parameters for API calls. Examples: + - ``encoding_format`` (str): Format of embeddings, "float" or "base64". + - ``user`` (str): User identifier for tracking. + + Raises: + ValueError: If API key is not provided and not in environment. + """ + # Initialize base class for API connection + OpenAIFunctionBase.__init__( + self, model=model, api_key=api_key, base_url=base_url + ) + + # Store dimension configuration + self._custom_dimension = dimension + + # Determine actual dimension + if dimension is None: + # Use model default dimension + self._dimension = self._MODEL_DIMENSIONS.get(model, 1536) + else: + self._dimension = dimension + + # Store dense-specific attributes + self._extra_params = kwargs + + @property + def dimension(self) -> int: + """int: The expected dimensionality of the embedding vector.""" + return self._dimension + + @property + def extra_params(self) -> dict: + """dict: Extra parameters for model-specific customization.""" + return self._extra_params + + def __call__(self, input: TEXT) -> DenseVectorType: + """Make the embedding function callable.""" + return self.embed(input) + + @lru_cache(maxsize=10) + def embed(self, input: TEXT) -> DenseVectorType: + """Generate dense embedding vector for the input text. + + This method calls the OpenAI Embeddings API to convert input text + into a dense vector representation. Results are cached to improve + performance for repeated inputs. + + Args: + input (TEXT): Input text string to embed. Must be non-empty after + stripping whitespace. Maximum length is 8191 tokens for most models. + + Returns: + DenseVectorType: A list of floats representing the embedding vector. + Length equals ``self.dimension``. Example: + ``[0.123, -0.456, 0.789, ...]`` + + Raises: + TypeError: If ``input`` is not a string. + ValueError: If input is empty/whitespace-only, or if the API returns + an error or malformed response. + RuntimeError: If network connectivity issues or OpenAI service + errors occur. + + Examples: + >>> emb = OpenAIDenseEmbedding() + >>> vector = emb.embed("Natural language processing") + >>> len(vector) + 1536 + >>> isinstance(vector[0], float) + True + + >>> # Error: empty input + >>> emb.embed(" ") + ValueError: Input text cannot be empty or whitespace only + + >>> # Error: non-string input + >>> emb.embed(123) + TypeError: Expected 'input' to be str, got int + + Note: + - This method is cached (maxsize=10). Identical inputs return cached results. + - The cache is based on exact string match (case-sensitive). + - Consider pre-processing text (lowercasing, normalization) for better caching. + """ + if not isinstance(input, TEXT): + raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}") + + input = input.strip() + if not input: + raise ValueError("Input text cannot be empty or whitespace only") + + # Call API + embedding_vector = self._call_text_embedding_api( + input=input, + dimension=self._custom_dimension, + ) + + # Verify dimension + if len(embedding_vector) != self.dimension: + raise ValueError( + f"Dimension mismatch: expected {self.dimension}, " + f"got {len(embedding_vector)}" + ) + + return embedding_vector diff --git a/python/zvec/extension/openai_function.py b/python/zvec/extension/openai_function.py new file mode 100644 index 0000000..d3f4de2 --- /dev/null +++ b/python/zvec/extension/openai_function.py @@ -0,0 +1,149 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import os +from typing import ClassVar, Optional + +from ..common.constants import TEXT +from ..tool import require_module + + +class OpenAIFunctionBase: + """Base class for OpenAI functions. + + This base class provides common functionality for calling OpenAI APIs + and handling responses. It supports embeddings (dense) operations. + + This class is not meant to be used directly. Use concrete implementations: + - ``OpenAIDenseEmbedding`` for dense embeddings + + Args: + model (str): OpenAI model identifier. + api_key (Optional[str]): OpenAI API authentication key. + base_url (Optional[str]): Custom API base URL. + + Note: + - This is an internal base class for code reuse across OpenAI features + - Subclasses should inherit from appropriate Protocol + - Provides unified API connection and response handling + """ + + # Model default dimensions + _MODEL_DIMENSIONS: ClassVar[dict[str, int]] = { + "text-embedding-3-small": 1536, + "text-embedding-3-large": 3072, + "text-embedding-ada-002": 1536, + } + + def __init__( + self, + model: str, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + ): + """Initialize the base OpenAI functionality. + + Args: + model (str): OpenAI model name. + api_key (Optional[str]): API key or None to use environment variable. + base_url (Optional[str]): Custom API base URL or None for default. + + Raises: + ValueError: If API key is not provided and not in environment. + """ + self._model = model + self._api_key = api_key or os.environ.get("OPENAI_API_KEY") + self._base_url = base_url + + if not self._api_key: + raise ValueError( + "OpenAI API key is required. Please provide 'api_key' parameter " + "or set the 'OPENAI_API_KEY' environment variable." + ) + + @property + def model(self) -> str: + """str: The OpenAI model name currently in use.""" + return self._model + + def _get_client(self): + """Get OpenAI client instance. + + Returns: + OpenAI: Configured OpenAI client. + + Raises: + ImportError: If openai package is not installed. + """ + openai = require_module("openai") + + if self._base_url: + return openai.OpenAI(api_key=self._api_key, base_url=self._base_url) + return openai.OpenAI(api_key=self._api_key) + + def _call_text_embedding_api( + self, + input: TEXT, + dimension: Optional[int] = None, + ) -> list: + """Call OpenAI Embeddings API. + + Args: + input (TEXT): Input text to embed. + dimension (Optional[int]): Target dimension (for models that support it). + + Returns: + list: Embedding vector as list of floats. + + Raises: + RuntimeError: If API call fails. + ValueError: If API returns error response. + """ + try: + client = self._get_client() + + # Prepare embedding parameters + params = {"model": self.model, "input": input} + + # Add dimension parameter for models that support it + if dimension is not None: + params["dimensions"] = dimension + + # Call OpenAI API + response = client.embeddings.create(**params) + + except Exception as e: + # Check if it's an OpenAI API error + openai = require_module("openai") + if isinstance(e, (openai.APIError, openai.APIConnectionError)): + raise RuntimeError(f"Failed to call OpenAI API: {e!s}") from e + raise RuntimeError(f"Unexpected error during API call: {e!s}") from e + + # Extract embedding from response + try: + if not response.data: + raise ValueError("Invalid API response: no embedding data returned") + + embedding_vector = response.data[0].embedding + + if not isinstance(embedding_vector, list): + raise ValueError( + "Invalid API response: embedding is not a list of numbers" + ) + + return embedding_vector + + except (AttributeError, IndexError, TypeError) as e: + raise ValueError(f"Failed to parse API response: {e!s}") from e diff --git a/python/zvec/extension/qwen_embedding_function.py b/python/zvec/extension/qwen_embedding_function.py new file mode 100644 index 0000000..7bdb69b --- /dev/null +++ b/python/zvec/extension/qwen_embedding_function.py @@ -0,0 +1,537 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from functools import lru_cache +from typing import Optional + +from ..common.constants import TEXT, DenseVectorType, SparseVectorType +from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction +from .qwen_function import QwenFunctionBase + + +class QwenDenseEmbedding(QwenFunctionBase, DenseEmbeddingFunction[TEXT]): + """Dense text embedding function using Qwen (DashScope) API. + + This class provides text-to-vector embedding capabilities using Alibaba Cloud's + DashScope service and Qwen embedding models. It inherits from + ``DenseEmbeddingFunction`` and implements dense text embedding. + + The implementation supports various Qwen embedding models with configurable + dimensions and includes automatic result caching for improved performance. + + Args: + dimension (int): Desired output embedding dimension. Common values: + - 512: Balanced performance and accuracy + - 1024: Higher accuracy, larger storage + - 1536: Maximum accuracy for supported models + model (str, optional): DashScope embedding model identifier. + Defaults to ``"text-embedding-v4"``. Other options include: + - ``"text-embedding-v3"`` + - ``"text-embedding-v2"`` + - ``"text-embedding-v1"`` + api_key (Optional[str], optional): DashScope API authentication key. + If ``None``, reads from ``DASHSCOPE_API_KEY`` environment variable. + Obtain your key from: https://dashscope.console.aliyun.com/ + **kwargs: Additional DashScope API parameters. Supported options: + - ``text_type`` (str): Specifies the text role in retrieval tasks. + Options: ``"query"`` (search query) or ``"document"`` (indexed content). + This parameter optimizes embeddings for asymmetric search scenarios. + + Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api + + Attributes: + dimension (int): The embedding vector dimension. + data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation. + model (str): The DashScope model name being used. + + Raises: + ValueError: If API key is not provided and not found in environment, + or if API returns an error response. + TypeError: If input to ``embed()`` is not a string. + RuntimeError: If network error or DashScope service error occurs. + + Note: + - Requires Python 3.10, 3.11, or 3.12 + - Requires the ``dashscope`` package: ``pip install dashscope`` + - Embedding results are cached (LRU cache, maxsize=10) to reduce API calls + - Network connectivity to DashScope API endpoints is required + - API usage may incur costs based on your DashScope subscription plan + + **Parameter Guidelines:** + + - Use ``text_type="query"`` for search queries and ``text_type="document"`` + for indexed content to optimize asymmetric retrieval tasks. + - For detailed API specifications and parameter usage, refer to: + https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api + + Examples: + >>> # Basic usage with default model + >>> from zvec.extension import QwenDenseEmbedding + >>> import os + >>> os.environ["DASHSCOPE_API_KEY"] = "your-api-key" + >>> + >>> emb_func = QwenDenseEmbedding(dimension=1024) + >>> vector = emb_func.embed("Hello, world!") + >>> len(vector) + 1024 + + >>> # Using specific model with explicit API key + >>> emb_func = QwenDenseEmbedding( + ... dimension=512, + ... model="text-embedding-v3", + ... api_key="sk-xxxxx" + ... ) + >>> vector = emb_func("Machine learning is fascinating") + >>> isinstance(vector, list) + True + + >>> # Using with custom parameters (text_type) + >>> # For search queries - optimize for query-document matching + >>> emb_func = QwenDenseEmbedding( + ... dimension=1024, + ... text_type="query" + ... ) + >>> query_vector = emb_func.embed("What is machine learning?") + >>> + >>> # For document embeddings - optimize for being matched by queries + >>> doc_emb_func = QwenDenseEmbedding( + ... dimension=1024, + ... text_type="document" + ... ) + >>> doc_vector = doc_emb_func.embed( + ... "Machine learning is a subset of artificial intelligence..." + ... ) + + >>> # Batch processing with caching benefit + >>> texts = ["First text", "Second text", "First text"] + >>> vectors = [emb_func.embed(text) for text in texts] + >>> # Third call uses cached result for "First text" + + >>> # Error handling + >>> try: + ... emb_func.embed("") # Empty string + ... except ValueError as e: + ... print(f"Error: {e}") + Error: Input text cannot be empty or whitespace only + + See Also: + - ``DenseEmbeddingFunction``: Base class for dense embeddings + - ``SparseEmbeddingFunction``: Base class for sparse embeddings + """ + + def __init__( + self, + dimension: int, + model: str = "text-embedding-v4", + api_key: Optional[str] = None, + **kwargs, + ): + """Initialize the Qwen dense embedding function. + + Args: + dimension (int): Target embedding dimension. + model (str): DashScope model name. Defaults to "text-embedding-v4". + api_key (Optional[str]): API key or None to use environment variable. + **kwargs: Additional DashScope API parameters. Supported options: + - ``text_type`` (str): Text role in asymmetric retrieval. + * ``"query"``: Optimize for search queries (short, question-like). + * ``"document"``: Optimize for indexed documents (longer content). + Using appropriate text_type improves retrieval accuracy by + optimizing the embedding space for query-document matching. + + For detailed API documentation, see: + https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api + + Raises: + ValueError: If API key is not provided and not in environment. + """ + # Initialize base class for API connection + QwenFunctionBase.__init__(self, model=model, api_key=api_key) + + # Store dense-specific attributes + self._dimension = dimension + self._extra_params = kwargs + + @property + def dimension(self) -> int: + """int: The expected dimensionality of the embedding vector.""" + return self._dimension + + @property + def extra_params(self) -> dict: + """dict: Extra parameters for model-specific customization.""" + return self._extra_params + + def __call__(self, input: TEXT) -> DenseVectorType: + """Make the embedding function callable.""" + return self.embed(input) + + @lru_cache(maxsize=10) + def embed(self, input: TEXT) -> DenseVectorType: + """Generate dense embedding vector for the input text. + + This method calls the DashScope TextEmbedding API to convert input text + into a dense vector representation. Results are cached to improve + performance for repeated inputs. + + Args: + input (TEXT): Input text string to embed. Must be non-empty after + stripping whitespace. Maximum length depends on the model used + (typically 2048-8192 tokens). + + Returns: + DenseVectorType: A list of floats representing the embedding vector. + Length equals ``self.dimension``. Example: + ``[0.123, -0.456, 0.789, ...]`` + + Raises: + TypeError: If ``input`` is not a string. + ValueError: If input is empty/whitespace-only, or if the API returns + an error or malformed response. + RuntimeError: If network connectivity issues or DashScope service + errors occur. + + Examples: + >>> emb = QwenDenseEmbedding(dimension=1024) + >>> vector = emb.embed("Natural language processing") + >>> len(vector) + 1024 + >>> isinstance(vector[0], float) + True + + >>> # Error: empty input + >>> emb.embed(" ") + ValueError: Input text cannot be empty or whitespace only + + >>> # Error: non-string input + >>> emb.embed(123) + TypeError: Expected 'input' to be str, got int + + Note: + - This method is cached (maxsize=10). Identical inputs return cached results. + - The cache is based on exact string match (case-sensitive). + - Consider pre-processing text (lowercasing, normalization) for better caching. + """ + if not isinstance(input, TEXT): + raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}") + + input = input.strip() + if not input: + raise ValueError("Input text cannot be empty or whitespace only") + + # Call API with dense output type + output = self._call_text_embedding_api( + input=input, + dimension=self.dimension, + output_type="dense", + text_type=self.extra_params.get("text_type"), + ) + + embeddings = output.get("embeddings") + if not isinstance(embeddings, list): + raise ValueError( + "Invalid API response: 'embeddings' field is missing or not a list" + ) + + if len(embeddings) != 1: + raise ValueError( + f"Expected exactly 1 embedding in response, got {len(embeddings)}" + ) + + first_emb = embeddings[0] + if not isinstance(first_emb, dict): + raise ValueError("Invalid API response: embedding item is not a dictionary") + + embedding_vector = first_emb.get("embedding") + if not isinstance(embedding_vector, list): + raise ValueError( + "Invalid API response: 'embedding' field is missing or not a list" + ) + + if len(embedding_vector) != self.dimension: + raise ValueError( + f"Dimension mismatch: expected {self.dimension}, " + f"got {len(embedding_vector)}" + ) + + return list(embedding_vector) + + +class QwenSparseEmbedding(QwenFunctionBase, SparseEmbeddingFunction[TEXT]): + """Sparse text embedding function using Qwen (DashScope) API. + + This class provides text-to-sparse-vector embedding capabilities using + Alibaba Cloud's DashScope service and Qwen embedding models. It generates + sparse keyword-weighted vectors suitable for lexical matching and BM25-style + retrieval scenarios. + + Sparse embeddings are particularly useful for: + - Keyword-based search and exact matching + - Hybrid retrieval (combining with dense embeddings) + - Interpretable search results (weights show term importance) + + Args: + dimension (int): Desired output embedding dimension. Common values: + - 512: Balanced performance and accuracy + - 1024: Higher accuracy, larger storage + - 1536: Maximum accuracy for supported models + model (str, optional): DashScope embedding model identifier. + Defaults to ``"text-embedding-v4"``. Other options include: + - ``"text-embedding-v3"`` + - ``"text-embedding-v2"`` + api_key (Optional[str], optional): DashScope API authentication key. + If ``None``, reads from ``DASHSCOPE_API_KEY`` environment variable. + Obtain your key from: https://dashscope.console.aliyun.com/ + **kwargs: Additional DashScope API parameters. Supported options: + - ``encoding_type`` (Literal["query", "document"]): Encoding type. + * ``"query"``: Optimize for search queries (default). + * ``"document"``: Optimize for indexed documents. + This distinction is important for asymmetric retrieval tasks. + + Attributes: + model (str): The DashScope model name being used. + encoding_type (str): The encoding type ("query" or "document"). + + Raises: + ValueError: If API key is not provided and not found in environment, + or if API returns an error response. + TypeError: If input to ``embed()`` is not a string. + RuntimeError: If network error or DashScope service error occurs. + + Note: + - Requires Python 3.10, 3.11, or 3.12 + - Requires the ``dashscope`` package: ``pip install dashscope`` + - Embedding results are cached (LRU cache, maxsize=10) to reduce API calls + - Network connectivity to DashScope API endpoints is required + - API usage may incur costs based on your DashScope subscription plan + - Sparse vectors have only non-zero dimensions stored as dict + - Output is sorted by indices (keys) in ascending order + + **Parameter Guidelines:** + + - Use ``encoding_type="query"`` for search queries and + ``encoding_type="document"`` for indexed content to optimize + asymmetric retrieval tasks. + - For detailed API specifications, refer to: + https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api + + Examples: + >>> # Basic usage for query embedding + >>> from zvec.extension import QwenSparseEmbedding + >>> import os + >>> os.environ["DASHSCOPE_API_KEY"] = "your-api-key" + >>> + >>> query_emb = QwenSparseEmbedding(dimension=1024, encoding_type="query") + >>> query_vec = query_emb.embed("machine learning") + >>> type(query_vec) + + >>> len(query_vec) # Only non-zero dimensions + 156 + + >>> # Document embedding + >>> doc_emb = QwenSparseEmbedding(dimension=1024, encoding_type="document") + >>> doc_vec = doc_emb.embed("Machine learning is a subset of AI") + >>> isinstance(doc_vec, dict) + True + + >>> # Asymmetric retrieval example + >>> query_vec = query_emb.embed("what causes aging fast") + >>> doc_vec = doc_emb.embed( + ... "UV-A light causes tanning, skin aging, and cataracts..." + ... ) + >>> + >>> # Calculate similarity (dot product for sparse vectors) + >>> similarity = sum( + ... query_vec.get(k, 0) * doc_vec.get(k, 0) + ... for k in set(query_vec) | set(doc_vec) + ... ) + + >>> # Output is sorted by indices + >>> list(query_vec.items())[:5] # First 5 dimensions (by index) + [(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)] + + >>> # Hybrid retrieval (combining dense + sparse) + >>> from zvec.extension import QwenDenseEmbedding + >>> dense_emb = QwenDenseEmbedding(dimension=1024) + >>> sparse_emb = QwenSparseEmbedding(dimension=1024) + >>> + >>> query = "deep learning neural networks" + >>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...] + >>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...} + + >>> # Error handling + >>> try: + ... sparse_emb.embed("") # Empty string + ... except ValueError as e: + ... print(f"Error: {e}") + Error: Input text cannot be empty or whitespace only + + See Also: + - ``SparseEmbeddingFunction``: Base class for sparse embeddings + - ``QwenDenseEmbedding``: Dense embedding using Qwen API + - ``DefaultSparseEmbedding``: Sparse embedding with SPLADE model + """ + + def __init__( + self, + dimension: int, + model: str = "text-embedding-v4", + api_key: Optional[str] = None, + **kwargs, + ): + """Initialize the Qwen sparse embedding function. + + Args: + dimension (int): Target embedding dimension. + model (str): DashScope model name. Defaults to "text-embedding-v4". + api_key (Optional[str]): API key or None to use environment variable. + **kwargs: Additional DashScope API parameters. Supported options: + - ``encoding_type`` (Literal["query", "document"]): Encoding type. + * ``"query"``: Optimize for search queries (default). + * ``"document"``: Optimize for indexed documents. + This distinction is important for asymmetric retrieval tasks. + + Raises: + ValueError: If API key is not provided and not in environment. + """ + # Initialize base class for API connection + QwenFunctionBase.__init__(self, model=model, api_key=api_key) + + self._dimension = dimension + self._extra_params = kwargs + + @property + def extra_params(self) -> dict: + """dict: Extra parameters for model-specific customization.""" + return self._extra_params + + def __call__(self, input: TEXT) -> SparseVectorType: + """Make the embedding function callable.""" + return self.embed(input) + + @lru_cache(maxsize=10) + def embed(self, input: TEXT) -> SparseVectorType: + """Generate sparse embedding vector for the input text. + + This method calls the DashScope TextEmbedding API with sparse output type + to convert input text into a sparse vector representation. The result is + a dictionary where keys are dimension indices and values are importance + weights (only non-zero values included). + + The embedding is optimized based on the ``encoding_type`` specified during + initialization: "query" for search queries or "document" for indexed content. + + Args: + input (TEXT): Input text string to embed. Must be non-empty after + stripping whitespace. Maximum length depends on the model used + (typically 2048-8192 tokens). + + Returns: + SparseVectorType: A dictionary mapping dimension index to weight. + Only non-zero dimensions are included. The dictionary is sorted + by indices (keys) in ascending order for consistent output. + Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}`` + + Raises: + TypeError: If ``input`` is not a string. + ValueError: If input is empty/whitespace-only, or if the API returns + an error or malformed response. + RuntimeError: If network connectivity issues or DashScope service + errors occur. + + Examples: + >>> emb = QwenSparseEmbedding(dimension=1024, encoding_type="query") + >>> sparse_vec = emb.embed("machine learning") + >>> isinstance(sparse_vec, dict) + True + >>> + >>> # Verify sorted output + >>> keys = list(sparse_vec.keys()) + >>> keys == sorted(keys) + True + + >>> # Error: empty input + >>> emb.embed(" ") + ValueError: Input text cannot be empty or whitespace only + + >>> # Error: non-string input + >>> emb.embed(123) + TypeError: Expected 'input' to be str, got int + + Note: + - This method is cached (maxsize=10). Identical inputs return cached results. + - The cache is based on exact string match (case-sensitive). + - Output dictionary is always sorted by indices for consistency. + """ + if not isinstance(input, TEXT): + raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}") + + input = input.strip() + if not input: + raise ValueError("Input text cannot be empty or whitespace only") + + # Call API with sparse output type + output = self._call_text_embedding_api( + input=input, + dimension=self._dimension, + output_type="sparse", + text_type=self.extra_params.get("encoding_type", "query"), + ) + + embeddings = output.get("embeddings") + if not isinstance(embeddings, list): + raise ValueError( + "Invalid API response: 'embeddings' field is missing or not a list" + ) + + if len(embeddings) != 1: + raise ValueError( + f"Expected exactly 1 embedding in response, got {len(embeddings)}" + ) + + first_emb = embeddings[0] + if not isinstance(first_emb, dict): + raise ValueError("Invalid API response: embedding item is not a dictionary") + + sparse_embedding = first_emb.get("sparse_embedding") + if not isinstance(sparse_embedding, list): + raise ValueError( + "Invalid API response: 'sparse_embedding' field is missing or not a list" + ) + + # Parse sparse embedding: convert array of {index, value, token} to dict + sparse_dict = {} + for item in sparse_embedding: + if not isinstance(item, dict): + raise ValueError( + "Invalid API response: sparse_embedding item is not a dictionary" + ) + + index = item.get("index") + value = item.get("value") + + if index is None or value is None: + raise ValueError( + "Invalid API response: sparse_embedding item missing 'index' or 'value'" + ) + + # Convert to int and float, filter positive values + idx = int(index) + val = float(value) + if val > 0: + sparse_dict[idx] = val + + # Sort by indices (keys) to ensure consistent ordering + return dict(sorted(sparse_dict.items())) diff --git a/python/zvec/extension/qwen_function.py b/python/zvec/extension/qwen_function.py new file mode 100644 index 0000000..b15ee4b --- /dev/null +++ b/python/zvec/extension/qwen_function.py @@ -0,0 +1,186 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import os +from http import HTTPStatus +from typing import Optional + +from ..common.constants import TEXT +from ..tool import require_module + + +class QwenFunctionBase: + """Base class for Qwen (DashScope) functions. + + This base class provides common functionality for calling DashScope APIs + and handling responses. It supports embeddings (dense and sparse) and + re-ranking operations. + + This class is not meant to be used directly. Use concrete implementations: + - ``QwenDenseEmbedding`` for dense embeddings + - ``QwenSparseEmbedding`` for sparse embeddings + - ``QwenReRanker`` for semantic re-ranking + + Args: + model (str): DashScope model identifier. + api_key (Optional[str]): DashScope API authentication key. + + Note: + - This is an internal base class for code reuse across Qwen features + - Subclasses should inherit from appropriate Protocol/ABC + - Provides unified API connection and response handling + """ + + def __init__( + self, + model: str, + api_key: Optional[str] = None, + ): + """Initialize the base Qwen embedding functionality. + + Args: + model (str): DashScope model name. + api_key (Optional[str]): API key or None to use environment variable. + + Raises: + ValueError: If API key is not provided and not in environment. + """ + self._model = model + self._api_key = api_key or os.environ.get("DASHSCOPE_API_KEY") + if not self._api_key: + raise ValueError( + "DashScope API key is required. Please provide 'api_key' parameter " + "or set the 'DASHSCOPE_API_KEY' environment variable." + ) + + @property + def model(self) -> str: + """str: The DashScope embedding model name currently in use.""" + return self._model + + def _get_connection(self): + """Establish connection to DashScope API. + + Returns: + module: The dashscope module with API key configured. + + Raises: + ImportError: If dashscope package is not installed. + """ + dashscope = require_module("dashscope") + dashscope.api_key = self._api_key + return dashscope + + def _call_text_embedding_api( + self, + input: TEXT, + dimension: int, + output_type: str, + text_type: Optional[str] = None, + ) -> dict: + """Call DashScope TextEmbedding API. + + Args: + input (TEXT): Input text to embed. + dimension (int): Target embedding dimension. + output_type (str): Output type ("dense" or "sparse"). + text_type (Optional[str]): Text type ("query" or "document"). + + Returns: + dict: API response output field. + + Raises: + RuntimeError: If API call fails. + ValueError: If API returns error response. + """ + try: + # Prepare API call parameters + call_params = { + "model": self.model, + "input": input, + "dimension": dimension, + "output_type": output_type, + } + + # Add optional text_type parameter if provided + if text_type is not None: + call_params["text_type"] = text_type + + resp = self._get_connection().TextEmbedding.call(**call_params) + except Exception as e: + raise RuntimeError(f"Failed to call DashScope API: {e!s}") from e + + if resp.status_code != HTTPStatus.OK: + error_msg = getattr(resp, "message", "Unknown error") + error_code = getattr(resp, "code", "N/A") + raise ValueError( + f"DashScope API error: [Code={error_code}, " + f"Status={resp.status_code}] {error_msg}" + ) + + output = getattr(resp, "output", None) + if not isinstance(output, dict): + raise ValueError( + "Invalid API response: missing or malformed 'output' field" + ) + + return output + + def _call_rerank_api( + self, + query: str, + documents: list[str], + top_n: int, + ) -> dict: + """Call DashScope TextReRank API. + + Args: + query (str): Query text for semantic matching. + documents (list[str]): List of document texts to re-rank. + top_n (int): Maximum number of documents to return. + + Returns: + dict: API response output field containing re-ranked results. + + Raises: + RuntimeError: If API call fails. + ValueError: If API returns error response. + """ + try: + resp = self._get_connection().TextReRank.call( + model=self.model, + query=query, + documents=documents, + top_n=top_n, + return_documents=False, + ) + except Exception as e: + raise RuntimeError(f"Failed to call DashScope API: {e!s}") from e + + if resp.status_code != HTTPStatus.OK: + error_msg = getattr(resp, "message", "Unknown error") + error_code = getattr(resp, "code", "N/A") + raise ValueError( + f"DashScope API error: [Code={error_code}, " + f"Status={resp.status_code}] {error_msg}" + ) + + output = getattr(resp, "output", None) + if not isinstance(output, dict): + raise ValueError( + "Invalid API response: missing or malformed 'output' field" + ) + + return output diff --git a/python/zvec/extension/qwen_rerank_function.py b/python/zvec/extension/qwen_rerank_function.py new file mode 100644 index 0000000..0a06eb2 --- /dev/null +++ b/python/zvec/extension/qwen_rerank_function.py @@ -0,0 +1,177 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from ..model.doc import Doc, DocList +from .qwen_function import QwenFunctionBase +from .rerank_function import RerankFunction + +if TYPE_CHECKING: + from ..model.schema import FieldSchema, VectorSchema + + +class QwenReRanker(QwenFunctionBase, RerankFunction): + """Re-ranker using Qwen (DashScope) cross-encoder API for semantic re-ranking. + + This re-ranker leverages DashScope's TextReRank service to perform + cross-encoder style re-ranking. It sends query and document pairs to the + API and receives relevance scores based on deep semantic understanding. + + The re-ranker is suitable for single-vector or multi-vector search scenarios + where semantic relevance to a specific query is required. + + Args: + query (str): Query text for semantic re-ranking. **Required**. + rerank_field (str): Document field name to use as re-ranking input text. + **Required** (e.g., "content", "title", "body"). + model (str, optional): DashScope re-ranking model identifier. + Defaults to ``"gte-rerank-v2"``. + api_key (Optional[str], optional): DashScope API authentication key. + If not provided, reads from ``DASHSCOPE_API_KEY`` environment variable. + + Raises: + ValueError: If ``query`` is empty/None, ``rerank_field`` is None, + or API key is not available. + + Note: + - Requires ``dashscope`` Python package installed + - Documents without valid content in ``rerank_field`` are skipped + - API rate limits and quotas apply per DashScope subscription + + Example: + >>> reranker = QwenReRanker( + ... query="machine learning algorithms", + ... rerank_field="content", + ... model="gte-rerank-v2", + ... api_key="your-api-key" + ... ) + >>> # Use in collection.query(reranker=reranker) + """ + + def __init__( + self, + query: Optional[str] = None, + rerank_field: Optional[str] = None, + model: str = "gte-rerank-v2", + api_key: Optional[str] = None, + ): + """Initialize QwenReRanker with query and configuration. + + Args: + query (Optional[str]): Query text for semantic matching. Required. + rerank_field (Optional[str]): Document field for re-ranking input. + model (str): DashScope model name. + api_key (Optional[str]): API key or None to use environment variable. + + Raises: + ValueError: If query is empty or API key is unavailable. + """ + QwenFunctionBase.__init__(self, model=model, api_key=api_key) + self._rerank_field = rerank_field + + if not query: + raise ValueError("Query is required for QwenReRanker") + self._query = query + + @property + def rerank_field(self) -> Optional[str]: + """Optional[str]: Field name used as re-ranking input.""" + return self._rerank_field + + @property + def query(self) -> str: + """str: Query text used for semantic re-ranking.""" + return self._query + + def rerank( + self, + query_results: list[list[Doc]], + topn: int = 10, + *, + fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002 + ) -> DocList: + """Re-rank documents using Qwen's TextReRank API. + + Sends document texts to DashScope TextReRank service along with the query. + Returns documents sorted by relevance scores from the cross-encoder model. + + Args: + query_results (list[list[Doc]]): Per-sub-query lists of retrieved + documents. Documents from all lists are deduplicated and + re-ranked together. + topn (int): Maximum number of documents to return. + fields: Unused; present for interface compatibility. + + Returns: + list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score`` + fields containing relevance scores from the API. + + Raises: + ValueError: If no valid documents are found or API call fails. + + Note: + - Duplicate documents (same ID) across lists are processed once + - Documents with empty/missing ``rerank_field`` content are skipped + - Returned scores are relevance scores from the cross-encoder model + """ + if not query_results: + return [] + + # Accept both dict (legacy) and list formats + if isinstance(query_results, dict): + query_results = list(query_results.values()) + + # Collect and deduplicate documents + id_to_doc: dict[str, Doc] = {} + doc_ids: list[str] = [] + contents: list[str] = [] + + for query_result in query_results: + for doc in query_result: + doc_id = doc.id + if doc_id in id_to_doc: + continue + + # Extract text content from specified field + field_value = doc.field(self.rerank_field) + rank_content = str(field_value).strip() if field_value else "" + if not rank_content: + continue + + id_to_doc[doc_id] = doc + doc_ids.append(doc_id) + contents.append(rank_content) + + if not contents: + raise ValueError("No documents to rerank") + + # Call DashScope TextReRank API + output = self._call_rerank_api( + query=self.query, + documents=contents, + top_n=topn, + ) + + # Build result list with updated scores + results: DocList = [] + for item in output["results"]: + idx = item["index"] + doc_id = doc_ids[idx] + doc = id_to_doc[doc_id] + new_doc = doc._replace(score=item["relevance_score"]) + results.append(new_doc) + + return results diff --git a/python/zvec/extension/rerank_function.py b/python/zvec/extension/rerank_function.py new file mode 100644 index 0000000..09a26c4 --- /dev/null +++ b/python/zvec/extension/rerank_function.py @@ -0,0 +1,56 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from ..model.doc import Doc, DocList + +if TYPE_CHECKING: + from ..model.schema import FieldSchema, VectorSchema + + +class RerankFunction(ABC): + """Abstract base class for reranker parameter containers. + + Subclasses define rerank parameters and implement _to_cpp_params() + for conversion to C++ parameter structs (used by collection fast path). + Each subclass also provides a standalone rerank() implementation. + """ + + def _to_cpp_params(self): + """Return C++ reranker params. Override in subclasses that use C++ path.""" + raise NotImplementedError + + @abstractmethod + def rerank( + self, + query_results: list[list[Doc]], + topn: int = 10, + *, + fields: list[FieldSchema | VectorSchema] | None = None, + ) -> DocList: + """Execute rerank on sub-query results. + + Args: + query_results: List of per-sub-query document lists. + topn: Maximum number of results to return. + fields: Per-sub-query Python FieldSchema/VectorSchema objects + (required for WeightedReRanker score normalization). + + Returns: + Re-ranked document list. + """ + ... diff --git a/python/zvec/extension/sentence_transformer_embedding_function.py b/python/zvec/extension/sentence_transformer_embedding_function.py new file mode 100644 index 0000000..032f02e --- /dev/null +++ b/python/zvec/extension/sentence_transformer_embedding_function.py @@ -0,0 +1,839 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import ClassVar, Literal, Optional + +import numpy as np + +from ..common.constants import TEXT, DenseVectorType, SparseVectorType +from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction +from .sentence_transformer_function import SentenceTransformerFunctionBase + + +class DefaultLocalDenseEmbedding( + SentenceTransformerFunctionBase, DenseEmbeddingFunction[TEXT] +): + """Default local dense embedding using all-MiniLM-L6-v2 model. + + This is the default implementation for dense text embedding that uses the + ``all-MiniLM-L6-v2`` model from Hugging Face by default. This model provides + a good balance between speed and quality for general-purpose text embedding. + + The class provides text-to-vector dense embedding capabilities using the + sentence-transformers library. It supports models from Hugging Face Hub and + ModelScope, runs locally without API calls, and supports CPU/GPU acceleration. + + The model produces 384-dimensional embeddings and is optimized for semantic + similarity tasks. It runs locally without requiring API keys. + + Args: + model_source (Literal["huggingface", "modelscope"], optional): Model source. + - ``"huggingface"``: Use Hugging Face Hub (default, for international users) + - ``"modelscope"``: Use ModelScope (recommended for users in China) + Defaults to ``"huggingface"``. + device (Optional[str], optional): Device to run the model on. + Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None`` + for automatic detection. Defaults to ``None``. + normalize_embeddings (bool, optional): Whether to normalize embeddings to + unit length (L2 normalization). Useful for cosine similarity. + Defaults to ``True``. + batch_size (int, optional): Batch size for encoding. Defaults to ``32``. + **kwargs: Additional parameters for future extension. + + Attributes: + dimension (int): Always 384 for both models. + model_name (str): "all-MiniLM-L6-v2" (HF) or "iic/nlp_gte_sentence-embedding_chinese-small" (MS). + model_source (str): The model source being used. + device (str): The device the model is running on. + + Raises: + ValueError: If the model cannot be loaded or input is invalid. + TypeError: If input to ``embed()`` is not a string. + RuntimeError: If model inference fails. + + Note: + - Requires Python 3.10, 3.11, or 3.12 + - Requires the ``sentence-transformers`` package: + ``pip install sentence-transformers`` + - For ModelScope, also requires: ``pip install modelscope`` + - First run downloads the model (~50-80MB) from chosen source + - Hugging Face cache: ``~/.cache/torch/sentence_transformers/`` + - ModelScope cache: ``~/.cache/modelscope/hub/`` + - No API keys or network required after initial download + - Inference speed: ~1000 sentences/sec on CPU, ~10000 on GPU + + **For users in China:** + + If you encounter Hugging Face access issues, use ModelScope instead: + + .. code-block:: python + + # Recommended for users in China + emb = DefaultLocalDenseEmbedding(model_source="modelscope") + + Alternatively, use Hugging Face mirror: + + .. code-block:: bash + + export HF_ENDPOINT=https://hf-mirror.com + # Then use default Hugging Face mode + + Examples: + >>> # Basic usage with Hugging Face (default) + >>> from zvec.extension import DefaultLocalDenseEmbedding + >>> + >>> emb_func = DefaultLocalDenseEmbedding() + >>> vector = emb_func.embed("Hello, world!") + >>> len(vector) + 384 + >>> isinstance(vector, list) + True + + >>> # Recommended for users in China (uses ModelScope) + >>> emb_func = DefaultLocalDenseEmbedding(model_source="modelscope") + >>> vector = emb_func.embed("你好，世界！") # Works well with Chinese text + >>> len(vector) + 384 + + >>> # Alternative for China users: Use Hugging Face mirror + >>> import os + >>> os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" + >>> emb_func = DefaultLocalDenseEmbedding() # Uses HF mirror + >>> vector = emb_func.embed("Hello, world!") + + >>> # Using GPU for faster inference + >>> emb_func = DefaultLocalDenseEmbedding(device="cuda") + >>> vector = emb_func("Machine learning is fascinating") + >>> # Normalized vector has unit length + >>> import numpy as np + >>> np.linalg.norm(vector) + 1.0 + + >>> # Batch processing + >>> texts = ["First text", "Second text", "Third text"] + >>> vectors = [emb_func.embed(text) for text in texts] + >>> len(vectors) + 3 + >>> all(len(v) == 384 for v in vectors) + True + + >>> # Semantic similarity + >>> v1 = emb_func.embed("The cat sits on the mat") + >>> v2 = emb_func.embed("A feline rests on a rug") + >>> v3 = emb_func.embed("Python programming") + >>> similarity_high = np.dot(v1, v2) # Similar sentences + >>> similarity_low = np.dot(v1, v3) # Different topics + >>> similarity_high > similarity_low + True + + >>> # Error handling + >>> try: + ... emb_func.embed("") # Empty string + ... except ValueError as e: + ... print(f"Error: {e}") + Error: Input text cannot be empty or whitespace only + + See Also: + - ``DenseEmbeddingFunction``: Base class for dense embeddings + - ``DefaultLocalSparseEmbedding``: Sparse embedding with SPLADE + - ``QwenDenseEmbedding``: Alternative using Qwen API + """ + + def __init__( + self, + model_source: Literal["huggingface", "modelscope"] = "huggingface", + device: Optional[str] = None, + normalize_embeddings: bool = True, + batch_size: int = 32, + **kwargs, + ): + """Initialize with all-MiniLM-L6-v2 model. + + Args: + model_source (Literal["huggingface", "modelscope"]): Model source. + Defaults to "huggingface". + device (Optional[str]): Target device ("cpu", "cuda", "mps", or None). + Defaults to None (automatic detection). + normalize_embeddings (bool): Whether to L2-normalize output vectors. + Defaults to True. + batch_size (int): Batch size for encoding. Defaults to 32. + **kwargs: Additional parameters for future extension. + + Raises: + ImportError: If sentence-transformers or modelscope is not installed. + ValueError: If model cannot be loaded. + """ + # Use different models based on source + if model_source == "modelscope": + # Use Chinese-optimized model for ModelScope (better for Chinese text) + model_name = "iic/nlp_gte_sentence-embedding_chinese-small" + else: + model_name = "all-MiniLM-L6-v2" + + # Initialize base class for model loading + SentenceTransformerFunctionBase.__init__( + self, model_name=model_name, model_source=model_source, device=device + ) + + self._normalize_embeddings = normalize_embeddings + self._batch_size = batch_size + + # Load model and get dimension + model = self._get_model() + self._dimension = model.get_sentence_embedding_dimension() + + # Store extra parameters + self._extra_params = kwargs + + @property + def dimension(self) -> int: + """int: The expected dimensionality of the embedding vector.""" + return self._dimension + + @property + def extra_params(self) -> dict: + """dict: Extra parameters for model-specific customization.""" + return self._extra_params + + def __call__(self, input: str) -> DenseVectorType: + """Make the embedding function callable.""" + return self.embed(input) + + def embed(self, input: str) -> DenseVectorType: + """Generate dense embedding vector for the input text. + + This method uses the Sentence Transformer model to convert input text + into a dense vector representation. The model runs locally without + requiring API calls. + + Args: + input (str): Input text string to embed. Must be non-empty after + stripping whitespace. Maximum length depends on the model used + (typically 128-512 tokens for most models). + + Returns: + DenseVectorType: A list of floats representing the embedding vector. + Length equals ``self.dimension``. If ``normalize_embeddings=True``, + the vector has unit length. Example: + ``[0.123, -0.456, 0.789, ...]`` + + Raises: + TypeError: If ``input`` is not a string. + ValueError: If input is empty or whitespace-only. + RuntimeError: If model inference fails. + + Examples: + >>> emb = DefaultLocalDenseEmbedding() + >>> vector = emb.embed("Natural language processing") + >>> len(vector) + 384 + >>> isinstance(vector[0], float) + True + + >>> # Normalized vectors have unit length + >>> import numpy as np + >>> emb = DefaultLocalDenseEmbedding(normalize_embeddings=True) + >>> vector = emb.embed("Test sentence") + >>> np.linalg.norm(vector) + 1.0 + + >>> # Error: empty input + >>> emb.embed(" ") + ValueError: Input text cannot be empty or whitespace only + + >>> # Error: non-string input + >>> emb.embed(123) + TypeError: Expected 'input' to be str, got int + + >>> # Semantic similarity example + >>> v1 = emb.embed("The cat sits on the mat") + >>> v2 = emb.embed("A feline rests on a rug") + >>> similarity = np.dot(v1, v2) # High similarity due to semantic meaning + >>> similarity > 0.7 + True + + Note: + - First call may be slower due to model loading + - Subsequent calls are much faster as the model stays in memory + - For batch processing, consider encoding multiple texts together + (though this method handles single texts only) + - GPU acceleration provides 5-10x speedup over CPU + """ + if not isinstance(input, str): + raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}") + + input = input.strip() + if not input: + raise ValueError("Input text cannot be empty or whitespace only") + + try: + model = self._get_model() + embedding = model.encode( + input, + convert_to_numpy=True, + normalize_embeddings=self._normalize_embeddings, + batch_size=self._batch_size, + ) + + # Convert numpy array to list + if isinstance(embedding, np.ndarray): + embedding_list = embedding.tolist() + else: + embedding_list = list(embedding) + + # Validate dimension + if len(embedding_list) != self.dimension: + raise ValueError( + f"Dimension mismatch: expected {self.dimension}, " + f"got {len(embedding_list)}" + ) + + return embedding_list + + except Exception as e: + if isinstance(e, (TypeError, ValueError)): + raise + raise RuntimeError(f"Failed to generate embedding: {e!s}") from e + + +class DefaultLocalSparseEmbedding( + SentenceTransformerFunctionBase, SparseEmbeddingFunction[TEXT] +): + """Default local sparse embedding using SPLADE model. + + This class provides sparse vector embedding using the SPLADE (SParse Lexical + AnD Expansion) model. SPLADE generates sparse, interpretable representations + where each dimension corresponds to a vocabulary term with learned importance + weights. It's ideal for lexical matching, BM25-style retrieval, and hybrid + search scenarios. + + The default model is ``naver/splade-cocondenser-ensembledistil``, which is + publicly available without authentication. It produces sparse vectors with + thousands of dimensions but only hundreds of non-zero values, making them + efficient for storage and retrieval while maintaining strong lexical matching. + + **Model Caching:** + + This class uses class-level caching to share the SPLADE model across all instances + with the same configuration (model_source, device). This significantly reduces + memory usage when creating multiple instances for different encoding types + (query vs document). + + **Cache Management:** + + The class provides methods to manage the model cache: + + - ``clear_cache()``: Clear all cached models to free memory + - ``get_cache_info()``: Get information about cached models + - ``remove_from_cache(model_source, device)``: Remove a specific model from cache + + .. note:: + **Why not use splade-v3?** + + The newer ``naver/splade-v3`` model is gated (requires access approval). + We use ``naver/splade-cocondenser-ensembledistil`` instead. + + **To use splade-v3 (if you have access):** + + 1. Request access at https://huggingface.co/naver/splade-v3 + 2. Get your Hugging Face token from https://huggingface.co/settings/tokens + 3. Set environment variable: + + .. code-block:: bash + + export HF_TOKEN="your_huggingface_token" + + 4. Or login programmatically: + + .. code-block:: python + + from huggingface_hub import login + login(token="your_huggingface_token") + + 5. To use a custom SPLADE model, you can subclass this class and override + the model_name in ``__init__``, or create your own implementation + inheriting from ``SentenceTransformerFunctionBase`` and + ``SparseEmbeddingFunction``. + + Args: + model_source (Literal["huggingface", "modelscope"], optional): Model source. + Defaults to ``"huggingface"``. ModelScope support may vary for SPLADE models. + device (Optional[str], optional): Device to run the model on. + Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None`` + for automatic detection. Defaults to ``None``. + encoding_type (Literal["query", "document"], optional): Encoding type. + - ``"query"``: Optimize for search queries (default) + - ``"document"``: Optimize for indexed documents + **kwargs: Additional parameters (currently unused, for future extension). + + Attributes: + model_name (str): Model identifier. + model_source (str): The model source being used. + device (str): The device the model is running on. + + Raises: + ValueError: If the model cannot be loaded or input is invalid. + TypeError: If input to ``embed()`` is not a string. + RuntimeError: If model inference fails. + + Note: + - Requires Python 3.10, 3.11, or 3.12 + - Requires the ``sentence-transformers`` package: + ``pip install sentence-transformers`` + - First run downloads the model (~100MB) from Hugging Face + - Cache location: ``~/.cache/torch/sentence_transformers/`` + - No API keys or authentication required + - Sparse vectors have ~30k dimensions but only ~100-200 non-zero values + - Best combined with dense embeddings for hybrid retrieval + + **SPLADE vs Dense Embeddings:** + + - **Dense**: Continuous semantic vectors, good for semantic similarity + - **Sparse**: Lexical keyword-based, interpretable, good for exact matching + - **Hybrid**: Combine both for best retrieval performance + + Examples: + >>> # Memory-efficient: both instances share the same model (~200MB) + >>> from zvec.extension import DefaultLocalSparseEmbedding + >>> + >>> # Query embedding + >>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query") + >>> query_vec = query_emb.embed("machine learning algorithms") + >>> type(query_vec) + + >>> len(query_vec) # Only non-zero dimensions + 156 + + >>> # Document embedding (shares model with query_emb) + >>> doc_emb = DefaultLocalSparseEmbedding(encoding_type="document") + >>> doc_vec = doc_emb.embed("Machine learning is a subset of AI") + >>> # Total memory: ~200MB (not 400MB) thanks to model caching + + >>> # Asymmetric retrieval example + >>> query_vec = query_emb.embed("what causes aging fast") + >>> doc_vec = doc_emb.embed( + ... "UV-A light causes tanning, skin aging, and cataracts..." + ... ) + >>> + >>> # Calculate similarity (dot product for sparse vectors) + >>> similarity = sum( + ... query_vec.get(k, 0) * doc_vec.get(k, 0) + ... for k in set(query_vec) | set(doc_vec) + ... ) + + >>> # Batch processing + >>> queries = ["query 1", "query 2", "query 3"] + >>> query_vecs = [query_emb.embed(q) for q in queries] + >>> + >>> documents = ["doc 1", "doc 2", "doc 3"] + >>> doc_vecs = [doc_emb.embed(d) for d in documents] + + >>> # Inspecting sparse dimensions (output is sorted by indices) + >>> query_vec = query_emb.embed("machine learning") + >>> list(query_vec.items())[:5] # First 5 dimensions (by index) + [(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)] + >>> + >>> # Sort by weight to find most important terms + >>> sorted_by_weight = sorted(query_vec.items(), key=lambda x: x[1], reverse=True) + >>> top_5 = sorted_by_weight[:5] # Top 5 most important terms + >>> top_5 + [(1023, 1.45), (245, 1.23), (8901, 0.98), (5678, 0.87), (12034, 0.76)] + + >>> # Using GPU for faster inference + >>> sparse_emb = DefaultLocalSparseEmbedding(device="cuda") + >>> vector = sparse_emb.embed("natural language processing") + + >>> # Hybrid retrieval example (combining dense + sparse) + >>> from zvec.extension import DefaultDenseEmbedding + >>> dense_emb = DefaultDenseEmbedding() + >>> sparse_emb = DefaultLocalSparseEmbedding() + >>> + >>> query = "deep learning neural networks" + >>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...] + >>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...} + + >>> # Error handling + >>> try: + ... sparse_emb.embed("") # Empty string + ... except ValueError as e: + ... print(f"Error: {e}") + Error: Input text cannot be empty or whitespace only + + >>> # Cache management + >>> # Check cache status + >>> info = DefaultLocalSparseEmbedding.get_cache_info() + >>> print(f"Cached models: {info['cached_models']}") + Cached models: 1 + >>> + >>> # Clear cache to free memory + >>> DefaultLocalSparseEmbedding.clear_cache() + >>> info = DefaultLocalSparseEmbedding.get_cache_info() + >>> print(f"Cached models: {info['cached_models']}") + Cached models: 0 + >>> + >>> # Remove specific model from cache + >>> query_emb = DefaultLocalSparseEmbedding() # Creates CPU model + >>> cuda_emb = DefaultLocalSparseEmbedding(device="cuda") # Creates CUDA model + >>> info = DefaultLocalSparseEmbedding.get_cache_info() + >>> print(f"Cached models: {info['cached_models']}") + Cached models: 2 + >>> + >>> # Remove only CPU model + >>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device=None) + >>> print(f"Removed: {removed}") + True + >>> info = DefaultLocalSparseEmbedding.get_cache_info() + >>> print(f"Cached models: {info['cached_models']}") + Cached models: 1 + + See Also: + - ``SparseEmbeddingFunction``: Base class for sparse embeddings + - ``DefaultDenseEmbedding``: Dense embedding with all-MiniLM-L6-v2 + - ``QwenDenseEmbedding``: Alternative using Qwen API + + References: + - SPLADE Paper: https://arxiv.org/abs/2109.10086 + - Model: https://huggingface.co/naver/splade-cocondenser-ensembledistil + """ + + # Class-level model cache: {(model_name, model_source, device): model} + # Shared across all DefaultLocalSparseEmbedding instances to save memory + _model_cache: ClassVar[dict] = {} + + @classmethod + def clear_cache(cls) -> None: + """Clear all cached SPLADE models from memory. + + This is useful for: + - Freeing memory when models are no longer needed + - Forcing a fresh model reload + - Testing and debugging + Examples: + >>> # Clear cache to free memory + >>> DefaultLocalSparseEmbedding.clear_cache() + + >>> # Or in tests to ensure fresh model loading + >>> def test_something(): + ... DefaultLocalSparseEmbedding.clear_cache() + ... emb = DefaultLocalSparseEmbedding() + ... # Test with fresh model + """ + cls._model_cache.clear() + + @classmethod + def get_cache_info(cls) -> dict: + """Get information about currently cached models. + + Returns: + dict: Dictionary with cache statistics: + - cached_models (int): Number of cached model instances + - cache_keys (list): List of cache keys (model_name, model_source, device) + + Examples: + >>> info = DefaultLocalSparseEmbedding.get_cache_info() + >>> print(f"Cached models: {info['cached_models']}") + Cached models: 2 + >>> print(f"Cache keys: {info['cache_keys']}") + Cache keys: [('naver/splade-cocondenser-ensembledistil', 'huggingface', None), + ('naver/splade-cocondenser-ensembledistil', 'huggingface', 'cuda')] + """ + return { + "cached_models": len(cls._model_cache), + "cache_keys": list(cls._model_cache.keys()), + } + + @classmethod + def remove_from_cache( + cls, model_source: str = "huggingface", device: Optional[str] = None + ) -> bool: + """Remove a specific model from cache. + + Args: + model_source (str): Model source ("huggingface" or "modelscope"). + Defaults to "huggingface". + device (Optional[str]): Device identifier. Defaults to None. + + Returns: + bool: True if model was found and removed, False otherwise. + + Examples: + >>> # Remove CPU model from cache + >>> removed = DefaultLocalSparseEmbedding.remove_from_cache() + >>> print(f"Removed: {removed}") + True + + >>> # Remove CUDA model from cache + >>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device="cuda") + >>> print(f"Removed: {removed}") + True + """ + model_name = "naver/splade-cocondenser-ensembledistil" + cache_key = (model_name, model_source, device) + + if cache_key in cls._model_cache: + del cls._model_cache[cache_key] + return True + return False + + def __init__( + self, + model_source: Literal["huggingface", "modelscope"] = "huggingface", + device: Optional[str] = None, + encoding_type: Literal["query", "document"] = "query", + **kwargs, + ): + """Initialize with SPLADE model. + + Args: + model_source (Literal["huggingface", "modelscope"]): Model source. + Defaults to "huggingface". + device (Optional[str]): Target device ("cpu", "cuda", "mps", or None). + Defaults to None (automatic detection). + encoding_type (Literal["query", "document"]): Encoding type for embeddings. + - "query": Optimize for search queries (default) + - "document": Optimize for indexed documents + This distinction is important for asymmetric retrieval tasks. + **kwargs: Additional parameters (reserved for future use). + + Raises: + ImportError: If sentence-transformers is not installed. + ValueError: If model cannot be loaded. + + Note: + Multiple instances with the same (model_source, device) configuration + will share the same underlying model to save memory. Different + instances can use different encoding_type settings while sharing + the model. + + **Model Selection:** + + Uses ``naver/splade-cocondenser-ensembledistil`` instead of the newer + ``naver/splade-v3`` because splade-v3 is a gated model requiring + Hugging Face authentication. The cocondenser-ensembledistil variant: + + - Does not require authentication or API tokens + - Is immediately available for all users + - Provides comparable retrieval performance (~2% difference) + - Avoids "Access to model is restricted" errors + + If you need splade-v3 and have obtained access, you can subclass + this class and override the model_name parameter. + + Examples: + >>> # Both instances share the same model (saves memory) + >>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query") + >>> doc_emb = DefaultLocalSparseEmbedding(encoding_type="document") + >>> # Only one model is loaded in memory + """ + # Use publicly available SPLADE model (no gated access required) + # Note: naver/splade-v3 requires authentication, so we use the + # cocondenser-ensembledistil variant which is publicly accessible + model_name = "naver/splade-cocondenser-ensembledistil" + + # Initialize base class for model loading + SentenceTransformerFunctionBase.__init__( + self, model_name=model_name, model_source=model_source, device=device + ) + + self._encoding_type = encoding_type + self._extra_params = kwargs + + # Create cache key for this model configuration + self._cache_key = (model_name, model_source, device) + + # Load model to ensure it's available (will use cache if exists) + self._get_model() + + @property + def extra_params(self) -> dict: + """dict: Extra parameters for model-specific customization.""" + return self._extra_params + + def __call__(self, input: str) -> SparseVectorType: + """Make the embedding function callable.""" + return self.embed(input) + + def embed(self, input: str) -> SparseVectorType: + """Generate sparse embedding vector for the input text. + + This method uses the SPLADE model to convert input text into a sparse + vector representation. The result is a dictionary where keys are dimension + indices and values are importance weights (only non-zero values included). + + The embedding is optimized based on the ``encoding_type`` specified during + initialization: "query" for search queries or "document" for indexed content. + + Args: + input (str): Input text string to embed. Must be non-empty after + stripping whitespace. + + Returns: + SparseVectorType: A dictionary mapping dimension index to weight. + Only non-zero dimensions are included. The dictionary is sorted + by indices (keys) in ascending order for consistent output. + Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}`` + + Raises: + TypeError: If ``input`` is not a string. + ValueError: If input is empty or whitespace-only. + RuntimeError: If model inference fails. + + Examples: + >>> # Query embedding + >>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query") + >>> query_vec = query_emb.embed("machine learning") + >>> isinstance(query_vec, dict) + True + + Note: + - First call may be slower due to model loading + - Subsequent calls are much faster as the model stays in memory + - GPU acceleration provides significant speedup + - Sparse vectors are memory-efficient (only store non-zero values) + """ + if not isinstance(input, str): + raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}") + + input = input.strip() + if not input: + raise ValueError("Input text cannot be empty or whitespace only") + + try: + model = self._get_model() + + # Use appropriate encoding method based on type + if self._encoding_type == "document" and hasattr(model, "encode_document"): + # Use document encoding + sparse_matrix = model.encode_document([input]) + elif hasattr(model, "encode_query"): + # Use query encoding (default) + sparse_matrix = model.encode_query([input]) + else: + # Fallback: manual implementation for older sentence-transformers + return self._manual_sparse_encode(input) + + # Convert sparse matrix to dictionary + # SPLADE returns shape [1, vocab_size] for single input + + # Check if it's a sparse matrix (duck typing - has toarray method) + if hasattr(sparse_matrix, "toarray"): + # Sparse matrix (CSR/CSC/etc.) - convert to dense array + sparse_array = sparse_matrix[0].toarray().flatten() + sparse_dict = { + int(idx): float(val) + for idx, val in enumerate(sparse_array) + if val > 0 + } + else: + # Dense array format (numpy array or similar) + if isinstance(sparse_matrix, np.ndarray): + sparse_array = sparse_matrix[0] + else: + sparse_array = sparse_matrix + + sparse_dict = { + int(idx): float(val) + for idx, val in enumerate(sparse_array) + if val > 0 + } + + # Sort by indices (keys) to ensure consistent ordering + return dict(sorted(sparse_dict.items())) + + except Exception as e: + if isinstance(e, (TypeError, ValueError)): + raise + raise RuntimeError(f"Failed to generate sparse embedding: {e!s}") from e + + def _manual_sparse_encode(self, input: str) -> SparseVectorType: + """Fallback manual SPLADE encoding for older sentence-transformers. + + Args: + input (str): Input text to encode. + + Returns: + SparseVectorType: Sparse vector as dictionary. + """ + import torch + + model = self._get_model() + + # Tokenize input + features = model.tokenize([input]) + + # Move to correct device + features = {k: v.to(model.device) for k, v in features.items()} + + # Forward pass with no gradient + with torch.no_grad(): + embeddings = model.forward(features) + + # Get logits from model output + # SPLADE models typically output 'token_embeddings' + if isinstance(embeddings, dict) and "token_embeddings" in embeddings: + logits = embeddings["token_embeddings"][0] # First batch item + elif hasattr(embeddings, "token_embeddings"): + logits = embeddings.token_embeddings[0] + # Fallback: try to get first value + elif isinstance(embeddings, dict): + logits = next(iter(embeddings.values()))[0] + else: + logits = embeddings[0] + + # Apply SPLADE activation: log(1 + relu(x)) + relu_log = torch.log(1 + torch.relu(logits)) + + # Max pooling over token dimension (reduce to vocab size) + if relu_log.dim() > 1: + sparse_vec, _ = torch.max(relu_log, dim=0) + else: + sparse_vec = relu_log + + # Convert to sparse dictionary (only non-zero values) + sparse_vec_np = sparse_vec.cpu().numpy() + sparse_dict = { + int(idx): float(val) for idx, val in enumerate(sparse_vec_np) if val > 0 + } + + # Sort by indices (keys) to ensure consistent ordering + return dict(sorted(sparse_dict.items())) + + def _get_model(self): + """Load or retrieve the SPLADE model from class-level cache. + + Returns: + SentenceTransformer: The loaded SPLADE model instance. + + Raises: + ImportError: If required packages are not installed. + ValueError: If model cannot be loaded. + + Note: + Models are cached at class level and shared across all instances + with the same (model_name, model_source, device) configuration. + This allows memory-efficient usage when creating multiple instances + with different encoding_type settings. + """ + # Check class-level cache first + if self._cache_key in self._model_cache: + return self._model_cache[self._cache_key] + + # Use parent class method to load model + model = super()._get_model() + + # Cache the model at class level + self._model_cache[self._cache_key] = model + + return model diff --git a/python/zvec/extension/sentence_transformer_function.py b/python/zvec/extension/sentence_transformer_function.py new file mode 100644 index 0000000..1ba1662 --- /dev/null +++ b/python/zvec/extension/sentence_transformer_function.py @@ -0,0 +1,150 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import Literal, Optional + +from ..tool import require_module + + +class SentenceTransformerFunctionBase: + """Base class for Sentence Transformer functions (both dense and sparse). + + This base class provides common functionality for loading and managing + sentence-transformers models from Hugging Face or ModelScope. It supports + both dense models (e.g., all-MiniLM-L6-v2) and sparse models (e.g., SPLADE). + + This class is not meant to be used directly. Use concrete implementations: + - ``SentenceTransformerEmbeddingFunction`` for dense embeddings + - ``SentenceTransformerSparseEmbeddingFunction`` for sparse embeddings + - ``DefaultDenseEmbedding`` for default dense embeddings + - ``DefaultSparseEmbedding`` for default sparse embeddings + + Args: + model_name (str): Model identifier or local path. + model_source (Literal["huggingface", "modelscope"]): Model source. + device (Optional[str]): Device to run the model on. + + Note: + - This is an internal base class for code reuse + - Subclasses should inherit from appropriate Protocol (Dense/Sparse) + - Provides model loading and management functionality + """ + + def __init__( + self, + model_name: str, + model_source: Literal["huggingface", "modelscope"] = "huggingface", + device: Optional[str] = None, + ): + """Initialize the base Sentence Transformer functionality. + + Args: + model_name (str): Model identifier or local path. + model_source (Literal["huggingface", "modelscope"]): Model source. + device (Optional[str]): Device to run the model on. + + Raises: + ValueError: If model_source is invalid. + """ + # Validate model_source + if model_source not in ("huggingface", "modelscope"): + raise ValueError( + f"Invalid model_source: '{model_source}'. " + "Must be 'huggingface' or 'modelscope'." + ) + + self._model_name = model_name + self._model_source = model_source + self._device = device + self._model = None + + @property + def model_name(self) -> str: + """str: The Sentence Transformer model name currently in use.""" + return self._model_name + + @property + def model_source(self) -> str: + """str: The model source being used ("huggingface" or "modelscope").""" + return self._model_source + + @property + def device(self) -> str: + """str: The device the model is running on.""" + model = self._get_model() + if model is not None: + return str(model.device) + return self._device or "cpu" + + def _get_model(self): + """Load or retrieve the Sentence Transformer model. + + Returns: + SentenceTransformer or SparseEncoder: The loaded model instance. + + Raises: + ImportError: If required packages are not installed. + ValueError: If model cannot be loaded. + """ + # Return cached model if exists + if self._model is not None: + return self._model + + # Load model + try: + sentence_transformers = require_module("sentence_transformers") + + if self._model_source == "modelscope": + # Load from ModelScope + require_module("modelscope") + from modelscope.hub.snapshot_download import snapshot_download + + # Download model to cache + model_dir = snapshot_download(self._model_name) + + # Load from local path + self._model = sentence_transformers.SentenceTransformer( + model_dir, device=self._device, trust_remote_code=True + ) + else: + # Load from Hugging Face (default) + self._model = sentence_transformers.SentenceTransformer( + self._model_name, device=self._device, trust_remote_code=True + ) + + return self._model + + except ImportError as e: + if "modelscope" in str(e) and self._model_source == "modelscope": + raise ImportError( + "ModelScope support requires the 'modelscope' package. " + "Please install it with: pip install modelscope" + ) from e + raise + except Exception as e: + raise ValueError( + f"Failed to load Sentence Transformer model '{self._model_name}' " + f"from {self._model_source}: {e!s}" + ) from e + + def _is_sparse_model(self) -> bool: + """Check if the loaded model is a sparse encoder (e.g., SPLADE). + + Returns: + bool: True if model supports sparse encoding. + """ + model = self._get_model() + # Check if model has sparse encoding methods + return hasattr(model, "encode_query") or hasattr(model, "encode_document") diff --git a/python/zvec/extension/sentence_transformer_rerank_function.py b/python/zvec/extension/sentence_transformer_rerank_function.py new file mode 100644 index 0000000..bc84242 --- /dev/null +++ b/python/zvec/extension/sentence_transformer_rerank_function.py @@ -0,0 +1,396 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal, Optional + +from ..model.doc import Doc, DocList +from ..tool import require_module +from .rerank_function import RerankFunction +from .sentence_transformer_function import SentenceTransformerFunctionBase + +if TYPE_CHECKING: + from ..model.schema import FieldSchema, VectorSchema + + +class DefaultLocalReRanker(SentenceTransformerFunctionBase, RerankFunction): + """Re-ranker using Sentence Transformer cross-encoder models for semantic re-ranking. + + This re-ranker leverages pre-trained cross-encoder models to perform deep semantic + re-ranking of search results. It runs locally without API calls, supports GPU + acceleration, and works with models from Hugging Face or ModelScope. + + Cross-encoder models evaluate query-document pairs jointly, providing more + accurate relevance scores than bi-encoder (embedding-based) similarity. + + Args: + query (str): Query text for semantic re-ranking. **Required**. + rerank_field (Optional[str], optional): Document field name to use as + re-ranking input text. **Required** (e.g., "content", "title", "body"). + model_name (str, optional): Cross-encoder model identifier or local path. + Defaults to ``"cross-encoder/ms-marco-MiniLM-L6-v2"`` (MS MARCO MiniLM). + Common options: + - ``"cross-encoder/ms-marco-MiniLM-L6-v2"``: Lightweight, fast (~80MB, recommended) + - ``"cross-encoder/ms-marco-MiniLM-L12-v2"``: Better accuracy (~120MB) + - ``"BAAI/bge-reranker-base"``: BGE Reranker Base (~280MB) + - ``"BAAI/bge-reranker-large"``: BGE Reranker Large (highest quality, ~560MB) + model_source (Literal["huggingface", "modelscope"], optional): Model source. + Defaults to ``"huggingface"``. + - ``"huggingface"``: Load from Hugging Face Hub + - ``"modelscope"``: Load from ModelScope (recommended for users in China) + device (Optional[str], optional): Device to run the model on. + Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None`` + for automatic detection. Defaults to ``None``. + batch_size (int, optional): Batch size for processing query-document pairs. + Larger values speed up processing but use more memory. Defaults to ``32``. + + Attributes: + query (str): The query text used for re-ranking. + rerank_field (Optional[str]): Field name used for re-ranking input. + model_name (str): The cross-encoder model being used. + model_source (str): The model source ("huggingface" or "modelscope"). + device (str): The device the model is running on. + + Raises: + ValueError: If ``query`` is empty/None, ``rerank_field`` is None, + or model cannot be loaded. + TypeError: If input types are invalid. + RuntimeError: If model inference fails. + + Note: + - Requires Python 3.10, 3.11, or 3.12 + - Requires ``sentence-transformers`` package: ``pip install sentence-transformers`` + - For ModelScope support, also requires: ``pip install modelscope`` + - First run downloads the model (~80-560MB depending on model) from chosen source + - No API keys or network required after initial download + - Cross-encoders are slower than bi-encoders but more accurate + - GPU acceleration provides significant speedup (5-10x) + + **MS MARCO MiniLM-L6-v2 Model (Default):** + + The default model ``cross-encoder/ms-marco-MiniLM-L6-v2`` is a lightweight and + efficient cross-encoder trained on MS MARCO dataset. It provides: + + - Fast inference speed (suitable for real-time applications) + - Small model size (~80MB, quick to download) + - Good balance between speed and accuracy + - Trained on 500K+ query-document pairs + - Public availability without authentication + + **For users in China:** + + If you encounter Hugging Face access issues, use ModelScope instead: + + .. code-block:: python + + # Recommended for users in China + reranker = SentenceTransformerReRanker( + query="机器学习算法", + rerank_field="content", + model_source="modelscope" + ) + + Alternatively, use Hugging Face mirror: + + .. code-block:: bash + + export HF_ENDPOINT=https://hf-mirror.com + + Examples: + >>> # Basic usage with default MS MARCO MiniLM model + >>> from zvec.extension import SentenceTransformerReRanker + >>> + >>> reranker = SentenceTransformerReRanker( + ... query="machine learning algorithms", + ... rerank_field="content" + ... ) + >>> + >>> # Use in collection.query() + >>> results = collection.query( + ... data={"vector_field": query_vector}, + ... reranker=reranker, + ... topk=20 + ... ) + + >>> # Using ModelScope for users in China + >>> reranker = SentenceTransformerReRanker( + ... query="深度学习", + ... rerank_field="content", + ... model_source="modelscope" + ... ) + + >>> # Using larger model for better quality + >>> reranker = SentenceTransformerReRanker( + ... query="neural networks", + ... rerank_field="content", + ... model_name="BAAI/bge-reranker-large", + ... device="cuda", + ... batch_size=64 + ... ) + + >>> # Direct rerank call (for testing) + >>> query_results = { + ... "vector1": [ + ... Doc(id="1", score=0.9, fields={"content": "Machine learning is..."}), + ... Doc(id="2", score=0.8, fields={"content": "Deep learning is..."}), + ... ] + ... } + >>> reranked = reranker.rerank(query_results) + >>> for doc in reranked: + ... print(f"ID: {doc.id}, Score: {doc.score:.4f}") + ID: 2, Score: 0.9234 + ID: 1, Score: 0.8567 + + See Also: + - ``RerankFunction``: Abstract base class for re-rankers + - ``QwenReRanker``: Re-ranker using Qwen API + - ``RrfReRanker``: Multi-vector re-ranker using RRF + - ``WeightedReRanker``: Multi-vector re-ranker using weighted scores + + References: + - MS MARCO Cross-Encoder: https://huggingface.co/cross-encoder/ms-marco-MiniLM-L6-v2 + - BGE Reranker: https://huggingface.co/BAAI/bge-reranker-base + - Cross-Encoder vs Bi-Encoder: https://www.sbert.net/examples/applications/cross-encoder/README.html + """ + + def __init__( + self, + query: Optional[str] = None, + rerank_field: Optional[str] = None, + model_name: str = "cross-encoder/ms-marco-MiniLM-L6-v2", + model_source: Literal["huggingface", "modelscope"] = "huggingface", + device: Optional[str] = None, + batch_size: int = 32, + ): + """Initialize SentenceTransformerReRanker with query and configuration. + + Args: + query (Optional[str]): Query text for semantic matching. Required. + rerank_field (Optional[str]): Document field for re-ranking input. + model_name (str): Cross-encoder model identifier. + model_source (Literal["huggingface", "modelscope"]): Model source. + device (Optional[str]): Target device ("cpu", "cuda", "mps", or None). + batch_size (int): Batch size for processing query-document pairs. + + Raises: + ValueError: If query is empty or model cannot be loaded. + """ + # Initialize base class for model loading + SentenceTransformerFunctionBase.__init__( + self, model_name=model_name, model_source=model_source, device=device + ) + + # Initialize rerank parameters + self._rerank_field = rerank_field + + # Validate query + if not query: + raise ValueError("Query is required for DefaultLocalReRanker") + self._query = query + self._batch_size = batch_size + + # Load and validate cross-encoder model + model = self._get_model() + if not hasattr(model, "predict"): + raise ValueError( + f"Model '{model_name}' does not appear to be a cross-encoder model. " + "Cross-encoder models should have a 'predict' method." + ) + self._model = model + + def _get_model(self): + """Load or retrieve the CrossEncoder model. + + This overrides the base class method to load CrossEncoder instead of + SentenceTransformer, as reranking requires cross-encoder models. + + Returns: + CrossEncoder: The loaded cross-encoder model instance. + + Raises: + ImportError: If required packages are not installed. + ValueError: If model cannot be loaded. + """ + # Return cached model if exists + if self._model is not None: + return self._model + + # Load cross-encoder model + try: + sentence_transformers = require_module("sentence_transformers") + + if self._model_source == "modelscope": + # Load from ModelScope + require_module("modelscope") + from modelscope.hub.snapshot_download import snapshot_download + + # Download model to cache + model_dir = snapshot_download(self._model_name) + + # Load CrossEncoder from local path + model = sentence_transformers.CrossEncoder( + model_dir, device=self._device + ) + else: + # Load CrossEncoder from Hugging Face (default) + model = sentence_transformers.CrossEncoder( + self._model_name, device=self._device + ) + + return model + + except ImportError as e: + if "modelscope" in str(e) and self._model_source == "modelscope": + raise ImportError( + "ModelScope support requires the 'modelscope' package. " + "Please install it with: pip install modelscope" + ) from e + raise + except Exception as e: + raise ValueError( + f"Failed to load CrossEncoder model '{self._model_name}' " + f"from {self._model_source}: {e!s}" + ) from e + + @property + def rerank_field(self) -> Optional[str]: + """Optional[str]: Field name used as re-ranking input.""" + return self._rerank_field + + @property + def query(self) -> str: + """str: Query text used for semantic re-ranking.""" + return self._query + + @property + def batch_size(self) -> int: + """int: Batch size for processing query-document pairs.""" + return self._batch_size + + def rerank( + self, + query_results: list[list[Doc]], + topn: int = 10, + *, + fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002 + ) -> DocList: + """Re-rank documents using Sentence Transformer cross-encoder model. + + Evaluates each query-document pair using the cross-encoder model to compute + relevance scores. Documents are then sorted by these scores and the top-k + results are returned. + + Args: + query_results (list[list[Doc]]): Per-sub-query lists of retrieved + documents. Documents from all lists are deduplicated and + re-ranked together. + topn (int): Maximum number of documents to return. + fields: Unused; present for interface compatibility. + + Returns: + list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score`` + fields containing relevance scores from the cross-encoder model. + + Raises: + ValueError: If no valid documents are found or model inference fails. + + Note: + - Duplicate documents (same ID) across fields are processed once + - Documents with empty/missing ``rerank_field`` content are skipped + - Returned scores are logits from the cross-encoder model + - Higher scores indicate higher relevance + - Processing time is O(n) where n is the number of documents + + Examples: + >>> reranker = SentenceTransformerReRanker( + ... query="machine learning", + ... topn=3, + ... rerank_field="content" + ... ) + >>> query_results = { + ... "vector1": [ + ... Doc(id="1", score=0.9, fields={"content": "ML basics"}), + ... Doc(id="2", score=0.8, fields={"content": "DL tutorial"}), + ... ] + ... } + >>> reranked = reranker.rerank(query_results) + >>> len(reranked) <= 3 + True + """ + if not query_results: + return [] + + # Accept both dict (legacy) and list formats + if isinstance(query_results, dict): + query_results = list(query_results.values()) + + # Collect and deduplicate documents + id_to_doc: dict[str, Doc] = {} + doc_ids: list[str] = [] + contents: list[str] = [] + + for query_result in query_results: + for doc in query_result: + doc_id = doc.id + if doc_id in id_to_doc: + continue + + # Extract text content from specified field + field_value = doc.field(self.rerank_field) + rank_content = str(field_value).strip() if field_value else "" + if not rank_content: + continue + + id_to_doc[doc_id] = doc + doc_ids.append(doc_id) + contents.append(rank_content) + + if not contents: + raise ValueError("No documents to rerank") + + try: + # Use standard cross-encoder predict method + pairs = [[self.query, content] for content in contents] + scores = self._model.predict( + pairs, + batch_size=self.batch_size, + show_progress_bar=False, + convert_to_numpy=True, + ) + + # Convert to float list if needed + if hasattr(scores, "tolist"): + scores = scores.tolist() + else: + scores = [float(s) for s in scores] + + except Exception as e: + raise RuntimeError(f"Failed to compute rerank scores: {e!s}") from e + + # Create scored documents + scored_docs = [ + (doc_ids[i], id_to_doc[doc_ids[i]], scores[i]) for i in range(len(doc_ids)) + ] + + # Sort by score (descending) and take top-k + scored_docs.sort(key=lambda x: x[2], reverse=True) + top_scored_docs = scored_docs[:topn] + + # Build result list with updated scores + results: DocList = [] + for _, doc, score in top_scored_docs: + new_doc = doc._replace(score=score) + results.append(new_doc) + + return results diff --git a/python/zvec/model/__init__.py b/python/zvec/model/__init__.py new file mode 100644 index 0000000..7d5b068 --- /dev/null +++ b/python/zvec/model/__init__.py @@ -0,0 +1,30 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from .collection import Collection +from .doc import Doc +from .param.query import Fts, Query, VectorQuery +from .schema.collection_schema import CollectionSchema +from .schema.field_schema import FieldSchema + +__all__ = [ + "Collection", + "CollectionSchema", + "Doc", + "FieldSchema", + "Fts", + "Query", + "VectorQuery", +] diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py new file mode 100644 index 0000000..b90caae --- /dev/null +++ b/python/zvec/model/collection.py @@ -0,0 +1,439 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import warnings +from typing import Optional, Union, overload + +from zvec._zvec import _Collection + +from ..executor import QueryContext, QueryExecutor +from ..extension import ReRanker +from ..typing import Status +from .convert import convert_to_cpp_doc, convert_to_py_doc +from .doc import Doc, DocList +from .param import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + FlatIndexParam, + FtsIndexParam, + HnswIndexParam, + HnswRabitqIndexParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + OptimizeOption, +) +from .param.query import Query +from .schema import CollectionSchema, CollectionStats, FieldSchema + +__all__ = ["Collection"] + + +class Collection: + """Represents an opened collection in Zvec. + + A `Collection` provides methods for data definition (DDL), data manipulation (DML), + and querying (DQL). It is obtained via `create_and_open()` or `open()`. + + This class is not meant to be instantiated directly; use factory functions instead. + """ + + def __init__(self, obj: _Collection): + self._obj = obj + self._schema = None + self._querier = None + + @classmethod + def _from_core(cls, core_collection: _Collection) -> Collection: + if not core_collection: + raise ValueError("Collection is None") + inst = cls.__new__(cls) + inst._obj = core_collection + schema = CollectionSchema._from_core(core_collection.Schema()) + inst._schema = schema + inst._querier = QueryExecutor(schema) + return inst + + @property + def path(self) -> str: + """str: The filesystem path of the collection.""" + return self._obj.Path() + + @property + def option(self) -> CollectionOption: + """CollectionOption: The options used to open the collection.""" + return self._obj.Options() + + @property + def schema(self) -> CollectionSchema: + """CollectionSchema: The schema defining the structure of the collection.""" + return self._schema + + @property + def stats(self) -> CollectionStats: + """CollectionStats: Runtime statistics about the collection (e.g., doc count, size).""" + return self._obj.Stats() + + # ========== Collection DDL Methods ========== + def destroy(self) -> None: + """Permanently delete the collection from disk. + + Warning: + This operation is irreversible. All data will be lost. + """ + self._obj.Destroy() + + def flush(self) -> None: + """Force all pending writes to disk. + + Ensures durability of recent inserts/updates. + """ + self._obj.Flush() + + # ========== Index DDL Methods ========== + def create_index( + self, + field_name: str, + index_param: Union[ + HnswIndexParam, + HnswRabitqIndexParam, + IVFIndexParam, + FlatIndexParam, + InvertIndexParam, + FtsIndexParam, + ], + option: IndexOption = IndexOption(), + ) -> None: + """Create an index on a field. + + Vector index types (HNSW, IVF, FLAT) can only be applied to vector fields. + Inverted index (`InvertIndexParam`) is for scalar fields. + FTS index (`FtsIndexParam`) is for full-text search on STRING fields. + + Args: + field_name (str): Name of the field to index. + index_param (Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam, FtsIndexParam]): + Index configuration. + option (Optional[IndexOption], optional): Index creation options. + Defaults to ``IndexOption()``. + + """ + self._obj.CreateIndex(field_name, index_param, option) + self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema + + def drop_index(self, field_name: str) -> None: + """Remove the index from a field. + + Args: + field_name (str): Name of the indexed field. + """ + self._obj.DropIndex(field_name) + self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema + + def optimize(self, option: OptimizeOption = OptimizeOption()) -> None: + """Optimize the collection (e.g., merge segments, rebuild index). + + Args: + option (Optional[OptimizeOption], optional): Optimization options. + Defaults to ``OptimizeOption()``. + """ + self._obj.Optimize(option) + + # ========== COLUMN DDL Methods ========== + def add_column( + self, + field_schema: FieldSchema, + expression: str = "", + option: AddColumnOption = AddColumnOption(), + ) -> None: + """Add a new column to the collection. + + The column is populated using the provided expression (e.g., SQL-like formula). + + Args: + field_schema (FieldSchema): Schema definition for the new column. + expression (str): Expression to compute values for existing documents. + option (Optional[AddColumnOption], optional): Options for the operation. + Defaults to ``AddColumnOption()``. + """ + self._obj.AddColumn(field_schema._get_object(), expression, option) + self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema + + def drop_column(self, field_name: str) -> None: + """Remove a column from the collection. + + Args: + field_name (str): Name of the column to drop. + """ + self._obj.DropColumn(field_name) + self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema + + def alter_column( + self, + old_name: str, + new_name: Optional[str] = None, + field_schema: Optional[FieldSchema] = None, + option: AlterColumnOption = AlterColumnOption(), + ) -> None: + """Rename a column, update its schema. + + This method supports three atomic operations: + 1. Rename only (when `field_schema` is None). + 2. Modify schema only (when `new_name` is None or empty string). + + Args: + old_name (str): The current name of the column to be altered. + new_name (Optional[str]): The new name for the column. + - If provided and non-empty, the column will be renamed. + - If `None` or empty string, no rename occurs. + field_schema (Optional[FieldSchema]): The new schema definition. + - If provided, the column's type, dimension, or other properties will be updated. + - If `None`, only renaming (if requested) is performed. + option (AlterColumnOption, optional): Options controlling the alteration behavior. + Defaults to ``AlterColumnOption()``. + + **Limitation**: This operation **only supports scalar numeric columns**. such as: + - `DOUBLE`, `FLOAT`, + - `INT32`, `INT64`, `UINT32`, `UINT64` + + Note: + - Schema modification may trigger data migration or index rebuild. + + Examples: + >>> # Rename column only + >>> results = collection.alter_column(old_name="id", new_name="doc_id") + + >>> # Modify schema only + >>> new_schema = FieldSchema(name="doc_id", dtype=DataType.INT64) + >>> collection.alter_column("id", field_schema=new_schema) + """ + self._obj.AlterColumn( + old_name, + new_name or "", + field_schema._get_object() if field_schema else None, + option, + ) + self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema + + # ========== Collection DDL Methods ========== + @overload + def insert(self, docs: Doc) -> Status: + pass + + @overload + def insert(self, docs: list[Doc]) -> list[Status]: + pass + + def insert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]: + """Insert new documents into the collection. + + Documents must have unique IDs and conform to the schema. + + Args: + docs (Union[Doc, list[Doc]]): One or more documents to insert. + + Returns: + Union[Status, list[Status]]: If a single Doc was given, returns its Status; + if a list was given, returns a list of Status objects. + """ + is_single = isinstance(docs, Doc) + doc_list = [docs] if is_single else docs + results = self._obj.Insert( + [convert_to_cpp_doc(doc, self.schema) for doc in doc_list] + ) + return results[0] if is_single else results + + @overload + def upsert(self, docs: Doc) -> Status: + pass + + @overload + def upsert(self, docs: list[Doc]) -> list[Status]: + pass + + def upsert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]: + """Insert new documents or update existing ones by ID. + + Args: + docs (Union[Doc, list[Doc]]): Documents to upsert. + + Returns: + Union[Status, list[Status]]: If a single Doc was given, returns its Status; + if a list was given, returns a list of Status objects. + """ + is_single = isinstance(docs, Doc) + doc_list = [docs] if is_single else docs + results = self._obj.Upsert( + [convert_to_cpp_doc(doc, self.schema) for doc in doc_list] + ) + return results[0] if is_single else results + + @overload + def update(self, docs: Doc) -> Status: + pass + + @overload + def update(self, docs: list[Doc]) -> list[Status]: + pass + + def update(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]: + """Update existing documents by ID. + + Only specified fields are updated; others remain unchanged. + + Args: + docs (Union[Doc, list[Doc]]): Documents containing updated fields. + + Returns: + Union[Status, list[Status]]: If a single Doc was given, returns its Status; + if a list was given, returns a list of Status objects. + """ + is_single = isinstance(docs, Doc) + doc_list = [docs] if is_single else docs + results = self._obj.Update( + [convert_to_cpp_doc(doc, self.schema) for doc in doc_list] + ) + return results[0] if is_single else results + + @overload + def delete(self, ids: str) -> Status: + pass + + @overload + def delete(self, ids: list[str]) -> list[Status]: + pass + + def delete(self, ids: Union[str, list[str]]) -> Union[Status, list[Status]]: + """Delete documents by ID. + + Args: + ids (Union[str, list[str]]): One or more document IDs to delete. + + Returns: + Union[Status, list[Status]]: If a single id was given, returns its Status; + if a list was given, returns a list of Status objects. + """ + is_single = isinstance(ids, str) + id_list = [ids] if isinstance(ids, str) else ids + results = self._obj.Delete(id_list) + return results[0] if is_single else results + + def delete_by_filter(self, filter: str) -> None: + """Delete documents matching a filter expression. + + Args: + filter (str): Boolean expression (e.g., ``"age > 30"``). + """ + self._obj.DeleteByFilter(filter) + + # ========== Collection DQL-fetch Methods ========== + def fetch( + self, + ids: Union[str, list[str]], + *, + output_fields: Optional[list[str]] = None, + include_vector: bool = True, + ) -> dict[str, Doc]: + """Retrieve documents by ID. + + Args: + ids (Union[str, list[str]]): Document IDs to fetch. + output_fields (Optional[list[str]], optional): Scalar fields to + include. If None, all fields are returned. Defaults to None. + include_vector (bool, optional): Whether to include vector data in + results. Defaults to True. + + Returns: + dict[str, Doc]: Mapping from ID to document. Missing IDs are omitted. + """ + ids = [ids] if isinstance(ids, str) else ids + docs = self._obj.Fetch(ids, output_fields, include_vector) + return { + doc_id: py_doc + for doc_id, core_doc in docs.items() + if (py_doc := convert_to_py_doc(core_doc, self.schema)) is not None + } + + # ========== Collection DQL-Query Methods ========== + + def query( + self, + queries: Optional[Union[Query, list[Query]]] = None, + *, + vectors: Optional[Union[Query, list[Query]]] = None, + topk: int = 10, + filter: Optional[str] = None, + include_vector: bool = False, + output_fields: Optional[list[str]] = None, + reranker: Optional[ReRanker] = None, + ) -> DocList: + """Perform vector similarity search with optional filtering and re-ranking. + + At least one `Query` must be provided via `queries`. + + Args: + queries (Optional[Union[Query, list[Query]]], optional): + One or more vector queries. Defaults to None. + vectors (Optional[Union[Query, list[Query]]], optional): + Deprecated. Use `queries` instead. + topk (int, optional): Number of nearest neighbors to return. + Defaults to 10. + filter (Optional[str], optional): Boolean expression to pre-filter candidates. + Defaults to None. + include_vector (bool, optional): Whether to include vector data in results. + Defaults to False. + output_fields (Optional[list[str]], optional): Scalar fields to include. + If None, all fields are returned. Defaults to None. + reranker (Optional[ReRanker], optional): Re-ranker to refine results. + Defaults to None. + + Returns: + DocList: Top-k matching documents, sorted by relevance score. + + Examples: + >>> from zvec import Query + >>> results = collection.query( + ... queries=Query(field_name="embedding", vector=[0.1, 0.2]), + ... topk=5, + ... filter="category == 'tech'", + ... output_fields=["title", "url"] + ... ) + """ + if vectors is not None: + warnings.warn( + "The 'vectors' parameter is deprecated and will be removed in a future version. " + "Use 'queries' instead.", + DeprecationWarning, + stacklevel=2, + ) + if queries is not None: + raise ValueError("Cannot specify both 'queries' and 'vectors'.") + queries = vectors + + ctx = QueryContext( + topk=topk, + filter=filter, + queries=[queries] if isinstance(queries, Query) else queries, + include_vector=include_vector, + output_fields=output_fields, + reranker=reranker, + ) + return self._querier.execute(ctx, self._obj) diff --git a/python/zvec/model/convert.py b/python/zvec/model/convert.py new file mode 100644 index 0000000..421bd17 --- /dev/null +++ b/python/zvec/model/convert.py @@ -0,0 +1,54 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from zvec._zvec import _Doc + +from .doc import Doc +from .schema import CollectionSchema + + +def convert_to_cpp_doc(doc: Doc, collection_schema: CollectionSchema) -> _Doc: + if not doc or not collection_schema: + return None + + _doc = _Doc() + + # set pk + _doc.set_pk(doc.id) + + # set scalar fields + for k, v in doc.fields.items(): + field_schema = collection_schema.field(k) + if not field_schema: + raise ValueError( + f"schema validate failed: {k} not found in collection schema" + ) + _doc.set_any(k, field_schema._get_object(), v) + + # set vector fields + for k, v in doc.vectors.items(): + vector_schema = collection_schema.vector(k) + if not vector_schema: + raise ValueError( + f"schema validate failed: {k} not found in collection schema" + ) + _doc.set_any(k, vector_schema._get_object(), v) + return _doc + + +def convert_to_py_doc(doc: _Doc, collection_schema: CollectionSchema) -> Doc: + if not doc or not collection_schema: + return None + + data_tuple = doc.get_all(collection_schema._get_object()) + return Doc._from_tuple(data_tuple) diff --git a/python/zvec/model/doc.py b/python/zvec/model/doc.py new file mode 100644 index 0000000..175c946 --- /dev/null +++ b/python/zvec/model/doc.py @@ -0,0 +1,178 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import json +from typing import Any, Optional + +from ..common import VectorType + +__all__ = [ + "Doc", + "DocList", +] + + +class Doc: + """Represents a retrieved document with optional metadata, fields, and vectors. + + This immutable data class encapsulates the result of a search or retrieval + operation. It includes the document ID, relevance score (if applicable), + scalar fields, and vector embeddings. + + During initialization, any `numpy.ndarray` in `vectors` is automatically + converted to a plain Python list for JSON serialization and immutability. + + Attributes: + id (str): Unique identifier of the document. + score (Optional[float], optional): Relevance score from search. + Defaults to None. + vectors (Optional[dict[str, VectorType]], optional): Named vector + embeddings associated with the document. Values are converted to + lists if originally `np.ndarray`. Defaults to None. + fields (Optional[dict[str, Any]], optional): Scalar metadata fields + (e.g., title, timestamp). Defaults to None. + + Examples: + >>> import numpy as np + >>> import zvec + >>> doc = zvec.Doc( + ... id="doc1", + ... score=0.95, + ... vectors={"emb": np.array([0.1, 0.2, 0.3])}, + ... fields={"title": "Hello World"} + ... ) + >>> print(doc.vector("emb")) + [0.1, 0.2, 0.3] + >>> print(doc.has_field("title")) + True + """ + + __slots__ = ("id", "score", "vectors", "fields") + + def __init__( + self, + id: str, + score: Optional[float] = None, + vectors: Optional[dict[str, VectorType]] = None, + fields: Optional[dict[str, Any]] = None, + ): + self.id = id + self.score = score + self.vectors = vectors or {} + self.fields = fields or {} + + def has_field(self, name: str) -> bool: + """Check if the document contains a scalar field with the given name. + + Args: + name (str): Name of the field to check. + + Returns: + bool: True if the field exists, False otherwise. + """ + return name in self.fields + + def has_vector(self, name: str) -> bool: + """Check if the document contains a vector with the given name. + + Args: + name (str): Name of the vector to check. + + Returns: + bool: True if the vector exists, False otherwise. + """ + return name in self.vectors + + def vector(self, name: str): + """Get a vector by name. + + Args: + name (str): Name of the vector. + + Returns: + Any: The vector (as a list) if it exists, otherwise None. + """ + return self.vectors and self.vectors.get(name) + + def field(self, name: str): + """Get a scalar field by name. + + Args: + name (str): Name of the field. + + Returns: + Any: The field value if it exists, otherwise None. + """ + return self.fields and self.fields.get(name) + + def vector_names(self) -> list[str]: + """Get the list of all vector names in this document. + + Returns: + list[str]: A list of vector field names. Empty if no vectors. + """ + return [] if not self.vectors else list(self.vectors.keys()) + + def field_names(self) -> list[str]: + """Get the list of all scalar field names in this document. + + Returns: + list[str]: A list of field names. Empty if no fields. + """ + return [] if not self.fields else list(self.fields.keys()) + + def __repr__(self) -> str: + try: + schema = { + "id": self.id, + "score": self.score, + "fields": self.fields, + "vectors": self.vectors, + } + return json.dumps(schema, indent=2, ensure_ascii=False) + except Exception as e: + return f"" + + def _replace(self, **changes): + new_tuple = ( + changes.get("id", self.id), + changes.get("score", self.score), + changes.get("fields", self.fields.copy() if self.fields else None), + changes.get("vectors", self.vectors.copy() if self.vectors else None), + ) + return type(self)._from_tuple(new_tuple) + + @classmethod + def _from_tuple( + cls, data_tuple: tuple[str, float, dict[str, Any], dict[str, VectorType]] + ): + obj = object.__new__(cls) + obj.id = data_tuple[0] + obj.score = data_tuple[1] + obj.fields = data_tuple[2] or {} + + vectors = data_tuple[3] + if vectors is not None: + obj.vectors = { + name: (vec.tolist() if hasattr(vec, "tolist") else vec) + for name, vec in vectors.items() + } + else: + obj.vectors = {} + return obj + + +#: Type alias for query results: a list of documents returned by a single query route. +DocList = list[Doc] diff --git a/python/zvec/model/param/__init__.py b/python/zvec/model/param/__init__.py new file mode 100644 index 0000000..bcf0f1d --- /dev/null +++ b/python/zvec/model/param/__init__.py @@ -0,0 +1,60 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from zvec._zvec.param import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + DiskAnnIndexParam, + DiskAnnQueryParam, + FlatIndexParam, + FtsIndexParam, + FtsQueryParam, + HnswIndexParam, + HnswQueryParam, + HnswRabitqIndexParam, + HnswRabitqQueryParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + IVFQueryParam, + OptimizeOption, + QuantizerParam, + VamanaIndexParam, + VamanaQueryParam, +) + +__all__ = [ + "AddColumnOption", + "AlterColumnOption", + "CollectionOption", + "DiskAnnIndexParam", + "DiskAnnQueryParam", + "FlatIndexParam", + "FtsIndexParam", + "FtsQueryParam", + "HnswIndexParam", + "HnswQueryParam", + "HnswRabitqIndexParam", + "HnswRabitqQueryParam", + "IVFIndexParam", + "IVFQueryParam", + "IndexOption", + "InvertIndexParam", + "OptimizeOption", + "QuantizerParam", + "VamanaIndexParam", + "VamanaQueryParam", +] diff --git a/python/zvec/model/param/__init__.pyi b/python/zvec/model/param/__init__.pyi new file mode 100644 index 0000000..2653d2c --- /dev/null +++ b/python/zvec/model/param/__init__.pyi @@ -0,0 +1,1206 @@ +""" +This module contains the params of Zvec +""" + +from __future__ import annotations + +import collections +import typing + +import zvec._zvec.typing + +__all__: list[str] = [ + "AddColumnOption", + "AlterColumnOption", + "CollectionOption", + "FlatIndexParam", + "FtsIndexParam", + "FtsQueryParam", + "HnswIndexParam", + "HnswQueryParam", + "HnswRabitqIndexParam", + "HnswRabitqQueryParam", + "IVFIndexParam", + "IVFQueryParam", + "IndexOption", + "IndexParam", + "InvertIndexParam", + "OptimizeOption", + "QuantizerParam", + "QueryParam", + "SegmentOption", + "VectorIndexParam", +] + +class AddColumnOption: + """ + + Options for adding a new column to a collection. + + Attributes: + concurrency (int): Number of threads to use when backfilling data + for the new column. If 0, auto-detect is used. Default is 0. + + Examples: + >>> opt = AddColumnOption(concurrency=1) + >>> print(opt.concurrency) + 1 + """ + + def __getstate__(self) -> tuple: ... + def __init__(self, concurrency: typing.SupportsInt = 0) -> None: + """ + Constructs an AddColumnOption instance. + + Args: + concurrency (int, optional): Number of threads for data backfill. + 0 means auto-detect. Defaults to 0. + """ + + def __setstate__(self, arg0: tuple) -> None: ... + @property + def concurrency(self) -> int: + """ + int: Number of threads used when adding a column (0 = auto). + """ + +class AlterColumnOption: + """ + + Options for altering an existing column (e.g., changing index settings). + + Attributes: + concurrency (int): Number of threads to use during the alteration process. + If 0, the system will choose an optimal value automatically. + Default is 0. + + Examples: + >>> opt = AlterColumnOption(concurrency=1) + >>> print(opt.concurrency) + 1 + """ + + def __getstate__(self) -> tuple: ... + def __init__(self, concurrency: typing.SupportsInt = 0) -> None: + """ + Constructs an AlterColumnOption instance. + + Args: + concurrency (int, optional): Number of threads for column alteration. + 0 means auto-detect. Defaults to 0. + """ + + def __setstate__(self, arg0: tuple) -> None: ... + @property + def concurrency(self) -> int: + """ + int: Number of threads used when altering a column (0 = auto). + """ + +class CollectionOption: + """ + + Options for opening or creating a collection. + + Attributes: + read_only (bool): Whether the collection is opened in read-only mode. + Default is False. + enable_mmap (bool): Whether to use memory-mapped I/O for data files. + Default is True. + + Examples: + >>> opt = CollectionOption(read_only=True, enable_mmap=False) + >>> print(opt.read_only) + True + """ + + def __getstate__(self) -> tuple: ... + def __init__(self, read_only: bool = False, enable_mmap: bool = True) -> None: + """ + Constructs a CollectionOption instance. + + Args: + read_only (bool, optional): Open collection in read-only mode. + Defaults to False. + enable_mmap (bool, optional): Enable memory-mapped I/O. + Defaults to True. + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def enable_mmap(self) -> bool: ... + @property + def read_only(self) -> bool: ... + +class FlatIndexParam(VectorIndexParam): + """ + + Parameters for configuring a flat (brute-force) index. + + A flat index performs exact nearest neighbor search by comparing the query + vector against all vectors in the collection. It is simple, accurate, and + suitable for small to medium datasets or as a baseline. + + Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Use ``QuantizeType.UNDEFINED`` to + disable quantization. Default is ``QuantizeType.UNDEFINED``. + quantizer_param (QuantizerParam): Optional quantizer parameters. See + ``QuantizerParam`` for available options. Default is ``QuantizerParam()``. + + Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = FlatIndexParam( + ... metric_type=MetricType.L2, + ... quantize_type=QuantizeType.FP16 + ... ) + >>> print(params) + {'metric_type': 'L2', 'quantize_type': 'FP16'} + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + metric_type: zvec._zvec.typing.MetricType = ..., + quantize_type: zvec._zvec.typing.QuantizeType = ..., + quantizer_param: QuantizerParam = ..., + ) -> None: + """ + Constructs a FlatIndexParam instance. + + Args: + metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP. + quantize_type (QuantizeType, optional): Vector quantization type. + Defaults to QuantizeType.UNDEFINED (no quantization). + quantizer_param (QuantizerParam, optional): Quantizer configuration. + Defaults to QuantizerParam(). + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + +class HnswIndexParam(VectorIndexParam): + """ + + Parameters for configuring an HNSW (Hierarchical Navigable Small World) index. + + HNSW is a graph-based approximate nearest neighbor search index. This class + encapsulates its construction hyperparameters. + + Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + m (int): Number of bi-directional links created for every new element + during construction. Higher values improve accuracy but increase + memory usage and construction time. Default is 50. + ef_construction (int): Size of the dynamic candidate list for nearest + neighbors during index construction. Larger values yield better + graph quality at the cost of slower build time. Default is 500. + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to + disable quantization. + use_contiguous_memory (bool): If True, the HNSW streamer allocates a + single contiguous memory arena for all graph nodes, improving cache + locality and search throughput at the cost of peak memory usage. + Default is False. + + Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = HnswIndexParam( + ... metric_type=MetricType.COSINE, + ... m=16, + ... ef_construction=200, + ... quantize_type=QuantizeType.INT8, + ... use_contiguous_memory=True, + ... ) + >>> print(params) + {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8', 'use_contiguous_memory': True} + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + metric_type: zvec._zvec.typing.MetricType = ..., + m: typing.SupportsInt = 50, + ef_construction: typing.SupportsInt = 500, + quantize_type: zvec._zvec.typing.QuantizeType = ..., + use_contiguous_memory: bool = False, + quantizer_param: QuantizerParam = ..., + ) -> None: ... + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def ef_construction(self) -> int: + """ + int: Candidate list size during index construction. + """ + + @property + def m(self) -> int: + """ + int: Maximum number of neighbors per node in upper layers. + """ + + @property + def use_contiguous_memory(self) -> bool: + """ + bool: Whether to allocate a single contiguous memory arena for all + HNSW graph nodes. Improves cache locality and search throughput at + the cost of peak memory usage. Defaults to False. + """ + +class HnswQueryParam(QueryParam): + """ + + Query parameters for HNSW (Hierarchical Navigable Small World) index. + + Controls the trade-off between search speed and accuracy via the `ef` parameter. + + Attributes: + type (IndexType): Always ``IndexType.HNSW``. + ef (int): Size of the dynamic candidate list during search. + Larger values improve recall but slow down search. + Default is 300. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + prefetch_offset (int, optional): Graph prefetch offset (PO) used by the + HNSW fast path. ``0`` disables prefetching. Default is ``8``. + Values are clamped to ``256``. + prefetch_lines (int, optional): Number of 64B cache lines to prefetch + per neighbour vector (PL). ``0`` (default) uses the auto-derived + value ``ceil(vector_size/64)``. Values are clamped to ``256``. + + Examples: + >>> params = HnswQueryParam(ef=300) + >>> print(params.ef) + 300 + >>> print(params.to_dict() if hasattr(params, 'to_dict') else params) + {"type":"HNSW", "ef":300} + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + ef: typing.SupportsInt = 300, + radius: typing.SupportsFloat = 0.0, + is_linear: bool = False, + is_using_refiner: bool = False, + extra_params: dict[str, int] = ..., + ) -> None: + """ + Constructs an HnswQueryParam instance. + + Args: + ef (int, optional): Search-time candidate list size. + Higher values improve accuracy. Defaults to 300. + radius (float, optional): Search radius for range queries. Default is 0.0. + is_linear (bool, optional): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + extra_params (dict, optional): Additional search parameters. Supported keys: + - ``prefetch_offset`` (int): Graph prefetch offset (PO). + ``0`` disables prefetching. Default is ``8``. + - ``prefetch_lines`` (int): Number of 64B cache lines to prefetch + per neighbour vector (PL). ``0`` (default) means auto-derive from vector size. + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def ef(self) -> int: + """ + int: Size of the dynamic candidate list during HNSW search. + """ + + @property + def prefetch_offset(self) -> int: + """ + int: Graph prefetch offset used by the HNSW fast path. + """ + + @property + def prefetch_lines(self) -> int: + """ + int: Override of prefetch cache lines per vector (0=auto). + """ + +class HnswRabitqIndexParam(VectorIndexParam): + """ + + Parameters for configuring an HNSW (Hierarchical Navigable Small World) index with RabitQ quantization. + + HNSW is a graph-based approximate nearest neighbor search index. RabitQ is a + quantization method that provides high compression with minimal accuracy loss. + + Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + total_bits (int): Total bits for RabitQ quantization. Default is 7. + num_clusters (int): Number of clusters for RabitQ. Default is 16. + m (int): Number of bi-directional links created for every new element + during construction. Higher values improve accuracy but increase + memory usage and construction time. Default is 50. + ef_construction (int): Size of the dynamic candidate list for nearest + neighbors during index construction. Larger values yield better + graph quality at the cost of slower build time. Default is 500. + sample_count (int): Sample count for RabitQ training. Default is 0. + + Examples: + >>> from zvec.typing import MetricType + >>> params = HnswRabitqIndexParam( + ... metric_type=MetricType.COSINE, + ... total_bits=8, + ... num_clusters=256, + ... m=16, + ... ef_construction=200, + ... sample_count=10000 + ... ) + >>> print(params) + {'metric_type': 'COSINE', 'total_bits': 8, 'num_clusters': 256, 'm': 16, 'ef_construction': 200, 'sample_count': 10000} + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + metric_type: zvec._zvec.typing.MetricType = ..., + total_bits: typing.SupportsInt = 7, + num_clusters: typing.SupportsInt = 16, + m: typing.SupportsInt = 50, + ef_construction: typing.SupportsInt = 500, + sample_count: typing.SupportsInt = 0, + ) -> None: ... + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def ef_construction(self) -> int: + """ + int: Candidate list size during index construction. + """ + + @property + def m(self) -> int: + """ + int: Maximum number of neighbors per node. + """ + + @property + def total_bits(self) -> int: + """ + int: Total bits for RabitQ quantization. + """ + + @property + def num_clusters(self) -> int: + """ + int: Number of clusters for RabitQ. + """ + + @property + def sample_count(self) -> int: + """ + int: Sample count for RabitQ training. + """ + +class HnswRabitqQueryParam(QueryParam): + """ + + Query parameters for HNSW index with RabitQ quantization. + + Controls the trade-off between search speed and accuracy via the `ef` parameter. + + Attributes: + type (IndexType): Always ``IndexType.HNSW_RABITQ``. + ef (int): Size of the dynamic candidate list during search. + Larger values improve recall but slow down search. + Default is 300. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + + Examples: + >>> params = HnswRabitqQueryParam(ef=300) + >>> print(params.ef) + 300 + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + ef: typing.SupportsInt = 300, + radius: typing.SupportsFloat = 0.0, + is_linear: bool = False, + is_using_refiner: bool = False, + ) -> None: + """ + Constructs an HnswRabitqQueryParam instance. + + Args: + ef (int, optional): Search-time candidate list size. + Higher values improve accuracy. Defaults to 300. + radius (float, optional): Search radius for range queries. Default is 0.0. + is_linear (bool, optional): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def ef(self) -> int: + """ + int: Size of the dynamic candidate list during HNSW search. + """ + +class IVFIndexParam(VectorIndexParam): + """ + + Parameters for configuring an IVF (Inverted File Index) index. + + IVF partitions the vector space into clusters (inverted lists). At query time, + only a subset of clusters is searched, providing a trade-off between speed + and accuracy. + + Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + n_list (int): Number of clusters (inverted lists) to partition the dataset into. + Default is 10. + n_iters (int): Number of iterations for k-means clustering during index training. + Higher values yield more stable centroids. Default is 10. + use_soar (bool): Whether to enable SOAR (Scalable Optimized Adaptive Routing) + for improved IVF search performance. Default is False. + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED``. + + Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = IVFIndexParam( + ... metric_type=MetricType.COSINE, + ... n_list=100, + ... n_iters=15, + ... use_soar=True, + ... quantize_type=QuantizeType.INT8 + ... ) + >>> print(params.n_list) + 100 + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + metric_type: zvec._zvec.typing.MetricType = ..., + n_list: typing.SupportsInt = 10, + n_iters: typing.SupportsInt = 10, + use_soar: bool = False, + quantize_type: zvec._zvec.typing.QuantizeType = ..., + quantizer_param: QuantizerParam = ..., + ) -> None: + """ + Constructs an IVFIndexParam instance. + + Args: + metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP. + n_list (int, optional): Number of inverted lists (clusters). + Defaults to 10. + n_iters (int, optional): Number of k-means iterations during training. + Defaults to 10. + use_soar (bool, optional): Enable SOAR optimization. Defaults to False. + quantize_type (QuantizeType, optional): Vector quantization type. + Defaults to QuantizeType.UNDEFINED. + quantizer_param (QuantizerParam, optional): Quantizer configuration. + Defaults to QuantizerParam(). + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def n_iters(self) -> int: + """ + int: Number of k-means iterations during training. + """ + + @property + def n_list(self) -> int: + """ + int: Number of inverted lists. + """ + + @property + def use_soar(self) -> bool: + """ + bool: Whether SOAR optimization is enabled. + """ + +class IVFQueryParam(QueryParam): + """ + + Query parameters for IVF (Inverted File Index) index. + + Controls how many inverted lists (`nprobe`) to visit during search. + + Attributes: + type (IndexType): Always ``IndexType.IVF``. + nprobe (int): Number of closest clusters (inverted lists) to search. + Higher values improve recall but increase latency. + Default is 10. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + + Examples: + >>> params = IVFQueryParam(nprobe=20) + >>> print(params.nprobe) + 20 + """ + + def __getstate__(self) -> tuple: ... + def __init__(self, nprobe: typing.SupportsInt = 10) -> None: + """ + Constructs an IVFQueryParam instance. + + Args: + nprobe (int, optional): Number of inverted lists to probe during search. + Higher values improve accuracy. Defaults to 10. + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def nprobe(self) -> int: + """ + int: Number of inverted lists to search during IVF query. + """ + +class VamanaIndexParam(VectorIndexParam): + """ + Parameters for configuring a Vamana (DiskANN) index. + + Attributes: + metric_type (MetricType): Distance metric. Default is ``MetricType.IP``. + max_degree (int): Maximum out-degree (R) of every node. Default is 64. + search_list_size (int): Candidate list size during construction. Default is 100. + alpha (float): RobustPrune alpha factor. Default is 1.2. + saturate_graph (bool): Force every node to reach max_degree. Default is False. + use_contiguous_memory (bool): Allocate contiguous memory arena. Default is False. + use_id_map (bool): Reserved flag for id remapping. Default is False. + quantize_type (QuantizeType): Vector quantization type. Default is ``QuantizeType.UNDEFINED``. + + Examples: + >>> params = VamanaIndexParam(metric_type=MetricType.COSINE, max_degree=64) + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + metric_type: zvec._zvec.typing.MetricType = ..., + max_degree: typing.SupportsInt = 64, + search_list_size: typing.SupportsInt = 100, + alpha: typing.SupportsFloat = 1.2, + saturate_graph: bool = False, + use_contiguous_memory: bool = False, + use_id_map: bool = False, + quantize_type: zvec._zvec.typing.QuantizeType = ..., + ) -> None: ... + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: ... + @property + def max_degree(self) -> int: + """int: Maximum out-degree (R) of every node in the Vamana graph.""" + + @property + def search_list_size(self) -> int: + """int: Candidate list size during Vamana graph construction.""" + + @property + def alpha(self) -> float: + """float: Vamana RobustPrune alpha factor.""" + + @property + def saturate_graph(self) -> bool: + """bool: Whether to saturate every node to max_degree neighbors.""" + + @property + def use_contiguous_memory(self) -> bool: + """bool: Whether to allocate a single contiguous memory arena.""" + + @property + def use_id_map(self) -> bool: + """bool: Reserved flag for engine-level id remapping.""" + +class VamanaQueryParam(QueryParam): + """ + Query parameters for the Vamana (DiskANN) index. + + Attributes: + type (IndexType): Always ``IndexType.VAMANA``. + ef_search (int): Size of the dynamic candidate list during search. Default is 200. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + is_using_refiner (bool): Whether to use refiner. Default is False. + prefetch_offset (int): Graph prefetch offset (PO). Default is 8. + prefetch_lines (int): Cache lines to prefetch per vector (PL). Default is 0 (auto). + + Examples: + >>> params = VamanaQueryParam(ef_search=200) + >>> print(params.ef_search) + 200 + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + ef_search: typing.SupportsInt = 200, + radius: typing.SupportsFloat = 0.0, + is_linear: bool = False, + is_using_refiner: bool = False, + extra_params: dict[str, int] = ..., + ) -> None: + """ + Constructs a VamanaQueryParam instance. + + Args: + ef_search (int, optional): Search-time candidate list size. Defaults to 200. + radius (float, optional): Search radius for range queries. Default is 0.0. + is_linear (bool, optional): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner. Default is False. + extra_params (dict, optional): Additional search parameters. Supported keys: + - ``prefetch_offset`` (int): Graph prefetch offset (PO). + ``0`` disables prefetching. Default is ``8``. + - ``prefetch_lines`` (int): Cache lines to prefetch per vector (PL). + ``0`` (default) means auto-derive from vector size. + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def ef_search(self) -> int: + """int: Size of the dynamic candidate list during Vamana search.""" + + @property + def prefetch_offset(self) -> int: + """int: Graph prefetch offset used by the Vamana fast path.""" + + @property + def prefetch_lines(self) -> int: + """int: Override of prefetch cache lines per vector (0=auto).""" + +class FtsIndexParam(IndexParam): + """ + + Parameters for configuring a full-text search (FTS) index. + + Controls the tokenizer pipeline used during indexing and querying. + + Attributes: + type (IndexType): Always ``IndexType.FTS``. + tokenizer_name (str): Name of the tokenizer (one of "standard", "jieba", + "whitespace"). + Default is "standard". + filters (list[str]): List of token filter names applied after tokenization. + Supported filters are "lowercase", "ascii_folding", and "stemmer". + Default is ["lowercase"]. + extra_params (str): Additional tokenizer/filter parameters as an empty + string or JSON object string. Supported keys are grouped by component: + Tokenizers: + standard: + - "max_token_length" (positive integer). + jieba: + - "jieba_dict_dir" (directory containing jieba.dict.utf8 and + hmm_model.utf8). + - "user_dict_path" (user dictionary path). + - "cut_mode" ("search", "mix", "full", or "hmm"; default + "search"). + whitespace: + - no extra_params. + Filters: + lowercase: + - no extra_params. + ascii_folding: + - no extra_params. + stemmer: + - "stemmer_lang" (Snowball language/algorithm; default + "english"), for example {"stemmer_lang":"porter"} for ES + behaviour. + Default is "". + + Examples: + >>> params = FtsIndexParam( + ... tokenizer_name="jieba", filters=["lowercase", "ascii_folding"] + ... ) + >>> print(params.tokenizer_name) + jieba + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + tokenizer_name: str = "standard", + filters: list[str] = ..., + extra_params: str = "", + ) -> None: + """ + Constructs an FtsIndexParam instance. + + Args: + tokenizer_name (str, optional): Tokenizer name. Defaults to "standard". + filters (list[str], optional): Token filter names. Supports + "lowercase", "ascii_folding", and "stemmer". Defaults to + ["lowercase"]. + extra_params (str, optional): Extra tokenizer/filter parameters as an + empty string or JSON object string. Supported keys: + Tokenizers: + standard: + - "max_token_length" (positive integer). + jieba: + - "jieba_dict_dir". + - "user_dict_path". + - "cut_mode" ("search", "mix", "full", or "hmm"; + default "search"). + whitespace: + - no extra_params. + Filters: + lowercase: + - no extra_params. + ascii_folding: + - no extra_params. + stemmer: + - "stemmer_lang" (Snowball language/algorithm; default + "english"). + Defaults to "". + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def tokenizer_name(self) -> str: + """ + str: Name of the tokenizer. + """ + + @property + def filters(self) -> list[str]: + """ + list[str]: Token filter names. + """ + + @property + def extra_params(self) -> str: + """ + str: Additional tokenizer parameters. + """ + +class FtsQueryParam(QueryParam): + """ + + Query parameters for full-text search (FTS) index. + + Controls the default boolean operator used to combine adjacent bare terms + in a query string. + + Attributes: + type (IndexType): Always ``IndexType.FTS``. + default_operator (str): Default boolean operator for adjacent bare terms. + Supported values (case-insensitive): "OR" (default), "AND". + + Examples: + >>> params = FtsQueryParam(default_operator="AND") + >>> print(params.default_operator) + AND + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + default_operator: str = "", + ) -> None: + """ + Constructs an FtsQueryParam instance. + + Args: + default_operator (str, optional): Default boolean operator for adjacent + bare terms. Supported: "OR", "AND". Defaults to "" (uses engine default). + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def default_operator(self) -> str: + """ + str: Default boolean operator for bare terms. + """ + +class IndexOption: + """ + + Options for creating an index. + + Attributes: + concurrency (int): Number of threads to use during index creation. + If 0, the system will choose an optimal value automatically. + Default is 0. + + Examples: + >>> opt = IndexOption(concurrency=4) + >>> print(opt.concurrency) + 4 + """ + + def __getstate__(self) -> tuple: ... + def __init__(self, concurrency: typing.SupportsInt = 0) -> None: + """ + Constructs an IndexOption instance. + + Args: + concurrency (int, optional): Number of concurrent threads. + 0 means auto-detect. Defaults to 0. + """ + + def __setstate__(self, arg0: tuple) -> None: ... + @property + def concurrency(self) -> int: + """ + int: Number of threads used for index creation (0 = auto). + """ + +class IndexParam: + """ + + Base class for all index parameter configurations. + + This abstract base class defines the common interface for index types. + It should not be instantiated directly; use derived classes instead. + + Attributes: + type (IndexType): The type of the index (e.g., HNSW, FLAT, INVERT). + """ + + __hash__: typing.ClassVar[None] = None + + def __eq__(self, arg0: typing.Any) -> bool: ... + def __getstate__(self) -> tuple: ... + def __setstate__(self, arg0: tuple) -> None: ... + def clone(self) -> IndexParam: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def type(self) -> zvec._zvec.typing.IndexType: + """ + IndexType: The type of the index. + """ + +class InvertIndexParam(IndexParam): + """ + + Parameters for configuring an invert index. + + This class controls whether range query + optimization is enabled for invert index structures. + + Attributes: + type (IndexType): Always `IndexType.INVERTED`. + enable_range_optimization (bool): Whether range optimization is enabled. + enable_extended_wildcard (bool): Whether extended wildcard (suffix and infix) search is enabled. + + Examples: + >>> params = InvertIndexParam(enable_range_optimization=True, enable_extended_wildcard=False) + >>> print(params.enable_range_optimization) + True + >>> print(params.enable_extended_wildcard) + False + >>> config = params.to_dict() + >>> print(config) + {'enable_range_optimization': True, 'enable_extended_wildcard': False} + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + enable_range_optimization: bool = False, + enable_extended_wildcard: bool = False, + ) -> None: + """ + Constructs an InvertIndexParam instance. + + Args: + enable_range_optimization (bool, optional): If True, enables range query + optimization for the invert index. Defaults to False. + enable_extended_wildcard (bool, optional): If True, enables extended wildcard + search including suffix and infix patterns. Defaults to False. + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def enable_extended_wildcard(self) -> bool: + """ + bool: Whether extended wildcard (suffix and infix) search is enabled. + Note: Prefix search is always enabled regardless of this setting. + """ + + @property + def enable_range_optimization(self) -> bool: + """ + bool: Whether range optimization is enabled for this inverted index. + """ + +class OptimizeOption: + """ + + Options for optimizing a collection (e.g., merging segments). + + Attributes: + concurrency (int): Number of threads to use during optimization. + If 0, the system will choose an optimal value automatically. + Default is 0. + + Examples: + >>> opt = OptimizeOption(concurrency=2) + >>> print(opt.concurrency) + 2 + """ + + def __getstate__(self) -> tuple: ... + def __init__(self, concurrency: typing.SupportsInt = 0) -> None: + """ + Constructs an OptimizeOption instance. + + Args: + concurrency (int, optional): Number of concurrent threads. + 0 means auto-detect. Defaults to 0. + """ + + def __setstate__(self, arg0: tuple) -> None: ... + @property + def concurrency(self) -> int: + """ + int: Number of threads used for optimization (0 = auto). + """ + +class QueryParam: + """ + + Base class for all query parameter configurations. + + This abstract base class defines common query settings such as search radius + and whether to force linear (brute-force) search. It should not be instantiated + directly; use derived classes like `HnswQueryParam` or `IVFQueryParam`. + + Attributes: + type (IndexType): The index type this query is configured for. + radius (float): Search radius for range queries. Used in combination with + top-k to filter results. Default is 0.0 (disabled). + is_linear (bool): If True, forces brute-force linear search instead of + using the index. Useful for debugging or small datasets. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + """ + + def __getstate__(self) -> tuple: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def is_linear(self) -> bool: + """ + bool: Whether to bypass the index and use brute-force linear search. + """ + + @property + def is_using_refiner(self) -> bool: + """ + bool: Whether to use refiner for the query. + """ + + @property + def radius(self) -> float: + """ + IndexType: The type of index this query targets. + """ + + @property + def type(self) -> zvec._zvec.typing.IndexType: + """ + IndexType: The type of index this query targets. + """ + +class SegmentOption: + """ + + Options for segment-level operations. + + Currently, this class mirrors CollectionOption and is used internally. + It supports read-only mode, memory mapping, and buffer configuration. + + Note: + This class is primarily for internal use. Most users should use + CollectionOption instead. + + Examples: + >>> opt = SegmentOption() + >>> print(opt.enable_mmap) + True + """ + + def __getstate__(self) -> tuple: ... + def __init__(self) -> None: + """ + Constructs a SegmentOption with default settings. + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def enable_mmap(self) -> bool: + """ + bool: Whether memory-mapped I/O is enabled. + """ + + @property + def max_buffer_size(self) -> int: + """ + int: Maximum buffer size in bytes (internal use). + """ + + @property + def read_only(self) -> bool: + """ + bool: Whether the segment is read-only. + """ + +class QuantizerParam: + """ + + Optional parameters for quantizer configuration. + + This class is only needed when customizing quantizer behavior (e.g., enabling + random rotation). It can be omitted for default quantization settings. + + Attributes: + enable_rotate (bool): Whether to apply random rotation before INT8/INT4 + quantization to reduce quantization error. + Only effective with quantize_type=INT8 or INT4. Defaults to False. + + Examples: + >>> qp = QuantizerParam(enable_rotate=True) + >>> print(qp.enable_rotate) + True + """ + + def __getstate__(self) -> tuple: ... + def __init__(self, enable_rotate: bool = False) -> None: + """ + Constructs a QuantizerParam instance. + + Args: + enable_rotate (bool, optional): Whether to apply random rotation + before INT8/INT4 quantization. Defaults to False. + """ + + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def __eq__(self, arg0: typing.Any) -> bool: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def enable_rotate(self) -> bool: + """ + bool: Whether random rotation is enabled before INT8/INT4 quantization. + """ + +class VectorIndexParam(IndexParam): + """ + + Base class for vector index parameter configurations. + + Encapsulates common settings for all vector index types. + + Attributes: + type (IndexType): The specific vector index type (e.g., HNSW, FLAT). + metric_type (MetricType): Distance metric used for similarity search. + quantize_type (QuantizeType): Optional vector quantization type. + quantizer_param (QuantizerParam): Optional quantizer parameters. + """ + + def __getstate__(self) -> tuple: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def metric_type(self) -> zvec._zvec.typing.MetricType: + """ + MetricType: Distance metric (e.g., IP, COSINE, L2). + """ + + @property + def quantize_type(self) -> zvec._zvec.typing.QuantizeType: + """ + QuantizeType: Vector quantization type (e.g., FP16, INT8). + """ + + @property + def quantizer_param(self) -> QuantizerParam: + """ + QuantizerParam: Quantizer configuration including enable_rotate. + """ + +class _SearchQuery: + field_name: str + filter: str + include_vector: bool + query_params: QueryParam + + def __getstate__(self) -> tuple: ... + def __init__(self) -> None: ... + def __setstate__(self, arg0: tuple) -> None: ... + def set_vector(self, field_schema: typing.Any, obj: typing.Any) -> None: + """ + Set the query vector. + + Dense vector source data must not be modified until the query finishes. + """ + + @property + def output_fields(self) -> list[str] | None: ... + @output_fields.setter + def output_fields(self, arg0: collections.abc.Sequence[str] | None) -> None: ... + @property + def topk(self) -> int: ... + @topk.setter + def topk(self, arg0: typing.SupportsInt) -> None: ... diff --git a/python/zvec/model/param/query.py b/python/zvec/model/param/query.py new file mode 100644 index 0000000..d0916b3 --- /dev/null +++ b/python/zvec/model/param/query.py @@ -0,0 +1,143 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import warnings +from dataclasses import dataclass +from typing import Optional, Union + +from ...common import VectorType +from . import FtsQueryParam, HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam + +__all__ = ["Fts", "Query", "VectorQuery"] + + +@dataclass(frozen=True) +class Fts: + """Full-text search query parameters. + + Attributes: + query_string (Optional[str]): FTS query expression + (e.g. '+vector -slow "exact phrase"'). Mutually exclusive with match_string. + match_string (Optional[str]): Natural language match string, + tokenized and combined using the default operator. + Mutually exclusive with query_string. + """ + + query_string: Optional[str] = None + match_string: Optional[str] = None + + +@dataclass(frozen=True) +class Query: + """Represents a search query for a specific field in a collection. + + A `Query` can be constructed for either vector search or full-text search, + but not both simultaneously. + + For vector search, provide `id` or `vector` (and optionally `param`). + For FTS, provide `fts`. + + Attributes: + field_name (str): Name of the field to query. + id (Optional[str], optional): Document ID to fetch vector from. Default is None. + vector (VectorType, optional): Explicit query vector. Default is None. + param (Optional[Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam, FtsQueryParam]], optional): + Index-specific query parameters. Default is None. + fts (Optional[Fts], optional): Full-text search parameters. Default is None. + + Examples: + >>> import zvec + >>> # Query by ID + >>> q1 = zvec.Query(field_name="embedding", id="doc123") + >>> # Query by vector + >>> q2 = zvec.Query( + ... field_name="embedding", + ... vector=[0.1, 0.2, 0.3], + ... param=HnswQueryParam(ef=300) + ... ) + >>> # FTS query + >>> q3 = zvec.Query( + ... field_name="content", + ... fts=Fts(match_string="machine learning") + ... ) + >>> # FTS query with custom operator + >>> q4 = zvec.Query( + ... field_name="content", + ... fts=Fts(match_string="machine learning"), + ... param=FtsQueryParam(default_operator="AND") + ... ) + """ + + field_name: str + id: Optional[str] = None + vector: VectorType = None + param: Optional[ + Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam, FtsQueryParam] + ] = None + fts: Optional[Fts] = None + + def has_id(self) -> bool: + """Check if the query is based on a document ID. + + Returns: + bool: True if `id` is set, False otherwise. + """ + return self.id is not None + + def has_vector(self) -> bool: + """Check if the query contains an explicit vector. + + Returns: + bool: True if `vector` is non-empty, False otherwise. + """ + return self.vector is not None and len(self.vector) > 0 + + def has_fts(self) -> bool: + """Check if the query contains an FTS (full-text search) condition. + + Returns: + bool: True if `fts` is set with a query_string or match_string. + """ + if self.fts is not None: + return bool(self.fts.query_string) or bool(self.fts.match_string) + return False + + def _validate(self) -> None: + if self.field_name is None: + raise ValueError("Field name cannot be empty") + if self.has_id() and self.has_vector(): + raise ValueError("Cannot provide both id and vector") + if self.has_fts() and (self.has_vector() or self.has_id()): + raise ValueError( + "Cannot combine fts with vector search fields (id/vector) in a single Query" + ) + if self.fts is not None and self.fts.query_string and self.fts.match_string: + raise ValueError( + "Cannot provide both query_string and match_string in Fts; " + "they are mutually exclusive" + ) + + +class VectorQuery(Query): + """Deprecated alias for Query. Use Query instead.""" + + def __new__(cls, *args, **kwargs): # noqa : ARG004 + warnings.warn( + "VectorQuery is deprecated and will be removed in a future version. " + "Use Query instead.", + DeprecationWarning, + stacklevel=2, + ) + return super().__new__(cls) diff --git a/python/zvec/model/schema/__init__.py b/python/zvec/model/schema/__init__.py new file mode 100644 index 0000000..b952454 --- /dev/null +++ b/python/zvec/model/schema/__init__.py @@ -0,0 +1,21 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from zvec._zvec.schema import CollectionStats + +from .collection_schema import CollectionSchema +from .field_schema import FieldSchema, VectorSchema + +__all__ = ["CollectionSchema", "CollectionStats", "FieldSchema", "VectorSchema"] diff --git a/python/zvec/model/schema/__init__.pyi b/python/zvec/model/schema/__init__.pyi new file mode 100644 index 0000000..0dfa80d --- /dev/null +++ b/python/zvec/model/schema/__init__.pyi @@ -0,0 +1,109 @@ +""" +This module contains the schema of Zvec +""" + +from __future__ import annotations + +import collections.abc +import typing + +import zvec._zvec.param +import zvec._zvec.typing + +from .collection_schema import CollectionSchema +from .field_schema import FieldSchema, VectorSchema + +__all__: list[str] = [ + "CollectionSchema", + "CollectionStats", + "FieldSchema", + "VectorSchema", +] + +class CollectionStats: + def __init__(self) -> None: ... + def __repr__(self) -> str: ... + @property + def doc_count(self) -> int: ... + @property + def index_completeness(self) -> dict[str, float]: ... + +class _CollectionSchema: + __hash__: typing.ClassVar[None] = None + + def __eq__(self, arg0: _CollectionSchema) -> bool: ... + def __init__( + self, name: str, fields: collections.abc.Sequence[_FieldSchema] + ) -> None: + """ + Construct with name and list of fields + """ + + def __ne__(self, arg0: _CollectionSchema) -> bool: ... + def fields(self) -> list[_FieldSchema]: + """ + Return list of all field schemas. + """ + + def forward_fields(self) -> list[_FieldSchema]: + """ + Return list of forward-indexed fields. + """ + + def get_field(self, field_name: str) -> _FieldSchema: + """ + Get field by name (const pointer), returns None if not found. + """ + + def get_forward_field(self, field_name: str) -> _FieldSchema: + """ + Get forward field (used for filtering). + """ + + def get_vector_field(self, field_name: str) -> _FieldSchema: + """ + Get vector field by name. + """ + + def has_field(self, field_name: str) -> bool: + """ + Check if a field exists. + """ + + def vector_fields(self) -> list[_FieldSchema]: + """ + Return list of vector fields. + """ + + @property + def name(self) -> str: ... + +class _FieldSchema: + __hash__: typing.ClassVar[None] = None + + def __eq__(self, arg0: _FieldSchema) -> bool: ... + def __init__( + self, + name: str, + data_type: zvec._zvec.typing.DataType, + nullable: bool = False, + dimension: typing.SupportsInt = 0, + index_param: zvec._zvec.param.IndexParam = None, + ) -> None: ... + def __ne__(self, arg0: _FieldSchema) -> bool: ... + @property + def data_type(self) -> zvec._zvec.typing.DataType: ... + @property + def dimension(self) -> int: ... + @property + def index_param(self) -> typing.Any: ... + @property + def index_type(self) -> zvec._zvec.typing.IndexType: ... + @property + def is_dense_vector(self) -> bool: ... + @property + def is_sparse_vector(self) -> bool: ... + @property + def name(self) -> str: ... + @property + def nullable(self) -> bool: ... diff --git a/python/zvec/model/schema/collection_schema.py b/python/zvec/model/schema/collection_schema.py new file mode 100644 index 0000000..3e89710 --- /dev/null +++ b/python/zvec/model/schema/collection_schema.py @@ -0,0 +1,215 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import json +from typing import Optional, Union + +from zvec._zvec.schema import _CollectionSchema, _FieldSchema + +from .field_schema import FieldSchema, VectorSchema + +__all__ = [ + "CollectionSchema", +] + + +class CollectionSchema: + """Defines the structure of a collection in Zvec. + + A collection schema specifies the name of the collection and its fields, + including both scalar fields (e.g., int, string) and vector fields. + Field names must be unique across both scalar and vector fields. + + Args: + name (str): Name of the collection. + fields (Optional[Union[FieldSchema, list[FieldSchema]]], optional): + One or more scalar field definitions. Defaults to None. + vectors (Optional[Union[VectorSchema, list[VectorSchema]]], optional): + One or more vector field definitions. Defaults to None. + + Raises: + TypeError: If `fields` or `vectors` are of unsupported types. + ValueError: If any field or vector name is duplicated. + + Examples: + >>> from zvec import FieldSchema, VectorSchema, DataType, IndexType + >>> id_field = FieldSchema("id", DataType.INT64, is_primary=True) + >>> emb_field = VectorSchema("embedding", dim=128, data_type=DataType.VECTOR_FP32) + >>> schema = CollectionSchema( + ... name="my_collection", + ... fields=id_field, + ... vectors=emb_field + ... ) + >>> print(schema.name) + my_collection + """ + + def __init__( + self, + name: str, + fields: Optional[Union[FieldSchema, list[FieldSchema]]] = None, + vectors: Optional[Union[VectorSchema, list[VectorSchema]]] = None, + ): + if name is None or not isinstance(name, str): + raise ValueError( + f"schema validate failed: collection name must be str, got {type(name).__name__}" + ) + + # handle fields + _fields_name: list[str] = [] + _fields_list: list[_FieldSchema] = [] + + self._check_fields(fields, _fields_name, _fields_list) + self._check_vectors(vectors, _fields_name, _fields_list) + + # init + self._cpp_obj = _CollectionSchema( + name=name, + fields=_fields_list, + ) + + def _check_fields( + self, + fields: Optional[Union[FieldSchema, list[FieldSchema]]], + _fields_name: list[str], + _fields_list: list[_FieldSchema], + ) -> None: + field_items = [] + + if isinstance(fields, FieldSchema): + field_items = [fields] + elif isinstance(fields, list): + field_items = fields + elif fields is None: + field_items = [] + else: + raise TypeError( + f"schema validate failed: invalid 'fields' type, expected FieldSchema or list[FieldSchema], " + f"got {type(fields).__name__}" + ) + + for idx, field in enumerate(field_items): + if not isinstance(field, FieldSchema): + raise TypeError( + f"schema validate failed: invalid field type in 'fields' list, expected FieldSchema, " + f"got {type(field).__name__} at index {idx}" + ) + + if field.name in _fields_name: + raise ValueError( + f"schema validate failed: duplicate field name '{field.name}': field names must be unique" + ) + _fields_name.append(field.name) + _fields_list.append(field._get_object()) + + def _check_vectors( + self, + vectors: Optional[Union[VectorSchema, list[VectorSchema]]], + _fields_name: list[str], + _fields_list: list[_FieldSchema], + ) -> None: + # handle vector + if isinstance(vectors, VectorSchema): + vectors_items = [vectors] + elif isinstance(vectors, list): + vectors_items = vectors + elif vectors is None: + vectors_items = [] + else: + raise TypeError( + f"schema validate failed: invalid 'vectors' type, expected VectorSchema or list[VectorSchema], " + f"got {type(vectors).__name__}" + ) + + for idx, vector in enumerate(vectors_items): + if not isinstance(vector, VectorSchema): + raise TypeError( + f"schema validate failed: invalid vector type in 'vectors' list, expected VectorSchema, " + f"got {type(vector).__name__} at index {idx}" + ) + + if vector.name in _fields_name: + raise ValueError( + f"schema validate failed: duplicate vector name '{vector.name}', vector names must be unique " + f"(conflicts with existing field or vector)" + ) + _fields_name.append(vector.name) + _fields_list.append(vector._get_object()) + + @classmethod + def _from_core(cls, core_collection_schema: _CollectionSchema): + inst = cls.__new__(cls) + if not core_collection_schema: + raise ValueError("schema validate failed: schema is null") + inst._cpp_obj = core_collection_schema + return inst + + @property + def name(self) -> str: + """str: The name of the collection.""" + return self._cpp_obj.name + + def field(self, name: str) -> Optional[FieldSchema]: + """Retrieve a scalar field by name. + + Args: + name (str): Name of the field. + + Returns: + Optional[FieldSchema]: The field if found, otherwise None. + """ + _field = self._cpp_obj.get_forward_field(name) + return FieldSchema._from_core(_field) if _field else None + + def vector(self, name: str) -> Optional[VectorSchema]: + """Retrieve a vector field by name. + + Args: + name (str): Name of the vector field. + + Returns: + Optional[VectorSchema]: The vector field if found, otherwise None. + """ + _field = self._cpp_obj.get_vector_field(name) + return VectorSchema._from_core(_field) if _field else None + + @property + def fields(self) -> list[FieldSchema]: + """list[FieldSchema]: All scalar (non-vector) fields in the schema.""" + _fields = self._cpp_obj.forward_fields() + return [FieldSchema._from_core(_field) for _field in _fields] + + @property + def vectors(self) -> list[VectorSchema]: + """list[VectorSchema]: All vector fields in the schema.""" + _vectors = self._cpp_obj.vector_fields() + return [VectorSchema._from_core(_vector) for _vector in _vectors] + + def _get_object(self) -> _CollectionSchema: + return self._cpp_obj + + def __repr__(self) -> str: + try: + schema = { + "name": self.name, + "fields": {field.name: field.__dict__() for field in self.fields}, + "vectors": {vector.name: vector.__dict__() for vector in self.vectors}, + } + return json.dumps(schema, indent=2, ensure_ascii=False) + except Exception as e: + return f"" + + def __str__(self) -> str: + return self.__repr__() diff --git a/python/zvec/model/schema/field_schema.py b/python/zvec/model/schema/field_schema.py new file mode 100644 index 0000000..9b757d6 --- /dev/null +++ b/python/zvec/model/schema/field_schema.py @@ -0,0 +1,310 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import json +from typing import Any, Optional, Union + +from zvec._zvec.schema import _FieldSchema +from zvec.model.param import ( + FlatIndexParam, + FtsIndexParam, + HnswIndexParam, + HnswRabitqIndexParam, + InvertIndexParam, + IVFIndexParam, +) +from zvec.typing import DataType + +__all__ = [ + "FieldSchema", + "VectorSchema", +] + +SUPPORT_VECTOR_DATA_TYPE = [ + DataType.VECTOR_FP16, + DataType.VECTOR_FP32, + DataType.VECTOR_FP64, + DataType.VECTOR_INT8, + DataType.SPARSE_VECTOR_FP16, + DataType.SPARSE_VECTOR_FP32, +] + +SUPPORT_SCALAR_DATA_TYPE = [ + DataType.INT32, + DataType.INT64, + DataType.UINT32, + DataType.UINT64, + DataType.FLOAT, + DataType.DOUBLE, + DataType.STRING, + DataType.BOOL, + DataType.ARRAY_INT32, + DataType.ARRAY_INT64, + DataType.ARRAY_UINT32, + DataType.ARRAY_UINT64, + DataType.ARRAY_FLOAT, + DataType.ARRAY_DOUBLE, + DataType.ARRAY_STRING, + DataType.ARRAY_BOOL, +] + + +class FieldSchema: + """Represents a scalar (non-vector) field in a collection schema. + + A `FieldSchema` defines the name, data type, nullability, and optional + inverted index configuration for a regular field (e.g., ID, timestamp, category). + + Args: + name (str): Name of the field. Must be unique within the collection. + data_type (DataType): Data type of the field (e.g., INT64, STRING). + nullable (bool, optional): Whether the field can contain null values. + Defaults to False. + index_param (Optional[Union[InvertIndexParam, FtsIndexParam]], optional): + Index parameters for this field. Use ``InvertIndexParam`` for scalar + inverted indexing, or ``FtsIndexParam`` for full-text search indexing + on STRING fields. Defaults to None. + + Examples: + >>> from zvec.typing import DataType + >>> from zvec.model.param import InvertIndexParam, FtsIndexParam + >>> id_field = FieldSchema( + ... name="id", + ... data_type=DataType.INT64, + ... nullable=False, + ... index_param=InvertIndexParam(enable_range_optimization=True) + ... ) + >>> content_field = FieldSchema( + ... name="content", + ... data_type=DataType.STRING, + ... nullable=False, + ... index_param=FtsIndexParam(tokenizer_name="standard") + ... ) + """ + + def __init__( + self, + name: str, + data_type: DataType, + nullable: bool = False, + index_param: Optional[Union[InvertIndexParam, FtsIndexParam]] = None, + ): + if name is None or not isinstance(name, str): + raise ValueError( + f"schema validate failed: field name must be str, got {type(name).__name__}" + ) + + if data_type not in SUPPORT_SCALAR_DATA_TYPE: + raise ValueError( + f"schema validate failed: scalar_field's data_type must be one of " + f"{', '.join(str(dt) for dt in SUPPORT_SCALAR_DATA_TYPE)}, " + f"but field[{name}]'s data_type is {data_type}" + ) + + self._cpp_obj = _FieldSchema( + name=name, + data_type=data_type, + dimension=0, + nullable=nullable, + index_param=index_param, + ) + + @classmethod + def _from_core(cls, core_field_schema: _FieldSchema): + if core_field_schema is None: + raise ValueError("schema validate failed: field schema is None") + inst = cls.__new__(cls) + inst._cpp_obj = core_field_schema + return inst + + def _get_object(self) -> _FieldSchema: + return self._cpp_obj + + @property + def name(self) -> str: + """str: The name of the field.""" + return self._cpp_obj.name + + @property + def data_type(self) -> DataType: + """DataType: The data type of the field (e.g., INT64, STRING).""" + return self._cpp_obj.data_type + + @property + def nullable(self) -> bool: + """bool: Whether the field allows null values.""" + return self._cpp_obj.nullable + + @property + def index_param(self) -> Optional[Union[InvertIndexParam, FtsIndexParam]]: + """Optional[Union[InvertIndexParam, FtsIndexParam]]: Index configuration, if any.""" + return self._cpp_obj.index_param + + def __dict__(self) -> dict[str, Any]: + return { + "name": self.name, + "data_type": ( + self.data_type.name + if hasattr(self.data_type, "name") + else str(self.data_type) + ), + "nullable": self.nullable, + "index_param": ( + self.index_param.to_dict() if self.index_param is not None else None + ), + } + + def __repr__(self) -> str: + try: + schema = self.__dict__() + return json.dumps(schema, indent=2, ensure_ascii=False) + except Exception as e: + return f"" + + def __str__(self) -> str: + return self.__repr__() + + def __eq__(self, other: object) -> bool: + if not isinstance(other, FieldSchema): + return False + return self._cpp_obj == other._cpp_obj + + def __hash__(self) -> int: + return hash((self.name, self.data_type, self.nullable)) + + +class VectorSchema: + """Represents a vector field in a collection schema. + + A `VectorSchema` defines the name, data type, dimensionality, and index + configuration for a vector field used in similarity search. + + Args: + name (str): Name of the vector field. Must be unique within the collection. + data_type (DataType): Vector data type (e.g., VECTOR_FP32, VECTOR_INT8). + dimension (int, optional): Dimensionality of the vector. Must be > 0 for dense vectors; + may be `None` for sparse vectors. + index_param (Union[HnswIndexParam, IVFIndexParam, FlatIndexParam], optional): + Index configuration for this vector field. Defaults to + ``HnswIndexParam()``. + + Examples: + >>> from zvec.typing import DataType + >>> from zvec.model.param import HnswIndexParam + >>> emb_field = VectorSchema( + ... name="embedding", + ... data_type=DataType.VECTOR_FP32, + ... dimension=128, + ... index_param=HnswIndexParam(ef_construction=200, m=16) + ... ) + """ + + def __init__( + self, + name: str, + data_type: DataType, + dimension: Optional[int] = 0, + index_param: Optional[ + Union[HnswIndexParam, HnswRabitqIndexParam, FlatIndexParam, IVFIndexParam] + ] = None, + ): + if name is None or not isinstance(name, str): + raise ValueError( + f"schema validate failed: field name must be str, got {type(name).__name__}" + ) + + if not isinstance(dimension, int) or dimension < 0: + raise ValueError("schema validate failed: vector's dimension must be >= 0") + + if data_type not in SUPPORT_VECTOR_DATA_TYPE: + raise ValueError( + f"schema validate failed: vector's data_type must be one of " + f"{', '.join(str(dt) for dt in SUPPORT_VECTOR_DATA_TYPE)}, " + f"but field[{name}]'s data_type is {data_type}" + ) + + if index_param is None: + index_param = FlatIndexParam() + + self._cpp_obj = _FieldSchema( + name=name, + data_type=data_type, + dimension=dimension, + nullable=False, + index_param=index_param, + ) + + @classmethod + def _from_core(cls, core_field_schema: _FieldSchema): + inst = cls.__new__(cls) + inst._cpp_obj = core_field_schema + return inst + + def _get_object(self) -> _FieldSchema: + return self._cpp_obj + + @property + def name(self) -> str: + """str: The name of the vector field.""" + return self._cpp_obj.name + + @property + def data_type(self) -> DataType: + """DataType: The vector data type (e.g., VECTOR_FP32).""" + return self._cpp_obj.data_type + + @property + def dimension(self) -> int: + """int: The dimensionality of the vector.""" + return self._cpp_obj.dimension + + @property + def index_param( + self, + ) -> Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]: + """Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector.""" + return self._cpp_obj.index_param + + def __dict__(self) -> dict[str, Any]: + return { + "name": self.name, + "data_type": ( + self.data_type.name + if hasattr(self.data_type, "name") + else str(self.data_type) + ), + "dimension": self.dimension, + "index_param": ( + self.index_param.to_dict() if self.index_param is not None else None + ), + } + + def __repr__(self) -> str: + try: + schema = self.__dict__() + return json.dumps(schema, indent=2, ensure_ascii=False) + except Exception as e: + return f"" + + def __str__(self) -> str: + return self.__repr__() + + def __eq__(self, other: object) -> bool: + if not isinstance(other, VectorSchema): + return False + return self._cpp_obj == other._cpp_obj + + def __hash__(self) -> int: + return hash((self.name, self.data_type, self.dimension)) diff --git a/python/zvec/py.typed b/python/zvec/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/python/zvec/tool/__init__.py b/python/zvec/tool/__init__.py new file mode 100644 index 0000000..00e3bfe --- /dev/null +++ b/python/zvec/tool/__init__.py @@ -0,0 +1,18 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from .util import require_module + +__all__ = ["require_module"] diff --git a/python/zvec/tool/util.py b/python/zvec/tool/util.py new file mode 100644 index 0000000..409a4d5 --- /dev/null +++ b/python/zvec/tool/util.py @@ -0,0 +1,63 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import importlib +from typing import Any, Optional + + +def require_module(module: str, mitigation: Optional[str] = None) -> Any: + """Import a Python module and raise a user-friendly error if it is not available. + + This utility helps provide actionable error messages when optional dependencies + are missing. It attempts to import the given module and, on failure, suggests + a `pip install` command based on either the module name or an optional + mitigation package name. + + Args: + module (str): The full module name to import (e.g., ``"numpy"``, ``"pandas.io.parquet"``). + mitigation (Optional[str], optional): The package name to suggest for installation + if the import fails. If not provided, the top-level package of `module` + will be used (e.g., ``"pandas"`` for ``"pandas.io.parquet"``). + + Returns: + Any: The imported module object. + + Raises: + ImportError: If the module cannot be imported, with a clear installation hint. + + Examples: + >>> import zvec + >>> np = zvec.require_module("numpy") + >>> pq = zvec.require_module("pyarrow.parquet", mitigation="pyarrow") + + Note: + This function is intended for lazy-loading optional dependencies + with helpful error messages, not for core dependencies. + """ + try: + return importlib.import_module(module) + except ImportError as e: + package = mitigation or module + msg = f"Required package '{package}' is not installed. " + if "." in module: + top_level = module.split(".", maxsplit=1)[0] + msg += f"Module '{module}' is part of '{top_level}', " + if mitigation: + msg += f"please pip install '{mitigation}'." + else: + msg += f"please pip install '{top_level}'." + else: + msg += f"Please pip install '{package}'." + raise ImportError(msg) from e diff --git a/python/zvec/typing/__init__.py b/python/zvec/typing/__init__.py new file mode 100644 index 0000000..0faa19d --- /dev/null +++ b/python/zvec/typing/__init__.py @@ -0,0 +1,32 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from zvec._zvec.typing import ( + DataType, + IndexType, + MetricType, + QuantizeType, + Status, + StatusCode, +) + +__all__ = [ + "DataType", + "IndexType", + "MetricType", + "QuantizeType", + "Status", + "StatusCode", +] diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi new file mode 100644 index 0000000..db3a201 --- /dev/null +++ b/python/zvec/typing/__init__.pyi @@ -0,0 +1,404 @@ +""" +This module contains the basic data types of Zvec +""" + +from __future__ import annotations + +import typing + +__all__: list[str] = [ + "DataType", + "IndexType", + "MetricType", + "QuantizeType", + "Status", + "StatusCode", +] + +class DataType: + """ + + Enumeration of supported data types in Zvec. + + Includes scalar types, dense/sparse vector types, and array types. + + Examples: + >>> import zvec + >>> print(zvec.DataType.FLOAT) + DataType.FLOAT + >>> print(zvec.DataType.VECTOR_FP32) + DataType.VECTOR_FP32 + + + Members: + + STRING + + BOOL + + INT32 + + INT64 + + FLOAT + + DOUBLE + + UINT32 + + UINT64 + + VECTOR_FP16 + + VECTOR_FP32 + + VECTOR_FP64 + + VECTOR_INT8 + + SPARSE_VECTOR_FP32 + + SPARSE_VECTOR_FP16 + + ARRAY_STRING + + ARRAY_INT32 + + ARRAY_INT64 + + ARRAY_FLOAT + + ARRAY_DOUBLE + + ARRAY_BOOL + + ARRAY_UINT32 + + ARRAY_UINT64 + """ + + ARRAY_BOOL: typing.ClassVar[DataType] # value = + ARRAY_DOUBLE: typing.ClassVar[DataType] # value = + ARRAY_FLOAT: typing.ClassVar[DataType] # value = + ARRAY_INT32: typing.ClassVar[DataType] # value = + ARRAY_INT64: typing.ClassVar[DataType] # value = + ARRAY_STRING: typing.ClassVar[DataType] # value = + ARRAY_UINT32: typing.ClassVar[DataType] # value = + ARRAY_UINT64: typing.ClassVar[DataType] # value = + BOOL: typing.ClassVar[DataType] # value = + DOUBLE: typing.ClassVar[DataType] # value = + FLOAT: typing.ClassVar[DataType] # value = + INT32: typing.ClassVar[DataType] # value = + INT64: typing.ClassVar[DataType] # value = + SPARSE_VECTOR_FP16: typing.ClassVar[ + DataType + ] # value = + SPARSE_VECTOR_FP32: typing.ClassVar[ + DataType + ] # value = + STRING: typing.ClassVar[DataType] # value = + UINT32: typing.ClassVar[DataType] # value = + UINT64: typing.ClassVar[DataType] # value = + VECTOR_FP16: typing.ClassVar[DataType] # value = + VECTOR_FP32: typing.ClassVar[DataType] # value = + VECTOR_FP64: typing.ClassVar[DataType] # value = + VECTOR_INT8: typing.ClassVar[DataType] # value = + __members__: typing.ClassVar[ + dict[str, DataType] + ] # value = {'STRING': , 'BOOL': , 'INT32': , 'INT64': , 'FLOAT': , 'DOUBLE': , 'UINT32': , 'UINT64': , 'VECTOR_FP16': , 'VECTOR_FP32': , 'VECTOR_FP64': , 'VECTOR_INT8': , 'SPARSE_VECTOR_FP32': , 'SPARSE_VECTOR_FP16': , 'ARRAY_STRING': , 'ARRAY_INT32': , 'ARRAY_INT64': , 'ARRAY_FLOAT': , 'ARRAY_DOUBLE': , 'ARRAY_BOOL': , 'ARRAY_UINT32': , 'ARRAY_UINT64': } + + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + def __init__(self, value: typing.SupportsInt) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: typing.SupportsInt) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... + +class IndexType: + """ + + Enumeration of supported index types in Zvec. + + Examples: + >>> import zvec + >>> print(zvec.IndexType.HNSW) + IndexType.HNSW + + + Members: + + UNDEFINED + + HNSW + + IVF + + FLAT + + INVERT + """ + + FLAT: typing.ClassVar[IndexType] # value = + HNSW: typing.ClassVar[IndexType] # value = + INVERT: typing.ClassVar[IndexType] # value = + IVF: typing.ClassVar[IndexType] # value = + UNDEFINED: typing.ClassVar[IndexType] # value = + __members__: typing.ClassVar[ + dict[str, IndexType] + ] # value = {'UNDEFINED': , 'HNSW': , 'IVF': , 'FLAT': , 'INVERT': } + + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + def __init__(self, value: typing.SupportsInt) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: typing.SupportsInt) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... + +class MetricType: + """ + + Enumeration of supported distance/similarity metrics. + + - COSINE: Cosine similarity. + - IP: Inner product (dot product). + - L2: Euclidean distance (L2 norm). + + Examples: + >>> import zvec + >>> print(zvec.MetricType.COSINE) + MetricType.COSINE + + + Members: + + COSINE + + IP + + L2 + """ + + COSINE: typing.ClassVar[MetricType] # value = + IP: typing.ClassVar[MetricType] # value = + L2: typing.ClassVar[MetricType] # value = + __members__: typing.ClassVar[ + dict[str, MetricType] + ] # value = {'COSINE': , 'IP': , 'L2': } + + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + def __init__(self, value: typing.SupportsInt) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: typing.SupportsInt) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... + +class QuantizeType: + """ + + Enumeration of supported quantization types for vector compression. + + Examples: + >>> import zvec + >>> print(zvec.QuantizeType.INT8) + QuantizeType.INT8 + + + Members: + + UNDEFINED + + FP16 + + INT8 + + INT4 + """ + + FP16: typing.ClassVar[QuantizeType] # value = + INT4: typing.ClassVar[QuantizeType] # value = + INT8: typing.ClassVar[QuantizeType] # value = + UNDEFINED: typing.ClassVar[QuantizeType] # value = + __members__: typing.ClassVar[ + dict[str, QuantizeType] + ] # value = {'UNDEFINED': , 'FP16': , 'INT8': , 'INT4': } + + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + def __init__(self, value: typing.SupportsInt) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: typing.SupportsInt) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... + +class Status: + """ + + Represents the outcome of a Zvec operation. + + A `Status` object is either OK (success) or carries an error code and message. + + Examples: + >>> from zvec.typing import Status, StatusCode + >>> s = Status() + >>> print(s.ok()) + True + >>> s = Status(StatusCode.INVALID_ARGUMENT, "Field not found") + >>> print(s.code() == StatusCode.INVALID_ARGUMENT) + True + >>> print(s.message()) + Field not found + """ + + __hash__: typing.ClassVar[None] = None + + @staticmethod + def AlreadyExists(message: str) -> Status: ... + @staticmethod + def InternalError(message: str) -> Status: ... + @staticmethod + def InvalidArgument(message: str) -> Status: ... + @staticmethod + def NotFound(message: str) -> Status: ... + @staticmethod + def OK() -> Status: + """ + Create an OK status. + """ + + @staticmethod + def PermissionDenied(message: str) -> Status: ... + def __eq__(self, arg0: Status) -> bool: ... + @typing.overload + def __init__(self) -> None: ... + @typing.overload + def __init__(self, code: StatusCode, message: str = "") -> None: + """ + Construct a status with the given code and optional message. + + Args: + code (StatusCode): The status code. + message (str, optional): Error message. Defaults to empty string. + """ + + def __ne__(self, arg0: Status) -> bool: ... + def __repr__(self) -> str: ... + def code(self) -> StatusCode: + """ + StatusCode: Returns the status code. + """ + + def message(self) -> str: + """ + str: Returns the error message (may be empty). + """ + + def ok(self) -> bool: + """ + bool: Returns True if the status is OK. + """ + +class StatusCode: + """ + + Enumeration of possible status codes for Zvec operations. + + Used by the `Status` class to indicate success or failure reason. + + + Members: + + OK + + NOT_FOUND + + ALREADY_EXISTS + + INVALID_ARGUMENT + + PERMISSION_DENIED + + FAILED_PRECONDITION + + RESOURCE_EXHAUSTED + + UNAVAILABLE + + INTERNAL_ERROR + + NOT_SUPPORTED + + UNKNOWN + """ + + ALREADY_EXISTS: typing.ClassVar[ + StatusCode + ] # value = + FAILED_PRECONDITION: typing.ClassVar[ + StatusCode + ] # value = + INTERNAL_ERROR: typing.ClassVar[ + StatusCode + ] # value = + INVALID_ARGUMENT: typing.ClassVar[ + StatusCode + ] # value = + NOT_FOUND: typing.ClassVar[StatusCode] # value = + NOT_SUPPORTED: typing.ClassVar[StatusCode] # value = + OK: typing.ClassVar[StatusCode] # value = + PERMISSION_DENIED: typing.ClassVar[ + StatusCode + ] # value = + RESOURCE_EXHAUSTED: typing.ClassVar[ + StatusCode + ] # value = + UNAVAILABLE: typing.ClassVar[StatusCode] # value = + UNKNOWN: typing.ClassVar[StatusCode] # value = + __members__: typing.ClassVar[ + dict[str, StatusCode] + ] # value = {'OK': , 'NOT_FOUND': , 'ALREADY_EXISTS': , 'INVALID_ARGUMENT': , 'PERMISSION_DENIED': , 'FAILED_PRECONDITION': , 'RESOURCE_EXHAUSTED': , 'UNAVAILABLE': , 'INTERNAL_ERROR': , 'NOT_SUPPORTED': , 'UNKNOWN': } + + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + def __init__(self, value: typing.SupportsInt) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: typing.SupportsInt) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... diff --git a/python/zvec/typing/enum.py b/python/zvec/typing/enum.py new file mode 100644 index 0000000..3283f4e --- /dev/null +++ b/python/zvec/typing/enum.py @@ -0,0 +1,62 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from enum import IntEnum + +__all__ = ["LogLevel", "LogType"] + + +class LogLevel(IntEnum): + """Enumeration of logging severity levels, ordered from lowest to highest priority. + + Used to control verbosity and filtering of log messages. Higher numeric values + indicate more severe conditions. + + Note: + ``WARNING`` is an alias for ``WARN`` to match Python's built-in :mod:`logging` + module convention. + + Attributes: + DEBUG (int): Detailed information, typically of interest only when diagnosing problems. + INFO (int): Confirmation that things are working as expected. + WARN (int): An indication that something unexpected happened, or indicative of + potential future problems. (Alias: ``WARNING``) + WARNING (int): Same as ``WARN``. + ERROR (int): Due to a more serious problem, the software has not been able + to perform some function. + FATAL (int): A serious error, indicating that the program itself may be unable + to continue running. + """ + + DEBUG = 0 + INFO = 1 + WARN = 2 + WARNING = 2 + ERROR = 3 + FATAL = 4 + + +class LogType(IntEnum): + """Enumeration of supported log output destinations. + + Specifies where log messages should be written. + + Attributes: + CONSOLE (int): Output logs to standard output/error (e.g., terminal or IDE console). + FILE (int): Write logs to a persistent file on disk. + """ + + CONSOLE = 0 + FILE = 1 diff --git a/python/zvec/zvec.py b/python/zvec/zvec.py new file mode 100644 index 0000000..0046a49 --- /dev/null +++ b/python/zvec/zvec.py @@ -0,0 +1,246 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import Optional + +from zvec._zvec import Initialize, _Collection + +from .model import Collection +from .model.param import CollectionOption +from .model.schema import CollectionSchema + +__all__ = ["create_and_open", "init", "open"] + +from .typing.enum import LogLevel, LogType + + +def init( + *, + log_type: Optional[LogType] = None, + log_level: Optional[LogLevel] = None, + log_dir: Optional[str] = "./logs", + log_basename: Optional[str] = "zvec.log", + log_file_size: Optional[int] = 2048, + log_overdue_days: Optional[int] = 7, + query_threads: Optional[int] = None, + optimize_threads: Optional[int] = None, + invert_to_forward_scan_ratio: Optional[float] = None, + brute_force_by_keys_ratio: Optional[float] = None, + fts_brute_force_by_keys_ratio: Optional[float] = None, + memory_limit_mb: Optional[int] = None, + jieba_dict_dir: Optional[str] = None, +) -> None: + """Initialize Zvec with configuration options. + + This function must be called before any other operation. + It can only be called once — subsequent calls raise a ``RuntimeError``. + + Parameters set to ``None`` are **omitted** from the configuration and + fall back to Zvec's internal defaults, which may be derived from the runtime + environment (e.g., cgroup CPU/memory limits). Explicitly provided values + always override defaults. + + Args: + log_type (Optional[LogType], optional): Logger destination. + - ``LogType.CONSOLE`` (default if omitted or set to this) + - ``LogType.FILE`` + - If ``None``, uses internal default (currently ``CONSOLE``). + log_level (Optional[LogLevel], optional): Minimum log severity. + Default: ``LogLevel.WARN``. + Accepted values: ``DEBUG``, ``INFO``, ``WARN``, ``ERROR``, ``FATAL``. + If ``None``, uses internal default (``WARN``). + log_dir (Optional[str], optional): + Directory for log files (only used when ``log_type=FILE``). + Parent directories are **not** created automatically. + Default: ``"./logs"``. + If ``None``, internal default is used. + log_basename (Optional[str], optional): + Base name for rotated log files (e.g., ``zvec.log.1``, ``zvec.log.2``). + Default: ``"zvec.log"``. + log_file_size (Optional[int], optional): + Max size per log file in **MB** before rotation. + Default: ``2048`` MB (2 GB). + log_overdue_days (Optional[int], optional): + Days to retain rotated log files before deletion. + Default: ``7`` days. + query_threads (Optional[int], optional): + Number of threads for query execution. + If ``None`` (default), inferred from available CPU cores (via cgroup). + Must be ≥ 1 if provided. + optimize_threads (Optional[int], optional): + Threads for background tasks (e.g., compaction, indexing). + If ``None``, defaults to same as ``query_threads`` or CPU count. + invert_to_forward_scan_ratio (Optional[float], optional): + Threshold to switch from inverted index to full forward scan. + Range: [0.0, 1.0]. Higher → more aggressive index skipping. + Default: ``0.9`` (if omitted). + brute_force_by_keys_ratio (Optional[float], optional): + Threshold to use brute-force key lookup over index. + Lower → prefer index; higher → prefer brute-force. + Range: [0.0, 1.0]. Default: ``0.1``. + fts_brute_force_by_keys_ratio (Optional[float], optional): + Threshold to switch FTS scan from posting-driven to + candidate-driven (brute-force) when the invert filter is + highly selective. Independent from ``brute_force_by_keys_ratio`` + because per-candidate FTS cost is higher. + Range: [0.0, 1.0]. Default: ``0.05``. + memory_limit_mb (Optional[int], optional): + Soft memory cap in MB. Zvec may throttle or fail operations + approaching this limit. + If ``None``, inferred from cgroup memory limit * 0.8 (e.g., in Docker). + Must be > 0 if provided. + jieba_dict_dir (Optional[str], optional): + Override the default directory containing ``jieba.dict.utf8`` and + ``hmm_model.utf8`` for the jieba FTS tokenizer. When ``None``, the + value previously registered by ``zvec.set_default_jieba_dict_dir`` + (called automatically on ``import zvec`` to point at the wheel's + bundled dict) is preserved. JiebaTokenizer also honors the + ``ZVEC_JIEBA_DICT_DIR`` environment variable and per-field + ``FtsIndexParam.extra_params.jieba_dict_dir`` ahead of this value. + + Raises: + RuntimeError: If Zvec is already initialized. + ValueError: On invalid values (e.g., negative thread count, log level out of range). + TypeError: If a value has incorrect type (e.g., string for ``query_threads``). + + Note: + - All ``None`` arguments are **excluded** from the configuration payload, + allowing the core library to apply environment-aware defaults. + - This design ensures container-friendliness: in Kubernetes/Docker, + omitting ``memory_limit_mb`` and thread counts lets Zvec auto-adapt. + + Examples: + Initialize with defaults (log to console, auto-detect resources): + >>> import zvec + >>> zvec.init() + + Customize logging to file with rotation: + >>> zvec.init( + ... log_type=LogType.FILE, + ... log_dir="/var/log/zvec", + ... log_file_size=1024, + ... log_overdue_days=30 + ... ) + + Limit resources explicitly: + >>> zvec.init( + ... memory_limit_mb=2048, + ... query_threads=4, + ... optimize_threads=2 + ... ) + + Fine-tune query heuristics: + >>> zvec.init( + ... invert_to_forward_scan_ratio=0.95, + ... brute_force_by_keys_ratio=0.05 + ... ) + """ + # Build config dict, skipping None values + config_dict = {} + if log_type is not None: + if not isinstance(log_type, LogType): + raise TypeError("log_type must be LogType") + config_dict["log_type"] = log_type.name + if log_level is not None: + if not isinstance(log_level, LogLevel): + raise TypeError("log_level must be LogLevel") + config_dict["log_level"] = log_level.name + if log_dir is not None: + config_dict["log_dir"] = log_dir + if log_basename is not None: + config_dict["log_basename"] = log_basename + if log_file_size is not None: + config_dict["log_file_size"] = log_file_size + if log_overdue_days is not None: + config_dict["log_overdue_days"] = log_overdue_days + if query_threads is not None: + config_dict["query_threads"] = query_threads + if optimize_threads is not None: + config_dict["optimize_threads"] = optimize_threads + if invert_to_forward_scan_ratio is not None: + config_dict["invert_to_forward_scan_ratio"] = invert_to_forward_scan_ratio + if brute_force_by_keys_ratio is not None: + config_dict["brute_force_by_keys_ratio"] = brute_force_by_keys_ratio + if fts_brute_force_by_keys_ratio is not None: + config_dict["fts_brute_force_by_keys_ratio"] = fts_brute_force_by_keys_ratio + if memory_limit_mb is not None: + config_dict["memory_limit_mb"] = memory_limit_mb + if jieba_dict_dir is not None: + config_dict["jieba_dict_dir"] = jieba_dict_dir + + Initialize(config_dict) + + +def create_and_open( + path: str, + schema: CollectionSchema, + option: Optional[CollectionOption] = None, +) -> Collection: + """Create a new collection and open it for use. + + If a collection already exists at the given path, it may raise an error + depending on the underlying implementation. + + Args: + path (str): Path or name of the collection to create. + schema (CollectionSchema): Schema defining the structure of the collection. + option (Optional[CollectionOption]): Configuration options + for opening the collection. Defaults to a default-constructed + ``CollectionOption()`` if not provided. + + Returns: + Collection: An opened collection instance ready for operations. + + Examples: + >>> import zvec + >>> schema = zvec.CollectionSchema( + ... name="my_collection", + ... fields=[zvec.FieldSchema("id", zvec.DataType.INT64, nullable=True)] + ... ) + >>> coll = create_and_open("./my_collection", schema) + """ + if not isinstance(path, str): + raise TypeError("path must be a string") + if not isinstance(schema, CollectionSchema): + raise TypeError("schema must be a CollectionSchema") + + option = option or CollectionOption() + if not isinstance(option, CollectionOption): + raise TypeError("option must be a CollectionOption") + + _collection = _Collection.CreateAndOpen(path, schema._get_object(), option) + return Collection._from_core(_collection) + + +def open(path: str, option: CollectionOption = CollectionOption()) -> Collection: + """Open an existing collection from disk. + + The collection must have been previously created with ``create_and_open``. + + Args: + path (str): Path or name of the existing collection. + option (CollectionOption, optional): Configuration options + for opening the collection. Defaults to a default-constructed + ``CollectionOption()`` if not provided. + + Returns: + Collection: An opened collection instance. + + Examples: + >>> import zvec + >>> coll = zvec.open("./my_collection") + """ + _collection = _Collection.Open(path, option) + return Collection._from_core(_collection) diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 0000000..e69de29 diff --git a/scripts/build_android.sh b/scripts/build_android.sh new file mode 100755 index 0000000..a1785a4 --- /dev/null +++ b/scripts/build_android.sh @@ -0,0 +1,395 @@ +#!/bin/bash +set -e + +####################################################################### +# build_android.sh — cross-compile zvec for Android arm64-v8a, +# then run all C/C++ unit tests inside an emulator. +# +# Usage: +# ./scripts/build_android.sh [api_level] [build_type] +# +# Examples: +# ./scripts/build_android.sh # defaults: API 36, Release +# ./scripts/build_android.sh 36 Debug +####################################################################### + +CURRENT_DIR=$(pwd) +ABI="arm64-v8a" +API_LEVEL=${1:-35} +BUILD_TYPE=${2:-"Release"} +CORE_COUNT=$(sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4) + +# ── Android SDK paths (set later, after host protoc build) ─────────── +ANDROID_SDK_ROOT=${ANDROID_SDK_ROOT:-$HOME/Library/Android/sdk} +ANDROID_NDK_HOME=${ANDROID_NDK_HOME:-$(ls -d "$ANDROID_SDK_ROOT/ndk/"* 2>/dev/null | sort -V | tail -1)} + +echo "============================================================" +echo " Android Cross Build & Test" +echo "============================================================" +echo " ABI : $ABI" +echo " API Level : $API_LEVEL" +echo " Build Type : $BUILD_TYPE" +echo " NDK : $ANDROID_NDK_HOME" +echo " CPU cores : $CORE_COUNT" +echo "============================================================" + +if [ ! -d "$ANDROID_NDK_HOME" ]; then + echo "ERROR: ANDROID_NDK_HOME not found at $ANDROID_NDK_HOME" + echo "Please install the NDK via Android Studio or sdkmanager." + exit 1 +fi + +# ── Step 1: build host protoc (using HOST compiler, NOT NDK) ───────── +echo "" +echo ">>> Step 1: Building protoc for host..." +HOST_BUILD_DIR="build_host" + +git submodule foreach --recursive 'git stash --include-untracked' 2>/dev/null || true + +if [ ! -f "$CURRENT_DIR/$HOST_BUILD_DIR/bin/protoc" ]; then + # Explicitly avoid NDK toolchain for host build + cmake -S . -B "$HOST_BUILD_DIR" \ + -DCMAKE_BUILD_TYPE="$BUILD_TYPE" \ + -DCMAKE_TOOLCHAIN_FILE="" \ + -G Ninja + cmake --build "$HOST_BUILD_DIR" --target protoc -j"$CORE_COUNT" +else + echo " (cached — skipping)" +fi +PROTOC_EXECUTABLE=$CURRENT_DIR/$HOST_BUILD_DIR/bin/protoc +echo ">>> Step 1: Done (protoc=$PROTOC_EXECUTABLE)" + +# ── Now export Android env vars for cross-compilation ──────────────── +export ANDROID_SDK_ROOT +export ANDROID_HOME=$ANDROID_SDK_ROOT +export ANDROID_NDK_HOME +export CMAKE_TOOLCHAIN_FILE=$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake + +export PATH=$PATH:$ANDROID_SDK_ROOT/cmdline-tools/latest/bin +export PATH=$PATH:$ANDROID_SDK_ROOT/platform-tools +export PATH=$PATH:$ANDROID_SDK_ROOT/emulator +export PATH=$PATH:$ANDROID_NDK_HOME + +# ── Step 2: cross-compile zvec + tests for Android ─────────────────── +echo "" +echo ">>> Step 2: Cross-compiling zvec for Android ($ABI, API $API_LEVEL)..." + +# reset thirdparty so the cross toolchain can patch cleanly +git submodule foreach --recursive 'git stash --include-untracked' 2>/dev/null || true + +BUILD_DIR="build_android_${ABI}" + +# Force CMake reconfigure to pick up any bazel.cmake changes +rm -f "$BUILD_DIR/CMakeCache.txt" + +cmake -S . -B "$BUILD_DIR" -G Ninja \ + -DANDROID_NDK="$ANDROID_NDK_HOME" \ + -DCMAKE_TOOLCHAIN_FILE="$CMAKE_TOOLCHAIN_FILE" \ + -DANDROID_ABI="$ABI" \ + -DANDROID_NATIVE_API_LEVEL="$API_LEVEL" \ + -DANDROID_STL="c++_static" \ + -DCMAKE_BUILD_TYPE="$BUILD_TYPE" \ + -DBUILD_PYTHON_BINDINGS=OFF \ + -DBUILD_TOOLS=OFF \ + -DENABLE_NATIVE=OFF \ + -DAUTO_DETECT_ARCH=OFF \ + -DCMAKE_INSTALL_PREFIX="$BUILD_DIR/install" \ + -DGLOBAL_CC_PROTOBUF_PROTOC="$PROTOC_EXECUTABLE" + +echo " Building library..." +cmake --build "$BUILD_DIR" -j"$CORE_COUNT" + +# Discover test target names from ctest metadata (before building) +echo " Discovering test targets from ctest..." +TEST_NAMES=() +while IFS= read -r line; do + name=$(echo "$line" | sed -n 's/.*Test[[:space:]]*#[0-9]*:[[:space:]]*//p') + if [ -n "$name" ]; then + TEST_NAMES+=("$name") + fi +done < <(cd "$BUILD_DIR" && ctest --show-only 2>/dev/null || true) + +# Fallback: find *_test targets from CMake cache +if [ ${#TEST_NAMES[@]} -eq 0 ]; then + echo " ctest listing unavailable, scanning ninja targets..." + while IFS= read -r line; do + name=$(echo "$line" | sed -n 's/^\([a-zA-Z0-9_]*_test\): .*/\1/p') + if [ -n "$name" ]; then + TEST_NAMES+=("$name") + fi + done < <(ninja -C "$BUILD_DIR" -t targets all 2>/dev/null || true) +fi + +echo " Found ${#TEST_NAMES[@]} test targets." +if [ ${#TEST_NAMES[@]} -eq 0 ]; then + echo "WARNING: No test targets found. Skipping emulator step." + exit 0 +fi + +# Build all test executables at once (without running ctest) +echo " Building ${#TEST_NAMES[@]} test executables..." +ninja -C "$BUILD_DIR" -j"$CORE_COUNT" "${TEST_NAMES[@]}" + +echo ">>> Step 2: Done" + +# ── Step 3: collect test binaries ──────────────────────────────────── +echo "" +echo ">>> Step 3: Collecting test binaries..." + +TEST_BINS=() +for name in "${TEST_NAMES[@]}"; do + bin_path=$(find "$BUILD_DIR" -name "$name" -type f -perm +111 2>/dev/null | head -1) + if [ -n "$bin_path" ]; then + TEST_BINS+=("$bin_path") + else + echo " WARNING: binary not found for target '$name'" + fi +done + +echo " Found ${#TEST_BINS[@]} test binaries." +if [ ${#TEST_BINS[@]} -eq 0 ]; then + echo "WARNING: No test binaries found. Skipping emulator step." + exit 0 +fi + +# ── Step 4: start emulator ─────────────────────────────────────────── +echo "" +echo ">>> Step 4: Starting Android emulator..." + +EMULATOR_BIN="$ANDROID_SDK_ROOT/emulator/emulator" +ADB_BIN="$ANDROID_SDK_ROOT/platform-tools/adb" +AVD_NAME="zvec_test_avd" + +# Find a usable system image +SYS_IMG="" +for candidate in \ + "$ANDROID_SDK_ROOT/system-images/android-${API_LEVEL}/default/arm64-v8a" \ + "$ANDROID_SDK_ROOT/system-images/android-${API_LEVEL}/google_apis/arm64-v8a" \ + "$ANDROID_SDK_ROOT/system-images/android-${API_LEVEL}/google_apis_playstore/arm64-v8a"; do + if [ -d "$candidate" ]; then + SYS_IMG="$candidate" + break + fi +done + +# If the exact API level isn't found, try any available arm64-v8a image +if [ -z "$SYS_IMG" ]; then + echo " No system image for API $API_LEVEL, searching for any arm64-v8a image..." + SYS_IMG=$(find "$ANDROID_SDK_ROOT/system-images" -type d -name "arm64-v8a" 2>/dev/null | head -1) +fi + +if [ -z "$SYS_IMG" ]; then + echo "ERROR: No arm64-v8a system image found." + echo "Install one via: sdkmanager 'system-images;android-${API_LEVEL};default;arm64-v8a'" + exit 1 +fi +echo " Using system image: $SYS_IMG" + +# Create AVD without avdmanager (write INI files directly) +AVD_DIR="$HOME/.android/avd/${AVD_NAME}.avd" +AVD_INI="$HOME/.android/avd/${AVD_NAME}.ini" + +cleanup_emulator() { + echo "" + echo ">>> Cleaning up emulator..." + $ADB_BIN emu kill 2>/dev/null || true + sleep 2 + # Kill any lingering emulator processes for our AVD + pkill -f "emulator.*${AVD_NAME}" 2>/dev/null || true +} +trap cleanup_emulator EXIT + +mkdir -p "$AVD_DIR" + +# Write the top-level .ini +cat > "$AVD_INI" << EOF +avd.ini.encoding=UTF-8 +path=${AVD_DIR} +path.rel=avd/${AVD_NAME}.avd +target=android-${API_LEVEL} +EOF + +# Write the AVD config +cat > "$AVD_DIR/config.ini" << EOF +AvdId=${AVD_NAME} +PlayStore.enabled=false +abi.type=arm64-v8a +avd.ini.displayname=${AVD_NAME} +avd.ini.encoding=UTF-8 +disk.dataPartition.size=8G +hw.accelerator.isConfigured=true +hw.cpu.arch=arm64 +hw.cpu.ncore=4 +hw.lcd.density=420 +hw.lcd.height=1920 +hw.lcd.width=1080 +hw.ramSize=4096 +image.sysdir.1=${SYS_IMG}/ +tag.display=Default +tag.id=default +EOF + +echo " Created AVD: $AVD_NAME" + +# Kill any existing emulator +$ADB_BIN emu kill 2>/dev/null || true +sleep 1 + +# Start emulator +echo " Launching emulator..." +$EMULATOR_BIN -avd "$AVD_NAME" \ + -no-window -no-audio -no-boot-anim \ + -gpu swiftshader_indirect \ + -netdelay none -netspeed full \ + -memory 4096 \ + -no-snapshot \ + -wipe-data \ + 2>&1 | sed 's/^/ [emulator] /' & +EMULATOR_PID=$! + +# Wait for the device to become reachable +echo " Waiting for device..." +$ADB_BIN wait-for-device + +# Wait for boot_completed +echo " Waiting for boot to complete..." +BOOT_TIMEOUT=300 +BOOT_ELAPSED=0 +while true; do + BOOTED=$($ADB_BIN shell getprop sys.boot_completed 2>/dev/null | tr -d '\r\n ' || true) + if [ "$BOOTED" = "1" ]; then + echo " Emulator is ready! (took ${BOOT_ELAPSED}s)" + break + fi + if [ $BOOT_ELAPSED -ge $BOOT_TIMEOUT ]; then + echo "ERROR: Emulator failed to boot within ${BOOT_TIMEOUT}s" + exit 1 + fi + sleep 3 + BOOT_ELAPSED=$((BOOT_ELAPSED + 3)) +done + +# Print device info +echo "" +echo " Device ABI : $($ADB_BIN shell getprop ro.product.cpu.abi | tr -d '\r')" +echo " ABI list : $($ADB_BIN shell getprop ro.product.cpu.abilist | tr -d '\r')" +echo " CPU info :" +$ADB_BIN shell 'cat /proc/cpuinfo | grep -E "^(Features|flags|processor)"' 2>/dev/null | head -4 | sed 's/^/ /' + +# ── Step 5: run tests on emulator ──────────────────────────────────── +echo "" +echo ">>> Step 5: Running ${#TEST_BINS[@]} unit tests on emulator..." + +DEVICE_TEST_DIR="/data/local/tmp/zvec_tests" +DEVICE_LIB_DIR="$DEVICE_TEST_DIR/lib" +$ADB_BIN shell "mkdir -p $DEVICE_TEST_DIR $DEVICE_LIB_DIR" 2>/dev/null + +# Push all shared libraries that tests may depend on +echo " Pushing shared libraries..." +SO_COUNT=0 +while IFS= read -r so_file; do + $ADB_BIN push "$so_file" "$DEVICE_LIB_DIR/$(basename "$so_file")" > /dev/null 2>&1 + SO_COUNT=$((SO_COUNT + 1)) +done < <(find "$BUILD_DIR/lib" -name "*.so" -type f 2>/dev/null) +echo " Pushed $SO_COUNT shared libraries." + +# Push non-test helper binaries (e.g. data_generator, collection_optimizer) +# These are needed by crash_recovery tests which fork+exec them at runtime. +echo " Pushing helper binaries..." +HELPER_COUNT=0 +for helper_name in data_generator collection_optimizer; do + helper_path=$(find "$BUILD_DIR" -name "$helper_name" -type f -perm +111 ! -name "*_test" 2>/dev/null | head -1) + if [ -n "$helper_path" ]; then + $ADB_BIN push "$helper_path" "$DEVICE_TEST_DIR/$helper_name" > /dev/null 2>&1 + $ADB_BIN shell "chmod 755 $DEVICE_TEST_DIR/$helper_name" 2>/dev/null + HELPER_COUNT=$((HELPER_COUNT + 1)) + fi +done +echo " Pushed $HELPER_COUNT helper binaries." + +TOTAL=${#TEST_BINS[@]} +PASSED=0 +FAILED=0 +FAILED_NAMES=() +IDX=0 + +for test_bin in "${TEST_BINS[@]}"; do + IDX=$((IDX + 1)) + test_name=$(basename "$test_bin") + device_path="$DEVICE_TEST_DIR/$test_name" + # Give each test its own working directory to avoid name collisions + # (e.g. flat_builder_test binary vs flat_builder_test/ directory) + WORK_DIR="$DEVICE_TEST_DIR/workdir_${test_name}" + + echo "" + echo "────────────────────────────────────────" + echo " [$IDX/$TOTAL] $test_name" + echo "────────────────────────────────────────" + + set +e + # Create isolated working directory + $ADB_BIN shell "mkdir -p $WORK_DIR" 2>/dev/null + + # Copy helper binaries into working directory so crash_recovery tests + # (which fork+exec data_generator / collection_optimizer) can find them. + $ADB_BIN shell "for h in $DEVICE_TEST_DIR/data_generator $DEVICE_TEST_DIR/collection_optimizer; do [ -f \$h ] && cp \$h $WORK_DIR/; done" 2>/dev/null + + # Push binary + $ADB_BIN push "$test_bin" "$device_path" > /dev/null 2>&1 + $ADB_BIN shell "chmod 755 $device_path" 2>/dev/null + + # Run test from its own working directory with LD_LIBRARY_PATH + OUTPUT=$($ADB_BIN shell "cd $WORK_DIR && LD_LIBRARY_PATH=$DEVICE_LIB_DIR $device_path 2>&1; echo EXIT_CODE=\$?" 2>&1) + + # Extract exit code from the output + EXIT_CODE=$(echo "$OUTPUT" | grep -o 'EXIT_CODE=[0-9]*' | tail -1 | cut -d= -f2) + set -e + + # Print test output (without the EXIT_CODE line) + echo "$OUTPUT" | grep -v 'EXIT_CODE=' | sed 's/^/ /' + + if [ "$EXIT_CODE" = "0" ]; then + echo " >>> PASSED" + PASSED=$((PASSED + 1)) + else + # Detect "segfault-on-exit" pattern: all gtest assertions passed but + # process crashed during teardown (common with cross-compiled gmock tests). + GTEST_PASSED_LINE=$(echo "$OUTPUT" | grep '\[ PASSED \]' | tail -1) + GTEST_FAILED_LINE=$(echo "$OUTPUT" | grep '\[ FAILED \]' | head -1) + if [ -n "$GTEST_PASSED_LINE" ] && [ -z "$GTEST_FAILED_LINE" ] && \ + { [ "$EXIT_CODE" = "139" ] || [ "$EXIT_CODE" = "134" ] || [ "$EXIT_CODE" = "135" ]; }; then + echo " >>> PASSED (with crash-on-exit, exit code: $EXIT_CODE — ignored)" + PASSED=$((PASSED + 1)) + else + echo " >>> FAILED (exit code: $EXIT_CODE)" + FAILED=$((FAILED + 1)) + FAILED_NAMES+=("$test_name") + fi + fi + + # Clean up binary and working directory to reclaim disk space + $ADB_BIN shell "rm -rf $device_path $WORK_DIR" 2>/dev/null || true +done + +# ── Summary ────────────────────────────────────────────────────────── +echo "" +echo "============================================================" +echo " Test Summary" +echo "============================================================" +echo " Total : $TOTAL" +echo " Passed : $PASSED" +echo " Failed : $FAILED" +if [ $FAILED -gt 0 ]; then + echo "" + echo " Failed tests:" + for name in "${FAILED_NAMES[@]}"; do + echo " - $name" + done +fi +echo "============================================================" + +if [ $FAILED -gt 0 ]; then + exit 1 +fi + +echo "All tests passed!" diff --git a/scripts/build_ios.sh b/scripts/build_ios.sh new file mode 100755 index 0000000..6a63405 --- /dev/null +++ b/scripts/build_ios.sh @@ -0,0 +1,178 @@ +#!/bin/bash +set -e +CURRENT_DIR=$(pwd) + +# Platform options: OS (arm64 device), SIMULATOR64 (x86_64 sim), SIMULATORARM64 (arm64 sim) +PLATFORM=${1:-"OS"} +BUILD_TYPE=${2:-"Release"} +IOS_DEPLOYMENT_TARGET="13.0" + +# Determine architecture based on platform +case "$PLATFORM" in + "OS") + ARCH="arm64" + ;; + "SIMULATOR64") + ARCH="x86_64" + ;; + "SIMULATORARM64") + ARCH="arm64" + ;; + *) + echo "error: Unknown platform '$PLATFORM'" + echo "Usage: $0 [OS|SIMULATOR64|SIMULATORARM64] [Release|Debug]" + echo " OS - Build for iOS device (arm64)" + echo " SIMULATOR64 - Build for iOS Simulator (x86_64)" + echo " SIMULATORARM64- Build for iOS Simulator (arm64, Apple Silicon)" + exit 1 + ;; +esac + +echo "Building zvec for iOS" +echo " Platform: $PLATFORM" +echo " Architecture: $ARCH" +echo " Build Type: $BUILD_TYPE" +echo " iOS Deployment Target: $IOS_DEPLOYMENT_TARGET" + +# step1: use host env to compile protoc +echo "step1: building protoc for host..." + +git submodule foreach --recursive 'git stash --include-untracked' + +HOST_BUILD_DIR="build_host" +mkdir -p $HOST_BUILD_DIR +cd $HOST_BUILD_DIR + +cmake -DCMAKE_BUILD_TYPE="$BUILD_TYPE" .. +make -j protoc +PROTOC_EXECUTABLE=$CURRENT_DIR/$HOST_BUILD_DIR/bin/protoc +cd $CURRENT_DIR + +echo "step1: Done!!!" + +# step2: cross build zvec for iOS +echo "step2: building zvec for iOS..." + +# reset thirdparty directory +git submodule foreach --recursive 'git stash --include-untracked' + +BUILD_DIR="build_ios_${PLATFORM}" +mkdir -p $BUILD_DIR +cd $BUILD_DIR + +# Determine SDK and additional flags based on platform +if [ "$PLATFORM" = "OS" ]; then + SDK_NAME="iphoneos" +else + SDK_NAME="iphonesimulator" +fi + +SDK_PATH=$(xcrun --sdk $SDK_NAME --show-sdk-path) + +echo "configure CMake..." +cmake \ + -DCMAKE_SYSTEM_NAME=iOS \ + -DCMAKE_OSX_DEPLOYMENT_TARGET="$IOS_DEPLOYMENT_TARGET" \ + -DCMAKE_OSX_ARCHITECTURES="$ARCH" \ + -DCMAKE_OSX_SYSROOT="$SDK_PATH" \ + -DCMAKE_BUILD_TYPE="$BUILD_TYPE" \ + -DBUILD_PYTHON_BINDINGS=OFF \ + -DBUILD_TOOLS=OFF \ + -DCMAKE_INSTALL_PREFIX="./install" \ + -DGLOBAL_CC_PROTOBUF_PROTOC=$PROTOC_EXECUTABLE \ + -DIOS=ON \ + ../ + +echo "building..." +CORE_COUNT=$(sysctl -n hw.ncpu) +make -j$CORE_COUNT + +echo "step2: Done!!!" + +# step3: build and run all unit tests on simulator +if [ "$PLATFORM" != "OS" ]; then + echo "step3: building and running unit tests on simulator..." + + make -j$CORE_COUNT unittest + + # Boot simulator + DEVICE_ID=$(xcrun simctl list devices available -j \ + | python3 -c " +import json, sys +data = json.load(sys.stdin) +for runtime, devices in data['devices'].items(): + if 'iOS' in runtime: + for d in devices: + if 'iPhone' in d['name'] and d['isAvailable']: + print(d['udid']) + sys.exit(0) +sys.exit(1) +") + xcrun simctl boot "$DEVICE_ID" + echo "Booted simulator: $DEVICE_ID" + + # Run all test .app bundles + FAILED_TESTS="" + PASSED=0 + TOTAL=0 + + for APP in bin/*_test.app; do + [ -d "$APP" ] || continue + TEST_NAME=$(basename "$APP" .app) + BUNDLE_ID="com.zvec.${TEST_NAME}" + TOTAL=$((TOTAL + 1)) + + echo "--- Running ${TEST_NAME} ---" + xcrun simctl install "$DEVICE_ID" "$APP" + set +e + for attempt in 1 2 3; do + xcrun simctl launch --console "$DEVICE_ID" "$BUNDLE_ID" 2>&1 | tee /tmp/${TEST_NAME}.log + LAUNCH_EXIT=${PIPESTATUS[0]} + if ! grep -q "unknown to FrontBoard" /tmp/${TEST_NAME}.log; then + break + fi + echo "Attempt ${attempt}/3: FrontBoard has not registered ${TEST_NAME} yet, retrying in 3s..." + sleep 3 + done + set -e + + if grep -q '\[ FAILED \]' /tmp/${TEST_NAME}.log; then + echo "FAIL: ${TEST_NAME}" + FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" + elif grep -q '\[ PASSED \]' /tmp/${TEST_NAME}.log; then + PASSED=$((PASSED + 1)) + elif grep -qE 'Failed: 0$' /tmp/${TEST_NAME}.log; then + # c_api_test uses a custom test framework (not GTest) + PASSED=$((PASSED + 1)) + elif [ "$LAUNCH_EXIT" -eq 0 ]; then + echo "WARN: ${TEST_NAME} exited 0 but produced no recognisable test summary" + PASSED=$((PASSED + 1)) + else + echo "FAIL: ${TEST_NAME} exited ${LAUNCH_EXIT} with no test summary" + FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" + fi + done + + # Shutdown simulator + xcrun simctl shutdown "$DEVICE_ID" || true + + echo "" + echo "Test summary: ${PASSED}/${TOTAL} passed" + if [ -n "$FAILED_TESTS" ]; then + echo "Failed tests:${FAILED_TESTS}" + exit 1 + fi + + echo "step3: Done!!!" +else + echo "Skipping tests (device build cannot run on simulator)" +fi + +echo "" +echo "Build completed successfully!" +echo "Output directory: $CURRENT_DIR/$BUILD_DIR" + +# Test On MacOS15 +# 1: xcrun simctl boot "iPhone 16" +# 2: cd $BUILD_DIR +# 3: xcrun simctl launch --console booted com.zvec.collection_test diff --git a/scripts/gcov.sh b/scripts/gcov.sh new file mode 100755 index 0000000..64e810b --- /dev/null +++ b/scripts/gcov.sh @@ -0,0 +1,41 @@ +#!/bin/bash + +project_name=proxima-zvec +gcov_tool=gcov +zip_html=false +output_name=html +keep_info=false + +script_dir=$(cd "$(dirname "$0")"; pwd) +source_base=$(dirname "$script_dir") +filter_list="'*/tests/*' '*/thirdparty/*' '*/deps/*' '*/proto/*' '*/external/*' '*/sqlengine/antlr/gen/*'" + +while getopts t:p:o:zk option; do + case "$option" in + t) + gcov_tool=$OPTARG;; + p) + project_name=$OPTARG;; + o) + output_name=$OPTARG;; + z) + zip_html=true;; + k) + keep_info=true;; + esac +done + +# Process sources +lcov -c -b "$source_base" -d . -o $project_name.lcov.info --gcov-tool=$gcov_tool --no-external --keep-going --ignore-errors unused,mismatch || exit 1 +eval $(echo lcov -r $project_name.lcov.info -o $project_name-filtered.lcov.info $filter_list --ignore-errors unused,mismatch) || exit 1 + +# Gather HTML files +genhtml -t "$project_name" -o $output_name $project_name-filtered.lcov.info || exit 1 +if [ "$keep_info" = false ]; then + rm -rf *.lcov.info +fi + +# Zip HTML files +if $zip_html ; then + zip -r $output_name.zip $output_name/ +fi diff --git a/scripts/generate_standard_tokenizer_unicode.py b/scripts/generate_standard_tokenizer_unicode.py new file mode 100755 index 0000000..1944af3 --- /dev/null +++ b/scripts/generate_standard_tokenizer_unicode.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +# +# Generate Unicode lookup tables for the standard tokenizer. +# +# Usage: +# scripts/generate_standard_tokenizer_unicode.py \ +# --ucd-dir /path/to/Public/17.0.0/ucd \ +# --out src/db/index/column/fts_column/tokenizer/standard_tokenizer_unicode.inc +from __future__ import annotations + +import argparse +import pathlib + +WORD_BREAK_CLASSES = { + "ALetter": "ALetter", + "CR": "CR", + "Double_Quote": "DoubleQuote", + "Extend": "Extend", + "ExtendNumLet": "ExtendNumLet", + "Format": "Format", + "Hebrew_Letter": "HebrewLetter", + "Katakana": "Katakana", + "LF": "LF", + "MidLetter": "MidLetter", + "MidNum": "MidNum", + "MidNumLet": "MidNumLet", + "Newline": "Newline", + "Numeric": "Numeric", + "Regional_Indicator": "RegionalIndicator", + "Single_Quote": "SingleQuote", + "WSegSpace": "WSegSpace", + "ZWJ": "ZWJ", +} + +SCRIPT_CLASSES = { + "Han": "Ideographic", + "Hangul": "Hangul", + "Hiragana": "Hiragana", +} + + +LINE_BREAK_COMPLEX_CONTEXT = { + "Complex_Context", + "SA", +} + + +def parse_codepoint_range(field): + if ".." in field: + start, end = field.split("..", 1) + return int(start, 16), int(end, 16) + cp = int(field, 16) + return cp, cp + + +def parse_property_file(path, accepted_properties): + ranges = [] + with path.open("r", encoding="utf-8") as fin: + for raw_line in fin: + line = raw_line.split("#", 1)[0].strip() + if not line: + continue + fields = [field.strip() for field in line.split(";")] + if len(fields) < 2: + continue + prop = fields[1] + if prop not in accepted_properties: + continue + start, end = parse_codepoint_range(fields[0]) + ranges.append((start, end, accepted_properties[prop])) + return merge_class_ranges(ranges) + + +def parse_range_properties(path, accepted_properties): + ranges = [] + with path.open("r", encoding="utf-8") as fin: + for raw_line in fin: + line = raw_line.split("#", 1)[0].strip() + if not line: + continue + fields = [field.strip() for field in line.split(";")] + if len(fields) < 2: + continue + if fields[1] not in accepted_properties: + continue + ranges.append(parse_codepoint_range(fields[0])) + return merge_ranges(ranges) + + +def merge_class_ranges(ranges): + merged = [] + for start, end, cls in sorted(ranges): + if merged and merged[-1][2] == cls and merged[-1][1] + 1 == start: + merged[-1] = (merged[-1][0], end, cls) + else: + merged.append((start, end, cls)) + return merged + + +def merge_ranges(ranges): + merged = [] + for start, end in sorted(ranges): + if merged and merged[-1][1] + 1 >= start: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((start, end)) + return merged + + +def parse_binary_property(path, property_name): + return parse_range_properties(path, {property_name}) + + +def write_class_table(out, name, ranges): + out.write(f"constexpr UnicodeClassRange {name}[] = {{\n") + for start, end, cls in ranges: + out.write(f" {{0x{start:04X}, 0x{end:04X}, WordBreakClass::{cls}}},\n") + out.write("};\n\n") + + +def write_range_table(out, name, ranges): + out.write(f"constexpr UnicodeRange {name}[] = {{\n") + for start, end in ranges: + out.write(f" {{0x{start:04X}, 0x{end:04X}}},\n") + out.write("};\n\n") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ucd-dir", required=True, type=pathlib.Path) + parser.add_argument("--out", required=True, type=pathlib.Path) + args = parser.parse_args() + + word_break_path = args.ucd_dir / "auxiliary" / "WordBreakProperty.txt" + emoji_data_path = args.ucd_dir / "emoji" / "emoji-data.txt" + line_break_path = args.ucd_dir / "LineBreak.txt" + scripts_path = args.ucd_dir / "Scripts.txt" + + word_break_ranges = parse_property_file(word_break_path, WORD_BREAK_CLASSES) + script_ranges = parse_property_file(scripts_path, SCRIPT_CLASSES) + extended_pictographic_ranges = parse_binary_property( + emoji_data_path, "Extended_Pictographic" + ) + emoji_modifier_base_ranges = parse_binary_property( + emoji_data_path, "Emoji_Modifier_Base" + ) + emoji_modifier_ranges = parse_binary_property(emoji_data_path, "Emoji_Modifier") + line_break_complex_context_ranges = parse_range_properties( + line_break_path, LINE_BREAK_COMPLEX_CONTEXT + ) + + args.out.parent.mkdir(parents=True, exist_ok=True) + with args.out.open("w", encoding="utf-8", newline="\n") as fout: + fout.write("// Generated by scripts/generate_standard_tokenizer_unicode.py\n") + fout.write( + "// Source: Unicode 17.0.0 WordBreakProperty, emoji-data, LineBreak, Scripts.\n" + ) + fout.write( + "// Derived from Unicode data files licensed under Unicode License V3; see NOTICE.\n" + ) + fout.write("// Do not edit by hand.\n\n") + fout.write("// clang-format off\n\n") + write_class_table(fout, "kWordBreakRanges", word_break_ranges) + write_class_table(fout, "kScriptClassRanges", script_ranges) + write_range_table( + fout, "kExtendedPictographicRanges", extended_pictographic_ranges + ) + write_range_table(fout, "kEmojiModifierBaseRanges", emoji_modifier_base_ranges) + write_range_table(fout, "kEmojiModifierRanges", emoji_modifier_ranges) + write_range_table( + fout, "kLineBreakComplexContextRanges", line_break_complex_context_ranges + ) + fout.write("// clang-format on\n") + + +if __name__ == "__main__": + main() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt new file mode 100644 index 0000000..5fc4e34 --- /dev/null +++ b/src/CMakeLists.txt @@ -0,0 +1,171 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +# Retrieve version from git repository +git_version(ZVEC_VERSION ${CMAKE_CURRENT_SOURCE_DIR}) + +# Add repository +cc_directory(ailego) +cc_directory(turbo) +cc_directory(core) +cc_directory(db) +cc_directory(binding) + +# ============================================================================= +# Build ALL-IN-ONE C++ Shared Libraries +# ============================================================================= +# Merges zvec internal static libraries into shared libraries while preserving +# C++ symbols for direct C++ linking. + +include(GNUInstallDirs) +find_package(Threads REQUIRED) + +function(zvec_add_all_in_one_shared TARGET_NAME OUTPUT_NAME) + cmake_parse_arguments(ZVEC_ALLIN "" "" "LIBS" ${ARGN}) + if(NOT ZVEC_ALLIN_LIBS) + message(FATAL_ERROR "zvec_add_all_in_one_shared requires LIBS") + endif() + + foreach(ZVEC_ALLIN_LIB ${ZVEC_ALLIN_LIBS}) + if(NOT TARGET ${ZVEC_ALLIN_LIB}) + message(FATAL_ERROR "Target ${ZVEC_ALLIN_LIB} is required by ${TARGET_NAME}") + endif() + endforeach() + + file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/${TARGET_NAME}_stub.cc + "// Auto-generated stub for ${TARGET_NAME}\n" + ) + + add_library(${TARGET_NAME} SHARED + ${CMAKE_CURRENT_BINARY_DIR}/${TARGET_NAME}_stub.cc + ) + + set_target_properties(${TARGET_NAME} PROPERTIES + OUTPUT_NAME "${OUTPUT_NAME}" + POSITION_INDEPENDENT_CODE ON + WINDOWS_EXPORT_ALL_SYMBOLS ON + ) + if(WIN32) + set_target_properties(${TARGET_NAME} PROPERTIES + ARCHIVE_OUTPUT_NAME "${TARGET_NAME}" + ) + endif() + target_compile_features(${TARGET_NAME} PUBLIC cxx_std_17) + + if(MSVC) + foreach(ZVEC_ALLIN_LIB ${ZVEC_ALLIN_LIBS}) + target_sources(${TARGET_NAME} PRIVATE + $ + ) + endforeach() + + target_link_libraries(${TARGET_NAME} + PRIVATE + ${ZVEC_ALLIN_LIBS} + Threads::Threads + ) + elseif(APPLE) + foreach(ZVEC_ALLIN_LIB ${ZVEC_ALLIN_LIBS}) + list(APPEND ZVEC_ALLIN_WA_OPTIONS + -Wl,-force_load,$ + ) + endforeach() + + target_link_options(${TARGET_NAME} PRIVATE ${ZVEC_ALLIN_WA_OPTIONS}) + target_link_libraries(${TARGET_NAME} + PRIVATE + ${ZVEC_ALLIN_LIBS} + Threads::Threads + ${CMAKE_DL_LIBS} + ) + else() + target_link_libraries(${TARGET_NAME} + PRIVATE + -Wl,--whole-archive + ${ZVEC_ALLIN_LIBS} + -Wl,--no-whole-archive + Threads::Threads + ${CMAKE_DL_LIBS} + ) + + # Try static linking of libstdc++ and libgcc for better portability. + set(ALLIN_STATIC_LIBS_FOUND FALSE) + execute_process( + COMMAND ${CMAKE_CXX_COMPILER} -print-file-name=libstdc++.a + OUTPUT_VARIABLE ALLIN_LIBSTDCPP_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + execute_process( + COMMAND ${CMAKE_CXX_COMPILER} -print-file-name=libgcc.a + OUTPUT_VARIABLE ALLIN_LIBGCC_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + if(EXISTS "${ALLIN_LIBSTDCPP_PATH}" AND EXISTS "${ALLIN_LIBGCC_PATH}") + if(IS_ABSOLUTE "${ALLIN_LIBSTDCPP_PATH}" AND IS_ABSOLUTE "${ALLIN_LIBGCC_PATH}") + set(ALLIN_STATIC_LIBS_FOUND TRUE) + endif() + endif() + if(ALLIN_STATIC_LIBS_FOUND) + target_link_options(${TARGET_NAME} PRIVATE -static-libstdc++ -static-libgcc) + message(STATUS "${TARGET_NAME}: Using static libstdc++ and libgcc") + else() + message(STATUS "${TARGET_NAME}: Using dynamic linking for libstdc++ and libgcc") + endif() + endif() + + target_include_directories(${TARGET_NAME} + PUBLIC + $ + $ + $ + PRIVATE + ${PROJECT_SOURCE_DIR}/src + ) + + # Strip symbols in release builds to reduce library size. + if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") + if(UNIX AND NOT APPLE) + # add_custom_command(TARGET ${TARGET_NAME} POST_BUILD + # COMMAND ${CMAKE_STRIP} $ + # COMMENT "Stripping symbols from $" + # ) + elseif(APPLE) + add_custom_command(TARGET ${TARGET_NAME} POST_BUILD + COMMAND /usr/bin/strip -x $ + COMMENT "Stripping symbols from $" + ) + endif() + endif() + + install(TARGETS ${TARGET_NAME} + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) +endfunction() + +if(BUILD_ZVEC_AILEGO_SHARED) + zvec_add_all_in_one_shared(zvec_ailego_shared zvec_ailego + LIBS + zvec_ailego + ) +endif() + +if(BUILD_ZVEC_CORE_SHARED) + zvec_add_all_in_one_shared(zvec_core_shared zvec_core + LIBS + zvec_core + zvec_ailego + zvec_turbo + ) +endif() + +if(BUILD_ZVEC_SHARED) + zvec_add_all_in_one_shared(zvec_shared zvec + LIBS + zvec + zvec_core + zvec_ailego + zvec_turbo + ) +endif() diff --git a/src/ailego/CMakeLists.txt b/src/ailego/CMakeLists.txt new file mode 100644 index 0000000..29cf22c --- /dev/null +++ b/src/ailego/CMakeLists.txt @@ -0,0 +1,125 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +find_package(Threads REQUIRED) + +if(UNIX AND NOT APPLE AND NOT ANDROID) + find_library(LIB_RT NAMES rt) +else() + set(LIB_RT "") +endif() + +git_version(GIT_SRCS_VER ${CMAKE_CURRENT_SOURCE_DIR}) +file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h) + +set(EXTRA_LIBS ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS}) + +if(UNIX AND NOT APPLE) + list(APPEND EXTRA_LIBS ${LIB_RT}) +endif() + +if(NOT ANDROID AND AUTO_DETECT_ARCH) + if(HOST_ARCH MATCHES "^(x86|x64)$") + setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512 MATH_MARCH_FLAG_AVX512FP16) + message(STATUS "best compiler march - sse: " ${MATH_MARCH_FLAG_SSE} ", avx2: " ${MATH_MARCH_FLAG_AVX2} ", avx512: " ${MATH_MARCH_FLAG_AVX512} ", avx512fp16: " ${MATH_MARCH_FLAG_AVX512FP16}) + + file(GLOB_RECURSE MATH_FILES_SSE + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_sse.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_sse.c + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_sse.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_sse.c + ) + + file(GLOB_RECURSE MATH_FILES_AVX2 + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx2.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx2.c + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx2.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx2.c + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx.c + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx.c + ) + + file(GLOB_RECURSE MATH_FILES_AVX512 + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512.c + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512.c + ) + + file(GLOB_RECURSE MATH_FILES_AVX512FP16 + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.c + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512fp16.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512fp16.c + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.c + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512fp16.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512fp16.c + ) + + foreach(MATH_FILE ${MATH_FILES_SSE}) + set_source_files_properties( + ${MATH_FILE} + PROPERTIES + COMPILE_FLAGS "${MATH_MARCH_FLAG_SSE}" + ) + endforeach() + + foreach(MATH_FILE ${MATH_FILES_AVX2}) + set_source_files_properties( + ${MATH_FILE} + PROPERTIES + COMPILE_FLAGS "${MATH_MARCH_FLAG_AVX2}" + ) + endforeach() + + foreach(MATH_FILE ${MATH_FILES_AVX512}) + set_source_files_properties( + ${MATH_FILE} + PROPERTIES + COMPILE_FLAGS "${MATH_MARCH_FLAG_AVX512}" + ) + endforeach() + + foreach(MATH_FILE ${MATH_FILES_AVX512FP16}) + set_source_files_properties( + ${MATH_FILE} + PROPERTIES + COMPILE_FLAGS "${MATH_MARCH_FLAG_AVX512FP16}" + ) + endforeach() + elseif (HOST_ARCH MATCHES "^(arm|arm64)$") + if(MSVC) + return() + endif() + set(MATH_MARCH_FLAG_NEON "-march=armv8-a") + + file(GLOB_RECURSE MATH_FILES_NEON + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.c + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.c + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_neon.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math/*_neon.c + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_neon.cc + ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_neon.c + ) + + foreach(MATH_FILE ${MATH_FILES_NEON}) + set_source_files_properties( + ${MATH_FILE} + PROPERTIES + COMPILE_FLAGS "${MATH_MARCH_FLAG_NEON}" + ) + endforeach() + endif() +endif() + +cc_library( + NAME zvec_ailego STATIC STRICT PACKED + SRCS ${ALL_SRCS} + LIBS ${EXTRA_LIBS} + VERSION "${GIT_SRCS_VER}" +) diff --git a/src/ailego/algorithm/binary_quantizer.cc b/src/ailego/algorithm/binary_quantizer.cc new file mode 100644 index 0000000..3a255aa --- /dev/null +++ b/src/ailego/algorithm/binary_quantizer.cc @@ -0,0 +1,77 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "binary_quantizer.h" +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +//! Feed the training data +bool BinaryQuantizer::feed(const float *vec, size_t dim) { + for (size_t i = 0; i < dim; ++i) { + data_.emplace_back(vec[i]); + } + return true; +} + +//! Train the quantizer +bool BinaryQuantizer::train(void) { + return true; +} + +//! Quantize data: encode the float input to uint32_t output +void BinaryQuantizer::encode(const float *in, size_t dim, uint32_t *out) const { + for (size_t i = 0; i < dim; i += 32) { + size_t remain = i + 32 <= dim ? 32 : dim - i; + uint32_t data = 0; + uint32_t mask = 1; + + for (size_t j = 0; j < remain; j++) { + if (in[i + j] >= threshold_) { + data |= mask; + } + + mask <<= 1; + } + + *out = data; + out++; + } +} + +//! De-quantize data: decode the input uint32_t to float output +//! bit value 1 will be mapped to 1.0 +//! bit value 0 will be mapped to -1.0 +void BinaryQuantizer::decode(const uint32_t *in, size_t dim, float *out) const { + for (size_t i = 0; i < dim; ++i) { + uint8_t bit = (in[i >> 5] >> (i & 31)) & 0x01; + + if (bit == 1) { + out[i] = 1.0f; + } else { + out[i] = -1.0f; + } + + // std::cout << "dim: " << i << ", value: " << (size_t)bit << std::endl; + } +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/algorithm/binary_quantizer.h b/src/ailego/algorithm/binary_quantizer.h new file mode 100644 index 0000000..f2dd847 --- /dev/null +++ b/src/ailego/algorithm/binary_quantizer.h @@ -0,0 +1,69 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec { +namespace ailego { + +/*! Binary Quantization Algorithm + */ +class BinaryQuantizer { + public: + //! Constructor + BinaryQuantizer(void) {} + + //! Feed the training data + bool feed(const float *vec, size_t dim); + + //! Train the quantizer + bool train(void); + + //! Quantize data: encode the float input to uint32_t output + void encode(const float *in, size_t dim, uint32_t *out) const; + + //! De-quantize data: decode the input uint32_t to float output + void decode(const uint32_t *in, size_t dim, float *out) const; + + //! Get encoded elements in type of uint32_t + static size_t EncodedSizeInBinary32(size_t dim) { + return (dim + 31) / 32; + } + + //! Set quantization threshold + void set_threshold(float threshold) { + threshold_ = threshold; + } + + //! Get quantization threshold + float threshold(void) const { + return threshold_; + } + + private: + //! Disable them + BinaryQuantizer(const BinaryQuantizer &) = delete; + BinaryQuantizer &operator=(const BinaryQuantizer &) = delete; + + private: + //! Members + std::vector data_{}; + float threshold_{0.0f}; +}; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/algorithm/integer_quantizer.cc b/src/ailego/algorithm/integer_quantizer.cc new file mode 100644 index 0000000..a637306 --- /dev/null +++ b/src/ailego/algorithm/integer_quantizer.cc @@ -0,0 +1,466 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "integer_quantizer.h" +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +//! Make smooth the distribution to eliminate zero in hist +static inline void MakeSmooth(std::vector &dist) { + constexpr float epsilon = std::numeric_limits::epsilon(); + + // L1 Normalize first + float norm = 1.0f; + Normalizer::L1(dist.data(), dist.size(), &norm); + + size_t zero_count = std::count_if(dist.begin(), dist.end(), [](float val) { + return (std::abs(val) < std::numeric_limits::epsilon()); + }); + size_t nonzero_count = dist.size() - zero_count; + + // Double check + if (nonzero_count == 0 || zero_count == 0) { + return; + } + + float y = epsilon * zero_count / static_cast(nonzero_count); + for (auto &it : dist) { + if (std::abs(it) < epsilon) { + it += epsilon; + } else { + it -= y; + } + } // end of for +} + +//! Compute the Entropy of distribution p/q by Kullback-Leibler Divergence +static inline double ComputeKlDivergence(const std::vector &p, + const std::vector &q) { + if (p.size() != q.size() || p.size() == 0) { + return std::numeric_limits::max(); + } + + double v = 0.0f; + for (size_t i = 0; i != p.size(); ++i) { + if (p[i] == 0 || q[i] == 0) { + return std::numeric_limits::max(); + } + v += p[i] * std::log(static_cast(p[i]) / static_cast(q[i])); + } + return v; +} + +//! Expand the quantization distribution to origin distribution in +//! [-threshold, threshold] +static inline void ExpandCandidateDistribution( + const std::vector &distribution, + const std::vector &quantized_distribution, size_t threshold, + std::vector *expand_distribution) { + expand_distribution->resize(threshold * 2, 0); + float merged_cnt = static_cast(expand_distribution->size()) / + quantized_distribution.size(); + size_t left_boundary = distribution.size() / 2 - threshold; + + for (size_t i = 0; i < quantized_distribution.size(); ++i) { + float start = i * merged_cnt; + float end = start + merged_cnt; + const size_t start_ceil = static_cast(std::ceil(start)); + const size_t end_floor = static_cast(std::floor(end)); + float left_ratio = static_cast(start_ceil) - start; + float right_ratio = end - static_cast(end_floor); + float nonzero_count = 0; + + //! Count the non-zeros bins, if the histogram bin is partially included, + //! non-zero bins is also partially counted + if (left_ratio > 0 && left_boundary + start_ceil > 0) { + if (distribution[left_boundary + start_ceil - 1] != 0) { + nonzero_count += left_ratio; + } + } + if (right_ratio > 0 && left_boundary + end_floor < distribution.size()) { + if (distribution[left_boundary + end_floor] != 0) { + nonzero_count += right_ratio; + } + } + for (size_t j = start_ceil; j < end_floor; j++) { + nonzero_count += distribution[left_boundary + j] != 0; + } + if (nonzero_count == 0) { + continue; + } + + //! expand the quantized value + float value = quantized_distribution[i] / nonzero_count; + if (left_ratio > 0 && start_ceil > 0) { + (*expand_distribution)[start_ceil - 1] += value * left_ratio; + } + if (right_ratio > 0 && end_floor < expand_distribution->size()) { + (*expand_distribution)[end_floor] += value * right_ratio; + } + for (size_t j = start_ceil; j < end_floor; j++) { + if (distribution[left_boundary + j] != 0) { + (*expand_distribution)[j] = value; + } + } // end of for + } // end of for +} + +/*! Compute quantization threshold bins + * Implement Int8 Quantization Algorithm ref: + * http://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf + */ +static inline size_t ComputeThreshold(const std::vector &hist, + const size_t target_bins) { + std::vector P_distribution(hist.size()); + size_t zero_point_index = hist.size() / 2; + + size_t start_bin = target_bins / 2; + size_t end_bin = hist.size() / 2; + size_t negative_outliers_count = 0; + size_t positive_outliers_count = 0; + double min_divergence = std::numeric_limits::max(); + size_t target_threshold = end_bin; + + for (size_t threshold = start_bin; threshold <= end_bin; ++threshold) { + negative_outliers_count += hist[zero_point_index - threshold]; + positive_outliers_count += hist[zero_point_index + threshold - 1]; + } + + //! for each zero-axised quantization range: [-threshold, threshold], search + //! the best solution + for (size_t threshold = start_bin; threshold <= end_bin; ++threshold) { + P_distribution.resize(threshold * 2); + auto p_hist = &hist[zero_point_index - threshold]; + for (size_t i = 0; i != P_distribution.size(); ++i) { + P_distribution[i] = static_cast(p_hist[i]); + } + + negative_outliers_count -= hist[zero_point_index - threshold]; + positive_outliers_count -= hist[zero_point_index + threshold - 1]; + P_distribution[0] += negative_outliers_count; + P_distribution[P_distribution.size() - 1] += positive_outliers_count; + + //! Quantize the bins in range [-threshold, threshold] to target_bins + std::vector Q_distribution(target_bins, 0); + float merged_cnt = static_cast(threshold * 2) / target_bins; + size_t left_boundary = zero_point_index - threshold; + for (size_t i = 0; i < target_bins; ++i) { + float start = i * merged_cnt; + float end = start + merged_cnt; + const size_t start_ceil = static_cast(std::ceil(start)); + const size_t end_floor = static_cast(std::floor(end)); + if (left_boundary + start_ceil > 0) { + Q_distribution[i] += + ((float)start_ceil - start) * hist[left_boundary + start_ceil - 1]; + } + if (left_boundary + end_floor < hist.size()) { + Q_distribution[i] += + (end - (float)end_floor) * hist[left_boundary + end_floor]; + } + + for (size_t j = start_ceil; j < end_floor; j++) { + Q_distribution[i] += hist[left_boundary + j]; + } + } + std::vector Q_expand_distribution; + ExpandCandidateDistribution(hist, Q_distribution, threshold, + &Q_expand_distribution); + + //! Compute Kullback-Leibler Divergence, normalize the smooth the data + //! first. Ref: http://hanj.cs.illinois.edu/cs412/bk3/KL-divergence.pdf + MakeSmooth(P_distribution); + MakeSmooth(Q_expand_distribution); + double divergence = + ComputeKlDivergence(P_distribution, Q_expand_distribution); + + if (divergence < min_divergence) { + min_divergence = divergence; + target_threshold = threshold; + } + } + return target_threshold; +} + +// Quantize the value in range +template +static inline float QuantizeValue(float val, float scale, float bias) { + val = (val + bias) * scale; + + if (val > RANGE_MAX) { + val = RANGE_MAX; + } else if (val < RANGE_MIN) { + val = RANGE_MIN; + } + return val; +} + +// Init the historgram params +#define INIT_HISTOGRAM() \ + { \ + if (histogram_bins_ == 0) { \ + size_t range = non_bias_ \ + ? std::max(std::abs(MIN_VALUE), std::abs(MAX_VALUE)) \ + : (MAX_VALUE - MIN_VALUE); \ + histogram_bins_ = std::max(4096u, range * 8); \ + } \ + histogram_.resize((histogram_bins_ + 1) >> 1 << 1); \ + if (non_bias_) { \ + bias_ = 0.0f; \ + auto val = std::max(std::abs(max_), std::abs(min_)); \ + left_boundary_ = -val; \ + hist_interval_ = (val * 2) / static_cast(histogram_.size()); \ + } else { \ + bias_ = -static_cast(min_ + (max_ - min_) * 0.5); \ + left_boundary_ = min_; \ + hist_interval_ = (max_ - min_) / static_cast(histogram_.size()); \ + } \ + } + +// Feed vector and update the historgram +#define UPDATE_HISTOGRAM(vec, dim) \ + { \ + if (max_ < min_) { \ + return false; \ + } \ + if (histogram_.size() == 0) { \ + INIT_HISTOGRAM() \ + } \ + for (size_t i = 0; i < dim; ++i) { \ + ssize_t index = 0; \ + if (hist_interval_ > 0.0) { \ + index = \ + static_cast((vec[i] - left_boundary_) / hist_interval_); \ + } \ + if (index < 0) { \ + index = 0; \ + } else if ((size_t)index >= histogram_.size()) { \ + index = histogram_.size() - 1; \ + } \ + ailego_assert_with((size_t)index < histogram_.size(), "Invalid index"); \ + histogram_[index] += 1; \ + } \ + return true; \ + } + +// Train the quantizer +#define TRAIN_QUANTIZER() \ + { \ + auto sum = std::accumulate(histogram_.begin(), histogram_.end(), 0); \ + if (sum == 0) { \ + return false; \ + } \ + size_t target_bins = \ + ailego_align(static_cast(MAX_VALUE - MIN_VALUE), 2); \ + auto threshold_bins = ComputeThreshold(histogram_, target_bins); \ + auto threshold = \ + (static_cast(threshold_bins) + 0.5f) * hist_interval_; \ + scale_ = target_bins / 2 / threshold; \ + if (!non_bias_) { \ + bias_ += (MAX_VALUE + MIN_VALUE) * 0.5f / scale_; \ + } \ + scale_reciprocal_ = 1 / scale_; \ + return true; \ + } + +// Feed the INT16 quantizer +bool EntropyInt16Quantizer::feed(const float *vec, size_t dim) { + UPDATE_HISTOGRAM(vec, dim) +} + +// Train the INT16 quantizer +bool EntropyInt16Quantizer::train(void) { + TRAIN_QUANTIZER() +} + +// Encode to INT16 +void EntropyInt16Quantizer::encode(const float *in, size_t dim, + int16_t *out) const { + for (size_t i = 0; i < dim; ++i) { + out[i] = static_cast( + std::round(QuantizeValue(in[i], scale_, bias_))); + } +} + +// Decode from INT16 +void EntropyInt16Quantizer::decode(const int16_t *in, size_t dim, + float *out) const { + for (size_t i = 0; i < dim; ++i) { + out[i] = in[i] * this->scale_reciprocal() - this->bias(); + } +} + +// Feed the UINT16 quantizer +bool EntropyUInt16Quantizer::feed(const float *vec, size_t dim) { + UPDATE_HISTOGRAM(vec, dim) +} + +// Train the UINT16 quantizer +bool EntropyUInt16Quantizer::train(void) { + TRAIN_QUANTIZER() +} + +// Encode to UINT16 +void EntropyUInt16Quantizer::encode(const float *in, size_t dim, + uint16_t *out) const { + for (size_t i = 0; i < dim; ++i) { + out[i] = static_cast_from_float_to_uint16( + std::round(QuantizeValue(in[i], scale_, bias_))); + } +} + +// Decode from INT16 +void EntropyUInt16Quantizer::decode(const uint16_t *in, size_t dim, + float *out) const { + for (size_t i = 0; i < dim; ++i) { + out[i] = in[i] * this->scale_reciprocal() - this->bias(); + } +} + +// Feed the INT8 quantizer +bool EntropyInt8Quantizer::feed(const float *vec, size_t dim) { + UPDATE_HISTOGRAM(vec, dim) +} + +// Train the INT8 quantizer +bool EntropyInt8Quantizer::train(void) { + TRAIN_QUANTIZER() +} + +// Encode to INT8 +void EntropyInt8Quantizer::encode(const float *in, size_t dim, + int8_t *out) const { + for (size_t i = 0; i < dim; ++i) { + out[i] = static_cast( + std::round(QuantizeValue(in[i], scale_, bias_))); + } +} + +// Decode from INT8 +void EntropyInt8Quantizer::decode(const int8_t *in, size_t dim, + float *out) const { + for (size_t i = 0; i < dim; ++i) { + out[i] = in[i] * this->scale_reciprocal() - this->bias(); + } +} + +// Feed the UINT8 quantizer +bool EntropyUInt8Quantizer::feed(const float *vec, size_t dim) { + UPDATE_HISTOGRAM(vec, dim) +} + +// Train the UINT8 quantizer +bool EntropyUInt8Quantizer::train(void) { + TRAIN_QUANTIZER() +} + +// Encode to INT8 +void EntropyUInt8Quantizer::encode(const float *in, size_t dim, + uint8_t *out) const { + for (size_t i = 0; i < dim; ++i) { + out[i] = static_cast_from_float_to_uint8( + std::round(QuantizeValue(in[i], scale_, bias_))); + } +} + +// Decode from UINT8 +void EntropyUInt8Quantizer::decode(const uint8_t *in, size_t dim, + float *out) const { + for (size_t i = 0; i < dim; ++i) { + out[i] = in[i] * this->scale_reciprocal() - this->bias(); + } +} + +// Feed the INT4 quantizer +bool EntropyInt4Quantizer::feed(const float *vec, size_t dim) { + UPDATE_HISTOGRAM(vec, dim) +} + +// Train the INT4 quantizer +bool EntropyInt4Quantizer::train(void) { + TRAIN_QUANTIZER() +} + +// Encode to INT4 +void EntropyInt4Quantizer::encode(const float *in, size_t dim, + uint8_t *out) const { + ailego_assert_with(dim % 2 == 0, "Dimension must be aligned with 2"); + + for (size_t i = 0; i < dim; i += 2) { + float lo = QuantizeValue(in[i], scale_, bias_); + float hi = QuantizeValue(in[i + 1], scale_, bias_); + out[i / 2] = (static_cast_from_float_to_uint8(std::round(hi)) << 4) | + (static_cast_from_float_to_uint8(std::round(lo)) & 0xF); + } +} + +// Decode from INT4 +void EntropyInt4Quantizer::decode(const uint8_t *in, size_t dim, + float *out) const { + ailego_assert_with(dim % 2 == 0, "Dimension must be aligned with 2"); + + size_t size = dim / 2; + for (size_t i = 0; i < size; i += 1) { + uint8_t v = in[i]; + int8_t lo = (static_cast(v << 4) >> 4); + int8_t hi = (static_cast(v & 0xf0) >> 4); + out[2 * i] = lo * this->scale_reciprocal() - this->bias(); + out[2 * i + 1] = hi * this->scale_reciprocal() - this->bias(); + } +} + +// Feed the UINT4 quantizer +bool EntropyUInt4Quantizer::feed(const float *vec, size_t dim) { + UPDATE_HISTOGRAM(vec, dim) +} + +// Train the UINT4 quantizer +bool EntropyUInt4Quantizer::train(void) { + TRAIN_QUANTIZER() +} + +// Encode to INT4 +void EntropyUInt4Quantizer::encode(const float *in, size_t dim, + uint8_t *out) const { + ailego_assert_with(dim % 2 == 0, "Dimension must be aligned with 2"); + + for (size_t i = 0; i < dim; i += 2) { + float lo = QuantizeValue(in[i], scale_, bias_); + float hi = QuantizeValue(in[i + 1], scale_, bias_); + out[i / 2] = (static_cast_from_float_to_uint8(std::round(hi)) << 4) | + (static_cast_from_float_to_uint8(std::round(lo)) & 0xF); + } +} + +// Decode from INT4 +void EntropyUInt4Quantizer::decode(const uint8_t *in, size_t dim, + float *out) const { + ailego_assert_with(dim % 2 == 0, "Dimension must be aligned with 2"); + + size_t size = dim / 2; + for (size_t i = 0; i < size; i += 1) { + uint8_t v = in[i]; + out[2 * i] = (v & 0xf) * this->scale_reciprocal() - this->bias(); + out[2 * i + 1] = (v >> 4) * this->scale_reciprocal() - this->bias(); + } +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/algorithm/integer_quantizer.h b/src/ailego/algorithm/integer_quantizer.h new file mode 100644 index 0000000..cdaf38a --- /dev/null +++ b/src/ailego/algorithm/integer_quantizer.h @@ -0,0 +1,238 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! Entropy-based Integer Quantization Algorithm + */ +template +class EntropyIntegerQuantizer { + public: + //! Primitive Built-in Types to store the quantized data + using ValueType = typename std::remove_cv::type; + + //! Constants + constexpr static int MIN_VALUE = RANGE_MIN; + constexpr static int MAX_VALUE = RANGE_MAX; + + // Check supporting type + static_assert(std::is_integral::value, "ValueType must be integral"); + + // Check template values + static_assert(RANGE_MIN < RANGE_MAX, "Invalid value range"); + + //! Constructor + EntropyIntegerQuantizer(void) {} + + //! Set histogram bins in train + void set_histogram_bins(size_t bins) { + if (bins > (RANGE_MAX - RANGE_MIN)) { + histogram_bins_ = bins; + } + } + + //! Set quantization params scale + void set_scale(float val) { + if (val > 0.0f) { + scale_ = val; + scale_reciprocal_ = 1 / scale_; + } + } + + //! Set quantization params bias + void set_bias(float val) { + bias_ = val; + } + + //! Set quantization params max + void set_max(float val) { + max_ = val; + } + + //! Set quantization params min + void set_min(float val) { + min_ = val; + } + + //! Set quantization params non bias + void set_non_bias(bool val) { + non_bias_ = val; + } + + //! Get histogram bins in train + size_t histogram_bins(void) const { + return histogram_bins_; + } + + //! Get quantization params scale + float scale(void) const { + return scale_; + } + + //! Get quantization params bias + float bias(void) const { + return bias_; + } + + //! Get quantization params max + float max(void) const { + return max_; + } + + //! Get quantization params min + float min(void) const { + return min_; + } + + //! Get quantization params non bias + bool non_bias(void) const { + return non_bias_; + } + + //! Retrieve the scale reciprocal for decoding + float scale_reciprocal(void) const { + return scale_reciprocal_; + } + + protected: + //! Disable them + EntropyIntegerQuantizer(const EntropyIntegerQuantizer &) = delete; + EntropyIntegerQuantizer &operator=(const EntropyIntegerQuantizer &) = delete; + + //! Members + size_t histogram_bins_{0}; + float hist_interval_{1.0f}; + float max_{std::numeric_limits::min()}; + float min_{std::numeric_limits::max()}; + float bias_{0.0f}; + float scale_{0.0f}; + float scale_reciprocal_{0.0f}; + float left_boundary_{0.0f}; + bool non_bias_{false}; + std::vector histogram_{}; +}; + +/*! INT16 Quantizer + */ +class EntropyInt16Quantizer + : public EntropyIntegerQuantizer { + public: + //! Feed the training data + bool feed(const float *vec, size_t dim); + + //! Train the quantizer + bool train(void); + + //! Encode float vector to int16 + void encode(const float *in, size_t dim, ValueType *out) const; + + //! Decode to float vector from int16 + void decode(const ValueType *in, size_t dim, float *out) const; +}; + +/*! UINT16 Quantizer + */ +class EntropyUInt16Quantizer + : public EntropyIntegerQuantizer { + public: + //! Feed the training data + bool feed(const float *vec, size_t dim); + + //! Train the quantizer + bool train(void); + + //! Encode float vector to uint16 + void encode(const float *in, size_t dim, ValueType *out) const; + + //! Decode to float vector from uint16 + void decode(const ValueType *in, size_t dim, float *out) const; +}; + +/*! INT8 Quantizer + */ +class EntropyInt8Quantizer : public EntropyIntegerQuantizer { + public: + //! Feed the training data + bool feed(const float *vec, size_t dim); + + //! Train the quantizer + bool train(void); + + //! Encode float vector to int8 + void encode(const float *in, size_t dim, ValueType *out) const; + + //! Decode to float vector from int8 + void decode(const ValueType *in, size_t dim, float *out) const; +}; + +/*! UINT8 Quantizer + */ +class EntropyUInt8Quantizer : public EntropyIntegerQuantizer { + public: + //! Feed the training data + bool feed(const float *vec, size_t dim); + + //! Train the quantizer + bool train(void); + + //! Encode float vector to uint8 + void encode(const float *in, size_t dim, ValueType *out) const; + + //! Decode to float vector from uint8 + void decode(const ValueType *in, size_t dim, float *out) const; +}; + +/*! INT4 Quantizer + */ +class EntropyInt4Quantizer : public EntropyIntegerQuantizer { + public: + //! Feed the training data + bool feed(const float *vec, size_t dim); + + //! Train the quantizer + bool train(void); + + //! Encode float vector to int4 + void encode(const float *in, size_t dim, ValueType *out) const; + + //! Decode to float vector from int4 + void decode(const ValueType *in, size_t dim, float *out) const; +}; + +/*! UINT4 Quantizer + */ +class EntropyUInt4Quantizer : public EntropyIntegerQuantizer { + public: + //! Feed the training data + bool feed(const float *vec, size_t dim); + + //! Train the quantizer + bool train(void); + + //! Encode float vector to uint4 + void encode(const float *in, size_t dim, ValueType *out) const; + + //! Decode to float vector from uint4 + void decode(const ValueType *in, size_t dim, float *out) const; +}; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/algorithm/kmeans.h b/src/ailego/algorithm/kmeans.h new file mode 100644 index 0000000..1fa5b2f --- /dev/null +++ b/src/ailego/algorithm/kmeans.h @@ -0,0 +1,1055 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "lloyd_cluster.h" + +namespace zvec { +namespace ailego { + +/*! K-MC2 Centroids Generator + */ +template +class Kmc2CentroidsGenerator { + public: + //! Type of values + using OwnerType = typename std::decay::type; + using ContainerType = typename OwnerType::ContainerType; + using ContextType = typename OwnerType::ContextType; + using ValueType = typename OwnerType::ValueType; + using StoreType = typename OwnerType::StoreType; + using ThreadPoolType = TPool; + + //! constexpr variables + constexpr static size_t BatchCount = OwnerType::BatchCount; + + //! Generate centroids + void operator()(OwnerType *owner, ThreadPoolType &pool) const { + if (chain_length_ == 0) { + this->init_centroids_random(owner); + } else if (!assumption_free_) { + this->init_centroids_kmc2(owner, pool); + } else { + this->init_centroids_afkmc2(owner, pool); + } + } + + //! Retrieve the markov chain length + size_t chain_length(void) const { + return chain_length_; + } + + //! Set the mutable markov chain length + void set_chain_length(size_t len) { + chain_length_ = len; + } + + //! Retrieve assumption free option + bool assumption_free(void) const { + return assumption_free_; + } + + //! Set the assumption free option + void set_assumption_free(bool val) { + assumption_free_ = val; + } + + protected: + //! Initialize centroids randomly + void init_centroids_random(OwnerType *owner) const { + RandomSelectBenches(owner->feature_cache(), owner->feature_matrix(), + owner->k_value(), owner->mutable_centroids()); + } + + //! Initialize centroids with K-MC2 + void init_centroids_kmc2(OwnerType *owner, ThreadPoolType &pool) const { + const auto &matrix = owner->feature_matrix(); + const auto &cache = owner->feature_cache(); + auto *centroids = owner->mutable_centroids(); + + std::mt19937 mt((std::random_device())()); + + std::uniform_real_distribution dist(0.0, 1.0); + + ContainerType benches(cache.dimension()); + std::vector scores; + + // Sample first center uniformly + RandomSelectBenches(cache, matrix, 1, centroids); + + // Make a thread group + auto group = pool.make_group(); + + for (size_t i = 1, k = owner->k_value(); i < k; ++i) { + RandomSelectBenches(cache, matrix, chain_length_, &benches); + + // Update bench scores + scores.resize(benches.count()); + for (size_t j = 0; j != scores.size(); ++j) { + group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateBenchScores, + centroids, benches[j], &scores[j])); + } + group->wait_finish(); + + //! Select the better centroid randomly + float x = scores[0]; + size_t xj = 0; + for (size_t j = 1; j != scores.size(); ++j) { + float y = scores[j]; + + if (x == 0.0f || x * dist(mt) < y) { + x = y; + xj = j; + } + } + centroids->append(benches[xj], benches.dimension()); + } // end of for + } + + //! Initialize centroids with K-MC2 + void init_centroids_afkmc2(OwnerType *owner, ThreadPoolType &pool) const { + const auto &matrix = owner->feature_matrix(); + const auto &cache = owner->feature_cache(); + + // Probability + std::vector probs(matrix.count() + cache.count()); + + // Sample first center uniformly + RandomSelectBenches(cache, matrix, 1, owner->mutable_centroids()); + + // Make a thread group + auto group = pool.make_group(); + if (!matrix.empty()) { + size_t n = matrix.count() / BatchCount; + size_t c = std::max(n / pool.count() / 2u, 1u); + size_t m = n / c * c; + + for (size_t i = 0; i != m; i += c) { + group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateMatrixScores, + owner, i, i + c, &probs[0])); + } + for (size_t i = m; i != n; i += 1) { + group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateMatrixScores, + owner, i, i + 1, &probs[0])); + } + } + if (!cache.empty()) { + group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateCacheScores, + owner, &probs[matrix.count()])); + } + group->wait_finish(); + + // Update probabilities + double p_sum = std::accumulate(probs.begin(), probs.end(), 0.0); + for (auto it = probs.begin(); it != probs.end(); ++it) { + *it = static_cast((*it / p_sum + 1.0 / probs.size()) * 0.5); + } + + std::mt19937 mt((std::random_device())()); + std::uniform_real_distribution dist(0.0, 1.0); + ContainerType benches(cache.dimension()); + std::vector scores; + std::vector bench_probs; + + for (size_t i = 1; i < owner->k_value(); ++i) { + RandomSelectBenches(cache, matrix, chain_length_, probs, &benches, + &bench_probs); + + // Update bench scores + scores.resize(benches.count()); + for (size_t j = 0; j != scores.size(); ++j) { + group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateBenchScores, + owner->mutable_centroids(), benches[j], + &scores[j])); + } + group->wait_finish(); + + // Update scores with probabilities + for (size_t j = 0; j != scores.size(); ++j) { + scores[j] /= bench_probs[j]; + } + + //! Select the better centroid randomly + float x = scores[0]; + size_t xj = 0; + for (size_t j = 1; j != scores.size(); ++j) { + float y = scores[j]; + + if (x == 0.0f || x * dist(mt) < y) { + x = y; + xj = j; + } + } + owner->mutable_centroids()->append(benches[xj], benches.dimension()); + } // end of for + } + + //! Update matrix score + static void UpdateMatrixScores(const OwnerType *owner, size_t first, + size_t last, float *out) { + const auto &matrix = owner->feature_matrix(); + const auto *bench = owner->centroids().data(); + + for (size_t i = first * BatchCount; i != last * BatchCount; + i += BatchCount) { + ContextType::template BatchDistance<1>(matrix[i], bench, + matrix.dimension(), &out[i]); + } + } + + //! Update cache score + static void UpdateCacheScores(const OwnerType *owner, float *out) { + const auto &cache = owner->feature_cache(); + const auto *bench = owner->centroids().data(); + + for (size_t i = 0, n = cache.count(); i != n; ++i) { + ContextType::Distance(bench, cache[i], cache.dimension(), &out[i]); + } + } + + //! Update bench score + static void UpdateBenchScores(const ContainerType *benches, + const StoreType *feat, float *out) { + float min_score = std::numeric_limits::max(); + + for (size_t i = 0, c = benches->count(); i != c; ++i) { + float new_score; + ContextType::Distance(benches->at(i), feat, benches->dimension(), + &new_score); + + if (new_score < min_score) { + min_score = new_score; + } + } + *out = min_score; + } + + //! Select k benches randomly + static void RandomSelectBenches(const ContainerType &cache, + const ContainerType &matrix, size_t k, + ContainerType *benches) { + ContainerType rows(cache.dimension()); + size_t m = matrix.count(); + size_t n = m + cache.count(); + std::mt19937 mt((std::random_device())()); + + rows.resize(BatchCount); + benches->reset(cache.dimension()); + benches->reserve(k); + + for (size_t i = 0; k > 0 && i < n; ++i) { + if (mt() % (n - i) >= k) { + continue; + } + // Selected a feature + if (i < m) { + ContextType::MatrixReverseTranspose(matrix[i / BatchCount * BatchCount], + matrix.dimension(), rows.data()); + benches->append(rows[i & (BatchCount - 1u)], matrix.dimension()); + } else { + benches->append(cache[i - m], cache.dimension()); + } + --k; + } // end of for + } + + //! Select k benches randomly + static void RandomSelectBenches(const ContainerType &cache, + const ContainerType &matrix, size_t k, + const std::vector &probs, + ContainerType *benches, + std::vector *bench_probs) { + std::mt19937 mt((std::random_device())()); + std::uniform_real_distribution dist(0.0, 1.0); + + // Sample features + KeyValueHeap> samples(k); + for (size_t i = 0; i < probs.size(); ++i) { + samples.emplace(i, std::pow(dist(mt), 1.0 / probs[i])); + } + + ContainerType rows(cache.dimension()); + size_t matrix_count = matrix.count(); + + rows.resize(BatchCount); + benches->reset(cache.dimension()); + benches->reserve(k); + bench_probs->clear(); + bench_probs->reserve(k); + + for (const auto &it : samples) { + // Selected a feature + if (it.first < matrix_count) { + ContextType::MatrixReverseTranspose( + matrix[it.first / BatchCount * BatchCount], matrix.dimension(), + rows.data()); + benches->append(rows[it.first & (BatchCount - 1u)], matrix.dimension()); + } else { + benches->append(cache[it.first - matrix_count], cache.dimension()); + } + bench_probs->push_back(probs[it.first]); + } + } + + private: + size_t chain_length_{32}; + bool assumption_free_{false}; +}; + +/*! Numerical K-Means Context + */ +template +class NumericalKmeansContext { + public: + //! constexpr variables + constexpr static size_t BatchCount = BATCH_COUNT; + + //! Type of values + using ValueType = typename std::remove_cv::type; + using StoreType = typename std::remove_cv::type; + + // Check supporting type + static_assert(IsSignedArithmetic::value, + "ValueType must be signed arithmetic"); + + /*! K-Means Context Cluster + */ + class Cluster { + public: + //! Constructor + Cluster(size_t dim) : accum_(dim, 0.0) {} + + //! Constructor + Cluster(const Cluster &rhs) + : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {} + + //! Constructor + Cluster(Cluster &&rhs) + : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {} + + //! Assignment + Cluster &operator=(const Cluster &rhs) { + cost_ = rhs.cost_; + count_ = rhs.count_; + accum_ = rhs.accum_; + return *this; + } + + //! Assignment + Cluster &operator=(Cluster &&rhs) { + cost_ = rhs.cost_; + count_ = rhs.count_; + accum_ = std::move(rhs.accum_); + return *this; + } + + //! Append a vector + void append(const ValueType *vec, size_t dim, float dist) { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + mutex_.lock(); + cost_ += dist; + count_ += 1; + + for (size_t i = 0; i != dim; ++i) { + accum_[i] += vec[i]; + } + mutex_.unlock(); + } + + //! Retrieve the centroid of vectors + void centroid(ValueType *out, size_t dim) const { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + for (size_t i = 0; i != dim; ++i) { + out[i] = count_ == 0 ? FloatCast(NAN) + : FloatCast(accum_[i] / count_); + } + } + + //! Retrieve squared error + double cost(void) const { + return cost_; + } + + //! Retrieve feature count + size_t count(void) const { + return count_; + } + + protected: + //! Convert float type to another type + template + static auto FloatCast(const double val) -> + typename std::enable_if::value, U>::type { + return static_cast(val); + } + + //! Convert float type to another type + template + static auto FloatCast(const double val) -> + typename std::enable_if::value, U>::type { + return static_cast(std::round(val)); + } + + private: + SpinMutex mutex_{}; + double cost_{0.0}; + size_t count_{0u}; + std::vector accum_{}; + }; + + //! operator [] + const Cluster &operator[](size_t i) const { + return clusters_[i]; + } + + //! operator [] + Cluster &operator[](size_t i) { + return clusters_[i]; + } + + //! Clear the context + void clear(void) { + clusters_.clear(); + } + + //! Reset the context + void reset(size_t k_value, size_t dim) { + clusters_.clear(); + clusters_.resize(k_value, dim); + } + + //! Retrieve context of clusters + const std::vector &clusters(void) const { + return clusters_; + } + + //! Compute the distance between matrix and query (batch) + template + static void BatchDistance(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + SquaredEuclideanDistanceMatrix::Compute(m, q, dim, + out); + } + + //! Compute the distance between matrix and query (single) + static void Distance(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out); + } + + //! Transpose a matrix + template + static auto MatrixTranspose(const U *src, size_t dim, T *dst) -> + typename std::enable_if= 2>::type { + MatrixHelper::Transpose(src, dim, dst); + } + + //! Transpose a matrix + template + static auto MatrixTranspose(const U *src, size_t dim, U *dst) -> + typename std::enable_if::type { + MatrixHelper::Transpose(src, dim >> 2, dst); + } + + //! Reverse transpose a matrix + template + static auto MatrixReverseTranspose(const U *src, size_t dim, U *dst) -> + typename std::enable_if= 2>::type { + MatrixHelper::ReverseTranspose(src, dim, dst); + } + + //! Reverse transpose a matrix + template + static auto MatrixReverseTranspose(const U *src, size_t dim, U *dst) -> + typename std::enable_if::type { + MatrixHelper::ReverseTranspose(src, dim >> 2, dst); + } + + //! Compute Norm2 + template ::value>::type> + static void Norm2(ValueType *data, size_t dim, float *norm) { + Normalizer::L2(data, dim, norm); + } + + //! Compute Norm2, for non-float do nothing + static void Norm2(ValueType * /*data*/, size_t /*dim*/, float *norm) { + *norm = 0.0f; + } + + private: + //! Members + std::vector clusters_{}; +}; + +/*! Nibble K-Means Context (INT4) + */ +template +class NibbleKmeansContext { + public: + //! constexpr variables + constexpr static size_t BatchCount = BATCH_COUNT; + + //! Type of values + using ValueType = typename std::remove_cv::type; + using StoreType = typename std::make_unsigned::type; + + // Check supporting type + static_assert(std::is_same::value || + std::is_same::value, + "ValueType must be int32_t or int64_t"); + + /*! K-Means Context Cluster + */ + class Cluster { + public: + //! Constructor + Cluster(size_t dim) : accum_(dim, 0.0) {} + + //! Constructor + Cluster(const Cluster &rhs) + : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {} + + //! Constructor + Cluster(Cluster &&rhs) + : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {} + + //! Assignment + Cluster &operator=(const Cluster &rhs) { + cost_ = rhs.cost_; + count_ = rhs.count_; + accum_ = rhs.accum_; + return *this; + } + + //! Assignment + Cluster &operator=(Cluster &&rhs) { + cost_ = rhs.cost_; + count_ = rhs.count_; + accum_ = std::move(rhs.accum_); + return *this; + } + + //! Append a vector + void append(const StoreType *vec, size_t dim, float dist) { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + mutex_.lock(); + cost_ += dist; + count_ += 1; + + const uint8_t *arr = reinterpret_cast(vec); + dim = (dim >> 1) << 1; + for (size_t i = 0; i != dim; i += 2) { + uint8_t val = arr[i >> 1]; + accum_[i] += ((int8_t)(val << 4) >> 4); + accum_[i + 1] += ((int8_t)(val) >> 4); + } + mutex_.unlock(); + } + + //! Retrieve the centroid of vectors + void centroid(StoreType *out, size_t dim) const { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + uint8_t *arr = reinterpret_cast(out); + dim = (dim >> 1) << 1; + for (size_t i = 0; i != dim; i += 2) { + int lo = + count_ == 0 ? 0 : static_cast(std::round(accum_[i] / count_)); + int hi = count_ == 0 + ? 0 + : static_cast(std::round(accum_[i + 1] / count_)); + arr[i >> 1] = (uint8_t)((hi << 4) & 0xf0) | (uint8_t)(lo & 0xf); + } + } + + //! Retrieve squared error + double cost(void) const { + return cost_; + } + + //! Retrieve feature count + size_t count(void) const { + return count_; + } + + private: + SpinMutex mutex_{}; + double cost_{0.0}; + size_t count_{0u}; + std::vector accum_{}; + }; + + //! operator [] + const Cluster &operator[](size_t i) const { + return clusters_[i]; + } + + //! operator [] + Cluster &operator[](size_t i) { + return clusters_[i]; + } + + //! Clear the context + void clear(void) { + clusters_.clear(); + } + + //! Reset the context + void reset(size_t k_value, size_t dim) { + clusters_.clear(); + clusters_.resize(k_value, dim); + } + + //! Retrieve context of clusters + const std::vector &clusters(void) const { + return clusters_; + } + + //! Compute the distance between matrix and query (batch) + template + static void BatchDistance(const StoreType *m, const StoreType *q, size_t dim, + float *out) { + SquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, out); + } + + //! Compute the distance between matrix and query (single) + static void Distance(const StoreType *m, const StoreType *q, size_t dim, + float *out) { + SquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, out); + } + + //! Transpose a matrix + static void MatrixTranspose(const StoreType *src, size_t dim, + StoreType *dst) { + MatrixHelper::Transpose(src, dim >> 3, dst); + } + + //! Reverse transpose a matrix + static void MatrixReverseTranspose(const StoreType *src, size_t dim, + StoreType *dst) { + MatrixHelper::ReverseTranspose(src, dim >> 3, dst); + } + + //! Compute and do norm2 + static void Norm2(StoreType * /*data*/, size_t /*dim*/, float *norm) { + *norm = 0; + } + + private: + //! Members + std::vector clusters_{}; +}; + +/*! Numerical InnerProduct K-Means Context + */ +template +class NumericalInnerProductKmeansContext { + public: + //! constexpr variables + constexpr static size_t BatchCount = BATCH_COUNT; + + //! Type of values + using ValueType = typename std::remove_cv::type; + using StoreType = typename std::remove_cv::type; + + // Check supporting type + static_assert(IsSignedArithmetic::value, + "ValueType must be signed arithmetic"); + + /*! K-Means Context Cluster + */ + class Cluster { + public: + //! Constructor + Cluster(size_t dim) : accum_(dim, 0.0) {} + + //! Constructor + Cluster(const Cluster &rhs) + : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {} + + //! Constructor + Cluster(Cluster &&rhs) + : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {} + + //! Assignment + Cluster &operator=(const Cluster &rhs) { + cost_ = rhs.cost_; + count_ = rhs.count_; + accum_ = rhs.accum_; + return *this; + } + + //! Assignment + Cluster &operator=(Cluster &&rhs) { + cost_ = rhs.cost_; + count_ = rhs.count_; + accum_ = std::move(rhs.accum_); + return *this; + } + + //! Append a vector + void append(const ValueType *vec, size_t dim, float dist) { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + mutex_.lock(); + cost_ += dist; + count_ += 1; + + for (size_t i = 0; i != dim; ++i) { + accum_[i] += vec[i]; + } + mutex_.unlock(); + } + + //! Retrieve the centroid of vectors + void centroid(ValueType *out, size_t dim) const { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + for (size_t i = 0; i != dim; ++i) { + out[i] = count_ == 0 ? FloatCast(NAN) + : FloatCast(accum_[i] / count_); + } + } + + //! Retrieve squared error + double cost(void) const { + return cost_; + } + + //! Retrieve feature count + size_t count(void) const { + return count_; + } + + protected: + //! Convert float type to another type + template + static auto FloatCast(const double val) -> + typename std::enable_if::value, U>::type { + return static_cast(val); + } + + //! Convert float type to another type + template + static auto FloatCast(const double val) -> + typename std::enable_if::value, U>::type { + return static_cast(std::round(val)); + } + + private: + SpinMutex mutex_{}; + double cost_{0.0}; + size_t count_{0u}; + std::vector accum_{}; + }; + + //! operator [] + const Cluster &operator[](size_t i) const { + return clusters_[i]; + } + + //! operator [] + Cluster &operator[](size_t i) { + return clusters_[i]; + } + + //! Clear the context + void clear(void) { + clusters_.clear(); + } + + //! Reset the context + void reset(size_t k_value, size_t dim) { + clusters_.clear(); + clusters_.resize(k_value, dim); + } + + //! Retrieve context of clusters + const std::vector &clusters(void) const { + return clusters_; + } + + //! Compute the distance between matrix and query (batch) + template + static void BatchDistance(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + MinusInnerProductMatrix::Compute(m, q, dim, out); + } + + //! Compute the distance between matrix and query (single) + static void Distance(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + MinusInnerProductMatrix::Compute(m, q, dim, out); + } + + //! Transpose a matrix + template + static auto MatrixTranspose(const U *src, size_t dim, T *dst) -> + typename std::enable_if= 2>::type { + MatrixHelper::Transpose(src, dim, dst); + } + + //! Transpose a matrix + template + static auto MatrixTranspose(const U *src, size_t dim, U *dst) -> + typename std::enable_if::type { + MatrixHelper::Transpose(src, dim >> 2, dst); + } + + //! Reverse transpose a matrix + template + static auto MatrixReverseTranspose(const U *src, size_t dim, U *dst) -> + typename std::enable_if= 2>::type { + MatrixHelper::ReverseTranspose(src, dim, dst); + } + + //! Reverse transpose a matrix + template + static auto MatrixReverseTranspose(const U *src, size_t dim, U *dst) -> + typename std::enable_if::type { + MatrixHelper::ReverseTranspose(src, dim >> 2, dst); + } + + //! Compute Norm2 + template ::value>::type> + static void Norm2(ValueType *data, size_t dim, float *norm) { + Normalizer::L2(data, dim, norm); + } + + //! Compute Norm2, for non-float do nothing + static void Norm2(ValueType * /*data*/, size_t /*dim*/, float *norm) { + *norm = 0.0f; + } + + private: + //! Members + std::vector clusters_{}; +}; + +/*! Nibble InnerProduct K-Means Context (INT4) + */ +template +class NibbleInnerProductKmeansContext { + public: + //! constexpr variables + constexpr static size_t BatchCount = BATCH_COUNT; + + //! Type of values + using ValueType = typename std::remove_cv::type; + using StoreType = typename std::make_unsigned::type; + + // Check supporting type + static_assert(std::is_same::value || + std::is_same::value, + "ValueType must be int32_t or int64_t"); + + /*! K-Means Context Cluster + */ + class Cluster { + public: + //! Constructor + Cluster(size_t dim) : accum_(dim, 0.0) {} + + //! Constructor + Cluster(const Cluster &rhs) + : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {} + + //! Constructor + Cluster(Cluster &&rhs) + : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {} + + //! Assignment + Cluster &operator=(const Cluster &rhs) { + cost_ = rhs.cost_; + count_ = rhs.count_; + accum_ = rhs.accum_; + return *this; + } + + //! Assignment + Cluster &operator=(Cluster &&rhs) { + cost_ = rhs.cost_; + count_ = rhs.count_; + accum_ = std::move(rhs.accum_); + return *this; + } + + //! Append a vector + void append(const StoreType *vec, size_t dim, float dist) { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + mutex_.lock(); + cost_ += dist; + count_ += 1; + + const uint8_t *arr = reinterpret_cast(vec); + dim = (dim >> 1) << 1; + for (size_t i = 0; i != dim; i += 2) { + uint8_t val = arr[i >> 1]; + accum_[i] += ((int8_t)(val << 4) >> 4); + accum_[i + 1] += ((int8_t)(val) >> 4); + } + mutex_.unlock(); + } + + //! Retrieve the centroid of vectors + void centroid(StoreType *out, size_t dim) const { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + uint8_t *arr = reinterpret_cast(out); + dim = (dim >> 1) << 1; + for (size_t i = 0; i != dim; i += 2) { + int lo = + count_ == 0 ? 0 : static_cast(std::round(accum_[i] / count_)); + int hi = count_ == 0 + ? 0 + : static_cast(std::round(accum_[i + 1] / count_)); + arr[i >> 1] = (uint8_t)((hi << 4) & 0xf0) | (uint8_t)(lo & 0xf); + } + } + + //! Retrieve squared error + double cost(void) const { + return cost_; + } + + //! Retrieve feature count + size_t count(void) const { + return count_; + } + + private: + SpinMutex mutex_{}; + double cost_{0.0}; + size_t count_{0u}; + std::vector accum_{}; + }; + + //! operator [] + const Cluster &operator[](size_t i) const { + return clusters_[i]; + } + + //! operator [] + Cluster &operator[](size_t i) { + return clusters_[i]; + } + + //! Clear the context + void clear(void) { + clusters_.clear(); + } + + //! Reset the context + void reset(size_t k_value, size_t dim) { + clusters_.clear(); + clusters_.resize(k_value, dim); + } + + //! Retrieve context of clusters + const std::vector &clusters(void) const { + return clusters_; + } + + //! Compute the distance between matrix and query (batch) + template + static void BatchDistance(const StoreType *m, const StoreType *q, size_t dim, + float *out) { + MinusInnerProductMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, out); + } + + //! Compute the distance between matrix and query (single) + static void Distance(const StoreType *m, const StoreType *q, size_t dim, + float *out) { + MinusInnerProductMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, out); + } + + //! Transpose a matrix + static void MatrixTranspose(const StoreType *src, size_t dim, + StoreType *dst) { + MatrixHelper::Transpose(src, dim >> 3, dst); + } + + //! Reverse transpose a matrix + static void MatrixReverseTranspose(const StoreType *src, size_t dim, + StoreType *dst) { + MatrixHelper::ReverseTranspose(src, dim >> 3, dst); + } + + //! Compute Norm2 + static void Norm2(StoreType * /*data*/, size_t /*dim*/, float *norm) { + *norm = 0; + } + + private: + //! Members + std::vector clusters_{}; +}; + +/*! Numerical K-Means cluster algorithm + */ +template > +using NumericalKmeans = + LloydCluster>; + +/*! Nibble K-Means cluster algorithm + */ +template > +using NibbleKmeans = LloydCluster>; + +/*! Numerical K-Means cluster algorithm + */ +template > +using NumericalInnerProductKmeans = + LloydCluster>; + +/*! Nibble K-Means cluster algorithm + */ +template > +using NibbleInnerProductKmeans = + LloydCluster>; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/algorithm/lloyd_cluster.h b/src/ailego/algorithm/lloyd_cluster.h new file mode 100644 index 0000000..b130b52 --- /dev/null +++ b/src/ailego/algorithm/lloyd_cluster.h @@ -0,0 +1,364 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! Random Centroids Generator + */ +template +struct RandomCentroidsGenerator { + //! Type of values + using OwnerType = typename std::decay::type; + using ContainerType = typename OwnerType::ContainerType; + using ContextType = typename OwnerType::ContextType; + using ThreadPoolType = TPool; + + //! constexpr variables + constexpr static size_t BatchCount = OwnerType::BatchCount; + + //! Generate centroids + void operator()(OwnerType *owner, ThreadPoolType &) const { + const auto &matrix = owner->feature_matrix(); + const auto &cache = owner->feature_cache(); + auto *centroids = owner->mutable_centroids(); + + ContainerType rows(cache.dimension()); + size_t m = matrix.count(); + size_t n = m + cache.count(); + size_t k = owner->k_value(); + std::mt19937 mt((std::random_device())()); + + rows.resize(BatchCount); + centroids->reset(cache.dimension()); + centroids->reserve(k); + + for (size_t i = 0; k > 0 && i < n; ++i) { + if (mt() % (n - i) >= k) { + continue; + } + // Selected a feature + if (i < m) { + ContextType::MatrixReverseTranspose(matrix[i / BatchCount * BatchCount], + matrix.dimension(), rows.data()); + centroids->append(rows[i & (BatchCount - 1u)], matrix.dimension()); + } else { + centroids->append(cache[i - m], cache.dimension()); + } + --k; + } + } +}; + +/*! Lloyd's algorithm cluster + */ +template +class LloydCluster { + public: + //! constexpr variables + constexpr static size_t BatchCount = TContext::BatchCount; + + //! Type of values + using ThreadPoolType = TPool; + using ContainerType = TContainer; + using ContextType = TContext; + using ValueType = typename TContext::ValueType; + using StoreType = typename TContext::StoreType; + + //! Constructor + LloydCluster(size_t k, size_t dim) + : k_value_(k), + feature_cache_(dim), + feature_matrix_(dim), + centroids_matrix_(dim), + centroids_(dim) {} + + //! Constructor + LloydCluster(size_t k, size_t dim, bool spherical) + : k_value_(k), + feature_cache_(dim), + feature_matrix_(dim), + centroids_matrix_(dim), + centroids_(dim), + spherical_{spherical} {} + + //! Constructor + LloydCluster(void) {} + + //! Destructor + ~LloydCluster(void) {} + + //! Append a feature + void append(const StoreType *arr, size_t dim) { + feature_cache_.append(arr, dim); + + if (feature_cache_.count() == BatchCount) { + size_t pos = feature_matrix_.count(); + feature_matrix_.resize(pos + BatchCount); + ContextType::MatrixTranspose(feature_cache_.data(), dim, + feature_matrix_[pos]); + feature_cache_.clear(); + } + } + + //! Reset cluster + void reset(size_t k, size_t dim) { + k_value_ = k; + feature_cache_.reset(dim); + feature_matrix_.reset(dim); + centroids_.reset(dim); + centroids_matrix_.reset(dim); + context_.clear(); + } + + //! Reset cluster + void reset(size_t k, size_t dim, bool spherical) { + k_value_ = k; + feature_cache_.reset(dim); + feature_matrix_.reset(dim); + centroids_.reset(dim); + centroids_matrix_.reset(dim); + context_.clear(); + spherical_ = spherical; + } + + //! Initialize centroids + template > + void init_centroids(ThreadPoolType &pool, const G &g = G()) { + g(this, pool); + } + + //! Cluster one time + template + bool cluster_once(ThreadPoolType &pool, double *cost) { + if (centroids_.empty()) { + RandomCentroidsGenerator g; + this->init_centroids(pool, g); + } + if (centroids_.count() != k_value_) { + return false; + } + context_.reset(centroids_.count(), centroids_.dimension()); + + size_t count = centroids_.count() / BatchCount * BatchCount; + centroids_matrix_.resize(count); + for (size_t i = 0; i != count; i += BatchCount) { + ContextType::MatrixTranspose(centroids_[i], centroids_.dimension(), + centroids_matrix_[i]); + } + size_t remain = static_cast(centroids_.count() - count); + if (remain > 0) { + centroids_matrix_.append(centroids_[count], centroids_.dimension(), + remain); + } + + // Using thread pool + auto group = pool.make_group(); + if (!feature_matrix_.empty()) { + size_t n = feature_matrix_.count() / BatchCount; + size_t c = std::max(n / pool.count() / 2u, 1u); + size_t m = n / c * c; + + for (size_t i = 0; i != m; i += c) { + group->submit(Closure::New(this, &LloydCluster::cluster_matrix_features, + i, i + c)); + } + for (size_t i = m; i != n; i += 1) { + group->submit(Closure::New(this, &LloydCluster::cluster_matrix_features, + i, i + 1)); + } + } + if (!feature_cache_.empty()) { + group->submit(Closure::New(this, &LloydCluster::cluster_cache_features)); + } + group->wait_finish(); + + *cost = 0.0; + for (size_t i = 0, n = centroids_.count(); i != n; ++i) { + const auto &item = context_[i]; + item.centroid(centroids_[i], centroids_.dimension()); + *cost += item.cost(); + } + + if (spherical_) { + for (size_t i = 0, n = centroids_.count(); i != n; ++i) { + float norm; + ContextType::Norm2(centroids_[i], centroids_.dimension(), &norm); + } + } + + return true; + } + + //! Retrieve the controids + ContainerType *mutable_centroids(void) { + return ¢roids_; + } + + //! Retrieve the controids + const ContainerType ¢roids(void) const { + return centroids_; + } + + //! Retrieve the K value + size_t k_value(void) const { + return k_value_; + } + + //! Retrieve context + const ContextType &context(void) const { + return context_; + } + + //! Retrieve the feature cache + const ContainerType &feature_cache(void) const { + return feature_cache_; + } + + //! Retrieve the feature matrix + const ContainerType &feature_matrix(void) const { + return feature_matrix_; + } + + //! Reserve the feature matrix + void feature_matrix_reserve(size_t count) { + feature_matrix_.reserve(count); + } + + protected: + //! Cluster the cache features + void cluster_cache_features(void) { + std::array scores; + + for (size_t i = 0, n = feature_cache_.count(); i != n; ++i) { + size_t count = centroids_matrix_.count() / BatchCount * BatchCount; + const StoreType *feature = feature_cache_[i]; + float nearest_score = std::numeric_limits::max(); + size_t nearest_index = 0; + + for (size_t j = 0; j != count; j += BatchCount) { + ContextType::template BatchDistance<1>(centroids_matrix_[j], feature, + centroids_matrix_.dimension(), + scores.data()); + + for (size_t k = 0; k < BatchCount; ++k) { + if (scores[k] < nearest_score) { + nearest_score = scores[k]; + nearest_index = j + k; + } + } + } // end of for + + for (size_t j = count, total = centroids_matrix_.count(); j != total; + ++j) { + ContextType::Distance(centroids_matrix_[j], feature, + centroids_matrix_.dimension(), scores.data()); + + if (scores[0] < nearest_score) { + nearest_score = scores[0]; + nearest_index = j; + } + } + context_[nearest_index].append(feature, feature_cache_.dimension(), + nearest_score); + } // end of for + } + + //! Cluster the matrix features + void cluster_matrix_features(size_t first, size_t last) { + std::array scores; + ContainerType rows(centroids_matrix_.dimension()); + + auto comp = [](float i, float j) { + if (std::isnan(i)) return false; + if (std::isnan(j)) return true; + + return i < j; + }; + + std::array nearest_scores; + std::array nearest_indexes; + + rows.resize(BatchCount); + for (size_t i = first * BatchCount; i != last * BatchCount; + i += BatchCount) { + size_t count = centroids_matrix_.count() / BatchCount * BatchCount; + const StoreType *block = feature_matrix_[i]; + + std::fill(nearest_indexes.data(), nearest_indexes.data() + BatchCount, 0); + std::fill(nearest_scores.data(), nearest_scores.data() + BatchCount, + std::numeric_limits::max()); + + for (size_t j = 0; j != count; j += BatchCount) { + ContextType::template BatchDistance( + centroids_matrix_[j], block, centroids_matrix_.dimension(), + scores.data()); + + for (size_t k = 0; k < BatchCount; ++k) { + const float *start = &scores[k * BatchCount]; + const float *result = + std::min_element(start, start + BatchCount, comp); + if (*result < nearest_scores[k]) { + nearest_scores[k] = *result; + nearest_indexes[k] = j + (result - start); + } + } + } // end of for + + for (size_t j = count, total = centroids_matrix_.count(); j != total; + ++j) { + ContextType::template BatchDistance<1>(block, centroids_matrix_[j], + centroids_matrix_.dimension(), + scores.data()); + + for (size_t k = 0; k < BatchCount; ++k) { + float score = scores[k]; + if (score < nearest_scores[k]) { + nearest_scores[k] = score; + nearest_indexes[k] = j; + } + } + } // end of for + + ContextType::MatrixReverseTranspose(block, feature_matrix_.dimension(), + rows.data()); + for (size_t k = 0; k < BatchCount; ++k) { + context_[nearest_indexes[k]].append( + rows[k], feature_matrix_.dimension(), nearest_scores[k]); + } + } // end of for + } + + private: + //! Members + size_t k_value_{0u}; + ContainerType feature_cache_{}; + ContainerType feature_matrix_{}; + ContainerType centroids_matrix_{}; + ContainerType centroids_{}; + ContextType context_{}; + bool spherical_{false}; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/buffer/block_eviction_queue.cc b/src/ailego/buffer/block_eviction_queue.cc new file mode 100644 index 0000000..eff9301 --- /dev/null +++ b/src/ailego/buffer/block_eviction_queue.cc @@ -0,0 +1,139 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +namespace zvec { +namespace ailego { + +int BlockEvictionQueue::init() { + evict_batch_size_ = 512; + for (size_t i = 0; i < CACHE_QUEUE_NUM; i++) { + evict_queues_.push_back(ConcurrentQueue(evict_batch_size_ * 200)); + } + return 0; +} + +bool BlockEvictionQueue::evict_single_block(BlockType &item) { + bool found = false; + for (size_t i = 0; i < CACHE_QUEUE_NUM; i++) { + found = evict_queues_[i].try_dequeue(item); + if (found) { + break; + } + } + return found; +} + +bool BlockEvictionQueue::is_valid_and_alive(const BlockType &item) { + std::shared_lock lock(valid_owners_mutex_); + if (item.owner == nullptr || + valid_owners_.find(item.owner) == valid_owners_.end()) { + return false; + } + return !item.owner->is_dead_block(item.owner_key, item.version); +} + +bool BlockEvictionQueue::evict_block(BlockType &item) { + bool ok = false; + do { + ok = evict_single_block(item); + if (!ok) { + return false; + } + } while (!is_valid_and_alive(item)); + return ok; +} + +void BlockEvictionQueue::recycle() { + BlockType item; + while (MemoryLimitPool::get_instance().is_full() && evict_block(item)) { + std::shared_lock lock(valid_owners_mutex_); + if (item.owner != nullptr && + valid_owners_.find(item.owner) != valid_owners_.end()) { + item.owner->evict_block(item.owner_key); + } + } +} + +bool BlockEvictionQueue::add_single_block(const BlockType &block, + int queue_index) { + bool ok = evict_queues_[queue_index].enqueue(block); + if (!ok) { + LOG_ERROR("enqueue failed."); + return false; + } + return true; +} + +int MemoryLimitPool::init(size_t pool_size) { + pool_size_ = 0; + BlockEvictionQueue::get_instance().recycle(); + pool_size_ = pool_size; + LOG_INFO("MemoryLimitPool initialized with pool size: %lu", pool_size_); + return 0; +} + +bool MemoryLimitPool::try_acquire_buffer(const size_t buffer_size, + char *&buffer) { + size_t expected, desired; + do { + expected = used_size_.load(); + if (expected >= pool_size_) { + return false; + } + desired = expected + buffer_size; + } while (!used_size_.compare_exchange_weak(expected, desired)); + buffer = (char *)ailego_aligned_malloc(buffer_size, 4096); + if (!buffer) { + used_size_.fetch_sub(buffer_size); + return false; + } + return true; +} + +void MemoryLimitPool::charge_external(const size_t buffer_size) { + size_t expected, desired; + do { + expected = used_size_.load(); + desired = expected + buffer_size; + } while (!used_size_.compare_exchange_weak(expected, desired)); +} + +void MemoryLimitPool::release_buffer(char *buffer, const size_t buffer_size) { + size_t expected, desired; + do { + expected = used_size_.load(); + desired = expected - buffer_size; + assert(expected >= buffer_size); + } while (!used_size_.compare_exchange_weak(expected, desired)); + ailego_free(buffer); +} + +void MemoryLimitPool::release_external(const size_t buffer_size) { + size_t expected, desired; + do { + expected = used_size_.load(); + desired = expected - buffer_size; + assert(expected >= buffer_size); + } while (!used_size_.compare_exchange_weak(expected, desired)); +} + +bool MemoryLimitPool::is_full() { + return used_size_.load() >= pool_size_; +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/buffer/vector_page_table.cc b/src/ailego/buffer/vector_page_table.cc new file mode 100644 index 0000000..3318db1 --- /dev/null +++ b/src/ailego/buffer/vector_page_table.cc @@ -0,0 +1,621 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include + +#if defined(_MSC_VER) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +static ssize_t zvec_pread(int fd, void *buf, size_t count, size_t offset) { + HANDLE handle = reinterpret_cast(_get_osfhandle(fd)); + if (handle == INVALID_HANDLE_VALUE) return -1; + OVERLAPPED ov = {}; + ov.Offset = static_cast(offset & 0xFFFFFFFF); + ov.OffsetHigh = static_cast(offset >> 32); + DWORD bytes_read = 0; + if (!ReadFile(handle, buf, static_cast(count), &bytes_read, &ov)) { + return -1; + } + return static_cast(bytes_read); +} +static ssize_t zvec_pwrite(int fd, const void *buf, size_t count, + size_t offset) { + HANDLE handle = reinterpret_cast(_get_osfhandle(fd)); + if (handle == INVALID_HANDLE_VALUE) return -1; + OVERLAPPED ov = {}; + ov.Offset = static_cast(offset & 0xFFFFFFFF); + ov.OffsetHigh = static_cast(offset >> 32); + DWORD bytes_written = 0; + if (!WriteFile(handle, buf, static_cast(count), &bytes_written, &ov)) { + return -1; + } + return static_cast(bytes_written); +} +#else +#include +static inline ssize_t zvec_pread(int fd, void *buf, size_t count, + size_t offset) { + return ::pread(fd, buf, count, static_cast(offset)); +} +static inline ssize_t zvec_pwrite(int fd, const void *buf, size_t count, + size_t offset) { + return ::pwrite(fd, buf, count, static_cast(offset)); +} +#endif + +namespace zvec { +namespace ailego { + +const size_t kVectorPageSize = MemoryHelper::PageSize(); + +bool VectorPageTable::init(size_t entry_num) { + size_t need_segments = (entry_num + kSegmentSize - 1) / kSegmentSize; + if (need_segments > kMaxSegments) { + LOG_ERROR( + "VectorPageTable::init: entry_num=%zu exceeds capacity " + "(kMaxEntries=%zu, need_segments=%zu, kMaxSegments=%zu); " + "refusing to init.", + entry_num, kMaxEntries, need_segments, kMaxSegments); + return false; + } + // Free old segments if any. init() is only called from VecBufferPool::init + // which is single-threaded with respect to other accesses, so a relaxed + // load of segment_count_ is sufficient here. + size_t old_count = segment_count_.load(std::memory_order_relaxed); + for (size_t i = 0; i < old_count; ++i) { + delete[] segments_[i]; + segments_[i] = nullptr; + } + for (size_t s = 0; s < need_segments; ++s) { + segments_[s] = new Entry[kSegmentSize]; + for (size_t i = 0; i < kSegmentSize; ++i) { + segments_[s][i].ref_count.store(std::numeric_limits::min()); + segments_[s][i].in_evict_queue.store(false); + segments_[s][i].is_dirty.store(false); + segments_[s][i].buffer = nullptr; + segments_[s][i].file_offset = 0; + } + } + // Publish new segments to readers. segment_count_ is published first + // (release) so that a reader that acquire-loads segment_count_ before + // entry_num_ also sees a consistent segment table; entry_num_ is the + // primary synchronization point used by callers via entry_num(). + segment_count_.store(need_segments, std::memory_order_release); + entry_num_.store(entry_num, std::memory_order_release); + return true; +} + +bool VectorPageTable::extend(size_t new_entry_num) { + // Relaxed read is fine: extend() is serialized by the caller (extend_file + // is invoked under the BufferStorage write latch). No other writer races + // with us on entry_num_ / segment_count_. + if (new_entry_num <= entry_num_.load(std::memory_order_relaxed)) { + return true; + } + size_t new_segment_count = (new_entry_num + kSegmentSize - 1) / kSegmentSize; + if (new_segment_count > kMaxSegments) { + LOG_ERROR( + "VectorPageTable::extend: new_entry_num=%zu exceeds capacity " + "(kMaxEntries=%zu, new_segment_count=%zu, kMaxSegments=%zu); " + "refusing to extend.", + new_entry_num, kMaxEntries, new_segment_count, kMaxSegments); + return false; + } + size_t old_count = segment_count_.load(std::memory_order_relaxed); + for (size_t s = old_count; s < new_segment_count; ++s) { + segments_[s] = new Entry[kSegmentSize]; + for (size_t i = 0; i < kSegmentSize; ++i) { + segments_[s][i].ref_count.store(std::numeric_limits::min()); + segments_[s][i].in_evict_queue.store(false); + segments_[s][i].is_dirty.store(false); + segments_[s][i].buffer = nullptr; + segments_[s][i].file_offset = 0; + } + } + // Publish in the same order as init(): segment_count_ first, entry_num_ + // last. Both are release-stores so that the prior segment allocation / + // Entry initialization is visible to any reader that acquire-loads either + // counter (typically via entry_num()). + segment_count_.store(new_segment_count, std::memory_order_release); + entry_num_.store(new_entry_num, std::memory_order_release); + return true; +} + +char *VectorPageTable::acquire_block(block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_relaxed)); + Entry &e = entry_at(block_id); + while (true) { + int current_count = e.ref_count.load(std::memory_order_acquire); + if (current_count < 0) { + return nullptr; + } + if (e.ref_count.compare_exchange_weak(current_count, current_count + 1, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + return e.buffer; + } + } +} + +void VectorPageTable::release_block(block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_relaxed)); + Entry &e = entry_at(block_id); + + if (e.ref_count.fetch_sub(1, std::memory_order_release) == 1) { + std::atomic_thread_fence(std::memory_order_acquire); + bool expected = false; + if (e.in_evict_queue.compare_exchange_strong(expected, true, + std::memory_order_acq_rel, + std::memory_order_relaxed)) { + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = 0; + BlockEvictionQueue::get_instance().add_single_block(block, 0); + } + } +} + +void VectorPageTable::evict_block(block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_relaxed)); + Entry &e = entry_at(block_id); + int expected = 0; + // Two-phase eviction to prevent data race on e.buffer with + // set_block_acquired. We first CAS to kEvicting (-1), which causes + // set_block_acquired to spin-wait; then do the actual work (flush, free, + // null buffer); finally store INT_MIN ("evicted") which unblocks + // set_block_acquired. + static constexpr int kEvicting = -1; + if (e.ref_count.compare_exchange_strong(expected, kEvicting)) { + char *buffer = e.buffer; + if (buffer && e.is_dirty.load(std::memory_order_relaxed) && + flush_callback_) { + flush_callback_(block_id, buffer, kVectorPageSize, e.file_offset); + e.is_dirty.store(false, std::memory_order_relaxed); + } + if (buffer) { + e.buffer = nullptr; + MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); + } + // Transition to fully-evicted state. Use release so that the + // set_block_acquired acquire-load sees e.buffer == nullptr. + e.ref_count.store(std::numeric_limits::min(), + std::memory_order_release); + } + e.in_evict_queue.store(false, std::memory_order_relaxed); +} + +char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer, + size_t file_offset) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + Entry &e = entry_at(block_id); + // Diagnostics for the kEvicting wait. The wait itself never gives up: + // the only thread that can transition kEvicting -> INT_MIN is the + // evict_block() owner, so abandoning the spin here would orphan the + // entry in kEvicting forever. Instead, we use bounded backoff and emit + // tiered logs so a stuck eviction is observable. + using clock = std::chrono::steady_clock; + const auto wait_start = clock::now(); + auto last_log = wait_start; + unsigned spin_count = 0; + bool warned = false; + while (true) { + int current_count = e.ref_count.load(std::memory_order_acquire); + if (current_count >= 0) { + if (e.ref_count.compare_exchange_weak(current_count, current_count + 1, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); + return e.buffer; + } + } else if (current_count == std::numeric_limits::min()) { + // Fully evicted — safe to claim this entry for our new buffer. + e.buffer = buffer; + e.file_offset = file_offset; + e.in_evict_queue.store(false, std::memory_order_relaxed); + e.is_dirty.store(false, std::memory_order_relaxed); + e.ref_count.store(1, std::memory_order_release); + return e.buffer; + } else { + // kEvicting (-1): eviction is in progress on this entry. + // Tiered backoff: hot spin first, then short sleep, then longer sleep. + ++spin_count; + if (spin_count < 64) { + // Pure busy wait for the common ~μs case. + } else if (spin_count < 1024) { + std::this_thread::yield(); + } else if (spin_count < 8192) { + std::this_thread::sleep_for(std::chrono::microseconds(100)); + } else { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + // Tiered diagnostics: warn once after 100ms, error every 1s after 1s. + const auto now = clock::now(); + const auto elapsed = now - wait_start; + if (!warned && elapsed >= std::chrono::milliseconds(100)) { + LOG_WARN( + "set_block_acquired: long kEvicting wait on block_id=%zu " + "(>=100ms); evict_block may be slow", + static_cast(block_id)); + warned = true; + } + if (elapsed >= std::chrono::seconds(1) && + (now - last_log) >= std::chrono::seconds(1)) { + const auto secs = + std::chrono::duration_cast(elapsed).count(); + LOG_ERROR( + "set_block_acquired: stuck in kEvicting on block_id=%zu for " + "%lld s; evict_block owner may be hung or starved", + static_cast(block_id), static_cast(secs)); + last_log = now; + } + } + } +} + +VecBufferPool::VecBufferPool(const std::string &filename, bool writable) { + file_name_ = filename; + writable_ = writable; +#if defined(_MSC_VER) + int flags = writable_ ? (O_RDWR | _O_BINARY) : (O_RDONLY | _O_BINARY); + fd_ = _open(filename.c_str(), flags, 0644); +#else + int flags = writable_ ? O_RDWR : O_RDONLY; + fd_ = ::open(filename.c_str(), flags, 0644); +#endif + if (fd_ < 0) { + throw std::runtime_error("Failed to open file: " + filename); + } +#if defined(_MSC_VER) + struct _stat64 st; + if (_fstat64(fd_, &st) < 0) { + _close(fd_); +#else + struct stat st; + if (fstat(fd_, &st) < 0) { + ::close(fd_); +#endif + throw std::runtime_error("Failed to stat file: " + filename); + } + file_size_ = st.st_size; +} + +int VecBufferPool::init() { + size_t block_num = (file_size_ + kVectorPageSize - 1) / kVectorPageSize; + if (!page_table_.init(block_num)) { + LOG_ERROR( + "VecBufferPool::init: page_table_ init failed for file[%s], " + "file_size=%zu, block_num=%zu (exceeds " + "VectorPageTable::kMaxEntries=%zu)", + file_name_.c_str(), file_size_, block_num, + VectorPageTable::kMaxEntries); + return -1; + } + block_mutexes_ = + std::make_unique(VecBufferPool::kMutexBucketCount); + LOG_DEBUG("entry num: %zu, file_size: %zu", page_table_.entry_num(), + file_size_); + + // In writable mode, inject a flush callback into the page table so that + // evict_block()/flush_block()/flush_all() can pwrite dirty blocks back to + // the backing file without needing to know about fd_ directly. + if (writable_) { + int fd = fd_; + const std::string &name = file_name_; + page_table_.set_flush_callback([fd, &name](block_id_t /*block_id*/, + char *buf, size_t sz, + size_t off) -> int { + ssize_t w = zvec_pwrite(fd, buf, sz, off); + if (w != static_cast(sz)) { + LOG_ERROR( + "Buffer pool flush failed: file[%s], offset[%zu], " + "expected[%zu], got[%zd]", + name.c_str(), off, sz, w); + return -1; + } + return 0; + }); + } + return 0; +} + +VecBufferPoolHandle VecBufferPool::get_handle() { + return VecBufferPoolHandle(*this); +} + +char *VecBufferPool::acquire_buffer(block_id_t page_id, int retry) { + assert(page_id < page_table_.entry_num()); + char *buffer = page_table_.acquire_block(page_id); + if (buffer) { + return buffer; + } + std::lock_guard lock( + block_mutexes_[page_id % VecBufferPool::kMutexBucketCount]); + buffer = page_table_.acquire_block(page_id); + if (buffer) { + return buffer; + } + { + bool found = MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer); + if (!found) { + for (int i = 0; i < retry; i++) { + BlockEvictionQueue::get_instance().recycle(); + found = MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer); + if (found) { + break; + } + } + } + if (!found) { + LOG_ERROR("Buffer pool failed to get free buffer: file[%s], page_id[%zu]", + file_name_.c_str(), page_id); + return nullptr; + } + } + + size_t page_offset = page_id * kVectorPageSize; + size_t expected_bytes = std::min(kVectorPageSize, file_size_ - page_offset); + if (expected_bytes < kVectorPageSize) { + std::memset(buffer + expected_bytes, 0, kVectorPageSize - expected_bytes); + } + ssize_t read_bytes = zvec_pread(fd_, buffer, expected_bytes, page_offset); + if (read_bytes != static_cast(expected_bytes)) { + LOG_ERROR( + "Buffer pool failed to read file at offset: file[%s], page_id[%zu], " + "offset[%zu], expected[%zu], got[%zd]", + file_name_.c_str(), page_id, page_offset, expected_bytes, read_bytes); + MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); + return nullptr; + } + return page_table_.set_block_acquired(page_id, buffer, page_offset); +} + +int VecBufferPool::get_meta(size_t offset, size_t length, char *buffer) { + ssize_t read_bytes = zvec_pread(fd_, buffer, length, offset); + if (read_bytes != static_cast(length)) { + LOG_ERROR( + "Buffer pool failed to read file at offset: file[%s], offset[%zu], " + "length[%zu]", + file_name_.c_str(), offset, length); + return -1; + } + return 0; +} + +int VecBufferPool::write_range(size_t file_offset, size_t length, + const char *src) { + if (!writable_) { + LOG_ERROR("write_range called on read-only pool: file[%s]", + file_name_.c_str()); + return -1; + } + if (length == 0) { + return 0; + } + size_t first_page = file_offset / kVectorPageSize; + size_t last_page = (file_offset + length - 1) / kVectorPageSize; + size_t remaining = length; + size_t src_cursor = 0; + for (size_t pg = first_page; pg <= last_page; ++pg) { + // Loading the page ensures we do not clobber unrelated bytes within the + // same page when the write is not page-aligned. acquire_buffer() pre-fills + // from the backing file (or zero-pads beyond EOF). + char *page = this->acquire_buffer(pg, 50); + if (!page) { + LOG_ERROR("write_range acquire failed: file[%s], page[%zu]", + file_name_.c_str(), pg); + return -1; + } + size_t page_start = pg * kVectorPageSize; + size_t intra_offset = (pg == first_page) ? (file_offset - page_start) : 0; + size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); + std::memcpy(page + intra_offset, src + src_cursor, chunk); + page_table_.mark_dirty(pg); + page_table_.release_block(pg); + src_cursor += chunk; + remaining -= chunk; + } + return 0; +} + +int VecBufferPool::write_meta(size_t offset, size_t length, + const char *buffer) { + if (!writable_) { + LOG_ERROR("write_meta called on read-only pool: file[%s]", + file_name_.c_str()); + return -1; + } + ssize_t w = zvec_pwrite(fd_, buffer, length, offset); + if (w != static_cast(length)) { + LOG_ERROR( + "Buffer pool failed to write meta: file[%s], offset[%zu], " + "length[%zu], got[%zd]", + file_name_.c_str(), offset, length, w); + return -1; + } + return 0; +} + +int VecBufferPool::flush_all() { + if (!writable_) { + return 0; + } + int rc = 0; + size_t total_dirty = 0; + size_t fail_count = 0; + for (size_t i = 0; i < page_table_.entry_num(); ++i) { + if (page_table_.is_block_dirty(i)) { + ++total_dirty; + int r = page_table_.flush_block(i); + if (r != 0) { + rc = r; + ++fail_count; + } + } + } + if (fail_count != 0) { + // Aggregated diagnostic so that callers (notably ~VecBufferPool, which + // discards the return value) cannot silently lose dirty pages: any + // unflushed page at this point means the on-disk image is now stale. + LOG_ERROR( + "VecBufferPool::flush_all: %zu/%zu dirty page(s) failed to flush, " + "file[%s] last_rc=%d -- on-disk data may be stale.", + fail_count, total_dirty, file_name_.c_str(), rc); + } + return rc; +} + +bool VecBufferPool::extend_file(size_t new_size) { + if (!writable_) { + LOG_ERROR("extend_file called on read-only pool: file[%s]", + file_name_.c_str()); + return false; + } + if (new_size <= file_size_) { + return true; + } + // Pre-validate against the page table's static capacity BEFORE mutating + // any on-disk state. Otherwise a successful ftruncate followed by a + // failed page_table_.extend() would leave the file size and the page + // table out of sync (file grew, but no Entry slots cover the new range). + size_t new_entry_num = (new_size + kVectorPageSize - 1) / kVectorPageSize; + if (new_entry_num > VectorPageTable::kMaxEntries) { + LOG_ERROR( + "extend_file: requested new_size=%zu would require %zu page entries, " + "exceeding VectorPageTable::kMaxEntries=%zu (file=%s).", + new_size, new_entry_num, VectorPageTable::kMaxEntries, + file_name_.c_str()); + return false; + } +#if defined(_MSC_VER) + if (_chsize_s(fd_, static_cast(new_size)) != 0) { + LOG_ERROR("extend_file _chsize_s failed: file[%s], new_size[%zu]", + file_name_.c_str(), new_size); + return false; + } +#else + if (::ftruncate(fd_, static_cast(new_size)) != 0) { + LOG_ERROR("extend_file ftruncate failed: file[%s], new_size[%zu]", + file_name_.c_str(), new_size); + return false; + } +#endif + file_size_ = new_size; + // Extend the page table to cover the new file range. Existing entries + // stay at their original addresses so concurrent readers are unaffected. + // Capacity has already been validated above, so this should never fail; + // a failure here would indicate a programming error and is logged. + if (new_entry_num > page_table_.entry_num()) { + if (!page_table_.extend(new_entry_num)) { + LOG_ERROR( + "extend_file: page_table_.extend(%zu) failed unexpectedly after " + "capacity pre-check (file=%s, new_size=%zu).", + new_entry_num, file_name_.c_str(), new_size); + return false; + } + } + return true; +} + +char *VecBufferPoolHandle::get_single_page(size_t file_offset, size_t len, + size_t &out_page_id) { + size_t first_page = file_offset / kVectorPageSize; + assert(len == 0 || (file_offset + len - 1) / kVectorPageSize == first_page); + out_page_id = first_page; + char *page = pool_.acquire_buffer(first_page, 50); + if (!page) { + LOG_ERROR( + "VecBufferPoolHandle::get_single_page: acquire_buffer failed, " + "file_offset=%zu, len=%zu, page=%zu, page_size=%zu", + file_offset, len, first_page, kVectorPageSize); + return nullptr; + } + return page + (file_offset - first_page * kVectorPageSize); +} + +bool VecBufferPoolHandle::read_range(size_t file_offset, size_t len, + char *out) { + if (len == 0) { + return true; + } + size_t first_page = file_offset / kVectorPageSize; + size_t last_page = (file_offset + len - 1) / kVectorPageSize; + size_t remaining = len; + size_t dst_cursor = 0; + for (size_t pg = first_page; pg <= last_page; ++pg) { + char *page = pool_.acquire_buffer(pg, 50); + if (!page) { + LOG_ERROR( + "VecBufferPoolHandle::read_range: acquire_buffer failed, " + "file_offset=%zu, len=%zu, page=%zu, first_page=%zu, last_page=%zu, " + "page_size=%zu", + file_offset, len, pg, first_page, last_page, kVectorPageSize); + return false; + } + size_t page_start = pg * kVectorPageSize; + size_t intra_offset = (pg == first_page) ? (file_offset - page_start) : 0; + size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); + std::memcpy(out + dst_cursor, page + intra_offset, chunk); + pool_.page_table_.release_block(pg); + dst_cursor += chunk; + remaining -= chunk; + } + return true; +} + +int VecBufferPoolHandle::get_meta(size_t offset, size_t length, char *buffer) { + return pool_.get_meta(offset, length, buffer); +} + +int VecBufferPoolHandle::write_range(size_t file_offset, size_t len, + const char *src) { + return pool_.write_range(file_offset, len, src); +} + +int VecBufferPoolHandle::write_meta(size_t offset, size_t length, + const char *buffer) { + return pool_.write_meta(offset, length, buffer); +} + +int VecBufferPoolHandle::flush_all() { + return pool_.flush_all(); +} + +bool VecBufferPoolHandle::writable() const { + return pool_.writable(); +} + +void VecBufferPoolHandle::release_one(block_id_t block_id) { + pool_.page_table_.release_block(block_id); +} + +void VecBufferPoolHandle::acquire_one(block_id_t block_id) { + // The caller must guarantee the block is already loaded before calling + // acquire_one(). The return value of acquire_block() is intentionally + // ignored here, as a null return would indicate a contract violation. + pool_.page_table_.acquire_block(block_id); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/container/bitmap.cc b/src/ailego/container/bitmap.cc new file mode 100644 index 0000000..652b1f8 --- /dev/null +++ b/src/ailego/container/bitmap.cc @@ -0,0 +1,443 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "bitmap.h" + +namespace zvec { +namespace ailego { + +size_t Bitset::BitwiseAndCardinality(const Bitset &lhs, const Bitset &rhs) { + return BitsetHelper::BitwiseAndCardinality( + lhs.array_.data(), rhs.array_.data(), + std::min(lhs.array_.size(), rhs.array_.size())); +} + +size_t Bitset::BitwiseAndnotCardinality(const Bitset &lhs, const Bitset &rhs) { + size_t lsize = lhs.array_.size(); + size_t rsize = rhs.array_.size(); + + if (lsize > rsize) { + return ( + BitsetHelper::BitwiseAndnotCardinality(lhs.array_.data(), + rhs.array_.data(), rsize) + + BitsetHelper::Cardinality(lhs.array_.data() + rsize, lsize - rsize)); + } + return BitsetHelper::BitwiseAndnotCardinality(lhs.array_.data(), + rhs.array_.data(), lsize); +} + +size_t Bitset::BitwiseXorCardinality(const Bitset &lhs, const Bitset &rhs) { + size_t lsize = lhs.array_.size(); + size_t rsize = rhs.array_.size(); + + if (lsize < rsize) { + return ( + BitsetHelper::BitwiseXorCardinality(lhs.array_.data(), + rhs.array_.data(), lsize) + + BitsetHelper::Cardinality(rhs.array_.data() + lsize, rsize - lsize)); + } else if (lsize > rsize) { + return ( + BitsetHelper::BitwiseXorCardinality(lhs.array_.data(), + rhs.array_.data(), rsize) + + BitsetHelper::Cardinality(lhs.array_.data() + rsize, lsize - rsize)); + } + return BitsetHelper::BitwiseXorCardinality(lhs.array_.data(), + rhs.array_.data(), lsize); +} + +size_t Bitset::BitwiseOrCardinality(const Bitset &lhs, const Bitset &rhs) { + size_t lsize = lhs.array_.size(); + size_t rsize = rhs.array_.size(); + + if (lsize < rsize) { + return ( + BitsetHelper::BitwiseOrCardinality(lhs.array_.data(), rhs.array_.data(), + lsize) + + BitsetHelper::Cardinality(rhs.array_.data() + lsize, rsize - lsize)); + } else if (lsize > rsize) { + return ( + BitsetHelper::BitwiseOrCardinality(lhs.array_.data(), rhs.array_.data(), + rsize) + + BitsetHelper::Cardinality(lhs.array_.data() + rsize, lsize - rsize)); + } + return BitsetHelper::BitwiseOrCardinality(lhs.array_.data(), + rhs.array_.data(), lsize); +} + +void Bitmap::clear(void) { + for (std::vector::iterator iter = array_.begin(); + iter != array_.end(); ++iter) { + delete (*iter); + } + array_.clear(); +} + +void Bitmap::copy(const Bitmap &rhs) { + this->clear(); + + for (std::vector::const_iterator iter = rhs.array_.begin(); + iter != rhs.array_.end(); ++iter) { + Bucket *bucket = nullptr; + if (*iter) { + bucket = new Bucket(*(*iter)); + } + array_.push_back(bucket); + } +} + +void Bitmap::shrink_to_fit(void) { + size_t shrink_count = 0; + std::vector::reverse_iterator iter; + + for (iter = array_.rbegin(); iter != array_.rend(); ++iter) { + if (*iter) { + if (!(*iter)->test_none()) { + break; + } + delete (*iter); + *iter = nullptr; + } + ++shrink_count; + } + for (; iter != array_.rend(); ++iter) { + if ((*iter) && (*iter)->test_none()) { + delete (*iter); + *iter = nullptr; + } + } + if (shrink_count != 0) { + array_.resize(array_.size() - shrink_count); + } +} + +bool Bitmap::test(size_t num) const { + // High 16 bits + size_t offset = num >> 16; + + if (offset < array_.size()) { + const Bucket *bucket = array_[offset]; + if (bucket) { + // Low 16 bits + return bucket->test(static_cast(num)); + } + } + return false; +} + +void Bitmap::set(size_t num) { + // High 16 bits + size_t offset = num >> 16; + if (offset >= array_.size()) { + array_.resize(offset + 1, nullptr); + } + + Bucket *&bucket = array_[offset]; + if (!bucket) { + bucket = new Bucket; + } + // Low 16 bits + bucket->set(static_cast(num)); +} + +void Bitmap::reset(size_t num) { + // High 16 bits + size_t offset = num >> 16; + if (offset >= array_.size()) { + array_.resize(offset + 1, nullptr); + } + + if (offset < array_.size()) { + Bucket *bucket = array_[offset]; + if (bucket) { + // Low 16 bits + bucket->reset(static_cast(num)); + } + } +} + +void Bitmap::flip(size_t num) { + // High 16 bits + uint16_t offset = (uint16_t)(num >> 16); + if (offset >= array_.size()) { + array_.resize(offset + 1, nullptr); + } + + Bucket *&bucket = array_[offset]; + if (!bucket) { + bucket = new Bucket; + } + // Low 16 bits + bucket->flip(static_cast(num)); +} + +void Bitmap::bitwise_and(const Bitmap &rhs) { + size_t overlap = std::min(array_.size(), rhs.array_.size()); + + for (size_t i = 0; i < overlap; ++i) { + Bucket *&dst = array_[i]; + + if (dst) { + const Bucket *src = rhs.array_[i]; + if (src) { + dst->bitwise_and(*src); + } else { + delete dst; + dst = nullptr; + } + } + } + for (size_t i = overlap; i < array_.size(); ++i) { + Bucket *&dst = array_[i]; + delete dst; + dst = nullptr; + } +} + +void Bitmap::bitwise_andnot(const Bitmap &rhs) { + size_t overlap = std::min(array_.size(), rhs.array_.size()); + + for (size_t i = 0; i < overlap; ++i) { + Bucket *&dst = array_[i]; + + if (dst) { + const Bucket *src = rhs.array_[i]; + if (src) { + dst->bitwise_andnot(*src); + } + } + } +} + +void Bitmap::bitwise_or(const Bitmap &rhs) { + size_t overlap = std::min(array_.size(), rhs.array_.size()); + + for (size_t i = 0; i < overlap; ++i) { + const Bucket *src = rhs.array_[i]; + + if (src) { + Bucket *&dst = array_[i]; + + if (dst) { + dst->bitwise_or(*src); + } else { + dst = new Bucket(*src); + } + } + } + for (size_t i = overlap; i < rhs.array_.size(); ++i) { + const Bucket *src = rhs.array_[i]; + Bucket *bucket = nullptr; + + if (src) { + bucket = new Bucket(*src); + } + array_.push_back(bucket); + } +} + +void Bitmap::bitwise_xor(const Bitmap &rhs) { + size_t overlap = std::min(array_.size(), rhs.array_.size()); + + for (size_t i = 0; i < overlap; ++i) { + const Bucket *src = rhs.array_[i]; + + if (src) { + Bucket *&dst = array_[i]; + + if (dst) { + dst->bitwise_xor(*src); + } else { + dst = new Bucket(*src); + } + } + } + for (size_t i = overlap; i < rhs.array_.size(); ++i) { + const Bucket *src = rhs.array_[i]; + Bucket *bucket = nullptr; + + if (src) { + bucket = new Bucket(*src); + } + array_.push_back(bucket); + } +} + +void Bitmap::bitwise_not(void) { + for (std::vector::iterator iter = array_.begin(); + iter != array_.end(); ++iter) { + Bucket *&bucket = *iter; + if (!bucket) { + bucket = new Bucket; + } + bucket->bitwise_not(); + } +} + +bool Bitmap::test_all(void) const { + if (array_.empty()) { + return false; + } + for (std::vector::const_iterator iter = array_.begin(); + iter != array_.end(); ++iter) { + if (!(*iter) || !(*iter)->test_all()) { + return false; + } + } + return true; +} + +bool Bitmap::test_any(void) const { + for (std::vector::const_iterator iter = array_.begin(); + iter != array_.end(); ++iter) { + if (*iter && (*iter)->test_any()) { + return true; + } + } + return false; +} + +bool Bitmap::test_none(void) const { + for (std::vector::const_iterator iter = array_.begin(); + iter != array_.end(); ++iter) { + if (*iter && !(*iter)->test_none()) { + return false; + } + } + return true; +} + +size_t Bitmap::cardinality(void) const { + size_t result = 0; + for (std::vector::const_iterator iter = array_.begin(); + iter != array_.end(); ++iter) { + if (*iter) { + result += (*iter)->cardinality(); + } + } + return result; +} + +void Bitmap::extract(size_t base, std::vector *out) const { + for (std::vector::const_iterator iter = array_.begin(); + iter != array_.end(); ++iter) { + if (*iter) { + (*iter)->extract(base, out); + } + base += Bucket::MAX_SIZE; + } +} + +size_t Bitmap::BitwiseAndCardinality(const Bitmap &lhs, const Bitmap &rhs) { + size_t overlap = std::min(lhs.array_.size(), rhs.array_.size()); + size_t dist = 0; + + for (size_t i = 0; i < overlap; ++i) { + const Bucket *l = lhs.array_[i]; + const Bucket *r = rhs.array_[i]; + + if (l && r) { + dist += Bucket::BitwiseAndCardinality(*l, *r); + } + } + return dist; +} + +size_t Bitmap::BitwiseAndnotCardinality(const Bitmap &lhs, const Bitmap &rhs) { + size_t overlap = std::min(lhs.array_.size(), rhs.array_.size()); + size_t dist = 0; + + for (size_t i = 0; i < overlap; ++i) { + const Bucket *l = lhs.array_[i]; + if (l) { + const Bucket *r = rhs.array_[i]; + if (r) { + dist += Bucket::BitwiseAndnotCardinality(*l, *r); + } else { + dist += l->cardinality(); + } + } + } + for (size_t i = overlap; i < lhs.array_.size(); ++i) { + const Bucket *l = lhs.array_[i]; + if (l) { + dist += l->cardinality(); + } + } + return dist; +} + +size_t Bitmap::BitwiseXorCardinality(const Bitmap &lhs, const Bitmap &rhs) { + size_t overlap = std::min(lhs.array_.size(), rhs.array_.size()); + size_t dist = 0; + + for (size_t i = 0; i < overlap; ++i) { + const Bucket *l = lhs.array_[i]; + const Bucket *r = rhs.array_[i]; + + if (l && r) { + dist += Bucket::BitwiseXorCardinality(*l, *r); + } else if (l) { + dist += l->cardinality(); + } else if (r) { + dist += r->cardinality(); + } + } + for (size_t i = overlap; i < lhs.array_.size(); ++i) { + const Bucket *l = lhs.array_[i]; + if (l) { + dist += l->cardinality(); + } + } + for (size_t i = overlap; i < rhs.array_.size(); ++i) { + const Bucket *r = rhs.array_[i]; + if (r) { + dist += r->cardinality(); + } + } + return dist; +} + +size_t Bitmap::BitwiseOrCardinality(const Bitmap &lhs, const Bitmap &rhs) { + size_t overlap = std::min(lhs.array_.size(), rhs.array_.size()); + size_t dist = 0; + + for (size_t i = 0; i < overlap; ++i) { + const Bucket *l = lhs.array_[i]; + const Bucket *r = rhs.array_[i]; + + if (l && r) { + dist += Bucket::BitwiseOrCardinality(*l, *r); + } else if (l) { + dist += l->cardinality(); + } else if (r) { + dist += r->cardinality(); + } + } + for (size_t i = overlap; i < lhs.array_.size(); ++i) { + const Bucket *l = lhs.array_[i]; + if (l) { + dist += l->cardinality(); + } + } + for (size_t i = overlap; i < rhs.array_.size(); ++i) { + const Bucket *r = rhs.array_[i]; + if (r) { + dist += r->cardinality(); + } + } + return dist; +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/container/bitmap.h b/src/ailego/container/bitmap.h new file mode 100644 index 0000000..72d5b50 --- /dev/null +++ b/src/ailego/container/bitmap.h @@ -0,0 +1,504 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! Fixed Bitset + */ +template ::type> +class FixedBitset { + public: + enum { MAX_SIZE = N }; + + //! Constructor + FixedBitset(void) { + memset(array_, 0, sizeof(array_)); + } + + //! Constructor + FixedBitset(const FixedBitset &rhs) { + memcpy(array_, rhs.array_, sizeof(array_)); + } + + //! Destructor + ~FixedBitset(void) {} + + //! Assignment + FixedBitset &operator=(const FixedBitset &rhs) { + memcpy(array_, rhs.array_, sizeof(array_)); + return *this; + } + + //! Retrieve data pointer + uint32_t *data(void) { + return reinterpret_cast(array_); + } + + //! Retrieve data pointer + const uint32_t *data(void) const { + return reinterpret_cast(array_); + } + + //! Retrieve count of bits in set + constexpr size_t size(void) const { + return MAX_SIZE; + } + + // ！Clear the bitset + void clear(void) { + memset(array_, 0, sizeof(array_)); + } + + //! Test a bit in bitset + bool test(size_t num) const { + ailego_assert_with(N > num, "overflow argument"); + return ((array_[num >> 5] & (1u << (num & 0x1f))) != 0); + } + + //! Set a bit in bitset + void set(size_t num) { + ailego_assert_with(N > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] |= mask; + } + + //! Clear a bit in bitset + void reset(size_t num) { + ailego_assert_with(N > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] &= ~mask; + } + + //! Toggle a bit in bitset + void flip(size_t num) { + ailego_assert_with(N > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] ^= mask; + } + + //! Perform binary AND + void bitwise_and(const FixedBitset &rhs) { + BitsetHelper::BitwiseAnd(array_, rhs.array_, ((N + 0x1f) >> 5)); + } + + //! Perform binary AND NOT + void bitwise_andnot(const FixedBitset &rhs) { + BitsetHelper::BitwiseAndnot(array_, rhs.array_, ((N + 0x1f) >> 5)); + } + + //! Perform binary OR + void bitwise_or(const FixedBitset &rhs) { + BitsetHelper::BitwiseOr(array_, rhs.array_, ((N + 0x1f) >> 5)); + } + + //! Perform binary XOR + void bitwise_xor(const FixedBitset &rhs) { + BitsetHelper::BitwiseXor(array_, rhs.array_, ((N + 0x1f) >> 5)); + } + + //! Perform binary NOT + void bitwise_not(void) { + BitsetHelper::BitwiseNot(array_, ((N + 0x1f) >> 5)); + } + + //! Check if all bits are set to true + bool test_all(void) const { + return BitsetHelper::TestAll(array_, ((N + 0x1f) >> 5)); + } + + //! Check if any bits are set to true + bool test_any(void) const { + return BitsetHelper::TestAny(array_, ((N + 0x1f) >> 5)); + } + + //! Check if none of the bits are set to true + bool test_none(void) const { + return BitsetHelper::TestNone(array_, ((N + 0x1f) >> 5)); + } + + //! Compute the cardinality of a bitset + size_t cardinality(void) const { + return BitsetHelper::Cardinality(array_, ((N + 0x1f) >> 5)); + } + + //! Extract the bitset to an array + void extract(size_t base, std::vector *out) const { + const uint32_t *iter = array_; + const uint32_t *last = array_ + ((N + 0x1f) >> 5); + + for (; iter != last; ++iter) { + uint32_t w = *iter; + + while (w != 0) { + uint32_t c = ailego_ctz32(w); + w &= ~(1u << c); + out->push_back(base + c); + } + base += 32u; + } + } + + //! Extract the bitset to an array + void extract(std::vector *out) const { + this->extract(0, out); + } + + //! Compute the AND cardinality between two bitsets + static size_t BitwiseAndCardinality(const FixedBitset &lhs, + const FixedBitset &rhs) { + return BitsetHelper::BitwiseAndCardinality(lhs.array_, rhs.array_, + ((N + 0x1f) >> 5)); + } + + //! Compute the ANDNOT cardinality between two bitsets + static size_t BitwiseAndnotCardinality(const FixedBitset &lhs, + const FixedBitset &rhs) { + return BitsetHelper::BitwiseAndnotCardinality(lhs.array_, rhs.array_, + ((N + 0x1f) >> 5)); + } + + //! Compute the XOR cardinality between two bitsets + static size_t BitwiseXorCardinality(const FixedBitset &lhs, + const FixedBitset &rhs) { + return BitsetHelper::BitwiseXorCardinality(lhs.array_, rhs.array_, + ((N + 0x1f) >> 5)); + } + + //! Compute the OR cardinality between two bitsets + static size_t BitwiseOrCardinality(const FixedBitset &lhs, + const FixedBitset &rhs) { + return BitsetHelper::BitwiseOrCardinality(lhs.array_, rhs.array_, + ((N + 0x1f) >> 5)); + } + + //! Convert a array pointer to bitset pointer + static FixedBitset *Cast(uint32_t *arr) { + return reinterpret_cast *>(arr); + } + + //! Convert a array pointer to bitset pointer + static const FixedBitset *Cast(const uint32_t *arr) { + return reinterpret_cast *>(arr); + } + + //! Convert a array pointer to bitset pointer + static FixedBitset *Cast(uint64_t *arr) { + return reinterpret_cast *>(arr); + } + + //! Convert a array pointer to bitset pointer + static const FixedBitset *Cast(const uint64_t *arr) { + return reinterpret_cast *>(arr); + } + + private: + uint32_t array_[(N + 0x1f) >> 5]; +}; + +/*! Fixed Bitset (Special) + */ +template <> +class FixedBitset<0> { + public: + enum { MAX_SIZE = 0 }; + + //! Retrieve max size of bitset + constexpr size_t size(void) const { + return MAX_SIZE; + } +}; + +/*! Bitset + */ +class Bitset { + public: + //! Constructor + Bitset(void) : array_() {} + + //! Constructor + Bitset(size_t bits) : array_((bits + 0x1f) >> 5) {} + + //! Constructor + Bitset(const Bitset &rhs) : array_(rhs.array_) {} + + //! Constructor + Bitset(Bitset &&rhs) : array_(std::move(rhs.array_)) {} + + //! Destructor + ~Bitset(void) {} + + //! Assignment + Bitset &operator=(const Bitset &rhs) { + array_ = rhs.array_; + return *this; + } + + //! Assignment + Bitset &operator=(Bitset &&rhs) { + array_ = std::move(rhs.array_); + return *this; + } + + //! Retrieve data pointer + uint32_t *data(void) { + return array_.data(); + } + + //! Retrieve data pointer + const uint32_t *data(void) const { + return array_.data(); + } + + //! Retrieve count of bits in set + size_t size(void) const { + return (array_.size() << 5); + } + + //! Resize the bitset + void resize(size_t bits) { + array_.resize((bits + 0x1f) >> 5); + } + + // ！Clear the bitset + void clear(void) { + array_.clear(); + } + + //! Test a bit in bitset + bool test(size_t num) const { + ailego_assert_with(this->size() > num, "overflow argument"); + return ((array_[num >> 5] & (1u << (num & 0x1f))) != 0); + } + + //! Set a bit in bitset + void set(size_t num) { + ailego_assert_with(this->size() > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] |= mask; + } + + //! Clear a bit in bitset + void reset(size_t num) { + ailego_assert_with(this->size() > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] &= ~mask; + } + + //! Toggle a bit in bitset + void flip(size_t num) { + ailego_assert_with(this->size() > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] ^= mask; + } + + //! Perform binary AND + void bitwise_and(const Bitset &rhs) { + BitsetHelper::BitwiseAnd(array_.data(), rhs.array_.data(), + std::min(array_.size(), rhs.array_.size())); + } + + //! Perform binary AND NOT + void bitwise_andnot(const Bitset &rhs) { + BitsetHelper::BitwiseAndnot(array_.data(), rhs.array_.data(), + std::min(array_.size(), rhs.array_.size())); + } + + //! Perform binary OR + void bitwise_or(const Bitset &rhs) { + BitsetHelper::BitwiseOr(array_.data(), rhs.array_.data(), + std::min(array_.size(), rhs.array_.size())); + } + + //! Perform binary XOR + void bitwise_xor(const Bitset &rhs) { + BitsetHelper::BitwiseXor(array_.data(), rhs.array_.data(), + std::min(array_.size(), rhs.array_.size())); + } + + //! Perform binary NOT + void bitwise_not(void) { + BitsetHelper::BitwiseNot(array_.data(), array_.size()); + } + + //! Check if all bits are set to true + bool test_all(void) const { + return BitsetHelper::TestAll(array_.data(), array_.size()); + } + + //! Check if any bits are set to true + bool test_any(void) const { + return BitsetHelper::TestAny(array_.data(), array_.size()); + } + + //! Check if none of the bits are set to true + bool test_none(void) const { + return BitsetHelper::TestNone(array_.data(), array_.size()); + } + + //! Compute the cardinality of a bitset + size_t cardinality(void) const { + return BitsetHelper::Cardinality(array_.data(), array_.size()); + } + + //! Extract the bitset to an array + void extract(size_t base, std::vector *out) const { + const uint32_t *iter = array_.data(); + const uint32_t *last = array_.data() + array_.size(); + + for (; iter != last; ++iter) { + uint32_t w = *iter; + + while (w != 0) { + uint32_t c = ailego_ctz32(w); + w &= ~(1u << c); + out->push_back(base + c); + } + base += 32u; + } + } + + //! Extract the bitset to an array + void extract(std::vector *out) const { + this->extract(0, out); + } + + //! Compute the AND cardinality between two bitsets + static size_t BitwiseAndCardinality(const Bitset &lhs, const Bitset &rhs); + + //! Compute the ANDNOT cardinality between two bitsets + static size_t BitwiseAndnotCardinality(const Bitset &lhs, const Bitset &rhs); + + //! Compute the XOR cardinality between two bitsets + static size_t BitwiseXorCardinality(const Bitset &lhs, const Bitset &rhs); + + //! Compute the OR cardinality between two bitsets + static size_t BitwiseOrCardinality(const Bitset &lhs, const Bitset &rhs); + + private: + std::vector array_; +}; + +/*! Bitmap + */ +class Bitmap { + public: + typedef FixedBitset<65536u> Bucket; + + //! Constructor + Bitmap(void) : array_() {} + + //! Constructor + Bitmap(const Bitmap &rhs) { + this->copy(rhs); + } + + //! Destructor + ~Bitmap(void) { + this->clear(); + } + + //! Assignment + Bitmap &operator=(const Bitmap &rhs) { + this->copy(rhs); + return *this; + } + + //! Retrieve bucket size of bitmap + size_t bucket_size(void) const { + return array_.size(); + } + + // ！Clear the bitmap + void clear(void); + + //! Remove the none buckets + void shrink_to_fit(void); + + //! Test a bit in bitmap + bool test(size_t num) const; + + //! Set a bit in bitmap + void set(size_t num); + + //! Reset a bit in bitmap + void reset(size_t num); + + //! Toggle a bit in bitmap + void flip(size_t num); + + //! Perform binary AND + void bitwise_and(const Bitmap &rhs); + + //! Perform binary AND NOT + void bitwise_andnot(const Bitmap &rhs); + + //! Perform binary OR + void bitwise_or(const Bitmap &rhs); + + //! Perform binary XOR + void bitwise_xor(const Bitmap &rhs); + + //! Perform binary NOT (It will expand the whole map) + void bitwise_not(void); + + //! Check if all bits are set to true + bool test_all(void) const; + + //! Check if any bits are set to true + bool test_any(void) const; + + //! Check if none of the bits are set to true + bool test_none(void) const; + + //! Compute the cardinality of a bitmap + size_t cardinality(void) const; + + //! Extract the bitmap to an array + void extract(size_t base, std::vector *out) const; + + //! Extract the bitmap to an array + void extract(std::vector *out) const { + this->extract(0, out); + } + + //! Compute the AND cardinality between two bitmaps + static size_t BitwiseAndCardinality(const Bitmap &lhs, const Bitmap &rhs); + + //! Compute the ANDNOT cardinality between two bitmaps + static size_t BitwiseAndnotCardinality(const Bitmap &lhs, const Bitmap &rhs); + + //! Compute the XOR cardinality between two bitmaps + static size_t BitwiseXorCardinality(const Bitmap &lhs, const Bitmap &rhs); + + //! Compute the OR cardinality between two bitmaps + static size_t BitwiseOrCardinality(const Bitmap &lhs, const Bitmap &rhs); + + protected: + //! Copy the content from another bitmap + void copy(const Bitmap &rhs); + + private: + std::vector array_; +}; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/container/bloom_filter.h b/src/ailego/container/bloom_filter.h new file mode 100644 index 0000000..9a8cacd --- /dev/null +++ b/src/ailego/container/bloom_filter.h @@ -0,0 +1,279 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! Bloom Filter Calculator + */ +struct BloomFilterCalculator { + /** + * \brief Calculate probability of false positives + * \param n Number of items in the filter + * \param m Number of bits in the filter + * \param k Number of hash functions + * \return Probability of false positives + */ + static double Probability(size_t n, size_t m, size_t k) { + return std::pow(1.0 - std::exp(-((double)k / (double)m * (double)n)), k); + } + + /** + * \brief Calculate number of items in the filter + * \param m Number of bits in the filter + * \param k Number of hash functions + * \param p Probability of false positives + * \return Number of items in the filter + */ + static size_t NumberOfItems(size_t m, size_t k, double p) { + return (size_t)std::ceil( + -((double)m / (double)k * + std::log(1.0 - std::exp(std::log(p) / (double)k)))); + } + + /** + * \brief Calculate number of bits in the filter + * \param n Number of items in the filter + * \param p Probability of false positives + * \return Number of bits in the filter + */ + static size_t NumberOfBits(size_t n, double p) { + return (size_t)std::ceil((double)n * std::log(p) / + std::log(1.0 / std::pow(2.0, std::log(2.0)))); + } + + /** + * \brief Calculate number of bits in the filter + * \param n Number of items in the filter + * \param k Number of hash functions + * \param p Probability of false positives + * \return Number of bits in the filter + */ + static size_t NumberOfBits(size_t n, size_t k, double p) { + return (size_t)std::ceil(-((double)k * (double)n / + std::log(1.0 - std::pow(p, 1.0 / (double)k)))); + } + + /** + * \brief Calculate number of bytes in the filter + * \param n Number of items in the filter + * \param p Probability of false positives + * \return Number of bytes in the filter + */ + static size_t NumberOfBytes(size_t n, double p) { + return ((NumberOfBits(n, p) + 7) >> 3); + } + + /** + * \brief Calculate number of bits in the filter + * \param n Number of items in the filter + * \param k Number of hash functions + * \param p Probability of false positives + * \return Number of bits in the filter + */ + static size_t NumberOfBytes(size_t n, size_t k, double p) { + return ((NumberOfBits(n, k, p) + 7) >> 3); + } + + /** + * \brief Calculate number of hash functions + * \param n Number of items in the filter + * \param m Number of bits in the filter + * \return Number of hash functions + */ + static size_t NumberOfHash(size_t n, size_t m) { + return (size_t)std::round((double)m / (double)n * std::log(2.0)); + } +}; + +/*! Bloom Filter + */ +template +class BloomFilter { + public: + //! Constructor + BloomFilter(void) {} + + //! Constructor + BloomFilter(size_t n, double p) { + if (n > 0 && p > 0.0 && p < 1.0) { + capacity_ = n; + bits_count_ = BloomFilterCalculator::NumberOfBits(n, K, p); + bits_count_ = ((bits_count_ + 31) >> 5) << 5; + probability_ = BloomFilterCalculator::Probability(n, bits_count_, K); + bitset_ = new uint32_t[bits_count_ >> 5]; + memset(bitset_, 0, (bits_count_ >> 3)); + } + } + + //! Constructor + BloomFilter(BloomFilter &&rhs) + : bitset_(rhs.bitset_), + bits_count_(rhs.bits_count_), + capacity_(rhs.capacity_), + count_(rhs.count_), + probability_(rhs.probability_) { + rhs.bitset_ = nullptr; + rhs.bits_count_ = 0u; + rhs.capacity_ = 0u; + rhs.count_ = 0u; + rhs.probability_ = 0u; + } + + //! Destructor + ~BloomFilter(void) { + delete[] bitset_; + } + + //! Test if the filter is valid + bool is_valid(void) const { + return (bitset_ != nullptr); + } + + //! Reset the bloom filter + bool reset(size_t n, double p) { + if (n <= 0 || p <= 0.0 || p >= 1.0) { + return false; + } + delete[] bitset_; + capacity_ = n; + count_ = 0u; + bits_count_ = BloomFilterCalculator::NumberOfBits(n, K, p); + bits_count_ = ((bits_count_ + 31) >> 5) << 5; + probability_ = BloomFilterCalculator::Probability(n, bits_count_, K); + bitset_ = new (std::nothrow) uint32_t[bits_count_ >> 5]; + if (!bitset_) { + return false; + } + memset(bitset_, 0, (bits_count_ >> 3)); + return true; + } + + //! Clear the bloom filter + void clear(void) { + if (bitset_) { + memset(bitset_, 0, (bits_count_ >> 3)); + count_ = 0u; + } + } + + //! Insert a item into bloom filter + template ...>::value && + sizeof...(TArgs) == K>::type> + bool insert(TArgs... vals) { + if (count_ >= capacity_) { + return false; + } + this->set_bits(vals...); + ++count_; + return true; + } + + //! Force insert a item into bloom filter + template ...>::value && + sizeof...(TArgs) == K>::type> + void force_insert(TArgs... vals) { + this->set_bits(vals...); + ++count_; + } + + //! Insert a item into bloom filter + template ...>::value && + sizeof...(TArgs) == K>::type> + bool has(TArgs... vals) const { + return this->test_bits(vals...); + } + + //! Retrieve count of bits in bloom filter + size_t bits_count(void) const { + return bits_count_; + } + + //! Retrieve capacity of bloom filter + size_t capacity(void) const { + return capacity_; + } + + //! Retrieve count of items in bloom filter + size_t count(void) const { + return count_; + } + + //! Retrieve probability of false positives + double probability(void) const { + return probability_; + } + + protected: + //! Disable them + BloomFilter(const BloomFilter &) = delete; + BloomFilter &operator=(const BloomFilter &) = delete; + + //! Set bits in bloom filter + template + void set_bits(TArg val) { + size_t num = static_cast(val) % bits_count_; + bitset_[num >> 5] |= (1u << (num & 0x1f)); + } + + //! Set bits in bloom filter + template + void set_bits(TArg val, TArgs... vals) { + this->set_bits(val); + this->set_bits(vals...); + } + + //! Test bits in bloom filter + template + bool test_bits(TArg val) const { + size_t num = static_cast(val) % bits_count_; + return ((bitset_[num >> 5] & (1u << (num & 0x1f))) != 0); + } + + //! Test bits in bloom filter + template + bool test_bits(TArg val, TArgs... vals) const { + if (!this->test_bits(val)) { + return false; + } + return this->test_bits(vals...); + } + + private: + uint32_t *bitset_{nullptr}; + size_t bits_count_{0u}; + size_t capacity_{0u}; + size_t count_{0u}; + double probability_{0.0}; +}; + +/*! Bloom Filter (Special) + */ +template <> +struct BloomFilter<0> {}; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/container/params.cc b/src/ailego/container/params.cc new file mode 100644 index 0000000..33e731a --- /dev/null +++ b/src/ailego/container/params.cc @@ -0,0 +1,139 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +//! Global environ variable +extern char **environ; + +namespace zvec { +namespace ailego { + +static void ParseFromJsonObject(const ailego::JsonObject &obj, Params *params) { + for (ailego::JsonObject::const_iterator it = obj.begin(); it != obj.end(); + ++it) { + const ailego::JsonValue &val = it->value(); + + if (val.is_boolean()) { + params->set(it->key().as_stl_string(), val.as_bool()); + } else if (val.is_integer()) { + params->set(it->key().as_stl_string(), + static_cast(val.as_integer())); + } else if (val.is_float()) { + params->set(it->key().as_stl_string(), val.as_float()); + } else if (val.is_string()) { + params->set(it->key().as_stl_string(), + val.as_string().decode().as_stl_string()); + } else if (val.is_object()) { + Params subparams; + ParseFromJsonObject(val.as_object(), &subparams); + params->set(it->key().as_stl_string(), std::move(subparams)); + } + } +} + +bool Params::ParseFromBuffer(const std::string &buf, Params *params) { + ailego::JsonValue val; + ailego::JsonParser parser; + + parser.set_comment(true); + parser.set_simple(true); + parser.set_squote(true); + parser.set_unstrict(false); + if (!parser.parse(buf.c_str(), &val)) { + return false; + } + + if (!val.is_object()) { + return false; + } + ParseFromJsonObject(val.as_object(), params); + return true; +} + +void Params::ParseFromEnvironment(Params *params) { + // Dump all environ string + for (size_t i = 0; environ[i]; ++i) { + const char *env = environ[i]; + const char *p = std::strchr(env, '='); + if (p) { + params->set(std::string(env, p - env), std::string(p + 1)); + } + } +} + +static void SerializeToJsonObject(const Params ¶ms, + ailego::JsonObject *obj) { + for (const auto &it : params.hypercube().cubes()) { + const ailego::Cube &cube = it.second; + const char *key = it.first.c_str(); + + if (cube.compatible()) { + const auto &val = cube.unsafe_cast(); + ailego::JsonString str(val.data(), val.size()); + obj->set(key, ailego::JsonValue(str.encode())); + } else if (cube.compatible()) { + obj->set(key, + ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + obj->set(key, ailego::JsonValue(cube.unsafe_cast())); + } else if (cube.compatible()) { + ailego::JsonObject subobj; + SerializeToJsonObject(cube.unsafe_cast(), &subobj); + obj->set(key, ailego::JsonValue(subobj)); + } else { + LOG_WARN("Unsupported serializing \'%s\' <%s>.", key, cube.type().name()); + } + } +} + +void Params::SerializeToBuffer(const Params ¶ms, std::string *buf) { + if (buf != nullptr) { + ailego::JsonObject obj; + SerializeToJsonObject(params, &obj); + buf->assign(ailego::JsonValue(obj).as_json_string().as_stl_string()); + } +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/container/reservoir.h b/src/ailego/container/reservoir.h new file mode 100644 index 0000000..210d41a --- /dev/null +++ b/src/ailego/container/reservoir.h @@ -0,0 +1,140 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec { +namespace ailego { + +/*! Sampling Reservoir + */ +template > +class Reservoir { + public: + //! Constructor + Reservoir(size_t cnt) + : samples_(cnt), total_(0), mt_(std::random_device()()), pool_() { + pool_.reserve(samples_); + } + + //! Constructor + Reservoir(const Reservoir &rhs) + : samples_(rhs.samples_), + total_(rhs.total_), + mt_(std::random_device()()), + pool_(rhs.pool_) {} + + //! Constructor + Reservoir(Reservoir &&rhs) + : samples_(rhs.samples_), + total_(rhs.total_), + mt_(std::random_device()()), + pool_(std::move(rhs.pool_)) {} + + //! Destructor + ~Reservoir(void) {} + + //! Assignment + Reservoir &operator=(const Reservoir &rhs) { + samples_ = rhs.samples_; + total_ = rhs.total_; + pool_ = rhs.pool_; + return *this; + } + + //! Assignment + Reservoir &operator=(Reservoir &&rhs) { + samples_ = rhs.samples_; + total_ = rhs.total_; + pool_ = std::move(rhs.pool_); + return *this; + } + + //! Retrieve pool of reservoir + std::vector *mutable_pool(void) { + return &pool_; + } + + //! Retrieve pool of reservoir + const std::vector &pool(void) const { + return pool_; + } + + //! Retrieve count of samples + size_t samples(void) const { + return samples_; + } + + //! Retrieve total count of filling + size_t total(void) const { + return total_; + } + + //! Reset the reservoir + void reset(void) { + total_ = 0; + pool_.clear(); + pool_.reserve(samples_); + } + + //! Fill the reservoir + void fill(const T &item) { + if (samples_ > 0) { + if (pool_.size() >= samples_) { + std::uniform_int_distribution dt(0, total_); + size_t i = dt(mt_); + + if (i < samples_) { + pool_[i] = item; + } + } else { + pool_.push_back(item); + } + } + ++total_; + } + + //! Fill the reservoir + void fill(T &&item) { + if (samples_ > 0) { + if (pool_.size() >= samples_) { + std::uniform_int_distribution dt(0, total_); + size_t i = dt(mt_); + + if (i < samples_) { + pool_[i] = std::move(item); + } + } else { + pool_.push_back(std::move(item)); + } + } + ++total_; + } + + private: + //! Disable them + Reservoir(void) = delete; + + //! Members + size_t samples_; + size_t total_; + std::mt19937 mt_; + std::vector pool_; +}; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/container/vector_array.h b/src/ailego/container/vector_array.h new file mode 100644 index 0000000..fe7a914 --- /dev/null +++ b/src/ailego/container/vector_array.h @@ -0,0 +1,556 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec { +namespace ailego { + +/*! Numerical Vector Array + */ +template ::value>::type> +class NumericalVectorArray { + public: + //! Type of value + using ValueType = typename NumericalVector::ValueType; + + //! Constructor + NumericalVectorArray(void) {} + + //! Constructor + explicit NumericalVectorArray(size_t dim) : dimension_(dim) {} + + //! Constructor + NumericalVectorArray(const NumericalVectorArray &rhs) + : dimension_(rhs.dimension_), buffer_(rhs.buffer_) {} + + //! Constructor + NumericalVectorArray(NumericalVectorArray &&rhs) + : dimension_(rhs.dimension_), buffer_(std::move(rhs.buffer_)) {} + + //! Assignment + NumericalVectorArray &operator=(const NumericalVectorArray &rhs) { + dimension_ = rhs.dimension_; + buffer_ = rhs.buffer_; + return *this; + } + + //! Assignment + NumericalVectorArray &operator=(NumericalVectorArray &&rhs) { + dimension_ = rhs.dimension_; + buffer_ = std::move(rhs.buffer_); + return *this; + } + + //! Overloaded operator [] + ValueType *operator[](size_t i) { + return (reinterpret_cast(&buffer_[0]) + i * dimension_); + } + + //! Overloaded operator [] + const ValueType *operator[](size_t i) const { + return (reinterpret_cast(buffer_.data()) + + i * dimension_); + } + + //! Append a vector + void append(const ValueType *vec, size_t dim) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + buffer_.append(reinterpret_cast(vec), + dim * sizeof(ValueType)); + } + + //! Append vectors + void append(const ValueType *vec, size_t dim, size_t cnt) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + buffer_.append(reinterpret_cast(vec), + cnt * dim * sizeof(ValueType)); + } + + //! Append a vector + void append(const NumericalVector &vec) { + this->append(vec.data(), vec.dimension()); + } + + //! Replace a vector + void replace(size_t index, const ValueType *vec, size_t dim) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + size_t element_size = dim * sizeof(ValueType); + buffer_.replace(index * element_size, element_size, + reinterpret_cast(vec), element_size); + } + + //! Replace a vector + void replace(size_t index, const ValueType *vec, size_t dim, size_t cnt) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + size_t element_size = dim * sizeof(ValueType); + size_t total = element_size * cnt; + buffer_.replace(index * element_size, total, + reinterpret_cast(vec), total); + } + + //! Replace a vector + void replace(size_t index, const NumericalVector &vec) { + this->replace(index, vec.data(), vec.dimension()); + } + + //! Request a change in capacity + void reserve(size_t n) { + buffer_.reserve(n * dimension_ * sizeof(ValueType)); + } + + //! Resize the array to a length of n elements + void resize(size_t n) { + buffer_.resize(n * dimension_ * sizeof(ValueType)); + } + + //! Clear the vector array + void clear(void) { + buffer_.clear(); + } + + //! Reset the vector array + void reset(size_t dim) { + dimension_ = dim; + buffer_.clear(); + } + + //! Requests the removal of unused capacity. + void shrink_to_fit(void) { + buffer_.shrink_to_fit(); + } + + //! Retrieve pointer of data + ValueType *data(void) { + return reinterpret_cast(&buffer_[0]); + } + + //! Retrieve pointer of data + const ValueType *data(void) const { + return reinterpret_cast(buffer_.data()); + } + + //! Retrieve pointer of data + ValueType *at(size_t i) { + if (ailego_unlikely(i >= this->count())) { + throw std::out_of_range("Index overflow"); + } + return (reinterpret_cast(&buffer_[0]) + i * dimension_); + } + + //! Retrieve pointer of data + const ValueType *at(size_t i) const { + if (ailego_unlikely(i >= this->count())) { + throw std::out_of_range("Index overflow"); + } + return (reinterpret_cast(buffer_.data()) + + i * dimension_); + } + + //! Test if the array is empty + bool empty(void) const { + return buffer_.empty(); + } + + //! Retrieve count of vectors + size_t count(void) const { + return (dimension_ > 0 ? buffer_.size() / (dimension_ * sizeof(ValueType)) + : 0u); + } + + //! Retrieve dimension of vector + size_t dimension(void) const { + return dimension_; + } + + //! Retrieve size of array in bytes + size_t bytes(void) const { + return buffer_.size(); + } + + private: + size_t dimension_{0u}; + std::string buffer_{}; +}; + +/*! Nibble Vector Array + */ +template ::value>::type> +class NibbleVectorArray { + public: + //! Type of value + using ValueType = typename NibbleVector::ValueType; + using StoreType = typename NibbleVector::StoreType; + + //! Constructor + NibbleVectorArray(void) {} + + //! Constructor + explicit NibbleVectorArray(size_t dim) + : dimension_((dim + (sizeof(ValueType) << 1) - 1) / + (sizeof(ValueType) << 1) * sizeof(ValueType) + << 1) {} + + //! Constructor + NibbleVectorArray(const NibbleVectorArray &rhs) + : dimension_(rhs.dimension_), buffer_(rhs.buffer_) {} + + //! Constructor + NibbleVectorArray(NibbleVectorArray &&rhs) + : dimension_(rhs.dimension_), buffer_(std::move(rhs.buffer_)) {} + + //! Assignment + NibbleVectorArray &operator=(const NibbleVectorArray &rhs) { + dimension_ = rhs.dimension_; + buffer_ = rhs.buffer_; + return *this; + } + + //! Assignment + NibbleVectorArray &operator=(NibbleVectorArray &&rhs) { + dimension_ = rhs.dimension_; + buffer_ = std::move(rhs.buffer_); + return *this; + } + + //! Overloaded operator [] + StoreType *operator[](size_t i) { + return reinterpret_cast(&buffer_[0] + i * (dimension_ >> 1)); + } + + //! Overloaded operator [] + const StoreType *operator[](size_t i) const { + return reinterpret_cast(&buffer_[0] + + i * (dimension_ >> 1)); + } + + //! Append a vector + void append(const StoreType *vec, size_t dim) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + buffer_.append(reinterpret_cast(vec), dim >> 1); + } + + //! Append vectors + void append(const StoreType *vec, size_t dim, size_t cnt) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + buffer_.append(reinterpret_cast(vec), cnt * (dim >> 1)); + } + + //! Append a vector + void append(const NibbleVector &vec) { + this->append(vec.data(), vec.dimension()); + } + + //! Replace a vector + void replace(size_t index, const StoreType *vec, size_t dim) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + size_t element_size = (dim >> 1); + buffer_.replace(index * element_size, element_size, + reinterpret_cast(vec), element_size); + } + + //! Replace a vector + void replace(size_t index, const StoreType *vec, size_t dim, size_t cnt) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + size_t element_size = (dim >> 1); + size_t total = element_size * cnt; + buffer_.replace(index * element_size, total, + reinterpret_cast(vec), total); + } + + //! Replace a vector + void replace(size_t index, const NibbleVector &vec) { + this->replace(index, vec.data(), vec.dimension()); + } + + //! Request a change in capacity + void reserve(size_t n) { + buffer_.reserve(n * (dimension_ >> 1)); + } + + //! Resize the array to a length of n elements + void resize(size_t n) { + buffer_.resize(n * (dimension_ >> 1)); + } + + //! Clear the vector array + void clear(void) { + buffer_.clear(); + } + + //! Reset the vector array + void reset(size_t dim) { + dimension_ = (dim + (sizeof(ValueType) << 1) - 1) / + (sizeof(ValueType) << 1) * sizeof(ValueType) + << 1; + buffer_.clear(); + } + + //! Requests the removal of unused capacity. + void shrink_to_fit(void) { + buffer_.shrink_to_fit(); + } + + //! Retrieve pointer of data + StoreType *data(void) { + return reinterpret_cast(&buffer_[0]); + } + + //! Retrieve pointer of data + const StoreType *data(void) const { + return reinterpret_cast(buffer_.data()); + } + + //! Retrieve pointer of data + StoreType *at(size_t i) { + if (ailego_unlikely(i >= this->count())) { + throw std::out_of_range("Index overflow"); + } + return reinterpret_cast(&buffer_[0] + i * (dimension_ >> 1)); + } + + //! Retrieve pointer of data + const StoreType *at(size_t i) const { + if (ailego_unlikely(i >= this->count())) { + throw std::out_of_range("Index overflow"); + } + return reinterpret_cast(buffer_.data() + + i * (dimension_ >> 1)); + } + + //! Test if the array is empty + bool empty(void) const { + return buffer_.empty(); + } + + //! Retrieve count of vectors + size_t count(void) const { + return (dimension_ > 1 ? buffer_.size() / (dimension_ >> 1) : 0u); + } + + //! Retrieve dimension of vector + size_t dimension(void) const { + return dimension_; + } + + //! Retrieve size of array in bytes + size_t bytes(void) const { + return buffer_.size(); + } + + private: + size_t dimension_{0u}; + std::string buffer_{}; +}; + +/*! Binary Vector Array + */ +template ::value>::type> +class BinaryVectorArray { + public: + //! Type of value + using ValueType = typename BinaryVector::ValueType; + + //! Constructor + BinaryVectorArray(void) {} + + //! Constructor + explicit BinaryVectorArray(size_t dim) + : dimension_((dim + (sizeof(ValueType) << 3) - 1) / + (sizeof(ValueType) << 3) * (sizeof(ValueType) << 3)) {} + + //! Constructor + BinaryVectorArray(const BinaryVectorArray &rhs) + : dimension_(rhs.dimension_), buffer_(rhs.buffer_) {} + + //! Constructor + BinaryVectorArray(BinaryVectorArray &&rhs) + : dimension_(rhs.dimension_), buffer_(std::move(rhs.buffer_)) {} + + //! Assignment + BinaryVectorArray &operator=(const BinaryVectorArray &rhs) { + dimension_ = rhs.dimension_; + buffer_ = rhs.buffer_; + return *this; + } + + //! Assignment + BinaryVectorArray &operator=(BinaryVectorArray &&rhs) { + dimension_ = rhs.dimension_; + buffer_ = std::move(rhs.buffer_); + return *this; + } + + //! Overloaded operator [] + ValueType *operator[](size_t i) { + return reinterpret_cast(&buffer_[0] + i * (dimension_ >> 3)); + } + + //! Overloaded operator [] + const ValueType *operator[](size_t i) const { + return reinterpret_cast(buffer_.data() + + i * (dimension_ >> 3)); + } + + //! Append a vector + void append(const ValueType *vec, size_t dim) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + buffer_.append(reinterpret_cast(vec), (dim >> 3)); + } + + //! Append vectors + void append(const ValueType *vec, size_t dim, size_t cnt) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + buffer_.append(reinterpret_cast(vec), cnt * (dim >> 3)); + } + + //! Append a vector + void append(const BinaryVector &vec) { + this->append(vec.data(), vec.dimension()); + } + + //! Replace a vector + void replace(size_t index, const ValueType *vec, size_t dim) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + size_t element_size = (dim >> 3); + buffer_.replace(index * element_size, element_size, + reinterpret_cast(vec), element_size); + } + + //! Replace a vector + void replace(size_t index, const ValueType *vec, size_t dim, size_t cnt) { + if (ailego_unlikely(dim != dimension_)) { + throw std::length_error("Unmatched dimension"); + } + size_t element_size = (dim >> 3); + size_t total = element_size * cnt; + buffer_.replace(index * element_size, total, + reinterpret_cast(vec), total); + } + + //! Replace a vector + void replace(size_t index, const BinaryVector &vec) { + this->replace(index, vec.data(), vec.dimension()); + } + + //! Request a change in capacity + void reserve(size_t n) { + buffer_.reserve(n * (dimension_ >> 3)); + } + + //! Resize the array to a length of n elements + void resize(size_t n) { + buffer_.resize(n * (dimension_ >> 3)); + } + + //! Clear the vector array + void clear(void) { + buffer_.clear(); + } + + //! Reset the vector array + void reset(size_t dim) { + dimension_ = (dim + (sizeof(ValueType) << 3) - 1) / + (sizeof(ValueType) << 3) * (sizeof(ValueType) << 3); + buffer_.clear(); + } + + //! Requests the removal of unused capacity. + void shrink_to_fit(void) { + buffer_.shrink_to_fit(); + } + + //! Retrieve pointer of data + ValueType *data(void) { + return reinterpret_cast(&buffer_[0]); + } + + //! Retrieve pointer of data + const ValueType *data(void) const { + return reinterpret_cast(buffer_.data()); + } + + //! Retrieve pointer of data + ValueType *at(size_t i) { + if (ailego_unlikely(i >= this->count())) { + throw std::out_of_range("Index overflow"); + } + return reinterpret_cast(&buffer_[0] + i * (dimension_ >> 3)); + } + + //! Retrieve pointer of data + const ValueType *at(size_t i) const { + if (ailego_unlikely(i >= this->count())) { + throw std::out_of_range("Index overflow"); + } + return reinterpret_cast(buffer_.data() + + i * (dimension_ >> 3)); + } + + //! Test if the array is empty + bool empty(void) const { + return buffer_.empty(); + } + + //! Retrieve count of vectors + size_t count(void) const { + return (dimension_ > 0 ? buffer_.size() / (dimension_ >> 3) : 0u); + } + + //! Retrieve dimension of vector + size_t dimension(void) const { + return dimension_; + } + + //! Retrieve size of array in bytes + size_t bytes(void) const { + return buffer_.size(); + } + + private: + size_t dimension_{0u}; + std::string buffer_{}; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/encoding/json/mod_json.c b/src/ailego/encoding/json/mod_json.c new file mode 100644 index 0000000..ea6a110 --- /dev/null +++ b/src/ailego/encoding/json/mod_json.c @@ -0,0 +1,3591 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include + +#ifndef MOD_JSON_TOKEN_DEFOPTS +#define MOD_JSON_TOKEN_DEFOPTS 0 /* default options of token */ +#endif +#ifndef MOD_JSON_TOKEN_DEFOBJDEP +#define MOD_JSON_TOKEN_DEFOBJDEP 64 /* default objects depth of token */ +#endif +#ifndef MOD_JSON_TOKEN_DEFARRDEP +#define MOD_JSON_TOKEN_DEFARRDEP 64 /* default arrays depth of token */ +#endif +#ifndef MOD_JSON_STRING_DEFSIZE +#define MOD_JSON_STRING_DEFSIZE 32 /* default started size of string */ +#endif +#ifndef MOD_JSON_ARRAY_DEFSIZE +#define MOD_JSON_ARRAY_DEFSIZE 32 /* default started size of array */ +#endif +#ifndef MOD_JSON_OBJECT_DEFSIZE +#define MOD_JSON_OBJECT_DEFSIZE 32 /* default started size of object */ +#endif + +#ifndef mod_json_malloc +#define mod_json_malloc malloc +#endif +#ifndef mod_json_free +#define mod_json_free free +#endif + +#ifdef __GNUC__ +#define mod_json_likely(x) __builtin_expect(!!(x), 1) +#define mod_json_unlikely(x) __builtin_expect(!!(x), 0) +#else +#define mod_json_likely(x) (x) +#define mod_json_unlikely(x) (x) +#endif + +#define mod_json_minus_if_ne_zero(COND) \ + if (mod_json_unlikely((COND) != 0)) return (-1) + +#define mod_json_minus_if_false(COND) \ + if (mod_json_unlikely(!(COND))) return (-1) + +#define mod_json_null_if_ne_zero(COND) \ + if (mod_json_unlikely((COND) != 0)) return (NULL) + +#define mod_json_null_if_false(COND) \ + if (mod_json_unlikely(!(COND))) return (NULL) + +#if defined(_MSC_VER) +#pragma warning(disable : 4200) +#define strtoull _strtoui64 +#define snprintf(buf, size, format, ...) \ + _snprintf_s(buf, size, _TRUNCATE, format, ##__VA_ARGS__) +#endif +#define mod_json_utils_snprintf snprintf +#define mod_json_utils_strtoi strtoull +#define mod_json_utils_strtof strtod +#define mod_json_utils_strlen strlen + +/*! JSON Token + */ +struct mod_json_token { + mod_json_state_t state; + mod_json_error_t error; + mod_json_cchar_t *context; + mod_json_size_t options; + mod_json_size_t object_max_depth; + mod_json_size_t array_max_depth; + mod_json_size_t object_depth; + mod_json_size_t array_depth; + mod_json_event_t event_code; + mod_json_event_proc event_proc; + mod_json_void_t *param; + mod_json_char_t tags[0]; +}; + +typedef struct mod_json_parser mod_json_parser_t; + +/*! JSON Parser + */ +struct mod_json_parser { + mod_json_string_t *key; + mod_json_value_t *val_null; + mod_json_value_t *val_true; + mod_json_value_t *val_false; + mod_json_value_t *val_zero; + mod_json_value_t *val_zerof; + mod_json_value_t *val_empty; + mod_json_value_t *vals[0]; +}; + +static inline mod_json_size_t mod_json_utils_clp2(mod_json_size_t n) { + n = n - 1; + n = n | (n >> 1); + n = n | (n >> 2); + n = n | (n >> 4); + n = n | (n >> 8); + n = n | (n >> 16); + return (n + 1); +} + +static inline mod_json_size_t mod_json_utils_itostr(mod_json_char_t *buf, + mod_json_integer_t val) { + mod_json_char_t *pos, *first, *last; + + pos = buf; + if (val < 0) { + *pos++ = '-'; + val = -val; + } + + /* save pointer to first digit */ + first = pos; + + do { + /* convert to ASCII and store */ + *pos++ = (mod_json_char_t)(val % 10 + '0'); + + /* next digit */ + val /= 10; + + } while (val > 0); + + *pos = '\0'; + + /* save pointer to last digit */ + last = pos - 1; + + /* reverse digit string */ + while (first < last) { + mod_json_char_t temp = *first; + *first++ = *last; + *last-- = temp; + } + return (mod_json_size_t)(pos - buf); +} + +static inline mod_json_float_t mod_json_utils_pow10(int n) { + /* 1e-308...1e308: 617 * 8 bytes = 4936 bytes */ + static const mod_json_float_t etab[] = { + 1e-308, 1e-307, 1e-306, 1e-305, 1e-304, 1e-303, 1e-302, 1e-301, 1e-300, + 1e-299, 1e-298, 1e-297, 1e-296, 1e-295, 1e-294, 1e-293, 1e-292, 1e-291, + 1e-290, 1e-289, 1e-288, 1e-287, 1e-286, 1e-285, 1e-284, 1e-283, 1e-282, + 1e-281, 1e-280, 1e-279, 1e-278, 1e-277, 1e-276, 1e-275, 1e-274, 1e-273, + 1e-272, 1e-271, 1e-270, 1e-269, 1e-268, 1e-267, 1e-266, 1e-265, 1e-264, + 1e-263, 1e-262, 1e-261, 1e-260, 1e-259, 1e-258, 1e-257, 1e-256, 1e-255, + 1e-254, 1e-253, 1e-252, 1e-251, 1e-250, 1e-249, 1e-248, 1e-247, 1e-246, + 1e-245, 1e-244, 1e-243, 1e-242, 1e-241, 1e-240, 1e-239, 1e-238, 1e-237, + 1e-236, 1e-235, 1e-234, 1e-233, 1e-232, 1e-231, 1e-230, 1e-229, 1e-228, + 1e-227, 1e-226, 1e-225, 1e-224, 1e-223, 1e-222, 1e-221, 1e-220, 1e-219, + 1e-218, 1e-217, 1e-216, 1e-215, 1e-214, 1e-213, 1e-212, 1e-211, 1e-210, + 1e-209, 1e-208, 1e-207, 1e-206, 1e-205, 1e-204, 1e-203, 1e-202, 1e-201, + 1e-200, 1e-199, 1e-198, 1e-197, 1e-196, 1e-195, 1e-194, 1e-193, 1e-192, + 1e-191, 1e-190, 1e-189, 1e-188, 1e-187, 1e-186, 1e-185, 1e-184, 1e-183, + 1e-182, 1e-181, 1e-180, 1e-179, 1e-178, 1e-177, 1e-176, 1e-175, 1e-174, + 1e-173, 1e-172, 1e-171, 1e-170, 1e-169, 1e-168, 1e-167, 1e-166, 1e-165, + 1e-164, 1e-163, 1e-162, 1e-161, 1e-160, 1e-159, 1e-158, 1e-157, 1e-156, + 1e-155, 1e-154, 1e-153, 1e-152, 1e-151, 1e-150, 1e-149, 1e-148, 1e-147, + 1e-146, 1e-145, 1e-144, 1e-143, 1e-142, 1e-141, 1e-140, 1e-139, 1e-138, + 1e-137, 1e-136, 1e-135, 1e-134, 1e-133, 1e-132, 1e-131, 1e-130, 1e-129, + 1e-128, 1e-127, 1e-126, 1e-125, 1e-124, 1e-123, 1e-122, 1e-121, 1e-120, + 1e-119, 1e-118, 1e-117, 1e-116, 1e-115, 1e-114, 1e-113, 1e-112, 1e-111, + 1e-110, 1e-109, 1e-108, 1e-107, 1e-106, 1e-105, 1e-104, 1e-103, 1e-102, + 1e-101, 1e-100, 1e-99, 1e-98, 1e-97, 1e-96, 1e-95, 1e-94, 1e-93, + 1e-92, 1e-91, 1e-90, 1e-89, 1e-88, 1e-87, 1e-86, 1e-85, 1e-84, + 1e-83, 1e-82, 1e-81, 1e-80, 1e-79, 1e-78, 1e-77, 1e-76, 1e-75, + 1e-74, 1e-73, 1e-72, 1e-71, 1e-70, 1e-69, 1e-68, 1e-67, 1e-66, + 1e-65, 1e-64, 1e-63, 1e-62, 1e-61, 1e-60, 1e-59, 1e-58, 1e-57, + 1e-56, 1e-55, 1e-54, 1e-53, 1e-52, 1e-51, 1e-50, 1e-49, 1e-48, + 1e-47, 1e-46, 1e-45, 1e-44, 1e-43, 1e-42, 1e-41, 1e-40, 1e-39, + 1e-38, 1e-37, 1e-36, 1e-35, 1e-34, 1e-33, 1e-32, 1e-31, 1e-30, + 1e-29, 1e-28, 1e-27, 1e-26, 1e-25, 1e-24, 1e-23, 1e-22, 1e-21, + 1e-20, 1e-19, 1e-18, 1e-17, 1e-16, 1e-15, 1e-14, 1e-13, 1e-12, + 1e-11, 1e-10, 1e-9, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3, + 1e-2, 1e-1, 1e+0, 1e+1, 1e+2, 1e+3, 1e+4, 1e+5, 1e+6, + 1e+7, 1e+8, 1e+9, 1e+10, 1e+11, 1e+12, 1e+13, 1e+14, 1e+15, + 1e+16, 1e+17, 1e+18, 1e+19, 1e+20, 1e+21, 1e+22, 1e+23, 1e+24, + 1e+25, 1e+26, 1e+27, 1e+28, 1e+29, 1e+30, 1e+31, 1e+32, 1e+33, + 1e+34, 1e+35, 1e+36, 1e+37, 1e+38, 1e+39, 1e+40, 1e+41, 1e+42, + 1e+43, 1e+44, 1e+45, 1e+46, 1e+47, 1e+48, 1e+49, 1e+50, 1e+51, + 1e+52, 1e+53, 1e+54, 1e+55, 1e+56, 1e+57, 1e+58, 1e+59, 1e+60, + 1e+61, 1e+62, 1e+63, 1e+64, 1e+65, 1e+66, 1e+67, 1e+68, 1e+69, + 1e+70, 1e+71, 1e+72, 1e+73, 1e+74, 1e+75, 1e+76, 1e+77, 1e+78, + 1e+79, 1e+80, 1e+81, 1e+82, 1e+83, 1e+84, 1e+85, 1e+86, 1e+87, + 1e+88, 1e+89, 1e+90, 1e+91, 1e+92, 1e+93, 1e+94, 1e+95, 1e+96, + 1e+97, 1e+98, 1e+99, 1e+100, 1e+101, 1e+102, 1e+103, 1e+104, 1e+105, + 1e+106, 1e+107, 1e+108, 1e+109, 1e+110, 1e+111, 1e+112, 1e+113, 1e+114, + 1e+115, 1e+116, 1e+117, 1e+118, 1e+119, 1e+120, 1e+121, 1e+122, 1e+123, + 1e+124, 1e+125, 1e+126, 1e+127, 1e+128, 1e+129, 1e+130, 1e+131, 1e+132, + 1e+133, 1e+134, 1e+135, 1e+136, 1e+137, 1e+138, 1e+139, 1e+140, 1e+141, + 1e+142, 1e+143, 1e+144, 1e+145, 1e+146, 1e+147, 1e+148, 1e+149, 1e+150, + 1e+151, 1e+152, 1e+153, 1e+154, 1e+155, 1e+156, 1e+157, 1e+158, 1e+159, + 1e+160, 1e+161, 1e+162, 1e+163, 1e+164, 1e+165, 1e+166, 1e+167, 1e+168, + 1e+169, 1e+170, 1e+171, 1e+172, 1e+173, 1e+174, 1e+175, 1e+176, 1e+177, + 1e+178, 1e+179, 1e+180, 1e+181, 1e+182, 1e+183, 1e+184, 1e+185, 1e+186, + 1e+187, 1e+188, 1e+189, 1e+190, 1e+191, 1e+192, 1e+193, 1e+194, 1e+195, + 1e+196, 1e+197, 1e+198, 1e+199, 1e+200, 1e+201, 1e+202, 1e+203, 1e+204, + 1e+205, 1e+206, 1e+207, 1e+208, 1e+209, 1e+210, 1e+211, 1e+212, 1e+213, + 1e+214, 1e+215, 1e+216, 1e+217, 1e+218, 1e+219, 1e+220, 1e+221, 1e+222, + 1e+223, 1e+224, 1e+225, 1e+226, 1e+227, 1e+228, 1e+229, 1e+230, 1e+231, + 1e+232, 1e+233, 1e+234, 1e+235, 1e+236, 1e+237, 1e+238, 1e+239, 1e+240, + 1e+241, 1e+242, 1e+243, 1e+244, 1e+245, 1e+246, 1e+247, 1e+248, 1e+249, + 1e+250, 1e+251, 1e+252, 1e+253, 1e+254, 1e+255, 1e+256, 1e+257, 1e+258, + 1e+259, 1e+260, 1e+261, 1e+262, 1e+263, 1e+264, 1e+265, 1e+266, 1e+267, + 1e+268, 1e+269, 1e+270, 1e+271, 1e+272, 1e+273, 1e+274, 1e+275, 1e+276, + 1e+277, 1e+278, 1e+279, 1e+280, 1e+281, 1e+282, 1e+283, 1e+284, 1e+285, + 1e+286, 1e+287, 1e+288, 1e+289, 1e+290, 1e+291, 1e+292, 1e+293, 1e+294, + 1e+295, 1e+296, 1e+297, 1e+298, 1e+299, 1e+300, 1e+301, 1e+302, 1e+303, + 1e+304, 1e+305, 1e+306, 1e+307, 1e+308}; + return (n < -308 ? 0.0 : etab[n + 308]); +} + +static inline mod_json_cchar_t *mod_json_utils_strskpb(mod_json_cchar_t *cstr) { + static const mod_json_char_t blanks[256] = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, '\t', '\n', '\v', '\f', '\r', 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ' ', 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; + + while (*(blanks + *cstr)) { + ++cstr; + } + return cstr; +} + +static inline mod_json_cchar_t *mod_json_utils_strskpc1( + mod_json_cchar_t *cstr) { + mod_json_char_t c; + + while ((c = *cstr++) != '\0') { + if (c == '\r' || c == '\n') { + return mod_json_utils_strskpb(cstr); + } + } + return (cstr - 1); +} + +static inline mod_json_cchar_t *mod_json_utils_strskpc2( + mod_json_cchar_t *cstr) { + mod_json_char_t c; + + while ((c = *cstr++) != '\0') { + /* asterisk, slash */ + if (c == '*' && *cstr == '/') { + return mod_json_utils_strskpb(cstr + 1); + } + } + return (cstr - 1); +} + +static inline mod_json_cchar_t *mod_json_utils_strskp(mod_json_cchar_t *cstr) { + cstr = mod_json_utils_strskpb(cstr); + + /* treat it as comments? */ + while (*cstr == '/') { + mod_json_char_t c = *(cstr + 1); /* second char */ + + if (c == '/') { + /* two slashes */ + cstr = mod_json_utils_strskpc1(cstr + 2); + } else if (c == '*') { + /* slash, asterisk */ + cstr = mod_json_utils_strskpc2(cstr + 2); + } else { + /* invalid format */ + break; + } + } + return cstr; +} + +static inline int mod_json_utils_char2hex(mod_json_char_t ch) { + static const mod_json_char_t char2hex[256] = { + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 0, 1, 2, 3, 4, 5, + 6, 7, 8, 9, 16, 16, 16, 16, 16, 16, 16, 10, 11, 12, 13, 14, 15, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 10, 11, 12, 13, 14, 15, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, + 16, 16, 16, 16}; + return *(char2hex + ch); +} + +static inline mod_json_cchar_t *mod_json_utils_strfquo(mod_json_cchar_t *cstr, + mod_json_char_t quo) { + mod_json_char_t c; + + for (c = *cstr; c != quo; c = *(++cstr)) { + if ((mod_json_uchar_t)c <= 0x1f) { + return NULL; + } + + if (c != '\\') { + continue; + } + + /* next char */ + switch (*(++cstr)) { + case '\"': + case '/': + case 'b': + case 'f': + case '\\': + case 'n': + case 'r': + case 't': + /* ignore next char */ + break; + + case 'u': + if (mod_json_utils_char2hex(*(cstr + 1)) > 15) { + return NULL; + } + if (mod_json_utils_char2hex(*(cstr + 2)) > 15) { + return NULL; + } + if (mod_json_utils_char2hex(*(cstr + 3)) > 15) { + return NULL; + } + if (mod_json_utils_char2hex(*(cstr + 4)) > 15) { + return NULL; + } + cstr += 4; + break; + + default: + /* invalid */ + return NULL; + } + } + /* found it */ + return cstr; +} + +static inline mod_json_cchar_t *mod_json_utils_strfquo2(mod_json_cchar_t *cstr, + mod_json_char_t quo) { + mod_json_char_t c; + + for (c = *cstr; c; c = *(++cstr)) { + if (c == quo) { + /* found it */ + return cstr; + } + + if (c == '\\') { + /* ignore next char */ + if (*(++cstr) == '\0') { + break; + } + } + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_utils_strfsep(mod_json_cchar_t *cstr) { + mod_json_char_t c; + + while ((c = *cstr++) != '\0') { + switch (c) { + case ':': + case ' ': + case '\t': + case '\r': + case '\n': + case '\f': + case '\v': + return (cstr - 1); + } + } + return (cstr - 1); +} + +static inline mod_json_cchar_t *mod_json_utils_strfsep2( + mod_json_cchar_t *cstr) { + mod_json_char_t c; + + while ((c = *cstr++) != '\0') { + switch (c) { + case ':': + case ' ': + case '\t': + case '\r': + case '\n': + case '\f': + case '\v': + return (cstr - 1); + + case '/': + if (*cstr == '/' || *cstr == '*') { + return (cstr - 1); + } + } + } + return (cstr - 1); +} + +static inline mod_json_char_t *mod_json_utils_uni2utf8(mod_json_char_t *buf, + mod_json_size_t size, + mod_json_uchar_t high, + mod_json_uchar_t low) { + /* convert to UTF-8 */ + if (high >= 0x8) { + /* 0800 - FFFF | 1110xxxx 10xxxxxx 10xxxxxx */ + if (size >= 3) { + *buf++ = (mod_json_char_t)(0xE0 | (high >> 4)); + *buf++ = (mod_json_char_t)(0x80 | ((high & 0xF) << 2) | (low >> 6)); + *buf++ = (mod_json_char_t)(0x80 | (low & 0x3F)); + return buf; + } + } else if (high > 0 || low >= 0x80) { + /* 0080 - 07FF | 110xxxxx 10xxxxxx */ + if (size >= 2) { + *buf++ = (mod_json_char_t)(0xC0 | (high << 2) | (low >> 6)); + *buf++ = (mod_json_char_t)(0x80 | (low & 0x3F)); + return buf; + } + } else { + /* 0000 - 007F | 0xxxxxxx */ + if (size >= 1) { + *buf++ = (mod_json_char_t)(low); + return buf; + } + } + return (mod_json_char_t *)0; +} + +mod_json_value_t *mod_json_value_set_null(void) { + mod_json_value_t *val; + + /* create a value */ + val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t)); + mod_json_null_if_false(val); + + val->refer = 1; + val->type = mod_json_type_null; + val->data.c_int = 0; + return val; +} + +mod_json_value_t *mod_json_value_set_object(mod_json_object_t *obj) { + mod_json_value_t *val; + + /* create a value */ + val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t)); + mod_json_null_if_false(val); + + val->refer = 1; + val->type = mod_json_type_object; + val->data.c_obj = obj ? mod_json_object_grab(obj) : NULL; + return val; +} + +mod_json_value_t *mod_json_value_set_array(mod_json_array_t *arr) { + mod_json_value_t *val; + + /* create a value */ + val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t)); + mod_json_null_if_false(val); + + val->refer = 1; + val->type = mod_json_type_array; + val->data.c_arr = arr ? mod_json_array_grab(arr) : NULL; + return val; +} + +mod_json_value_t *mod_json_value_set_string(mod_json_string_t *str) { + mod_json_value_t *val; + + /* create a value */ + val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t)); + mod_json_null_if_false(val); + + val->refer = 1; + val->type = mod_json_type_string; + val->data.c_str = str ? mod_json_string_grab(str) : NULL; + return val; +} + +mod_json_value_t *mod_json_value_set_buffer(mod_json_cchar_t *buf, + mod_json_size_t len) { + mod_json_value_t *val; + mod_json_string_t *str; + + /* create a value */ + val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t)); + mod_json_null_if_false(val); + + /* create a string */ + str = mod_json_string_set(buf, len); + if (mod_json_unlikely(!str)) { + mod_json_free(val); + return NULL; + } + + val->refer = 1; + val->type = mod_json_type_string; + val->data.c_str = str; + return val; +} + +mod_json_value_t *mod_json_value_set_integer(mod_json_integer_t num) { + mod_json_value_t *val; + + /* create a value */ + val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t)); + mod_json_null_if_false(val); + + val->refer = 1; + val->type = mod_json_type_integer; + val->data.c_int = num; + return val; +} + +mod_json_value_t *mod_json_value_set_float(mod_json_float_t dbl) { + mod_json_value_t *val; + + /* create a value */ + val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t)); + mod_json_null_if_false(val); + + val->refer = 1; + val->type = mod_json_type_float; + val->data.c_float = dbl; + return val; +} + +mod_json_value_t *mod_json_value_set_boolean(mod_json_boolean_t bol) { + mod_json_value_t *val; + + /* create a value */ + val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t)); + mod_json_null_if_false(val); + + val->refer = 1; + val->type = mod_json_type_boolean; + val->data.c_bool = bol ? MOD_JSON_TRUE : MOD_JSON_FALSE; + return val; +} + +static inline void mod_json_value_clear(mod_json_value_t *val) { + switch (val->type) { + case mod_json_type_object: + mod_json_object_unset(val->data.c_obj); + break; + + case mod_json_type_array: + mod_json_array_unset(val->data.c_arr); + break; + + case mod_json_type_string: + mod_json_string_unset(val->data.c_str); + break; + + default: + break; + } +} + +void mod_json_value_assign_null(mod_json_value_t *val) { + if (val) { + mod_json_value_clear(val); + val->type = mod_json_type_null; + val->data.c_int = 0; + } +} + +void mod_json_value_assign_object(mod_json_value_t *val, + mod_json_object_t *obj) { + if (val) { + mod_json_value_clear(val); + val->type = mod_json_type_object; + val->data.c_obj = obj ? mod_json_object_grab(obj) : NULL; + } +} + +void mod_json_value_assign_array(mod_json_value_t *val, mod_json_array_t *arr) { + if (val) { + mod_json_value_clear(val); + val->type = mod_json_type_array; + val->data.c_arr = arr ? mod_json_array_grab(arr) : NULL; + } +} + +void mod_json_value_assign_string(mod_json_value_t *val, + mod_json_string_t *str) { + if (val) { + mod_json_value_clear(val); + val->type = mod_json_type_string; + val->data.c_str = str ? mod_json_string_grab(str) : NULL; + } +} + +void mod_json_value_assign_integer(mod_json_value_t *val, + mod_json_integer_t num) { + if (val) { + mod_json_value_clear(val); + val->type = mod_json_type_integer; + val->data.c_int = num; + } +} + +void mod_json_value_assign_float(mod_json_value_t *val, mod_json_float_t dbl) { + if (val) { + mod_json_value_clear(val); + val->type = mod_json_type_float; + val->data.c_float = dbl; + } +} + +void mod_json_value_assign_boolean(mod_json_value_t *val, + mod_json_boolean_t bol) { + if (val) { + mod_json_value_clear(val); + val->type = mod_json_type_boolean; + val->data.c_bool = bol ? MOD_JSON_TRUE : MOD_JSON_FALSE; + } +} + +void mod_json_value_assign(mod_json_value_t *dst, mod_json_value_t *src) { + if (!dst || dst == src) { + return; + } + + if (!src) { + /* treat as JSON null */ + mod_json_value_assign_null(dst); + return; + } + + switch (src->type) { + case mod_json_type_boolean: + mod_json_value_assign_boolean(dst, src->data.c_bool); + break; + + case mod_json_type_integer: + mod_json_value_assign_integer(dst, src->data.c_int); + break; + + case mod_json_type_float: + mod_json_value_assign_float(dst, src->data.c_float); + break; + + case mod_json_type_string: + mod_json_value_assign_string(dst, src->data.c_str); + break; + + case mod_json_type_array: + mod_json_value_assign_array(dst, src->data.c_arr); + break; + + case mod_json_type_object: + mod_json_value_assign_object(dst, src->data.c_obj); + break; + + default: + mod_json_value_assign_null(dst); + break; + } +} + +static inline int mod_json_value_merge_array(mod_json_value_t *val, + mod_json_array_t *arr) { + if (val->type != mod_json_type_array || !val->data.c_arr) { + mod_json_value_assign_array(val, arr); + return 0; + } + + if (arr) { + if (mod_json_array_is_shared(val->data.c_arr)) { + mod_json_array_put(val->data.c_arr); + val->data.c_arr = mod_json_array_clone(val->data.c_arr); + } + return mod_json_array_merge(val->data.c_arr, arr); + } + return 0; +} + +static inline int mod_json_value_merge_object(mod_json_value_t *val, + mod_json_object_t *obj) { + if (val->type != mod_json_type_object || !val->data.c_obj) { + mod_json_value_assign_object(val, obj); + return 0; + } + + if (obj) { + if (mod_json_object_is_shared(val->data.c_obj)) { + mod_json_object_put(val->data.c_obj); + val->data.c_obj = mod_json_object_clone(val->data.c_obj); + } + return mod_json_object_merge(val->data.c_obj, obj); + } + return 0; +} + +int mod_json_value_merge(mod_json_value_t *dst, mod_json_value_t *src) { + mod_json_minus_if_false(dst && dst != src); + + if (!src) { + mod_json_value_assign_null(dst); + return 0; + } + + switch (src->type) { + case mod_json_type_boolean: + mod_json_value_assign_boolean(dst, src->data.c_bool); + break; + + case mod_json_type_integer: + mod_json_value_assign_integer(dst, src->data.c_int); + break; + + case mod_json_type_float: + mod_json_value_assign_float(dst, src->data.c_float); + break; + + case mod_json_type_string: + mod_json_value_assign_string(dst, src->data.c_str); + break; + + case mod_json_type_array: + return mod_json_value_merge_array(dst, src->data.c_arr); + + case mod_json_type_object: + return mod_json_value_merge_object(dst, src->data.c_obj); + + default: + mod_json_value_assign_null(dst); + break; + } + return 0; +} + +mod_json_object_t *mod_json_value_object(mod_json_value_t *val) { + if (val && val->type == mod_json_type_object) { + return (val->data.c_obj); + } + return NULL; +} + +mod_json_array_t *mod_json_value_array(mod_json_value_t *val) { + if (val && val->type == mod_json_type_array) { + return (val->data.c_arr); + } + return NULL; +} + +mod_json_string_t *mod_json_value_string(mod_json_value_t *val) { + if (val && val->type == mod_json_type_string) { + return (val->data.c_str); + } + return NULL; +} + +mod_json_cchar_t *mod_json_value_cstring(mod_json_value_t *val) { + if (val && val->type == mod_json_type_string) { + return mod_json_string_cstr(val->data.c_str); + } + return NULL; +} + +mod_json_float_t mod_json_value_float(mod_json_value_t *val) { + if (val) { + switch (val->type) { + case mod_json_type_boolean: + return (val->data.c_bool ? 1.0 : 0.0); + + case mod_json_type_integer: + return (mod_json_float_t)(val->data.c_int); + + case mod_json_type_float: + return (val->data.c_float); + + case mod_json_type_string: + return mod_json_string_float(val->data.c_str); + + default: + break; + } + } + return (0.0); +} + +mod_json_boolean_t mod_json_value_boolean(mod_json_value_t *val) { + if (val) { + switch (val->type) { + case mod_json_type_null: + return MOD_JSON_FALSE; + + case mod_json_type_object: + return (mod_json_object_count(val->data.c_obj) != 0); + + case mod_json_type_array: + return (mod_json_array_count(val->data.c_arr) != 0); + + case mod_json_type_string: + return (mod_json_string_length(val->data.c_str) != 0); + + case mod_json_type_integer: + return (val->data.c_int != 0); + + case mod_json_type_float: + return (val->data.c_float != 0); + + case mod_json_type_boolean: + return (val->data.c_bool); + + default: + break; + } + } + return MOD_JSON_FALSE; +} + +mod_json_integer_t mod_json_value_integer(mod_json_value_t *val) { + if (val) { + switch (val->type) { + case mod_json_type_boolean: + return (val->data.c_bool ? 1 : 0); + + case mod_json_type_integer: + return (val->data.c_int); + + case mod_json_type_float: + return (mod_json_integer_t)(val->data.c_float); + + case mod_json_type_string: + return mod_json_string_integer(val->data.c_str); + + default: + break; + } + } + return (0); +} + +mod_json_value_t *mod_json_value_clone(mod_json_value_t *val) { + if (val) { + switch (val->type) { + case mod_json_type_null: + return mod_json_value_set_null(); + + case mod_json_type_object: + return mod_json_value_set_object(val->data.c_obj); + + case mod_json_type_array: + return mod_json_value_set_array(val->data.c_arr); + + case mod_json_type_string: + return mod_json_value_set_string(val->data.c_str); + + case mod_json_type_integer: + return mod_json_value_set_integer(val->data.c_int); + + case mod_json_type_float: + return mod_json_value_set_float(val->data.c_float); + + case mod_json_type_boolean: + return mod_json_value_set_boolean(val->data.c_bool); + + default: + break; + } + } + return NULL; +} + +static inline mod_json_boolean_t mod_json_value_is_equal_float( + mod_json_float_t lhs, mod_json_float_t rhs) { + mod_json_float_t diff = lhs - rhs; + return ((diff < DBL_EPSILON) && (diff > -DBL_EPSILON)); +} + +mod_json_boolean_t mod_json_value_is_equal(mod_json_value_t *lhs, + mod_json_value_t *rhs) { + if (lhs == rhs) { + /* The same pointer */ + return MOD_JSON_TRUE; + } + + if (lhs && rhs && lhs->type == rhs->type) { + switch (lhs->type) { + case mod_json_type_null: + return MOD_JSON_TRUE; + + case mod_json_type_object: + return mod_json_object_is_equal(lhs->data.c_obj, rhs->data.c_obj); + + case mod_json_type_array: + return mod_json_array_is_equal(lhs->data.c_arr, rhs->data.c_arr); + + case mod_json_type_string: + return (mod_json_string_compare(lhs->data.c_str, rhs->data.c_str) == 0); + + case mod_json_type_integer: + return (lhs->data.c_int == rhs->data.c_int); + + case mod_json_type_float: + return mod_json_value_is_equal_float(lhs->data.c_float, + rhs->data.c_float); + + case mod_json_type_boolean: + return ((!lhs->data.c_bool) == (!rhs->data.c_bool)); + + default: + break; + } + } + return MOD_JSON_FALSE; +} + +void mod_json_value_unset(mod_json_value_t *val) { + if (val && mod_json_value_put(val) <= 0) { + mod_json_value_clear(val); + mod_json_free(val); + } +} + +static inline int mod_json_string_expand(mod_json_string_t *str, + mod_json_size_t size) { + mod_json_char_t *cstr; + mod_json_size_t len; + + size = mod_json_utils_clp2(size); + if (size < MOD_JSON_STRING_DEFSIZE) { + size = MOD_JSON_STRING_DEFSIZE; + } + mod_json_minus_if_false(size > str->size); + + cstr = (mod_json_char_t *)mod_json_malloc(size * sizeof(mod_json_char_t)); + mod_json_minus_if_false(cstr); + + len = (mod_json_size_t)(str->last - str->first); + if (len != 0) { + memcpy(cstr, str->first, len + 1); + } else { + *cstr = '\0'; /* terminal character */ + } + mod_json_free(str->first); + str->first = cstr; + str->last = cstr + len; + str->size = size; + + /* success */ + return 0; +} + +int mod_json_string_reserve(mod_json_string_t *str, mod_json_size_t n) { + mod_json_minus_if_false(str); + + if (str->size >= n + 1) { + /* needn't grow */ + return 0; + } + return mod_json_string_expand(str, n + 1); +} + +static inline mod_json_string_t *mod_json_string_malloc(mod_json_size_t size) { + mod_json_string_t *str; + mod_json_char_t *buf; + + buf = (mod_json_char_t *)mod_json_malloc(size * sizeof(mod_json_char_t)); + mod_json_null_if_false(buf); + + str = (mod_json_string_t *)mod_json_malloc(sizeof(mod_json_string_t)); + if (mod_json_unlikely(!str)) { + mod_json_free(buf); + return NULL; + } + + str->refer = 1; + str->size = size; + str->first = buf; + str->last = buf; + *buf = '\0'; + return str; +} + +int mod_json_string_assign(mod_json_string_t *str, mod_json_cchar_t *cstr, + mod_json_size_t len) { + mod_json_string_reset(str); + mod_json_minus_if_ne_zero(mod_json_string_reserve(str, len)); + + if (cstr && len) { + memcpy(str->first, cstr, len); + } + str->last = str->first + len; + *(str->last) = '\0'; + + /* success */ + return 0; +} + +static inline mod_json_string_t *mod_json_string_set_empty(void) { + return mod_json_string_malloc(MOD_JSON_STRING_DEFSIZE); +} + +static inline mod_json_string_t *mod_json_string_set_cstr( + mod_json_cchar_t *cstr, mod_json_size_t len) { + mod_json_string_t *str; + + str = mod_json_string_malloc(mod_json_utils_clp2(len + 1)); + mod_json_null_if_false(str); + + str->last = str->first + len; + memcpy(str->first, cstr, len); + *(str->last) = '\0'; + return str; +} + +mod_json_string_t *mod_json_string_set(mod_json_cchar_t *cstr, + mod_json_size_t len) { + return ((cstr && len) ? mod_json_string_set_cstr(cstr, len) + : mod_json_string_set_empty()); +} + +void mod_json_string_unset(mod_json_string_t *str) { + if (str && mod_json_string_put(str) <= 0) { + mod_json_free(str->first); + mod_json_free(str); + } +} + +void mod_json_string_reset(mod_json_string_t *str) { + if (str) { + str->last = str->first; + *(str->first) = '\0'; + } +} + +static inline int mod_json_string_add_char(mod_json_string_t *str, + mod_json_char_t ch) { + mod_json_size_t need; + + need = (mod_json_size_t)(str->last - str->first) + 2; + if (need > str->size) { + mod_json_minus_if_ne_zero(mod_json_string_expand(str, need)); + } + + /* append to string */ + *(str->last++) = ch; + *(str->last) = '\0'; + + /* success */ + return 0; +} + +static inline int mod_json_string_add_cstr(mod_json_string_t *str, + mod_json_cchar_t *cstr, + mod_json_size_t len) { + if (cstr && len) { + mod_json_size_t need; + + need = len + (mod_json_size_t)(str->last - str->first) + 1; + if (need > str->size) { + mod_json_minus_if_ne_zero(mod_json_string_expand(str, need)); + } + + /* append to string */ + memcpy(str->last, cstr, len); + str->last += len; + *(str->last) = '\0'; + } + + /* success */ + return 0; +} + +static inline int mod_json_string_add_jstr(mod_json_string_t *str, + mod_json_string_t *val) { + return mod_json_string_add_cstr(str, val->first, + (mod_json_size_t)(val->last - val->first)); +} + +int mod_json_string_add(mod_json_string_t *str, mod_json_string_t *val) { + return mod_json_string_add_jstr(str, val); +} + +int mod_json_string_append(mod_json_string_t *str, mod_json_cchar_t *cstr, + mod_json_size_t len) { + return mod_json_string_add_cstr(str, cstr, len); +} + +mod_json_size_t mod_json_string_hash(mod_json_string_t *str) { + mod_json_size_t hash = 1; + + if (str) { + mod_json_cchar_t *iter = str->first; + mod_json_cchar_t *last = str->last; + + for (; iter != last; ++iter) { + mod_json_size_t c = (mod_json_size_t)(*iter); + hash = hash * 131 + c; + } + } + return hash; +} + +int mod_json_string_compare(mod_json_string_t *str1, mod_json_string_t *str2) { + mod_json_size_t len1 = 0, len2 = 0; + + if (str1 == str2) { + /* The same pointer */ + return 0; + } + + if (str1) { + len1 = (mod_json_size_t)(str1->last - str1->first); + if (str2) { + len2 = (mod_json_size_t)(str2->last - str2->first); + if (len1 == len2) { + return memcmp(str1->first, str2->first, len1); + } + } + } else { + /* The first string is null, and the second string it not null. */ + len2 = (mod_json_size_t)(str2->last - str2->first); + } + return (int)(len1 - len2); +} + +mod_json_integer_t mod_json_string_integer(mod_json_string_t *str) { + return (str ? (mod_json_integer_t)mod_json_utils_strtoi(str->first, NULL, 0) + : 0); +} + +mod_json_float_t mod_json_string_float(mod_json_string_t *str) { + return (str ? mod_json_utils_strtof(str->first, NULL) : 0.0); +} + +static inline int mod_json_string_flat(mod_json_string_t *dst, + mod_json_string_t *src) { + static mod_json_cchar_t *flattab[32] = { + "\\u0000", "\\u0001", "\\u0002", "\\u0003", "\\u0004", "\\u0005", + "\\u0006", "\\u0007", "\\b", "\\t", "\\n", "\\u000b", + "\\f", "\\r", "\\u000e", "\\u000f", "\\u0010", "\\u0011", + "\\u0012", "\\u0013", "\\u0014", "\\u0015", "\\u0016", "\\u0017", + "\\u0018", "\\u0019", "\\u001a", "\\u001b", "\\u001c", "\\u001d", + "\\u001e", "\\u001f"}; + + /* length of items in flat table */ + static const mod_json_uchar_t flatlen[32] = {6, 6, 6, 6, 6, 6, 6, 6, 2, 2, 2, + 6, 2, 2, 6, 6, 6, 6, 6, 6, 6, 6, + 6, 6, 6, 6, 6, 6, 6, 6, 6, 6}; + + mod_json_cchar_t *first = src->first; + mod_json_cchar_t *iter = src->first; + mod_json_cchar_t *last = src->last; + + /* the whole string */ + for (; iter != last; ++iter) { + int c = *iter; + + if ((mod_json_uchar_t)c <= 0x1f) { + if (iter > first) { + mod_json_minus_if_ne_zero(mod_json_string_add_cstr( + dst, first, (mod_json_size_t)(iter - first))); + } + mod_json_minus_if_ne_zero( + mod_json_string_add_cstr(dst, flattab[c], flatlen[c])); + + /* skip current character */ + first = iter + 1; + } else if (c == '\"' || c == '\\') { + if (iter > first) { + mod_json_minus_if_ne_zero(mod_json_string_add_cstr( + dst, first, (mod_json_size_t)(iter - first))); + } + mod_json_minus_if_ne_zero(mod_json_string_add_char(dst, '\\')); + + /* don't skip current character */ + first = iter; + } + } + + if (iter > first) { + mod_json_minus_if_ne_zero( + mod_json_string_add_cstr(dst, first, (mod_json_size_t)(iter - first))); + } + + /* success */ + return 0; +} + +static inline int mod_json_string_unflat(mod_json_string_t *dst, + mod_json_string_t *src) { + enum { + state_normal, + state_rev_slash, + state_digit_1, + state_digit_2, + state_digit_3, + state_digit_4 + } state; + + mod_json_char_t *pbuf = dst->first; + mod_json_char_t *pend = dst->first + dst->size; + mod_json_cchar_t *iter = src->first; + mod_json_cchar_t *last = src->last; + mod_json_uchar_t high = 0; + mod_json_uchar_t low = 0; + + /* the whole string */ + for (state = state_normal; iter != last; ++iter) { + int c = *iter; + + switch (state) { + case state_normal: + if (c != '\\') { + mod_json_minus_if_false(pbuf < pend); + *pbuf++ = (mod_json_char_t)c; + } else { + /* '\\' in process */ + state = state_rev_slash; + } + break; + + case state_rev_slash: + mod_json_minus_if_false(pbuf < pend); + + switch (c) { + case '\"': + state = state_normal; + *pbuf++ = '\"'; + break; + case '/': + state = state_normal; + *pbuf++ = '/'; + break; + case 'b': + state = state_normal; + *pbuf++ = '\b'; + break; + case 'f': + state = state_normal; + *pbuf++ = '\f'; + break; + case '\\': + state = state_normal; + *pbuf++ = '\\'; + break; + case 'n': + state = state_normal; + *pbuf++ = '\n'; + break; + case 'r': + state = state_normal; + *pbuf++ = '\r'; + break; + case 't': + state = state_normal; + *pbuf++ = '\t'; + break; + case 'u': + state = state_digit_1; + break; + default: + return -1; + } + break; + + case state_digit_1: + if ((c = mod_json_utils_char2hex((mod_json_char_t)c)) > 15) { + /* invalid character */ + return -1; + } + high = (mod_json_uchar_t)(c << 4); + state = state_digit_2; + break; + + case state_digit_2: + if ((c = mod_json_utils_char2hex((mod_json_char_t)c)) > 15) { + /* invalid character */ + return -1; + } + high |= (mod_json_uchar_t)c; + state = state_digit_3; + break; + + case state_digit_3: + if ((c = mod_json_utils_char2hex((mod_json_char_t)c)) > 15) { + /* invalid character */ + return -1; + } + low = (mod_json_uchar_t)(c << 4); + state = state_digit_4; + break; + + case state_digit_4: + if ((c = mod_json_utils_char2hex((mod_json_char_t)c)) > 15) { + /* invalid character */ + return -1; + } + low |= (mod_json_uchar_t)c; + + /* decode as a UTF-8 string */ + pbuf = mod_json_utils_uni2utf8(pbuf, (mod_json_size_t)(pend - pbuf), + high, low); + if (!pbuf) { + /* lack of buffer */ + return -1; + } + state = state_normal; + break; + } + } + + if (state != state_normal) { + /* uncompleted state */ + return -1; + } + mod_json_minus_if_false(pbuf < pend); + + /* update the last pointer */ + *(dst->last = pbuf) = '\0'; + + /* success */ + return 0; +} + +mod_json_string_t *mod_json_string_encode(mod_json_string_t *src) { + mod_json_string_t *dst; + mod_json_null_if_false(src); + + dst = mod_json_string_malloc( + mod_json_utils_clp2((mod_json_size_t)(src->last - src->first) + 1)); + mod_json_null_if_false(dst); + + if (mod_json_unlikely(mod_json_string_flat(dst, src) != 0)) { + mod_json_string_unset(dst); + return NULL; + } + return dst; +} + +mod_json_string_t *mod_json_string_decode(mod_json_string_t *src) { + mod_json_string_t *dst; + mod_json_null_if_false(src); + + dst = mod_json_string_malloc( + mod_json_utils_clp2((mod_json_size_t)(src->last - src->first) + 1)); + mod_json_null_if_false(dst); + + if (mod_json_unlikely(mod_json_string_unflat(dst, src) != 0)) { + mod_json_string_unset(dst); + return NULL; + } + return dst; +} + +mod_json_array_t *mod_json_array_set(mod_json_size_t size) { + mod_json_array_t *arr; + mod_json_value_t **buf; + + size = (size ? mod_json_utils_clp2(size) : MOD_JSON_ARRAY_DEFSIZE); + buf = (mod_json_value_t **)mod_json_malloc(size * sizeof(mod_json_value_t *)); + mod_json_null_if_false(buf); + + /* create an array */ + arr = (mod_json_array_t *)mod_json_malloc(sizeof(mod_json_array_t)); + if (mod_json_unlikely(!arr)) { + mod_json_free(buf); + return NULL; + } + + arr->refer = 1; + arr->size = size; + arr->first = buf; + arr->last = buf; + return arr; +} + +mod_json_array_t *mod_json_array_clone(mod_json_array_t *arr) { + mod_json_array_t *arr2 = NULL; + + if (arr) { + arr2 = mod_json_array_set((mod_json_size_t)(arr->last - arr->first)); + if (arr2) { + mod_json_value_t **iter = arr->first; + + /* clone items */ + for (; iter != arr->last; ++iter) { + *arr2->last++ = *iter ? mod_json_value_grab(*iter) : NULL; + } + } + } + return arr2; +} + +mod_json_boolean_t mod_json_array_is_equal(mod_json_array_t *lhs, + mod_json_array_t *rhs) { + mod_json_value_t **itl, **itr; + + if (lhs == rhs) { + return MOD_JSON_TRUE; + } + + if (!lhs || !rhs || ((lhs->last - lhs->first) != (rhs->last - rhs->first))) { + return MOD_JSON_FALSE; + } + + /* compare items */ + for (itl = lhs->first, itr = rhs->first; itl != lhs->last; ++itl, ++itr) { + if (!mod_json_value_is_equal(*itl, *itr)) { + return MOD_JSON_FALSE; + } + } + return MOD_JSON_TRUE; +} + +void mod_json_array_unset(mod_json_array_t *arr) { + if (arr && mod_json_array_put(arr) <= 0) { + mod_json_value_t **iter = arr->first; + + for (; iter != arr->last; ++iter) { + mod_json_value_unset(*iter); + } + mod_json_free(arr->first); + mod_json_free(arr); + } +} + +void mod_json_array_reset(mod_json_array_t *arr) { + if (arr) { + mod_json_value_t **iter = arr->first; + + for (; iter != arr->last; ++iter) { + mod_json_value_unset(*iter); + } + arr->last = arr->first; + } +} + +static inline void mod_json_array_migrate(mod_json_array_t *arr, + mod_json_value_t **buf, + mod_json_size_t size) { + mod_json_size_t count = (mod_json_size_t)(arr->last - arr->first); + if (count > 0) { + memcpy(buf, arr->first, count * sizeof(mod_json_value_t *)); + } + mod_json_free(arr->first); + + arr->first = buf; + arr->last = buf + count; + arr->size = size; +} + +static inline int mod_json_array_expand(mod_json_array_t *arr, + mod_json_size_t n) { + mod_json_size_t size; + mod_json_value_t **vals; + + size = mod_json_utils_clp2(n); + if (size < MOD_JSON_ARRAY_DEFSIZE) { + size = MOD_JSON_ARRAY_DEFSIZE; + } + mod_json_minus_if_false(size > arr->size); + + vals = + (mod_json_value_t **)mod_json_malloc(size * sizeof(mod_json_value_t *)); + mod_json_minus_if_false(vals); + + /* use new buffer */ + mod_json_array_migrate(arr, vals, size); + + /* success */ + return 0; +} + +int mod_json_array_reserve(mod_json_array_t *arr, mod_json_size_t n) { + mod_json_minus_if_false(arr); + + if (arr->size >= n) { + /* needn't grow */ + return 0; + } + return mod_json_array_expand(arr, n); +} + +void mod_json_array_reverse(mod_json_array_t *arr) { + if (arr) { + mod_json_value_t **first = arr->first; + mod_json_value_t **last = arr->last - 1; + + while (first < last) { + mod_json_value_t *temp = *first; + *first++ = *last; + *last-- = temp; + } + } +} + +int mod_json_array_push(mod_json_array_t *arr, mod_json_value_t *val) { + mod_json_size_t count; + mod_json_minus_if_false(arr); + + count = (mod_json_size_t)(arr->last - arr->first); + if (count >= arr->size) { + mod_json_minus_if_ne_zero(mod_json_array_expand(arr, count + 1)); + } + + *arr->last++ = val ? mod_json_value_grab(val) : NULL; + return 0; +} + +void mod_json_array_pop(mod_json_array_t *arr) { + if (arr && arr->first != arr->last) { + mod_json_value_unset(*(--arr->last)); + } +} + +void mod_json_array_shift(mod_json_array_t *arr) { + if (arr && arr->first != arr->last) { + mod_json_value_t **it = arr->first; + mod_json_value_t **last = --arr->last; + + mod_json_value_unset(*it++); + for (; it <= last; ++it) { + *(it - 1) = *it; + } + } +} + +mod_json_value_t *mod_json_array_at(mod_json_array_t *arr, mod_json_size_t id) { + if (arr && ((arr->first + id) < arr->last)) { + return (arr->first[id]); + } + return NULL; +} + +int mod_json_array_merge(mod_json_array_t *dst, mod_json_array_t *src) { + long count, len1, len2; + + mod_json_minus_if_false(dst && src && dst != src); + + /* update length of array */ + len1 = (mod_json_size_t)(src->last - src->first); + len2 = (mod_json_size_t)(dst->last - dst->first); + mod_json_minus_if_false(len1 >= 0 && len2 >= 0); + + /* append empty values */ + count = len1 - len2; + for (; count > 0; --count) { + mod_json_array_push(dst, NULL); + } + + /* It must be assigned again. */ + len2 = (mod_json_size_t)(dst->last - dst->first); + count = (len1 < len2 ? len1 : len2); + + while ((count--) > 0) { + mod_json_value_t **iter1 = src->first + count; + mod_json_value_t **iter2 = dst->first + count; + + if (!(*iter2)) { + *iter2 = *iter1 ? mod_json_value_grab(*iter1) : NULL; + continue; + } + + if (mod_json_value_is_shared(*iter2)) { + mod_json_value_put(*iter2); + *iter2 = mod_json_value_clone(*iter2); + } + mod_json_value_merge(*iter2, *iter1); + } + + /* success */ + return 0; +} + +int mod_json_array_resize(mod_json_array_t *arr, mod_json_size_t n, + mod_json_value_t *val) { + mod_json_size_t orig; + + /* check input */ + mod_json_minus_if_false(arr); + + /* original count of array */ + orig = (mod_json_size_t)(arr->last - arr->first); + + if (orig < n) { + mod_json_value_t **iter; + + if (arr->size < n) { + mod_json_minus_if_ne_zero(mod_json_array_expand(arr, n)); + } + + iter = arr->last; + arr->last = arr->first + n; + + /* grab the first one, but get the others */ + *iter++ = val = val ? mod_json_value_grab(val) : NULL; + for (; iter != arr->last; ++iter) { + *iter = val ? mod_json_value_get(val) : NULL; + } + } else if (orig > n) { + mod_json_value_t **iter = arr->first + n; + + for (; iter != arr->last; ++iter) { + mod_json_value_unset(*iter); + *iter = NULL; + } + arr->last = arr->first + n; + } + + /* success */ + return 0; +} + +static inline void mod_json_pair_init(mod_json_pair_t *pair, + mod_json_string_t *key, + mod_json_value_t *val) { + pair->key = mod_json_string_grab(key); + pair->val = val ? mod_json_value_grab(val) : NULL; +} + +static inline void mod_json_pair_cleanup(mod_json_pair_t *pair) { + mod_json_string_unset(pair->key); + mod_json_value_unset(pair->val); + pair->key = NULL; + pair->val = NULL; +} + +mod_json_object_t *mod_json_object_set(mod_json_size_t size) { + mod_json_object_t *obj; + mod_json_pair_t *buf; + + size = (size ? mod_json_utils_clp2(size) : MOD_JSON_OBJECT_DEFSIZE); + buf = (mod_json_pair_t *)mod_json_malloc(size * sizeof(mod_json_pair_t)); + mod_json_null_if_false(buf); + + /* create a object */ + obj = (mod_json_object_t *)mod_json_malloc(sizeof(mod_json_object_t)); + if (mod_json_unlikely(!obj)) { + mod_json_free(buf); + return NULL; + } + + obj->refer = 1; + obj->size = size; + obj->first = buf; + obj->last = buf; + return obj; +} + +void mod_json_object_unset(mod_json_object_t *obj) { + if (obj && mod_json_object_put(obj) <= 0) { + mod_json_pair_t *iter = obj->first; + + for (; iter != obj->last; ++iter) { + mod_json_pair_cleanup(iter); + } + mod_json_free(obj->first); + mod_json_free(obj); + } +} + +void mod_json_object_reset(mod_json_object_t *obj) { + if (obj) { + mod_json_pair_t *iter = obj->first; + + for (; iter != obj->last; ++iter) { + mod_json_pair_cleanup(iter); + } + obj->last = obj->first; + } +} + +static inline void mod_json_object_migrate(mod_json_object_t *obj, + mod_json_pair_t *buf, + mod_json_size_t size) { + mod_json_size_t count = (mod_json_size_t)(obj->last - obj->first); + if (count > 0) { + memcpy(buf, obj->first, count * sizeof(mod_json_pair_t)); + } + mod_json_free(obj->first); + + obj->first = buf; + obj->last = buf + count; + obj->size = size; +} + +static inline int mod_json_object_expand(mod_json_object_t *obj, + mod_json_size_t n) { + mod_json_size_t size; + mod_json_pair_t *buf; + + size = mod_json_utils_clp2(n); + if (size < MOD_JSON_OBJECT_DEFSIZE) { + size = MOD_JSON_OBJECT_DEFSIZE; + } + mod_json_minus_if_false(size > obj->size); + + buf = (mod_json_pair_t *)mod_json_malloc(size * sizeof(mod_json_pair_t)); + mod_json_minus_if_false(buf); + + /* use new buffer */ + mod_json_object_migrate(obj, buf, size); + + /* success */ + return 0; +} + +static inline mod_json_pair_t *mod_json_object_find_pair(mod_json_object_t *obj, + mod_json_string_t *key, + mod_json_size_t *out) { + mod_json_pair_t *first = obj->first; + mod_json_pair_t *last = obj->last; + + while (first < last) { + mod_json_pair_t *middle = first + ((last - first) >> 2); + int diff = mod_json_string_compare(middle->key, key); + + if (diff < 0) { + first = middle + 1; + } else if (diff > 0) { + last = middle; + } else /*if (diff == 0)*/ + { + *out = (mod_json_size_t)(middle - obj->first); + return middle; + } + } + *out = (mod_json_size_t)(first - obj->first); + return NULL; +} + +mod_json_pair_t *mod_json_object_insert_force(mod_json_object_t *obj, + mod_json_size_t npos, + mod_json_string_t *key, + mod_json_value_t *val) { + mod_json_pair_t *iter, *pos; + mod_json_size_t count; + + count = (mod_json_size_t)(obj->last - obj->first); + if (count >= obj->size) { + mod_json_null_if_ne_zero(mod_json_object_expand(obj, count + 1)); + } + + pos = obj->first + npos; + iter = obj->last++; + for (; iter != pos; --iter) { + mod_json_pair_t *prev = iter - 1; + iter->key = prev->key; + iter->val = prev->val; + } + mod_json_pair_init(pos, key, val); + return pos; +} + +mod_json_pair_t *mod_json_object_insert(mod_json_object_t *obj, + mod_json_string_t *key, + mod_json_value_t *val) { + mod_json_size_t npos; + mod_json_null_if_false(obj && key); + + if (mod_json_object_find_pair(obj, key, &npos)) { + /* One in object */ + return NULL; + } + return mod_json_object_insert_force(obj, npos, key, val); +} + +mod_json_pair_t *mod_json_object_assign(mod_json_object_t *obj, + mod_json_string_t *key, + mod_json_value_t *val) { + mod_json_pair_t *elem = NULL; + + if (obj && key) { + mod_json_size_t npos; + + elem = mod_json_object_find_pair(obj, key, &npos); + if (elem) { + if (!elem->val) { + elem->val = val ? mod_json_value_grab(val) : NULL; + } else { + /* overwrite the old value */ + mod_json_value_assign(elem->val, val); + } + } else { + /* insert a new one */ + elem = mod_json_object_insert_force(obj, npos, key, val); + } + } + return elem; +} + +mod_json_pair_t *mod_json_object_touch(mod_json_object_t *obj, + mod_json_cchar_t *key) { + mod_json_pair_t *elem = NULL; + + if (obj && key) { + mod_json_string_t str; + mod_json_size_t npos; + + str.first = (mod_json_char_t *)key; + str.last = str.first + mod_json_utils_strlen(key); + + elem = mod_json_object_find_pair(obj, &str, &npos); + if (!elem) { + mod_json_string_t *jkey; + + /* insert a new one */ + jkey = + mod_json_string_set(key, (mod_json_size_t)mod_json_utils_strlen(key)); + elem = mod_json_object_insert_force(obj, npos, jkey, NULL); + mod_json_string_unset(jkey); + } + } + return elem; +} + +mod_json_object_t *mod_json_object_clone(mod_json_object_t *obj) { + mod_json_object_t *obj2 = NULL; + + if (obj) { + obj2 = mod_json_object_set((mod_json_size_t)(obj->last - obj->first)); + if (obj2) { + mod_json_pair_t *iter = obj->first; + + /* clone items */ + for (; iter != obj->last; ++iter) { + mod_json_pair_init(obj2->last++, iter->key, iter->val); + } + } + } + return obj2; +} + +mod_json_boolean_t mod_json_object_is_equal(mod_json_object_t *lhs, + mod_json_object_t *rhs) { + mod_json_pair_t *itl, *itr; + + if (lhs == rhs) { + /* The same pointer */ + return MOD_JSON_TRUE; + } + + if (!lhs || !rhs || ((lhs->last - lhs->first) != (rhs->last - rhs->first))) { + return MOD_JSON_FALSE; + } + + /* compare items */ + for (itl = lhs->first, itr = rhs->first; itl != lhs->last; ++itl, ++itr) { + if ((mod_json_string_compare(itl->key, itr->key) != 0) || + (!mod_json_value_is_equal(itl->val, itr->val))) { + return MOD_JSON_FALSE; + } + } + return MOD_JSON_TRUE; +} + +void mod_json_object_erase(mod_json_object_t *obj, mod_json_cchar_t *key) { + if (obj && key) { + mod_json_string_t str; + mod_json_pair_t *iter; + mod_json_size_t npos; + + str.first = (mod_json_char_t *)key; + str.last = str.first + mod_json_utils_strlen(key); + + iter = mod_json_object_find_pair(obj, &str, &npos); + if (iter) { + mod_json_pair_cleanup(iter++); + + for (; iter != obj->last; ++iter) { + mod_json_pair_t *prev = iter - 1; + prev->key = iter->key; + prev->val = iter->val; + } + --obj->last; + } + } +} + +mod_json_value_t *mod_json_object_at(mod_json_object_t *obj, + mod_json_cchar_t *key) { + if (obj && key) { + mod_json_string_t str; + mod_json_pair_t *elem; + mod_json_size_t npos; + + str.first = (mod_json_char_t *)key; + str.last = str.first + mod_json_utils_strlen(key); + + elem = mod_json_object_find_pair(obj, &str, &npos); + if (elem) { + return (elem->val); + } + } + return NULL; +} + +mod_json_pair_t *mod_json_object_find(mod_json_object_t *obj, + mod_json_cchar_t *key) { + if (obj && key) { + mod_json_string_t str; + mod_json_size_t npos; + + str.first = (mod_json_char_t *)key; + str.last = str.first + mod_json_utils_strlen(key); + + return mod_json_object_find_pair(obj, &str, &npos); + } + return NULL; +} + +int mod_json_object_merge(mod_json_object_t *dst, mod_json_object_t *src) { + mod_json_pair_t *iter; + + mod_json_minus_if_false(dst && src && dst != src); + + for (iter = src->first; iter != src->last; ++iter) { + mod_json_pair_t *elem; + mod_json_size_t npos; + + elem = mod_json_object_find_pair(dst, iter->key, &npos); + if (!elem) { + /* insert a new one */ + mod_json_object_insert_force(dst, npos, iter->key, iter->val); + continue; + } + + if (!elem->val) { + elem->val = iter->val ? mod_json_value_grab(iter->val) : NULL; + continue; + } + + if (mod_json_value_is_shared(elem->val)) { + mod_json_value_put(elem->val); + elem->val = mod_json_value_clone(elem->val); + } + mod_json_value_merge(elem->val, iter->val); + } + return 0; +} + +static inline mod_json_cchar_t *mod_json_token_strskp(mod_json_token_t *tok, + mod_json_cchar_t *cstr) { + if ((tok->options & MOD_JSON_COMMENT) == 0) { + return mod_json_utils_strskpb(cstr); + } + return mod_json_utils_strskp(cstr); +} + +static inline mod_json_cchar_t *mod_json_token_strfquo(mod_json_token_t *tok, + mod_json_cchar_t *cstr, + mod_json_char_t quo) { + if ((tok->options & MOD_JSON_UNSTRICT) == 0) { + return mod_json_utils_strfquo(cstr, quo); + } + return mod_json_utils_strfquo2(cstr, quo); +} + +static inline mod_json_cchar_t *mod_json_token_strfsep(mod_json_token_t *tok, + mod_json_cchar_t *cstr) { + if ((tok->options & MOD_JSON_COMMENT) == 0) { + return mod_json_utils_strfsep(cstr); + } + return mod_json_utils_strfsep2(cstr); +} + +mod_json_token_t *mod_json_token_create(mod_json_option_t *opt) { + mod_json_token_t *tok; + mod_json_size_t opts = MOD_JSON_TOKEN_DEFOPTS; + mod_json_size_t mobj = MOD_JSON_TOKEN_DEFOBJDEP; + mod_json_size_t marr = MOD_JSON_TOKEN_DEFARRDEP; + + if (opt) { + opts = opt->options; + + if (opt->object_depth > 0) { + mobj = opt->object_depth; + } + if (opt->array_depth > 0) { + marr = opt->array_depth; + } + } + + tok = (mod_json_token_t *)mod_json_malloc( + (mobj + marr) * sizeof(mod_json_char_t) + sizeof(mod_json_token_t)); + mod_json_null_if_false(tok); + + memset(tok, 0, sizeof(mod_json_token_t)); + tok->state = mod_json_state_null; + tok->error = mod_json_error_null; + tok->options = opts; + tok->object_max_depth = mobj; + tok->array_max_depth = marr; + return tok; +} + +void mod_json_token_destroy(mod_json_token_t *tok) { + mod_json_free(tok); +} + +static inline void mod_json_token_set_tag(mod_json_token_t *tok, + mod_json_char_t tag) { + mod_json_size_t depth = tok->object_depth + tok->array_depth; + if (depth != 0) { + tok->tags[depth - 1] = tag; + } +} + +static inline mod_json_char_t mod_json_token_tag(mod_json_token_t *tok) { + mod_json_size_t depth = tok->object_depth + tok->array_depth; + + /* type of current depth */ + return (depth ? tok->tags[depth - 1] : (mod_json_char_t)-1); +} + +mod_json_error_t mod_json_token_error(mod_json_token_t *tok) { + return (tok->error); +} + +mod_json_cchar_t *mod_json_token_context(mod_json_token_t *tok) { + return (tok->context); +} + +mod_json_state_t mod_json_token_state(mod_json_token_t *tok) { + return (tok->state); +} + +mod_json_size_t mod_json_token_object_depth(mod_json_token_t *tok) { + return (tok->object_depth); +} + +mod_json_size_t mod_json_token_array_depth(mod_json_token_t *tok) { + return (tok->array_depth); +} + +mod_json_size_t mod_json_token_depth(mod_json_token_t *tok) { + return (tok->object_depth + tok->array_depth); +} + +mod_json_size_t mod_json_token_max_object_depth(mod_json_token_t *tok) { + return (tok->object_max_depth); +} + +mod_json_size_t mod_json_token_max_array_depth(mod_json_token_t *tok) { + return (tok->array_max_depth); +} + +mod_json_size_t mod_json_token_max_depth(mod_json_token_t *tok) { + return (tok->object_max_depth + tok->array_max_depth); +} + +mod_json_void_t *mod_json_token_param(mod_json_token_t *tok) { + return (tok->param); +} + +void mod_json_token_set_param(mod_json_token_t *tok, mod_json_void_t *param) { + tok->param = param; +} + +void mod_json_token_set_event(mod_json_token_t *tok, mod_json_event_proc proc) { + tok->event_proc = proc; +} + +mod_json_event_t mod_json_token_event(mod_json_token_t *tok) { + return (tok->event_code); +} + +static inline int mod_json_token_invoke_field(mod_json_token_t *tok, + mod_json_cchar_t *val, + mod_json_size_t len) { + mod_json_event_proc invoke = tok->event_proc; + if (invoke) { + tok->event_code = mod_json_event_field; + return invoke(tok, (mod_json_void_t *)val, len); + } + return 0; +} + +static inline int mod_json_token_invoke_object(mod_json_token_t *tok) { + mod_json_event_proc invoke = tok->event_proc; + if (invoke) { + tok->event_code = mod_json_event_object; + return invoke(tok, NULL, 0); + } + return 0; +} + +static inline int mod_json_token_invoke_array(mod_json_token_t *tok) { + mod_json_event_proc invoke = tok->event_proc; + if (invoke) { + tok->event_code = mod_json_event_array; + return invoke(tok, NULL, 0); + } + return 0; +} + +static inline int mod_json_token_invoke_null(mod_json_token_t *tok) { + mod_json_event_proc invoke = tok->event_proc; + if (invoke) { + tok->event_code = mod_json_event_null; + return invoke(tok, NULL, 0); + } + return 0; +} + +static inline int mod_json_token_invoke_boolean(mod_json_token_t *tok, + mod_json_boolean_t val) { + mod_json_event_proc invoke = tok->event_proc; + if (invoke) { + tok->event_code = mod_json_event_boolean; + return invoke(tok, &val, sizeof(val)); + } + return 0; +} + +static inline int mod_json_token_invoke_integer(mod_json_token_t *tok, + mod_json_integer_t val) { + mod_json_event_proc invoke = tok->event_proc; + if (invoke) { + tok->event_code = mod_json_event_integer; + return invoke(tok, &val, sizeof(val)); + } + return 0; +} + +static inline int mod_json_token_invoke_float(mod_json_token_t *tok, + mod_json_float_t val) { + mod_json_event_proc invoke = tok->event_proc; + if (invoke) { + tok->event_code = mod_json_event_float; + return invoke(tok, &val, sizeof(val)); + } + return 0; +} + +static inline int mod_json_token_invoke_string(mod_json_token_t *tok, + mod_json_cchar_t *val, + mod_json_size_t len) { + mod_json_event_proc invoke = tok->event_proc; + if (invoke) { + tok->event_code = mod_json_event_string; + return invoke(tok, (mod_json_void_t *)val, len); + } + return 0; +} + +static inline mod_json_cchar_t *mod_json_token_start(mod_json_token_t *tok, + mod_json_cchar_t *cstr) { + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case '{': + tok->state = mod_json_state_object_start; + return (cstr + 1); + + case '[': + tok->state = mod_json_state_array_start; + return (cstr + 1); + + case '\0': + tok->error = mod_json_error_empty; + tok->context = cstr; + break; + + default: + tok->error = mod_json_error_start; + tok->context = cstr; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_value_null( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + mod_json_char_t c1 = *(cstr + 1); + mod_json_char_t c2 = *(cstr + 2); + mod_json_char_t c3 = *(cstr + 3); + + if ((c1 != 'u' && c1 != 'U') || (c2 != 'l' && c2 != 'L') || + (c3 != 'l' && c3 != 'L')) { + tok->error = mod_json_error_value; + tok->context = cstr; + return NULL; + } + + if (mod_json_token_invoke_null(tok) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + return (cstr + 4); +} + +static inline mod_json_cchar_t *mod_json_token_value_true( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + mod_json_char_t c1 = *(cstr + 1); + mod_json_char_t c2 = *(cstr + 2); + mod_json_char_t c3 = *(cstr + 3); + + if ((c1 != 'r' && c1 != 'R') || (c2 != 'u' && c2 != 'U') || + (c3 != 'e' && c3 != 'E')) { + tok->error = mod_json_error_value; + tok->context = cstr; + return NULL; + } + + if (mod_json_token_invoke_boolean(tok, MOD_JSON_TRUE) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + return (cstr + 4); +} + +static inline mod_json_cchar_t *mod_json_token_value_false( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + mod_json_char_t c1 = *(cstr + 1); + mod_json_char_t c2 = *(cstr + 2); + mod_json_char_t c3 = *(cstr + 3); + mod_json_char_t c4 = *(cstr + 4); + + if ((c1 != 'a' && c1 != 'A') || (c2 != 'l' && c2 != 'L') || + (c3 != 's' && c3 != 'S') || (c4 != 'e' && c4 != 'E')) { + tok->error = mod_json_error_value; + tok->context = cstr; + return NULL; + } + + if (mod_json_token_invoke_boolean(tok, MOD_JSON_FALSE) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + return (cstr + 5); +} + +static inline mod_json_cchar_t *mod_json_token_value_infinity( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + mod_json_char_t c1 = *(cstr + 1); + mod_json_char_t c2 = *(cstr + 2); + + if ((c1 != 'n' && c1 != 'N') || (c2 != 'f' && c2 != 'F')) { + tok->error = mod_json_error_value; + tok->context = cstr; + return NULL; + } + + if (mod_json_token_invoke_float(tok, MOD_JSON_INFINITY) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + return (cstr + 3); +} + +static inline mod_json_cchar_t *mod_json_token_value_string( + mod_json_token_t *tok, mod_json_cchar_t *cstr, mod_json_char_t quo) { + mod_json_cchar_t *cstr2 = mod_json_token_strfquo(tok, ++cstr, quo); + if (!cstr2) { + tok->error = mod_json_error_quote; + tok->context = cstr; + return NULL; + } + + if (mod_json_token_invoke_string(tok, cstr, + (mod_json_size_t)(cstr2 - cstr)) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + return (cstr2 + 1); +} + +static inline mod_json_cchar_t *mod_json_token_value_number( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + enum { number_integer, number_float } num_type = number_integer; + + mod_json_float_t dbl = 0.0; + uint32_t dig = 0; + uint64_t u64 = 0; + int32_t minus = 0; + int32_t exp_frac = 0, exp = 0; + + /* Parse minus */ + minus = *cstr; + if (minus == '-' || minus == '+') { + ++cstr; + } + + /* The first digit */ + if ((dig = (uint32_t)(*cstr - '0')) > 9) { + return NULL; + } + + /* Save the first digit */ + u64 = dig; + + /* Parse as 64bit integer */ + if (minus != '-') { + while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) { + if (u64 >= 1844674407370955161uLL) { + /* 2^64 - 1 = 18446744073709551615 */ + if (u64 != 1844674407370955161uLL || dig > 5) { + dbl = (mod_json_float_t)u64 * 10 + dig; + num_type = number_float; + break; + } + } + u64 = u64 * 10 + dig; + } + } else { + while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) { + /* 2^63 = 9223372036854775808 */ + if (u64 >= 922337203685477580uLL) { + if (u64 != 922337203685477580uLL || dig > 8) { + dbl = (mod_json_float_t)u64 * 10 + dig; + num_type = number_float; + break; + } + } + u64 = u64 * 10 + dig; + } + } + + /* Force double for big integer */ + if (num_type == number_float) { + while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) { + if (dbl >= 1E307) { + /* Number too big to store in double */ + return NULL; + } + dbl = dbl * 10 + dig; + } + } + + /* Parse frac = decimal-point 1*DIGIT */ + if (*cstr == '.') { + if (num_type != number_float) { + dbl = (mod_json_float_t)u64; + num_type = number_float; + } + + if ((dig = (uint32_t)(*(++cstr) - '0')) > 9) { + /* At least one digit in fraction part */ + return NULL; + } + + dbl = dbl * 10 + dig; + --exp_frac; + + while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) { + if (exp_frac > -16) { + dbl = dbl * 10 + dig; + --exp_frac; + } + } + } + + /* Parse exp = e [ minus / plus ] 1*DIGIT */ + if (*cstr == 'e' || *cstr == 'E') { + int32_t exp_minus = 0; + + if (num_type != number_float) { + dbl = (mod_json_float_t)u64; + num_type = number_float; + } + + exp_minus = *(++cstr); + if (exp_minus == '-' || exp_minus == '+') { + ++cstr; + } + + /* The first number char after 'e/E' */ + if ((dig = (uint32_t)(*cstr - '0')) > 9) { + return NULL; + } + exp = (int32_t)dig; + + while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) { + exp = exp * 10 + (int32_t)dig; + if (exp > 308) { + /* Number too big to store in double */ + return NULL; + } + } + + if (exp_minus == '-') { + exp = -exp; + } + } + + /* Finish parsing, call event according to the type of number. */ + if (num_type == number_float) { + dbl *= mod_json_utils_pow10(exp + exp_frac); + if (minus == '-') { + dbl = -dbl; + } + if (mod_json_token_invoke_float(tok, dbl) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + } else { + if (minus == '-') { + u64 = (uint64_t)(-(int64_t)u64); + } + if (mod_json_token_invoke_integer(tok, (mod_json_integer_t)u64) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + } + return cstr; +} + +static inline mod_json_cchar_t *mod_json_token_array_start( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + if (tok->array_depth < tok->array_max_depth) { + /* callback */ + if (mod_json_token_invoke_array(tok) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + + /* increase depth */ + ++tok->array_depth; + + /* push current tag */ + mod_json_token_set_tag(tok, '['); + + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case '[': + tok->state = mod_json_state_array_start; + return (cstr + 1); + + case ']': + tok->state = mod_json_state_array_finish; + return (cstr + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + break; + + default: + tok->state = mod_json_state_array_half; + return (cstr); + } + } else { + tok->error = mod_json_error_depth; + tok->context = cstr; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_array_half( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case ',': + tok->state = mod_json_state_array_half; + return (cstr + 1); + + case '[': + tok->state = mod_json_state_array_start; + return (cstr + 1); + + case ']': + tok->state = mod_json_state_array_finish; + return (cstr + 1); + + case '{': + tok->state = mod_json_state_object_start; + return (cstr + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + return NULL; + + /* value in array */ + case 't': + case 'T': + cstr = mod_json_token_value_true(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case 'f': + case 'F': + cstr = mod_json_token_value_false(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case 'n': + case 'N': + cstr = mod_json_token_value_null(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case '0': + case '1': + case '2': + case '3': + case '4': + case '5': + case '6': + case '7': + case '8': + case '9': + case '+': + case '-': + cstr = mod_json_token_value_number(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case '\"': + cstr = mod_json_token_value_string(tok, cstr, '\"'); + if (!cstr) { + return NULL; + } + break; + + case '\'': + if (tok->options & MOD_JSON_SQUOTE) { + cstr = mod_json_token_value_string(tok, cstr, '\''); + if (!cstr) { + return NULL; + } + break; + } + /* FALLTHRU */ + + default: + tok->error = mod_json_error_value; + tok->context = cstr; + return NULL; + } + + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case ',': + tok->state = mod_json_state_array_half; + return (cstr + 1); + + case ']': + tok->state = mod_json_state_array_finish; + return (cstr + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + break; + + default: + tok->error = mod_json_error_value; + tok->context = cstr; + break; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_array_finish( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + if (tok->array_depth) { + /* decrease depth */ + --tok->array_depth; + + /* callback */ + if (mod_json_token_invoke_array(tok) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case ']': + tok->state = mod_json_state_array_finish; + return (cstr + 1); + + case '}': + tok->state = mod_json_state_object_finish; + return (cstr + 1); + + case '\0': + if (tok->object_depth || tok->array_depth) { + tok->error = mod_json_error_trunc; + tok->context = cstr; + } else { + tok->state = mod_json_state_finish; + } + break; + + case ',': + if (tok->object_depth || tok->array_depth) { + mod_json_char_t tag = mod_json_token_tag(tok); + + if (tag == '{') { + tok->state = mod_json_state_object_half1; + return (cstr + 1); + } else if (tag == '[') { + tok->state = mod_json_state_array_half; + return (cstr + 1); + } + } + /* FALLTHRU */ + + default: + tok->error = mod_json_error_array; + tok->context = cstr; + } + } else { + tok->error = mod_json_error_depth; + tok->context = cstr; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_object_start( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + if (tok->object_depth < tok->object_max_depth) { + /* callback */ + if (mod_json_token_invoke_object(tok) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + + /* increase depth */ + ++tok->object_depth; + + /* push current tag */ + mod_json_token_set_tag(tok, '{'); + + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case '}': + tok->state = mod_json_state_object_finish; + return (cstr + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + break; + + default: + tok->state = mod_json_state_object_half1; + return (cstr); + } + } else { + tok->error = mod_json_error_depth; + tok->context = cstr; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_object_quotekey( + mod_json_token_t *tok, mod_json_cchar_t *cstr, mod_json_char_t quo) { + mod_json_cchar_t *cstr2 = mod_json_token_strfquo(tok, ++cstr, quo); + if (cstr2) { + /* callback */ + if (mod_json_token_invoke_field(tok, cstr, + (mod_json_size_t)(cstr2 - cstr)) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + + cstr2 = mod_json_token_strskp(tok, ++cstr2); + switch (*cstr2) { + case ':': + tok->state = mod_json_state_object_half2; + return (cstr2 + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + break; + + default: + tok->error = mod_json_error_key; + tok->context = cstr2; + break; + } + } else { + tok->error = mod_json_error_quote; + tok->context = cstr; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_object_simplekey( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + mod_json_cchar_t *cstr2 = mod_json_token_strfsep(tok, cstr); + if (cstr2 != cstr) { + /* callback */ + if (mod_json_token_invoke_field(tok, cstr, + (mod_json_size_t)(cstr2 - cstr)) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + + cstr2 = mod_json_token_strskp(tok, cstr2); + switch (*cstr2) { + case ':': + tok->state = mod_json_state_object_half2; + return (cstr2 + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + break; + + default: + tok->error = mod_json_error_key; + tok->context = cstr2; + break; + } + } else { + tok->error = mod_json_error_key; + tok->context = cstr; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_object_half1( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case ',': + tok->state = mod_json_state_object_half1; + return (cstr + 1); + + case '}': + tok->state = mod_json_state_object_finish; + return (cstr + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + break; + + case '\"': + /* The key with double quotes */ + return mod_json_token_object_quotekey(tok, cstr, '\"'); + + case '\'': + if (tok->options & MOD_JSON_SQUOTE) { + /* The key with single quotes */ + return mod_json_token_object_quotekey(tok, cstr, '\''); + } + /* FALLTHRU */ + + default: + /* support simple format? */ + if (tok->options & MOD_JSON_SIMPLE) { + return mod_json_token_object_simplekey(tok, cstr); + } else { + tok->error = mod_json_error_quote; + tok->context = cstr; + } + break; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_object_half2( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case '{': + tok->state = mod_json_state_object_start; + return (cstr + 1); + + case '[': + tok->state = mod_json_state_array_start; + return (cstr + 1); + + case ',': + tok->state = mod_json_state_object_half1; + return (cstr + 1); + + case '}': + tok->state = mod_json_state_object_finish; + return (cstr + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + return NULL; + + case 't': + case 'T': + cstr = mod_json_token_value_true(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case 'f': + case 'F': + cstr = mod_json_token_value_false(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case 'i': + case 'I': + cstr = mod_json_token_value_infinity(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case 'n': + case 'N': + cstr = mod_json_token_value_null(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case '0': + case '1': + case '2': + case '3': + case '4': + case '5': + case '6': + case '7': + case '8': + case '9': + case '+': + case '-': + cstr = mod_json_token_value_number(tok, cstr); + if (!cstr) { + return NULL; + } + break; + + case '\"': + cstr = mod_json_token_value_string(tok, cstr, '\"'); + if (!cstr) { + return NULL; + } + break; + + case '\'': + if (tok->options & MOD_JSON_SQUOTE) { + cstr = mod_json_token_value_string(tok, cstr, '\''); + if (!cstr) { + return NULL; + } + break; + } + /* FALLTHRU */ + + default: + tok->error = mod_json_error_value; + tok->context = cstr; + return NULL; + } + + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case ',': + tok->state = mod_json_state_object_half1; + return (cstr + 1); + + case '}': + tok->state = mod_json_state_object_finish; + return (cstr + 1); + + case '\0': + tok->error = mod_json_error_trunc; + tok->context = cstr; + break; + + default: + tok->error = mod_json_error_value; + tok->context = cstr; + break; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_object_finish( + mod_json_token_t *tok, mod_json_cchar_t *cstr) { + if (tok->object_depth) { + /* decrease depth */ + --tok->object_depth; + + /* callback */ + if (mod_json_token_invoke_object(tok) != 0) { + tok->error = mod_json_error_break; + tok->context = cstr; + return NULL; + } + + cstr = mod_json_token_strskp(tok, cstr); + switch (*cstr) { + case '}': + tok->state = mod_json_state_object_finish; + return (cstr + 1); + + case ']': + tok->state = mod_json_state_array_finish; + return (cstr + 1); + + case '\0': + if (tok->object_depth || tok->array_depth) { + tok->error = mod_json_error_trunc; + tok->context = cstr; + } else { + tok->state = mod_json_state_finish; + } + break; + + case ',': + if (tok->object_depth || tok->array_depth) { + mod_json_char_t tag = mod_json_token_tag(tok); + + if (tag == '{') { + tok->state = mod_json_state_object_half1; + return (cstr + 1); + } else if (tag == '[') { + tok->state = mod_json_state_array_half; + return (cstr + 1); + } + } + /* FALLTHRU */ + + default: + tok->error = mod_json_error_object; + tok->context = cstr; + } + } else { + tok->error = mod_json_error_depth; + tok->context = cstr; + } + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_null(mod_json_token_t *tok, + mod_json_cchar_t *cstr) { + if (!cstr || *cstr == '\0') { + tok->error = mod_json_error_invalid; + tok->context = cstr; + return NULL; + } + + tok->state = mod_json_state_start; + return cstr; +} + +static inline mod_json_cchar_t *mod_json_token_finish(mod_json_token_t *tok, + mod_json_cchar_t *cstr) { + tok->error = mod_json_error_null; + (void)cstr; + return NULL; +} + +static inline mod_json_cchar_t *mod_json_token_default(mod_json_token_t *tok, + mod_json_cchar_t *cstr) { + tok->error = mod_json_error_state; + tok->context = cstr; + return NULL; +} + +int mod_json_token_parse(mod_json_token_t *tok, mod_json_cchar_t *cstr) { + while (cstr) { + switch (tok->state) { + case mod_json_state_start: + cstr = mod_json_token_start(tok, cstr); + break; + + case mod_json_state_array_start: + cstr = mod_json_token_array_start(tok, cstr); + break; + + case mod_json_state_array_half: + cstr = mod_json_token_array_half(tok, cstr); + break; + + case mod_json_state_array_finish: + cstr = mod_json_token_array_finish(tok, cstr); + break; + + case mod_json_state_object_start: + cstr = mod_json_token_object_start(tok, cstr); + break; + + case mod_json_state_object_half1: + cstr = mod_json_token_object_half1(tok, cstr); + break; + + case mod_json_state_object_half2: + cstr = mod_json_token_object_half2(tok, cstr); + break; + + case mod_json_state_object_finish: + cstr = mod_json_token_object_finish(tok, cstr); + break; + + case mod_json_state_null: + cstr = mod_json_token_null(tok, cstr); + break; + + case mod_json_state_finish: + cstr = mod_json_token_finish(tok, cstr); + break; + + default: + cstr = mod_json_token_default(tok, cstr); + break; + } + } + return (tok->error == mod_json_error_null ? 0 : -1); +} + +static inline int mod_json_parser_insert(mod_json_parser_t *par, + mod_json_size_t depth, + mod_json_value_t *val) { + if (depth > 0) { + mod_json_value_t *cur = par->vals[depth - 1]; + + switch (cur->type) { + case mod_json_type_object: + return (mod_json_object_insert(cur->data.c_obj, par->key, val) ? 0 + : -1); + + case mod_json_type_array: + return mod_json_array_push(cur->data.c_arr, val); + + default: + break; + } + } + return -1; +} + +static inline int mod_json_parser_insert_object(mod_json_parser_t *par, + mod_json_size_t depth) { + mod_json_object_t *obj; + mod_json_value_t *jval; + + obj = mod_json_object_set_default(); + mod_json_minus_if_false(obj); + + jval = mod_json_value_set_object(obj); + mod_json_object_unset(obj); + mod_json_minus_if_false(jval); + + if (depth > 0) { + int ret = mod_json_parser_insert(par, depth, jval); + if (ret == 0) { + par->vals[depth] = jval; + } + mod_json_value_unset(jval); + return ret; + } else { + /* It's the root, save the pointer. Don't unset it. */ + par->vals[0] = jval; + } + return 0; +} + +static inline int mod_json_parser_insert_array(mod_json_parser_t *par, + mod_json_size_t depth) { + mod_json_array_t *arr; + mod_json_value_t *jval; + + arr = mod_json_array_set_default(); + mod_json_minus_if_false(arr); + + jval = mod_json_value_set_array(arr); + mod_json_array_unset(arr); + mod_json_minus_if_false(jval); + + if (depth > 0) { + int ret = mod_json_parser_insert(par, depth, jval); + if (ret == 0) { + par->vals[depth] = jval; + } + mod_json_value_unset(jval); + return ret; + } else { + /* It's the root, save the pointer. Don't unset it. */ + par->vals[0] = jval; + } + return 0; +} + +static inline void mod_json_token_set_parser(mod_json_token_t *tok, + mod_json_parser_t *par) { + mod_json_token_set_param(tok, par); +} + +static inline mod_json_parser_t *mod_json_token_parser(mod_json_token_t *tok) { + return (mod_json_parser_t *)mod_json_token_param(tok); +} + +static inline int mod_json_parser_event_field(mod_json_token_t *tok, + mod_json_cchar_t *val, + mod_json_size_t len) { + mod_json_parser_t *parser; + + /* get information */ + parser = mod_json_token_parser(tok); + + /* unset previous one */ + mod_json_string_unset(parser->key); + + parser->key = mod_json_string_set(val, len); + return (parser->key ? 0 : -1); +} + +static inline int mod_json_parser_event_array(mod_json_token_t *tok) { + switch (mod_json_token_state(tok)) { + case mod_json_state_array_finish: + /* continue */ + return 0; + + case mod_json_state_array_start: + return mod_json_parser_insert_array(mod_json_token_parser(tok), + mod_json_token_depth(tok)); + + default: + break; + } + return -1; +} + +static inline int mod_json_parser_event_object(mod_json_token_t *tok) { + switch (mod_json_token_state(tok)) { + case mod_json_state_object_finish: + /* continue */ + return 0; + + case mod_json_state_object_start: + return mod_json_parser_insert_object(mod_json_token_parser(tok), + mod_json_token_depth(tok)); + + default: + break; + } + return -1; +} + +static inline int mod_json_parser_event_null(mod_json_token_t *tok) { + mod_json_parser_t *parser; + + /* get information */ + parser = mod_json_token_parser(tok); + + if (!parser->val_null) { + parser->val_null = mod_json_value_set_null(); + mod_json_minus_if_false(parser->val_null); + } + return mod_json_parser_insert(parser, mod_json_token_depth(tok), + parser->val_null); +} + +static inline int mod_json_parser_event_true(mod_json_token_t *tok) { + mod_json_parser_t *parser; + + /* get information */ + parser = mod_json_token_parser(tok); + + if (!parser->val_true) { + parser->val_true = mod_json_value_set_boolean(MOD_JSON_TRUE); + mod_json_minus_if_false(parser->val_true); + } + return mod_json_parser_insert(parser, mod_json_token_depth(tok), + parser->val_true); +} + +static inline int mod_json_parser_event_false(mod_json_token_t *tok) { + mod_json_parser_t *parser; + + /* get information */ + parser = mod_json_token_parser(tok); + + if (!parser->val_false) { + parser->val_false = mod_json_value_set_boolean(MOD_JSON_FALSE); + mod_json_minus_if_false(parser->val_false); + } + return mod_json_parser_insert(parser, mod_json_token_depth(tok), + parser->val_false); +} + +static inline int mod_json_parser_event_boolean(mod_json_token_t *tok, + mod_json_boolean_t val) { + if (!val) { + return mod_json_parser_event_false(tok); + } + return mod_json_parser_event_true(tok); +} + +static inline int mod_json_parser_event_zero(mod_json_token_t *tok) { + mod_json_parser_t *parser; + + /* get information */ + parser = mod_json_token_parser(tok); + + if (!parser->val_zero) { + parser->val_zero = mod_json_value_set_integer(0); + mod_json_minus_if_false(parser->val_zero); + } + return mod_json_parser_insert(parser, mod_json_token_depth(tok), + parser->val_zero); +} + +static inline int mod_json_parser_event_integer(mod_json_token_t *tok, + mod_json_integer_t val) { + int ret = -1; + + if (val != 0) { + mod_json_value_t *jval; + + jval = mod_json_value_set_integer(val); + if (jval) { + ret = mod_json_parser_insert(mod_json_token_parser(tok), + mod_json_token_depth(tok), jval); + mod_json_value_unset(jval); + } + } else { + /* zero event */ + ret = mod_json_parser_event_zero(tok); + } + return ret; +} + +static inline int mod_json_parser_event_zerof(mod_json_token_t *tok) { + mod_json_parser_t *parser; + + /* get information */ + parser = mod_json_token_parser(tok); + + if (!parser->val_zerof) { + parser->val_zerof = mod_json_value_set_float(0.0); + mod_json_minus_if_false(parser->val_zerof); + } + return mod_json_parser_insert(parser, mod_json_token_depth(tok), + parser->val_zerof); +} + +static inline int mod_json_parser_event_float(mod_json_token_t *tok, + mod_json_float_t val) { + int ret = -1; + + if (val != 0.0) { + mod_json_value_t *jval; + + jval = mod_json_value_set_float(val); + if (jval) { + ret = mod_json_parser_insert(mod_json_token_parser(tok), + mod_json_token_depth(tok), jval); + mod_json_value_unset(jval); + } + } else { + /* zero event */ + ret = mod_json_parser_event_zerof(tok); + } + return ret; +} + +static inline int mod_json_parser_event_empty(mod_json_token_t *tok) { + mod_json_parser_t *parser; + + /* get information */ + parser = mod_json_token_parser(tok); + + if (!parser->val_empty) { + mod_json_string_t *str; + + str = mod_json_string_set("", 0); + mod_json_minus_if_false(str); + + parser->val_empty = mod_json_value_set_string(str); + mod_json_string_unset(str); + mod_json_minus_if_false(parser->val_empty); + } + return mod_json_parser_insert(parser, mod_json_token_depth(tok), + parser->val_empty); +} + +static inline int mod_json_parser_event_string(mod_json_token_t *tok, + mod_json_cchar_t *val, + mod_json_size_t len) { + int ret = -1; + + if (len > 0) { + mod_json_string_t *str; + mod_json_value_t *jval; + + str = mod_json_string_set(val, len); + if (str) { + jval = mod_json_value_set_string(str); + } else { + jval = NULL; + } + mod_json_string_unset(str); + + if (jval) { + ret = mod_json_parser_insert(mod_json_token_parser(tok), + mod_json_token_depth(tok), jval); + mod_json_value_unset(jval); + } + } else { + /* empty event */ + ret = mod_json_parser_event_empty(tok); + } + return ret; +} + +static int mod_json_parser_event(mod_json_token_t *tok, mod_json_void_t *val, + mod_json_size_t len) { + switch (tok->event_code) { + case mod_json_event_field: + return mod_json_parser_event_field(tok, (mod_json_cchar_t *)val, len); + + case mod_json_event_object: + return mod_json_parser_event_object(tok); + + case mod_json_event_array: + return mod_json_parser_event_array(tok); + + case mod_json_event_null: + return mod_json_parser_event_null(tok); + + case mod_json_event_boolean: + return mod_json_parser_event_boolean(tok, *(mod_json_boolean_t *)val); + + case mod_json_event_integer: + return mod_json_parser_event_integer(tok, *(mod_json_integer_t *)val); + + case mod_json_event_float: + return mod_json_parser_event_float(tok, *(mod_json_float_t *)val); + + case mod_json_event_string: + return mod_json_parser_event_string(tok, (mod_json_cchar_t *)val, len); + + default: + break; + } + return -1; +} + +static inline mod_json_parser_t *mod_json_parser_create(mod_json_size_t depth) { + mod_json_parser_t *parser; + mod_json_null_if_false(depth > 0); + + parser = (mod_json_parser_t *)mod_json_malloc( + depth * sizeof(mod_json_value_t *) + sizeof(mod_json_parser_t)); + mod_json_null_if_false(parser); + + memset(parser, 0, sizeof(mod_json_parser_t)); + parser->vals[0] = NULL; + return parser; +} + +static inline void mod_json_parser_destroy(mod_json_parser_t *par) { + mod_json_value_unset(par->val_null); + mod_json_value_unset(par->val_true); + mod_json_value_unset(par->val_false); + mod_json_value_unset(par->val_zero); + mod_json_value_unset(par->val_zerof); + mod_json_value_unset(par->val_empty); + mod_json_string_unset(par->key); + mod_json_free(par); +} + +mod_json_value_t *mod_json_parse(mod_json_token_t *tok, + mod_json_cchar_t *cstr) { + mod_json_parser_t *parser; + mod_json_value_t *root; + mod_json_null_if_false(tok && cstr && *cstr); + + parser = mod_json_parser_create(mod_json_token_max_depth(tok)); + mod_json_null_if_false(parser); + + mod_json_token_set_parser(tok, parser); + mod_json_token_set_event(tok, mod_json_parser_event); + + if (mod_json_token_parse(tok, cstr) == 0) { + root = parser->vals[0]; + } else { + /* error occur */ + root = NULL; + mod_json_value_unset(parser->vals[0]); + } + + /* clean up */ + mod_json_parser_destroy(parser); + + /* success? */ + return root; +} + +mod_json_value_t *mod_json_parse_simply(mod_json_cchar_t *cstr, + mod_json_size_t opts) { + mod_json_value_t *val; + mod_json_token_t *tok; + mod_json_option_t opt; + + opt.options = opts; + opt.object_depth = 0; /* Use default object depth */ + opt.array_depth = 0; /* Use default array depth */ + + tok = mod_json_token_create(&opt); + mod_json_null_if_false(tok); + + val = mod_json_parse(tok, cstr); + mod_json_token_destroy(tok); + + /* value of root */ + return val; +} + +static inline int mod_json_dump_null(mod_json_string_t *str) { + return mod_json_string_add_cstr(str, "null", 4); +} + +static inline int mod_json_dump_boolean(mod_json_string_t *str, + mod_json_boolean_t bol) { + if (!bol) { + return mod_json_string_add_cstr(str, "false", 5); + } + return mod_json_string_add_cstr(str, "true", 4); +} + +static inline int mod_json_dump_integer(mod_json_string_t *str, + mod_json_integer_t num) { + mod_json_char_t buf[32]; + + return mod_json_string_add_cstr(str, buf, mod_json_utils_itostr(buf, num)); +} + +static inline int mod_json_dump_float(mod_json_string_t *str, + mod_json_float_t dbl) { + mod_json_char_t buf[32]; + + return mod_json_string_add_cstr( + str, buf, + (mod_json_size_t)mod_json_utils_snprintf(buf, sizeof(buf), "%g", dbl)); +} + +static inline int mod_json_dump_string(mod_json_string_t *str, + mod_json_string_t *val) { + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '\"')); + + if (val) { + mod_json_minus_if_ne_zero(mod_json_string_add_jstr(str, val)); + } + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '\"')); + return 0; +} + +static inline int mod_json_dump_value(mod_json_string_t *str, + mod_json_value_t *val); + +static inline int mod_json_dump_array(mod_json_string_t *str, + mod_json_array_t *arr) { + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '[')); + + if (arr) { + mod_json_value_t **iter = arr->first; + + for (; iter != arr->last; ++iter) { + mod_json_minus_if_ne_zero(mod_json_dump_value(str, *iter)); + if (iter + 1 != arr->last) { + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, ',')); + } + } + } + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, ']')); + return 0; +} + +static inline int mod_json_dump_key(mod_json_string_t *str, + mod_json_string_t *key) { + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '\"')); + mod_json_minus_if_ne_zero(mod_json_string_add_jstr(str, key)); + mod_json_minus_if_ne_zero(mod_json_string_add_cstr(str, "\":", 2)); + return 0; +} + +static inline int mod_json_dump_object(mod_json_string_t *str, + mod_json_object_t *obj) { + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '{')); + + if (obj) { + mod_json_pair_t *iter = obj->first; + + for (; iter != obj->last; ++iter) { + mod_json_minus_if_ne_zero(mod_json_dump_key(str, iter->key)); + mod_json_minus_if_ne_zero(mod_json_dump_value(str, iter->val)); + + if (iter + 1 != obj->last) { + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, ',')); + } + } + } + mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '}')); + return 0; +} + +static inline int mod_json_dump_value(mod_json_string_t *str, + mod_json_value_t *val) { + if (val) { + switch (val->type) { + case mod_json_type_null: + return mod_json_dump_null(str); + + case mod_json_type_boolean: + return mod_json_dump_boolean(str, val->data.c_bool); + + case mod_json_type_integer: + return mod_json_dump_integer(str, val->data.c_int); + + case mod_json_type_float: + return mod_json_dump_float(str, val->data.c_float); + + case mod_json_type_string: + return mod_json_dump_string(str, val->data.c_str); + + case mod_json_type_array: + return mod_json_dump_array(str, val->data.c_arr); + + case mod_json_type_object: + return mod_json_dump_object(str, val->data.c_obj); + + default: + return -1; + } + } + return mod_json_dump_null(str); +} + +mod_json_string_t *mod_json_dump(mod_json_value_t *val) { + mod_json_string_t *str = mod_json_string_set("", 0); + mod_json_null_if_false(str); + + if (mod_json_unlikely(mod_json_dump_value(str, val) != 0)) { + /* error occur */ + mod_json_string_unset(str); + return NULL; + } + return str; +} diff --git a/src/ailego/hash/crc32c.cc b/src/ailego/hash/crc32c.cc new file mode 100644 index 0000000..bb5ab58 --- /dev/null +++ b/src/ailego/hash/crc32c.cc @@ -0,0 +1,637 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +#if !defined(__SSE4_2__) && !defined(__ARM_FEATURE_CRC32) +/** + * The following CRC lookup table was generated automagically + * using the following model parameters: + * + * Generator Polynomial = ................. 0x1EDC6F41 + * Generator Polynomial Length = .......... 32 bits + * Reflected Bits = ....................... TRUE + * Table Generation Offset = .............. 32 bits + * Number of Slices = ..................... 8 slices + * Slice Lengths = ........................ 8 8 8 8 8 8 8 8 + */ +static uint32_t crc_tableil8_o32[256] = { + 0x00000000, 0xF26B8303, 0xE13B70F7, 0x1350F3F4, 0xC79A971F, 0x35F1141C, + 0x26A1E7E8, 0xD4CA64EB, 0x8AD958CF, 0x78B2DBCC, 0x6BE22838, 0x9989AB3B, + 0x4D43CFD0, 0xBF284CD3, 0xAC78BF27, 0x5E133C24, 0x105EC76F, 0xE235446C, + 0xF165B798, 0x030E349B, 0xD7C45070, 0x25AFD373, 0x36FF2087, 0xC494A384, + 0x9A879FA0, 0x68EC1CA3, 0x7BBCEF57, 0x89D76C54, 0x5D1D08BF, 0xAF768BBC, + 0xBC267848, 0x4E4DFB4B, 0x20BD8EDE, 0xD2D60DDD, 0xC186FE29, 0x33ED7D2A, + 0xE72719C1, 0x154C9AC2, 0x061C6936, 0xF477EA35, 0xAA64D611, 0x580F5512, + 0x4B5FA6E6, 0xB93425E5, 0x6DFE410E, 0x9F95C20D, 0x8CC531F9, 0x7EAEB2FA, + 0x30E349B1, 0xC288CAB2, 0xD1D83946, 0x23B3BA45, 0xF779DEAE, 0x05125DAD, + 0x1642AE59, 0xE4292D5A, 0xBA3A117E, 0x4851927D, 0x5B016189, 0xA96AE28A, + 0x7DA08661, 0x8FCB0562, 0x9C9BF696, 0x6EF07595, 0x417B1DBC, 0xB3109EBF, + 0xA0406D4B, 0x522BEE48, 0x86E18AA3, 0x748A09A0, 0x67DAFA54, 0x95B17957, + 0xCBA24573, 0x39C9C670, 0x2A993584, 0xD8F2B687, 0x0C38D26C, 0xFE53516F, + 0xED03A29B, 0x1F682198, 0x5125DAD3, 0xA34E59D0, 0xB01EAA24, 0x42752927, + 0x96BF4DCC, 0x64D4CECF, 0x77843D3B, 0x85EFBE38, 0xDBFC821C, 0x2997011F, + 0x3AC7F2EB, 0xC8AC71E8, 0x1C661503, 0xEE0D9600, 0xFD5D65F4, 0x0F36E6F7, + 0x61C69362, 0x93AD1061, 0x80FDE395, 0x72966096, 0xA65C047D, 0x5437877E, + 0x4767748A, 0xB50CF789, 0xEB1FCBAD, 0x197448AE, 0x0A24BB5A, 0xF84F3859, + 0x2C855CB2, 0xDEEEDFB1, 0xCDBE2C45, 0x3FD5AF46, 0x7198540D, 0x83F3D70E, + 0x90A324FA, 0x62C8A7F9, 0xB602C312, 0x44694011, 0x5739B3E5, 0xA55230E6, + 0xFB410CC2, 0x092A8FC1, 0x1A7A7C35, 0xE811FF36, 0x3CDB9BDD, 0xCEB018DE, + 0xDDE0EB2A, 0x2F8B6829, 0x82F63B78, 0x709DB87B, 0x63CD4B8F, 0x91A6C88C, + 0x456CAC67, 0xB7072F64, 0xA457DC90, 0x563C5F93, 0x082F63B7, 0xFA44E0B4, + 0xE9141340, 0x1B7F9043, 0xCFB5F4A8, 0x3DDE77AB, 0x2E8E845F, 0xDCE5075C, + 0x92A8FC17, 0x60C37F14, 0x73938CE0, 0x81F80FE3, 0x55326B08, 0xA759E80B, + 0xB4091BFF, 0x466298FC, 0x1871A4D8, 0xEA1A27DB, 0xF94AD42F, 0x0B21572C, + 0xDFEB33C7, 0x2D80B0C4, 0x3ED04330, 0xCCBBC033, 0xA24BB5A6, 0x502036A5, + 0x4370C551, 0xB11B4652, 0x65D122B9, 0x97BAA1BA, 0x84EA524E, 0x7681D14D, + 0x2892ED69, 0xDAF96E6A, 0xC9A99D9E, 0x3BC21E9D, 0xEF087A76, 0x1D63F975, + 0x0E330A81, 0xFC588982, 0xB21572C9, 0x407EF1CA, 0x532E023E, 0xA145813D, + 0x758FE5D6, 0x87E466D5, 0x94B49521, 0x66DF1622, 0x38CC2A06, 0xCAA7A905, + 0xD9F75AF1, 0x2B9CD9F2, 0xFF56BD19, 0x0D3D3E1A, 0x1E6DCDEE, 0xEC064EED, + 0xC38D26C4, 0x31E6A5C7, 0x22B65633, 0xD0DDD530, 0x0417B1DB, 0xF67C32D8, + 0xE52CC12C, 0x1747422F, 0x49547E0B, 0xBB3FFD08, 0xA86F0EFC, 0x5A048DFF, + 0x8ECEE914, 0x7CA56A17, 0x6FF599E3, 0x9D9E1AE0, 0xD3D3E1AB, 0x21B862A8, + 0x32E8915C, 0xC083125F, 0x144976B4, 0xE622F5B7, 0xF5720643, 0x07198540, + 0x590AB964, 0xAB613A67, 0xB831C993, 0x4A5A4A90, 0x9E902E7B, 0x6CFBAD78, + 0x7FAB5E8C, 0x8DC0DD8F, 0xE330A81A, 0x115B2B19, 0x020BD8ED, 0xF0605BEE, + 0x24AA3F05, 0xD6C1BC06, 0xC5914FF2, 0x37FACCF1, 0x69E9F0D5, 0x9B8273D6, + 0x88D28022, 0x7AB90321, 0xAE7367CA, 0x5C18E4C9, 0x4F48173D, 0xBD23943E, + 0xF36E6F75, 0x0105EC76, 0x12551F82, 0xE03E9C81, 0x34F4F86A, 0xC69F7B69, + 0xD5CF889D, 0x27A40B9E, 0x79B737BA, 0x8BDCB4B9, 0x988C474D, 0x6AE7C44E, + 0xBE2DA0A5, 0x4C4623A6, 0x5F16D052, 0xAD7D5351}; + +/** + * The following CRC lookup table was generated automagically + * using the following model parameters: + * + * Generator Polynomial = ................. 0x1EDC6F41 + * Generator Polynomial Length = .......... 32 bits + * Reflected Bits = ....................... TRUE + * Table Generation Offset = .............. 32 bits + * Number of Slices = ..................... 8 slices + * Slice Lengths = ........................ 8 8 8 8 8 8 8 8 + */ +static uint32_t crc_tableil8_o40[256] = { + 0x00000000, 0x13A29877, 0x274530EE, 0x34E7A899, 0x4E8A61DC, 0x5D28F9AB, + 0x69CF5132, 0x7A6DC945, 0x9D14C3B8, 0x8EB65BCF, 0xBA51F356, 0xA9F36B21, + 0xD39EA264, 0xC03C3A13, 0xF4DB928A, 0xE7790AFD, 0x3FC5F181, 0x2C6769F6, + 0x1880C16F, 0x0B225918, 0x714F905D, 0x62ED082A, 0x560AA0B3, 0x45A838C4, + 0xA2D13239, 0xB173AA4E, 0x859402D7, 0x96369AA0, 0xEC5B53E5, 0xFFF9CB92, + 0xCB1E630B, 0xD8BCFB7C, 0x7F8BE302, 0x6C297B75, 0x58CED3EC, 0x4B6C4B9B, + 0x310182DE, 0x22A31AA9, 0x1644B230, 0x05E62A47, 0xE29F20BA, 0xF13DB8CD, + 0xC5DA1054, 0xD6788823, 0xAC154166, 0xBFB7D911, 0x8B507188, 0x98F2E9FF, + 0x404E1283, 0x53EC8AF4, 0x670B226D, 0x74A9BA1A, 0x0EC4735F, 0x1D66EB28, + 0x298143B1, 0x3A23DBC6, 0xDD5AD13B, 0xCEF8494C, 0xFA1FE1D5, 0xE9BD79A2, + 0x93D0B0E7, 0x80722890, 0xB4958009, 0xA737187E, 0xFF17C604, 0xECB55E73, + 0xD852F6EA, 0xCBF06E9D, 0xB19DA7D8, 0xA23F3FAF, 0x96D89736, 0x857A0F41, + 0x620305BC, 0x71A19DCB, 0x45463552, 0x56E4AD25, 0x2C896460, 0x3F2BFC17, + 0x0BCC548E, 0x186ECCF9, 0xC0D23785, 0xD370AFF2, 0xE797076B, 0xF4359F1C, + 0x8E585659, 0x9DFACE2E, 0xA91D66B7, 0xBABFFEC0, 0x5DC6F43D, 0x4E646C4A, + 0x7A83C4D3, 0x69215CA4, 0x134C95E1, 0x00EE0D96, 0x3409A50F, 0x27AB3D78, + 0x809C2506, 0x933EBD71, 0xA7D915E8, 0xB47B8D9F, 0xCE1644DA, 0xDDB4DCAD, + 0xE9537434, 0xFAF1EC43, 0x1D88E6BE, 0x0E2A7EC9, 0x3ACDD650, 0x296F4E27, + 0x53028762, 0x40A01F15, 0x7447B78C, 0x67E52FFB, 0xBF59D487, 0xACFB4CF0, + 0x981CE469, 0x8BBE7C1E, 0xF1D3B55B, 0xE2712D2C, 0xD69685B5, 0xC5341DC2, + 0x224D173F, 0x31EF8F48, 0x050827D1, 0x16AABFA6, 0x6CC776E3, 0x7F65EE94, + 0x4B82460D, 0x5820DE7A, 0xFBC3FAF9, 0xE861628E, 0xDC86CA17, 0xCF245260, + 0xB5499B25, 0xA6EB0352, 0x920CABCB, 0x81AE33BC, 0x66D73941, 0x7575A136, + 0x419209AF, 0x523091D8, 0x285D589D, 0x3BFFC0EA, 0x0F186873, 0x1CBAF004, + 0xC4060B78, 0xD7A4930F, 0xE3433B96, 0xF0E1A3E1, 0x8A8C6AA4, 0x992EF2D3, + 0xADC95A4A, 0xBE6BC23D, 0x5912C8C0, 0x4AB050B7, 0x7E57F82E, 0x6DF56059, + 0x1798A91C, 0x043A316B, 0x30DD99F2, 0x237F0185, 0x844819FB, 0x97EA818C, + 0xA30D2915, 0xB0AFB162, 0xCAC27827, 0xD960E050, 0xED8748C9, 0xFE25D0BE, + 0x195CDA43, 0x0AFE4234, 0x3E19EAAD, 0x2DBB72DA, 0x57D6BB9F, 0x447423E8, + 0x70938B71, 0x63311306, 0xBB8DE87A, 0xA82F700D, 0x9CC8D894, 0x8F6A40E3, + 0xF50789A6, 0xE6A511D1, 0xD242B948, 0xC1E0213F, 0x26992BC2, 0x353BB3B5, + 0x01DC1B2C, 0x127E835B, 0x68134A1E, 0x7BB1D269, 0x4F567AF0, 0x5CF4E287, + 0x04D43CFD, 0x1776A48A, 0x23910C13, 0x30339464, 0x4A5E5D21, 0x59FCC556, + 0x6D1B6DCF, 0x7EB9F5B8, 0x99C0FF45, 0x8A626732, 0xBE85CFAB, 0xAD2757DC, + 0xD74A9E99, 0xC4E806EE, 0xF00FAE77, 0xE3AD3600, 0x3B11CD7C, 0x28B3550B, + 0x1C54FD92, 0x0FF665E5, 0x759BACA0, 0x663934D7, 0x52DE9C4E, 0x417C0439, + 0xA6050EC4, 0xB5A796B3, 0x81403E2A, 0x92E2A65D, 0xE88F6F18, 0xFB2DF76F, + 0xCFCA5FF6, 0xDC68C781, 0x7B5FDFFF, 0x68FD4788, 0x5C1AEF11, 0x4FB87766, + 0x35D5BE23, 0x26772654, 0x12908ECD, 0x013216BA, 0xE64B1C47, 0xF5E98430, + 0xC10E2CA9, 0xD2ACB4DE, 0xA8C17D9B, 0xBB63E5EC, 0x8F844D75, 0x9C26D502, + 0x449A2E7E, 0x5738B609, 0x63DF1E90, 0x707D86E7, 0x0A104FA2, 0x19B2D7D5, + 0x2D557F4C, 0x3EF7E73B, 0xD98EEDC6, 0xCA2C75B1, 0xFECBDD28, 0xED69455F, + 0x97048C1A, 0x84A6146D, 0xB041BCF4, 0xA3E32483}; + +/** + * The following CRC lookup table was generated automagically + * using the following model parameters: + * + * Generator Polynomial = ................. 0x1EDC6F41 + * Generator Polynomial Length = .......... 32 bits + * Reflected Bits = ....................... TRUE + * Table Generation Offset = .............. 32 bits + * Number of Slices = ..................... 8 slices + * Slice Lengths = ........................ 8 8 8 8 8 8 8 8 + */ +static uint32_t crc_tableil8_o48[256] = { + 0x00000000, 0xA541927E, 0x4F6F520D, 0xEA2EC073, 0x9EDEA41A, 0x3B9F3664, + 0xD1B1F617, 0x74F06469, 0x38513EC5, 0x9D10ACBB, 0x773E6CC8, 0xD27FFEB6, + 0xA68F9ADF, 0x03CE08A1, 0xE9E0C8D2, 0x4CA15AAC, 0x70A27D8A, 0xD5E3EFF4, + 0x3FCD2F87, 0x9A8CBDF9, 0xEE7CD990, 0x4B3D4BEE, 0xA1138B9D, 0x045219E3, + 0x48F3434F, 0xEDB2D131, 0x079C1142, 0xA2DD833C, 0xD62DE755, 0x736C752B, + 0x9942B558, 0x3C032726, 0xE144FB14, 0x4405696A, 0xAE2BA919, 0x0B6A3B67, + 0x7F9A5F0E, 0xDADBCD70, 0x30F50D03, 0x95B49F7D, 0xD915C5D1, 0x7C5457AF, + 0x967A97DC, 0x333B05A2, 0x47CB61CB, 0xE28AF3B5, 0x08A433C6, 0xADE5A1B8, + 0x91E6869E, 0x34A714E0, 0xDE89D493, 0x7BC846ED, 0x0F382284, 0xAA79B0FA, + 0x40577089, 0xE516E2F7, 0xA9B7B85B, 0x0CF62A25, 0xE6D8EA56, 0x43997828, + 0x37691C41, 0x92288E3F, 0x78064E4C, 0xDD47DC32, 0xC76580D9, 0x622412A7, + 0x880AD2D4, 0x2D4B40AA, 0x59BB24C3, 0xFCFAB6BD, 0x16D476CE, 0xB395E4B0, + 0xFF34BE1C, 0x5A752C62, 0xB05BEC11, 0x151A7E6F, 0x61EA1A06, 0xC4AB8878, + 0x2E85480B, 0x8BC4DA75, 0xB7C7FD53, 0x12866F2D, 0xF8A8AF5E, 0x5DE93D20, + 0x29195949, 0x8C58CB37, 0x66760B44, 0xC337993A, 0x8F96C396, 0x2AD751E8, + 0xC0F9919B, 0x65B803E5, 0x1148678C, 0xB409F5F2, 0x5E273581, 0xFB66A7FF, + 0x26217BCD, 0x8360E9B3, 0x694E29C0, 0xCC0FBBBE, 0xB8FFDFD7, 0x1DBE4DA9, + 0xF7908DDA, 0x52D11FA4, 0x1E704508, 0xBB31D776, 0x511F1705, 0xF45E857B, + 0x80AEE112, 0x25EF736C, 0xCFC1B31F, 0x6A802161, 0x56830647, 0xF3C29439, + 0x19EC544A, 0xBCADC634, 0xC85DA25D, 0x6D1C3023, 0x8732F050, 0x2273622E, + 0x6ED23882, 0xCB93AAFC, 0x21BD6A8F, 0x84FCF8F1, 0xF00C9C98, 0x554D0EE6, + 0xBF63CE95, 0x1A225CEB, 0x8B277743, 0x2E66E53D, 0xC448254E, 0x6109B730, + 0x15F9D359, 0xB0B84127, 0x5A968154, 0xFFD7132A, 0xB3764986, 0x1637DBF8, + 0xFC191B8B, 0x595889F5, 0x2DA8ED9C, 0x88E97FE2, 0x62C7BF91, 0xC7862DEF, + 0xFB850AC9, 0x5EC498B7, 0xB4EA58C4, 0x11ABCABA, 0x655BAED3, 0xC01A3CAD, + 0x2A34FCDE, 0x8F756EA0, 0xC3D4340C, 0x6695A672, 0x8CBB6601, 0x29FAF47F, + 0x5D0A9016, 0xF84B0268, 0x1265C21B, 0xB7245065, 0x6A638C57, 0xCF221E29, + 0x250CDE5A, 0x804D4C24, 0xF4BD284D, 0x51FCBA33, 0xBBD27A40, 0x1E93E83E, + 0x5232B292, 0xF77320EC, 0x1D5DE09F, 0xB81C72E1, 0xCCEC1688, 0x69AD84F6, + 0x83834485, 0x26C2D6FB, 0x1AC1F1DD, 0xBF8063A3, 0x55AEA3D0, 0xF0EF31AE, + 0x841F55C7, 0x215EC7B9, 0xCB7007CA, 0x6E3195B4, 0x2290CF18, 0x87D15D66, + 0x6DFF9D15, 0xC8BE0F6B, 0xBC4E6B02, 0x190FF97C, 0xF321390F, 0x5660AB71, + 0x4C42F79A, 0xE90365E4, 0x032DA597, 0xA66C37E9, 0xD29C5380, 0x77DDC1FE, + 0x9DF3018D, 0x38B293F3, 0x7413C95F, 0xD1525B21, 0x3B7C9B52, 0x9E3D092C, + 0xEACD6D45, 0x4F8CFF3B, 0xA5A23F48, 0x00E3AD36, 0x3CE08A10, 0x99A1186E, + 0x738FD81D, 0xD6CE4A63, 0xA23E2E0A, 0x077FBC74, 0xED517C07, 0x4810EE79, + 0x04B1B4D5, 0xA1F026AB, 0x4BDEE6D8, 0xEE9F74A6, 0x9A6F10CF, 0x3F2E82B1, + 0xD50042C2, 0x7041D0BC, 0xAD060C8E, 0x08479EF0, 0xE2695E83, 0x4728CCFD, + 0x33D8A894, 0x96993AEA, 0x7CB7FA99, 0xD9F668E7, 0x9557324B, 0x3016A035, + 0xDA386046, 0x7F79F238, 0x0B899651, 0xAEC8042F, 0x44E6C45C, 0xE1A75622, + 0xDDA47104, 0x78E5E37A, 0x92CB2309, 0x378AB177, 0x437AD51E, 0xE63B4760, + 0x0C158713, 0xA954156D, 0xE5F54FC1, 0x40B4DDBF, 0xAA9A1DCC, 0x0FDB8FB2, + 0x7B2BEBDB, 0xDE6A79A5, 0x3444B9D6, 0x91052BA8}; + +/** + * The following CRC lookup table was generated automagically + * using the following model parameters: + * + * Generator Polynomial = ................. 0x1EDC6F41 + * Generator Polynomial Length = .......... 32 bits + * Reflected Bits = ....................... TRUE + * Table Generation Offset = .............. 32 bits + * Number of Slices = ..................... 8 slices + * Slice Lengths = ........................ 8 8 8 8 8 8 8 8 + */ +static uint32_t crc_tableil8_o56[256] = { + 0x00000000, 0xDD45AAB8, 0xBF672381, 0x62228939, 0x7B2231F3, 0xA6679B4B, + 0xC4451272, 0x1900B8CA, 0xF64463E6, 0x2B01C95E, 0x49234067, 0x9466EADF, + 0x8D665215, 0x5023F8AD, 0x32017194, 0xEF44DB2C, 0xE964B13D, 0x34211B85, + 0x560392BC, 0x8B463804, 0x924680CE, 0x4F032A76, 0x2D21A34F, 0xF06409F7, + 0x1F20D2DB, 0xC2657863, 0xA047F15A, 0x7D025BE2, 0x6402E328, 0xB9474990, + 0xDB65C0A9, 0x06206A11, 0xD725148B, 0x0A60BE33, 0x6842370A, 0xB5079DB2, + 0xAC072578, 0x71428FC0, 0x136006F9, 0xCE25AC41, 0x2161776D, 0xFC24DDD5, + 0x9E0654EC, 0x4343FE54, 0x5A43469E, 0x8706EC26, 0xE524651F, 0x3861CFA7, + 0x3E41A5B6, 0xE3040F0E, 0x81268637, 0x5C632C8F, 0x45639445, 0x98263EFD, + 0xFA04B7C4, 0x27411D7C, 0xC805C650, 0x15406CE8, 0x7762E5D1, 0xAA274F69, + 0xB327F7A3, 0x6E625D1B, 0x0C40D422, 0xD1057E9A, 0xABA65FE7, 0x76E3F55F, + 0x14C17C66, 0xC984D6DE, 0xD0846E14, 0x0DC1C4AC, 0x6FE34D95, 0xB2A6E72D, + 0x5DE23C01, 0x80A796B9, 0xE2851F80, 0x3FC0B538, 0x26C00DF2, 0xFB85A74A, + 0x99A72E73, 0x44E284CB, 0x42C2EEDA, 0x9F874462, 0xFDA5CD5B, 0x20E067E3, + 0x39E0DF29, 0xE4A57591, 0x8687FCA8, 0x5BC25610, 0xB4868D3C, 0x69C32784, + 0x0BE1AEBD, 0xD6A40405, 0xCFA4BCCF, 0x12E11677, 0x70C39F4E, 0xAD8635F6, + 0x7C834B6C, 0xA1C6E1D4, 0xC3E468ED, 0x1EA1C255, 0x07A17A9F, 0xDAE4D027, + 0xB8C6591E, 0x6583F3A6, 0x8AC7288A, 0x57828232, 0x35A00B0B, 0xE8E5A1B3, + 0xF1E51979, 0x2CA0B3C1, 0x4E823AF8, 0x93C79040, 0x95E7FA51, 0x48A250E9, + 0x2A80D9D0, 0xF7C57368, 0xEEC5CBA2, 0x3380611A, 0x51A2E823, 0x8CE7429B, + 0x63A399B7, 0xBEE6330F, 0xDCC4BA36, 0x0181108E, 0x1881A844, 0xC5C402FC, + 0xA7E68BC5, 0x7AA3217D, 0x52A0C93F, 0x8FE56387, 0xEDC7EABE, 0x30824006, + 0x2982F8CC, 0xF4C75274, 0x96E5DB4D, 0x4BA071F5, 0xA4E4AAD9, 0x79A10061, + 0x1B838958, 0xC6C623E0, 0xDFC69B2A, 0x02833192, 0x60A1B8AB, 0xBDE41213, + 0xBBC47802, 0x6681D2BA, 0x04A35B83, 0xD9E6F13B, 0xC0E649F1, 0x1DA3E349, + 0x7F816A70, 0xA2C4C0C8, 0x4D801BE4, 0x90C5B15C, 0xF2E73865, 0x2FA292DD, + 0x36A22A17, 0xEBE780AF, 0x89C50996, 0x5480A32E, 0x8585DDB4, 0x58C0770C, + 0x3AE2FE35, 0xE7A7548D, 0xFEA7EC47, 0x23E246FF, 0x41C0CFC6, 0x9C85657E, + 0x73C1BE52, 0xAE8414EA, 0xCCA69DD3, 0x11E3376B, 0x08E38FA1, 0xD5A62519, + 0xB784AC20, 0x6AC10698, 0x6CE16C89, 0xB1A4C631, 0xD3864F08, 0x0EC3E5B0, + 0x17C35D7A, 0xCA86F7C2, 0xA8A47EFB, 0x75E1D443, 0x9AA50F6F, 0x47E0A5D7, + 0x25C22CEE, 0xF8878656, 0xE1873E9C, 0x3CC29424, 0x5EE01D1D, 0x83A5B7A5, + 0xF90696D8, 0x24433C60, 0x4661B559, 0x9B241FE1, 0x8224A72B, 0x5F610D93, + 0x3D4384AA, 0xE0062E12, 0x0F42F53E, 0xD2075F86, 0xB025D6BF, 0x6D607C07, + 0x7460C4CD, 0xA9256E75, 0xCB07E74C, 0x16424DF4, 0x106227E5, 0xCD278D5D, + 0xAF050464, 0x7240AEDC, 0x6B401616, 0xB605BCAE, 0xD4273597, 0x09629F2F, + 0xE6264403, 0x3B63EEBB, 0x59416782, 0x8404CD3A, 0x9D0475F0, 0x4041DF48, + 0x22635671, 0xFF26FCC9, 0x2E238253, 0xF36628EB, 0x9144A1D2, 0x4C010B6A, + 0x5501B3A0, 0x88441918, 0xEA669021, 0x37233A99, 0xD867E1B5, 0x05224B0D, + 0x6700C234, 0xBA45688C, 0xA345D046, 0x7E007AFE, 0x1C22F3C7, 0xC167597F, + 0xC747336E, 0x1A0299D6, 0x782010EF, 0xA565BA57, 0xBC65029D, 0x6120A825, + 0x0302211C, 0xDE478BA4, 0x31035088, 0xEC46FA30, 0x8E647309, 0x5321D9B1, + 0x4A21617B, 0x9764CBC3, 0xF54642FA, 0x2803E842}; + +/** + * The following CRC lookup table was generated automagically + * using the following model parameters: + * + * Generator Polynomial = ................. 0x1EDC6F41 + * Generator Polynomial Length = .......... 32 bits + * Reflected Bits = ....................... TRUE + * Table Generation Offset = .............. 32 bits + * Number of Slices = ..................... 8 slices + * Slice Lengths = ........................ 8 8 8 8 8 8 8 8 + */ +static uint32_t crc_tableil8_o64[256] = { + 0x00000000, 0x38116FAC, 0x7022DF58, 0x4833B0F4, 0xE045BEB0, 0xD854D11C, + 0x906761E8, 0xA8760E44, 0xC5670B91, 0xFD76643D, 0xB545D4C9, 0x8D54BB65, + 0x2522B521, 0x1D33DA8D, 0x55006A79, 0x6D1105D5, 0x8F2261D3, 0xB7330E7F, + 0xFF00BE8B, 0xC711D127, 0x6F67DF63, 0x5776B0CF, 0x1F45003B, 0x27546F97, + 0x4A456A42, 0x725405EE, 0x3A67B51A, 0x0276DAB6, 0xAA00D4F2, 0x9211BB5E, + 0xDA220BAA, 0xE2336406, 0x1BA8B557, 0x23B9DAFB, 0x6B8A6A0F, 0x539B05A3, + 0xFBED0BE7, 0xC3FC644B, 0x8BCFD4BF, 0xB3DEBB13, 0xDECFBEC6, 0xE6DED16A, + 0xAEED619E, 0x96FC0E32, 0x3E8A0076, 0x069B6FDA, 0x4EA8DF2E, 0x76B9B082, + 0x948AD484, 0xAC9BBB28, 0xE4A80BDC, 0xDCB96470, 0x74CF6A34, 0x4CDE0598, + 0x04EDB56C, 0x3CFCDAC0, 0x51EDDF15, 0x69FCB0B9, 0x21CF004D, 0x19DE6FE1, + 0xB1A861A5, 0x89B90E09, 0xC18ABEFD, 0xF99BD151, 0x37516AAE, 0x0F400502, + 0x4773B5F6, 0x7F62DA5A, 0xD714D41E, 0xEF05BBB2, 0xA7360B46, 0x9F2764EA, + 0xF236613F, 0xCA270E93, 0x8214BE67, 0xBA05D1CB, 0x1273DF8F, 0x2A62B023, + 0x625100D7, 0x5A406F7B, 0xB8730B7D, 0x806264D1, 0xC851D425, 0xF040BB89, + 0x5836B5CD, 0x6027DA61, 0x28146A95, 0x10050539, 0x7D1400EC, 0x45056F40, + 0x0D36DFB4, 0x3527B018, 0x9D51BE5C, 0xA540D1F0, 0xED736104, 0xD5620EA8, + 0x2CF9DFF9, 0x14E8B055, 0x5CDB00A1, 0x64CA6F0D, 0xCCBC6149, 0xF4AD0EE5, + 0xBC9EBE11, 0x848FD1BD, 0xE99ED468, 0xD18FBBC4, 0x99BC0B30, 0xA1AD649C, + 0x09DB6AD8, 0x31CA0574, 0x79F9B580, 0x41E8DA2C, 0xA3DBBE2A, 0x9BCAD186, + 0xD3F96172, 0xEBE80EDE, 0x439E009A, 0x7B8F6F36, 0x33BCDFC2, 0x0BADB06E, + 0x66BCB5BB, 0x5EADDA17, 0x169E6AE3, 0x2E8F054F, 0x86F90B0B, 0xBEE864A7, + 0xF6DBD453, 0xCECABBFF, 0x6EA2D55C, 0x56B3BAF0, 0x1E800A04, 0x269165A8, + 0x8EE76BEC, 0xB6F60440, 0xFEC5B4B4, 0xC6D4DB18, 0xABC5DECD, 0x93D4B161, + 0xDBE70195, 0xE3F66E39, 0x4B80607D, 0x73910FD1, 0x3BA2BF25, 0x03B3D089, + 0xE180B48F, 0xD991DB23, 0x91A26BD7, 0xA9B3047B, 0x01C50A3F, 0x39D46593, + 0x71E7D567, 0x49F6BACB, 0x24E7BF1E, 0x1CF6D0B2, 0x54C56046, 0x6CD40FEA, + 0xC4A201AE, 0xFCB36E02, 0xB480DEF6, 0x8C91B15A, 0x750A600B, 0x4D1B0FA7, + 0x0528BF53, 0x3D39D0FF, 0x954FDEBB, 0xAD5EB117, 0xE56D01E3, 0xDD7C6E4F, + 0xB06D6B9A, 0x887C0436, 0xC04FB4C2, 0xF85EDB6E, 0x5028D52A, 0x6839BA86, + 0x200A0A72, 0x181B65DE, 0xFA2801D8, 0xC2396E74, 0x8A0ADE80, 0xB21BB12C, + 0x1A6DBF68, 0x227CD0C4, 0x6A4F6030, 0x525E0F9C, 0x3F4F0A49, 0x075E65E5, + 0x4F6DD511, 0x777CBABD, 0xDF0AB4F9, 0xE71BDB55, 0xAF286BA1, 0x9739040D, + 0x59F3BFF2, 0x61E2D05E, 0x29D160AA, 0x11C00F06, 0xB9B60142, 0x81A76EEE, + 0xC994DE1A, 0xF185B1B6, 0x9C94B463, 0xA485DBCF, 0xECB66B3B, 0xD4A70497, + 0x7CD10AD3, 0x44C0657F, 0x0CF3D58B, 0x34E2BA27, 0xD6D1DE21, 0xEEC0B18D, + 0xA6F30179, 0x9EE26ED5, 0x36946091, 0x0E850F3D, 0x46B6BFC9, 0x7EA7D065, + 0x13B6D5B0, 0x2BA7BA1C, 0x63940AE8, 0x5B856544, 0xF3F36B00, 0xCBE204AC, + 0x83D1B458, 0xBBC0DBF4, 0x425B0AA5, 0x7A4A6509, 0x3279D5FD, 0x0A68BA51, + 0xA21EB415, 0x9A0FDBB9, 0xD23C6B4D, 0xEA2D04E1, 0x873C0134, 0xBF2D6E98, + 0xF71EDE6C, 0xCF0FB1C0, 0x6779BF84, 0x5F68D028, 0x175B60DC, 0x2F4A0F70, + 0xCD796B76, 0xF56804DA, 0xBD5BB42E, 0x854ADB82, 0x2D3CD5C6, 0x152DBA6A, + 0x5D1E0A9E, 0x650F6532, 0x081E60E7, 0x300F0F4B, 0x783CBFBF, 0x402DD013, + 0xE85BDE57, 0xD04AB1FB, 0x9879010F, 0xA0686EA3}; + +/** + * The following CRC lookup table was generated automagically + * using the following model parameters: + * + * Generator Polynomial = ................. 0x1EDC6F41 + * Generator Polynomial Length = .......... 32 bits + * Reflected Bits = ....................... TRUE + * Table Generation Offset = .............. 32 bits + * Number of Slices = ..................... 8 slices + * Slice Lengths = ........................ 8 8 8 8 8 8 8 8 + */ +static uint32_t crc_tableil8_o72[256] = { + 0x00000000, 0xEF306B19, 0xDB8CA0C3, 0x34BCCBDA, 0xB2F53777, 0x5DC55C6E, + 0x697997B4, 0x8649FCAD, 0x6006181F, 0x8F367306, 0xBB8AB8DC, 0x54BAD3C5, + 0xD2F32F68, 0x3DC34471, 0x097F8FAB, 0xE64FE4B2, 0xC00C303E, 0x2F3C5B27, + 0x1B8090FD, 0xF4B0FBE4, 0x72F90749, 0x9DC96C50, 0xA975A78A, 0x4645CC93, + 0xA00A2821, 0x4F3A4338, 0x7B8688E2, 0x94B6E3FB, 0x12FF1F56, 0xFDCF744F, + 0xC973BF95, 0x2643D48C, 0x85F4168D, 0x6AC47D94, 0x5E78B64E, 0xB148DD57, + 0x370121FA, 0xD8314AE3, 0xEC8D8139, 0x03BDEA20, 0xE5F20E92, 0x0AC2658B, + 0x3E7EAE51, 0xD14EC548, 0x570739E5, 0xB83752FC, 0x8C8B9926, 0x63BBF23F, + 0x45F826B3, 0xAAC84DAA, 0x9E748670, 0x7144ED69, 0xF70D11C4, 0x183D7ADD, + 0x2C81B107, 0xC3B1DA1E, 0x25FE3EAC, 0xCACE55B5, 0xFE729E6F, 0x1142F576, + 0x970B09DB, 0x783B62C2, 0x4C87A918, 0xA3B7C201, 0x0E045BEB, 0xE13430F2, + 0xD588FB28, 0x3AB89031, 0xBCF16C9C, 0x53C10785, 0x677DCC5F, 0x884DA746, + 0x6E0243F4, 0x813228ED, 0xB58EE337, 0x5ABE882E, 0xDCF77483, 0x33C71F9A, + 0x077BD440, 0xE84BBF59, 0xCE086BD5, 0x213800CC, 0x1584CB16, 0xFAB4A00F, + 0x7CFD5CA2, 0x93CD37BB, 0xA771FC61, 0x48419778, 0xAE0E73CA, 0x413E18D3, + 0x7582D309, 0x9AB2B810, 0x1CFB44BD, 0xF3CB2FA4, 0xC777E47E, 0x28478F67, + 0x8BF04D66, 0x64C0267F, 0x507CEDA5, 0xBF4C86BC, 0x39057A11, 0xD6351108, + 0xE289DAD2, 0x0DB9B1CB, 0xEBF65579, 0x04C63E60, 0x307AF5BA, 0xDF4A9EA3, + 0x5903620E, 0xB6330917, 0x828FC2CD, 0x6DBFA9D4, 0x4BFC7D58, 0xA4CC1641, + 0x9070DD9B, 0x7F40B682, 0xF9094A2F, 0x16392136, 0x2285EAEC, 0xCDB581F5, + 0x2BFA6547, 0xC4CA0E5E, 0xF076C584, 0x1F46AE9D, 0x990F5230, 0x763F3929, + 0x4283F2F3, 0xADB399EA, 0x1C08B7D6, 0xF338DCCF, 0xC7841715, 0x28B47C0C, + 0xAEFD80A1, 0x41CDEBB8, 0x75712062, 0x9A414B7B, 0x7C0EAFC9, 0x933EC4D0, + 0xA7820F0A, 0x48B26413, 0xCEFB98BE, 0x21CBF3A7, 0x1577387D, 0xFA475364, + 0xDC0487E8, 0x3334ECF1, 0x0788272B, 0xE8B84C32, 0x6EF1B09F, 0x81C1DB86, + 0xB57D105C, 0x5A4D7B45, 0xBC029FF7, 0x5332F4EE, 0x678E3F34, 0x88BE542D, + 0x0EF7A880, 0xE1C7C399, 0xD57B0843, 0x3A4B635A, 0x99FCA15B, 0x76CCCA42, + 0x42700198, 0xAD406A81, 0x2B09962C, 0xC439FD35, 0xF08536EF, 0x1FB55DF6, + 0xF9FAB944, 0x16CAD25D, 0x22761987, 0xCD46729E, 0x4B0F8E33, 0xA43FE52A, + 0x90832EF0, 0x7FB345E9, 0x59F09165, 0xB6C0FA7C, 0x827C31A6, 0x6D4C5ABF, + 0xEB05A612, 0x0435CD0B, 0x308906D1, 0xDFB96DC8, 0x39F6897A, 0xD6C6E263, + 0xE27A29B9, 0x0D4A42A0, 0x8B03BE0D, 0x6433D514, 0x508F1ECE, 0xBFBF75D7, + 0x120CEC3D, 0xFD3C8724, 0xC9804CFE, 0x26B027E7, 0xA0F9DB4A, 0x4FC9B053, + 0x7B757B89, 0x94451090, 0x720AF422, 0x9D3A9F3B, 0xA98654E1, 0x46B63FF8, + 0xC0FFC355, 0x2FCFA84C, 0x1B736396, 0xF443088F, 0xD200DC03, 0x3D30B71A, + 0x098C7CC0, 0xE6BC17D9, 0x60F5EB74, 0x8FC5806D, 0xBB794BB7, 0x544920AE, + 0xB206C41C, 0x5D36AF05, 0x698A64DF, 0x86BA0FC6, 0x00F3F36B, 0xEFC39872, + 0xDB7F53A8, 0x344F38B1, 0x97F8FAB0, 0x78C891A9, 0x4C745A73, 0xA344316A, + 0x250DCDC7, 0xCA3DA6DE, 0xFE816D04, 0x11B1061D, 0xF7FEE2AF, 0x18CE89B6, + 0x2C72426C, 0xC3422975, 0x450BD5D8, 0xAA3BBEC1, 0x9E87751B, 0x71B71E02, + 0x57F4CA8E, 0xB8C4A197, 0x8C786A4D, 0x63480154, 0xE501FDF9, 0x0A3196E0, + 0x3E8D5D3A, 0xD1BD3623, 0x37F2D291, 0xD8C2B988, 0xEC7E7252, 0x034E194B, + 0x8507E5E6, 0x6A378EFF, 0x5E8B4525, 0xB1BB2E3C}; + +/** + * The following CRC lookup table was generated automagically + * using the following model parameters: + * + * Generator Polynomial = ................. 0x1EDC6F41 + * Generator Polynomial Length = .......... 32 bits + * Reflected Bits = ....................... TRUE + * Table Generation Offset = .............. 32 bits + * Number of Slices = ..................... 8 slices + * Slice Lengths = ........................ 8 8 8 8 8 8 8 8 + */ +static uint32_t crc_tableil8_o80[256] = { + 0x00000000, 0x68032CC8, 0xD0065990, 0xB8057558, 0xA5E0C5D1, 0xCDE3E919, + 0x75E69C41, 0x1DE5B089, 0x4E2DFD53, 0x262ED19B, 0x9E2BA4C3, 0xF628880B, + 0xEBCD3882, 0x83CE144A, 0x3BCB6112, 0x53C84DDA, 0x9C5BFAA6, 0xF458D66E, + 0x4C5DA336, 0x245E8FFE, 0x39BB3F77, 0x51B813BF, 0xE9BD66E7, 0x81BE4A2F, + 0xD27607F5, 0xBA752B3D, 0x02705E65, 0x6A7372AD, 0x7796C224, 0x1F95EEEC, + 0xA7909BB4, 0xCF93B77C, 0x3D5B83BD, 0x5558AF75, 0xED5DDA2D, 0x855EF6E5, + 0x98BB466C, 0xF0B86AA4, 0x48BD1FFC, 0x20BE3334, 0x73767EEE, 0x1B755226, + 0xA370277E, 0xCB730BB6, 0xD696BB3F, 0xBE9597F7, 0x0690E2AF, 0x6E93CE67, + 0xA100791B, 0xC90355D3, 0x7106208B, 0x19050C43, 0x04E0BCCA, 0x6CE39002, + 0xD4E6E55A, 0xBCE5C992, 0xEF2D8448, 0x872EA880, 0x3F2BDDD8, 0x5728F110, + 0x4ACD4199, 0x22CE6D51, 0x9ACB1809, 0xF2C834C1, 0x7AB7077A, 0x12B42BB2, + 0xAAB15EEA, 0xC2B27222, 0xDF57C2AB, 0xB754EE63, 0x0F519B3B, 0x6752B7F3, + 0x349AFA29, 0x5C99D6E1, 0xE49CA3B9, 0x8C9F8F71, 0x917A3FF8, 0xF9791330, + 0x417C6668, 0x297F4AA0, 0xE6ECFDDC, 0x8EEFD114, 0x36EAA44C, 0x5EE98884, + 0x430C380D, 0x2B0F14C5, 0x930A619D, 0xFB094D55, 0xA8C1008F, 0xC0C22C47, + 0x78C7591F, 0x10C475D7, 0x0D21C55E, 0x6522E996, 0xDD279CCE, 0xB524B006, + 0x47EC84C7, 0x2FEFA80F, 0x97EADD57, 0xFFE9F19F, 0xE20C4116, 0x8A0F6DDE, + 0x320A1886, 0x5A09344E, 0x09C17994, 0x61C2555C, 0xD9C72004, 0xB1C40CCC, + 0xAC21BC45, 0xC422908D, 0x7C27E5D5, 0x1424C91D, 0xDBB77E61, 0xB3B452A9, + 0x0BB127F1, 0x63B20B39, 0x7E57BBB0, 0x16549778, 0xAE51E220, 0xC652CEE8, + 0x959A8332, 0xFD99AFFA, 0x459CDAA2, 0x2D9FF66A, 0x307A46E3, 0x58796A2B, + 0xE07C1F73, 0x887F33BB, 0xF56E0EF4, 0x9D6D223C, 0x25685764, 0x4D6B7BAC, + 0x508ECB25, 0x388DE7ED, 0x808892B5, 0xE88BBE7D, 0xBB43F3A7, 0xD340DF6F, + 0x6B45AA37, 0x034686FF, 0x1EA33676, 0x76A01ABE, 0xCEA56FE6, 0xA6A6432E, + 0x6935F452, 0x0136D89A, 0xB933ADC2, 0xD130810A, 0xCCD53183, 0xA4D61D4B, + 0x1CD36813, 0x74D044DB, 0x27180901, 0x4F1B25C9, 0xF71E5091, 0x9F1D7C59, + 0x82F8CCD0, 0xEAFBE018, 0x52FE9540, 0x3AFDB988, 0xC8358D49, 0xA036A181, + 0x1833D4D9, 0x7030F811, 0x6DD54898, 0x05D66450, 0xBDD31108, 0xD5D03DC0, + 0x8618701A, 0xEE1B5CD2, 0x561E298A, 0x3E1D0542, 0x23F8B5CB, 0x4BFB9903, + 0xF3FEEC5B, 0x9BFDC093, 0x546E77EF, 0x3C6D5B27, 0x84682E7F, 0xEC6B02B7, + 0xF18EB23E, 0x998D9EF6, 0x2188EBAE, 0x498BC766, 0x1A438ABC, 0x7240A674, + 0xCA45D32C, 0xA246FFE4, 0xBFA34F6D, 0xD7A063A5, 0x6FA516FD, 0x07A63A35, + 0x8FD9098E, 0xE7DA2546, 0x5FDF501E, 0x37DC7CD6, 0x2A39CC5F, 0x423AE097, + 0xFA3F95CF, 0x923CB907, 0xC1F4F4DD, 0xA9F7D815, 0x11F2AD4D, 0x79F18185, + 0x6414310C, 0x0C171DC4, 0xB412689C, 0xDC114454, 0x1382F328, 0x7B81DFE0, + 0xC384AAB8, 0xAB878670, 0xB66236F9, 0xDE611A31, 0x66646F69, 0x0E6743A1, + 0x5DAF0E7B, 0x35AC22B3, 0x8DA957EB, 0xE5AA7B23, 0xF84FCBAA, 0x904CE762, + 0x2849923A, 0x404ABEF2, 0xB2828A33, 0xDA81A6FB, 0x6284D3A3, 0x0A87FF6B, + 0x17624FE2, 0x7F61632A, 0xC7641672, 0xAF673ABA, 0xFCAF7760, 0x94AC5BA8, + 0x2CA92EF0, 0x44AA0238, 0x594FB2B1, 0x314C9E79, 0x8949EB21, 0xE14AC7E9, + 0x2ED97095, 0x46DA5C5D, 0xFEDF2905, 0x96DC05CD, 0x8B39B544, 0xE33A998C, + 0x5B3FECD4, 0x333CC01C, 0x60F48DC6, 0x08F7A10E, 0xB0F2D456, 0xD8F1F89E, + 0xC5144817, 0xAD1764DF, 0x15121187, 0x7D113D4F}; + +/** + * The following CRC lookup table was generated automagically + * using the following model parameters: + * + * Generator Polynomial = ................. 0x1EDC6F41 + * Generator Polynomial Length = .......... 32 bits + * Reflected Bits = ....................... TRUE + * Table Generation Offset = .............. 32 bits + * Number of Slices = ..................... 8 slices + * Slice Lengths = ........................ 8 8 8 8 8 8 8 8 + */ +static uint32_t crc_tableil8_o88[256] = { + 0x00000000, 0x493C7D27, 0x9278FA4E, 0xDB448769, 0x211D826D, 0x6821FF4A, + 0xB3657823, 0xFA590504, 0x423B04DA, 0x0B0779FD, 0xD043FE94, 0x997F83B3, + 0x632686B7, 0x2A1AFB90, 0xF15E7CF9, 0xB86201DE, 0x847609B4, 0xCD4A7493, + 0x160EF3FA, 0x5F328EDD, 0xA56B8BD9, 0xEC57F6FE, 0x37137197, 0x7E2F0CB0, + 0xC64D0D6E, 0x8F717049, 0x5435F720, 0x1D098A07, 0xE7508F03, 0xAE6CF224, + 0x7528754D, 0x3C14086A, 0x0D006599, 0x443C18BE, 0x9F789FD7, 0xD644E2F0, + 0x2C1DE7F4, 0x65219AD3, 0xBE651DBA, 0xF759609D, 0x4F3B6143, 0x06071C64, + 0xDD439B0D, 0x947FE62A, 0x6E26E32E, 0x271A9E09, 0xFC5E1960, 0xB5626447, + 0x89766C2D, 0xC04A110A, 0x1B0E9663, 0x5232EB44, 0xA86BEE40, 0xE1579367, + 0x3A13140E, 0x732F6929, 0xCB4D68F7, 0x827115D0, 0x593592B9, 0x1009EF9E, + 0xEA50EA9A, 0xA36C97BD, 0x782810D4, 0x31146DF3, 0x1A00CB32, 0x533CB615, + 0x8878317C, 0xC1444C5B, 0x3B1D495F, 0x72213478, 0xA965B311, 0xE059CE36, + 0x583BCFE8, 0x1107B2CF, 0xCA4335A6, 0x837F4881, 0x79264D85, 0x301A30A2, + 0xEB5EB7CB, 0xA262CAEC, 0x9E76C286, 0xD74ABFA1, 0x0C0E38C8, 0x453245EF, + 0xBF6B40EB, 0xF6573DCC, 0x2D13BAA5, 0x642FC782, 0xDC4DC65C, 0x9571BB7B, + 0x4E353C12, 0x07094135, 0xFD504431, 0xB46C3916, 0x6F28BE7F, 0x2614C358, + 0x1700AEAB, 0x5E3CD38C, 0x857854E5, 0xCC4429C2, 0x361D2CC6, 0x7F2151E1, + 0xA465D688, 0xED59ABAF, 0x553BAA71, 0x1C07D756, 0xC743503F, 0x8E7F2D18, + 0x7426281C, 0x3D1A553B, 0xE65ED252, 0xAF62AF75, 0x9376A71F, 0xDA4ADA38, + 0x010E5D51, 0x48322076, 0xB26B2572, 0xFB575855, 0x2013DF3C, 0x692FA21B, + 0xD14DA3C5, 0x9871DEE2, 0x4335598B, 0x0A0924AC, 0xF05021A8, 0xB96C5C8F, + 0x6228DBE6, 0x2B14A6C1, 0x34019664, 0x7D3DEB43, 0xA6796C2A, 0xEF45110D, + 0x151C1409, 0x5C20692E, 0x8764EE47, 0xCE589360, 0x763A92BE, 0x3F06EF99, + 0xE44268F0, 0xAD7E15D7, 0x572710D3, 0x1E1B6DF4, 0xC55FEA9D, 0x8C6397BA, + 0xB0779FD0, 0xF94BE2F7, 0x220F659E, 0x6B3318B9, 0x916A1DBD, 0xD856609A, + 0x0312E7F3, 0x4A2E9AD4, 0xF24C9B0A, 0xBB70E62D, 0x60346144, 0x29081C63, + 0xD3511967, 0x9A6D6440, 0x4129E329, 0x08159E0E, 0x3901F3FD, 0x703D8EDA, + 0xAB7909B3, 0xE2457494, 0x181C7190, 0x51200CB7, 0x8A648BDE, 0xC358F6F9, + 0x7B3AF727, 0x32068A00, 0xE9420D69, 0xA07E704E, 0x5A27754A, 0x131B086D, + 0xC85F8F04, 0x8163F223, 0xBD77FA49, 0xF44B876E, 0x2F0F0007, 0x66337D20, + 0x9C6A7824, 0xD5560503, 0x0E12826A, 0x472EFF4D, 0xFF4CFE93, 0xB67083B4, + 0x6D3404DD, 0x240879FA, 0xDE517CFE, 0x976D01D9, 0x4C2986B0, 0x0515FB97, + 0x2E015D56, 0x673D2071, 0xBC79A718, 0xF545DA3F, 0x0F1CDF3B, 0x4620A21C, + 0x9D642575, 0xD4585852, 0x6C3A598C, 0x250624AB, 0xFE42A3C2, 0xB77EDEE5, + 0x4D27DBE1, 0x041BA6C6, 0xDF5F21AF, 0x96635C88, 0xAA7754E2, 0xE34B29C5, + 0x380FAEAC, 0x7133D38B, 0x8B6AD68F, 0xC256ABA8, 0x19122CC1, 0x502E51E6, + 0xE84C5038, 0xA1702D1F, 0x7A34AA76, 0x3308D751, 0xC951D255, 0x806DAF72, + 0x5B29281B, 0x1215553C, 0x230138CF, 0x6A3D45E8, 0xB179C281, 0xF845BFA6, + 0x021CBAA2, 0x4B20C785, 0x906440EC, 0xD9583DCB, 0x613A3C15, 0x28064132, + 0xF342C65B, 0xBA7EBB7C, 0x4027BE78, 0x091BC35F, 0xD25F4436, 0x9B633911, + 0xA777317B, 0xEE4B4C5C, 0x350FCB35, 0x7C33B612, 0x866AB316, 0xCF56CE31, + 0x14124958, 0x5D2E347F, 0xE54C35A1, 0xAC704886, 0x7734CFEF, 0x3E08B2C8, + 0xC451B7CC, 0x8D6DCAEB, 0x56294D82, 0x1F1530A5}; + +/** + * Implementations adapted from Intel's Slicing By 8 Sourceforge Project + * http://sourceforge.net/projects/slicing-by-8/ + * http://www.evanjones.ca/crc32c.html + */ +static inline uint32_t crc32c_slicing8(const void *data, size_t len, + uint32_t crc) { + const uint8_t *p_buf = (const uint8_t *)data; + + /* Handle leading misaligned bytes */ + size_t init_bytes = + (sizeof(int32_t) - (intptr_t)p_buf) & (sizeof(int32_t) - 1); + if (len < init_bytes) { + init_bytes = len; + } + for (size_t li = 0; li < init_bytes; li++) { + crc = crc_tableil8_o32[(crc ^ *p_buf++) & 0x000000FF] ^ (crc >> 8); + } + + len -= init_bytes; + size_t running_length = len & ~(sizeof(uint64_t) - 1); + size_t end_bytes = len - running_length; + + for (size_t li = 0; li < running_length / 8; li++) { + uint32_t term1, term2; + + crc ^= *(uint32_t *)p_buf; + p_buf += 4; + term1 = crc_tableil8_o88[crc & 0x000000FF] ^ + crc_tableil8_o80[(crc >> 8) & 0x000000FF]; + term2 = crc >> 16; + crc = term1 ^ crc_tableil8_o72[term2 & 0x000000FF] ^ + crc_tableil8_o64[(term2 >> 8) & 0x000000FF]; + term1 = crc_tableil8_o56[(*(uint32_t *)p_buf) & 0x000000FF] ^ + crc_tableil8_o48[((*(uint32_t *)p_buf) >> 8) & 0x000000FF]; + + term2 = (*(uint32_t *)p_buf) >> 16; + crc = crc ^ term1 ^ crc_tableil8_o40[term2 & 0x000000FF] ^ + crc_tableil8_o32[(term2 >> 8) & 0x000000FF]; + p_buf += 4; + } + + for (size_t li = 0; li < end_bytes; li++) { + crc = crc_tableil8_o32[(crc ^ *p_buf++) & 0x000000FF] ^ (crc >> 8); + } + return crc; +} +#endif // !__SSE4_2__ + +#if defined(__SSE4_2__) +#if defined(AILEGO_M64) +static inline uint32_t crc32c_sse42(const void *data, size_t len, + uint32_t crc) { + const uint8_t *first = (const uint8_t *)data; + const uint8_t *last = first + ((len >> 3) << 3); + + for (; first != last; first += 8) { + crc = (uint32_t)_mm_crc32_u64(crc, *(uint64_t *)first); + } + switch (((uint8_t *)data + len) - last) { + case 1: + crc = _mm_crc32_u8(crc, *last); + break; + case 2: + crc = _mm_crc32_u16(crc, *(uint16_t *)last); + break; + case 3: + crc = _mm_crc32_u16(crc, *(uint16_t *)last); + crc = _mm_crc32_u8(crc, *(last + 2)); + break; + case 4: + crc = _mm_crc32_u32(crc, *(uint32_t *)last); + break; + case 5: + crc = _mm_crc32_u32(crc, *(uint32_t *)last); + crc = _mm_crc32_u8(crc, *(last + 4)); + break; + case 6: + crc = _mm_crc32_u32(crc, *(uint32_t *)last); + crc = _mm_crc32_u16(crc, *(uint16_t *)(last + 4)); + break; + case 7: + crc = _mm_crc32_u32(crc, *(uint32_t *)last); + crc = _mm_crc32_u16(crc, *(uint16_t *)(last + 4)); + crc = _mm_crc32_u8(crc, *(last + 6)); + break; + } + return crc; +} +#else +static inline uint32_t crc32c_sse42(const void *data, size_t len, + uint32_t crc) { + const uint8_t *first = (const uint8_t *)data; + const uint8_t *last = first + ((len >> 2) << 2); + + for (; first != last; first += 4) { + crc = _mm_crc32_u32(crc, *(uint32_t *)first); + } + switch (((uint8_t *)data + len) - last) { + case 1: + crc = _mm_crc32_u8(crc, *last); + break; + case 2: + crc = _mm_crc32_u16(crc, *(uint16_t *)last); + break; + case 3: + crc = _mm_crc32_u16(crc, *(uint16_t *)last); + crc = _mm_crc32_u8(crc, *(last + 2)); + break; + } + return crc; +} +#endif // AILEGO_M64 +#endif // __SSE4_2__ + +#if defined(__ARM_FEATURE_CRC32) +static inline uint32_t crc32c_neon(const void *data, size_t len, uint32_t crc) { + const uint8_t *first = (const uint8_t *)data; + const uint8_t *last = first + ((len >> 3) << 3); + + for (; first != last; first += 8) { + crc = __crc32cd(crc, *(uint64_t *)first); + } + switch (((uint8_t *)data + len) - last) { + case 1: + crc = __crc32cb(crc, *last); + break; + case 2: + crc = __crc32ch(crc, *(uint16_t *)last); + break; + case 3: + crc = __crc32ch(crc, *(uint16_t *)last); + crc = __crc32cb(crc, *(last + 2)); + break; + case 4: + crc = __crc32cw(crc, *(uint32_t *)last); + break; + case 5: + crc = __crc32cw(crc, *(uint32_t *)last); + crc = __crc32cb(crc, *(last + 4)); + break; + case 6: + crc = __crc32cw(crc, *(uint32_t *)last); + crc = __crc32ch(crc, *(uint16_t *)(last + 4)); + break; + case 7: + crc = __crc32cw(crc, *(uint32_t *)last); + crc = __crc32ch(crc, *(uint16_t *)(last + 4)); + crc = __crc32cb(crc, *(last + 6)); + break; + } + return crc; +} +#endif // __ARM_FEATURE_CRC32 + +namespace zvec { +namespace ailego { + +uint32_t Crc32c::Hash(const void *data, size_t len, uint32_t crc) { +#if defined(__SSE4_2__) + return crc32c_sse42(data, len, crc); +#elif defined(__ARM_FEATURE_CRC32) + return crc32c_neon(data, len, crc); +#else + return crc32c_slicing8(data, len, crc); +#endif +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/internal/cpu_features.cc b/src/ailego/internal/cpu_features.cc new file mode 100644 index 0000000..066b0a6 --- /dev/null +++ b/src/ailego/internal/cpu_features.cc @@ -0,0 +1,427 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "cpu_features.h" +#include + +#if defined(_MSC_VER) +#include +#endif + +#if (defined(__x86_64__) || defined(__i386__)) && !defined(_MSC_VER) +#include +#endif + +namespace zvec { +namespace ailego { +namespace internal { + +// +// REFER: https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/ +// tree/arch/x86/include/asm/cpufeatures.h +// https://software.intel.com/sites/default/files/managed/c5/15/ +// architecture-instruction-set-extensions-programming-reference.pdf +// + +CpuFeatures::CpuFlags CpuFeatures::flags_; + +#if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)) +CpuFeatures::CpuFlags::CpuFlags(void) + : L1_ECX(0), L1_EDX(0), L7_EBX(0), L7_ECX(0), L7_EDX(0) { + int l1[4] = {0, 0, 0, 0}; + int l7[4] = {0, 0, 0, 0}; + + __cpuidex(l1, 1, 0); + __cpuidex(l7, 7, 0); + L1_ECX = l1[2]; + L1_EDX = l1[3]; + L7_EBX = l7[1]; + L7_ECX = l7[2]; + L7_EDX = l7[3]; +} +#elif defined(__x86_64__) || defined(__i386__) +CpuFeatures::CpuFlags::CpuFlags(void) + : L1_ECX(0), L1_EDX(0), L7_EBX(0), L7_ECX(0), L7_EDX(0) { + uint32_t eax, ebx, ecx, edx; + + if (__get_cpuid(1, &eax, &ebx, &ecx, &edx)) { + L1_ECX = ecx; + L1_EDX = edx; + } + if (__get_cpuid_max(0, nullptr) >= 7) { + __cpuid_count(7, 0, eax, ebx, ecx, edx); + L7_EBX = ebx; + L7_ECX = ecx; + L7_EDX = edx; + } +} +#else +CpuFeatures::CpuFlags::CpuFlags(void) + : L1_ECX(0), L1_EDX(0), L7_EBX(0), L7_ECX(0), L7_EDX(0) {} +#endif + +//! 16-bit FP conversions +bool CpuFeatures::F16C(void) { + return !!(flags_.L1_ECX & (1u << 29)); +} + +//! Multimedia Extensions +bool CpuFeatures::MMX(void) { + return !!(flags_.L1_EDX & (1u << 23)); +} + +//! Streaming SIMD Extensions +bool CpuFeatures::SSE(void) { + return !!(flags_.L1_EDX & (1u << 25)); +} + +//! Streaming SIMD Extensions 2 +bool CpuFeatures::SSE2(void) { + return !!(flags_.L1_EDX & (1u << 26)); +} + +//! Streaming SIMD Extensions 3 +bool CpuFeatures::SSE3(void) { + return !!(flags_.L1_ECX & (1u << 0)); +} + +//! Supplemental Streaming SIMD Extensions 3 +bool CpuFeatures::SSSE3(void) { + return !!(flags_.L1_ECX & (1u << 9)); +} + +//! Streaming SIMD Extensions 4.1 +bool CpuFeatures::SSE4_1(void) { + return !!(flags_.L1_ECX & (1u << 19)); +} + +//! Streaming SIMD Extensions 4.2 +bool CpuFeatures::SSE4_2(void) { + return !!(flags_.L1_ECX & (1u << 20)); +} + +//! Advanced Vector Extensions +bool CpuFeatures::AVX(void) { + return !!(flags_.L1_ECX & (1u << 28)); +} + +//! Advanced Vector Extensions 2 +bool CpuFeatures::AVX2(void) { + return !!(flags_.L7_EBX & (1u << 5)); +} + +//! AVX-512 Foundation +bool CpuFeatures::AVX512F(void) { + return !!(flags_.L7_EBX & (1u << 16)); +} + +//! AVX-512 DQ (Double/Quad granular) Instructions +bool CpuFeatures::AVX512DQ(void) { + return !!(flags_.L7_EBX & (1u << 17)); +} + +//! AVX-512 Prefetch +bool CpuFeatures::AVX512PF(void) { + return !!(flags_.L7_EBX & (1u << 26)); +} + +//! AVX-512 Exponential and Reciprocal +bool CpuFeatures::AVX512ER(void) { + return !!(flags_.L7_EBX & (1u << 27)); +} + +//! AVX-512 Conflict Detection +bool CpuFeatures::AVX512CD(void) { + return !!(flags_.L7_EBX & (1u << 28)); +} + +//! AVX-512 BW (Byte/Word granular) Instructions +bool CpuFeatures::AVX512BW(void) { + return !!(flags_.L7_EBX & (1u << 30)); +} + +//! AVX-512 VL (128/256 Vector Length) Extensions +bool CpuFeatures::AVX512VL(void) { + return !!(flags_.L7_EBX & (1u << 31)); +} + +//! AVX-512 Integer Fused Multiply-Add instructions +bool CpuFeatures::AVX512_IFMA(void) { + return !!(flags_.L7_EBX & (1u << 21)); +} + +//! AVX512 Vector Bit Manipulation instructions +bool CpuFeatures::AVX512_VBMI(void) { + return !!(flags_.L7_ECX & (1u << 1)); +} + +//! Additional AVX512 Vector Bit Manipulation Instructions +bool CpuFeatures::AVX512_VBMI2(void) { + return !!(flags_.L7_ECX & (1u << 6)); +} + +//! Vector Neural Network Instructions +bool CpuFeatures::AVX512_VNNI(void) { + return !!(flags_.L7_ECX & (1u << 11)); +} + +//! Support for VPOPCNT[B,W] and VPSHUF-BITQMB instructions +bool CpuFeatures::AVX512_BITALG(void) { + return !!(flags_.L7_ECX & (1u << 12)); +} + +//! POPCNT for vectors of DW/QW +bool CpuFeatures::AVX512_VPOPCNTDQ(void) { + return !!(flags_.L7_ECX & (1u << 14)); +} + +//! AVX-512 Neural Network Instructions +bool CpuFeatures::AVX512_4VNNIW(void) { + return !!(flags_.L7_EDX & (1u << 2)); +} + +//! AVX-512 Multiply Accumulation Single precision +bool CpuFeatures::AVX512_4FMAPS(void) { + return !!(flags_.L7_EDX & (1u << 3)); +} + +//! AVX-512 FP16 instructions +bool CpuFeatures::AVX512_FP16(void) { + return !!(flags_.L7_EDX & (1u << 23)); +} + +//! CMPXCHG8 instruction +bool CpuFeatures::CX8(void) { + return !!(flags_.L1_EDX & (1u << 8)); +} + +//! CMPXCHG16B instruction +bool CpuFeatures::CX16(void) { + return !!(flags_.L1_ECX & (1u << 13)); +} + +//! PCLMULQDQ instruction +bool CpuFeatures::PCLMULQDQ(void) { + return !!(flags_.L1_ECX & (1u << 1)); +} + +//! Carry-Less Multiplication Double Quadword +bool CpuFeatures::VPCLMULQDQ(void) { + return !!(flags_.L7_ECX & (1u << 10)); +} + +//! CMOV instructions (plus FCMOVcc, FCOMI with FPU) +bool CpuFeatures::CMOV(void) { + return !!(flags_.L1_EDX & (1u << 15)); +} + +//! MOVBE instruction +bool CpuFeatures::MOVBE(void) { + return !!(flags_.L1_ECX & (1u << 22)); +} + +//! Enhanced REP MOVSB/STOSB instructions +bool CpuFeatures::ERMS(void) { + return !!(flags_.L7_EBX & (1u << 9)); +} + +//! POPCNT instruction +bool CpuFeatures::POPCNT(void) { + return !!(flags_.L1_ECX & (1u << 23)); +} + +//! XSAVE/XRSTOR/XSETBV/XGETBV instructions +bool CpuFeatures::XSAVE(void) { + return !!(flags_.L1_ECX & (1u << 26)); +} + +//! Fused multiply-add +bool CpuFeatures::FMA(void) { + return !!(flags_.L1_ECX & (1u << 12)); +} + +//! ADCX and ADOX instructions +bool CpuFeatures::ADX(void) { + return !!(flags_.L7_EBX & (1u << 19)); +} + +//! Galois Field New Instructions +bool CpuFeatures::GFNI(void) { + return !!(flags_.L7_ECX & (1u << 8)); +} + +//! AES instructions +bool CpuFeatures::AES(void) { + return !!(flags_.L1_ECX & (1u << 25)); +} + +//! Vector AES +bool CpuFeatures::VAES(void) { + return !!(flags_.L7_ECX & (1u << 9)); +} + +//! RDSEED instruction +bool CpuFeatures::RDSEED(void) { + return !!(flags_.L7_EBX & (1u << 18)); +} + +//! RDRAND instruction +bool CpuFeatures::RDRAND(void) { + return !!(flags_.L1_ECX & (1u << 30)); +} + +//! SHA1/SHA256 Instruction Extensions +bool CpuFeatures::SHA(void) { + return !!(flags_.L7_EBX & (1u << 29)); +} + +//! 1st group bit manipulation extensions +bool CpuFeatures::BMI1(void) { + return !!(flags_.L7_EBX & (1u << 3)); +} + +//! 2nd group bit manipulation extensions +bool CpuFeatures::BMI2(void) { + return !!(flags_.L7_EBX & (1u << 8)); +} + +//! CLFLUSH instruction +bool CpuFeatures::CLFLUSH(void) { + return !!(flags_.L1_EDX & (1u << 19)); +} + +//! CLFLUSHOPT instruction +bool CpuFeatures::CLFLUSHOPT(void) { + return !!(flags_.L7_EBX & (1u << 23)); +} + +//! CLWB instruction +bool CpuFeatures::CLWB(void) { + return !!(flags_.L7_EBX & (1u << 24)); +} + +//! RDPID instruction +bool CpuFeatures::RDPID(void) { + return !!(flags_.L7_ECX & (1u << 22)); +} + +//! Onboard FPU +bool CpuFeatures::FPU(void) { + return !!(flags_.L1_EDX & (1u << 0)); +} + +//! Hyper-Threading +bool CpuFeatures::HT(void) { + return !!(flags_.L1_EDX & (1u << 28)); +} + +//! Hardware virtualization +bool CpuFeatures::VMX(void) { + return !!(flags_.L1_ECX & (1u << 5)); +} + +// ！Running on a hypervisor +bool CpuFeatures::HYPERVISOR(void) { + return !!(flags_.L1_ECX & (1u << 31)); +} + +const char *CpuFeatures::Intrinsics(void) { + return "" +#if defined(__ARM_NEON) + "Neon" +#if defined(__ARM_FEATURE_CRC32) + "+CRC" +#endif +#if defined(__ARM_FEATURE_FP16_SCALAR_ARITHMETIC) || \ + defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) + "+FP16" +#endif +#elif defined(__AVX512F__) + "AVX512F" +#if defined(__AVX512VL__) + "+AVX512VL" +#endif +#if defined(__AVX512BW__) + "+AVX512BW" +#endif +#if defined(__AVX512DQ__) + "+AVX512DQ" +#endif +#if defined(__AVX512CD__) + "+AVX512CD" +#endif +#if defined(__AVX512ER__) + "+AVX512ER" +#endif +#if defined(__AVX512PF__) + "+AVX512PF" +#endif +#if defined(__AVX512IFMA__) + "+AVX512IFMA" +#endif +#if defined(__AVX512VBMI__) + "+AVX512VBMI" +#endif +#if defined(__AVX512VBMI2__) + "+AVX512VBMI2" +#endif +#if defined(__AVX512VNNI__) + "+AVX512VNNI" +#endif +#if defined(__AVX512BITALG__) + "+AVX512BITALG" +#endif +#if defined(__AVX512VPOPCNTDQ__) + "+AVX512VPOPCNTDQ" +#endif +#if defined(__AVX512FP16__) + "+AVX512FP16" +#endif +#elif defined(__AVX2__) + "AVX2" +#elif defined(__AVX__) + "AVX" +#elif defined(__SSE4_2__) + "SSE4.2" +#elif defined(__SSE4_1__) + "SSE4.1" +#elif defined(__SSSE3__) + "SSSE3" +#elif defined(__SSE3__) + "SSE3" +#elif defined(__SSE2__) + "SSE2" +#elif defined(__SSE__) + "SSE" +#elif defined(__MMX__) + "MMX" +#endif +#if defined(__FMA__) + "+FMA" +#endif +#if defined(__BMI2__) + "+BMI2" +#elif defined(__BMI__) + "+BMI" +#endif +#if defined(__F16C__) + "+F16C" +#endif + ; +} + +CpuFeatures::StaticFlags CpuFeatures::static_flags_; +} // namespace internal +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/internal/cpu_features.h b/src/ailego/internal/cpu_features.h new file mode 100644 index 0000000..3db1dee --- /dev/null +++ b/src/ailego/internal/cpu_features.h @@ -0,0 +1,370 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +namespace zvec { +namespace ailego { +namespace internal { + +/*! Cpu Features + */ +class CpuFeatures { + public: + //! 16-bit FP conversions + static bool F16C(void); + + //! Multimedia Extensions + static bool MMX(void); + + //! Streaming SIMD Extensions + static bool SSE(void); + + //! Streaming SIMD Extensions 2 + static bool SSE2(void); + + //! Streaming SIMD Extensions 3 + static bool SSE3(void); + + //! Supplemental Streaming SIMD Extensions 3 + static bool SSSE3(void); + + //! Streaming SIMD Extensions 4.1 + static bool SSE4_1(void); + + //! Streaming SIMD Extensions 4.2 + static bool SSE4_2(void); + + //! Advanced Vector Extensions + static bool AVX(void); + + //! Advanced Vector Extensions 2 + static bool AVX2(void); + + //! AVX-512 Foundation + static bool AVX512F(void); + + //! AVX-512 DQ (Double/Quad granular) Instructions + static bool AVX512DQ(void); + + //! AVX-512 Prefetch + static bool AVX512PF(void); + + //! AVX-512 Exponential and Reciprocal + static bool AVX512ER(void); + + //! AVX-512 Conflict Detection + static bool AVX512CD(void); + + //! AVX-512 BW (Byte/Word granular) Instructions + static bool AVX512BW(void); + + //! AVX-512 VL (128/256 Vector Length) Extensions + static bool AVX512VL(void); + + //! AVX-512 Integer Fused Multiply-Add instructions + static bool AVX512_IFMA(void); + + //! AVX512 Vector Bit Manipulation instructions + static bool AVX512_VBMI(void); + + //! Additional AVX512 Vector Bit Manipulation Instructions + static bool AVX512_VBMI2(void); + + //! Vector Neural Network Instructions + static bool AVX512_VNNI(void); + + //! Support for VPOPCNT[B,W] and VPSHUF-BITQMB instructions + static bool AVX512_BITALG(void); + + //! POPCNT for vectors of DW/QW + static bool AVX512_VPOPCNTDQ(void); + + //! AVX-512 Neural Network Instructions + static bool AVX512_4VNNIW(void); + + //! AVX-512 Multiply Accumulation Single precision + static bool AVX512_4FMAPS(void); + + //! AVX-512 FP16 instructions + static bool AVX512_FP16(void); + + //! CMPXCHG8 instruction + static bool CX8(void); + + //! CMPXCHG16B instruction + static bool CX16(void); + + //! PCLMULQDQ instruction + static bool PCLMULQDQ(void); + + //! Carry-Less Multiplication Double Quadword + static bool VPCLMULQDQ(void); + + //! CMOV instructions (plus FCMOVcc, FCOMI with FPU) + static bool CMOV(void); + + //! MOVBE instruction + static bool MOVBE(void); + + //! Enhanced REP MOVSB/STOSB instructions + static bool ERMS(void); + + //! POPCNT instruction + static bool POPCNT(void); + + //! XSAVE/XRSTOR/XSETBV/XGETBV instructions + static bool XSAVE(void); + + //! Fused multiply-add + static bool FMA(void); + + //! ADCX and ADOX instructions + static bool ADX(void); + + //! Galois Field New Instructions + static bool GFNI(void); + + //! AES instructions + static bool AES(void); + + //! Vector AES + static bool VAES(void); + + //! RDSEED instruction + static bool RDSEED(void); + + //! RDRAND instruction + static bool RDRAND(void); + + //! SHA1/SHA256 Instruction Extensions + static bool SHA(void); + + //! 1st group bit manipulation extensions + static bool BMI1(void); + + //! 2nd group bit manipulation extensions + static bool BMI2(void); + + //! CLFLUSH instruction + static bool CLFLUSH(void); + + //! CLFLUSHOPT instruction + static bool CLFLUSHOPT(void); + + //! CLWB instruction + static bool CLWB(void); + + //! RDPID instruction + static bool RDPID(void); + + //! Onboard FPU + static bool FPU(void); + + //! Hyper-Threading + static bool HT(void); + + //! Hardware virtualization + static bool VMX(void); + + // ！Running on a hypervisor + static bool HYPERVISOR(void); + + //! Intrinsics of compiling + static const char *Intrinsics(void); + + private: + struct CpuFlags { + //! Constructor + CpuFlags(void); + + //! Members + uint32_t L1_ECX; + uint32_t L1_EDX; + uint32_t L7_EBX; + uint32_t L7_ECX; + uint32_t L7_EDX; + }; + + //! Static Members + static CpuFlags flags_; + + public: + struct StaticFlags { + //! 16-bit FP conversions + bool F16C = CpuFeatures::F16C(); + + //! Multimedia Extensions + bool MMX = CpuFeatures::MMX(); + + //! Streaming SIMD Extensions + bool SSE = CpuFeatures::SSE(); + + //! Streaming SIMD Extensions 2 + bool SSE2 = CpuFeatures::SSE2(); + + //! Streaming SIMD Extensions 3 + bool SSE3 = CpuFeatures::SSE3(); + + //! Supplemental Streaming SIMD Extensions 3 + bool SSSE3 = CpuFeatures::SSSE3(); + + //! Streaming SIMD Extensions 4.1 + bool SSE4_1 = CpuFeatures::SSE4_1(); + + //! Streaming SIMD Extensions 4.2 + bool SSE4_2 = CpuFeatures::SSE4_2(); + + //! Advanced Vector Extensions + bool AVX = CpuFeatures::AVX(); + + //! Advanced Vector Extensions 2 + bool AVX2 = CpuFeatures::AVX2(); + + //! AVX-512 Foundation + bool AVX512F = CpuFeatures::AVX512F(); + + //! AVX-512 DQ (Double/Quad granular) Instructions + bool AVX512DQ = CpuFeatures::AVX512DQ(); + + //! AVX-512 Prefetch + bool AVX512PF = CpuFeatures::AVX512PF(); + + //! AVX-512 Exponential and Reciprocal + bool AVX512ER = CpuFeatures::AVX512ER(); + + //! AVX-512 Conflict Detection + bool AVX512CD = CpuFeatures::AVX512CD(); + + //! AVX-512 BW (Byte/Word granular) Instructions + bool AVX512BW = CpuFeatures::AVX512BW(); + + //! AVX-512 VL (128/256 Vector Length) Extensions + bool AVX512VL = CpuFeatures::AVX512VL(); + + //! AVX-512 Integer Fused Multiply-Add instructions + bool AVX512_IFMA = CpuFeatures::AVX512_IFMA(); + + //! AVX512 Vector Bit Manipulation instructions + bool AVX512_VBMI = CpuFeatures::AVX512_VBMI(); + + //! Additional AVX512 Vector Bit Manipulation Instructions + bool AVX512_VBMI2 = CpuFeatures::AVX512_VBMI2(); + + //! Vector Neural Network Instructions + bool AVX512_VNNI = CpuFeatures::AVX512_VNNI(); + + //! Support for VPOPCNT[B,W] and VPSHUF-BITQMB instructions + bool AVX512_BITALG = CpuFeatures::AVX512_BITALG(); + + //! POPCNT for vectors of DW/QW + bool AVX512_VPOPCNTDQ = CpuFeatures::AVX512_VPOPCNTDQ(); + + //! AVX-512 Neural Network Instructions + bool AVX512_4VNNIW = CpuFeatures::AVX512_4VNNIW(); + + //! AVX-512 Multiply Accumulation Single precision + bool AVX512_4FMAPS = CpuFeatures::AVX512_4FMAPS(); + + //! AVX-512 FP16 instructions + bool AVX512_FP16 = CpuFeatures::AVX512_FP16(); + + //! CMPXCHG8 instruction + bool CX8 = CpuFeatures::CX8(); + + //! CMPXCHG16B instruction + bool CX16 = CpuFeatures::CX16(); + + //! PCLMULQDQ instruction + bool PCLMULQDQ = CpuFeatures::PCLMULQDQ(); + + //! Carry-Less Multiplication Double Quadword + bool VPCLMULQDQ = CpuFeatures::VPCLMULQDQ(); + + //! CMOV instructions (plus FCMOVcc, FCOMI with FPU) + bool CMOV = CpuFeatures::CMOV(); + + //! MOVBE instruction + bool MOVBE = CpuFeatures::MOVBE(); + + //! Enhanced REP MOVSB/STOSB instructions + bool ERMS = CpuFeatures::ERMS(); + + //! POPCNT instruction + bool POPCNT = CpuFeatures::POPCNT(); + + //! XSAVE/XRSTOR/XSETBV/XGETBV instructions + bool XSAVE = CpuFeatures::XSAVE(); + + //! Fused multiply-add + bool FMA = CpuFeatures::FMA(); + + //! ADCX and ADOX instructions + bool ADX = CpuFeatures::ADX(); + + //! Galois Field New Instructions + bool GFNI = CpuFeatures::GFNI(); + + //! AES instructions + bool AES = CpuFeatures::AES(); + + //! Vector AES + bool VAES = CpuFeatures::VAES(); + + //! RDSEED instruction + bool RDSEED = CpuFeatures::RDSEED(); + + //! RDRAND instruction + bool RDRAND = CpuFeatures::RDRAND(); + + //! SHA1/SHA256 Instruction Extensions + bool SHA = CpuFeatures::SHA(); + + //! 1st group bit manipulation extensions + bool BMI1 = CpuFeatures::BMI1(); + + //! 2nd group bit manipulation extensions + bool BMI2 = CpuFeatures::BMI2(); + + //! CLFLUSH instruction + bool CLFLUSH = CpuFeatures::CLFLUSH(); + + //! CLFLUSHOPT instruction + bool CLFLUSHOPT = CpuFeatures::CLFLUSHOPT(); + + //! CLWB instruction + bool CLWB = CpuFeatures::CLWB(); + + //! RDPID instruction + bool RDPID = CpuFeatures::RDPID(); + + //! Onboard FPU + bool FPU = CpuFeatures::FPU(); + + //! Hyper-Threading + bool HT = CpuFeatures::HT(); + + //! Hardware virtualization + bool VMX = CpuFeatures::VMX(); + + // ！Running on a hypervisor + bool HYPERVISOR = CpuFeatures::HYPERVISOR(); + }; + static StaticFlags static_flags_; +}; + +} // namespace internal +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/io/file.cc b/src/ailego/io/file.cc new file mode 100644 index 0000000..9b23d68 --- /dev/null +++ b/src/ailego/io/file.cc @@ -0,0 +1,780 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#if !defined(_WIN64) && !defined(_WIN32) +#include +#include +#include +#include +#include +#include +#include +#else +#include +#include +#include +#endif + +namespace zvec { +namespace ailego { + +#if !defined(_WIN64) && !defined(_WIN32) + +static inline int OpenSafely(const char *path, int flags) { + int fd = open(path, flags, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH); + while (fd == -1 && errno == EINTR) { + fd = open(path, flags, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH); + } + return fd; +} + +static inline void CloseSafely(int fd) { + int ret = close(fd); + while (ret == -1 && errno == EINTR) { + ret = close(fd); + } +} + +static inline ssize_t ReadSafely(int fd, void *buf, size_t count) { + ssize_t ret = read(fd, buf, count); + while (ret == -1 && errno == EINTR) { + ret = read(fd, buf, count); + } + return ret; +} + +static inline ssize_t PreadSafely(int fd, void *buf, size_t count, + ssize_t offset) { + ssize_t ret = pread(fd, buf, count, offset); + while (ret == -1 && errno == EINTR) { + ret = pread(fd, buf, count, offset); + } + return ret; +} + +static inline ssize_t WriteSafely(int fd, const void *buf, size_t count) { + ssize_t ret = write(fd, buf, count); + while (ret == -1 && errno == EINTR) { + ret = write(fd, buf, count); + } + return ret; +} + +static inline ssize_t PwriteSafely(int fd, const void *buf, size_t count, + ssize_t offset) { + ssize_t ret = pwrite(fd, buf, count, offset); + while (ret == -1 && errno == EINTR) { + ret = pwrite(fd, buf, count, offset); + } + return ret; +} + +static inline size_t ReadAll(int fd, void *buf, size_t count) { + size_t rdlen = 0; + while (rdlen < count) { + ssize_t ret = ReadSafely(fd, (char *)buf + rdlen, count - rdlen); + if (ret <= 0) { + break; + } + rdlen += ret; + } + return rdlen; +} + +static inline size_t PreadAll(int fd, void *buf, size_t count, ssize_t offset) { + size_t rdlen = 0; + while (rdlen < count) { + ssize_t ret = + PreadSafely(fd, (char *)buf + rdlen, count - rdlen, offset + rdlen); + if (ret <= 0) { + break; + } + rdlen += ret; + } + return rdlen; +} + +static inline size_t WriteAll(int fd, const void *buf, size_t count) { + size_t wrlen = 0; + while (wrlen < count) { + ssize_t ret = WriteSafely(fd, (const char *)buf + wrlen, count - wrlen); + if (ret <= 0) { + break; + } + wrlen += ret; + } + return wrlen; +} + +static inline size_t PwriteAll(int fd, const void *buf, size_t count, + ssize_t offset) { + size_t wrlen = 0; + while (wrlen < count) { + ssize_t ret = PwriteSafely(fd, (const char *)buf + wrlen, count - wrlen, + offset + wrlen); + if (ret <= 0) { + break; + } + wrlen += ret; + } + return wrlen; +} + +bool File::create(const char *path, size_t len, bool direct) { + ailego_false_if_false(native_handle_ == File::InvalidHandle && path); + + // Try opening or creating a file + int flags = O_RDWR | O_CREAT; +#ifdef O_DIRECT + if (direct) { + flags |= O_DIRECT; + } +#else + (void)direct; +#endif + + int fd = OpenSafely(path, flags); + ailego_false_if_lt_zero(fd); + +#ifdef F_NOCACHE + // Direct IO canonical solution for Mac OSX + if (direct) { + ailego_false_if_ne_zero(fcntl(fd, F_NOCACHE, 1)); + } +#endif + + // Truncate the file to the specified size + ailego_do_if_ne_zero(ftruncate(fd, len)) { + CloseSafely(fd); + return false; + } + + read_only_ = false; + native_handle_ = fd; + return true; +} + +bool File::open(const char *path, bool rdonly, bool direct) { + ailego_false_if_false(native_handle_ == File::InvalidHandle && path); + + // Try opening the file + int flags = rdonly ? O_RDONLY : O_RDWR; +#ifdef O_DIRECT + if (direct) { + flags |= O_DIRECT; + } +#else + (void)direct; +#endif + + int fd = OpenSafely(path, flags); + ailego_false_if_lt_zero(fd); + +#ifdef F_NOCACHE + // Direct IO canonical solution for Mac OSX + if (direct) { + ailego_false_if_ne_zero(fcntl(fd, F_NOCACHE, 1)); + } +#endif + + read_only_ = rdonly; + native_handle_ = fd; + return true; +} + +void File::close(void) { + ailego_return_if_false(native_handle_ != File::InvalidHandle); + CloseSafely(native_handle_); + native_handle_ = File::InvalidHandle; +} + +void File::reset(void) { + ailego_return_if_false(native_handle_ != File::InvalidHandle); + lseek(native_handle_, 0, SEEK_SET); +} + +size_t File::write(const void *data, size_t len) { + const size_t block_size = 0x40000000u; + size_t total = 0u; + + for (; len >= block_size; len -= block_size) { + size_t wrlen = + WriteAll(native_handle_, (const uint8_t *)data + total, block_size); + if (wrlen != block_size) { + return (total + wrlen); + } + total += block_size; + } + if (len > 0) { + total += WriteAll(native_handle_, (const uint8_t *)data + total, len); + } + return total; +} + +size_t File::write(ssize_t off, const void *data, size_t len) { + const size_t block_size = 0x40000000u; + size_t total = 0u; + + for (; len >= block_size; len -= block_size) { + size_t wrlen = PwriteAll(native_handle_, (const uint8_t *)data + total, + block_size, off + total); + if (wrlen != block_size) { + return (total + wrlen); + } + total += block_size; + } + if (len > 0) { + total += PwriteAll(native_handle_, (const uint8_t *)data + total, len, + off + total); + } + return total; +} + +size_t File::read(void *buf, size_t len) { + const size_t block_size = 0x40000000u; + size_t total = 0u; + + for (; len >= block_size; len -= block_size) { + size_t rdlen = ReadAll(native_handle_, (uint8_t *)buf + total, block_size); + if (rdlen != block_size) { + return (total + rdlen); + } + total += block_size; + } + if (len > 0) { + total += ReadAll(native_handle_, (uint8_t *)buf + total, len); + } + return total; +} + +size_t File::read(ssize_t off, void *buf, size_t len) { + const size_t block_size = 0x40000000u; + size_t total = 0u; + + for (; len >= block_size; len -= block_size) { + size_t rdlen = PreadAll(native_handle_, (uint8_t *)buf + total, block_size, + off + total); + if (rdlen != block_size) { + return (total + rdlen); + } + total += block_size; + } + if (len > 0) { + total += PreadAll(native_handle_, (uint8_t *)buf + total, len, off + total); + } + return total; +} + +bool File::flush(void) { + ailego_false_if_false(native_handle_ != File::InvalidHandle); + return (fsync(native_handle_) == 0); +} + +bool File::seek(ssize_t off, Origin origin) { + ailego_false_if_false(native_handle_ != File::InvalidHandle); + ailego_false_if_false(lseek(native_handle_, off, (int)origin) != (off_t)-1); + return true; +} + +bool File::truncate(size_t len) { + ailego_false_if_false(native_handle_ != File::InvalidHandle); + ailego_false_if_ne_zero(ftruncate(native_handle_, (off_t)len)); + return true; +} + +size_t File::size(void) const { + struct stat fs; + ailego_zero_if_false(native_handle_ != File::InvalidHandle && + fstat(native_handle_, &fs) == 0); + return (fs.st_size); +} + +ssize_t File::offset(void) const { + off_t off; + ailego_zero_if_false(native_handle_ != File::InvalidHandle && + (off = lseek(native_handle_, 0, SEEK_CUR)) != -1); + return off; +} + +void *File::MemoryMap(NativeHandle handle, ssize_t off, size_t len, int opts) { + int prot = + ((opts & File::MMAP_READONLY) ? PROT_READ : PROT_READ | PROT_WRITE); + int flags = (opts & File::MMAP_SHARED) ? MAP_SHARED : MAP_PRIVATE; + +#if defined(MAP_POPULATE) + if (opts & File::MMAP_POPULATE) { + flags |= MAP_POPULATE; + } +#endif + +#if defined(MAP_HUGETLB) + if (opts & File::MMAP_HUGE_PAGE) { + flags |= MAP_HUGETLB; + } +#endif + + void *addr = mmap(nullptr, len, prot, flags, handle, off); + ailego_null_if_false(addr != MAP_FAILED); + + if (opts & File::MMAP_LOCKED) { + mlock(addr, len); + } + if (opts & File::MMAP_WARMUP) { + File::MemoryWarmup(addr, len); + } + return addr; +} + +#if !defined(MAP_ANONYMOUS) && defined(MAP_ANON) +#define MAP_ANONYMOUS MAP_ANON +#endif + +void *File::MemoryMap(size_t len, int opts) { +#if defined(MAP_ANONYMOUS) + int prot = + ((opts & File::MMAP_READONLY) ? PROT_READ : PROT_READ | PROT_WRITE); + int flags = (opts & File::MMAP_SHARED) ? MAP_SHARED | MAP_ANONYMOUS + : MAP_PRIVATE | MAP_ANONYMOUS; +#if defined(MAP_POPULATE) + if (opts & File::MMAP_POPULATE) { + flags |= MAP_POPULATE; + } +#endif +#if defined(MAP_HUGETLB) + if (opts & File::MMAP_HUGE_PAGE) { + flags |= MAP_HUGETLB; + } +#endif + void *addr = mmap(nullptr, len, prot, flags, -1, 0); + ailego_null_if_false(addr != MAP_FAILED); + return addr; +#else + (void)len; + (void)opts; + return nullptr; +#endif // MAP_ANONYMOUS +} + +void *File::MemoryRemap(void *oldptr, size_t oldsize, void *newptr, + size_t newsize) { +#if defined(__linux) || defined(__linux__) + return newptr ? mremap(oldptr, oldsize, newsize, MREMAP_FIXED, newptr) + : mremap(oldptr, oldsize, newsize, MREMAP_MAYMOVE); +#elif defined(__NetBSD__) + return newptr ? mremap(oldptr, oldsize, newptr, newsize, MAP_FIXED) + : mremap(oldptr, oldsize, nullptr, newsize, 0); +#else + (void)oldptr; + (void)oldsize; + (void)newptr; + (void)newsize; + errno = ENOTSUP; + return nullptr; +#endif +} + +void File::MemoryUnmap(void *addr, size_t len) { + ailego_return_if_false(addr); + munmap(addr, len); +} + +bool File::MemoryFlush(void *addr, size_t len) { + ailego_false_if_false(addr); + return (msync(addr, len, MS_ASYNC) == 0); +} + +bool File::MemoryLock(void *addr, size_t len) { + ailego_false_if_false(addr && len); + return (mlock(addr, len) == 0); +} + +bool File::MemoryUnlock(void *addr, size_t len) { + ailego_false_if_false(addr && len); + return (munlock(addr, len) == 0); +} + +#else + +namespace { + +bool Utf8PathOk(const char *path, const std::wstring &wide) { + return path && path[0] != '\0' && !wide.empty(); +} + +} // namespace + +//! Create a local file +bool File::create(const char *path, size_t len, bool direct) { + ailego_false_if_false(native_handle_ == File::InvalidHandle && path); + + const std::wstring wpath = FileHelper::Utf8ToWide(path); + ailego_false_if_false(Utf8PathOk(path, wpath)); + + // Try opening or creating the file + HANDLE file_handle = + ::CreateFileW(wpath.c_str(), GENERIC_WRITE | GENERIC_READ, + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, + nullptr, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr); + ailego_false_if_false(file_handle != INVALID_HANDLE_VALUE); + + // Truncate the file to the specified size + LARGE_INTEGER file_size; + file_size.QuadPart = len; + ailego_do_if_false( + SetFilePointerEx(file_handle, file_size, nullptr, FILE_BEGIN) && + SetEndOfFile(file_handle)) { + CloseHandle(file_handle); + return false; + } + + if (!direct) { + // Reset the file pointer + SetFilePointer(file_handle, 0, nullptr, FILE_BEGIN); + } else { + // Close and reopen file + CloseHandle(file_handle); + file_handle = ::CreateFileW( + wpath.c_str(), GENERIC_WRITE | GENERIC_READ, + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, nullptr, + OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING, nullptr); + ailego_false_if_false(file_handle != INVALID_HANDLE_VALUE); + } + + read_only_ = false; + native_handle_ = file_handle; + return true; +} + +//! Open a local file +bool File::open(const char *path, bool rdonly, bool direct) { + ailego_false_if_false(native_handle_ == File::InvalidHandle && path); + + const std::wstring wpath = FileHelper::Utf8ToWide(path); + ailego_false_if_false(Utf8PathOk(path, wpath)); + + // Try opening the file + DWORD flags = FILE_ATTRIBUTE_NORMAL; + if (direct) { + flags |= FILE_FLAG_NO_BUFFERING; + } + HANDLE file_handle = ::CreateFileW( + wpath.c_str(), (rdonly ? GENERIC_READ : GENERIC_READ | GENERIC_WRITE), + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, nullptr, + OPEN_EXISTING, flags, nullptr); + ailego_false_if_false(file_handle != INVALID_HANDLE_VALUE); + + read_only_ = rdonly; + native_handle_ = file_handle; + return true; +} + +void File::close(void) { + ailego_return_if_false(native_handle_ != File::InvalidHandle); + CloseHandle(native_handle_); + native_handle_ = File::InvalidHandle; +} + +void File::reset(void) { + ailego_return_if_false(native_handle_ != File::InvalidHandle); + SetFilePointer(native_handle_, 0, nullptr, FILE_BEGIN); +} + +size_t File::write(const void *data, size_t len) { + const DWORD block_size = 0x40000000u; + DWORD wrlen = 0u; + size_t total = 0u; + + for (; len >= block_size; len -= block_size) { + if (!WriteFile(native_handle_, (const uint8_t *)data + total, block_size, + &wrlen, nullptr)) { + return total; + } + if (wrlen != block_size) { + return (total + wrlen); + } + total += block_size; + } + if (len > 0 && WriteFile(native_handle_, (const uint8_t *)data + total, + (DWORD)len, &wrlen, nullptr)) { + total += wrlen; + } + return total; +} + +size_t File::write(ssize_t off, const void *data, size_t len) { + const DWORD block_size = 0x40000000u; + DWORD wrlen = 0u; + size_t total = 0u; + + OVERLAPPED overlapped; + memset(&overlapped, 0, sizeof(OVERLAPPED)); + + for (; len >= block_size; len -= block_size) { + uint64_t current = off + total; + overlapped.OffsetHigh = (DWORD)(current >> 32); + overlapped.Offset = (DWORD)(current & 0xffffffffu); + + if (!WriteFile(native_handle_, (const uint8_t *)data + total, block_size, + &wrlen, &overlapped)) { + return total; + } + if (wrlen != block_size) { + return (total + wrlen); + } + total += block_size; + } + if (len > 0) { + uint64_t current = off + total; + overlapped.OffsetHigh = (DWORD)(current >> 32); + overlapped.Offset = (DWORD)(current & 0xffffffffu); + + if (WriteFile(native_handle_, (const uint8_t *)data + total, (DWORD)len, + &wrlen, &overlapped)) { + total += wrlen; + } + } + return total; +} + +size_t File::read(void *buf, size_t len) { + const DWORD block_size = 0x40000000u; + DWORD rdlen = 0u; + size_t total = 0u; + + for (; len >= block_size; len -= block_size) { + if (!ReadFile(native_handle_, (uint8_t *)buf + total, block_size, &rdlen, + nullptr)) { + return total; + } + if (rdlen != block_size) { + return (total + rdlen); + } + total += block_size; + } + if (len > 0 && ReadFile(native_handle_, (uint8_t *)buf + total, (DWORD)len, + &rdlen, nullptr)) { + total += rdlen; + } + return total; +} + +size_t File::read(ssize_t off, void *buf, size_t len) { + const DWORD block_size = 0x40000000u; + DWORD rdlen = 0u; + size_t total = 0u; + + OVERLAPPED overlapped; + memset(&overlapped, 0, sizeof(OVERLAPPED)); + + for (; len >= block_size; len -= block_size) { + uint64_t current = off + total; + overlapped.OffsetHigh = (DWORD)(current >> 32); + overlapped.Offset = (DWORD)(current & 0xffffffffu); + + if (!ReadFile(native_handle_, (uint8_t *)buf + total, block_size, &rdlen, + &overlapped)) { + return total; + } + if (rdlen != block_size) { + return (total + rdlen); + } + total += block_size; + } + if (len > 0) { + uint64_t current = off + total; + overlapped.OffsetHigh = (DWORD)(current >> 32); + overlapped.Offset = (DWORD)(current & 0xffffffffu); + + if (ReadFile(native_handle_, (uint8_t *)buf + total, (DWORD)len, &rdlen, + &overlapped)) { + total += rdlen; + } + } + return total; +} + +bool File::flush(void) { + ailego_false_if_false(native_handle_ != File::InvalidHandle); + return (!!FlushFileBuffers(native_handle_)); +} + +bool File::seek(ssize_t off, Origin origin) { + ailego_false_if_false(native_handle_ != File::InvalidHandle); + + LARGE_INTEGER file_offset; + file_offset.QuadPart = off; + ailego_false_if_false(SetFilePointerEx(native_handle_, file_offset, nullptr, + (DWORD)origin) != 0); + return true; +} + +bool File::truncate(size_t len) { + ailego_false_if_false(native_handle_ != File::InvalidHandle); + + LARGE_INTEGER file_size, orig_file_size; + file_size.QuadPart = 0; + orig_file_size.QuadPart = 0; + ailego_false_if_false(SetFilePointerEx(native_handle_, file_size, + &orig_file_size, FILE_CURRENT)); + + // Truncate the file to the specified size + file_size.QuadPart = len; + ailego_false_if_false( + SetFilePointerEx(native_handle_, file_size, nullptr, FILE_BEGIN) && + SetEndOfFile(native_handle_)); + + // Reset the file pointer + SetFilePointerEx(native_handle_, orig_file_size, nullptr, FILE_BEGIN); + return true; +} + +size_t File::size(void) const { + LARGE_INTEGER file_size; + ailego_zero_if_false(native_handle_ != File::InvalidHandle && + GetFileSizeEx(native_handle_, &file_size)); + return (size_t)file_size.QuadPart; +} + +ssize_t File::offset(void) const { + LARGE_INTEGER file_size; + LARGE_INTEGER file_size_new; + file_size.QuadPart = 0; + ailego_zero_if_false(native_handle_ != File::InvalidHandle && + SetFilePointerEx(native_handle_, file_size, + &file_size_new, FILE_CURRENT)); + return (size_t)file_size_new.QuadPart; +} + +void *File::MemoryMap(NativeHandle handle, ssize_t off, size_t len, int opts) { + // Root cause: Windows MapViewOfFile requires the file offset to be aligned to + // the allocation granularity (64 KB), but segment offsets were only + // page-aligned (4 KB). Also, CreateFileMapping was using len instead of + // off + len as the max size. + // + // Fix: Align the view offset down to allocation granularity, adjust the map + // length, and return base + excess. MemoryUnmap recovers the base by rounding + // down to granularity. + + SYSTEM_INFO si; + GetSystemInfo(&si); + DWORD granularity = si.dwAllocationGranularity; + ssize_t aligned_off = (off / (ssize_t)granularity) * (ssize_t)granularity; + size_t excess = (size_t)(off - aligned_off); + + LARGE_INTEGER max_size; + max_size.QuadPart = off + len; + + HANDLE file_mapping = CreateFileMapping( + handle, nullptr, + ((opts & File::MMAP_READONLY) ? PAGE_READONLY : PAGE_READWRITE), + max_size.HighPart, max_size.LowPart, nullptr); + ailego_null_if_false(file_mapping != nullptr); + + DWORD desired_access = FILE_MAP_READ; + if (!(opts & File::MMAP_READONLY)) { + desired_access |= FILE_MAP_WRITE; + } + if (!(opts & File::MMAP_SHARED)) { + desired_access |= FILE_MAP_COPY; + } + + LARGE_INTEGER view_offset; + view_offset.QuadPart = aligned_off; + size_t view_len = len + excess; + + void *base = MapViewOfFile(file_mapping, desired_access, view_offset.HighPart, + view_offset.LowPart, view_len); + CloseHandle(file_mapping); + + ailego_null_if_false(base); + void *addr = (char *)base + excess; + if (opts & File::MMAP_LOCKED) { + VirtualLock(addr, len); + } + if (opts & File::MMAP_WARMUP) { + File::MemoryWarmup(addr, len); + } + return addr; +} + +void *File::MemoryMap(size_t len, int opts) { + void *addr = + VirtualAlloc(nullptr, len, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + ailego_null_if_false(addr); + if (opts & File::MMAP_LOCKED) { + VirtualLock(addr, len); + } + if (opts & File::MMAP_WARMUP) { + File::MemoryWarmup(addr, len); + } + return addr; +} + +void *File::MemoryRemap(void *, size_t, void *, size_t) { + return nullptr; +} + +void File::MemoryUnmap(void *addr, size_t /*len*/) { + ailego_return_if_false(addr); + MEMORY_BASIC_INFORMATION mbi; + if (VirtualQuery(addr, &mbi, sizeof(mbi))) { + if (mbi.Type == MEM_MAPPED) { + UnmapViewOfFile(mbi.AllocationBase); + } else { + VirtualFree(mbi.AllocationBase, 0, MEM_RELEASE); + } + } +} + +bool File::MemoryFlush(void *addr, size_t /*len*/) { + ailego_false_if_false(addr); + return (!!FlushViewOfFile(addr, 0)); +} + +bool File::MemoryLock(void *addr, size_t len) { + ailego_false_if_false(addr && len); + return (!!VirtualLock(addr, len)); +} + +bool File::MemoryUnlock(void *addr, size_t len) { + ailego_false_if_false(addr && len); + return (!!VirtualUnlock(addr, len)); +} + +static inline int getpagesize(void) { + SYSTEM_INFO info; + GetSystemInfo(&info); + return info.dwPageSize; +} +#endif + +void File::MemoryWarmup(void *addr, size_t len) { + static int page_size = getpagesize(); + + if (addr && len) { + uint8_t *p = reinterpret_cast(addr); + uint8_t *end = p + len; + volatile uint8_t tmp = 0; + + while (p < end) { + tmp ^= *p; + p += page_size; + } + } +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/io/file_lock.cc b/src/ailego/io/file_lock.cc new file mode 100644 index 0000000..50a62c4 --- /dev/null +++ b/src/ailego/io/file_lock.cc @@ -0,0 +1,80 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "file_lock.h" + +#if !defined(_WIN64) && !defined(_WIN32) +#include +#else +#include +#endif + +namespace zvec { +namespace ailego { + +#if !defined(_WIN64) && !defined(_WIN32) +bool FileLock::Lock(int fd) { + return (flock(fd, LOCK_EX) == 0); +} + +bool FileLock::TryLock(int fd) { + return (flock(fd, LOCK_EX | LOCK_NB) == 0); +} + +bool FileLock::LockShared(int fd) { + return (flock(fd, LOCK_SH) == 0); +} + +bool FileLock::TryLockShared(int fd) { + return (flock(fd, LOCK_SH | LOCK_NB) == 0); +} + +bool FileLock::Unlock(int fd) { + return (flock(fd, LOCK_UN) == 0); +} + +#else +bool FileLock::Lock(HANDLE handle) { + OVERLAPPED ol = {0}; + return (!!LockFileEx(handle, LOCKFILE_EXCLUSIVE_LOCK, 0, MAXDWORD, MAXDWORD, + &ol)); +} + +bool FileLock::TryLock(HANDLE handle) { + OVERLAPPED ol = {0}; + return (!!LockFileEx(handle, + LOCKFILE_EXCLUSIVE_LOCK | LOCKFILE_FAIL_IMMEDIATELY, 0, + MAXDWORD, MAXDWORD, &ol)); +} + +bool FileLock::LockShared(HANDLE handle) { + OVERLAPPED ol = {0}; + return (!!LockFileEx(handle, 0, 0, MAXDWORD, MAXDWORD, &ol)); +} + +bool FileLock::TryLockShared(HANDLE handle) { + OVERLAPPED ol = {0}; + return (!!LockFileEx(handle, LOCKFILE_FAIL_IMMEDIATELY, 0, MAXDWORD, MAXDWORD, + &ol)); +} + +bool FileLock::Unlock(HANDLE handle) { + OVERLAPPED ol = {0}; + return (!!UnlockFileEx(handle, 0, MAXDWORD, MAXDWORD, &ol)); +} + +#endif // !_WIN64 && !_WIN32 + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/io/file_lock.h b/src/ailego/io/file_lock.h new file mode 100644 index 0000000..a571d54 --- /dev/null +++ b/src/ailego/io/file_lock.h @@ -0,0 +1,84 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +/*! File Utility + */ +class FileLock { + public: + //! Constructor + FileLock(const File &file) : native_handle_(file.native_handle()) {} + + //! Constructor + FileLock(File::NativeHandle handle) : native_handle_(handle) {} + + //! Locking + bool lock(void) const { + return FileLock::Lock(native_handle_); + } + + //! Try locking + bool try_lock(void) const { + return FileLock::TryLock(native_handle_); + } + + //! Locking (shared) + bool lock_shared(void) const { + return FileLock::LockShared(native_handle_); + } + + //! Try locking (shared) + bool try_lock_shared(void) const { + return FileLock::TryLockShared(native_handle_); + } + + //! Unlocking + bool unlock(void) const { + return FileLock::Unlock(native_handle_); + } + + //! Locking + static bool Lock(File::NativeHandle handle); + + //! Try locking + static bool TryLock(File::NativeHandle handle); + + //! Locking (shared) + static bool LockShared(File::NativeHandle handle); + + //! Try locking (shared) + static bool TryLockShared(File::NativeHandle handle); + + //! Unlocking + static bool Unlock(File::NativeHandle handle); + + private: + //! Disable them + FileLock(const FileLock &) = delete; + FileLock(FileLock &&) = delete; + FileLock &operator=(const FileLock &) = delete; + FileLock &operator=(FileLock &&) = delete; + + //! Members + File::NativeHandle native_handle_; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/io/file_writer.h b/src/ailego/io/file_writer.h new file mode 100644 index 0000000..2448dcc --- /dev/null +++ b/src/ailego/io/file_writer.h @@ -0,0 +1,127 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include "file.h" + +namespace zvec { +namespace ailego { + +/*! File Stream Writer + */ +class FileWriter { + public: + //! Constructor + FileWriter(void) {} + + //! Constructor + FileWriter(FileWriter &&rhs) : file_(std::move(rhs.file_)) {} + + //! Destructor + ~FileWriter(void) {} + + //! Assignment + FileWriter &operator=(FileWriter &&rhs) { + file_ = std::move(rhs.file_); + return *this; + } + + //! Output to writer + FileWriter &operator<<(const char *str) { + size_t len = std::strlen(str); + if (file_.write(str, len) != len) { + throw std::ios_base::failure("Write error"); + } + return *this; + } + + //! Output to writer + FileWriter &operator<<(const std::string &str) { + if (file_.write(str.data(), str.size()) != str.size()) { + throw std::ios_base::failure("Write error"); + } + return *this; + } + + //! Output to writer + FileWriter &operator<<(char c) { + if (file_.write(&c, 1) != 1) { + throw std::ios_base::failure("Write error"); + } + return *this; + } + + //! Test if the file is valid + bool is_valid(void) const { + return file_.is_valid(); + } + + //! Create a local file + bool create(const char *path) { + return file_.create(path, 0, false); + } + + //! Open a local file + bool open(const char *path) { + return file_.open(path, false, false); + } + + //! Close the local file + void close(void) { + file_.close(); + } + + //! Write data into the file + size_t write(const void *data, size_t len) { + return file_.write(data, len); + } + + //! Synchronize memory with physical storage + bool flush(void) { + return file_.flush(); + } + + //! Output with format + void print(const char *format, va_list args) { + char buf[8192]; + std::vsnprintf(buf, sizeof(buf), format, args); + (*this) << buf; + } + + //! Output with format +#if defined(__GNUC__) + void print(const char *format, ...) __attribute__((format(printf, 2, 3))) { +#else + void print(const char *format, ...) { +#endif + va_list args; + va_start(args, format); + this->print(format, args); + va_end(args); + } + + private: + //! Disable them + FileWriter(const FileWriter &) = delete; + FileWriter &operator=(const FileWriter &) = delete; + + //! Members + File file_; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/logger/logger.cc b/src/ailego/logger/logger.cc new file mode 100644 index 0000000..deecb88 --- /dev/null +++ b/src/ailego/logger/logger.cc @@ -0,0 +1,76 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +const int Logger::LEVEL_DEBUG = 0; +const int Logger::LEVEL_INFO = 1; +const int Logger::LEVEL_WARN = 2; +const int Logger::LEVEL_ERROR = 3; +const int Logger::LEVEL_FATAL = 4; + +/*! Console Logger + */ +struct ConsoleLogger : public Logger { + //! Initialize Logger + int init(const Params &) override { + return 0; + } + + //! Cleanup Logger + int cleanup(void) override { + return 0; + } + + //! Log Message + void log(int level, const char *file, int line, const char *format, + va_list args) override { + char buffer[8192]; + std::ostringstream stream; + + ailego::Realtime::Localtime(buffer, sizeof(buffer)); + stream << '[' << LevelString(level) << ' ' << buffer << ' ' + << std::this_thread::get_id() << ' ' << ailego::File::BaseName(file) + << ':' << line << "] "; + + vsnprintf(buffer, sizeof(buffer), format, args); + stream << buffer << '\n'; + + if (level <= LEVEL_INFO) { + std::cout << stream.str() << std::flush; + } else { + std::cerr << stream.str() << std::flush; + } + } +}; + +//! Logger Level +int LoggerBroker::logger_level_ = Logger::LEVEL_WARN; + +//! Logger +Logger::Pointer LoggerBroker::logger_(new ConsoleLogger); + +//! Register Console Logger in Factory +FACTORY_REGISTER_LOGGER(ConsoleLogger); + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/cosine_distance_matrix.h b/src/ailego/math/cosine_distance_matrix.h new file mode 100644 index 0000000..2f5ab51 --- /dev/null +++ b/src/ailego/math/cosine_distance_matrix.h @@ -0,0 +1,90 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +/*! Cosine Distance Matrix + */ +template +struct CosineDistanceMatrix; + +/*! Cosine Distance Matrix (M=1, N=1) + */ +template +struct CosineDistanceMatrix< + T, 1, 1, typename std::enable_if::value>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + constexpr size_t extra_dim = sizeof(float) / sizeof(ValueType); + size_t d = dim - extra_dim; + + float ip; + InnerProductMatrix::Compute(m, q, d, &ip); + + *out = 1 - ip; + } +}; + +/*! Cosine Distance Matrix + */ +template +struct CosineDistanceMatrix< + T, M, N, + typename std::enable_if::value && M >= 2 && + N >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType * /*m*/, const ValueType * /*q*/, + size_t /*dim*/, float *out) { + // ailego_assert(m && q && dim && out); + + *out = 0.0f; + } +}; + +/*! Cosine Distance Matrix (N=1) + */ +template +struct CosineDistanceMatrix< + T, M, 1, + typename std::enable_if::value && M >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType * /*m*/, const ValueType * /*q*/, + size_t /*dim*/, float *out) { + // ailego_assert(m && q && dim && out); + + *out = 0.0f; + } +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/distance.h b/src/ailego/math/distance.h new file mode 100644 index 0000000..a52bc4b --- /dev/null +++ b/src/ailego/math/distance.h @@ -0,0 +1,254 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "distance_matrix.h" + +namespace zvec { +namespace ailego { + +/*! Distance module + */ +struct Distance { + //! Compute the squared euclidean distance between two vectors (FP32) + static float SquaredEuclidean(const float *lhs, const float *rhs, + size_t dim) { + float result; + SquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, + &result); + return result; + } + + //! Compute the squared euclidean distance between two vectors (FP16) + static float SquaredEuclidean(const Float16 *lhs, const Float16 *rhs, + size_t dim) { + float result; + SquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, + &result); + return result; + } + + //! Compute the squared euclidean distance between two vectors (INT8) + static float SquaredEuclidean(const int8_t *lhs, const int8_t *rhs, + size_t dim) { + float result; + SquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, + &result); + return result; + } + + //! Compute the squared euclidean distance between two vectors (INT4) + static float SquaredEuclidean(const uint8_t *lhs, const uint8_t *rhs, + size_t dim) { + float result; + SquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, + &result); + return result; + } + + //! Compute the euclidean distance between two vectors (FP32) + static float Euclidean(const float *lhs, const float *rhs, size_t dim) { + float result; + EuclideanDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the euclidean distance between two vectors (FP16) + static float Euclidean(const Float16 *lhs, const Float16 *rhs, size_t dim) { + float result; + EuclideanDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the euclidean distance between two vectors (INT8) + static float Euclidean(const int8_t *lhs, const int8_t *rhs, size_t dim) { + float result; + EuclideanDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the euclidean distance between two vectors (INT4) + static float Euclidean(const uint8_t *lhs, const uint8_t *rhs, size_t dim) { + float result; + EuclideanDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the inner product between two vectors (FP32) + static float InnerProduct(const float *lhs, const float *rhs, size_t dim) { + float result; + InnerProductMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the inner product between two vectors (FP16) + static float InnerProduct(const Float16 *lhs, const Float16 *rhs, + size_t dim) { + float result; + InnerProductMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the inner product between two vectors (INT8) + static float InnerProduct(const int8_t *lhs, const int8_t *rhs, size_t dim) { + float result; + InnerProductMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the minus inner product between two vectors (INT4) + static float InnerProduct(const uint8_t *lhs, const uint8_t *rhs, + size_t dim) { + float result; + InnerProductMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the minus inner product between two vectors (FP32) + static float MinusInnerProduct(const float *lhs, const float *rhs, + size_t dim) { + float result; + MinusInnerProductMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the minus inner product between two vectors (FP16) + static float MinusInnerProduct(const Float16 *lhs, const Float16 *rhs, + size_t dim) { + float result; + MinusInnerProductMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the minus inner product between two vectors (INT8) + static float MinusInnerProduct(const int8_t *lhs, const int8_t *rhs, + size_t dim) { + float result; + MinusInnerProductMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the minus inner product between two vectors (INT4) + static float MinusInnerProduct(const uint8_t *lhs, const uint8_t *rhs, + size_t dim) { + float result; + MinusInnerProductMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the mips squared L2 distance between two vectors + //! (FP32, RepeatedQuadraticInjection) + static float MipsSquaredEuclidean(const float *lhs, const float *rhs, + size_t dim, size_t m, float eta) { + float result; + MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, m, + eta, &result); + return result; + } + + //! Compute the mips squared L2 distance between two vectors + //! (FP16, RepeatedQuadraticInjection) + static float MipsSquaredEuclidean(const Float16 *lhs, const Float16 *rhs, + size_t dim, size_t m, float eta) { + float result; + MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, m, + eta, &result); + return result; + } + + //! Compute the mips squared L2 distance between two vectors + //! (INT8, RepeatedQuadraticInjection) + static float MipsSquaredEuclidean(const int8_t *lhs, const int8_t *rhs, + size_t dim, size_t m, float eta) { + float result; + MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, m, + eta, &result); + return result; + } + + //! Compute the mips squared L2 distance between two vectors + //! (INT4, RepeatedQuadraticInjection) + static float MipsSquaredEuclidean(const uint8_t *lhs, const uint8_t *rhs, + size_t dim, size_t m, float eta) { + float result; + MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, m, + eta, &result); + return result; + } + + //! Compute the mips squared L2 distance between two vectors + //! (FP32, SphericalInjection) + static float MipsSquaredEuclidean(const float *lhs, const float *rhs, + size_t dim, float eta) { + float result; + MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, eta, + &result); + return result; + } + + //! Compute the mips squared L2 distance between two vectors + //! (FP16, SphericalInjection) + static float MipsSquaredEuclidean(const Float16 *lhs, const Float16 *rhs, + size_t dim, float eta) { + float result; + MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, + eta, &result); + return result; + } + + //! Compute the mips squared L2 distance between two vectors + //! (INT8, SphericalInjection) + static float MipsSquaredEuclidean(const int8_t *lhs, const int8_t *rhs, + size_t dim, float eta) { + float result; + MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, + eta, &result); + return result; + } + + //! Compute the mips squared L2 distance between two vectors + //! (INT4, SphericalInjection) + static float MipsSquaredEuclidean(const uint8_t *lhs, const uint8_t *rhs, + size_t dim, float eta) { + float result; + MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, + eta, &result); + return result; + } + + //! Compute the cosine distance between two vectors (FP32) + static float Cosine(const float *lhs, const float *rhs, size_t dim) { + float result; + CosineDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the cosine distance between two vectors (FP16) + static float Cosine(const Float16 *lhs, const Float16 *rhs, size_t dim) { + float result; + CosineDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + + //! Compute the cosine distance between two vectors (FP16) + static float Cosine(const int8_t *lhs, const int8_t *rhs, size_t dim) { + float result; + CosineDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/distance_matrix.h b/src/ailego/math/distance_matrix.h new file mode 100644 index 0000000..23b4644 --- /dev/null +++ b/src/ailego/math/distance_matrix.h @@ -0,0 +1,20 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "cosine_distance_matrix.h" +#include "euclidean_distance_matrix.h" +#include "inner_product_matrix.h" +#include "mips_euclidean_distance_matrix.h" diff --git a/src/ailego/math/distance_matrix_accum_fp16.i b/src/ailego/math/distance_matrix_accum_fp16.i new file mode 100644 index 0000000..6f94b06 --- /dev/null +++ b/src/ailego/math/distance_matrix_accum_fp16.i @@ -0,0 +1,879 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_fp16.i" +#include "matrix_utility.i" + +#if !defined(__FMA__) +#define _mm_fmadd_ps(a, b, c) _mm_add_ps(_mm_mul_ps((a), (b)), (c)) +#define _mm256_fmadd_ps(a, b, c) _mm256_add_ps(_mm256_mul_ps((a), (b)), (c)) +#endif // !__FMA__ + +#if defined(__AVX512F__) && !defined(__AVX512DQ__) +#define _mm512_and_ps(a, b) \ + _mm512_castsi512_ps( \ + _mm512_and_epi32(_mm512_castps_si512(a), _mm512_castps_si512(b))) +#define _mm512_mask_and_ps(src, k, a, b) \ + _mm512_castsi512_ps(_mm512_mask_and_epi32(_mm512_castps_si512(src), (k), \ + _mm512_castps_si512(a), \ + _mm512_castps_si512(b))) +#endif // __AVX512DQ__ + +//! Compute the distance between matrix and query (FP16, M=1, N=1) +#define ACCUM_FP16_1X1_AVX(m, q, dim, out, _MASK, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + const Float16 *qe = q + dim; \ + const Float16 *qe_aligned = q + ((dim >> 4) << 4); \ + if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \ + for (; q != qe_aligned; m += 16, q += 16) { \ + MATRIX_FP16_ITER_1X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + if (qe >= qe_aligned + 8) { \ + __m256 ymm_m = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)m)); \ + __m256 ymm_q = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)q)); \ + ACCUM_FP32_STEP_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + m += 8; \ + q += 8; \ + } \ + } else { \ + for (; q != qe_aligned; m += 16, q += 16) { \ + MATRIX_FP16_ITER_1X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + if (qe >= qe_aligned + 8) { \ + __m256 ymm_m = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)m)); \ + __m256 ymm_q = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q)); \ + ACCUM_FP32_STEP_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + m += 8; \ + q += 8; \ + } \ + } \ + MATRIX_FP16_MASK_AVX(m, q, (qe - q), _MASK, ymm_sum, ACCUM_FP32_STEP_AVX) \ + *out = _NORM(HorizontalAdd_FP32_V256(ymm_sum_0_0)); + +//! Compute the distance between matrix and query (FP16, M=2, N=1) +#define ACCUM_FP16_2X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + const Float16 *qe_aligned = q + ((dim >> 2) << 2); \ + const Float16 *qe = q + dim; \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_aligned; m += 8, q += 4) { \ + MATRIX_FP16_ITER_2X1_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (; q != qe_aligned; m += 8, q += 4) { \ + MATRIX_FP16_ITER_2X1_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + if (qe >= qe_aligned + 2) { \ + __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64x(*(const long long *)(m))); \ + __m128 xmm_q = _mm_cvtph_ps( \ + _mm_shufflelo_epi16(_mm_broadcast_si32(q), _MM_SHUFFLE(1, 1, 0, 0))); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + m += 4; \ + q += 2; \ + } \ + xmm_sum_0_0 = \ + _mm_add_ps(xmm_sum_0_0, _mm_movehl_ps(xmm_sum_0_0, xmm_sum_0_0)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_cvtph_ps( \ + _mm_shufflelo_epi16(_mm_broadcast_si32(m), _MM_SHUFFLE(0, 0, 1, 0))); \ + __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi16(*(const short *)(q))); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0)); + +//! Compute the distance between matrix and query (FP16, M=2, N=2) +#define ACCUM_FP16_2X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + const Float16 *qe_aligned = q + ((dim >> 2) << 3); \ + const Float16 *qe = q + (dim << 1); \ + if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \ + for (; q != qe_aligned; m += 8, q += 8) { \ + MATRIX_FP16_ITER_2X2_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (; q != qe_aligned; m += 8, q += 8) { \ + MATRIX_FP16_ITER_2X2_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \ + _mm256_extractf128_ps(ymm_sum_0_1, 1)); \ + if (qe >= qe_aligned + 4) { \ + __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64x(*(const long long *)(m))); \ + __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi64x(*(const long long *)(q))); \ + __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(2, 2, 0, 0)); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_0) \ + xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 1, 1)); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_1) \ + m += 4; \ + q += 4; \ + } \ + xmm_sum_0_0 = _mm_add_ps(_mm_movelh_ps(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_movehl_ps(xmm_sum_0_1, xmm_sum_0_0)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_cvtph_ps( \ + _mm_shufflelo_epi16(_mm_broadcast_si32(m), _MM_SHUFFLE(1, 0, 1, 0))); \ + __m128 xmm_q = _mm_cvtph_ps( \ + _mm_shufflelo_epi16(_mm_broadcast_si32(q), _MM_SHUFFLE(1, 1, 0, 0))); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=4, N=1) +#define ACCUM_FP16_4X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + const Float16 *qe = q + dim; \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const Float16 *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \ + m += 8, q += 2) { \ + MATRIX_FP16_ITER_4X1_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \ + m += 8, q += 2) { \ + MATRIX_FP16_ITER_4X1_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64x(*(const long long *)(m))); \ + __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi16(*(const short *)(q))); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=4, N=2) +#define ACCUM_FP16_4X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + const Float16 *qe = q + (dim << 1); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const Float16 *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \ + m += 8, q += 4) { \ + MATRIX_FP16_ITER_4X2_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \ + m += 8, q += 4) { \ + MATRIX_FP16_ITER_4X2_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \ + _mm256_extractf128_ps(ymm_sum_0_1, 1)); \ + if (q != qe) { \ + __m128 xmm_q_0 = _mm_cvtph_ps(_mm_set1_epi16(*(const short *)(q + 0))); \ + __m128 xmm_q_1 = _mm_cvtph_ps(_mm_set1_epi16(*(const short *)(q + 1))); \ + __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64x(*(const long long *)(m))); \ + MATRIX_VAR_PROC(1, 2, 0, xmm_m, xmm_q, xmm_sum, ACCUM_FP32_STEP_SSE) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=4, N=4) +#define ACCUM_FP16_4X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256, ymm_sum, _mm256_setzero_ps()) \ + const Float16 *qe = q + (dim << 2); \ + if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \ + for (const Float16 *qe_aligned = q + ((dim >> 1) << 3); q != qe_aligned; \ + m += 8, q += 8) { \ + MATRIX_FP16_ITER_4X4_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe_aligned = q + ((dim >> 1) << 3); q != qe_aligned; \ + m += 8, q += 8) { \ + MATRIX_FP16_ITER_4X4_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \ + _mm256_extractf128_ps(ymm_sum_0_1, 1)); \ + __m128 xmm_sum_0_2 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_2), \ + _mm256_extractf128_ps(ymm_sum_0_2, 1)); \ + __m128 xmm_sum_0_3 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_3), \ + _mm256_extractf128_ps(ymm_sum_0_3, 1)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64x(*(const long long *)(m))); \ + __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi64x(*(const long long *)(q))); \ + __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(0, 0, 0, 0)); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_0) \ + xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(1, 1, 1, 1)); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_1) \ + xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(2, 2, 2, 2)); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_2) \ + xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 3, 3)); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_3) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=8, N=1) +#define ACCUM_FP16_8X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const Float16 *qe = q + dim; q != qe; m += 8, ++q) { \ + MATRIX_FP16_ITER_8X1_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + dim; q != qe; m += 8, ++q) { \ + MATRIX_FP16_ITER_8X1_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=8, N=2) +#define ACCUM_FP16_8X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \ + MATRIX_FP16_ITER_8X2_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \ + MATRIX_FP16_ITER_8X2_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=8, N=4) +#define ACCUM_FP16_8X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \ + MATRIX_FP16_ITER_8X4_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \ + MATRIX_FP16_ITER_8X4_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=8, N=8) +#define ACCUM_FP16_8X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 8, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \ + MATRIX_FP16_ITER_8X8_AVX(m, q, ymm_sum, _mm_load_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \ + MATRIX_FP16_ITER_8X8_AVX(m, q, ymm_sum, _mm_loadu_si128, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=1) +#define ACCUM_FP16_16X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP16_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP16_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=2) +#define ACCUM_FP16_16X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP16_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP16_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=4) +#define ACCUM_FP16_16X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP16_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP16_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=8) +#define ACCUM_FP16_16X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP16_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP16_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=16) +#define ACCUM_FP16_16X16_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 16, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP16_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP16_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=1) +#define ACCUM_FP16_32X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP16_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP16_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=2) +#define ACCUM_FP16_32X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP16_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP16_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=4) +#define ACCUM_FP16_32X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP16_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP16_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=8) +#define ACCUM_FP16_32X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP16_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP16_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=16) +#define ACCUM_FP16_32X16_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP16_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP16_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=32) +#define ACCUM_FP16_32X32_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 32, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP16_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP16_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=1, N=1) +#define ACCUM_FP16_1X1_AVX512(m, q, dim, out, _MASK, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m512, zmm_sum, _mm512_setzero_ps()) \ + const Float16 *qe = q + dim; \ + const Float16 *qe_aligned = q + ((dim >> 5) << 5); \ + if (((uintptr_t)m & 0x3f) == 0 && ((uintptr_t)q & 0x3f) == 0) { \ + for (; q != qe_aligned; m += 32, q += 32) { \ + MATRIX_FP16_ITER_1X1_AVX512(m, q, zmm_sum, _mm512_load_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + if (qe >= qe_aligned + 16) { \ + __m512 zmm_m = _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)m)); \ + __m512 zmm_q = _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)q)); \ + ACCUM_FP32_STEP_AVX512(zmm_m, zmm_q, zmm_sum_0_0) \ + m += 16; \ + q += 16; \ + } \ + } else { \ + for (; q != qe_aligned; m += 32, q += 32) { \ + MATRIX_FP16_ITER_1X1_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + if (qe >= qe_aligned + 16) { \ + __m512 zmm_m = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)m)); \ + __m512 zmm_q = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)q)); \ + ACCUM_FP32_STEP_AVX512(zmm_m, zmm_q, zmm_sum_0_0) \ + m += 16; \ + q += 16; \ + } \ + } \ + __m256 ymm_sum_0_0 = _mm256_add_ps(_mm512_castps512_ps256(zmm_sum_0_0), \ + _mm256_castpd_ps(_mm512_extractf64x4_pd( \ + _mm512_castps_pd(zmm_sum_0_0), 1))); \ + if (qe >= q + 8) { \ + __m256 ymm_m = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)m)); \ + __m256 ymm_q = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q)); \ + ACCUM_FP32_STEP_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + m += 8; \ + q += 8; \ + } \ + MATRIX_FP16_MASK_AVX(m, q, (qe - q), _MASK, ymm_sum, ACCUM_FP32_STEP_AVX) \ + *out = _NORM(HorizontalAdd_FP32_V256(ymm_sum_0_0)); + +//! Compute the distance between matrix and query (FP16, M=16, N=1) +#define ACCUM_FP16_16X1_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP16_ITER_16X1_AVX512(m, q, zmm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP16_ITER_16X1_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 1, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=2) +#define ACCUM_FP16_16X2_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP16_ITER_16X2_AVX512(m, q, zmm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP16_ITER_16X2_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 2, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=4) +#define ACCUM_FP16_16X4_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP16_ITER_16X4_AVX512(m, q, zmm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP16_ITER_16X4_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 4, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=8) +#define ACCUM_FP16_16X8_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 8, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP16_ITER_16X8_AVX512(m, q, zmm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP16_ITER_16X8_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 8, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 8, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=16, N=16) +#define ACCUM_FP16_16X16_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 16, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \ + for (const Float16 *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP16_ITER_16X16_AVX512(m, q, zmm_sum, _mm256_load_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP16_ITER_16X16_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 16, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 16, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=1) +#define ACCUM_FP16_32X1_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const Float16 *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP16_ITER_32X1_AVX512(m, q, zmm_sum, _mm512_load_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP16_ITER_32X1_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 1, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=2) +#define ACCUM_FP16_32X2_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP16_ITER_32X2_AVX512(m, q, zmm_sum, _mm512_load_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP16_ITER_32X2_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 2, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=4) +#define ACCUM_FP16_32X4_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP16_ITER_32X4_AVX512(m, q, zmm_sum, _mm512_load_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP16_ITER_32X4_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 4, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=8) +#define ACCUM_FP16_32X8_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP16_ITER_32X8_AVX512(m, q, zmm_sum, _mm512_load_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP16_ITER_32X8_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 8, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=16) +#define ACCUM_FP16_32X16_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 16, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const Float16 *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP16_ITER_32X16_AVX512(m, q, zmm_sum, _mm512_load_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP16_ITER_32X16_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 16, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 16, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP16, M=32, N=32) +#define ACCUM_FP16_32X32_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 32, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const Float16 *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP16_ITER_32X32_AVX512(m, q, zmm_sum, _mm512_load_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const Float16 *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP16_ITER_32X32_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 32, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 32, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) +//! Compute the distance between matrix and query (FP16, M=1, N=1) +#define ACCUM_FP16_1X1_NEON(m, q, dim, out, _MASK, _NORM) \ + MATRIX_VAR_INIT(1, 1, float16x8_t, v_sum, vdupq_n_f16(0)) \ + const Float16 *qe = q + dim; \ + const Float16 *qe_aligned = q + ((dim >> 3) << 3); \ + for (; q != qe_aligned; m += 8, q += 8) { \ + MATRIX_FP16_ITER_1X1_NEON(m, q, v_sum, ACCUM_FP16_STEP_NEON) \ + } \ + if (qe >= qe_aligned + 4) { \ + float16x8_t v_m = \ + vcombine_f16(vld1_f16((const float16_t *)m), \ + vreinterpret_f16_u64(vdup_n_u64((uint64_t)(_MASK)))); \ + float16x8_t v_q = \ + vcombine_f16(vld1_f16((const float16_t *)q), \ + vreinterpret_f16_u64(vdup_n_u64((uint64_t)(_MASK)))); \ + ACCUM_FP16_STEP_NEON(v_m, v_q, v_sum_0_0) \ + m += 4; \ + q += 4; \ + } \ + float result = vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(v_sum_0_0)), \ + vcvt_high_f32_f16(v_sum_0_0))); \ + switch (qe - q) { \ + case 3: \ + ACCUM_FP16_STEP_GENERAL(m[2], q[2], result) \ + /* FALLTHRU */ \ + case 2: \ + ACCUM_FP16_STEP_GENERAL(m[1], q[1], result) \ + /* FALLTHRU */ \ + case 1: \ + ACCUM_FP16_STEP_GENERAL(m[0], q[0], result) \ + } \ + *out = _NORM(result); + +#else +//! Compute the distance between matrix and query (FP16, M=1, N=1) +#define ACCUM_FP16_1X1_NEON(m, q, dim, out, _MASK, _NORM) \ + MATRIX_VAR_INIT(1, 1, float32x4_t, v_sum, vdupq_n_f32(0)) \ + const Float16 *qe = q + dim; \ + const Float16 *qe_aligned = q + ((dim >> 3) << 3); \ + for (; q != qe_aligned; m += 8, q += 8) { \ + MATRIX_FP16_ITER_1X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + if (qe >= qe_aligned + 4) { \ + float32x4_t v_m = vcvt_f32_f16(vld1_f16((const float16_t *)m)); \ + float32x4_t v_q = vcvt_f32_f16(vld1_f16((const float16_t *)q)); \ + ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum_0_0) \ + m += 4; \ + q += 4; \ + } \ + float result = vaddvq_f32(v_sum_0_0); \ + switch (qe - q) { \ + case 3: \ + ACCUM_FP16_STEP_GENERAL(m[2], q[2], result) \ + /* FALLTHRU */ \ + case 2: \ + ACCUM_FP16_STEP_GENERAL(m[1], q[1], result) \ + /* FALLTHRU */ \ + case 1: \ + ACCUM_FP16_STEP_GENERAL(m[0], q[0], result) \ + } \ + *out = _NORM(result); + +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC diff --git a/src/ailego/math/distance_matrix_accum_fp32.i b/src/ailego/math/distance_matrix_accum_fp32.i new file mode 100644 index 0000000..c186492 --- /dev/null +++ b/src/ailego/math/distance_matrix_accum_fp32.i @@ -0,0 +1,1406 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_fp32.i" +#include "matrix_utility.i" + +#if !defined(__FMA__) +#define _mm_fmadd_ps(a, b, c) _mm_add_ps(_mm_mul_ps((a), (b)), (c)) +#define _mm256_fmadd_ps(a, b, c) _mm256_add_ps(_mm256_mul_ps((a), (b)), (c)) +#endif // !__FMA__ + +#if defined(__AVX512F__) && !defined(__AVX512DQ__) +#define _mm512_and_ps(a, b) \ + _mm512_castsi512_ps( \ + _mm512_and_epi32(_mm512_castps_si512(a), _mm512_castps_si512(b))) +#define _mm512_mask_and_ps(src, k, a, b) \ + _mm512_castsi512_ps(_mm512_mask_and_epi32(_mm512_castps_si512(src), (k), \ + _mm512_castps_si512(a), \ + _mm512_castps_si512(b))) +#endif // __AVX512DQ__ + +#if defined(__ARM_NEON) && !defined(__aarch64__) +static inline float32_t vaddvq_f32(float32x4_t v) { + float32x2_t s = vadd_f32(vget_low_f32(v), vget_high_f32(v)); + return vget_lane_f32(vpadd_f32(s, s), 0); +} + +static inline int32_t vaddvq_s32(int32x4_t v) { + int32x2_t s = vadd_s32(vget_low_s32(v), vget_high_s32(v)); + return vget_lane_s32(vpadd_s32(s, s), 0); +} +#endif //__ARM_NEON && !__aarch64__ + +#if defined(__aarch64__) +#define ACCUM_FP32_2X1_NEON ACCUM_FP32_2X1_NEON_A64 +#else +#define ACCUM_FP32_2X1_NEON ACCUM_FP32_2X1_NEON_A32 +#endif // __aarch64__ + +//! Compute the distance between matrix and query (FP32, M=2, N=1) +#define ACCUM_FP32_2X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128, xmm_sum, _mm_setzero_ps()) \ + const float *qe_aligned = q + ((dim >> 2) << 2); \ + const float *qe = q + dim; \ + if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \ + for (; q != qe_aligned; m += 8, q += 4) { \ + MATRIX_FP32_ITER_2X1_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + if (qe >= qe_aligned + 2) { \ + __m128 xmm_m = _mm_load_ps(m); \ + __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + m += 4; \ + q += 2; \ + } \ + } else { \ + for (; q != qe_aligned; m += 8, q += 4) { \ + MATRIX_FP32_ITER_2X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + if (qe >= qe_aligned + 2) { \ + __m128 xmm_m = _mm_loadu_ps(m); \ + __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + m += 4; \ + q += 2; \ + } \ + } \ + xmm_sum_0_0 = _mm_add_ps(xmm_sum_0_0, xmm_sum_0_1); \ + xmm_sum_0_0 = \ + _mm_add_ps(xmm_sum_0_0, _mm_movehl_ps(xmm_sum_0_0, xmm_sum_0_0)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_set_ps(0.0f, 0.0f, m[1], m[0]); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0)); + +//! Compute the distance between matrix and query (FP32, M=2, N=2) +#define ACCUM_FP32_2X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128, xmm_sum, _mm_setzero_ps()) \ + const float *qe = q + (dim << 1); \ + if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \ + for (const float *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \ + m += 4, q += 4) { \ + MATRIX_FP32_ITER_2X2_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \ + m += 4, q += 4) { \ + MATRIX_FP32_ITER_2X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + xmm_sum_0_0 = _mm_add_ps(_mm_movelh_ps(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_movehl_ps(xmm_sum_0_1, xmm_sum_0_0)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_set_ps(m[1], m[0], m[1], m[0]); \ + __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=4, N=1) +#define ACCUM_FP32_4X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128, xmm_sum, _mm_setzero_ps()) \ + const float *qe = q + dim; \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \ + m += 8, q += 2) { \ + MATRIX_FP32_ITER_4X1_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + if (q != qe) { \ + __m128 xmm_m = _mm_load_ps(m); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + } else { \ + for (const float *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \ + m += 8, q += 2) { \ + MATRIX_FP32_ITER_4X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + if (q != qe) { \ + __m128 xmm_m = _mm_loadu_ps(m); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + } \ + xmm_sum_0_0 = _mm_add_ps(xmm_sum_0_0, xmm_sum_0_1); \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=4, N=2) +#define ACCUM_FP32_4X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 4, q += 2) { \ + MATRIX_FP32_ITER_4X2_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 4, q += 2) { \ + MATRIX_FP32_ITER_4X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=4, N=4) +#define ACCUM_FP32_4X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 4, q += 4) { \ + MATRIX_FP32_ITER_4X4_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 4, q += 4) { \ + MATRIX_FP32_ITER_4X4_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=8, N=1) +#define ACCUM_FP32_8X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + dim; q != qe; m += 8, ++q) { \ + MATRIX_FP32_ITER_8X1_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + dim; q != qe; m += 8, ++q) { \ + MATRIX_FP32_ITER_8X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=8, N=2) +#define ACCUM_FP32_8X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \ + MATRIX_FP32_ITER_8X2_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \ + MATRIX_FP32_ITER_8X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=8, N=4) +#define ACCUM_FP32_8X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \ + MATRIX_FP32_ITER_8X4_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \ + MATRIX_FP32_ITER_8X4_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=8, N=8) +#define ACCUM_FP32_8X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \ + MATRIX_FP32_ITER_8X8_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \ + MATRIX_FP32_ITER_8X8_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=1) +#define ACCUM_FP32_16X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP32_ITER_16X1_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP32_ITER_16X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=2) +#define ACCUM_FP32_16X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP32_ITER_16X2_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP32_ITER_16X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=4) +#define ACCUM_FP32_16X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP32_ITER_16X4_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP32_ITER_16X4_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=8) +#define ACCUM_FP32_16X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP32_ITER_16X8_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP32_ITER_16X8_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=16) +#define ACCUM_FP32_16X16_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP32_ITER_16X16_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP32_ITER_16X16_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=1) +#define ACCUM_FP32_32X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 1, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP32_ITER_32X1_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP32_ITER_32X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=2) +#define ACCUM_FP32_32X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 2, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP32_ITER_32X2_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP32_ITER_32X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=4) +#define ACCUM_FP32_32X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 4, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP32_ITER_32X4_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP32_ITER_32X4_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=8) +#define ACCUM_FP32_32X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 8, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP32_ITER_32X8_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP32_ITER_32X8_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=16) +#define ACCUM_FP32_32X16_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 16, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP32_ITER_32X16_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP32_ITER_32X16_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=32) +#define ACCUM_FP32_32X32_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 32, __m128, xmm_sum, _mm_setzero_ps()) \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP32_ITER_32X32_SSE(m, q, xmm_sum, _mm_load_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } else { \ + for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP32_ITER_32X32_SSE(m, q, xmm_sum, _mm_loadu_ps, \ + ACCUM_FP32_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=2, N=1) +#define ACCUM_FP32_2X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + const float *qe_aligned = q + ((dim >> 2) << 2); \ + const float *qe = q + dim; \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_aligned; m += 8, q += 4) { \ + MATRIX_FP32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (; q != qe_aligned; m += 8, q += 4) { \ + MATRIX_FP32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + if (qe >= qe_aligned + 2) { \ + __m128 xmm_m = _mm_loadu_ps(m); \ + __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + m += 4; \ + q += 2; \ + } \ + xmm_sum_0_0 = \ + _mm_add_ps(xmm_sum_0_0, _mm_movehl_ps(xmm_sum_0_0, xmm_sum_0_0)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_set_ps(0.0f, 0.0f, m[1], m[0]); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0)); + +//! Compute the distance between matrix and query (FP32, M=2, N=2) +#define ACCUM_FP32_2X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + const float *qe_aligned = q + ((dim >> 2) << 3); \ + const float *qe = q + (dim << 1); \ + if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \ + for (; q != qe_aligned; m += 8, q += 8) { \ + MATRIX_FP32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (; q != qe_aligned; m += 8, q += 8) { \ + MATRIX_FP32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \ + _mm256_extractf128_ps(ymm_sum_0_1, 1)); \ + if (qe >= qe_aligned + 4) { \ + __m128 xmm_q = _mm_loadu_ps(q); \ + __m128 xmm_m = _mm_loadu_ps(m); \ + __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(2, 2, 0, 0)); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_0) \ + xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 1, 1)); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_1) \ + m += 4; \ + q += 4; \ + } \ + xmm_sum_0_0 = _mm_add_ps(_mm_movelh_ps(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_movehl_ps(xmm_sum_0_1, xmm_sum_0_0)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_set_ps(m[1], m[0], m[1], m[0]); \ + __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=4, N=1) +#define ACCUM_FP32_4X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + const float *qe = q + dim; \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \ + m += 8, q += 2) { \ + MATRIX_FP32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \ + m += 8, q += 2) { \ + MATRIX_FP32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_loadu_ps(m); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=4, N=2) +#define ACCUM_FP32_4X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + const float *qe = q + (dim << 1); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \ + m += 8, q += 4) { \ + MATRIX_FP32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \ + m += 8, q += 4) { \ + MATRIX_FP32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \ + _mm256_extractf128_ps(ymm_sum_0_1, 1)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_loadu_ps(m); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_1) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=4, N=4) +#define ACCUM_FP32_4X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256, ymm_sum, _mm256_setzero_ps()) \ + const float *qe = q + (dim << 2); \ + if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \ + for (const float *qe_aligned = q + ((dim >> 1) << 3); q != qe_aligned; \ + m += 8, q += 8) { \ + MATRIX_FP32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe_aligned = q + ((dim >> 1) << 3); q != qe_aligned; \ + m += 8, q += 8) { \ + MATRIX_FP32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \ + _mm256_extractf128_ps(ymm_sum_0_0, 1)); \ + __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \ + _mm256_extractf128_ps(ymm_sum_0_1, 1)); \ + __m128 xmm_sum_0_2 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_2), \ + _mm256_extractf128_ps(ymm_sum_0_2, 1)); \ + __m128 xmm_sum_0_3 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_3), \ + _mm256_extractf128_ps(ymm_sum_0_3, 1)); \ + if (q != qe) { \ + __m128 xmm_m = _mm_loadu_ps(m); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_1) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_2) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_3) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=8, N=1) +#define ACCUM_FP32_8X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + dim; q != qe; m += 8, ++q) { \ + MATRIX_FP32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + dim; q != qe; m += 8, ++q) { \ + MATRIX_FP32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=8, N=2) +#define ACCUM_FP32_8X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \ + MATRIX_FP32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \ + MATRIX_FP32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=8, N=4) +#define ACCUM_FP32_8X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \ + MATRIX_FP32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \ + MATRIX_FP32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=8, N=8) +#define ACCUM_FP32_8X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 8, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \ + MATRIX_FP32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \ + MATRIX_FP32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=1) +#define ACCUM_FP32_16X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=2) +#define ACCUM_FP32_16X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=4) +#define ACCUM_FP32_16X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=8) +#define ACCUM_FP32_16X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=16) +#define ACCUM_FP32_16X16_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 16, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=1) +#define ACCUM_FP32_32X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=2) +#define ACCUM_FP32_32X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=4) +#define ACCUM_FP32_32X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=8) +#define ACCUM_FP32_32X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=16) +#define ACCUM_FP32_32X16_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=32) +#define ACCUM_FP32_32X32_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 32, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } else { \ + for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_ps, \ + ACCUM_FP32_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=1) +#define ACCUM_FP32_16X1_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP32_ITER_16X1_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP32_ITER_16X1_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 1, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=2) +#define ACCUM_FP32_16X2_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP32_ITER_16X2_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP32_ITER_16X2_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 2, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=4) +#define ACCUM_FP32_16X4_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP32_ITER_16X4_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP32_ITER_16X4_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 4, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=8) +#define ACCUM_FP32_16X8_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 8, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP32_ITER_16X8_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP32_ITER_16X8_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 8, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 8, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=16, N=16) +#define ACCUM_FP32_16X16_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 16, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP32_ITER_16X16_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP32_ITER_16X16_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(1, 16, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 16, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=1) +#define ACCUM_FP32_32X1_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP32_ITER_32X1_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP32_ITER_32X1_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 1, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=2) +#define ACCUM_FP32_32X2_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP32_ITER_32X2_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP32_ITER_32X2_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 2, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=4) +#define ACCUM_FP32_32X4_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP32_ITER_32X4_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP32_ITER_32X4_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 4, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=8) +#define ACCUM_FP32_32X8_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP32_ITER_32X8_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP32_ITER_32X8_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 8, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=16) +#define ACCUM_FP32_32X16_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 16, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP32_ITER_32X16_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP32_ITER_32X16_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 16, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 16, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=32, N=32) +#define ACCUM_FP32_32X32_AVX512(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 32, __m512, zmm_sum, _mm512_setzero_ps()) \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP32_ITER_32X32_AVX512(m, q, zmm_sum, _mm512_load_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } else { \ + for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP32_ITER_32X32_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \ + ACCUM_FP32_STEP_AVX512) \ + } \ + } \ + if (((uintptr_t)out & 0x3f) == 0) { \ + MATRIX_VAR_STORE(2, 32, 16, zmm_sum, out, _mm512_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 32, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (FP32, M=2, N=1) on A64 +#define ACCUM_FP32_2X1_NEON_A64(m, q, dim, out, _NORM) \ + float32x4_t v_sum = vdupq_n_f32(0); \ + const float *qe_aligned = q + ((dim >> 1) << 1); \ + const float *qe = q + dim; \ + for (; q != qe_aligned; m += 4, q += 2) { \ + MATRIX_FP32_ITER_2X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + v_sum = vaddq_f32( \ + vreinterpretq_f32_u64(vdupq_laneq_u64(vreinterpretq_u64_f32(v_sum), 1)), \ + v_sum); \ + if (q != qe) { \ + float32x4_t v_m = vreinterpretq_f32_u64( \ + vdupq_lane_u64(vld1_u64((const uint64_t *)m), 0)); \ + float32x4_t v_q = vld1q_dup_f32(q); \ + ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum) \ + } \ + vst1_f32(out, _NORM(vget_low_f32(v_sum))); + +//! Compute the distance between matrix and query (FP32, M=2, N=1) on A32 +#define ACCUM_FP32_2X1_NEON_A32(m, q, dim, out, _NORM) \ + float32x4_t v_sum = vdupq_n_f32(0); \ + const float *qe_aligned = q + ((dim >> 1) << 1); \ + const float *qe = q + dim; \ + for (; q != qe_aligned; m += 4, q += 2) { \ + MATRIX_FP32_ITER_2X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + float32x2_t sum = vadd_f32(vget_low_f32(v_sum), vget_high_f32(v_sum)); \ + v_sum = vcombine_f32(sum, sum); \ + if (q != qe) { \ + float32x4_t v_m = vreinterpretq_f32_u64( \ + vdupq_lane_u64(vld1_u64((const uint64_t *)m), 0)); \ + float32x4_t v_q = vld1q_dup_f32(q); \ + ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum) \ + } \ + vst1_f32(out, _NORM(vget_low_f32(v_sum))); + +//! Compute the distance between matrix and query (FP32, M=2, N=2) +#define ACCUM_FP32_2X2_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \ + const float *qe_aligned = q + ((dim >> 1) << 2); \ + const float *qe = q + (dim << 1); \ + for (; q != qe_aligned; m += 4, q += 4) { \ + MATRIX_FP32_ITER_2X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + v_sum_0_0 = vaddq_f32( \ + vcombine_f32(vget_low_f32(v_sum_0_0), vget_low_f32(v_sum_0_1)), \ + vcombine_f32(vget_high_f32(v_sum_0_0), vget_high_f32(v_sum_0_1))); \ + if (q != qe) { \ + float32x2_t v_m_0 = vld1_f32(m); \ + float32x2_t v_q_0 = vld1_f32(q); \ + float32x4_t v_m = vcombine_f32(v_m_0, v_m_0); \ + float32x4_t v_q = \ + vcombine_f32(vdup_lane_f32(v_q_0, 0), vdup_lane_f32(v_q_0, 1)); \ + ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum_0_0) \ + } \ + MATRIX_VAR_STORE(1, 1, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=4, N=1) +#define ACCUM_FP32_4X1_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \ + const float *qe_aligned = q + ((dim >> 1) << 1); \ + const float *qe = q + dim; \ + for (; q != qe_aligned; m += 8, q += 2) { \ + MATRIX_FP32_ITER_4X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + if (q != qe) { \ + float32x4_t v_m = vld1q_f32(m); \ + float32x4_t v_q = vld1q_dup_f32(q); \ + ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum_0_0) \ + } \ + v_sum_0_0 = vaddq_f32(v_sum_0_0, v_sum_0_1); \ + MATRIX_VAR_STORE(1, 1, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=4, N=2) +#define ACCUM_FP32_4X2_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 1); q != qe; m += 4, q += 2) { \ + MATRIX_FP32_ITER_4X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(1, 2, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=4, N=4) +#define ACCUM_FP32_4X4_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 2); q != qe; m += 4, q += 4) { \ + MATRIX_FP32_ITER_4X4_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(1, 4, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=8, N=1) +#define ACCUM_FP32_8X1_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + dim; q != qe; m += 8, ++q) { \ + MATRIX_FP32_ITER_8X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(2, 1, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=8, N=2) +#define ACCUM_FP32_8X2_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \ + MATRIX_FP32_ITER_8X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(2, 2, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=8, N=4) +#define ACCUM_FP32_8X4_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \ + MATRIX_FP32_ITER_8X4_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(2, 4, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=8, N=8) +#define ACCUM_FP32_8X8_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \ + MATRIX_FP32_ITER_8X8_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(2, 8, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=16, N=1) +#define ACCUM_FP32_16X1_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + dim; q != qe; m += 16, ++q) { \ + MATRIX_FP32_ITER_16X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(4, 1, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=16, N=2) +#define ACCUM_FP32_16X2_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \ + MATRIX_FP32_ITER_16X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(4, 2, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=16, N=4) +#define ACCUM_FP32_16X4_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \ + MATRIX_FP32_ITER_16X4_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(4, 4, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=16, N=8) +#define ACCUM_FP32_16X8_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \ + MATRIX_FP32_ITER_16X8_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(4, 8, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=16, N=16) +#define ACCUM_FP32_16X16_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \ + MATRIX_FP32_ITER_16X16_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(4, 16, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=32, N=1) +#define ACCUM_FP32_32X1_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 1, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + dim; q != qe; m += 32, ++q) { \ + MATRIX_FP32_ITER_32X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(8, 1, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=32, N=2) +#define ACCUM_FP32_32X2_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \ + MATRIX_FP32_ITER_32X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(8, 2, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=32, N=4) +#define ACCUM_FP32_32X4_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 4, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \ + MATRIX_FP32_ITER_32X4_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(8, 4, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=32, N=8) +#define ACCUM_FP32_32X8_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 8, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \ + MATRIX_FP32_ITER_32X8_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(8, 8, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=32, N=16) +#define ACCUM_FP32_32X16_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 16, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \ + MATRIX_FP32_ITER_32X16_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(8, 16, 4, v_sum, out, vst1q_f32, _NORM) + +//! Compute the distance between matrix and query (FP32, M=32, N=32) +#define ACCUM_FP32_32X32_NEON(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 32, float32x4_t, v_sum, vdupq_n_f32(0)) \ + for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \ + MATRIX_FP32_ITER_32X32_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \ + } \ + MATRIX_VAR_STORE(8, 32, 4, v_sum, out, vst1q_f32, _NORM) diff --git a/src/ailego/math/distance_matrix_accum_int4.i b/src/ailego/math/distance_matrix_accum_int4.i new file mode 100644 index 0000000..3f377a8 --- /dev/null +++ b/src/ailego/math/distance_matrix_accum_int4.i @@ -0,0 +1,1084 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_int32.i" +#include "matrix_utility.i" + +//! Compute the distance between matrix and query (INT4, M=2, N=1) +#define ACCUM_INT4_2X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe_aligned = qi + ((dim >> 5) << 2); \ + const uint32_t *qe = qi + (dim >> 3); \ + if (((uintptr_t)mi & 0xf) == 0 && ((uintptr_t)qi & 0xf) == 0) { \ + for (; qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + if (qe >= qe_aligned + 2) { \ + __m128i xmm_mi = _mm_load_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + mi += 4; \ + qi += 2; \ + } \ + } else { \ + for (; qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + if (qe >= qe_aligned + 2) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + mi += 4; \ + qi += 2; \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_0_1); \ + xmm_sum_0_0 = _mm_add_epi32( \ + xmm_sum_0_0, _mm_shuffle_epi32(xmm_sum_0_0, _MM_SHUFFLE(0, 0, 3, 2))); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_set_epi32(0, 0, mi[1], mi[0]); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0)); + +//! Compute the distance between matrix and query (INT4, M=2, N=2) +#define ACCUM_INT4_2X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + ((dim >> 3) << 1); \ + if (((uintptr_t)mi & 0xf) == 0 && ((uintptr_t)qi & 0xf) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \ + qi != qe_aligned; mi += 4, qi += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \ + qi != qe_aligned; mi += 4, qi += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_set_epi32(mi[1], mi[0], mi[1], mi[0]); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=4, N=1) +#define ACCUM_INT4_4X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + (dim >> 3); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 1); \ + qi != qe_aligned; mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_load_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 1); \ + qi != qe_aligned; mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_1_0); \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=4, N=2) +#define ACCUM_INT4_4X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 4, qi += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 4, qi += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=4, N=4) +#define ACCUM_INT4_4X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 4, qi += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 4, qi += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=8, N=1) +#define ACCUM_INT4_8X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 8, ++qi) { \ + MATRIX_INT32_ITER_8X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 8, ++qi) { \ + MATRIX_INT32_ITER_8X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=8, N=2) +#define ACCUM_INT4_8X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=8, N=4) +#define ACCUM_INT4_8X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=8, N=8) +#define ACCUM_INT4_8X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=1) +#define ACCUM_INT4_16X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 16, ++qi) { \ + MATRIX_INT32_ITER_16X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 16, ++qi) { \ + MATRIX_INT32_ITER_16X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=2) +#define ACCUM_INT4_16X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 16, qi += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 16, qi += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=4) +#define ACCUM_INT4_16X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 16, qi += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 16, qi += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=8) +#define ACCUM_INT4_16X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 16, qi += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 16, qi += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=16) +#define ACCUM_INT4_16X16_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \ + mi += 16, qi += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \ + mi += 16, qi += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=1) +#define ACCUM_INT4_32X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 32, ++qi) { \ + MATRIX_INT32_ITER_32X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 32, ++qi) { \ + MATRIX_INT32_ITER_32X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=2) +#define ACCUM_INT4_32X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 32, qi += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 32, qi += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=4) +#define ACCUM_INT4_32X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 32, qi += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 32, qi += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=8) +#define ACCUM_INT4_32X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 32, qi += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 32, qi += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=16) +#define ACCUM_INT4_32X16_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 16, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \ + mi += 32, qi += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \ + mi += 32, qi += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=32) +#define ACCUM_INT4_32X32_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 32, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 5); qi != qe; \ + mi += 32, qi += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 5); qi != qe; \ + mi += 32, qi += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT4_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=2, N=1) +#define ACCUM_INT4_2X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe_aligned = qi + ((dim >> 5) << 2); \ + const uint32_t *qe = qi + (dim >> 3); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (; qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (; qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0 = _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + if (qe >= qe_aligned + 2) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0) \ + mi += 4; \ + qi += 2; \ + } \ + xmm_sum_0 = _mm_add_epi32( \ + xmm_sum_0, _mm_shuffle_epi32(xmm_sum_0, _MM_SHUFFLE(0, 0, 3, 2))); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_set_epi32(0, 0, mi[1], mi[0]); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0)); + +//! Compute the distance between matrix and query (INT4, M=2, N=2) +#define ACCUM_INT4_2X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe_aligned = qi + ((dim >> 5) << 3); \ + const uint32_t *qe = qi + ((dim >> 3) << 1); \ + if (((uintptr_t)mi & 0x1f) == 0 && ((uintptr_t)qi & 0x1f) == 0) { \ + for (; qi != qe_aligned; mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (; qi != qe_aligned; mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + if (qe >= qe_aligned + 4) { \ + __m128i xmm_qi = _mm_loadu_si128((const __m128i *)(qi)); \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_pi, xmm_sum_0_0) \ + xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_pi, xmm_sum_0_1) \ + mi += 4; \ + qi += 4; \ + } \ + xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_set_epi32(mi[1], mi[0], mi[1], mi[0]); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=4, N=1) +#define ACCUM_INT4_4X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + (dim >> 3); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 1); \ + qi != qe_aligned; mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 1); \ + qi != qe_aligned; mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=4, N=2) +#define ACCUM_INT4_4X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + ((dim >> 3) << 1); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \ + qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \ + qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_1) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=4, N=4) +#define ACCUM_INT4_4X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + ((dim >> 3) << 2); \ + if (((uintptr_t)mi & 0x1f) == 0 && ((uintptr_t)qi & 0x1f) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 3); \ + qi != qe_aligned; mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 3); \ + qi != qe_aligned; mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + __m128i xmm_sum_0_2 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_2), \ + _mm256_extracti128_si256(ymm_sum_0_2, 1)); \ + __m128i xmm_sum_0_3 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_3), \ + _mm256_extracti128_si256(ymm_sum_0_3, 1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_1) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_2) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_3) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=8, N=1) +#define ACCUM_INT4_8X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 8, ++qi) { \ + MATRIX_INT32_ITER_8X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 8, ++qi) { \ + MATRIX_INT32_ITER_8X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=8, N=2) +#define ACCUM_INT4_8X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=8, N=4) +#define ACCUM_INT4_8X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=8, N=8) +#define ACCUM_INT4_8X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=1) +#define ACCUM_INT4_16X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 16, ++qi) { \ + MATRIX_INT32_ITER_16X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 16, ++qi) { \ + MATRIX_INT32_ITER_16X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=2) +#define ACCUM_INT4_16X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 16, qi += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 16, qi += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=4) +#define ACCUM_INT4_16X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 16, qi += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 16, qi += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=8) +#define ACCUM_INT4_16X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 16, qi += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 16, qi += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=16, N=16) +#define ACCUM_INT4_16X16_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \ + mi += 16, qi += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \ + mi += 16, qi += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=1) +#define ACCUM_INT4_32X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 32, ++qi) { \ + MATRIX_INT32_ITER_32X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 32, ++qi) { \ + MATRIX_INT32_ITER_32X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=2) +#define ACCUM_INT4_32X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 32, qi += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \ + mi += 32, qi += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=4) +#define ACCUM_INT4_32X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 32, qi += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \ + mi += 32, qi += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=8) +#define ACCUM_INT4_32X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 32, qi += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \ + mi += 32, qi += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=16) +#define ACCUM_INT4_32X16_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \ + mi += 32, qi += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \ + mi += 32, qi += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT4, M=32, N=32) +#define ACCUM_INT4_32X32_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 32, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 5); qi != qe; \ + mi += 32, qi += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 3) << 5); qi != qe; \ + mi += 32, qi += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT4_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } diff --git a/src/ailego/math/distance_matrix_accum_int8.i b/src/ailego/math/distance_matrix_accum_int8.i new file mode 100644 index 0000000..b6d178f --- /dev/null +++ b/src/ailego/math/distance_matrix_accum_int8.i @@ -0,0 +1,1084 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_int32.i" +#include "matrix_utility.i" + +//! Compute the distance between matrix and query (INT8, M=2, N=1) +#define ACCUM_INT8_2X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \ + const uint32_t *qe = qi + (dim >> 2); \ + if (((uintptr_t)mi & 0xf) == 0 && ((uintptr_t)qi & 0xf) == 0) { \ + for (; qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + if (qe >= qe_aligned + 2) { \ + __m128i xmm_mi = _mm_load_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + mi += 4; \ + qi += 2; \ + } \ + } else { \ + for (; qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + if (qe >= qe_aligned + 2) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + mi += 4; \ + qi += 2; \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_0_1); \ + xmm_sum_0_0 = _mm_add_epi32( \ + xmm_sum_0_0, _mm_shuffle_epi32(xmm_sum_0_0, _MM_SHUFFLE(0, 0, 3, 2))); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_set_epi32(0, 0, mi[1], mi[0]); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0)); + +//! Compute the distance between matrix and query (INT8, M=2, N=2) +#define ACCUM_INT8_2X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + ((dim >> 2) << 1); \ + if (((uintptr_t)mi & 0xf) == 0 && ((uintptr_t)qi & 0xf) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 2); \ + qi != qe_aligned; mi += 4, qi += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 2); \ + qi != qe_aligned; mi += 4, qi += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_set_epi32(mi[1], mi[0], mi[1], mi[0]); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=4, N=1) +#define ACCUM_INT8_4X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + (dim >> 2); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 1); \ + qi != qe_aligned; mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_load_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 1); \ + qi != qe_aligned; mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_1_0); \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=4, N=2) +#define ACCUM_INT8_4X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 4, qi += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 4, qi += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=4, N=4) +#define ACCUM_INT8_4X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 4, qi += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 4, qi += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=8, N=1) +#define ACCUM_INT8_8X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 8, ++qi) { \ + MATRIX_INT32_ITER_8X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 8, ++qi) { \ + MATRIX_INT32_ITER_8X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=8, N=2) +#define ACCUM_INT8_8X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=8, N=4) +#define ACCUM_INT8_8X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=8, N=8) +#define ACCUM_INT8_8X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=1) +#define ACCUM_INT8_16X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 16, ++qi) { \ + MATRIX_INT32_ITER_16X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 16, ++qi) { \ + MATRIX_INT32_ITER_16X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=2) +#define ACCUM_INT8_16X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 16, qi += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 16, qi += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=4) +#define ACCUM_INT8_16X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 16, qi += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 16, qi += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=8) +#define ACCUM_INT8_16X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 16, qi += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 16, qi += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=16) +#define ACCUM_INT8_16X16_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \ + mi += 16, qi += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \ + mi += 16, qi += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=1) +#define ACCUM_INT8_32X1_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 32, ++qi) { \ + MATRIX_INT32_ITER_32X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 32, ++qi) { \ + MATRIX_INT32_ITER_32X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=2) +#define ACCUM_INT8_32X2_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 32, qi += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 32, qi += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=4) +#define ACCUM_INT8_32X4_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 32, qi += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 32, qi += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=8) +#define ACCUM_INT8_32X8_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 32, qi += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 32, qi += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=16) +#define ACCUM_INT8_32X16_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 16, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \ + mi += 32, qi += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \ + mi += 32, qi += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=32) +#define ACCUM_INT8_32X32_SSE(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(8, 32, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0xf) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 5); qi != qe; \ + mi += 32, qi += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(mi, qi, xmm_sum, _mm_load_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 5); qi != qe; \ + mi += 32, qi += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \ + ACCUM_INT8_STEP_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=2, N=1) +#define ACCUM_INT8_2X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \ + const uint32_t *qe = qi + (dim >> 2); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (; qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (; qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0 = _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + if (qe >= qe_aligned + 2) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0) \ + mi += 4; \ + qi += 2; \ + } \ + xmm_sum_0 = _mm_add_epi32( \ + xmm_sum_0, _mm_shuffle_epi32(xmm_sum_0, _MM_SHUFFLE(0, 0, 3, 2))); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_set_epi32(0, 0, mi[1], mi[0]); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0)); + +//! Compute the distance between matrix and query (INT8, M=2, N=2) +#define ACCUM_INT8_2X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe_aligned = qi + ((dim >> 4) << 3); \ + const uint32_t *qe = qi + ((dim >> 2) << 1); \ + if (((uintptr_t)mi & 0x1f) == 0 && ((uintptr_t)qi & 0x1f) == 0) { \ + for (; qi != qe_aligned; mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (; qi != qe_aligned; mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + if (qe >= qe_aligned + 4) { \ + __m128i xmm_qi = _mm_loadu_si128((const __m128i *)(qi)); \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_pi, xmm_sum_0_0) \ + xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_pi, xmm_sum_0_1) \ + mi += 4; \ + qi += 4; \ + } \ + xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_set_epi32(mi[1], mi[0], mi[1], mi[0]); \ + __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=4, N=1) +#define ACCUM_INT8_4X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + (dim >> 2); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 1); \ + qi != qe_aligned; mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 1); \ + qi != qe_aligned; mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=4, N=2) +#define ACCUM_INT8_4X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + ((dim >> 2) << 1); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 2); \ + qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 2); \ + qi != qe_aligned; mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_1) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=4, N=4) +#define ACCUM_INT8_4X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + const uint32_t *qe = qi + ((dim >> 2) << 2); \ + if (((uintptr_t)mi & 0x1f) == 0 && ((uintptr_t)qi & 0x1f) == 0) { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 3); \ + qi != qe_aligned; mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 3); \ + qi != qe_aligned; mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + __m128i xmm_sum_0_2 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_2), \ + _mm256_extracti128_si256(ymm_sum_0_2, 1)); \ + __m128i xmm_sum_0_3 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_3), \ + _mm256_extracti128_si256(ymm_sum_0_3, 1)); \ + if (qi != qe) { \ + __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_1) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_2) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_3) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=8, N=1) +#define ACCUM_INT8_8X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 8, ++qi) { \ + MATRIX_INT32_ITER_8X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 8, ++qi) { \ + MATRIX_INT32_ITER_8X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=8, N=2) +#define ACCUM_INT8_8X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 8, qi += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=8, N=4) +#define ACCUM_INT8_8X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 8, qi += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=8, N=8) +#define ACCUM_INT8_8X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 8, qi += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=1) +#define ACCUM_INT8_16X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 16, ++qi) { \ + MATRIX_INT32_ITER_16X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 16, ++qi) { \ + MATRIX_INT32_ITER_16X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=2) +#define ACCUM_INT8_16X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 16, qi += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 16, qi += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=4) +#define ACCUM_INT8_16X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 16, qi += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 16, qi += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=8) +#define ACCUM_INT8_16X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 16, qi += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 16, qi += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=16, N=16) +#define ACCUM_INT8_16X16_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(2, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \ + mi += 16, qi += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \ + mi += 16, qi += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=1) +#define ACCUM_INT8_32X1_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 32, ++qi) { \ + MATRIX_INT32_ITER_32X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 32, ++qi) { \ + MATRIX_INT32_ITER_32X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=2) +#define ACCUM_INT8_32X2_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 32, qi += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \ + mi += 32, qi += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=4) +#define ACCUM_INT8_32X4_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 32, qi += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \ + mi += 32, qi += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=8) +#define ACCUM_INT8_32X8_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 32, qi += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \ + mi += 32, qi += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=16) +#define ACCUM_INT8_32X16_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \ + mi += 32, qi += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \ + mi += 32, qi += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (INT8, M=32, N=32) +#define ACCUM_INT8_32X32_AVX(m, q, dim, out, _NORM) \ + MATRIX_VAR_INIT(4, 32, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qi = reinterpret_cast(q); \ + const uint32_t *mi = reinterpret_cast(m); \ + if (((uintptr_t)mi & 0x1f) == 0) { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 5); qi != qe; \ + mi += 32, qi += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(mi, qi, ymm_sum, _mm256_load_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } else { \ + for (const uint32_t *qe = qi + ((dim >> 2) << 5); qi != qe; \ + mi += 32, qi += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \ + ACCUM_INT8_STEP_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } diff --git a/src/ailego/math/distance_matrix_euclidean_utility.i b/src/ailego/math/distance_matrix_euclidean_utility.i new file mode 100644 index 0000000..b0b8937 --- /dev/null +++ b/src/ailego/math/distance_matrix_euclidean_utility.i @@ -0,0 +1,253 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Calculate sum of squared difference (GENERAL) +#define SSD_FP32_GENERAL(m, q, sum) \ + { \ + float x = m - q; \ + sum += (x * x); \ + } + +//! Calculate sum of squared difference (SSE) +#define SSD_FP32_SSE(xmm_m, xmm_q, xmm_sum) \ + { \ + __m128 xmm_d = _mm_sub_ps(xmm_m, xmm_q); \ + xmm_sum = _mm_fmadd_ps(xmm_d, xmm_d, xmm_sum); \ + } + +//! Calculate sum of squared difference (AVX) +#define SSD_FP32_AVX(ymm_m, ymm_q, ymm_sum) \ + { \ + __m256 ymm_d = _mm256_sub_ps(ymm_m, ymm_q); \ + ymm_sum = _mm256_fmadd_ps(ymm_d, ymm_d, ymm_sum); \ + } + +//! Calculate sum of squared difference (NEON) +#define SSD_FP32_NEON(v_m, v_q, v_sum) \ + { \ + float32x4_t v_d = vsubq_f32(v_m, v_q); \ + v_sum = vfmaq_f32(v_sum, v_d, v_d); \ + } + +//! Calculate sum of squared difference (GENERAL) +#define SSD_FP16_GENERAL(m, q, sum) \ + { \ + float x = m - q; \ + sum += (x * x); \ + } + +//! Calculate sum of squared difference (NEON) +#define SSD_FP16_NEON(v_m, v_q, v_sum) \ + { \ + float16x8_t v_d = vsubq_f16(v_m, v_q); \ + v_sum = vfmaq_f16(v_sum, v_d, v_d); \ + } + +//! Calculate sum of squared difference (AVX512) +#define SSD_FP32_AVX512(zmm_m, zmm_q, zmm_sum) \ + { \ + __m512 zmm_d = _mm512_sub_ps(zmm_m, zmm_q); \ + zmm_sum = _mm512_fmadd_ps(zmm_d, zmm_d, zmm_sum); \ + } + +//! Calculate sum of squared difference (GENERAL) +#define SSD_INT4_GENERAL(m, q, sum) \ + sum += Int4SquaredDiffTable[(((m) << 4) & 0xf0) | (((q) >> 0) & 0xf)] + \ + Int4SquaredDiffTable[(((m) >> 0) & 0xf0) | (((q) >> 4) & 0xf)]; + + +#if defined(__SSE4_1__) +static const __m128i MASK_INT4_SSE = _mm_set1_epi32(0xf0f0f0f0); +static const __m128i ONES_INT16_SSE = _mm_set1_epi32(0x00010001); +#endif // __SSE4_1__ + +//! Compute the square root of value (SSE) +#define SQRT_FP32_SSE(v, ...) _mm_sqrt_ps(_mm_cvtepi32_ps(v)) + +#if defined(__AVX2__) +static const __m256i MASK_INT4_AVX = _mm256_set1_epi32(0xf0f0f0f0); +static const __m256i ONES_INT16_AVX = _mm256_set1_epi32(0x00010001); +#endif // __AVX2__ + +//! Calculate sum of squared difference (SSE) +#define SSD_INT4_SSE(xmm_m, xmm_q, xmm_sum) \ + { \ + __m128i xmm_lhs = \ + _mm_and_si128(_mm_slli_epi32((xmm_m), 4), MASK_INT4_SSE); \ + __m128i xmm_rhs = \ + _mm_and_si128(_mm_slli_epi32((xmm_q), 4), MASK_INT4_SSE); \ + xmm_lhs = _mm_srli_epi32(_mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs), \ + _mm_min_epi8(xmm_lhs, xmm_rhs)), \ + 4); \ + xmm_sum = _mm_add_epi32( \ + _mm_madd_epi16(_mm_maddubs_epi16(xmm_lhs, xmm_lhs), ONES_INT16_SSE), \ + xmm_sum); \ + xmm_lhs = _mm_and_si128((xmm_m), MASK_INT4_SSE); \ + xmm_rhs = _mm_and_si128((xmm_q), MASK_INT4_SSE); \ + xmm_lhs = _mm_srli_epi32(_mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs), \ + _mm_min_epi8(xmm_lhs, xmm_rhs)), \ + 4); \ + xmm_sum = _mm_add_epi32( \ + _mm_madd_epi16(_mm_maddubs_epi16(xmm_lhs, xmm_lhs), ONES_INT16_SSE), \ + xmm_sum); \ + } + +//! Compute the distance between matrix and query +#define SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) \ + { \ + __m128i xmm_lhs_0 = \ + _mm_and_si128(_mm_slli_epi32((xmm_lhs), 4), MASK_INT4_SSE); \ + __m128i xmm_rhs_0 = \ + _mm_and_si128(_mm_slli_epi32((xmm_rhs), 4), MASK_INT4_SSE); \ + __m128i xmm_lhs_1 = _mm_and_si128((xmm_lhs), MASK_INT4_SSE); \ + __m128i xmm_rhs_1 = _mm_and_si128((xmm_rhs), MASK_INT4_SSE); \ + xmm_lhs_0 = \ + _mm_srli_epi32(_mm_sub_epi8(_mm_max_epi8(xmm_lhs_0, xmm_rhs_0), \ + _mm_min_epi8(xmm_lhs_0, xmm_rhs_0)), \ + 4); \ + xmm_rhs_0 = \ + _mm_srli_epi32(_mm_sub_epi8(_mm_max_epi8(xmm_lhs_1, xmm_rhs_1), \ + _mm_min_epi8(xmm_lhs_1, xmm_rhs_1)), \ + 4); \ + xmm_lhs_0 = _mm_madd_epi16(_mm_maddubs_epi16(xmm_lhs_0, xmm_lhs_0), \ + ONES_INT16_SSE); \ + xmm_rhs_0 = _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_0, xmm_rhs_0), \ + ONES_INT16_SSE); \ + xmm_sum = _mm_add_epi32(_mm_add_epi32(xmm_lhs_0, xmm_rhs_0), xmm_sum); \ + } + +//! Calculate sum of squared difference (AVX) +#define SSD_INT4_AVX(ymm_m, ymm_q, ymm_sum) \ + { \ + __m256i ymm_lhs = \ + _mm256_and_si256(_mm256_slli_epi32((ymm_m), 4), MASK_INT4_AVX); \ + __m256i ymm_rhs = \ + _mm256_and_si256(_mm256_slli_epi32((ymm_q), 4), MASK_INT4_AVX); \ + ymm_lhs = \ + _mm256_srli_epi32(_mm256_sub_epi8(_mm256_max_epi8(ymm_lhs, ymm_rhs), \ + _mm256_min_epi8(ymm_lhs, ymm_rhs)), \ + 4); \ + ymm_sum = _mm256_add_epi32( \ + _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_lhs, ymm_lhs), \ + ONES_INT16_AVX), \ + ymm_sum); \ + ymm_lhs = _mm256_and_si256((ymm_m), MASK_INT4_AVX); \ + ymm_rhs = _mm256_and_si256((ymm_q), MASK_INT4_AVX); \ + ymm_lhs = \ + _mm256_srli_epi32(_mm256_sub_epi8(_mm256_max_epi8(ymm_lhs, ymm_rhs), \ + _mm256_min_epi8(ymm_lhs, ymm_rhs)), \ + 4); \ + ymm_sum = _mm256_add_epi32( \ + _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_lhs, ymm_lhs), \ + ONES_INT16_AVX), \ + ymm_sum); \ + } + +//! Compute the distance between matrix and query +#define SSD_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum) \ + { \ + __m256i ymm_lhs_0 = \ + _mm256_and_si256(_mm256_slli_epi32((ymm_lhs), 4), MASK_INT4_AVX); \ + __m256i ymm_rhs_0 = \ + _mm256_and_si256(_mm256_slli_epi32((ymm_rhs), 4), MASK_INT4_AVX); \ + __m256i ymm_lhs_1 = _mm256_and_si256((ymm_lhs), MASK_INT4_AVX); \ + __m256i ymm_rhs_1 = _mm256_and_si256((ymm_rhs), MASK_INT4_AVX); \ + ymm_lhs_0 = _mm256_srli_epi32( \ + _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_0, ymm_rhs_0), \ + _mm256_min_epi8(ymm_lhs_0, ymm_rhs_0)), \ + 4); \ + ymm_rhs_0 = _mm256_srli_epi32( \ + _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_1, ymm_rhs_1), \ + _mm256_min_epi8(ymm_lhs_1, ymm_rhs_1)), \ + 4); \ + ymm_lhs_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_lhs_0, ymm_lhs_0), \ + ONES_INT16_AVX); \ + ymm_rhs_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_0, ymm_rhs_0), \ + ONES_INT16_AVX); \ + ymm_sum = \ + _mm256_add_epi32(_mm256_add_epi32(ymm_lhs_0, ymm_rhs_0), ymm_sum); \ + } + +//! Calculate sum of squared difference (GENERAL) +#define SSD_INT8_GENERAL(m, q, sum) \ + { \ + int32_t x = m - q; \ + sum += static_cast(x * x); \ + } + +//! Calculate sum of squared difference (SSE) +#define SSD_INT8_SSE(xmm_m, xmm_q, xmm_sum) \ + { \ + xmm_sum = _mm_add_epi32( \ + _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_m), \ + _mm_sign_epi8(xmm_m, xmm_m)), \ + ONES_INT16_SSE), \ + xmm_sum); \ + xmm_sum = _mm_add_epi32( \ + _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_q), \ + _mm_sign_epi8(xmm_q, xmm_q)), \ + ONES_INT16_SSE), \ + xmm_sum); \ + xmm_sum = _mm_sub_epi32( \ + xmm_sum, \ + _mm_slli_epi32( \ + _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_q), \ + _mm_sign_epi8(xmm_m, xmm_q)), \ + ONES_INT16_SSE), \ + 1)); \ + } + +//! Calculate sum of squared difference (AVX) +#define SSD_INT8_AVX(ymm_m, ymm_q, ymm_sum) \ + { \ + ymm_sum = _mm256_add_epi32( \ + _mm256_madd_epi16( \ + _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_m), \ + _mm256_sign_epi8(ymm_m, ymm_m)), \ + ONES_INT16_AVX), \ + ymm_sum); \ + ymm_sum = _mm256_add_epi32( \ + _mm256_madd_epi16( \ + _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_q), \ + _mm256_sign_epi8(ymm_q, ymm_q)), \ + ONES_INT16_AVX), \ + ymm_sum); \ + ymm_sum = _mm256_sub_epi32( \ + ymm_sum, _mm256_slli_epi32( \ + _mm256_madd_epi16( \ + _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_q), \ + _mm256_sign_epi8(ymm_m, ymm_q)), \ + ONES_INT16_AVX), \ + 1)); \ + } + +//! Compute the square root of value (AVX) +#define SQRT_FP32_AVX(v, ...) _mm256_sqrt_ps(_mm256_cvtepi32_ps(v)) + +//! Compute the square root of value (AVX512) +#define SQRT_FP32_AVX512(v, ...) _mm512_sqrt_ps(_mm512_cvtepi32_ps(v)) + +#define ACCUM_FP32_STEP_SSE SSD_FP32_SSE +#define ACCUM_FP32_STEP_AVX SSD_FP32_AVX + +#define ACCUM_FP32_STEP_AVX512 SSD_FP32_AVX512 +#define ACCUM_FP16_STEP_GENERAL SSD_FP16_GENERAL + +#define ACCUM_FP16_STEP_NEON SSD_FP16_NEON +#define ACCUM_FP32_STEP_NEON SSD_FP32_NEON + +#define ACCUM_INT4_STEP_SSE SSD_INT4_SSE +#define ACCUM_INT4_STEP_AVX SSD_INT4_AVX +#define ACCUM_INT8_STEP_SSE SSD_INT8_SSE +#define ACCUM_INT8_STEP_AVX SSD_INT8_AVX \ No newline at end of file diff --git a/src/ailego/math/distance_matrix_fp16.i b/src/ailego/math/distance_matrix_fp16.i new file mode 100644 index 0000000..11f05eb --- /dev/null +++ b/src/ailego/math/distance_matrix_fp16.i @@ -0,0 +1,861 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "matrix_define.i" +#include +#if !defined(__AVX__) +#define _mm_broadcast_si32(a) _mm_castps_si128(_mm_load1_ps((const float *)(a))) +#else +#define _mm_broadcast_si32(a) \ + _mm_castps_si128(_mm_broadcast_ss((const float *)(a))) +#define _mm256_broadcast_si32(a) \ + _mm256_castps_si256(_mm256_broadcast_ss((const float *)(a))) +#endif // !__AVX__ + +//! Mask process of computing distance (FP16) +#define MATRIX_FP16_MASK_AVX(lhs, rhs, cnt, _MASK, _RES, _PROC) \ + switch (cnt) { \ + case 7: { \ + __m256 ymm_lhs = _mm256_cvtph_ps(_mm_set_epi16( \ + (short)(_MASK), *((const short *)(lhs) + 6), \ + *((const short *)(lhs) + 5), *((const short *)(lhs) + 4), \ + *((const short *)(lhs) + 3), *((const short *)(lhs) + 2), \ + *((const short *)(lhs) + 1), *((const short *)(lhs)))); \ + __m256 ymm_rhs = _mm256_cvtph_ps(_mm_set_epi16( \ + (short)(_MASK), *((const short *)(rhs) + 6), \ + *((const short *)(rhs) + 5), *((const short *)(rhs) + 4), \ + *((const short *)(rhs) + 3), *((const short *)(rhs) + 2), \ + *((const short *)(rhs) + 1), *((const short *)(rhs)))); \ + _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \ + break; \ + } \ + case 6: { \ + __m256 ymm_lhs = _mm256_cvtph_ps( \ + _mm_set_epi32((int)(_MASK), *((const int *)(lhs) + 2), \ + *((const int *)(lhs) + 1), *((const int *)(lhs)))); \ + __m256 ymm_rhs = _mm256_cvtph_ps( \ + _mm_set_epi32((int)(_MASK), *((const int *)(rhs) + 2), \ + *((const int *)(rhs) + 1), *((const int *)(rhs)))); \ + _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \ + break; \ + } \ + case 5: { \ + __m256 ymm_lhs = _mm256_cvtph_ps(_mm_set_epi16( \ + (short)(_MASK), (short)(_MASK), (short)(_MASK), \ + *((const short *)(lhs) + 4), *((const short *)(lhs) + 3), \ + *((const short *)(lhs) + 2), *((const short *)(lhs) + 1), \ + *((const short *)(lhs)))); \ + __m256 ymm_rhs = _mm256_cvtph_ps(_mm_set_epi16( \ + (short)(_MASK), (short)(_MASK), (short)(_MASK), \ + *((const short *)(rhs) + 4), *((const short *)(rhs) + 3), \ + *((const short *)(rhs) + 2), *((const short *)(rhs) + 1), \ + *((const short *)(rhs)))); \ + _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \ + break; \ + } \ + case 4: { \ + __m256 ymm_lhs = _mm256_cvtph_ps( \ + _mm_set_epi64x((long long)(_MASK), *(const long long *)(lhs))); \ + __m256 ymm_rhs = _mm256_cvtph_ps( \ + _mm_set_epi64x((long long)(_MASK), *(const long long *)(rhs))); \ + _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \ + break; \ + } \ + case 3: { \ + __m256 ymm_lhs = _mm256_cvtph_ps(_mm_set_epi16( \ + (short)(_MASK), (short)(_MASK), (short)(_MASK), (short)(_MASK), \ + (short)(_MASK), *((const short *)(lhs) + 2), \ + *((const short *)(lhs) + 1), *((const short *)(lhs)))); \ + __m256 ymm_rhs = _mm256_cvtph_ps(_mm_set_epi16( \ + (short)(_MASK), (short)(_MASK), (short)(_MASK), (short)(_MASK), \ + (short)(_MASK), *((const short *)(rhs) + 2), \ + *((const short *)(rhs) + 1), *((const short *)(rhs)))); \ + _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \ + break; \ + } \ + case 2: { \ + __m256 ymm_lhs = _mm256_cvtph_ps(_mm_set_epi32( \ + (int)(_MASK), (int)(_MASK), (int)(_MASK), *((const int *)(lhs)))); \ + __m256 ymm_rhs = _mm256_cvtph_ps(_mm_set_epi32( \ + (int)(_MASK), (int)(_MASK), (int)(_MASK), *((const int *)(rhs)))); \ + _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \ + break; \ + } \ + case 1: { \ + __m256 ymm_lhs = _mm256_cvtph_ps( \ + _mm_set_epi16(*((const short *)(lhs)), (short)(_MASK), \ + (short)(_MASK), (short)(_MASK), (short)(_MASK), \ + (short)(_MASK), (short)(_MASK), (short)(_MASK))); \ + __m256 ymm_rhs = _mm256_cvtph_ps( \ + _mm_set_epi16(*((const short *)(rhs)), (short)(_MASK), \ + (short)(_MASK), (short)(_MASK), (short)(_MASK), \ + (short)(_MASK), (short)(_MASK), (short)(_MASK))); \ + _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \ + break; \ + } \ + } + +//! Iterative process of computing distance (FP16, M=1, N=1) +#define MATRIX_FP16_ITER_1X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)m); \ + __m256i ymm_qi = _LOAD((const __m256i *)q); \ + __m256 ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_q = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_qi)); \ + _PROC(ymm_m, ymm_q, _RES##_0_0); \ + ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + ymm_q = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_qi, 1)); \ + _PROC(ymm_m, ymm_q, _RES##_0_0); \ + } + +//! Iterative process of computing distance (FP16, M=2, N=1) +#define MATRIX_FP16_ITER_2X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m256 ymm_q = _mm256_cvtph_ps(_mm_shufflehi_epi16( \ + _mm_shufflelo_epi16(_mm_set1_epi64x(*(const long long *)(q)), \ + _MM_SHUFFLE(1, 1, 0, 0)), \ + _MM_SHUFFLE(3, 3, 2, 2))); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + } + +//! Iterative process of computing distance (FP16, M=2, N=2) +#define MATRIX_FP16_ITER_2X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_q = _mm256_cvtph_ps(_LOAD((const __m128i *)(q))); \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m256 ymm_p = _mm256_moveldup_ps(ymm_q); \ + _PROC(ymm_m, ymm_p, _RES##_0_0) \ + ymm_p = _mm256_movehdup_ps(ymm_q); \ + _PROC(ymm_m, ymm_p, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP16, M=4, N=1) +#define MATRIX_FP16_ITER_4X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m256 ymm_q = _mm256_cvtph_ps( \ + _mm_shufflehi_epi16(_mm_shufflelo_epi16(_mm_broadcast_si32(q), 0), \ + _MM_SHUFFLE(1, 1, 1, 1))); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + } + +//! Iterative process of computing distance (FP16, M=4, N=2) +#define MATRIX_FP16_ITER_4X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi = _mm_set1_epi64x(*(const long long *)(q)); \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m256 ymm_q_0 = _mm256_cvtph_ps(_mm_shufflehi_epi16( \ + _mm_shufflelo_epi16(xmm_qi, _MM_SHUFFLE(0, 0, 0, 0)), \ + _MM_SHUFFLE(2, 2, 2, 2))); \ + __m256 ymm_q_1 = _mm256_cvtph_ps(_mm_shufflehi_epi16( \ + _mm_shufflelo_epi16(xmm_qi, _MM_SHUFFLE(1, 1, 1, 1)), \ + _MM_SHUFFLE(3, 3, 3, 3))); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=4, N=4) +#define MATRIX_FP16_ITER_4X4_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m256 ymm_q = _mm256_cvtph_ps(_LOAD((const __m128i *)(q))); \ + __m256 ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(0, 0, 0, 0)); \ + _PROC(ymm_m, ymm_p, _RES##_0_0) \ + ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(1, 1, 1, 1)); \ + _PROC(ymm_m, ymm_p, _RES##_0_1) \ + ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(2, 2, 2, 2)); \ + _PROC(ymm_m, ymm_p, _RES##_0_2) \ + ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(3, 3, 3, 3)); \ + _PROC(ymm_m, ymm_p, _RES##_0_3) \ + } + +//! Iterative process of computing distance (FP16, M=8, N=1) +#define MATRIX_FP16_ITER_8X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m256 ymm_q = _mm256_cvtph_ps(_mm_set1_epi16(*(const short *)(q))); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + } + +//! Iterative process of computing distance (FP16, M=8, N=2) +#define MATRIX_FP16_ITER_8X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \ + __m256 ymm_q_0 = _mm256_set1_ps(xmm_p[0]); \ + __m256 ymm_q_1 = _mm256_set1_ps(xmm_p[1]); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=8, N=4) +#define MATRIX_FP16_ITER_8X4_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \ + __m256 ymm_q = _mm256_set1_ps(xmm_p[0]); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + ymm_q = _mm256_set1_ps(xmm_p[1]); \ + _PROC(ymm_m, ymm_q, _RES##_0_1) \ + ymm_q = _mm256_set1_ps(xmm_p[2]); \ + _PROC(ymm_m, ymm_q, _RES##_0_2) \ + ymm_q = _mm256_set1_ps(xmm_p[3]); \ + _PROC(ymm_m, ymm_q, _RES##_0_3) \ + } + +//! Iterative process of computing distance (FP16, M=8, N=8) +#define MATRIX_FP16_ITER_8X8_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \ + __m256 ymm_p = _mm256_cvtph_ps(_LOAD((const __m128i *)(q))); \ + __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + _PROC(ymm_m, ymm_q, _RES##_0_1) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + _PROC(ymm_m, ymm_q, _RES##_0_2) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + _PROC(ymm_m, ymm_q, _RES##_0_3) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + _PROC(ymm_m, ymm_q, _RES##_0_4) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + _PROC(ymm_m, ymm_q, _RES##_0_5) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + _PROC(ymm_m, ymm_q, _RES##_0_6) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + _PROC(ymm_m, ymm_q, _RES##_0_7) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=1) +#define MATRIX_FP16_ITER_16X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_q = _mm256_cvtph_ps(_mm_set1_epi16(*(const short *)q)); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=2) +#define MATRIX_FP16_ITER_16X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \ + __m256 ymm_q_0 = _mm256_set1_ps(xmm_p[0]); \ + __m256 ymm_q_1 = _mm256_set1_ps(xmm_p[1]); \ + __m256 ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + MATRIX_VAR_PROC(1, 2, 1, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=4) +#define MATRIX_FP16_ITER_16X4_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \ + __m256 ymm_q = _mm256_set1_ps(xmm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(xmm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(xmm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(xmm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=8) +#define MATRIX_FP16_ITER_16X8_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \ + __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(2, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(2, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(2, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(2, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=16) +#define MATRIX_FP16_ITER_16X16_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \ + __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(2, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(2, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(2, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(2, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 8))); \ + ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 8, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 9, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 10, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 11, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(2, 1, 12, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(2, 1, 13, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(2, 1, 14, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(2, 1, 15, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=1) +#define MATRIX_FP16_ITER_32X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + ymm_mi = _LOAD((const __m256i *)(m + 16)); \ + __m256 ymm_m_2 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_3 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + __m256 ymm_q = _mm256_cvtph_ps(_mm_set1_epi16(*(const short *)q)); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=2) +#define MATRIX_FP16_ITER_32X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \ + __m256 ymm_q_0 = _mm256_set1_ps(xmm_p[0]); \ + __m256 ymm_q_1 = _mm256_set1_ps(xmm_p[1]); \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + MATRIX_VAR_PROC(1, 2, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(m + 16)); \ + ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + MATRIX_VAR_PROC(1, 2, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + MATRIX_VAR_PROC(1, 2, 3, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=4) +#define MATRIX_FP16_ITER_32X4_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \ + __m256 ymm_q_0 = _mm256_set1_ps(xmm_p[0]); \ + __m256 ymm_q_1 = _mm256_set1_ps(xmm_p[1]); \ + __m256 ymm_q_2 = _mm256_set1_ps(xmm_p[2]); \ + __m256 ymm_q_3 = _mm256_set1_ps(xmm_p[3]); \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + MATRIX_VAR_PROC(1, 4, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + MATRIX_VAR_PROC(1, 4, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(m + 16)); \ + ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + MATRIX_VAR_PROC(1, 4, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + MATRIX_VAR_PROC(1, 4, 3, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=8) +#define MATRIX_FP16_ITER_32X8_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + ymm_mi = _LOAD((const __m256i *)(m + 16)); \ + __m256 ymm_m_2 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_3 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \ + __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=16) +#define MATRIX_FP16_ITER_32X16_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + ymm_mi = _LOAD((const __m256i *)(m + 16)); \ + __m256 ymm_m_2 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_3 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \ + __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 8))); \ + ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(4, 1, 8, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(4, 1, 9, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(4, 1, 10, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(4, 1, 11, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(4, 1, 12, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(4, 1, 13, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(4, 1, 14, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(4, 1, 15, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=32) +#define MATRIX_FP16_ITER_32X32_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(m)); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + ymm_mi = _LOAD((const __m256i *)(m + 16)); \ + __m256 ymm_m_2 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_3 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \ + __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 8))); \ + ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(4, 1, 8, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(4, 1, 9, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(4, 1, 10, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(4, 1, 11, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(4, 1, 12, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(4, 1, 13, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(4, 1, 14, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(4, 1, 15, ymm_m, ymm_q, _RES, _PROC) \ + ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 16))); \ + ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(4, 1, 16, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(4, 1, 17, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(4, 1, 18, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(4, 1, 19, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(4, 1, 20, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(4, 1, 21, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(4, 1, 22, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(4, 1, 23, ymm_m, ymm_q, _RES, _PROC) \ + ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 24))); \ + ymm_q = _mm256_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(4, 1, 24, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(4, 1, 25, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(4, 1, 26, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(4, 1, 27, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(4, 1, 28, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(4, 1, 29, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(4, 1, 30, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(4, 1, 31, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=1, N=1) +#define MATRIX_FP16_ITER_1X1_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512i zmm_mi = _LOAD((const __m512i *)m); \ + __m512i zmm_qi = _LOAD((const __m512i *)q); \ + __m512 zmm_m = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_q = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_qi)); \ + _PROC(zmm_m, zmm_q, _RES##_0_0); \ + zmm_m = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + zmm_q = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_qi, 1)); \ + _PROC(zmm_m, zmm_q, _RES##_0_0); \ + } + +//! Iterative process of computing distance (FP16, M=16, N=1) +#define MATRIX_FP16_ITER_16X1_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \ + __m512 zmm_q = _mm512_cvtph_ps(_mm256_set1_epi16(*(const short *)q)); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=2) +#define MATRIX_FP16_ITER_16X2_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \ + __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \ + __m512 zmm_q_0 = _mm512_set1_ps(xmm_p[0]); \ + __m512 zmm_q_1 = _mm512_set1_ps(xmm_p[1]); \ + MATRIX_VAR_PROC(1, 2, 0, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=4) +#define MATRIX_FP16_ITER_16X4_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \ + __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \ + __m512 zmm_q = _mm512_set1_ps(xmm_p[0]); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + zmm_q = _mm512_set1_ps(xmm_p[1]); \ + _PROC(zmm_m, zmm_q, _RES##_0_1) \ + zmm_q = _mm512_set1_ps(xmm_p[2]); \ + _PROC(zmm_m, zmm_q, _RES##_0_2) \ + zmm_q = _mm512_set1_ps(xmm_p[3]); \ + _PROC(zmm_m, zmm_q, _RES##_0_3) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=8) +#define MATRIX_FP16_ITER_16X8_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \ + __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \ + __m512 zmm_q = _mm512_set1_ps(ymm_p[0]); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + zmm_q = _mm512_set1_ps(ymm_p[1]); \ + _PROC(zmm_m, zmm_q, _RES##_0_1) \ + zmm_q = _mm512_set1_ps(ymm_p[2]); \ + _PROC(zmm_m, zmm_q, _RES##_0_2) \ + zmm_q = _mm512_set1_ps(ymm_p[3]); \ + _PROC(zmm_m, zmm_q, _RES##_0_3) \ + zmm_q = _mm512_set1_ps(ymm_p[4]); \ + _PROC(zmm_m, zmm_q, _RES##_0_4) \ + zmm_q = _mm512_set1_ps(ymm_p[5]); \ + _PROC(zmm_m, zmm_q, _RES##_0_5) \ + zmm_q = _mm512_set1_ps(ymm_p[6]); \ + _PROC(zmm_m, zmm_q, _RES##_0_6) \ + zmm_q = _mm512_set1_ps(ymm_p[7]); \ + _PROC(zmm_m, zmm_q, _RES##_0_7) \ + } + +//! Iterative process of computing distance (FP16, M=16, N=16) +#define MATRIX_FP16_ITER_16X16_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \ + __m512 zmm_p = _mm512_cvtph_ps(_LOAD((const __m256i *)(q))); \ + __m512 zmm_q = _mm512_set1_ps(zmm_p[0]); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + zmm_q = _mm512_set1_ps(zmm_p[1]); \ + _PROC(zmm_m, zmm_q, _RES##_0_1) \ + zmm_q = _mm512_set1_ps(zmm_p[2]); \ + _PROC(zmm_m, zmm_q, _RES##_0_2) \ + zmm_q = _mm512_set1_ps(zmm_p[3]); \ + _PROC(zmm_m, zmm_q, _RES##_0_3) \ + zmm_q = _mm512_set1_ps(zmm_p[4]); \ + _PROC(zmm_m, zmm_q, _RES##_0_4) \ + zmm_q = _mm512_set1_ps(zmm_p[5]); \ + _PROC(zmm_m, zmm_q, _RES##_0_5) \ + zmm_q = _mm512_set1_ps(zmm_p[6]); \ + _PROC(zmm_m, zmm_q, _RES##_0_6) \ + zmm_q = _mm512_set1_ps(zmm_p[7]); \ + _PROC(zmm_m, zmm_q, _RES##_0_7) \ + zmm_q = _mm512_set1_ps(zmm_p[8]); \ + _PROC(zmm_m, zmm_q, _RES##_0_8) \ + zmm_q = _mm512_set1_ps(zmm_p[9]); \ + _PROC(zmm_m, zmm_q, _RES##_0_9) \ + zmm_q = _mm512_set1_ps(zmm_p[10]); \ + _PROC(zmm_m, zmm_q, _RES##_0_10) \ + zmm_q = _mm512_set1_ps(zmm_p[11]); \ + _PROC(zmm_m, zmm_q, _RES##_0_11) \ + zmm_q = _mm512_set1_ps(zmm_p[12]); \ + _PROC(zmm_m, zmm_q, _RES##_0_12) \ + zmm_q = _mm512_set1_ps(zmm_p[13]); \ + _PROC(zmm_m, zmm_q, _RES##_0_13) \ + zmm_q = _mm512_set1_ps(zmm_p[14]); \ + _PROC(zmm_m, zmm_q, _RES##_0_14) \ + zmm_q = _mm512_set1_ps(zmm_p[15]); \ + _PROC(zmm_m, zmm_q, _RES##_0_15) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=1) +#define MATRIX_FP16_ITER_32X1_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512i zmm_mi = _LOAD((const __m512i *)(m)); \ + __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + __m512 zmm_q = _mm512_cvtph_ps(_mm256_set1_epi16(*(const short *)q)); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=2) +#define MATRIX_FP16_ITER_32X2_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512i zmm_mi = _LOAD((const __m512i *)(m)); \ + __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \ + __m512 zmm_q = _mm512_set1_ps(xmm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(xmm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=4) +#define MATRIX_FP16_ITER_32X4_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512i zmm_mi = _LOAD((const __m512i *)(m)); \ + __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \ + __m512 zmm_q = _mm512_set1_ps(xmm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(xmm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(xmm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(xmm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=8) +#define MATRIX_FP16_ITER_32X8_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512i zmm_mi = _LOAD((const __m512i *)(m)); \ + __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \ + __m512 zmm_q = _mm512_set1_ps(ymm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(ymm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(ymm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(ymm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(ymm_p[4]); \ + MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(ymm_p[5]); \ + MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(ymm_p[6]); \ + MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(ymm_p[7]); \ + MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=16) +#define MATRIX_FP16_ITER_32X16_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512i zmm_mi = _LOAD((const __m512i *)(m)); \ + __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + __m512 zmm_p = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(q))); \ + __m512 zmm_q = _mm512_set1_ps(zmm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[4]); \ + MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[5]); \ + MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[6]); \ + MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[7]); \ + MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[8]); \ + MATRIX_VAR_PROC(2, 1, 8, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[9]); \ + MATRIX_VAR_PROC(2, 1, 9, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[10]); \ + MATRIX_VAR_PROC(2, 1, 10, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[11]); \ + MATRIX_VAR_PROC(2, 1, 11, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[12]); \ + MATRIX_VAR_PROC(2, 1, 12, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[13]); \ + MATRIX_VAR_PROC(2, 1, 13, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[14]); \ + MATRIX_VAR_PROC(2, 1, 14, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[15]); \ + MATRIX_VAR_PROC(2, 1, 15, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP16, M=32, N=32) +#define MATRIX_FP16_ITER_32X32_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512i zmm_mi = _LOAD((const __m512i *)(m)); \ + __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + __m512 zmm_p = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(q))); \ + __m512 zmm_q = _mm512_set1_ps(zmm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[4]); \ + MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[5]); \ + MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[6]); \ + MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[7]); \ + MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[8]); \ + MATRIX_VAR_PROC(2, 1, 8, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[9]); \ + MATRIX_VAR_PROC(2, 1, 9, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[10]); \ + MATRIX_VAR_PROC(2, 1, 10, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[11]); \ + MATRIX_VAR_PROC(2, 1, 11, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[12]); \ + MATRIX_VAR_PROC(2, 1, 12, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[13]); \ + MATRIX_VAR_PROC(2, 1, 13, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[14]); \ + MATRIX_VAR_PROC(2, 1, 14, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[15]); \ + MATRIX_VAR_PROC(2, 1, 15, zmm_m, zmm_q, _RES, _PROC) \ + zmm_p = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(q + 16))); \ + zmm_q = _mm512_set1_ps(zmm_p[0]); \ + MATRIX_VAR_PROC(2, 1, 16, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[1]); \ + MATRIX_VAR_PROC(2, 1, 17, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[2]); \ + MATRIX_VAR_PROC(2, 1, 18, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[3]); \ + MATRIX_VAR_PROC(2, 1, 19, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[4]); \ + MATRIX_VAR_PROC(2, 1, 20, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[5]); \ + MATRIX_VAR_PROC(2, 1, 21, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[6]); \ + MATRIX_VAR_PROC(2, 1, 22, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[7]); \ + MATRIX_VAR_PROC(2, 1, 23, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[8]); \ + MATRIX_VAR_PROC(2, 1, 24, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[9]); \ + MATRIX_VAR_PROC(2, 1, 25, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[10]); \ + MATRIX_VAR_PROC(2, 1, 26, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[11]); \ + MATRIX_VAR_PROC(2, 1, 27, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[12]); \ + MATRIX_VAR_PROC(2, 1, 28, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[13]); \ + MATRIX_VAR_PROC(2, 1, 29, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[14]); \ + MATRIX_VAR_PROC(2, 1, 30, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(zmm_p[15]); \ + MATRIX_VAR_PROC(2, 1, 31, zmm_m, zmm_q, _RES, _PROC) \ + } + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) +//! Iterative process of computing distance (FP16, M=1, N=1) +#define MATRIX_FP16_ITER_1X1_NEON(m, q, _RES, _PROC) \ + { \ + float16x8_t v_m = vld1q_f16((const float16_t *)m); \ + float16x8_t v_q = vld1q_f16((const float16_t *)q); \ + _PROC(v_m, v_q, _RES##_0_0) \ + } + +#else +//! Iterative process of computing distance (FP16, M=1, N=1) +#define MATRIX_FP16_ITER_1X1_NEON(m, q, _RES, _PROC) \ + { \ + float16x8_t v_m = vld1q_f16((const float16_t *)m); \ + float16x8_t v_q = vld1q_f16((const float16_t *)q); \ + float32x4_t v_m_0 = vcvt_f32_f16(vget_low_f16(v_m)); \ + float32x4_t v_q_0 = vcvt_f32_f16(vget_low_f16(v_q)); \ + _PROC(v_m_0, v_q_0, _RES##_0_0) \ + v_m_0 = vcvt_high_f32_f16(v_m); \ + v_q_0 = vcvt_high_f32_f16(v_q); \ + _PROC(v_m_0, v_q_0, _RES##_0_0) \ + } + +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC \ No newline at end of file diff --git a/src/ailego/math/distance_matrix_fp32.i b/src/ailego/math/distance_matrix_fp32.i new file mode 100644 index 0000000..a9ddcd0 --- /dev/null +++ b/src/ailego/math/distance_matrix_fp32.i @@ -0,0 +1,1599 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "matrix_define.i" + +#if !defined(__AVX__) +#undef _mm_permute_ps +#define _mm_permute_ps(a, b) _mm_shuffle_ps((a), (a), (b)) +#define _mm_broadcast_ss(a) _mm_load1_ps(a) +#endif // !__AVX__ + +#if defined(__AVX__) && defined(__GNUC__) +#define _mm256_set_m128(a, b) \ + _mm256_insertf128_ps(_mm256_castps128_ps256(b), (a), 1) +#endif // __AVX__ + +#if defined(__ARM_NEON) && !defined(__aarch64__) +#define vdupq_laneq_f32(a, b) vdupq_n_f32(vgetq_lane_f32(a, b)) +#endif // __ARM_NEON && __aarch64__ + +//! Iterative process of computing distance (FP32, M=2, N=1) +#define MATRIX_FP32_ITER_2X1_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_q = _LOAD(q); \ + __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(1, 1, 0, 0)); \ + _PROC(xmm_m_0, xmm_p, _RES##_0_0) \ + xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 2, 2)); \ + _PROC(xmm_m_1, xmm_p, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=2, N=2) +#define MATRIX_FP32_ITER_2X2_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_q = _LOAD(q); \ + __m128 xmm_m = _LOAD(m); \ + __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(2, 2, 0, 0)); \ + _PROC(xmm_m, xmm_p, _RES##_0_0) \ + xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 1, 1)); \ + _PROC(xmm_m, xmm_p, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=1) +#define MATRIX_FP32_ITER_4X1_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_q = _mm_broadcast_ss(q + 0); \ + _PROC(xmm_m_0, xmm_q, _RES##_0_0) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + _PROC(xmm_m_1, xmm_q, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=2) +#define MATRIX_FP32_ITER_4X2_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m = _LOAD(m); \ + __m128 xmm_q = _mm_broadcast_ss(q + 0); \ + _PROC(xmm_m, xmm_q, _RES##_0_0) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + _PROC(xmm_m, xmm_q, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=4) +#define MATRIX_FP32_ITER_4X4_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m = _LOAD(m); \ + __m128 xmm_q = _mm_broadcast_ss(q + 0); \ + _PROC(xmm_m, xmm_q, _RES##_0_0) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + _PROC(xmm_m, xmm_q, _RES##_0_1) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + _PROC(xmm_m, xmm_q, _RES##_0_2) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + _PROC(xmm_m, xmm_q, _RES##_0_3) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=1) +#define MATRIX_FP32_ITER_8X1_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + _PROC(xmm_m_0, xmm_q, _RES##_0_0) \ + _PROC(xmm_m_1, xmm_q, _RES##_1_0) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=2) +#define MATRIX_FP32_ITER_8X2_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_q = _mm_broadcast_ss(q + 0); \ + MATRIX_VAR_PROC(2, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(2, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=4) +#define MATRIX_FP32_ITER_8X4_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_q = _mm_broadcast_ss(q + 0); \ + MATRIX_VAR_PROC(2, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(2, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(2, 1, 2, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(2, 1, 3, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=8) +#define MATRIX_FP32_ITER_8X8_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + MATRIX_VAR_PROC(2, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(2, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(2, 1, 2, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(2, 1, 3, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(2, 1, 4, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(2, 1, 5, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(2, 1, 6, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(2, 1, 7, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=1) +#define MATRIX_FP32_ITER_16X1_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=2) +#define MATRIX_FP32_ITER_16X2_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + __m128 xmm_q = _mm_broadcast_ss(q + 0); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=4) +#define MATRIX_FP32_ITER_16X4_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + __m128 xmm_q = _mm_broadcast_ss(q + 0); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(4, 1, 2, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(4, 1, 3, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=8) +#define MATRIX_FP32_ITER_16X8_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(4, 1, 2, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(4, 1, 3, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(4, 1, 4, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(4, 1, 5, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(4, 1, 6, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(4, 1, 7, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=16) +#define MATRIX_FP32_ITER_16X16_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(4, 1, 2, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(4, 1, 3, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(4, 1, 4, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(4, 1, 5, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(4, 1, 6, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(4, 1, 7, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 8); \ + MATRIX_VAR_PROC(4, 1, 8, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 9); \ + MATRIX_VAR_PROC(4, 1, 9, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 10); \ + MATRIX_VAR_PROC(4, 1, 10, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 11); \ + MATRIX_VAR_PROC(4, 1, 11, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 12); \ + MATRIX_VAR_PROC(4, 1, 12, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 13); \ + MATRIX_VAR_PROC(4, 1, 13, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 14); \ + MATRIX_VAR_PROC(4, 1, 14, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 15); \ + MATRIX_VAR_PROC(4, 1, 15, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=1) +#define MATRIX_FP32_ITER_32X1_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + _PROC(xmm_m_0, xmm_q, _RES##_0_0) \ + _PROC(xmm_m_1, xmm_q, _RES##_1_0) \ + _PROC(xmm_m_2, xmm_q, _RES##_2_0) \ + _PROC(xmm_m_3, xmm_q, _RES##_3_0) \ + xmm_m_0 = _LOAD(m + 16); \ + xmm_m_1 = _LOAD(m + 20); \ + xmm_m_2 = _LOAD(m + 24); \ + xmm_m_3 = _LOAD(m + 28); \ + _PROC(xmm_m_0, xmm_q, _RES##_4_0) \ + _PROC(xmm_m_1, xmm_q, _RES##_5_0) \ + _PROC(xmm_m_2, xmm_q, _RES##_6_0) \ + _PROC(xmm_m_3, xmm_q, _RES##_7_0) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=2) +#define MATRIX_FP32_ITER_32X2_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_q_0 = _mm_broadcast_ss(q + 0); \ + __m128 xmm_q_1 = _mm_broadcast_ss(q + 1); \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + MATRIX_VAR_PROC(1, 2, 0, xmm_m_0, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 1, xmm_m_1, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 2, xmm_m_2, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 3, xmm_m_3, xmm_q, _RES, _PROC) \ + xmm_m_0 = _LOAD(m + 16); \ + xmm_m_1 = _LOAD(m + 20); \ + xmm_m_2 = _LOAD(m + 24); \ + xmm_m_3 = _LOAD(m + 28); \ + MATRIX_VAR_PROC(1, 2, 4, xmm_m_0, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 5, xmm_m_1, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 6, xmm_m_2, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 7, xmm_m_3, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=4) +#define MATRIX_FP32_ITER_32X4_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_q_0 = _mm_broadcast_ss(q + 0); \ + __m128 xmm_q_1 = _mm_broadcast_ss(q + 1); \ + __m128 xmm_q_2 = _mm_broadcast_ss(q + 2); \ + __m128 xmm_q_3 = _mm_broadcast_ss(q + 3); \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + MATRIX_VAR_PROC(1, 4, 0, xmm_m_0, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 1, xmm_m_1, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 2, xmm_m_2, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 3, xmm_m_3, xmm_q, _RES, _PROC) \ + xmm_m_0 = _LOAD(m + 16); \ + xmm_m_1 = _LOAD(m + 20); \ + xmm_m_2 = _LOAD(m + 24); \ + xmm_m_3 = _LOAD(m + 28); \ + MATRIX_VAR_PROC(1, 4, 4, xmm_m_0, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 5, xmm_m_1, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 6, xmm_m_2, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 7, xmm_m_3, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=8) +#define MATRIX_FP32_ITER_32X8_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_q_0 = _mm_broadcast_ss(q + 0); \ + __m128 xmm_q_1 = _mm_broadcast_ss(q + 1); \ + __m128 xmm_q_2 = _mm_broadcast_ss(q + 2); \ + __m128 xmm_q_3 = _mm_broadcast_ss(q + 3); \ + __m128 xmm_q_4 = _mm_broadcast_ss(q + 4); \ + __m128 xmm_q_5 = _mm_broadcast_ss(q + 5); \ + __m128 xmm_q_6 = _mm_broadcast_ss(q + 6); \ + __m128 xmm_q_7 = _mm_broadcast_ss(q + 7); \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + MATRIX_VAR_PROC(1, 8, 0, xmm_m_0, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 1, xmm_m_1, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 2, xmm_m_2, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 3, xmm_m_3, xmm_q, _RES, _PROC) \ + xmm_m_0 = _LOAD(m + 16); \ + xmm_m_1 = _LOAD(m + 20); \ + xmm_m_2 = _LOAD(m + 24); \ + xmm_m_3 = _LOAD(m + 28); \ + MATRIX_VAR_PROC(1, 8, 4, xmm_m_0, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 5, xmm_m_1, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 6, xmm_m_2, xmm_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 7, xmm_m_3, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=16) +#define MATRIX_FP32_ITER_32X16_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + __m128 xmm_m_4 = _LOAD(m + 16); \ + __m128 xmm_m_5 = _LOAD(m + 20); \ + __m128 xmm_m_6 = _LOAD(m + 24); \ + __m128 xmm_m_7 = _LOAD(m + 28); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + MATRIX_VAR_PROC(8, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(8, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(8, 1, 2, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(8, 1, 3, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(8, 1, 4, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(8, 1, 5, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(8, 1, 6, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(8, 1, 7, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 8); \ + MATRIX_VAR_PROC(8, 1, 8, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 9); \ + MATRIX_VAR_PROC(8, 1, 9, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 10); \ + MATRIX_VAR_PROC(8, 1, 10, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 11); \ + MATRIX_VAR_PROC(8, 1, 11, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 12); \ + MATRIX_VAR_PROC(8, 1, 12, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 13); \ + MATRIX_VAR_PROC(8, 1, 13, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 14); \ + MATRIX_VAR_PROC(8, 1, 14, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 15); \ + MATRIX_VAR_PROC(8, 1, 15, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=32) +#define MATRIX_FP32_ITER_32X32_SSE(m, q, _RES, _LOAD, _PROC) \ + { \ + __m128 xmm_m_0 = _LOAD(m + 0); \ + __m128 xmm_m_1 = _LOAD(m + 4); \ + __m128 xmm_m_2 = _LOAD(m + 8); \ + __m128 xmm_m_3 = _LOAD(m + 12); \ + __m128 xmm_m_4 = _LOAD(m + 16); \ + __m128 xmm_m_5 = _LOAD(m + 20); \ + __m128 xmm_m_6 = _LOAD(m + 24); \ + __m128 xmm_m_7 = _LOAD(m + 28); \ + __m128 xmm_q = _mm_broadcast_ss(q); \ + MATRIX_VAR_PROC(8, 1, 0, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(8, 1, 1, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(8, 1, 2, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(8, 1, 3, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(8, 1, 4, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(8, 1, 5, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(8, 1, 6, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(8, 1, 7, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 8); \ + MATRIX_VAR_PROC(8, 1, 8, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 9); \ + MATRIX_VAR_PROC(8, 1, 9, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 10); \ + MATRIX_VAR_PROC(8, 1, 10, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 11); \ + MATRIX_VAR_PROC(8, 1, 11, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 12); \ + MATRIX_VAR_PROC(8, 1, 12, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 13); \ + MATRIX_VAR_PROC(8, 1, 13, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 14); \ + MATRIX_VAR_PROC(8, 1, 14, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 15); \ + MATRIX_VAR_PROC(8, 1, 15, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 16); \ + MATRIX_VAR_PROC(8, 1, 16, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 17); \ + MATRIX_VAR_PROC(8, 1, 17, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 18); \ + MATRIX_VAR_PROC(8, 1, 18, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 19); \ + MATRIX_VAR_PROC(8, 1, 19, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 20); \ + MATRIX_VAR_PROC(8, 1, 20, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 21); \ + MATRIX_VAR_PROC(8, 1, 21, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 22); \ + MATRIX_VAR_PROC(8, 1, 22, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 23); \ + MATRIX_VAR_PROC(8, 1, 23, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 24); \ + MATRIX_VAR_PROC(8, 1, 24, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 25); \ + MATRIX_VAR_PROC(8, 1, 25, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 26); \ + MATRIX_VAR_PROC(8, 1, 26, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 27); \ + MATRIX_VAR_PROC(8, 1, 27, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 28); \ + MATRIX_VAR_PROC(8, 1, 28, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 29); \ + MATRIX_VAR_PROC(8, 1, 29, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 30); \ + MATRIX_VAR_PROC(8, 1, 30, xmm_m, xmm_q, _RES, _PROC) \ + xmm_q = _mm_broadcast_ss(q + 31); \ + MATRIX_VAR_PROC(8, 1, 31, xmm_m, xmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=2, N=1) +#define MATRIX_FP32_ITER_2X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_q = \ + _mm256_set_ps(q[3], q[3], q[2], q[2], q[1], q[1], q[0], q[0]); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + } + +//! Iterative process of computing distance (FP32, M=2, N=2) +#define MATRIX_FP32_ITER_2X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_q = _LOAD(q); \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_p = _mm256_moveldup_ps(ymm_q); \ + _PROC(ymm_m, ymm_p, _RES##_0_0) \ + ymm_p = _mm256_movehdup_ps(ymm_q); \ + _PROC(ymm_m, ymm_p, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=1) +#define MATRIX_FP32_ITER_4X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_q = \ + _mm256_set_m128(_mm_broadcast_ss(q + 1), _mm_broadcast_ss(q)); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=2) +#define MATRIX_FP32_ITER_4X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_q = \ + _mm256_set_m128(_mm_broadcast_ss(q + 2), _mm_broadcast_ss(q + 0)); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + ymm_q = _mm256_set_m128(_mm_broadcast_ss(q + 3), _mm_broadcast_ss(q + 1)); \ + _PROC(ymm_m, ymm_q, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=4) +#define MATRIX_FP32_ITER_4X4_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_q = _LOAD(q); \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(0, 0, 0, 0)); \ + _PROC(ymm_m, ymm_p, _RES##_0_0) \ + ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(1, 1, 1, 1)); \ + _PROC(ymm_m, ymm_p, _RES##_0_1) \ + ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(2, 2, 2, 2)); \ + _PROC(ymm_m, ymm_p, _RES##_0_2) \ + ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(3, 3, 3, 3)); \ + _PROC(ymm_m, ymm_p, _RES##_0_3) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=1) +#define MATRIX_FP32_ITER_8X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=2) +#define MATRIX_FP32_ITER_8X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + _PROC(ymm_m, ymm_q, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=4) +#define MATRIX_FP32_ITER_8X4_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + _PROC(ymm_m, ymm_q, _RES##_0_1) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + _PROC(ymm_m, ymm_q, _RES##_0_2) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + _PROC(ymm_m, ymm_q, _RES##_0_3) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=8) +#define MATRIX_FP32_ITER_8X8_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m = _LOAD(m); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + _PROC(ymm_m, ymm_q, _RES##_0_0) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + _PROC(ymm_m, ymm_q, _RES##_0_1) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + _PROC(ymm_m, ymm_q, _RES##_0_2) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + _PROC(ymm_m, ymm_q, _RES##_0_3) \ + ymm_q = _mm256_broadcast_ss(q + 4); \ + _PROC(ymm_m, ymm_q, _RES##_0_4) \ + ymm_q = _mm256_broadcast_ss(q + 5); \ + _PROC(ymm_m, ymm_q, _RES##_0_5) \ + ymm_q = _mm256_broadcast_ss(q + 6); \ + _PROC(ymm_m, ymm_q, _RES##_0_6) \ + ymm_q = _mm256_broadcast_ss(q + 7); \ + _PROC(ymm_m, ymm_q, _RES##_0_7) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=1) +#define MATRIX_FP32_ITER_16X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=2) +#define MATRIX_FP32_ITER_16X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=4) +#define MATRIX_FP32_ITER_16X4_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=8) +#define MATRIX_FP32_ITER_16X8_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(2, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(2, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(2, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(2, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=16) +#define MATRIX_FP32_ITER_16X16_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(2, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(2, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(2, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(2, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 8); \ + MATRIX_VAR_PROC(2, 1, 8, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 9); \ + MATRIX_VAR_PROC(2, 1, 9, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 10); \ + MATRIX_VAR_PROC(2, 1, 10, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 11); \ + MATRIX_VAR_PROC(2, 1, 11, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 12); \ + MATRIX_VAR_PROC(2, 1, 12, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 13); \ + MATRIX_VAR_PROC(2, 1, 13, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 14); \ + MATRIX_VAR_PROC(2, 1, 14, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 15); \ + MATRIX_VAR_PROC(2, 1, 15, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=1) +#define MATRIX_FP32_ITER_32X1_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_m_2 = _LOAD(m + 16); \ + __m256 ymm_m_3 = _LOAD(m + 24); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=2) +#define MATRIX_FP32_ITER_32X2_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_m_2 = _LOAD(m + 16); \ + __m256 ymm_m_3 = _LOAD(m + 24); \ + __m256 ymm_q = _mm256_broadcast_ss(q + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=4) +#define MATRIX_FP32_ITER_32X4_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_m_2 = _LOAD(m + 16); \ + __m256 ymm_m_3 = _LOAD(m + 24); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=8) +#define MATRIX_FP32_ITER_32X8_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_m_2 = _LOAD(m + 16); \ + __m256 ymm_m_3 = _LOAD(m + 24); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=16) +#define MATRIX_FP32_ITER_32X16_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_m_2 = _LOAD(m + 16); \ + __m256 ymm_m_3 = _LOAD(m + 24); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 8); \ + MATRIX_VAR_PROC(4, 1, 8, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 9); \ + MATRIX_VAR_PROC(4, 1, 9, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 10); \ + MATRIX_VAR_PROC(4, 1, 10, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 11); \ + MATRIX_VAR_PROC(4, 1, 11, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 12); \ + MATRIX_VAR_PROC(4, 1, 12, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 13); \ + MATRIX_VAR_PROC(4, 1, 13, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 14); \ + MATRIX_VAR_PROC(4, 1, 14, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 15); \ + MATRIX_VAR_PROC(4, 1, 15, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=32) +#define MATRIX_FP32_ITER_32X32_AVX(m, q, _RES, _LOAD, _PROC) \ + { \ + __m256 ymm_m_0 = _LOAD(m + 0); \ + __m256 ymm_m_1 = _LOAD(m + 8); \ + __m256 ymm_m_2 = _LOAD(m + 16); \ + __m256 ymm_m_3 = _LOAD(m + 24); \ + __m256 ymm_q = _mm256_broadcast_ss(q); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 4); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 5); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 6); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 7); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 8); \ + MATRIX_VAR_PROC(4, 1, 8, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 9); \ + MATRIX_VAR_PROC(4, 1, 9, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 10); \ + MATRIX_VAR_PROC(4, 1, 10, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 11); \ + MATRIX_VAR_PROC(4, 1, 11, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 12); \ + MATRIX_VAR_PROC(4, 1, 12, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 13); \ + MATRIX_VAR_PROC(4, 1, 13, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 14); \ + MATRIX_VAR_PROC(4, 1, 14, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 15); \ + MATRIX_VAR_PROC(4, 1, 15, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 16); \ + MATRIX_VAR_PROC(4, 1, 16, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 17); \ + MATRIX_VAR_PROC(4, 1, 17, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 18); \ + MATRIX_VAR_PROC(4, 1, 18, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 19); \ + MATRIX_VAR_PROC(4, 1, 19, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 20); \ + MATRIX_VAR_PROC(4, 1, 20, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 21); \ + MATRIX_VAR_PROC(4, 1, 21, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 22); \ + MATRIX_VAR_PROC(4, 1, 22, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 23); \ + MATRIX_VAR_PROC(4, 1, 23, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 24); \ + MATRIX_VAR_PROC(4, 1, 24, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 25); \ + MATRIX_VAR_PROC(4, 1, 25, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 26); \ + MATRIX_VAR_PROC(4, 1, 26, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 27); \ + MATRIX_VAR_PROC(4, 1, 27, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 28); \ + MATRIX_VAR_PROC(4, 1, 28, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 29); \ + MATRIX_VAR_PROC(4, 1, 29, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 30); \ + MATRIX_VAR_PROC(4, 1, 30, ymm_m, ymm_q, _RES, _PROC) \ + ymm_q = _mm256_broadcast_ss(q + 31); \ + MATRIX_VAR_PROC(4, 1, 31, ymm_m, ymm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=1) +#define MATRIX_FP32_ITER_16X1_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _LOAD(m); \ + __m512 zmm_q = _mm512_set1_ps(*q); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=2) +#define MATRIX_FP32_ITER_16X2_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _LOAD(m); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + zmm_q = _mm512_set1_ps(q[1]); \ + _PROC(zmm_m, zmm_q, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=4) +#define MATRIX_FP32_ITER_16X4_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _LOAD(m); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + zmm_q = _mm512_set1_ps(q[1]); \ + _PROC(zmm_m, zmm_q, _RES##_0_1) \ + zmm_q = _mm512_set1_ps(q[2]); \ + _PROC(zmm_m, zmm_q, _RES##_0_2) \ + zmm_q = _mm512_set1_ps(q[3]); \ + _PROC(zmm_m, zmm_q, _RES##_0_3) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=8) +#define MATRIX_FP32_ITER_16X8_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _LOAD(m); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + zmm_q = _mm512_set1_ps(q[1]); \ + _PROC(zmm_m, zmm_q, _RES##_0_1) \ + zmm_q = _mm512_set1_ps(q[2]); \ + _PROC(zmm_m, zmm_q, _RES##_0_2) \ + zmm_q = _mm512_set1_ps(q[3]); \ + _PROC(zmm_m, zmm_q, _RES##_0_3) \ + zmm_q = _mm512_set1_ps(q[4]); \ + _PROC(zmm_m, zmm_q, _RES##_0_4) \ + zmm_q = _mm512_set1_ps(q[5]); \ + _PROC(zmm_m, zmm_q, _RES##_0_5) \ + zmm_q = _mm512_set1_ps(q[6]); \ + _PROC(zmm_m, zmm_q, _RES##_0_6) \ + zmm_q = _mm512_set1_ps(q[7]); \ + _PROC(zmm_m, zmm_q, _RES##_0_7) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=16) +#define MATRIX_FP32_ITER_16X16_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m = _LOAD(m); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + zmm_q = _mm512_set1_ps(q[1]); \ + _PROC(zmm_m, zmm_q, _RES##_0_1) \ + zmm_q = _mm512_set1_ps(q[2]); \ + _PROC(zmm_m, zmm_q, _RES##_0_2) \ + zmm_q = _mm512_set1_ps(q[3]); \ + _PROC(zmm_m, zmm_q, _RES##_0_3) \ + zmm_q = _mm512_set1_ps(q[4]); \ + _PROC(zmm_m, zmm_q, _RES##_0_4) \ + zmm_q = _mm512_set1_ps(q[5]); \ + _PROC(zmm_m, zmm_q, _RES##_0_5) \ + zmm_q = _mm512_set1_ps(q[6]); \ + _PROC(zmm_m, zmm_q, _RES##_0_6) \ + zmm_q = _mm512_set1_ps(q[7]); \ + _PROC(zmm_m, zmm_q, _RES##_0_7) \ + zmm_q = _mm512_set1_ps(q[8]); \ + _PROC(zmm_m, zmm_q, _RES##_0_8) \ + zmm_q = _mm512_set1_ps(q[9]); \ + _PROC(zmm_m, zmm_q, _RES##_0_9) \ + zmm_q = _mm512_set1_ps(q[10]); \ + _PROC(zmm_m, zmm_q, _RES##_0_10) \ + zmm_q = _mm512_set1_ps(q[11]); \ + _PROC(zmm_m, zmm_q, _RES##_0_11) \ + zmm_q = _mm512_set1_ps(q[12]); \ + _PROC(zmm_m, zmm_q, _RES##_0_12) \ + zmm_q = _mm512_set1_ps(q[13]); \ + _PROC(zmm_m, zmm_q, _RES##_0_13) \ + zmm_q = _mm512_set1_ps(q[14]); \ + _PROC(zmm_m, zmm_q, _RES##_0_14) \ + zmm_q = _mm512_set1_ps(q[15]); \ + _PROC(zmm_m, zmm_q, _RES##_0_15) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=1) +#define MATRIX_FP32_ITER_32X1_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_q = _mm512_set1_ps(*q); \ + __m512 zmm_m = _LOAD(m); \ + _PROC(zmm_m, zmm_q, _RES##_0_0) \ + zmm_m = _LOAD(m + 16); \ + _PROC(zmm_m, zmm_q, _RES##_1_0) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=2) +#define MATRIX_FP32_ITER_32X2_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m_0 = _LOAD(m + 0); \ + __m512 zmm_m_1 = _LOAD(m + 16); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=4) +#define MATRIX_FP32_ITER_32X4_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m_0 = _LOAD(m + 0); \ + __m512 zmm_m_1 = _LOAD(m + 16); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[2]); \ + MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[3]); \ + MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=8) +#define MATRIX_FP32_ITER_32X8_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m_0 = _LOAD(m + 0); \ + __m512 zmm_m_1 = _LOAD(m + 16); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[2]); \ + MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[3]); \ + MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[4]); \ + MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[5]); \ + MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[6]); \ + MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[7]); \ + MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=16) +#define MATRIX_FP32_ITER_32X16_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m_0 = _LOAD(m + 0); \ + __m512 zmm_m_1 = _LOAD(m + 16); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[2]); \ + MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[3]); \ + MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[4]); \ + MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[5]); \ + MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[6]); \ + MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[7]); \ + MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[8]); \ + MATRIX_VAR_PROC(2, 1, 8, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[9]); \ + MATRIX_VAR_PROC(2, 1, 9, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[10]); \ + MATRIX_VAR_PROC(2, 1, 10, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[11]); \ + MATRIX_VAR_PROC(2, 1, 11, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[12]); \ + MATRIX_VAR_PROC(2, 1, 12, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[13]); \ + MATRIX_VAR_PROC(2, 1, 13, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[14]); \ + MATRIX_VAR_PROC(2, 1, 14, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[15]); \ + MATRIX_VAR_PROC(2, 1, 15, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=32) +#define MATRIX_FP32_ITER_32X32_AVX512(m, q, _RES, _LOAD, _PROC) \ + { \ + __m512 zmm_m_0 = _LOAD(m + 0); \ + __m512 zmm_m_1 = _LOAD(m + 16); \ + __m512 zmm_q = _mm512_set1_ps(q[0]); \ + MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[1]); \ + MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[2]); \ + MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[3]); \ + MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[4]); \ + MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[5]); \ + MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[6]); \ + MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[7]); \ + MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[8]); \ + MATRIX_VAR_PROC(2, 1, 8, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[9]); \ + MATRIX_VAR_PROC(2, 1, 9, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[10]); \ + MATRIX_VAR_PROC(2, 1, 10, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[11]); \ + MATRIX_VAR_PROC(2, 1, 11, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[12]); \ + MATRIX_VAR_PROC(2, 1, 12, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[13]); \ + MATRIX_VAR_PROC(2, 1, 13, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[14]); \ + MATRIX_VAR_PROC(2, 1, 14, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[15]); \ + MATRIX_VAR_PROC(2, 1, 15, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[16]); \ + MATRIX_VAR_PROC(2, 1, 16, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[17]); \ + MATRIX_VAR_PROC(2, 1, 17, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[18]); \ + MATRIX_VAR_PROC(2, 1, 18, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[19]); \ + MATRIX_VAR_PROC(2, 1, 19, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[20]); \ + MATRIX_VAR_PROC(2, 1, 20, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[21]); \ + MATRIX_VAR_PROC(2, 1, 21, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[22]); \ + MATRIX_VAR_PROC(2, 1, 22, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[23]); \ + MATRIX_VAR_PROC(2, 1, 23, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[24]); \ + MATRIX_VAR_PROC(2, 1, 24, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[25]); \ + MATRIX_VAR_PROC(2, 1, 25, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[26]); \ + MATRIX_VAR_PROC(2, 1, 26, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[27]); \ + MATRIX_VAR_PROC(2, 1, 27, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[28]); \ + MATRIX_VAR_PROC(2, 1, 28, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[29]); \ + MATRIX_VAR_PROC(2, 1, 29, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[30]); \ + MATRIX_VAR_PROC(2, 1, 30, zmm_m, zmm_q, _RES, _PROC) \ + zmm_q = _mm512_set1_ps(q[31]); \ + MATRIX_VAR_PROC(2, 1, 31, zmm_m, zmm_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=2, N=1) +#define MATRIX_FP32_ITER_2X1_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m = vld1q_f32(m); \ + float32x2_t v_q = vld1_f32(q); \ + float32x4_t v_p = \ + vcombine_f32(vdup_lane_f32(v_q, 0), vdup_lane_f32(v_q, 1)); \ + _PROC(v_m, v_p, _RES) \ + } + +//! Iterative process of computing distance (FP32, M=2, N=2) +#define MATRIX_FP32_ITER_2X2_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_q = vld1q_f32(q); \ + float32x4_t v_m = vld1q_f32(m); \ + float32x2_t v_q_0 = vget_low_f32(v_q); \ + float32x2_t v_q_1 = vget_high_f32(v_q); \ + v_q = vcombine_f32(vdup_lane_f32(v_q_0, 0), vdup_lane_f32(v_q_1, 0)); \ + _PROC(v_m, v_q, _RES##_0_0) \ + v_q = vcombine_f32(vdup_lane_f32(v_q_0, 1), vdup_lane_f32(v_q_1, 1)); \ + _PROC(v_m, v_q, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=1) +#define MATRIX_FP32_ITER_4X1_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x2_t v_p = vld1_f32(q); \ + float32x4_t v_q = vdupq_lane_f32(v_p, 0); \ + _PROC(v_m_0, v_q, _RES##_0_0) \ + v_q = vdupq_lane_f32(v_p, 1); \ + _PROC(v_m_1, v_q, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=2) +#define MATRIX_FP32_ITER_4X2_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m = vld1q_f32(m); \ + float32x2_t v_p = vld1_f32(q); \ + float32x4_t v_q = vdupq_lane_f32(v_p, 0); \ + _PROC(v_m, v_q, _RES##_0_0) \ + v_q = vdupq_lane_f32(v_p, 1); \ + _PROC(v_m, v_q, _RES##_0_1) \ + } + +//! Iterative process of computing distance (FP32, M=4, N=4) +#define MATRIX_FP32_ITER_4X4_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m = vld1q_f32(m); \ + float32x4_t v_p = vld1q_f32(q); \ + float32x4_t v_q = vdupq_laneq_f32(v_p, 0); \ + _PROC(v_m, v_q, _RES##_0_0) \ + v_q = vdupq_laneq_f32(v_p, 1); \ + _PROC(v_m, v_q, _RES##_0_1) \ + v_q = vdupq_laneq_f32(v_p, 2); \ + _PROC(v_m, v_q, _RES##_0_2) \ + v_q = vdupq_laneq_f32(v_p, 3); \ + _PROC(v_m, v_q, _RES##_0_3) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=1) +#define MATRIX_FP32_ITER_8X1_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_q = vld1q_dup_f32(q); \ + _PROC(v_m_0, v_q, _RES##_0_0) \ + _PROC(v_m_1, v_q, _RES##_1_0) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=2) +#define MATRIX_FP32_ITER_8X2_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x2_t v_p = vld1_f32(q); \ + float32x4_t v_q = vdupq_lane_f32(v_p, 0); \ + MATRIX_VAR_PROC(2, 1, 0, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_lane_f32(v_p, 1); \ + MATRIX_VAR_PROC(2, 1, 1, v_m, v_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=4) +#define MATRIX_FP32_ITER_8X4_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_p = vld1q_f32(q); \ + float32x4_t v_q = vdupq_laneq_f32(v_p, 0); \ + MATRIX_VAR_PROC(2, 1, 0, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 1); \ + MATRIX_VAR_PROC(2, 1, 1, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 2); \ + MATRIX_VAR_PROC(2, 1, 2, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 3); \ + MATRIX_VAR_PROC(2, 1, 3, v_m, v_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=8, N=8) +#define MATRIX_FP32_ITER_8X8_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_p = vld1q_f32(q + 0); \ + float32x4_t v_q = vdupq_laneq_f32(v_p, 0); \ + MATRIX_VAR_PROC(2, 1, 0, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 1); \ + MATRIX_VAR_PROC(2, 1, 1, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 2); \ + MATRIX_VAR_PROC(2, 1, 2, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 3); \ + MATRIX_VAR_PROC(2, 1, 3, v_m, v_q, _RES, _PROC) \ + v_p = vld1q_f32(q + 4); \ + v_q = vdupq_laneq_f32(v_p, 0); \ + MATRIX_VAR_PROC(2, 1, 4, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 1); \ + MATRIX_VAR_PROC(2, 1, 5, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 2); \ + MATRIX_VAR_PROC(2, 1, 6, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_laneq_f32(v_p, 3); \ + MATRIX_VAR_PROC(2, 1, 7, v_m, v_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=1) +#define MATRIX_FP32_ITER_16X1_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + float32x4_t v_q = vld1q_dup_f32(q); \ + MATRIX_VAR_PROC(4, 1, 0, v_m, v_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=2) +#define MATRIX_FP32_ITER_16X2_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + float32x2_t v_p = vld1_f32(q); \ + float32x4_t v_q = vdupq_lane_f32(v_p, 0); \ + MATRIX_VAR_PROC(4, 1, 0, v_m, v_q, _RES, _PROC) \ + v_q = vdupq_lane_f32(v_p, 1); \ + MATRIX_VAR_PROC(4, 1, 1, v_m, v_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=4) +#define MATRIX_FP32_ITER_16X4_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + float32x4_t v_q = vld1q_f32(q); \ + float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(4, 1, 0, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(4, 1, 1, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(4, 1, 2, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(4, 1, 3, v_m, v_p, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=8) +#define MATRIX_FP32_ITER_16X8_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + float32x4_t v_q = vld1q_f32(q + 0); \ + float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(4, 1, 0, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(4, 1, 1, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(4, 1, 2, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(4, 1, 3, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 4); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(4, 1, 4, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(4, 1, 5, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(4, 1, 6, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(4, 1, 7, v_m, v_p, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=16, N=16) +#define MATRIX_FP32_ITER_16X16_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + float32x4_t v_q = vld1q_f32(q + 0); \ + float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(4, 1, 0, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(4, 1, 1, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(4, 1, 2, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(4, 1, 3, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 4); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(4, 1, 4, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(4, 1, 5, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(4, 1, 6, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(4, 1, 7, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 8); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(4, 1, 8, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(4, 1, 9, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(4, 1, 10, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(4, 1, 11, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 12); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(4, 1, 12, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(4, 1, 13, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(4, 1, 14, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(4, 1, 15, v_m, v_p, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=1) +#define MATRIX_FP32_ITER_32X1_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_q = vld1q_dup_f32(q); \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + _PROC(v_m_0, v_q, _RES##_0_0) \ + _PROC(v_m_1, v_q, _RES##_1_0) \ + _PROC(v_m_2, v_q, _RES##_2_0) \ + _PROC(v_m_3, v_q, _RES##_3_0) \ + v_m_0 = vld1q_f32(m + 16); \ + v_m_1 = vld1q_f32(m + 20); \ + v_m_2 = vld1q_f32(m + 24); \ + v_m_3 = vld1q_f32(m + 28); \ + _PROC(v_m_0, v_q, _RES##_4_0) \ + _PROC(v_m_1, v_q, _RES##_5_0) \ + _PROC(v_m_2, v_q, _RES##_6_0) \ + _PROC(v_m_3, v_q, _RES##_7_0) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=2) +#define MATRIX_FP32_ITER_32X2_NEON(m, q, _RES, _PROC) \ + { \ + float32x2_t v_p = vld1_f32(q); \ + float32x4_t v_q_0 = vdupq_lane_f32(v_p, 0); \ + float32x4_t v_q_1 = vdupq_lane_f32(v_p, 1); \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + MATRIX_VAR_PROC(1, 2, 0, v_m_0, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 1, v_m_1, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 2, v_m_2, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 3, v_m_3, v_q, _RES, _PROC) \ + v_m_0 = vld1q_f32(m + 16); \ + v_m_1 = vld1q_f32(m + 20); \ + v_m_2 = vld1q_f32(m + 24); \ + v_m_3 = vld1q_f32(m + 28); \ + MATRIX_VAR_PROC(1, 2, 4, v_m_0, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 5, v_m_1, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 6, v_m_2, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 7, v_m_3, v_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=4) +#define MATRIX_FP32_ITER_32X4_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_p = vld1q_f32(q); \ + float32x4_t v_q_0 = vdupq_laneq_f32(v_p, 0); \ + float32x4_t v_q_1 = vdupq_laneq_f32(v_p, 1); \ + float32x4_t v_q_2 = vdupq_laneq_f32(v_p, 2); \ + float32x4_t v_q_3 = vdupq_laneq_f32(v_p, 3); \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + MATRIX_VAR_PROC(1, 4, 0, v_m_0, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 1, v_m_1, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 2, v_m_2, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 3, v_m_3, v_q, _RES, _PROC) \ + v_m_0 = vld1q_f32(m + 16); \ + v_m_1 = vld1q_f32(m + 20); \ + v_m_2 = vld1q_f32(m + 24); \ + v_m_3 = vld1q_f32(m + 28); \ + MATRIX_VAR_PROC(1, 4, 4, v_m_0, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 5, v_m_1, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 6, v_m_2, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 7, v_m_3, v_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=8) +#define MATRIX_FP32_ITER_32X8_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_p_0 = vld1q_f32(q + 0); \ + float32x4_t v_p_1 = vld1q_f32(q + 4); \ + float32x4_t v_q_0 = vdupq_laneq_f32(v_p_0, 0); \ + float32x4_t v_q_1 = vdupq_laneq_f32(v_p_0, 1); \ + float32x4_t v_q_2 = vdupq_laneq_f32(v_p_0, 2); \ + float32x4_t v_q_3 = vdupq_laneq_f32(v_p_0, 3); \ + float32x4_t v_q_4 = vdupq_laneq_f32(v_p_1, 0); \ + float32x4_t v_q_5 = vdupq_laneq_f32(v_p_1, 1); \ + float32x4_t v_q_6 = vdupq_laneq_f32(v_p_1, 2); \ + float32x4_t v_q_7 = vdupq_laneq_f32(v_p_1, 3); \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + MATRIX_VAR_PROC(1, 8, 0, v_m_0, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 1, v_m_1, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 2, v_m_2, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 3, v_m_3, v_q, _RES, _PROC) \ + v_m_0 = vld1q_f32(m + 16); \ + v_m_1 = vld1q_f32(m + 20); \ + v_m_2 = vld1q_f32(m + 24); \ + v_m_3 = vld1q_f32(m + 28); \ + MATRIX_VAR_PROC(1, 8, 4, v_m_0, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 5, v_m_1, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 6, v_m_2, v_q, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 7, v_m_3, v_q, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=16) +#define MATRIX_FP32_ITER_32X16_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + float32x4_t v_m_4 = vld1q_f32(m + 16); \ + float32x4_t v_m_5 = vld1q_f32(m + 20); \ + float32x4_t v_m_6 = vld1q_f32(m + 24); \ + float32x4_t v_m_7 = vld1q_f32(m + 28); \ + float32x4_t v_q = vld1q_f32(q + 0); \ + float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 0, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 1, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 2, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 3, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 4); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 4, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 5, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 6, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 7, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 8); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 8, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 9, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 10, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 11, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 12); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 12, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 13, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 14, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 15, v_m, v_p, _RES, _PROC) \ + } + +//! Iterative process of computing distance (FP32, M=32, N=32) +#define MATRIX_FP32_ITER_32X32_NEON(m, q, _RES, _PROC) \ + { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + float32x4_t v_m_2 = vld1q_f32(m + 8); \ + float32x4_t v_m_3 = vld1q_f32(m + 12); \ + float32x4_t v_m_4 = vld1q_f32(m + 16); \ + float32x4_t v_m_5 = vld1q_f32(m + 20); \ + float32x4_t v_m_6 = vld1q_f32(m + 24); \ + float32x4_t v_m_7 = vld1q_f32(m + 28); \ + float32x4_t v_q = vld1q_f32(q + 0); \ + float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 0, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 1, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 2, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 3, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 4); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 4, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 5, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 6, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 7, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 8); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 8, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 9, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 10, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 11, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 12); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 12, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 13, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 14, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 15, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 16); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 16, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 17, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 18, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 19, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 20); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 20, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 21, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 22, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 23, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 24); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 24, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 25, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 26, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 27, v_m, v_p, _RES, _PROC) \ + v_q = vld1q_f32(q + 28); \ + v_p = vdupq_laneq_f32(v_q, 0); \ + MATRIX_VAR_PROC(8, 1, 28, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 1); \ + MATRIX_VAR_PROC(8, 1, 29, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 2); \ + MATRIX_VAR_PROC(8, 1, 30, v_m, v_p, _RES, _PROC) \ + v_p = vdupq_laneq_f32(v_q, 3); \ + MATRIX_VAR_PROC(8, 1, 31, v_m, v_p, _RES, _PROC) \ + } diff --git a/src/ailego/math/distance_matrix_inner_product_utility.i b/src/ailego/math/distance_matrix_inner_product_utility.i new file mode 100644 index 0000000..3f28b15 --- /dev/null +++ b/src/ailego/math/distance_matrix_inner_product_utility.i @@ -0,0 +1,208 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__SSE4_1__) +//! Four-bits Convert Table +static const AILEGO_ALIGNED(32) int8_t Int4ConvertTable[32] = { + 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1, + 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1}; + +#define NEGZEROS_FP32_SSE _mm_set1_ps(-0.0f) +#define MASK_INT4_SSE _mm_set1_epi32(0x0f0f0f0f) +#define ONES_INT16_SSE _mm_set1_epi32(0x00010001) +#define INT4_LOOKUP_SSE _mm_load_si128((const __m128i *)Int4ConvertTable) +#endif // __SSE4_1__ + +#if defined(__AVX__) +// #define NEGZEROS_FP32_AVX _mm256_set1_ps(-0.0f) +#define MASK_INT4_AVX _mm256_set1_epi32(0x0f0f0f0f) +#define ONES_INT16_AVX _mm256_set1_epi32(0x00010001) +#define INT4_LOOKUP_AVX _mm256_load_si256((const __m256i *)Int4ConvertTable) +#endif // __AVX__ + +#if defined(__AVX512F__) && !defined(__AVX512DQ__) +#define _mm512_xor_ps(a, b) \ + _mm512_castsi512_ps( \ + _mm512_xor_epi32(_mm512_castps_si512(a), _mm512_castps_si512(b))) +#endif // __AVX512DQ__ + +//! Reverse sign of value (GENERAL) +#define NEGATE_FP32_GENERAL(v) -(v) + +//! Calculate Fused-Multiply-Add (SSE) +#define FMA_FP32_SSE(xmm_m, xmm_q, xmm_sum) \ + xmm_sum = _mm_fmadd_ps(xmm_m, xmm_q, xmm_sum); + +//! Calculate Fused-Multiply-Add (AVX) +#define FMA_FP32_AVX(ymm_m, ymm_q, ymm_sum) \ + ymm_sum = _mm256_fmadd_ps(ymm_m, ymm_q, ymm_sum); + +//! Calculate Fused-Multiply-Add (AVX512) +#define FMA_FP32_AVX512(zmm_m, zmm_q, zmm_sum) \ + zmm_sum = _mm512_fmadd_ps(zmm_m, zmm_q, zmm_sum); + +//! Calculate Fused-Multiply-Add (AVX512FP16) +#define FMA_FP16_AVX512FP16(zmm_m, zmm_q, zmm_sum) \ + zmm_sum = _mm512_fmadd_ph(zmm_m, zmm_q, zmm_sum); + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_FP16_GENERAL(m, q, sum) sum += (m * q); + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_FP32_GENERAL(m, q, sum) sum += (m * q); + +//! Calculate Fused-Multiply-Add (NEON) +#define FMA_FP16_NEON(v_m, v_q, v_sum) v_sum = vfmaq_f16(v_sum, v_m, v_q); + +//! Calculate Fused-Multiply-Add (NEON) +#define FMA_FP32_NEON(v_m, v_q, v_sum) v_sum = vfmaq_f32(v_sum, v_m, v_q); + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_INT4_GENERAL(m, q, sum) \ + sum += Int4MulTable[(((m) << 4) & 0xf0) | (((q) >> 0) & 0xf)] + \ + Int4MulTable[(((m) >> 0) & 0xf0) | (((q) >> 4) & 0xf)]; + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_INT8_GENERAL(m, q, sum) sum += static_cast(m * q); + +//! Calculate Fused-Multiply-Add (SSE) +#define FMA_INT8_SSE(xmm_m, xmm_q, xmm_sum) \ + xmm_sum = _mm_add_epi32( \ + _mm_madd_epi16( \ + _mm_maddubs_epi16(_mm_abs_epi8(xmm_q), _mm_sign_epi8(xmm_m, xmm_q)), \ + ONES_INT16_SSE), \ + xmm_sum); + +//! Calculate Fused-Multiply-Add (AVX) +#define FMA_INT8_AVX(ymm_m, ymm_q, ymm_sum) \ + ymm_sum = _mm256_add_epi32( \ + _mm256_madd_epi16(_mm256_maddubs_epi16(_mm256_abs_epi8(ymm_q), \ + _mm256_sign_epi8(ymm_m, ymm_q)), \ + ONES_INT16_AVX), \ + ymm_sum); + +//! Calculate Fused-Multiply-Add (SSE) +#define FMA_INT4_SSE(xmm_m, xmm_q, xmm_sum) \ + { \ + __m128i xmm_lhs = _mm_shuffle_epi8(INT4_LOOKUP_SSE, \ + _mm_and_si128((xmm_m), MASK_INT4_SSE)); \ + __m128i xmm_rhs = _mm_shuffle_epi8(INT4_LOOKUP_SSE, \ + _mm_and_si128((xmm_q), MASK_INT4_SSE)); \ + xmm_sum = _mm_add_epi32( \ + _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_rhs), \ + _mm_sign_epi8(xmm_lhs, xmm_rhs)), \ + ONES_INT16_SSE), \ + xmm_sum); \ + xmm_lhs = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, \ + _mm_and_si128(_mm_srli_epi32((xmm_m), 4), MASK_INT4_SSE)); \ + xmm_rhs = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, \ + _mm_and_si128(_mm_srli_epi32((xmm_q), 4), MASK_INT4_SSE)); \ + xmm_sum = _mm_add_epi32( \ + _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_rhs), \ + _mm_sign_epi8(xmm_lhs, xmm_rhs)), \ + ONES_INT16_SSE), \ + xmm_sum); \ + } + +//! Calculate Fused-Multiply-Add (AVX) +#define FMA_INT4_AVX(ymm_m, ymm_q, ymm_sum) \ + { \ + __m256i ymm_lhs = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, _mm256_and_si256((ymm_m), MASK_INT4_AVX)); \ + __m256i ymm_rhs = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, _mm256_and_si256((ymm_q), MASK_INT4_AVX)); \ + ymm_sum = _mm256_add_epi32( \ + _mm256_madd_epi16( \ + _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_rhs), \ + _mm256_sign_epi8(ymm_lhs, ymm_rhs)), \ + ONES_INT16_AVX), \ + ymm_sum); \ + ymm_lhs = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, \ + _mm256_and_si256(_mm256_srli_epi32((ymm_m), 4), MASK_INT4_AVX)); \ + ymm_rhs = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, \ + _mm256_and_si256(_mm256_srli_epi32((ymm_q), 4), MASK_INT4_AVX)); \ + ymm_sum = _mm256_add_epi32( \ + _mm256_madd_epi16( \ + _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_rhs), \ + _mm256_sign_epi8(ymm_lhs, ymm_rhs)), \ + ONES_INT16_AVX), \ + ymm_sum); \ + } + +//! Compute the distance between matrix and query +#define FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) \ + { \ + __m128i xmm_lhs_0 = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, _mm_and_si128((xmm_lhs), MASK_INT4_SSE)); \ + __m128i xmm_rhs_0 = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, _mm_and_si128((xmm_rhs), MASK_INT4_SSE)); \ + __m128i xmm_lhs_1 = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, \ + _mm_and_si128(_mm_srli_epi32((xmm_lhs), 4), MASK_INT4_SSE)); \ + __m128i xmm_rhs_1 = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, \ + _mm_and_si128(_mm_srli_epi32((xmm_rhs), 4), MASK_INT4_SSE)); \ + xmm_lhs_0 = _mm_sign_epi8(xmm_lhs_0, xmm_rhs_0); \ + xmm_lhs_1 = _mm_sign_epi8(xmm_lhs_1, xmm_rhs_1); \ + xmm_rhs_0 = _mm_abs_epi8(xmm_rhs_0); \ + xmm_rhs_1 = _mm_abs_epi8(xmm_rhs_1); \ + xmm_lhs_0 = _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_0, xmm_lhs_0), \ + ONES_INT16_SSE); \ + xmm_lhs_1 = _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_1, xmm_lhs_1), \ + ONES_INT16_SSE); \ + xmm_sum = _mm_add_epi32(_mm_add_epi32(xmm_lhs_0, xmm_lhs_1), xmm_sum); \ + } + +//! Compute the distance between matrix and query +#define FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum) \ + { \ + __m256i ymm_lhs_0 = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, _mm256_and_si256((ymm_lhs), MASK_INT4_AVX)); \ + __m256i ymm_rhs_0 = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, _mm256_and_si256((ymm_rhs), MASK_INT4_AVX)); \ + __m256i ymm_lhs_1 = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, \ + _mm256_and_si256(_mm256_srli_epi32((ymm_lhs), 4), MASK_INT4_AVX)); \ + __m256i ymm_rhs_1 = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, \ + _mm256_and_si256(_mm256_srli_epi32((ymm_rhs), 4), MASK_INT4_AVX)); \ + ymm_lhs_0 = _mm256_sign_epi8(ymm_lhs_0, ymm_rhs_0); \ + ymm_lhs_1 = _mm256_sign_epi8(ymm_lhs_1, ymm_rhs_1); \ + ymm_rhs_0 = _mm256_abs_epi8(ymm_rhs_0); \ + ymm_rhs_1 = _mm256_abs_epi8(ymm_rhs_1); \ + ymm_lhs_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_0, ymm_lhs_0), \ + ONES_INT16_AVX); \ + ymm_lhs_1 = _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_1, ymm_lhs_1), \ + ONES_INT16_AVX); \ + ymm_sum = \ + _mm256_add_epi32(_mm256_add_epi32(ymm_lhs_0, ymm_lhs_1), ymm_sum); \ + } + +#define ACCUM_FP16_STEP_GENERAL FMA_FP16_GENERAL +#define ACCUM_FP16_STEP_NEON FMA_FP16_NEON + +#define ACCUM_FP32_STEP_SSE FMA_FP32_SSE +#define ACCUM_FP32_STEP_AVX FMA_FP32_AVX +#define ACCUM_FP32_STEP_AVX512 FMA_FP32_AVX512 +#define ACCUM_FP32_STEP_NEON FMA_FP32_NEON + +#define ACCUM_INT4_STEP_SSE FMA_INT4_SSE +#define ACCUM_INT4_STEP_AVX FMA_INT4_AVX + +#define ACCUM_INT8_STEP_SSE FMA_INT8_SSE +#define ACCUM_INT8_STEP_AVX FMA_INT8_AVX diff --git a/src/ailego/math/distance_matrix_int32.i b/src/ailego/math/distance_matrix_int32.i new file mode 100644 index 0000000..fbbd5a4 --- /dev/null +++ b/src/ailego/math/distance_matrix_int32.i @@ -0,0 +1,859 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "matrix_define.i" + +#if defined(__AVX__) && defined(__GNUC__) +#define _mm256_set_m128i(a, b) \ + _mm256_inserti128_si256(_mm256_castsi128_si256(b), (a), 1) +#endif // __AVX__ + +#if !defined(__AVX__) +#define _mm_broadcast_si32(a) _mm_castps_si128(_mm_load1_ps((const float *)(a))) +#else +#define _mm_broadcast_si32(a) \ + _mm_castps_si128(_mm_broadcast_ss((const float *)(a))) +#define _mm256_broadcast_si32(a) \ + _mm256_castps_si256(_mm256_broadcast_ss((const float *)(a))) +#endif // !__AVX__ + +//! Iterative process of computing distance (INT32, M=2, N=1) +#define MATRIX_INT32_ITER_2X1_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi = _LOAD((const __m128i *)(qi)); \ + __m128i xmm_mi = _LOAD((const __m128i *)(mi)); \ + __m128i xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(1, 1, 0, 0)); \ + _PROC(xmm_mi, xmm_pi, _RES##_0_0) \ + xmm_mi = _LOAD((const __m128i *)(mi + 4)); \ + xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(3, 3, 2, 2)); \ + _PROC(xmm_mi, xmm_pi, _RES##_0_1) \ + } + +//! Iterative process of computing distance (INT32, M=2, N=2) +#define MATRIX_INT32_ITER_2X2_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi = _LOAD((const __m128i *)(qi)); \ + __m128i xmm_mi = _LOAD((const __m128i *)(mi)); \ + __m128i xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \ + _PROC(xmm_mi, xmm_pi, _RES##_0_0) \ + xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \ + _PROC(xmm_mi, xmm_pi, _RES##_0_1) \ + } + +//! Iterative process of computing distance (INT32, M=4, N=1) +#define MATRIX_INT32_ITER_4X1_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + _PROC(xmm_mi_0, xmm_qi, _RES##_0_0) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + _PROC(xmm_mi_1, xmm_qi, _RES##_1_0) \ + } + +//! Iterative process of computing distance (INT32, M=4, N=2) +#define MATRIX_INT32_ITER_4X2_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \ + __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \ + __m128i xmm_mi = _LOAD((const __m128i *)(mi)); \ + MATRIX_VAR_PROC(1, 2, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=4, N=4) +#define MATRIX_INT32_ITER_4X4_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi = _LOAD((const __m128i *)(mi)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + _PROC(xmm_mi, xmm_qi, _RES##_0_0) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + _PROC(xmm_mi, xmm_qi, _RES##_0_1) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + _PROC(xmm_mi, xmm_qi, _RES##_0_2) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + _PROC(xmm_mi, xmm_qi, _RES##_0_3) \ + } + +//! Iterative process of computing distance (INT32, M=8, N=1) +#define MATRIX_INT32_ITER_8X1_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + MATRIX_VAR_PROC(2, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=8, N=2) +#define MATRIX_INT32_ITER_8X2_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \ + __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \ + __m128i xmm_mi = _LOAD((const __m128i *)(mi + 0)); \ + MATRIX_VAR_PROC(1, 2, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_mi = _LOAD((const __m128i *)(mi + 4)); \ + MATRIX_VAR_PROC(1, 2, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=8, N=4) +#define MATRIX_INT32_ITER_8X4_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(2, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(2, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(2, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=8, N=8) +#define MATRIX_INT32_ITER_8X8_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(2, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(2, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(2, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(2, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(2, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(2, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(2, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=1) +#define MATRIX_INT32_ITER_16X1_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=2) +#define MATRIX_INT32_ITER_16X2_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=4) +#define MATRIX_INT32_ITER_16X4_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=8) +#define MATRIX_INT32_ITER_16X8_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(4, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(4, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(4, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(4, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=16) +#define MATRIX_INT32_ITER_16X16_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(4, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(4, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(4, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(4, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 8); \ + MATRIX_VAR_PROC(4, 1, 8, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 9); \ + MATRIX_VAR_PROC(4, 1, 9, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 10); \ + MATRIX_VAR_PROC(4, 1, 10, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 11); \ + MATRIX_VAR_PROC(4, 1, 11, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 12); \ + MATRIX_VAR_PROC(4, 1, 12, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 13); \ + MATRIX_VAR_PROC(4, 1, 13, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 14); \ + MATRIX_VAR_PROC(4, 1, 14, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 15); \ + MATRIX_VAR_PROC(4, 1, 15, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=1) +#define MATRIX_INT32_ITER_32X1_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi = _mm_broadcast_si32(qi); \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + _PROC(xmm_mi_0, xmm_qi, _RES##_0_0) \ + _PROC(xmm_mi_1, xmm_qi, _RES##_1_0) \ + _PROC(xmm_mi_2, xmm_qi, _RES##_2_0) \ + _PROC(xmm_mi_3, xmm_qi, _RES##_3_0) \ + xmm_mi_0 = _LOAD((const __m128i *)(mi + 16)); \ + xmm_mi_1 = _LOAD((const __m128i *)(mi + 20)); \ + xmm_mi_2 = _LOAD((const __m128i *)(mi + 24)); \ + xmm_mi_3 = _LOAD((const __m128i *)(mi + 28)); \ + _PROC(xmm_mi_0, xmm_qi, _RES##_4_0) \ + _PROC(xmm_mi_1, xmm_qi, _RES##_5_0) \ + _PROC(xmm_mi_2, xmm_qi, _RES##_6_0) \ + _PROC(xmm_mi_3, xmm_qi, _RES##_7_0) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=2) +#define MATRIX_INT32_ITER_32X2_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \ + __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + MATRIX_VAR_PROC(1, 2, 0, xmm_mi_0, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 1, xmm_mi_1, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 2, xmm_mi_2, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 3, xmm_mi_3, xmm_qi, _RES, _PROC) \ + xmm_mi_0 = _LOAD((const __m128i *)(mi + 16)); \ + xmm_mi_1 = _LOAD((const __m128i *)(mi + 20)); \ + xmm_mi_2 = _LOAD((const __m128i *)(mi + 24)); \ + xmm_mi_3 = _LOAD((const __m128i *)(mi + 28)); \ + MATRIX_VAR_PROC(1, 2, 4, xmm_mi_0, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 5, xmm_mi_1, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 6, xmm_mi_2, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 2, 7, xmm_mi_3, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=4) +#define MATRIX_INT32_ITER_32X4_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \ + __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \ + __m128i xmm_qi_2 = _mm_broadcast_si32(qi + 2); \ + __m128i xmm_qi_3 = _mm_broadcast_si32(qi + 3); \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + MATRIX_VAR_PROC(1, 4, 0, xmm_mi_0, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 1, xmm_mi_1, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 2, xmm_mi_2, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 3, xmm_mi_3, xmm_qi, _RES, _PROC) \ + xmm_mi_0 = _LOAD((const __m128i *)(mi + 16)); \ + xmm_mi_1 = _LOAD((const __m128i *)(mi + 20)); \ + xmm_mi_2 = _LOAD((const __m128i *)(mi + 24)); \ + xmm_mi_3 = _LOAD((const __m128i *)(mi + 28)); \ + MATRIX_VAR_PROC(1, 4, 4, xmm_mi_0, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 5, xmm_mi_1, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 6, xmm_mi_2, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 4, 7, xmm_mi_3, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=8) +#define MATRIX_INT32_ITER_32X8_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \ + __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \ + __m128i xmm_qi_2 = _mm_broadcast_si32(qi + 2); \ + __m128i xmm_qi_3 = _mm_broadcast_si32(qi + 3); \ + __m128i xmm_qi_4 = _mm_broadcast_si32(qi + 4); \ + __m128i xmm_qi_5 = _mm_broadcast_si32(qi + 5); \ + __m128i xmm_qi_6 = _mm_broadcast_si32(qi + 6); \ + __m128i xmm_qi_7 = _mm_broadcast_si32(qi + 7); \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + MATRIX_VAR_PROC(1, 8, 0, xmm_mi_0, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 1, xmm_mi_1, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 2, xmm_mi_2, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 3, xmm_mi_3, xmm_qi, _RES, _PROC) \ + xmm_mi_0 = _LOAD((const __m128i *)(mi + 16)); \ + xmm_mi_1 = _LOAD((const __m128i *)(mi + 20)); \ + xmm_mi_2 = _LOAD((const __m128i *)(mi + 24)); \ + xmm_mi_3 = _LOAD((const __m128i *)(mi + 28)); \ + MATRIX_VAR_PROC(1, 8, 4, xmm_mi_0, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 5, xmm_mi_1, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 6, xmm_mi_2, xmm_qi, _RES, _PROC) \ + MATRIX_VAR_PROC(1, 8, 7, xmm_mi_3, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=16) +#define MATRIX_INT32_ITER_32X16_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + __m128i xmm_mi_4 = _LOAD((const __m128i *)(mi + 16)); \ + __m128i xmm_mi_5 = _LOAD((const __m128i *)(mi + 20)); \ + __m128i xmm_mi_6 = _LOAD((const __m128i *)(mi + 24)); \ + __m128i xmm_mi_7 = _LOAD((const __m128i *)(mi + 28)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(8, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(8, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(8, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(8, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(8, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(8, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(8, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(8, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 8); \ + MATRIX_VAR_PROC(8, 1, 8, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 9); \ + MATRIX_VAR_PROC(8, 1, 9, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 10); \ + MATRIX_VAR_PROC(8, 1, 10, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 11); \ + MATRIX_VAR_PROC(8, 1, 11, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 12); \ + MATRIX_VAR_PROC(8, 1, 12, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 13); \ + MATRIX_VAR_PROC(8, 1, 13, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 14); \ + MATRIX_VAR_PROC(8, 1, 14, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 15); \ + MATRIX_VAR_PROC(8, 1, 15, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=32) +#define MATRIX_INT32_ITER_32X32_SSE(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \ + __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \ + __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \ + __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \ + __m128i xmm_mi_4 = _LOAD((const __m128i *)(mi + 16)); \ + __m128i xmm_mi_5 = _LOAD((const __m128i *)(mi + 20)); \ + __m128i xmm_mi_6 = _LOAD((const __m128i *)(mi + 24)); \ + __m128i xmm_mi_7 = _LOAD((const __m128i *)(mi + 28)); \ + __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(8, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(8, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(8, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(8, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(8, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(8, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(8, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(8, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 8); \ + MATRIX_VAR_PROC(8, 1, 8, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 9); \ + MATRIX_VAR_PROC(8, 1, 9, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 10); \ + MATRIX_VAR_PROC(8, 1, 10, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 11); \ + MATRIX_VAR_PROC(8, 1, 11, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 12); \ + MATRIX_VAR_PROC(8, 1, 12, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 13); \ + MATRIX_VAR_PROC(8, 1, 13, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 14); \ + MATRIX_VAR_PROC(8, 1, 14, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 15); \ + MATRIX_VAR_PROC(8, 1, 15, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 16); \ + MATRIX_VAR_PROC(8, 1, 16, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 17); \ + MATRIX_VAR_PROC(8, 1, 17, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 18); \ + MATRIX_VAR_PROC(8, 1, 18, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 19); \ + MATRIX_VAR_PROC(8, 1, 19, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 20); \ + MATRIX_VAR_PROC(8, 1, 20, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 21); \ + MATRIX_VAR_PROC(8, 1, 21, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 22); \ + MATRIX_VAR_PROC(8, 1, 22, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 23); \ + MATRIX_VAR_PROC(8, 1, 23, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 24); \ + MATRIX_VAR_PROC(8, 1, 24, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 25); \ + MATRIX_VAR_PROC(8, 1, 25, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 26); \ + MATRIX_VAR_PROC(8, 1, 26, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 27); \ + MATRIX_VAR_PROC(8, 1, 27, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 28); \ + MATRIX_VAR_PROC(8, 1, 28, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 29); \ + MATRIX_VAR_PROC(8, 1, 29, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 30); \ + MATRIX_VAR_PROC(8, 1, 30, xmm_mi, xmm_qi, _RES, _PROC) \ + xmm_qi = _mm_broadcast_si32(qi + 31); \ + MATRIX_VAR_PROC(8, 1, 31, xmm_mi, xmm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=2, N=1) +#define MATRIX_INT32_ITER_2X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_qi = _mm256_set_epi32(qi[3], qi[3], qi[2], qi[2], qi[1], \ + qi[1], qi[0], qi[0]); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + } + +//! Iterative process of computing distance (INT32, M=2, N=2) +#define MATRIX_INT32_ITER_2X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi = _LOAD((const __m256i *)(qi)); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_0) \ + ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_1) \ + } + +//! Iterative process of computing distance (INT32, M=4, N=1) +#define MATRIX_INT32_ITER_4X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_qi = \ + _mm256_set_m128i(_mm_broadcast_si32(qi + 1), _mm_broadcast_si32(qi)); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + } + +//! Iterative process of computing distance (INT32, M=4, N=2) +#define MATRIX_INT32_ITER_4X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_qi = _mm256_set_m128i(_mm_broadcast_si32(qi + 2), \ + _mm_broadcast_si32(qi + 0)); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + ymm_qi = _mm256_set_m128i(_mm_broadcast_si32(qi + 3), \ + _mm_broadcast_si32(qi + 1)); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_1) \ + } + +//! Iterative process of computing distance (INT32, M=4, N=4) +#define MATRIX_INT32_ITER_4X4_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi = _LOAD((const __m256i *)(qi)); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(0, 0, 0, 0)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_0) \ + ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(1, 1, 1, 1)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_1) \ + ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(2, 2, 2, 2)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_2) \ + ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(3, 3, 3, 3)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_3) \ + } + +//! Iterative process of computing distance (INT32, M=8, N=1) +#define MATRIX_INT32_ITER_8X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi = _mm256_broadcast_si32(qi); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + } + +//! Iterative process of computing distance (INT32, M=8, N=2) +#define MATRIX_INT32_ITER_8X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi_0 = _mm256_broadcast_si32(qi + 0); \ + __m256i ymm_qi_1 = _mm256_broadcast_si32(qi + 1); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=8, N=4) +#define MATRIX_INT32_ITER_8X4_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_1) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_2) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_3) \ + } + +//! Iterative process of computing distance (INT32, M=8, N=8) +#define MATRIX_INT32_ITER_8X8_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_1) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_2) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_3) \ + ymm_qi = _mm256_broadcast_si32(qi + 4); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_4) \ + ymm_qi = _mm256_broadcast_si32(qi + 5); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_5) \ + ymm_qi = _mm256_broadcast_si32(qi + 6); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_6) \ + ymm_qi = _mm256_broadcast_si32(qi + 7); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_7) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=1) +#define MATRIX_INT32_ITER_16X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=2) +#define MATRIX_INT32_ITER_16X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=4) +#define MATRIX_INT32_ITER_16X4_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=8) +#define MATRIX_INT32_ITER_16X8_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(2, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(2, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(2, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(2, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=16, N=16) +#define MATRIX_INT32_ITER_16X16_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(2, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(2, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(2, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(2, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 8); \ + MATRIX_VAR_PROC(2, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 9); \ + MATRIX_VAR_PROC(2, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 10); \ + MATRIX_VAR_PROC(2, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 11); \ + MATRIX_VAR_PROC(2, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 12); \ + MATRIX_VAR_PROC(2, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 13); \ + MATRIX_VAR_PROC(2, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 14); \ + MATRIX_VAR_PROC(2, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 15); \ + MATRIX_VAR_PROC(2, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=1) +#define MATRIX_INT32_ITER_32X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=2) +#define MATRIX_INT32_ITER_32X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=4) +#define MATRIX_INT32_ITER_32X4_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=8) +#define MATRIX_INT32_ITER_32X8_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=16) +#define MATRIX_INT32_ITER_32X16_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 8); \ + MATRIX_VAR_PROC(4, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 9); \ + MATRIX_VAR_PROC(4, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 10); \ + MATRIX_VAR_PROC(4, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 11); \ + MATRIX_VAR_PROC(4, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 12); \ + MATRIX_VAR_PROC(4, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 13); \ + MATRIX_VAR_PROC(4, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 14); \ + MATRIX_VAR_PROC(4, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 15); \ + MATRIX_VAR_PROC(4, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT32, M=32, N=32) +#define MATRIX_INT32_ITER_32X32_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \ + __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 4); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 5); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 6); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 7); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 8); \ + MATRIX_VAR_PROC(4, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 9); \ + MATRIX_VAR_PROC(4, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 10); \ + MATRIX_VAR_PROC(4, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 11); \ + MATRIX_VAR_PROC(4, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 12); \ + MATRIX_VAR_PROC(4, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 13); \ + MATRIX_VAR_PROC(4, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 14); \ + MATRIX_VAR_PROC(4, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 15); \ + MATRIX_VAR_PROC(4, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 16); \ + MATRIX_VAR_PROC(4, 1, 16, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 17); \ + MATRIX_VAR_PROC(4, 1, 17, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 18); \ + MATRIX_VAR_PROC(4, 1, 18, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 19); \ + MATRIX_VAR_PROC(4, 1, 19, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 20); \ + MATRIX_VAR_PROC(4, 1, 20, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 21); \ + MATRIX_VAR_PROC(4, 1, 21, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 22); \ + MATRIX_VAR_PROC(4, 1, 22, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 23); \ + MATRIX_VAR_PROC(4, 1, 23, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 24); \ + MATRIX_VAR_PROC(4, 1, 24, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 25); \ + MATRIX_VAR_PROC(4, 1, 25, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 26); \ + MATRIX_VAR_PROC(4, 1, 26, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 27); \ + MATRIX_VAR_PROC(4, 1, 27, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 28); \ + MATRIX_VAR_PROC(4, 1, 28, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 29); \ + MATRIX_VAR_PROC(4, 1, 29, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 30); \ + MATRIX_VAR_PROC(4, 1, 30, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si32(qi + 31); \ + MATRIX_VAR_PROC(4, 1, 31, ymm_mi, ymm_qi, _RES, _PROC) \ + } diff --git a/src/ailego/math/distance_matrix_int64.i b/src/ailego/math/distance_matrix_int64.i new file mode 100644 index 0000000..0364d1e --- /dev/null +++ b/src/ailego/math/distance_matrix_int64.i @@ -0,0 +1,472 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "matrix_define.i" + +#if defined(__AVX__) +#define _mm256_broadcast_si64(a) \ + _mm256_castpd_si256(_mm256_broadcast_sd((const double *)(a))) +#endif // __AVX__ + +//! Iterative process of computing distance (INT64, M=2, N=1) +#define MATRIX_INT64_ITER_2X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi = _LOAD((const __m256i *)(qi)); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_pi = \ + _mm256_permute4x64_epi64(ymm_qi, _MM_SHUFFLE(1, 1, 0, 0)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + ymm_pi = _mm256_permute4x64_epi64(ymm_qi, _MM_SHUFFLE(3, 3, 2, 2)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_1) \ + } + +//! Iterative process of computing distance (INT64, M=2, N=2) +#define MATRIX_INT64_ITER_2X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi = _LOAD((const __m256i *)(qi)); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_pi = \ + _mm256_permute4x64_epi64(ymm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_0) \ + ymm_pi = _mm256_permute4x64_epi64(ymm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \ + _PROC(ymm_mi, ymm_pi, _RES##_0_1) \ + } + +//! Iterative process of computing distance (INT64, M=4, N=1) +#define MATRIX_INT64_ITER_4X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + _PROC(ymm_mi, ymm_qi, _RES##_1_0) \ + } + +//! Iterative process of computing distance (INT64, M=4, N=2) +#define MATRIX_INT64_ITER_4X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \ + __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=4, N=4) +#define MATRIX_INT64_ITER_4X4_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_1) \ + ymm_qi = _mm256_broadcast_si64(qi + 2); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_2) \ + ymm_qi = _mm256_broadcast_si64(qi + 3); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_3) \ + } + +//! Iterative process of computing distance (INT64, M=8, N=1) +#define MATRIX_INT64_ITER_8X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi = _mm256_broadcast_si64(qi); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + _PROC(ymm_mi, ymm_qi, _RES##_1_0) \ + } + +//! Iterative process of computing distance (INT64, M=8, N=2) +#define MATRIX_INT64_ITER_8X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \ + __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + MATRIX_VAR_PROC(1, 2, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=8, N=4) +#define MATRIX_INT64_ITER_8X4_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 2); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 3); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=8, N=8) +#define MATRIX_INT64_ITER_8X8_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 2); \ + MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 3); \ + MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 4); \ + MATRIX_VAR_PROC(2, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 5); \ + MATRIX_VAR_PROC(2, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 6); \ + MATRIX_VAR_PROC(2, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 7); \ + MATRIX_VAR_PROC(2, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=16, N=1) +#define MATRIX_INT64_ITER_16X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi = _mm256_broadcast_si64(qi); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + _PROC(ymm_mi, ymm_qi, _RES##_1_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 8)); \ + _PROC(ymm_mi, ymm_qi, _RES##_2_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 12)); \ + _PROC(ymm_mi, ymm_qi, _RES##_3_0) \ + } + +//! Iterative process of computing distance (INT64, M=16, N=2) +#define MATRIX_INT64_ITER_16X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \ + __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + MATRIX_VAR_PROC(1, 2, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 8)); \ + MATRIX_VAR_PROC(1, 2, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 12)); \ + MATRIX_VAR_PROC(1, 2, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=16, N=4) +#define MATRIX_INT64_ITER_16X4_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=16, N=8) +#define MATRIX_INT64_ITER_16X8_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 4); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 5); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 6); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 7); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=16, N=16) +#define MATRIX_INT64_ITER_16X16_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 2); \ + MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 3); \ + MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 4); \ + MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 5); \ + MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 6); \ + MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 7); \ + MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 8); \ + MATRIX_VAR_PROC(4, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 9); \ + MATRIX_VAR_PROC(4, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 10); \ + MATRIX_VAR_PROC(4, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 11); \ + MATRIX_VAR_PROC(4, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 12); \ + MATRIX_VAR_PROC(4, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 13); \ + MATRIX_VAR_PROC(4, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 14); \ + MATRIX_VAR_PROC(4, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 15); \ + MATRIX_VAR_PROC(4, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=32, N=1) +#define MATRIX_INT64_ITER_32X1_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi = _mm256_broadcast_si64(qi); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + _PROC(ymm_mi, ymm_qi, _RES##_0_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + _PROC(ymm_mi, ymm_qi, _RES##_1_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 8)); \ + _PROC(ymm_mi, ymm_qi, _RES##_2_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 12)); \ + _PROC(ymm_mi, ymm_qi, _RES##_3_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 16)); \ + _PROC(ymm_mi, ymm_qi, _RES##_4_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 20)); \ + _PROC(ymm_mi, ymm_qi, _RES##_5_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 24)); \ + _PROC(ymm_mi, ymm_qi, _RES##_6_0) \ + ymm_mi = _LOAD((const __m256i *)(mi + 28)); \ + _PROC(ymm_mi, ymm_qi, _RES##_7_0) \ + } + +//! Iterative process of computing distance (INT64, M=32, N=2) +#define MATRIX_INT64_ITER_32X2_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \ + __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + MATRIX_VAR_PROC(1, 2, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 8)); \ + MATRIX_VAR_PROC(1, 2, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 12)); \ + MATRIX_VAR_PROC(1, 2, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 16)); \ + MATRIX_VAR_PROC(1, 2, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 20)); \ + MATRIX_VAR_PROC(1, 2, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 24)); \ + MATRIX_VAR_PROC(1, 2, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 28)); \ + MATRIX_VAR_PROC(1, 2, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=32, N=4) +#define MATRIX_INT64_ITER_32X4_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \ + __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \ + __m256i ymm_qi_2 = _mm256_broadcast_si64(qi + 2); \ + __m256i ymm_qi_3 = _mm256_broadcast_si64(qi + 3); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + MATRIX_VAR_PROC(1, 4, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + MATRIX_VAR_PROC(1, 4, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 8)); \ + MATRIX_VAR_PROC(1, 4, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 12)); \ + MATRIX_VAR_PROC(1, 4, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 16)); \ + MATRIX_VAR_PROC(1, 4, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 20)); \ + MATRIX_VAR_PROC(1, 4, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 24)); \ + MATRIX_VAR_PROC(1, 4, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 28)); \ + MATRIX_VAR_PROC(1, 4, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=32, N=8) +#define MATRIX_INT64_ITER_32X8_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \ + __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \ + __m256i ymm_qi_2 = _mm256_broadcast_si64(qi + 2); \ + __m256i ymm_qi_3 = _mm256_broadcast_si64(qi + 3); \ + __m256i ymm_qi_4 = _mm256_broadcast_si64(qi + 4); \ + __m256i ymm_qi_5 = _mm256_broadcast_si64(qi + 5); \ + __m256i ymm_qi_6 = _mm256_broadcast_si64(qi + 6); \ + __m256i ymm_qi_7 = _mm256_broadcast_si64(qi + 7); \ + __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \ + MATRIX_VAR_PROC(1, 8, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 4)); \ + MATRIX_VAR_PROC(1, 8, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 8)); \ + MATRIX_VAR_PROC(1, 8, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 12)); \ + MATRIX_VAR_PROC(1, 8, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 16)); \ + MATRIX_VAR_PROC(1, 8, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 20)); \ + MATRIX_VAR_PROC(1, 8, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 24)); \ + MATRIX_VAR_PROC(1, 8, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_mi = _LOAD((const __m256i *)(mi + 28)); \ + MATRIX_VAR_PROC(1, 8, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=32, N=16) +#define MATRIX_INT64_ITER_32X16_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \ + __m256i ymm_mi_4 = _LOAD((const __m256i *)(mi + 16)); \ + __m256i ymm_mi_5 = _LOAD((const __m256i *)(mi + 20)); \ + __m256i ymm_mi_6 = _LOAD((const __m256i *)(mi + 24)); \ + __m256i ymm_mi_7 = _LOAD((const __m256i *)(mi + 28)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + MATRIX_VAR_PROC(8, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + MATRIX_VAR_PROC(8, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 2); \ + MATRIX_VAR_PROC(8, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 3); \ + MATRIX_VAR_PROC(8, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 4); \ + MATRIX_VAR_PROC(8, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 5); \ + MATRIX_VAR_PROC(8, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 6); \ + MATRIX_VAR_PROC(8, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 7); \ + MATRIX_VAR_PROC(8, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 8); \ + MATRIX_VAR_PROC(8, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 9); \ + MATRIX_VAR_PROC(8, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 10); \ + MATRIX_VAR_PROC(8, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 11); \ + MATRIX_VAR_PROC(8, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 12); \ + MATRIX_VAR_PROC(8, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 13); \ + MATRIX_VAR_PROC(8, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 14); \ + MATRIX_VAR_PROC(8, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 15); \ + MATRIX_VAR_PROC(8, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \ + } + +//! Iterative process of computing distance (INT64, M=32, N=32) +#define MATRIX_INT64_ITER_32X32_AVX(mi, qi, _RES, _LOAD, _PROC) \ + { \ + __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \ + __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \ + __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \ + __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \ + __m256i ymm_mi_4 = _LOAD((const __m256i *)(mi + 16)); \ + __m256i ymm_mi_5 = _LOAD((const __m256i *)(mi + 20)); \ + __m256i ymm_mi_6 = _LOAD((const __m256i *)(mi + 24)); \ + __m256i ymm_mi_7 = _LOAD((const __m256i *)(mi + 28)); \ + __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \ + MATRIX_VAR_PROC(8, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 1); \ + MATRIX_VAR_PROC(8, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 2); \ + MATRIX_VAR_PROC(8, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 3); \ + MATRIX_VAR_PROC(8, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 4); \ + MATRIX_VAR_PROC(8, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 5); \ + MATRIX_VAR_PROC(8, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 6); \ + MATRIX_VAR_PROC(8, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 7); \ + MATRIX_VAR_PROC(8, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 8); \ + MATRIX_VAR_PROC(8, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 9); \ + MATRIX_VAR_PROC(8, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 10); \ + MATRIX_VAR_PROC(8, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 11); \ + MATRIX_VAR_PROC(8, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 12); \ + MATRIX_VAR_PROC(8, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 13); \ + MATRIX_VAR_PROC(8, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 14); \ + MATRIX_VAR_PROC(8, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 15); \ + MATRIX_VAR_PROC(8, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 16); \ + MATRIX_VAR_PROC(8, 1, 16, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 17); \ + MATRIX_VAR_PROC(8, 1, 17, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 18); \ + MATRIX_VAR_PROC(8, 1, 18, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 19); \ + MATRIX_VAR_PROC(8, 1, 19, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 20); \ + MATRIX_VAR_PROC(8, 1, 20, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 21); \ + MATRIX_VAR_PROC(8, 1, 21, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 22); \ + MATRIX_VAR_PROC(8, 1, 22, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 23); \ + MATRIX_VAR_PROC(8, 1, 23, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 24); \ + MATRIX_VAR_PROC(8, 1, 24, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 25); \ + MATRIX_VAR_PROC(8, 1, 25, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 26); \ + MATRIX_VAR_PROC(8, 1, 26, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 27); \ + MATRIX_VAR_PROC(8, 1, 27, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 28); \ + MATRIX_VAR_PROC(8, 1, 28, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 29); \ + MATRIX_VAR_PROC(8, 1, 29, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 30); \ + MATRIX_VAR_PROC(8, 1, 30, ymm_mi, ymm_qi, _RES, _PROC) \ + ymm_qi = _mm256_broadcast_si64(qi + 31); \ + MATRIX_VAR_PROC(8, 1, 31, ymm_mi, ymm_qi, _RES, _PROC) \ + } diff --git a/src/ailego/math/distance_matrix_mips_utility.i b/src/ailego/math/distance_matrix_mips_utility.i new file mode 100644 index 0000000..871fdaa --- /dev/null +++ b/src/ailego/math/distance_matrix_mips_utility.i @@ -0,0 +1,160 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Calculate Fused-Multiply-Add (AVX512) +#define FMA_FP32_AVX512(zmm_m, zmm_q, zmm_sum) \ + zmm_sum = _mm512_fmadd_ps(zmm_m, zmm_q, zmm_sum); + +#define FMA_MASK_FP32_AVX512(zmm_m, zmm_q, zmm_sum, mask) \ + zmm_sum = _mm512_mask3_fmadd_ps(zmm_m, zmm_q, zmm_sum, mask); + +#define HorizontalAdd_FP16_NEON(v) \ + vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(v)), vcvt_high_f32_f16(v))) + +#define HorizontalAdd_FP32_V512_TO_V256(zmm) \ + _mm256_add_ps( \ + _mm512_castps512_ps256(zmm), \ + _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(zmm), 1))) + +//! Calculate Fused-Multiply-Add (AVX, FP16) +#define FMA_FP16_GENERAL(lhs, rhs, sum, norm1, norm2) \ + { \ + float v1 = lhs; \ + float v2 = rhs; \ + sum += v1 * v2; \ + norm1 += v1 * v1; \ + norm2 += v2 * v2; \ + } + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_FP32_GENERAL(lhs, rhs, sum, norm1, norm2) \ + { \ + sum += (lhs) * (rhs); \ + norm1 += (lhs) * (lhs); \ + norm2 += (rhs) * (rhs); \ + } + +#if defined(__SSE4_1__) +//! Four-bits Convert Table +static const AILEGO_ALIGNED(32) int8_t Int4ConvertTable[32] = { + 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1, + 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1}; +#endif // __SSE4_1__ + +#if defined(__SSE4_1__) +static const __m128i MASK_INT4_SSE = _mm_set1_epi32(0x0f0f0f0f); +static const __m128i ONES_INT16_SSE = _mm_set1_epi32(0x00010001); +static const __m128i INT4_LOOKUP_SSE = + _mm_load_si128((const __m128i *)Int4ConvertTable); +#endif // __SSE4_1__ + +#if defined(__AVX2__) +static const __m256i MASK_INT4_AVX = _mm256_set1_epi32(0x0f0f0f0f); +static const __m256i ONES_INT16_AVX = _mm256_set1_epi32(0x00010001); +static const __m256i INT4_LOOKUP_AVX = + _mm256_load_si256((const __m256i *)Int4ConvertTable); +#endif // __AVX2__ + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_INT4_GENERAL(lhs, rhs, sum, norm1, norm2) \ + { \ + sum += Int4MulTable[(((lhs) << 4) & 0xf0) | (((rhs) >> 0) & 0xf)] + \ + Int4MulTable[(((lhs) >> 0) & 0xf0) | (((rhs) >> 4) & 0xf)]; \ + norm1 += static_cast( \ + ((int8_t)((lhs) << 4) >> 4) * ((int8_t)((lhs) << 4) >> 4) + \ + ((int8_t)((lhs) & 0xf0) >> 4) * ((int8_t)((lhs) & 0xf0) >> 4)); \ + norm2 += static_cast( \ + ((int8_t)((rhs) << 4) >> 4) * ((int8_t)((rhs) << 4) >> 4) + \ + ((int8_t)((rhs) & 0xf0) >> 4) * ((int8_t)((rhs) & 0xf0) >> 4)); \ + } + + +//! Compute the distance between matrix and query (SSE) +#define FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum_0, xmm_sum_norm1, \ + xmm_sum_norm2) \ + { \ + __m128i xmm_lhs_0 = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, _mm_and_si128((xmm_lhs), MASK_INT4_SSE)); \ + __m128i xmm_rhs_0 = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, _mm_and_si128((xmm_rhs), MASK_INT4_SSE)); \ + __m128i xmm_lhs_1 = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, \ + _mm_and_si128(_mm_srli_epi32((xmm_lhs), 4), MASK_INT4_SSE)); \ + __m128i xmm_rhs_1 = _mm_shuffle_epi8( \ + INT4_LOOKUP_SSE, \ + _mm_and_si128(_mm_srli_epi32((xmm_rhs), 4), MASK_INT4_SSE)); \ + FMA_INT8_SSE(xmm_lhs_0, xmm_rhs_0, xmm_sum_0); \ + FMA_INT8_SSE(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1); \ + FMA_INT8_SSE(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2); \ + FMA_INT8_SSE(xmm_lhs_1, xmm_rhs_1, xmm_sum_0); \ + FMA_INT8_SSE(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1); \ + FMA_INT8_SSE(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2); \ + } + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_INT8_GENERAL(lhs, rhs, sum, norm1, norm2) \ + { \ + sum += static_cast(lhs * rhs); \ + norm1 += static_cast(lhs * lhs); \ + norm2 += static_cast(rhs * rhs); \ + } + +//! Calculate Fused-Multiply-Add (SSE) +#define FMA_INT8_SSE(xmm_lhs, xmm_rhs, xmm_sum) \ + xmm_sum = _mm_add_epi32( \ + _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_rhs), \ + _mm_sign_epi8(xmm_lhs, xmm_rhs)), \ + ONES_INT16_SSE), \ + xmm_sum) + +//! Calculate Fused-Multiply-Add (AVX) +#define FMA_INT8_AVX(ymm_lhs, ymm_rhs, ymm_sum) \ + ymm_sum = _mm256_add_epi32( \ + _mm256_madd_epi16( \ + _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_rhs), \ + _mm256_sign_epi8(ymm_lhs, ymm_rhs)), \ + ONES_INT16_AVX), \ + ymm_sum) + +#define FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_rhs, ymm_sum) \ + ymm_sum = _mm256_add_epi32( \ + _mm256_set_m128i( \ + _mm_setzero_si128(), \ + _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_rhs), \ + _mm_sign_epi8(xmm_lhs, xmm_rhs)), \ + ONES_INT16_SSE)), \ + ymm_sum) + +//! Compute the distance between matrix and query (AVX) +#define FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum_0, ymm_sum1, \ + ymm_sum_norm1, ymm_sum_norm2) \ + { \ + __m256i ymm_lhs_0 = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, _mm256_and_si256((ymm_lhs), MASK_INT4_AVX)); \ + __m256i ymm_rhs_0 = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, _mm256_and_si256((ymm_rhs), MASK_INT4_AVX)); \ + __m256i ymm_lhs_1 = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, \ + _mm256_and_si256(_mm256_srli_epi32((ymm_lhs), 4), MASK_INT4_AVX)); \ + __m256i ymm_rhs_1 = _mm256_shuffle_epi8( \ + INT4_LOOKUP_AVX, \ + _mm256_and_si256(_mm256_srli_epi32((ymm_rhs), 4), MASK_INT4_AVX)); \ + FMA_INT8_AVX(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); \ + FMA_INT8_AVX(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); \ + FMA_INT8_AVX(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); \ + FMA_INT8_AVX(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1); \ + FMA_INT8_AVX(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); \ + FMA_INT8_AVX(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2); \ + } + diff --git a/src/ailego/math/distance_matrix_popcnt.i b/src/ailego/math/distance_matrix_popcnt.i new file mode 100644 index 0000000..bb6b4d9 --- /dev/null +++ b/src/ailego/math/distance_matrix_popcnt.i @@ -0,0 +1,2565 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_int32.i" +#include "distance_matrix_int64.i" +#include "matrix_utility.i" + +//! Calculate population count (UINT32 Permute 1 SSE) +#define POPCNT_UINT32_PERMUTE1_SSE(v, ...) \ + _mm_add_epi16(_mm_srli_epi16(v, 8), _mm_and_si128(v, _mm_set1_epi16(0xff))) + +//! Calculate population count (UINT32 Permute 2 SSE) +#define POPCNT_UINT32_PERMUTE2_SSE(v, ...) \ + _mm_add_epi32(_mm_srli_epi32(v, 16), _mm_and_si128(v, _mm_set1_epi32(0xffff))) + +//! Calculate population count (UINT32 Permute 1 AVX) +#define POPCNT_UINT32_PERMUTE1_AVX(v, ...) \ + _mm256_add_epi16(_mm256_srli_epi16(v, 8), \ + _mm256_and_si256(v, _mm256_set1_epi16(0xff))) + +//! Calculate population count (UINT32 Permute 2 AVX) +#define POPCNT_UINT32_PERMUTE2_AVX(v, ...) \ + _mm256_add_epi32(_mm256_srli_epi32(v, 16), \ + _mm256_and_si256(v, _mm256_set1_epi32(0xffff))) + +//! Calculate population count (UINT64 Permute AVX) +#define POPCNT_UINT64_PERMUTE_AVX(v, ...) \ + _mm256_sad_epu8(v, _mm256_setzero_si256()) + +//! Compute the distance between matrix and query (UINT32, M=2, N=1) +#define POPCNT_UINT32_2X1_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + ((cnt >> 2) << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 2) << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 2) << 2) : qe_0); \ + const uint32_t *qe_3 = q + cnt; \ + if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + if (qe_3 >= qe_0 + 2) { \ + __m128i xmm_m = _mm_load_si128((const __m128i *)(m)); \ + __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + m += 4; \ + q += 2; \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + if (qe_3 >= qe_0 + 2) { \ + __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \ + __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + m += 4; \ + q += 2; \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_0_1); \ + xmm_sum_0_0 = _mm_add_epi32( \ + xmm_sum_0_0, _mm_shuffle_epi32(xmm_sum_0_0, _MM_SHUFFLE(0, 0, 3, 2))); \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_set_epi32(0, 0, m[1], m[0]); \ + __m128i xmm_q = _mm_broadcast_si32(q); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0)); + +//! Compute the distance between matrix and query (UINT32, M=2, N=2) +#define POPCNT_UINT32_2X2_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + ((cnt >> 1) << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 2) : qe_0); \ + const uint32_t *qe_3 = q + (cnt << 1); \ + if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \ + for (; q != qe_1; m += 4, q += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 4, q += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 4, q += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 4, q += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 4, q += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 4, q += 4) { \ + MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_set_epi32(m[1], m[0], m[1], m[0]); \ + __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=4, N=1) +#define POPCNT_UINT32_4X1_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + ((cnt >> 1) << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 1) : qe_0); \ + const uint32_t *qe_3 = q + cnt; \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_load_si128((const __m128i *)(m)); \ + __m128i xmm_q = _mm_broadcast_si32(q); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \ + __m128i xmm_q = _mm_broadcast_si32(q); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + } \ + xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_1_0); \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=4, N=2) +#define POPCNT_UINT32_4X2_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 4, q += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 4, q += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 4, q += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 4, q += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 4, q += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 4, q += 2) { \ + MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=4, N=4) +#define POPCNT_UINT32_4X4_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 4, q += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 4, q += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 4, q += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 4, q += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 4, q += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 4, q += 4) { \ + MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=8, N=1) +#define POPCNT_UINT32_8X1_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + cnt; \ + const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=8, N=2) +#define POPCNT_UINT32_8X2_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=8, N=4) +#define POPCNT_UINT32_8X4_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=8, N=8) +#define POPCNT_UINT32_8X8_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 3); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=1) +#define POPCNT_UINT32_16X1_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + cnt; \ + const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=2) +#define POPCNT_UINT32_16X2_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=4) +#define POPCNT_UINT32_16X4_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=8) +#define POPCNT_UINT32_16X8_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 3); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=16) +#define POPCNT_UINT32_16X16_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 4); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 4) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=1) +#define POPCNT_UINT32_32X1_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 1, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + cnt; \ + const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=2) +#define POPCNT_UINT32_32X2_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 2, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=4) +#define POPCNT_UINT32_32X4_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 4, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=8) +#define POPCNT_UINT32_32X8_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 8, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 3); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=16) +#define POPCNT_UINT32_32X16_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 16, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 4); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 4) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 16, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 16, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 16, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 16, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=32) +#define POPCNT_UINT32_32X32_SSE(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 32, __m128i, xmm_sum, _mm_setzero_si128()) \ + const uint32_t *qe_0 = q + (cnt << 5); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 5) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 5) : qe_0); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; q != qe_1; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 32, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 32, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_load_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP1_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 32, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \ + for (; q != qe_2; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP2_SSE) \ + } \ + MATRIX_VAR_PERMUTE(8, 32, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \ + for (; q != qe_0; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_loadu_si128, \ + POPCNT_UINT32_STEP3_SSE) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=2, N=1) +#define POPCNT_UINT32_2X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + ((cnt >> 2) << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 2) << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 2) << 2) : qe_0); \ + const uint32_t *qe_3 = q + cnt; \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + __m128i xmm_sum_0 = _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + if (qe_3 >= qe_0 + 2) { \ + __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \ + __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0) \ + m += 4; \ + q += 2; \ + } \ + xmm_sum_0 = _mm_add_epi32( \ + xmm_sum_0, _mm_shuffle_epi32(xmm_sum_0, _MM_SHUFFLE(0, 0, 3, 2))); \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_set_epi32(0, 0, m[1], m[0]); \ + __m128i xmm_q = _mm_broadcast_si32(q); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0)); + +//! Compute the distance between matrix and query (UINT32, M=2, N=2) +#define POPCNT_UINT32_2X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + ((cnt >> 2) << 3); \ + const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 2) << 3) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 2) << 3) : qe_0); \ + const uint32_t *qe_3 = q + (cnt << 1); \ + if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + if (qe_3 >= qe_0 + 4) { \ + __m128i xmm_q = _mm_loadu_si128((const __m128i *)(q)); \ + __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \ + __m128i xmm_p = _mm_shuffle_epi32(xmm_q, _MM_SHUFFLE(2, 2, 0, 0)); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_p, xmm_sum_0_0) \ + xmm_p = _mm_shuffle_epi32(xmm_q, _MM_SHUFFLE(3, 3, 1, 1)); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_p, xmm_sum_0_1) \ + m += 4; \ + q += 4; \ + } \ + xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \ + _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_set_epi32(m[1], m[0], m[1], m[0]); \ + __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=4, N=1) +#define POPCNT_UINT32_4X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + ((cnt >> 1) << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 1) : qe_0); \ + const uint32_t *qe_3 = q + cnt; \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \ + __m128i xmm_q = _mm_broadcast_si32(q); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=4, N=2) +#define POPCNT_UINT32_4X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + ((cnt >> 1) << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 2) : qe_0); \ + const uint32_t *qe_3 = q + (cnt << 1); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \ + __m128i xmm_q = _mm_broadcast_si32(q); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + xmm_q = _mm_broadcast_si32(q + 1); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_1) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=4, N=4) +#define POPCNT_UINT32_4X4_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + ((cnt >> 1) << 3); \ + const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 3) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 3) : qe_0); \ + const uint32_t *qe_3 = q + (cnt << 2); \ + if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + __m128i xmm_sum_0_0 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \ + _mm256_extracti128_si256(ymm_sum_0_0, 1)); \ + __m128i xmm_sum_0_1 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \ + _mm256_extracti128_si256(ymm_sum_0_1, 1)); \ + __m128i xmm_sum_0_2 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_2), \ + _mm256_extracti128_si256(ymm_sum_0_2, 1)); \ + __m128i xmm_sum_0_3 = \ + _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_3), \ + _mm256_extracti128_si256(ymm_sum_0_3, 1)); \ + if (q != qe_3) { \ + __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \ + __m128i xmm_q = _mm_broadcast_si32(q); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \ + xmm_q = _mm_broadcast_si32(q + 1); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_1) \ + xmm_q = _mm_broadcast_si32(q + 2); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_2) \ + xmm_q = _mm_broadcast_si32(q + 3); \ + POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_3) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=8, N=1) +#define POPCNT_UINT32_8X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + cnt; \ + const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, ++q) { \ + MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=8, N=2) +#define POPCNT_UINT32_8X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=8, N=4) +#define POPCNT_UINT32_8X4_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=8, N=8) +#define POPCNT_UINT32_8X8_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 3); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=1) +#define POPCNT_UINT32_16X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + cnt; \ + const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, ++q) { \ + MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=2) +#define POPCNT_UINT32_16X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, q += 2) { \ + MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=4) +#define POPCNT_UINT32_16X4_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, q += 4) { \ + MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=8) +#define POPCNT_UINT32_16X8_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 3); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, q += 8) { \ + MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=16, N=16) +#define POPCNT_UINT32_16X16_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 4); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 4) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 16, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 16, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 16, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 16, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 16, q += 16) { \ + MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=1) +#define POPCNT_UINT32_32X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + cnt; \ + const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, ++q) { \ + MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=2) +#define POPCNT_UINT32_32X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 1); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 2) { \ + MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=4) +#define POPCNT_UINT32_32X4_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 2); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 4) { \ + MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=8) +#define POPCNT_UINT32_32X8_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 3); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 8) { \ + MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=16) +#define POPCNT_UINT32_32X16_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 4); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 4) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 16) { \ + MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT32, M=32, N=32) +#define POPCNT_UINT32_32X32_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 32, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint32_t *qe_0 = q + (cnt << 5); \ + const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 5) : qe_0); \ + const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 5) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 32, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 32, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 32, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \ + for (; q != qe_2; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP2_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 32, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \ + for (; q != qe_0; m += 32, q += 32) { \ + MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT32_STEP3_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0x1f) == 0) { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=2, N=1) +#define POPCNT_UINT64_2X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + ((cnt >> 2) << 2); \ + const uint64_t *qe_1 = (cnt > 31 ? q + ((31 >> 2) << 2) : qe_0); \ + const uint64_t *qe_2 = q + cnt; \ + if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT64_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT64_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + if (qe_2 >= qe_0 + 2) { \ + __m256i ymm_m = _mm256_load_si256((const __m256i *)(m)); \ + __m256i ymm_q = _mm256_set_epi64x(q[1], q[1], q[0], q[0]); \ + POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + m += 4; \ + q += 2; \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT64_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT64_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + if (qe_2 >= qe_0 + 2) { \ + __m256i ymm_m = _mm256_loadu_si256((const __m256i *)(m)); \ + __m256i ymm_q = _mm256_set_epi64x(q[1], q[1], q[0], q[0]); \ + POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + m += 4; \ + q += 2; \ + } \ + } \ + ymm_sum_0_0 = _mm256_add_epi64(ymm_sum_0_0, ymm_sum_0_1); \ + ymm_sum_0_0 = _mm256_add_epi64( \ + ymm_sum_0_0, \ + _mm256_permute4x64_epi64(ymm_sum_0_0, _MM_SHUFFLE(0, 0, 3, 2))); \ + if (q != qe_2) { \ + __m256i ymm_m = _mm256_set_epi64x(0, 0, m[1], m[0]); \ + __m256i ymm_q = _mm256_broadcast_si64(q); \ + POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + } \ + _mm_storel_pi((__m64 *)out, _NORM(ymm_sum_0_0)); + +//! Compute the distance between matrix and query (UINT64, M=2, N=2) +#define POPCNT_UINT64_2X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + ((cnt >> 1) << 2); \ + const uint64_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 2) : qe_0); \ + const uint64_t *qe_2 = q + (cnt << 1); \ + if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \ + for (; q != qe_1; m += 4, q += 4) { \ + MATRIX_INT64_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 4, q += 4) { \ + MATRIX_INT64_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 4, q += 4) { \ + MATRIX_INT64_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 4, q += 4) { \ + MATRIX_INT64_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + ymm_sum_0_0 = _mm256_add_epi64( \ + _mm256_inserti128_si256(ymm_sum_0_0, \ + _mm256_castsi256_si128(ymm_sum_0_1), 1), \ + _mm256_inserti128_si256(ymm_sum_0_1, \ + _mm256_extractf128_si256(ymm_sum_0_0, 1), 0)); \ + if (q != qe_2) { \ + __m256i ymm_m = _mm256_set_epi64x(m[1], m[0], m[1], m[0]); \ + __m256i ymm_q = _mm256_set_epi64x(q[1], q[1], q[0], q[0]); \ + POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=4, N=1) +#define POPCNT_UINT64_4X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + ((cnt >> 1) << 1); \ + const uint64_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 1) : qe_0); \ + const uint64_t *qe_2 = q + cnt; \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT64_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT64_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + if (q != qe_2) { \ + __m256i ymm_m = _mm256_load_si256((const __m256i *)(m)); \ + __m256i ymm_q = _mm256_broadcast_si64(q); \ + POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT64_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT64_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + if (q != qe_2) { \ + __m256i ymm_m = _mm256_loadu_si256((const __m256i *)(m)); \ + __m256i ymm_q = _mm256_broadcast_si64(q); \ + POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \ + } \ + } \ + ymm_sum_0_0 = _mm256_add_epi64(ymm_sum_0_0, ymm_sum_1_0); \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=4, N=2) +#define POPCNT_UINT64_4X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 1); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 4, q += 2) { \ + MATRIX_INT64_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 4, q += 2) { \ + MATRIX_INT64_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 4, q += 2) { \ + MATRIX_INT64_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 4, q += 2) { \ + MATRIX_INT64_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 2, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 2, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=4, N=4) +#define POPCNT_UINT64_4X4_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 2); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 4, q += 4) { \ + MATRIX_INT64_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 4, q += 4) { \ + MATRIX_INT64_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 4, q += 4) { \ + MATRIX_INT64_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 4, q += 4) { \ + MATRIX_INT64_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(1, 4, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(1, 4, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=8, N=1) +#define POPCNT_UINT64_8X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + cnt; \ + const uint64_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, ++q) { \ + MATRIX_INT64_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, ++q) { \ + MATRIX_INT64_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, ++q) { \ + MATRIX_INT64_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, ++q) { \ + MATRIX_INT64_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=8, N=2) +#define POPCNT_UINT64_8X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 1); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT64_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT64_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 2) { \ + MATRIX_INT64_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 2) { \ + MATRIX_INT64_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 2, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 2, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=8, N=4) +#define POPCNT_UINT64_8X4_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 2); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT64_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT64_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 4) { \ + MATRIX_INT64_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 4) { \ + MATRIX_INT64_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 4, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 4, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=8, N=8) +#define POPCNT_UINT64_8X8_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(2, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 3); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT64_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT64_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 8, q += 8) { \ + MATRIX_INT64_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 8, q += 8) { \ + MATRIX_INT64_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(2, 8, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(2, 8, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=16, N=1) +#define POPCNT_UINT64_16X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + cnt; \ + const uint64_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, ++q) { \ + MATRIX_INT64_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, ++q) { \ + MATRIX_INT64_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, ++q) { \ + MATRIX_INT64_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, ++q) { \ + MATRIX_INT64_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=16, N=2) +#define POPCNT_UINT64_16X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 1); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, q += 2) { \ + MATRIX_INT64_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, q += 2) { \ + MATRIX_INT64_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 2) { \ + MATRIX_INT64_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, q += 2) { \ + MATRIX_INT64_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 2, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 2, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=16, N=4) +#define POPCNT_UINT64_16X4_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 2); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, q += 4) { \ + MATRIX_INT64_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, q += 4) { \ + MATRIX_INT64_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 4) { \ + MATRIX_INT64_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, q += 4) { \ + MATRIX_INT64_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 4, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 4, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=16, N=8) +#define POPCNT_UINT64_16X8_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 3); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, q += 8) { \ + MATRIX_INT64_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, q += 8) { \ + MATRIX_INT64_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 8) { \ + MATRIX_INT64_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, q += 8) { \ + MATRIX_INT64_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 8, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 8, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=16, N=16) +#define POPCNT_UINT64_16X16_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(4, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 4); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 16, q += 16) { \ + MATRIX_INT64_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, q += 16) { \ + MATRIX_INT64_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 16, q += 16) { \ + MATRIX_INT64_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 16, q += 16) { \ + MATRIX_INT64_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(4, 16, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(4, 16, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=32, N=1) +#define POPCNT_UINT64_32X1_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + cnt; \ + const uint64_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, ++q) { \ + MATRIX_INT64_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, ++q) { \ + MATRIX_INT64_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, ++q) { \ + MATRIX_INT64_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, ++q) { \ + MATRIX_INT64_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=32, N=2) +#define POPCNT_UINT64_32X2_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 1); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 2) { \ + MATRIX_INT64_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 2) { \ + MATRIX_INT64_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 2) { \ + MATRIX_INT64_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 2) { \ + MATRIX_INT64_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 2, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 2, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=32, N=4) +#define POPCNT_UINT64_32X4_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 2); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 4) { \ + MATRIX_INT64_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 4) { \ + MATRIX_INT64_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 4) { \ + MATRIX_INT64_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 4) { \ + MATRIX_INT64_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 4, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 4, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=32, N=8) +#define POPCNT_UINT64_32X8_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 3); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 8) { \ + MATRIX_INT64_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 8) { \ + MATRIX_INT64_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 8) { \ + MATRIX_INT64_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 8) { \ + MATRIX_INT64_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 8, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 8, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=32, N=16) +#define POPCNT_UINT64_32X16_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 4); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 16) { \ + MATRIX_INT64_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 16, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 16) { \ + MATRIX_INT64_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 16) { \ + MATRIX_INT64_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 16, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 16) { \ + MATRIX_INT64_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 16, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 16, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } + +//! Compute the distance between matrix and query (UINT64, M=32, N=32) +#define POPCNT_UINT64_32X32_AVX(m, q, cnt, out, _NORM) \ + MATRIX_VAR_INIT(8, 32, __m256i, ymm_sum, _mm256_setzero_si256()) \ + const uint64_t *qe_0 = q + (cnt << 5); \ + const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 5) : qe_0); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; q != qe_1; m += 32, q += 32) { \ + MATRIX_INT64_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 32, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 32) { \ + MATRIX_INT64_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } else { \ + for (; q != qe_1; m += 32, q += 32) { \ + MATRIX_INT64_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP1_AVX) \ + } \ + MATRIX_VAR_PERMUTE(8, 32, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \ + for (; q != qe_0; m += 32, q += 32) { \ + MATRIX_INT64_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \ + POPCNT_UINT64_STEP2_AVX) \ + } \ + } \ + if (((uintptr_t)out & 0xf) == 0) { \ + MATRIX_VAR_STORE(8, 32, 4, ymm_sum, out, _mm_store_ps, _NORM) \ + } else { \ + MATRIX_VAR_STORE(8, 32, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \ + } diff --git a/src/ailego/math/distance_utility.h b/src/ailego/math/distance_utility.h new file mode 100644 index 0000000..b892e2a --- /dev/null +++ b/src/ailego/math/distance_utility.h @@ -0,0 +1,65 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +/*! Four-bits Squared Difference Table + */ +static const AILEGO_ALIGNED(64) uint8_t Int4SquaredDiffTable[256] = { + 0, 1, 4, 9, 16, 25, 36, 49, 64, 49, 36, 25, 16, 9, 4, 1, + 1, 0, 1, 4, 9, 16, 25, 36, 81, 64, 49, 36, 25, 16, 9, 4, + 4, 1, 0, 1, 4, 9, 16, 25, 100, 81, 64, 49, 36, 25, 16, 9, + 9, 4, 1, 0, 1, 4, 9, 16, 121, 100, 81, 64, 49, 36, 25, 16, + 16, 9, 4, 1, 0, 1, 4, 9, 144, 121, 100, 81, 64, 49, 36, 25, + 25, 16, 9, 4, 1, 0, 1, 4, 169, 144, 121, 100, 81, 64, 49, 36, + 36, 25, 16, 9, 4, 1, 0, 1, 196, 169, 144, 121, 100, 81, 64, 49, + 49, 36, 25, 16, 9, 4, 1, 0, 225, 196, 169, 144, 121, 100, 81, 64, + 64, 81, 100, 121, 144, 169, 196, 225, 0, 1, 4, 9, 16, 25, 36, 49, + 49, 64, 81, 100, 121, 144, 169, 196, 1, 0, 1, 4, 9, 16, 25, 36, + 36, 49, 64, 81, 100, 121, 144, 169, 4, 1, 0, 1, 4, 9, 16, 25, + 25, 36, 49, 64, 81, 100, 121, 144, 9, 4, 1, 0, 1, 4, 9, 16, + 16, 25, 36, 49, 64, 81, 100, 121, 16, 9, 4, 1, 0, 1, 4, 9, + 9, 16, 25, 36, 49, 64, 81, 100, 25, 16, 9, 4, 1, 0, 1, 4, + 4, 9, 16, 25, 36, 49, 64, 81, 36, 25, 16, 9, 4, 1, 0, 1, + 1, 4, 9, 16, 25, 36, 49, 64, 49, 36, 25, 16, 9, 4, 1, 0, +}; + +/*! Four-bits Integer Multiplication Table + */ +static const AILEGO_ALIGNED(64) int8_t Int4MulTable[256] = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1, + 0, 2, 4, 6, 8, 10, 12, 14, -16, -14, -12, -10, -8, -6, -4, -2, + 0, 3, 6, 9, 12, 15, 18, 21, -24, -21, -18, -15, -12, -9, -6, -3, + 0, 4, 8, 12, 16, 20, 24, 28, -32, -28, -24, -20, -16, -12, -8, -4, + 0, 5, 10, 15, 20, 25, 30, 35, -40, -35, -30, -25, -20, -15, -10, -5, + 0, 6, 12, 18, 24, 30, 36, 42, -48, -42, -36, -30, -24, -18, -12, -6, + 0, 7, 14, 21, 28, 35, 42, 49, -56, -49, -42, -35, -28, -21, -14, -7, + 0, -8, -16, -24, -32, -40, -48, -56, 64, 56, 48, 40, 32, 24, 16, 8, + 0, -7, -14, -21, -28, -35, -42, -49, 56, 49, 42, 35, 28, 21, 14, 7, + 0, -6, -12, -18, -24, -30, -36, -42, 48, 42, 36, 30, 24, 18, 12, 6, + 0, -5, -10, -15, -20, -25, -30, -35, 40, 35, 30, 25, 20, 15, 10, 5, + 0, -4, -8, -12, -16, -20, -24, -28, 32, 28, 24, 20, 16, 12, 8, 4, + 0, -3, -6, -9, -12, -15, -18, -21, 24, 21, 18, 15, 12, 9, 6, 3, + 0, -2, -4, -6, -8, -10, -12, -14, 16, 14, 12, 10, 8, 6, 4, 2, + 0, -1, -2, -3, -4, -5, -6, -7, 8, 7, 6, 5, 4, 3, 2, 1, +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/euclidean_distance_matrix.h b/src/ailego/math/euclidean_distance_matrix.h new file mode 100644 index 0000000..e774093 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix.h @@ -0,0 +1,649 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include "distance_utility.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +/*! Squared Euclidean Distance Matrix + */ +template +struct SquaredEuclideanDistanceMatrix; + +/*! Squared Euclidean Distance Matrix (M=1, N=1) + */ +template +struct SquaredEuclideanDistanceMatrix< + T, 1, 1, typename std::enable_if::value>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += MathHelper::SquaredDifference(m[i], q[i]); + } + *out = sum; + } +}; + +template <> +struct SquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct SquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct SquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = Float16; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct SquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = float; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Squared Euclidean Distance Matrix + */ +template +struct SquaredEuclideanDistanceMatrix< + T, M, N, + typename std::enable_if::value && sizeof(T) >= 2 && + M >= 2 && N >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + ValueType m_val = m[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = MathHelper::SquaredDifference(m_val, q[j]); + r += M; + } + } + m += M; + q += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + ValueType m_val = m[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r += MathHelper::SquaredDifference(m_val, q[j]); + r += M; + } + } + m += M; + q += N; + } + } +}; + +/*! Squared Euclidean Distance Matrix (N=1) + */ +template +struct SquaredEuclideanDistanceMatrix< + T, M, 1, + typename std::enable_if::value && sizeof(T) >= 2 && + M >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + const ValueType *q_end = q + dim; + if (q != q_end) { + ValueType q_val = *q++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = MathHelper::SquaredDifference(m[i], q_val); + } + m += M; + } + + while (q != q_end) { + ValueType q_val = *q++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) += MathHelper::SquaredDifference(m[i], q_val); + } + m += M; + } + } +}; + +/*! Squared Euclidean Distance Matrix (INT8) + */ +template +struct SquaredEuclideanDistanceMatrix< + int8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + + dim >>= 2; + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = SquaredDifference(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r += SquaredDifference(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + } + + protected: + //! Calculate the squared difference + static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) { + volatile int32_t sum = MathHelper::SquaredDifference( + (int8_t)(lhs >> 0), (int8_t)(rhs >> 0)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 8), (int8_t)(rhs >> 8)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 16), (int8_t)(rhs >> 16)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 24), (int8_t)(rhs >> 24)); + return static_cast(sum); + } +}; + +/*! Squared Euclidean Distance Matrix (INT8, N=1) + */ +template +struct SquaredEuclideanDistanceMatrix= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 2); + + if (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = SquaredDifference(m_it[i], q_val); + } + m_it += M; + } + + while (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) += SquaredDifference(m_it[i], q_val); + } + m_it += M; + } + } + + protected: + //! Calculate the squared difference + static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) { + volatile int32_t sum = MathHelper::SquaredDifference( + (int8_t)(lhs >> 0), (int8_t)(rhs >> 0)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 8), (int8_t)(rhs >> 8)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 16), (int8_t)(rhs >> 16)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 24), (int8_t)(rhs >> 24)); + return static_cast(sum); + } +}; + +/*! Squared Euclidean Distance Matrix (INT4) + */ +template +struct SquaredEuclideanDistanceMatrix< + uint8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 7) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + + dim >>= 3; + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = SquaredDifference(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r += SquaredDifference(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + } + + protected: + //! Calculate the squared difference + static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4SquaredDiffTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } +}; + +/*! Squared Euclidean Distance Matrix (INT4, N=1) + */ +template +struct SquaredEuclideanDistanceMatrix= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 7) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 3); + + if (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = SquaredDifference(m_it[i], q_val); + } + m_it += M; + } + + while (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) += SquaredDifference(m_it[i], q_val); + } + m_it += M; + } + } + + protected: + //! Calculate the squared difference + static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4SquaredDiffTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } +}; + +/*! Euclidean Distance Matrix + */ +template ::value || + std::is_same::value) && + M >= 1 && N >= 1>::type> +struct EuclideanDistanceMatrix { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out); + for (size_t i = 0; i < N * M; ++i) { + float val = *out; + *out++ = std::sqrt(val); + } + } +}; + +/*! Euclidean Distance Matrix (M=1, N=1) + */ +template +struct EuclideanDistanceMatrix< + T, 1, 1, typename std::enable_if::value>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += MathHelper::SquaredDifference(m[i], q[i]); + } + *out = std::sqrt(sum); + } +}; + +template <> +struct EuclideanDistanceMatrix { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct EuclideanDistanceMatrix { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct EuclideanDistanceMatrix { + //! Type of value + using ValueType = Float16; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct EuclideanDistanceMatrix { + //! Type of value + using ValueType = float; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + + +//-------------------------------------------------- +// Sparse +//-------------------------------------------------- +/*! Squared Euclidean Distance Sparse Matrix + */ +template +struct SquaredEuclideanSparseDistanceMatrix { + //! Type of value + using ValueType = typename std::remove_cv::type; + + static float ComputeSquaredEuclideanSparseDistanceInSegment( + uint32_t m_sparse_count, const uint16_t *m_sparse_index, + const ValueType *m_sparse_value, uint32_t q_sparse_count, + const uint16_t *q_sparse_index, const ValueType *q_sparse_value); + + //! Compute the distance between matrix and query + static inline void Compute(const void *m_sparse_data_in, + const void *q_sparse_data_in, float *out) { + ailego_assert(out); + + const uint8_t *m_sparse_data = + reinterpret_cast(m_sparse_data_in); + const uint8_t *q_sparse_data = + reinterpret_cast(q_sparse_data_in); + + const uint32_t m_sparse_count = + *reinterpret_cast(m_sparse_data); + const uint32_t q_sparse_count = + *reinterpret_cast(q_sparse_data); + + const uint32_t m_seg_count = + *reinterpret_cast(m_sparse_data + sizeof(uint32_t)); + const uint32_t q_seg_count = + *reinterpret_cast(q_sparse_data + sizeof(uint32_t)); + + const uint32_t *m_seg_id = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t)); + const uint32_t *q_seg_id = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t)); + + const uint32_t *m_seg_vec_cnt = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t) + m_seg_count * sizeof(uint32_t)); + const uint32_t *q_seg_vec_cnt = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t) + q_seg_count * sizeof(uint32_t)); + + const uint16_t *m_sparse_index = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t) + + m_seg_count * 2 * sizeof(uint32_t)); + const uint16_t *q_sparse_index = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t) + + q_seg_count * 2 * sizeof(uint32_t)); + + const ValueType *m_sparse_value = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t) + + m_seg_count * 2 * sizeof(uint32_t) + m_sparse_count * sizeof(uint16_t)); + const ValueType *q_sparse_value = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t) + + q_seg_count * 2 * sizeof(uint32_t) + q_sparse_count * sizeof(uint16_t)); + + float sum = 0.0f; + + size_t m_s = 0; + size_t q_s = 0; + + size_t m_count = 0; + size_t q_count = 0; + + while (m_s < m_seg_count && q_s < q_seg_count) { + if (m_seg_id[m_s] == q_seg_id[q_s]) { + sum += ComputeSquaredEuclideanSparseDistanceInSegment( + m_seg_vec_cnt[m_s], m_sparse_index + m_count, + m_sparse_value + m_count, q_seg_vec_cnt[q_s], + q_sparse_index + q_count, q_sparse_value + q_count); + + m_count += m_seg_vec_cnt[m_s]; + q_count += q_seg_vec_cnt[q_s]; + + ++m_s; + ++q_s; + } else if (m_seg_id[m_s] < q_seg_id[q_s]) { + for (size_t i = 0; i < m_seg_vec_cnt[m_s]; i++) { + float value = (m_sparse_value + m_count)[i]; + sum += value * value; + } + + m_count += m_seg_vec_cnt[m_s]; + + ++m_s; + } else { + for (size_t i = 0; i < q_seg_vec_cnt[q_s]; i++) { + float value = (q_sparse_value + q_count)[i]; + sum += value * value; + } + + q_count += q_seg_vec_cnt[q_s]; + ++q_s; + } + } + + for (; m_s < m_seg_count; m_s++) { + for (size_t i = 0; i < m_seg_vec_cnt[m_s]; i++) { + float diff = (m_sparse_value + m_count)[i]; + sum += diff * diff; + } + + m_count += m_seg_vec_cnt[m_s]; + } + + for (; q_s < q_seg_count; q_s++) { + for (size_t i = 0; i < q_seg_vec_cnt[q_s]; i++) { + float diff = (q_sparse_value + q_count)[i]; + sum += diff * diff; + } + + q_count += q_seg_vec_cnt[q_s]; + } + + *out = sum; + } +}; + +template +float SquaredEuclideanSparseDistanceMatrix:: + ComputeSquaredEuclideanSparseDistanceInSegment( + uint32_t m_sparse_count, const uint16_t *m_sparse_index, + const ValueType *m_sparse_value, uint32_t q_sparse_count, + const uint16_t *q_sparse_index, const ValueType *q_sparse_value) { + float sum = 0.0f; + + size_t m_i = 0; + size_t q_i = 0; + + while (m_i < m_sparse_count && q_i < q_sparse_count) { + if (m_sparse_index[m_i] == q_sparse_index[q_i]) { + float diff = m_sparse_value[m_i] - q_sparse_value[q_i]; + sum += diff * diff; + ++m_i; + ++q_i; + } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) { + float diff = m_sparse_value[m_i]; + sum += diff * diff; + ++m_i; + } else { + float diff = q_sparse_value[q_i]; + sum += diff * diff; + + ++q_i; + } + } + + for (; m_i < m_sparse_count; m_i++) { + float diff = m_sparse_value[m_i]; + sum += diff * diff; + } + + for (; q_i < q_sparse_count; q_i++) { + float diff = q_sparse_value[q_i]; + sum += diff * diff; + } + + return sum; +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/euclidean_distance_matrix_fp16_avx.cc b/src/ailego/math/euclidean_distance_matrix_fp16_avx.cc new file mode 100644 index 0000000..7258b25 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp16_avx.cc @@ -0,0 +1,36 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX__) + +float SquaredEuclideanDistanceFp16AVX(const Float16 *lhs, const Float16 *rhs, + size_t size) { + float score{0.0f}; + + ACCUM_FP16_1X1_AVX(lhs, rhs, size, &score, 0ull, ) + + return score; +} + +#endif // __AVX__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_fp16_avx512.cc b/src/ailego/math/euclidean_distance_matrix_fp16_avx512.cc new file mode 100644 index 0000000..df97f40 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp16_avx512.cc @@ -0,0 +1,33 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX512F__) +float SquaredEuclideanDistanceFp16AVX512(const Float16 *lhs, const Float16 *rhs, + size_t size) { + float score{0.0f}; + + ACCUM_FP16_1X1_AVX512(lhs, rhs, size, &score, 0ull, ) + + return score; +} +#endif +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_fp16_avx512fp16.cc b/src/ailego/math/euclidean_distance_matrix_fp16_avx512fp16.cc new file mode 100644 index 0000000..b0e862e --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp16_avx512fp16.cc @@ -0,0 +1,82 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX512FP16__) +//! Squared Euclidean Distance +float SquaredEuclideanDistanceFp16AVX512FP16(const Float16 *lhs, + const Float16 *rhs, size_t size) { + const Float16 *last = lhs + size; + const Float16 *last_aligned = lhs + ((size >> 6) << 6); + + __m512h zmm_sum_0 = _mm512_setzero_ph(); + __m512h zmm_sum_1 = _mm512_setzero_ph(); + + if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + __m512h zmm_d_0 = + _mm512_sub_ph(_mm512_load_ph(lhs + 0), _mm512_load_ph(rhs + 0)); + __m512h zmm_d_1 = + _mm512_sub_ph(_mm512_load_ph(lhs + 32), _mm512_load_ph(rhs + 32)); + zmm_sum_0 = _mm512_fmadd_ph(zmm_d_0, zmm_d_0, zmm_sum_0); + zmm_sum_1 = _mm512_fmadd_ph(zmm_d_1, zmm_d_1, zmm_sum_1); + } + + if (last >= last_aligned + 32) { + __m512h zmm_d = _mm512_sub_ph(_mm512_load_ph(lhs), _mm512_load_ph(rhs)); + zmm_sum_0 = _mm512_fmadd_ph(zmm_d, zmm_d, zmm_sum_0); + lhs += 32; + rhs += 32; + } + } else { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + __m512h zmm_d_0 = + _mm512_sub_ph(_mm512_loadu_ph(lhs + 0), _mm512_loadu_ph(rhs + 0)); + __m512h zmm_d_1 = + _mm512_sub_ph(_mm512_loadu_ph(lhs + 32), _mm512_loadu_ph(rhs + 32)); + zmm_sum_0 = _mm512_fmadd_ph(zmm_d_0, zmm_d_0, zmm_sum_0); + zmm_sum_1 = _mm512_fmadd_ph(zmm_d_1, zmm_d_1, zmm_sum_1); + } + + if (last >= last_aligned + 32) { + __m512h zmm_d = _mm512_sub_ph(_mm512_loadu_ph(lhs), _mm512_loadu_ph(rhs)); + zmm_sum_0 = _mm512_fmadd_ph(zmm_d, zmm_d, zmm_sum_0); + lhs += 32; + rhs += 32; + } + } + + zmm_sum_0 = _mm512_add_ph(zmm_sum_0, zmm_sum_1); + if (lhs != last) { + __mmask32 mask = (__mmask32)((1 << (last - lhs)) - 1); + __m512i zmm_undefined = _mm512_undefined_epi32(); + __m512h zmm_undefined_ph = _mm512_undefined_ph(); + __m512h zmm_d = _mm512_mask_sub_ph( + zmm_undefined_ph, mask, + _mm512_castsi512_ph(_mm512_mask_loadu_epi16(zmm_undefined, mask, lhs)), + _mm512_castsi512_ph(_mm512_mask_loadu_epi16(zmm_undefined, mask, rhs))); + zmm_sum_0 = _mm512_mask3_fmadd_ph(zmm_d, zmm_d, zmm_sum_0, mask); + } + + return HorizontalAdd_FP16_V512(zmm_sum_0); +} +#endif +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/euclidean_distance_matrix_fp16_dispatch.cc b/src/ailego/math/euclidean_distance_matrix_fp16_dispatch.cc new file mode 100644 index 0000000..fb14526 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp16_dispatch.cc @@ -0,0 +1,85 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) +float SquaredEuclideanDistanceFp16NEON(const Float16 *lhs, const Float16 *rhs, + size_t size); +#endif + +#if defined(__AVX512FP16__) +float SquaredEuclideanDistanceFp16AVX512FP16(const Float16 *lhs, + const Float16 *rhs, size_t size); +#endif + +#if defined(__AVX512F__) +float SquaredEuclideanDistanceFp16AVX512(const Float16 *lhs, const Float16 *rhs, + size_t size); +#endif + +#if defined(__AVX__) +float SquaredEuclideanDistanceFp16AVX(const Float16 *lhs, const Float16 *rhs, + size_t size); +#endif + +float SquaredEuclideanDistanceFp16Scalar(const Float16 *lhs, const Float16 *rhs, + size_t size); + +//! Compute the distance between matrix and query (FP16, M=1, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + *out = SquaredEuclideanDistanceFp16NEON(m, q, dim); +#else +#if defined(__AVX512FP16__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) { + *out = SquaredEuclideanDistanceFp16AVX512FP16(m, q, dim); + return; + } +#endif +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = SquaredEuclideanDistanceFp16AVX512(m, q, dim); + return; + } +#endif + +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = SquaredEuclideanDistanceFp16AVX(m, q, dim); + return; + } +#endif + *out = SquaredEuclideanDistanceFp16Scalar(m, q, dim); + +#endif //__ARM_NEON +} + +//! Compute the distance between matrix and query (FP16, M=1, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out); + *out = std::sqrt(*out); +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_fp16_neon.cc b/src/ailego/math/euclidean_distance_matrix_fp16_neon.cc new file mode 100644 index 0000000..3d3bf87 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp16_neon.cc @@ -0,0 +1,34 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) +float SquaredEuclideanDistanceFp16NEON(const Float16 *lhs, const Float16 *rhs, + size_t size) { + float score{0.0f}; + + ACCUM_FP16_1X1_NEON(lhs, rhs, size, &score, 0ull, ) + + return score; +} +#endif + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_fp32_avx.cc b/src/ailego/math/euclidean_distance_matrix_fp32_avx.cc new file mode 100644 index 0000000..c7f6f5b --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp32_avx.cc @@ -0,0 +1,106 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX__) +float SquaredEuclideanDistanceFp32SSEInternal(const float *lhs, + const float *rhs, size_t size); + +float SquaredEuclideanDistanceFp32AVXInternal(const float *lhs, + const float *rhs, size_t size) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 4) << 4); + + __m256 ymm_sum_0 = _mm256_setzero_ps(); + __m256 ymm_sum_1 = _mm256_setzero_ps(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m256 ymm_d_0 = + _mm256_sub_ps(_mm256_load_ps(lhs + 0), _mm256_load_ps(rhs + 0)); + __m256 ymm_d_1 = + _mm256_sub_ps(_mm256_load_ps(lhs + 8), _mm256_load_ps(rhs + 8)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_d_0, ymm_d_0, ymm_sum_0); + ymm_sum_1 = _mm256_fmadd_ps(ymm_d_1, ymm_d_1, ymm_sum_1); + } + + if (last >= last_aligned + 8) { + __m256 ymm_d = _mm256_sub_ps(_mm256_load_ps(lhs), _mm256_load_ps(rhs)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_d, ymm_d, ymm_sum_0); + lhs += 8; + rhs += 8; + } + } else { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m256 ymm_d_0 = + _mm256_sub_ps(_mm256_loadu_ps(lhs + 0), _mm256_loadu_ps(rhs + 0)); + __m256 ymm_d_1 = + _mm256_sub_ps(_mm256_loadu_ps(lhs + 8), _mm256_loadu_ps(rhs + 8)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_d_0, ymm_d_0, ymm_sum_0); + ymm_sum_1 = _mm256_fmadd_ps(ymm_d_1, ymm_d_1, ymm_sum_1); + } + + if (last >= last_aligned + 8) { + __m256 ymm_d = _mm256_sub_ps(_mm256_loadu_ps(lhs), _mm256_loadu_ps(rhs)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_d, ymm_d, ymm_sum_0); + lhs += 8; + rhs += 8; + } + } + float result = HorizontalAdd_FP32_V256(_mm256_add_ps(ymm_sum_0, ymm_sum_1)); + + switch (last - lhs) { + case 7: + SSD_FP32_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + SSD_FP32_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + SSD_FP32_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + SSD_FP32_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + SSD_FP32_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + SSD_FP32_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + SSD_FP32_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float SquaredEuclideanDistanceFp32AVX(const float *lhs, const float *rhs, + size_t size) { + if (size > 7) { + return SquaredEuclideanDistanceFp32AVXInternal(lhs, rhs, size); + } + + return SquaredEuclideanDistanceFp32SSEInternal(lhs, rhs, size); +} + +#endif // __AVX__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_fp32_avx512.cc b/src/ailego/math/euclidean_distance_matrix_fp32_avx512.cc new file mode 100644 index 0000000..3363a52 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp32_avx512.cc @@ -0,0 +1,100 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX512F__) +float SquaredEuclideanDistanceFp32SSEInternal(const float *lhs, + const float *rhs, size_t size); + +float SquaredEuclideanDistanceFp32AVXInternal(const float *lhs, + const float *rhs, size_t size); + +float SquaredEuclideanDistanceFp32AVX512Internal(const float *lhs, + const float *rhs, + size_t size) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 5) << 5); + + __m512 zmm_sum_0 = _mm512_setzero_ps(); + __m512 zmm_sum_1 = _mm512_setzero_ps(); + + if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m512 zmm_d_0 = + _mm512_sub_ps(_mm512_load_ps(lhs + 0), _mm512_load_ps(rhs + 0)); + __m512 zmm_d_1 = + _mm512_sub_ps(_mm512_load_ps(lhs + 16), _mm512_load_ps(rhs + 16)); + zmm_sum_0 = _mm512_fmadd_ps(zmm_d_0, zmm_d_0, zmm_sum_0); + zmm_sum_1 = _mm512_fmadd_ps(zmm_d_1, zmm_d_1, zmm_sum_1); + } + + if (last >= last_aligned + 16) { + __m512 zmm_d = _mm512_sub_ps(_mm512_load_ps(lhs), _mm512_load_ps(rhs)); + zmm_sum_0 = _mm512_fmadd_ps(zmm_d, zmm_d, zmm_sum_0); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m512 zmm_d_0 = + _mm512_sub_ps(_mm512_loadu_ps(lhs + 0), _mm512_loadu_ps(rhs + 0)); + __m512 zmm_d_1 = + _mm512_sub_ps(_mm512_loadu_ps(lhs + 16), _mm512_loadu_ps(rhs + 16)); + zmm_sum_0 = _mm512_fmadd_ps(zmm_d_0, zmm_d_0, zmm_sum_0); + zmm_sum_1 = _mm512_fmadd_ps(zmm_d_1, zmm_d_1, zmm_sum_1); + } + + if (last >= last_aligned + 16) { + __m512 zmm_d = _mm512_sub_ps(_mm512_loadu_ps(lhs), _mm512_loadu_ps(rhs)); + zmm_sum_0 = _mm512_fmadd_ps(zmm_d, zmm_d, zmm_sum_0); + lhs += 16; + rhs += 16; + } + } + + zmm_sum_0 = _mm512_add_ps(zmm_sum_0, zmm_sum_1); + if (lhs != last) { + __mmask16 mask = (__mmask16)((1 << (last - lhs)) - 1); + __m512 zmm_undefined = _mm512_undefined_ps(); + __m512 zmm_d = _mm512_mask_sub_ps( + zmm_undefined, mask, _mm512_mask_loadu_ps(zmm_undefined, mask, lhs), + _mm512_mask_loadu_ps(zmm_undefined, mask, rhs)); + zmm_sum_0 = _mm512_mask3_fmadd_ps(zmm_d, zmm_d, zmm_sum_0, mask); + } + return HorizontalAdd_FP32_V512(zmm_sum_0); +} + +float SquaredEuclideanDistanceFp32AVX512(const float *lhs, const float *rhs, + size_t size) { + if (size > 15) { + return SquaredEuclideanDistanceFp32AVX512Internal(lhs, rhs, size); + } + + if (size > 7) { + return SquaredEuclideanDistanceFp32AVXInternal(lhs, rhs, size); + } + + return SquaredEuclideanDistanceFp32SSEInternal(lhs, rhs, size); +} + +#endif + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_fp32_dispatch.cc b/src/ailego/math/euclidean_distance_matrix_fp32_dispatch.cc new file mode 100644 index 0000000..cc30443 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp32_dispatch.cc @@ -0,0 +1,90 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) +void SquaredEuclideanDistanceFp32NEON(const float *lhs, const float *rhs, + size_t size, float *out); +#endif + +#if defined(__AVX512F__) +float SquaredEuclideanDistanceFp32AVX512(const float *lhs, const float *rhs, + size_t size); +#endif + +#if defined(__AVX__) +float SquaredEuclideanDistanceFp32AVX(const float *lhs, const float *rhs, + size_t size); +#endif + +#if defined(__SSE__) +float SquaredEuclideanDistanceFp32SSE(const float *lhs, const float *rhs, + size_t size); +#endif + +float SquaredEuclideanDistanceFp32Scalar(const float *lhs, const float *rhs, + size_t size); + +//----------------------------------------------------------- +// SquaredEuclideanDistance +//----------------------------------------------------------- +//! Compute the distance between matrix and query (FP32, M=1, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + SquaredEuclideanDistanceFp32NEON(m, q, dim, out); +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = SquaredEuclideanDistanceFp32AVX512(m, q, dim); + return; + } +#endif // __AVX512F__ +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = SquaredEuclideanDistanceFp32AVX(m, q, dim); + return; + } +#endif // __AVX__ + +#if defined(__SSE__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) { + *out = SquaredEuclideanDistanceFp32SSE(m, q, dim); + return; + } +#endif // __SSE__ + *out = SquaredEuclideanDistanceFp32Scalar(m, q, dim); +#endif // __ARM_NEON +} + +//----------------------------------------------------------- +// EuclideanDistance +//----------------------------------------------------------- +//! Compute the distance between matrix and query (FP32, M=1, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out); + *out = std::sqrt(*out); +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_fp32_neon.cc b/src/ailego/math/euclidean_distance_matrix_fp32_neon.cc new file mode 100644 index 0000000..aa1694e --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp32_neon.cc @@ -0,0 +1,62 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) +//! Squared Euclidean Distance +void SquaredEuclideanDistanceFp32NEON(const float *lhs, const float *rhs, + size_t size, float *out) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 3) << 3); + + float32x4_t v_sum_0 = vdupq_n_f32(0); + float32x4_t v_sum_1 = vdupq_n_f32(0); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + float32x4_t v_d_0 = vsubq_f32(vld1q_f32(lhs + 0), vld1q_f32(rhs + 0)); + float32x4_t v_d_1 = vsubq_f32(vld1q_f32(lhs + 4), vld1q_f32(rhs + 4)); + v_sum_0 = vfmaq_f32(v_sum_0, v_d_0, v_d_0); + v_sum_1 = vfmaq_f32(v_sum_1, v_d_1, v_d_1); + } + if (last >= last_aligned + 4) { + float32x4_t v_d = vsubq_f32(vld1q_f32(lhs), vld1q_f32(rhs)); + v_sum_0 = vfmaq_f32(v_sum_0, v_d, v_d); + lhs += 4; + rhs += 4; + } + + float result = vaddvq_f32(vaddq_f32(v_sum_0, v_sum_1)); + switch (last - lhs) { + case 3: + SSD_FP32_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + SSD_FP32_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + SSD_FP32_GENERAL(lhs[0], rhs[0], result) + } + *out = result; +} + +#endif // __ARM_NEON + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_fp32_sse.cc b/src/ailego/math/euclidean_distance_matrix_fp32_sse.cc new file mode 100644 index 0000000..9574ed6 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp32_sse.cc @@ -0,0 +1,83 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE__) +float SquaredEuclideanDistanceFp32SSEInternal(const float *lhs, + const float *rhs, size_t size) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 3) << 3); + + __m128 xmm_sum_0 = _mm_setzero_ps(); + __m128 xmm_sum_1 = _mm_setzero_ps(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m128 xmm_d_0 = _mm_sub_ps(_mm_load_ps(lhs + 0), _mm_load_ps(rhs + 0)); + __m128 xmm_d_1 = _mm_sub_ps(_mm_load_ps(lhs + 4), _mm_load_ps(rhs + 4)); + xmm_sum_0 = _mm_fmadd_ps(xmm_d_0, xmm_d_0, xmm_sum_0); + xmm_sum_1 = _mm_fmadd_ps(xmm_d_1, xmm_d_1, xmm_sum_1); + } + + if (last >= last_aligned + 4) { + __m128 xmm_d = _mm_sub_ps(_mm_load_ps(lhs), _mm_load_ps(rhs)); + xmm_sum_0 = _mm_fmadd_ps(xmm_d, xmm_d, xmm_sum_0); + lhs += 4; + rhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m128 xmm_d_0 = _mm_sub_ps(_mm_loadu_ps(lhs + 0), _mm_loadu_ps(rhs + 0)); + __m128 xmm_d_1 = _mm_sub_ps(_mm_loadu_ps(lhs + 4), _mm_loadu_ps(rhs + 4)); + xmm_sum_0 = _mm_fmadd_ps(xmm_d_0, xmm_d_0, xmm_sum_0); + xmm_sum_1 = _mm_fmadd_ps(xmm_d_1, xmm_d_1, xmm_sum_1); + } + + if (last >= last_aligned + 4) { + __m128 xmm_d = _mm_sub_ps(_mm_loadu_ps(lhs), _mm_loadu_ps(rhs)); + xmm_sum_0 = _mm_fmadd_ps(xmm_d, xmm_d, xmm_sum_0); + lhs += 4; + rhs += 4; + } + } + float result = HorizontalAdd_FP32_V128(_mm_add_ps(xmm_sum_0, xmm_sum_1)); + + switch (last - lhs) { + case 3: + SSD_FP32_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + SSD_FP32_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + SSD_FP32_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float SquaredEuclideanDistanceFp32SSE(const float *lhs, const float *rhs, + size_t size) { + return SquaredEuclideanDistanceFp32SSEInternal(lhs, rhs, size); +} + +#endif // __SSE__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_int4_avx2.cc b/src/ailego/math/euclidean_distance_matrix_int4_avx2.cc new file mode 100644 index 0000000..dacb278 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_int4_avx2.cc @@ -0,0 +1,130 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int4.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX2__) +float SquaredEuclideanDistanceInt4SSEInternal(const uint8_t *lhs, + const uint8_t *rhs, size_t size); + +inline float SquaredEuclideanDistanceInt4AVX2Internal(const uint8_t *lhs, + const uint8_t *rhs, + size_t size) { + const uint8_t *last = lhs + size; + const uint8_t *last_aligned = lhs + ((size >> 5) << 5); + + __m256i ymm_sum = _mm256_setzero_si256(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m256i ymm_lhs = _mm256_load_si256((const __m256i *)(lhs)); + __m256i ymm_rhs = _mm256_load_si256((const __m256i *)(rhs)); + SSD_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum) + } + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs); + __m128i xmm_sum = _mm_setzero_si128(); + SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) + ymm_sum = _mm256_add_epi32(_mm256_set_m128i(_mm_setzero_si128(), xmm_sum), + ymm_sum); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)(lhs)); + __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)(rhs)); + SSD_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum) + } + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + __m128i xmm_sum = _mm_setzero_si128(); + SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) + ymm_sum = _mm256_add_epi32(_mm256_set_m128i(_mm_setzero_si128(), xmm_sum), + ymm_sum); + lhs += 16; + rhs += 16; + } + } + float result = static_cast(HorizontalAdd_INT32_V256(ymm_sum)); + + switch (last - lhs) { + case 15: + SSD_INT4_GENERAL(lhs[14], rhs[14], result) + /* FALLTHRU */ + case 14: + SSD_INT4_GENERAL(lhs[13], rhs[13], result) + /* FALLTHRU */ + case 13: + SSD_INT4_GENERAL(lhs[12], rhs[12], result) + /* FALLTHRU */ + case 12: + SSD_INT4_GENERAL(lhs[11], rhs[11], result) + /* FALLTHRU */ + case 11: + SSD_INT4_GENERAL(lhs[10], rhs[10], result) + /* FALLTHRU */ + case 10: + SSD_INT4_GENERAL(lhs[9], rhs[9], result) + /* FALLTHRU */ + case 9: + SSD_INT4_GENERAL(lhs[8], rhs[8], result) + /* FALLTHRU */ + case 8: + SSD_INT4_GENERAL(lhs[7], rhs[7], result) + /* FALLTHRU */ + case 7: + SSD_INT4_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + SSD_INT4_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + SSD_INT4_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + SSD_INT4_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + SSD_INT4_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + SSD_INT4_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + SSD_INT4_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float SquaredEuclideanDistanceInt4AVX2(const uint8_t *lhs, const uint8_t *rhs, + size_t size) { + if (size > 63) { + return SquaredEuclideanDistanceInt4AVX2Internal(lhs, rhs, size >> 1); + } + + return SquaredEuclideanDistanceInt4SSEInternal(lhs, rhs, size >> 1); +} + +#endif // __AVX2__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_int4_dispatch.cc b/src/ailego/math/euclidean_distance_matrix_int4_dispatch.cc new file mode 100644 index 0000000..d4ff74d --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_int4_dispatch.cc @@ -0,0 +1,65 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX2__) +float SquaredEuclideanDistanceInt4AVX2(const uint8_t *lhs, const uint8_t *rhs, + size_t size); +#endif + +#if defined(__SSE4_1__) +float SquaredEuclideanDistanceInt4SSE(const uint8_t *lhs, const uint8_t *rhs, + size_t size); +#endif + +float SquaredEuclideanDistanceInt4Scalar(const uint8_t *lhs, const uint8_t *rhs, + size_t size); + +//! Compute the distance between matrix and query (INT4, M=1, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = SquaredEuclideanDistanceInt4AVX2(m, q, dim); + return; + } +#endif // __AVX2__ + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = SquaredEuclideanDistanceInt4SSE(m, q, dim); + return; + } +#endif + + *out = SquaredEuclideanDistanceInt4Scalar(m, q, dim); +} + +//! Compute the distance between matrix and query (INT4, M=1, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out); + *out = std::sqrt(*out); +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_int4_sse.cc b/src/ailego/math/euclidean_distance_matrix_int4_sse.cc new file mode 100644 index 0000000..1e998ea --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_int4_sse.cc @@ -0,0 +1,102 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int4.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE4_1__) +float SquaredEuclideanDistanceInt4SSEInternal(const uint8_t *lhs, + const uint8_t *rhs, size_t size) { + const uint8_t *last = lhs + size; + const uint8_t *last_aligned = lhs + ((size >> 4) << 4); + + __m128i xmm_sum = _mm_setzero_si128(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)(lhs)); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)(rhs)); + SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) + } + } else { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)(lhs)); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)(rhs)); + SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) + } + } + float result = static_cast(HorizontalAdd_INT32_V128(xmm_sum)); + + switch (last - lhs) { + case 15: + SSD_INT4_GENERAL(lhs[14], rhs[14], result) + /* FALLTHRU */ + case 14: + SSD_INT4_GENERAL(lhs[13], rhs[13], result) + /* FALLTHRU */ + case 13: + SSD_INT4_GENERAL(lhs[12], rhs[12], result) + /* FALLTHRU */ + case 12: + SSD_INT4_GENERAL(lhs[11], rhs[11], result) + /* FALLTHRU */ + case 11: + SSD_INT4_GENERAL(lhs[10], rhs[10], result) + /* FALLTHRU */ + case 10: + SSD_INT4_GENERAL(lhs[9], rhs[9], result) + /* FALLTHRU */ + case 9: + SSD_INT4_GENERAL(lhs[8], rhs[8], result) + /* FALLTHRU */ + case 8: + SSD_INT4_GENERAL(lhs[7], rhs[7], result) + /* FALLTHRU */ + case 7: + SSD_INT4_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + SSD_INT4_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + SSD_INT4_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + SSD_INT4_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + SSD_INT4_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + SSD_INT4_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + SSD_INT4_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float SquaredEuclideanDistanceInt4SSE(const uint8_t *lhs, const uint8_t *rhs, + size_t size) { + return SquaredEuclideanDistanceInt4SSEInternal(lhs, rhs, size >> 1); +} + +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_int8_avx2.cc b/src/ailego/math/euclidean_distance_matrix_int8_avx2.cc new file mode 100644 index 0000000..ef46589 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_int8_avx2.cc @@ -0,0 +1,192 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int8.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX2__) +float SquaredEuclideanDistanceInt8SSEInternal(const int8_t *lhs, + const int8_t *rhs, size_t size); + +float SquaredEuclideanDistanceInt8AVX2Internal(const int8_t *lhs, + const int8_t *rhs, size_t size) { + const int8_t *last = lhs + size; + const int8_t *last_aligned = lhs + ((size >> 6) << 6); + float result = 0.0; + + __m256i ymm_sum_0 = _mm256_setzero_si256(); + __m256i ymm_sum_1 = _mm256_setzero_si256(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + __m256i ymm_lhs_0 = _mm256_load_si256((const __m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_load_si256((const __m256i *)(lhs + 32)); + __m256i ymm_rhs_0 = _mm256_load_si256((const __m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_load_si256((const __m256i *)(rhs + 32)); + + __m256i ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_0, ymm_rhs_0), + _mm256_min_epi8(ymm_lhs_0, ymm_rhs_0)); + ymm_lhs_0 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d)); + ymm_rhs_0 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1)); + ymm_sum_0 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs_0, ymm_lhs_0), ymm_sum_0); + ymm_sum_1 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs_0, ymm_rhs_0), ymm_sum_1); + + ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_1, ymm_rhs_1), + _mm256_min_epi8(ymm_lhs_1, ymm_rhs_1)); + ymm_lhs_1 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d)); + ymm_rhs_1 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1)); + ymm_sum_0 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs_1, ymm_lhs_1), ymm_sum_0); + ymm_sum_1 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs_1, ymm_rhs_1), ymm_sum_1); + } + + if (last >= last_aligned + 32) { + __m256i ymm_lhs = _mm256_load_si256((const __m256i *)lhs); + __m256i ymm_rhs = _mm256_load_si256((const __m256i *)rhs); + __m256i ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs, ymm_rhs), + _mm256_min_epi8(ymm_lhs, ymm_rhs)); + ymm_lhs = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d)); + ymm_rhs = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1)); + ymm_sum_0 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs, ymm_lhs), ymm_sum_0); + ymm_sum_1 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs, ymm_rhs), ymm_sum_1); + lhs += 32; + rhs += 32; + } + } else { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + __m256i ymm_lhs_0 = _mm256_loadu_si256((const __m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_loadu_si256((const __m256i *)(lhs + 32)); + __m256i ymm_rhs_0 = _mm256_loadu_si256((const __m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_loadu_si256((const __m256i *)(rhs + 32)); + + __m256i ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_0, ymm_rhs_0), + _mm256_min_epi8(ymm_lhs_0, ymm_rhs_0)); + ymm_lhs_0 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d)); + ymm_rhs_0 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1)); + ymm_sum_0 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs_0, ymm_lhs_0), ymm_sum_0); + ymm_sum_1 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs_0, ymm_rhs_0), ymm_sum_1); + + ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_1, ymm_rhs_1), + _mm256_min_epi8(ymm_lhs_1, ymm_rhs_1)); + ymm_lhs_1 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d)); + ymm_rhs_1 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1)); + ymm_sum_0 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs_1, ymm_lhs_1), ymm_sum_0); + ymm_sum_1 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs_1, ymm_rhs_1), ymm_sum_1); + } + + if (last >= last_aligned + 32) { + __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)lhs); + __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)rhs); + __m256i ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs, ymm_rhs), + _mm256_min_epi8(ymm_lhs, ymm_rhs)); + ymm_lhs = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d)); + ymm_rhs = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1)); + ymm_sum_0 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs, ymm_lhs), ymm_sum_0); + ymm_sum_1 = + _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs, ymm_rhs), ymm_sum_1); + lhs += 32; + rhs += 32; + } + } + result = static_cast( + HorizontalAdd_INT32_V256(_mm256_add_epi32(ymm_sum_0, ymm_sum_1))); + + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + __m128i xmm_sum = _mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs), + _mm_min_epi8(xmm_lhs, xmm_rhs)); + xmm_lhs = _mm_cvtepu8_epi16(xmm_sum); + xmm_rhs = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_sum, xmm_sum)); + xmm_sum = _mm_add_epi32(_mm_madd_epi16(xmm_lhs, xmm_lhs), + _mm_madd_epi16(xmm_rhs, xmm_rhs)); + result += static_cast(HorizontalAdd_INT32_V128(xmm_sum)); + lhs += 16; + rhs += 16; + } + switch (last - lhs) { + case 15: + SSD_INT8_GENERAL(lhs[14], rhs[14], result) + /* FALLTHRU */ + case 14: + SSD_INT8_GENERAL(lhs[13], rhs[13], result) + /* FALLTHRU */ + case 13: + SSD_INT8_GENERAL(lhs[12], rhs[12], result) + /* FALLTHRU */ + case 12: + SSD_INT8_GENERAL(lhs[11], rhs[11], result) + /* FALLTHRU */ + case 11: + SSD_INT8_GENERAL(lhs[10], rhs[10], result) + /* FALLTHRU */ + case 10: + SSD_INT8_GENERAL(lhs[9], rhs[9], result) + /* FALLTHRU */ + case 9: + SSD_INT8_GENERAL(lhs[8], rhs[8], result) + /* FALLTHRU */ + case 8: + SSD_INT8_GENERAL(lhs[7], rhs[7], result) + /* FALLTHRU */ + case 7: + SSD_INT8_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + SSD_INT8_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + SSD_INT8_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + SSD_INT8_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + SSD_INT8_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + SSD_INT8_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + SSD_INT8_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float SquaredEuclideanDistanceInt8AVX2(const int8_t *lhs, const int8_t *rhs, + size_t size) { + if (size > 31) { + return SquaredEuclideanDistanceInt8AVX2Internal(lhs, rhs, size); + } + + return SquaredEuclideanDistanceInt8SSEInternal(lhs, rhs, size); +} +#endif // __AVX2__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_int8_dispatch.cc b/src/ailego/math/euclidean_distance_matrix_int8_dispatch.cc new file mode 100644 index 0000000..d64ca1e --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_int8_dispatch.cc @@ -0,0 +1,65 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX2__) +float SquaredEuclideanDistanceInt8AVX2(const int8_t *lhs, const int8_t *rhs, + size_t size); +#endif + +#if defined(__SSE4_1__) +float SquaredEuclideanDistanceInt8SSE(const int8_t *lhs, const int8_t *rhs, + size_t size); +#endif + +float SquaredEuclideanDistanceInt8Scalar(const int8_t *lhs, const int8_t *rhs, + size_t size); + +//! Compute the distance between matrix and query (INT8, M=1, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = SquaredEuclideanDistanceInt8AVX2(m, q, dim); + return; + } +#endif // __AVX2__ + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = SquaredEuclideanDistanceInt8SSE(m, q, dim); + return; + } +#endif + + *out = SquaredEuclideanDistanceInt8Scalar(m, q, dim); +} + +//! Compute the distance between matrix and query (INT8, M=1, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out); + *out = std::sqrt(*out); +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_int8_sse.cc b/src/ailego/math/euclidean_distance_matrix_int8_sse.cc new file mode 100644 index 0000000..7fd7117 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_int8_sse.cc @@ -0,0 +1,170 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int8.i" +#include "distance_matrix_euclidean_utility.i" +#include "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE4_1__) +inline float SquaredEuclideanDistanceInt8SSEInternal(const int8_t *lhs, + const int8_t *rhs, + size_t size) { + const int8_t *last = lhs + size; + const int8_t *last_aligned = lhs + ((size >> 5) << 5); + + __m128i xmm_sum_0 = _mm_setzero_si128(); + __m128i xmm_sum_1 = _mm_setzero_si128(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m128i xmm_lhs_0 = _mm_load_si128((const __m128i *)(lhs + 0)); + __m128i xmm_lhs_1 = _mm_load_si128((const __m128i *)(lhs + 16)); + __m128i xmm_rhs_0 = _mm_load_si128((const __m128i *)(rhs + 0)); + __m128i xmm_rhs_1 = _mm_load_si128((const __m128i *)(rhs + 16)); + + __m128i xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs_0, xmm_rhs_0), + _mm_min_epi8(xmm_lhs_0, xmm_rhs_0)); + xmm_lhs_0 = _mm_cvtepu8_epi16(xmm_d); + xmm_rhs_0 = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d)); + xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs_1, xmm_rhs_1), + _mm_min_epi8(xmm_lhs_1, xmm_rhs_1)); + xmm_lhs_1 = _mm_cvtepu8_epi16(xmm_d); + xmm_rhs_1 = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d)); + + xmm_sum_0 = + _mm_add_epi32(_mm_madd_epi16(xmm_lhs_0, xmm_lhs_0), xmm_sum_0); + xmm_sum_1 = + _mm_add_epi32(_mm_madd_epi16(xmm_rhs_0, xmm_rhs_0), xmm_sum_1); + xmm_sum_0 = + _mm_add_epi32(_mm_madd_epi16(xmm_lhs_1, xmm_lhs_1), xmm_sum_0); + xmm_sum_1 = + _mm_add_epi32(_mm_madd_epi16(xmm_rhs_1, xmm_rhs_1), xmm_sum_1); + } + + if (last >= last_aligned + 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs); + __m128i xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs), + _mm_min_epi8(xmm_lhs, xmm_rhs)); + xmm_lhs = _mm_cvtepu8_epi16(xmm_d); + xmm_rhs = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d)); + xmm_sum_0 = _mm_add_epi32(_mm_madd_epi16(xmm_lhs, xmm_lhs), xmm_sum_0); + xmm_sum_1 = _mm_add_epi32(_mm_madd_epi16(xmm_rhs, xmm_rhs), xmm_sum_1); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m128i xmm_lhs_0 = _mm_loadu_si128((const __m128i *)(lhs + 0)); + __m128i xmm_lhs_1 = _mm_loadu_si128((const __m128i *)(lhs + 16)); + __m128i xmm_rhs_0 = _mm_loadu_si128((const __m128i *)(rhs + 0)); + __m128i xmm_rhs_1 = _mm_loadu_si128((const __m128i *)(rhs + 16)); + + __m128i xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs_0, xmm_rhs_0), + _mm_min_epi8(xmm_lhs_0, xmm_rhs_0)); + xmm_lhs_0 = _mm_cvtepu8_epi16(xmm_d); + xmm_rhs_0 = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d)); + xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs_1, xmm_rhs_1), + _mm_min_epi8(xmm_lhs_1, xmm_rhs_1)); + xmm_lhs_1 = _mm_cvtepu8_epi16(xmm_d); + xmm_rhs_1 = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d)); + + xmm_sum_0 = + _mm_add_epi32(_mm_madd_epi16(xmm_lhs_0, xmm_lhs_0), xmm_sum_0); + xmm_sum_1 = + _mm_add_epi32(_mm_madd_epi16(xmm_rhs_0, xmm_rhs_0), xmm_sum_1); + xmm_sum_0 = + _mm_add_epi32(_mm_madd_epi16(xmm_lhs_1, xmm_lhs_1), xmm_sum_0); + xmm_sum_1 = + _mm_add_epi32(_mm_madd_epi16(xmm_rhs_1, xmm_rhs_1), xmm_sum_1); + } + + if (last >= last_aligned + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + __m128i xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs), + _mm_min_epi8(xmm_lhs, xmm_rhs)); + xmm_lhs = _mm_cvtepu8_epi16(xmm_d); + xmm_rhs = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d)); + xmm_sum_0 = _mm_add_epi32(_mm_madd_epi16(xmm_lhs, xmm_lhs), xmm_sum_0); + xmm_sum_1 = _mm_add_epi32(_mm_madd_epi16(xmm_rhs, xmm_rhs), xmm_sum_1); + lhs += 16; + rhs += 16; + } + } + float result = static_cast( + HorizontalAdd_INT32_V128(_mm_add_epi32(xmm_sum_0, xmm_sum_1))); + + switch (last - lhs) { + case 15: + SSD_INT8_GENERAL(lhs[14], rhs[14], result) + /* FALLTHRU */ + case 14: + SSD_INT8_GENERAL(lhs[13], rhs[13], result) + /* FALLTHRU */ + case 13: + SSD_INT8_GENERAL(lhs[12], rhs[12], result) + /* FALLTHRU */ + case 12: + SSD_INT8_GENERAL(lhs[11], rhs[11], result) + /* FALLTHRU */ + case 11: + SSD_INT8_GENERAL(lhs[10], rhs[10], result) + /* FALLTHRU */ + case 10: + SSD_INT8_GENERAL(lhs[9], rhs[9], result) + /* FALLTHRU */ + case 9: + SSD_INT8_GENERAL(lhs[8], rhs[8], result) + /* FALLTHRU */ + case 8: + SSD_INT8_GENERAL(lhs[7], rhs[7], result) + /* FALLTHRU */ + case 7: + SSD_INT8_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + SSD_INT8_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + SSD_INT8_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + SSD_INT8_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + SSD_INT8_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + SSD_INT8_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + SSD_INT8_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +//! Squared Euclidean Distance +float SquaredEuclideanDistanceInt8SSE(const int8_t *lhs, const int8_t *rhs, + size_t size) { + return SquaredEuclideanDistanceInt8SSEInternal(lhs, rhs, size); +} + +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_scalar.cc b/src/ailego/math/euclidean_distance_matrix_scalar.cc new file mode 100644 index 0000000..0ab0516 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_scalar.cc @@ -0,0 +1,114 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "distance_utility.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +template +inline float SquaredEuclideanDistanceScalar(const T *m, const T *q, + size_t dim) { + ailego_assert(m && q && dim); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += MathHelper::SquaredDifference(m[i], q[i]); + } + + return sum; +} + +template +inline float EuclideanDistanceScalar(const T *m, const T *q, size_t dim) { + ailego_assert(m && q && dim); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += MathHelper::SquaredDifference(m[i], q[i]); + } + + return std::sqrt(sum); +} + +float SquaredEuclideanDistanceInt4Scalar(const uint8_t *m, const uint8_t *q, + size_t dim) { + ailego_assert(m && q && dim && !(dim & 1)); + + float sum = 0.0; + for (size_t i = 0; i < (dim >> 1); ++i) { + uint8_t m_val = m[i]; + uint8_t q_val = q[i]; + sum += Int4SquaredDiffTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] + + Int4SquaredDiffTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)]; + } + + return sum; +} + + +float EuclideanDistanceInt4Scalar(const uint8_t *m, const uint8_t *q, + size_t dim) { + ailego_assert(m && q && dim && !(dim & 1)); + + float sum = 0.0; + for (size_t i = 0; i < (dim >> 1); ++i) { + uint8_t m_val = m[i]; + uint8_t q_val = q[i]; + sum += Int4SquaredDiffTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] + + Int4SquaredDiffTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)]; + } + + return std::sqrt(sum); +} + + +float SquaredEuclideanDistanceInt8Scalar(const int8_t *m, const int8_t *q, + size_t dim) { + return SquaredEuclideanDistanceScalar(m, q, dim); +} + +float EuclideanDistanceInt8Scalar(const int8_t *m, const int8_t *q, + size_t dim) { + return EuclideanDistanceScalar(m, q, dim); +} + +float SquaredEuclideanDistanceFp16Scalar(const ailego::Float16 *m, + const ailego::Float16 *q, size_t dim) { + return SquaredEuclideanDistanceScalar(m, q, dim); +} + +float EuclideanDistanceFp16Scalar(const ailego::Float16 *m, + const ailego::Float16 *q, size_t dim) { + return EuclideanDistanceScalar(m, q, dim); +} + +float SquaredEuclideanDistanceFp32Scalar(const float *m, const float *q, + size_t dim) { + return SquaredEuclideanDistanceScalar(m, q, dim); +} + +float EuclideanDistanceFp32Scalar(const float *m, const float *q, size_t dim) { + return EuclideanDistanceScalar(m, q, dim); +} + + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/fht.h b/src/ailego/math/fht.h new file mode 100644 index 0000000..8cfc934 --- /dev/null +++ b/src/ailego/math/fht.h @@ -0,0 +1,38 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include + +namespace zvec { +namespace ailego { + +//! Flip the sign of elements based on a packed bit-array. +void fht_flip_sign(const uint8_t *flip, float *data, size_t dim); + +//! Kac random walk: butterfly add/sub between first and second halves. +void fht_kacs_walk(float *data, size_t len); + +//! Inverse Kac walk: undo butterfly add/sub with 0.5 factor. +void fht_inv_kacs_walk(float *data, size_t len); + +//! In-place Fast Hadamard Transform on a power-of-2 length array. +void fht_inplace(float *data, size_t n); + +//! Scale each element by a constant factor. +void fht_vec_rescale(float *data, size_t n, float factor); + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/fht_avx2.cc b/src/ailego/math/fht_avx2.cc new file mode 100644 index 0000000..b2a7f91 --- /dev/null +++ b/src/ailego/math/fht_avx2.cc @@ -0,0 +1,125 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__AVX2__) + +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~31u; + constexpr size_t kChunk = 32; + const __m256i bit_select = + _mm256_setr_epi32(0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80); + const __m256 sign_flip = _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000)); + for (size_t i = 0; i < simd_end; i += kChunk) { + uint32_t mask_bits; + std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits)); + for (int b = 0; b < 4; ++b) { + __m256i mb = _mm256_set1_epi32((mask_bits >> (b * 8)) & 0xFF); + __m256i test = _mm256_and_si256(mb, bit_select); + __m256i cmp = _mm256_cmpeq_epi32(test, bit_select); + __m256 xor_mask = _mm256_and_ps(_mm256_castsi256_ps(cmp), sign_flip); + __m256 v = _mm256_loadu_ps(&data[i + b * 8]); + v = _mm256_xor_ps(v, xor_mask); + _mm256_storeu_ps(&data[i + b * 8], v); + } + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_avx2(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~7u; + for (size_t i = 0; i < half_end; i += 8) { + __m256 x = _mm256_loadu_ps(&data[i]); + __m256 y = _mm256_loadu_ps(&data[i + offset]); + _mm256_storeu_ps(&data[i], _mm256_add_ps(x, y)); + _mm256_storeu_ps(&data[i + offset], _mm256_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_avx2(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~7u; + const __m256 half_fac = _mm256_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 8) { + __m256 a = _mm256_loadu_ps(&data[i]); + __m256 b = _mm256_loadu_ps(&data[i + offset]); + _mm256_storeu_ps(&data[i], _mm256_mul_ps(_mm256_add_ps(a, b), half_fac)); + _mm256_storeu_ps(&data[i + offset], + _mm256_mul_ps(_mm256_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace_avx2(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + size_t step = len << 1; + size_t simd_end = len & ~7u; + for (size_t i = 0; i < n; i += step) { + for (size_t j = 0; j < simd_end; j += 8) { + __m256 u = _mm256_loadu_ps(&data[i + j]); + __m256 v = _mm256_loadu_ps(&data[i + j + len]); + _mm256_storeu_ps(&data[i + j], _mm256_add_ps(u, v)); + _mm256_storeu_ps(&data[i + j + len], _mm256_sub_ps(u, v)); + } + for (size_t j = simd_end; j < len; ++j) { + float u = data[i + j]; + float v = data[i + j + len]; + data[i + j] = u + v; + data[i + j + len] = u - v; + } + } + } +} + +} // namespace ailego +} // namespace zvec + +#endif // __AVX2__ diff --git a/src/ailego/math/fht_avx512.cc b/src/ailego/math/fht_avx512.cc new file mode 100644 index 0000000..ba0fa32 --- /dev/null +++ b/src/ailego/math/fht_avx512.cc @@ -0,0 +1,130 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__AVX512F__) + +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~63u; + constexpr size_t kChunk = 64; + const __m512 sign_flip = _mm512_castsi512_ps(_mm512_set1_epi32(0x80000000)); + for (size_t i = 0; i < simd_end; i += kChunk) { + uint64_t mask_bits; + std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits)); + const __mmask16 m0 = _cvtu32_mask16(mask_bits & 0xFFFF); + const __mmask16 m1 = _cvtu32_mask16((mask_bits >> 16) & 0xFFFF); + const __mmask16 m2 = _cvtu32_mask16((mask_bits >> 32) & 0xFFFF); + const __mmask16 m3 = _cvtu32_mask16((mask_bits >> 48) & 0xFFFF); + __m512 v0 = _mm512_loadu_ps(&data[i]); + v0 = _mm512_mask_xor_ps(v0, m0, v0, sign_flip); + _mm512_storeu_ps(&data[i], v0); + __m512 v1 = _mm512_loadu_ps(&data[i + 16]); + v1 = _mm512_mask_xor_ps(v1, m1, v1, sign_flip); + _mm512_storeu_ps(&data[i + 16], v1); + __m512 v2 = _mm512_loadu_ps(&data[i + 32]); + v2 = _mm512_mask_xor_ps(v2, m2, v2, sign_flip); + _mm512_storeu_ps(&data[i + 32], v2); + __m512 v3 = _mm512_loadu_ps(&data[i + 48]); + v3 = _mm512_mask_xor_ps(v3, m3, v3, sign_flip); + _mm512_storeu_ps(&data[i + 48], v3); + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_avx512(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~15u; + for (size_t i = 0; i < half_end; i += 16) { + __m512 x = _mm512_loadu_ps(&data[i]); + __m512 y = _mm512_loadu_ps(&data[i + offset]); + _mm512_storeu_ps(&data[i], _mm512_add_ps(x, y)); + _mm512_storeu_ps(&data[i + offset], _mm512_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_avx512(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~15u; + const __m512 half_fac = _mm512_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 16) { + __m512 a = _mm512_loadu_ps(&data[i]); + __m512 b = _mm512_loadu_ps(&data[i + offset]); + _mm512_storeu_ps(&data[i], _mm512_mul_ps(_mm512_add_ps(a, b), half_fac)); + _mm512_storeu_ps(&data[i + offset], + _mm512_mul_ps(_mm512_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace_avx512(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + size_t step = len << 1; + size_t simd_end = len & ~15u; + for (size_t i = 0; i < n; i += step) { + for (size_t j = 0; j < simd_end; j += 16) { + __m512 u = _mm512_loadu_ps(&data[i + j]); + __m512 v = _mm512_loadu_ps(&data[i + j + len]); + _mm512_storeu_ps(&data[i + j], _mm512_add_ps(u, v)); + _mm512_storeu_ps(&data[i + j + len], _mm512_sub_ps(u, v)); + } + for (size_t j = simd_end; j < len; ++j) { + float u = data[i + j]; + float v = data[i + j + len]; + data[i + j] = u + v; + data[i + j + len] = u - v; + } + } + } +} + +} // namespace ailego +} // namespace zvec + +#endif // __AVX512F__ diff --git a/src/ailego/math/fht_dispatch.cc b/src/ailego/math/fht_dispatch.cc new file mode 100644 index 0000000..b1e9c4a --- /dev/null +++ b/src/ailego/math/fht_dispatch.cc @@ -0,0 +1,150 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "fht.h" + +namespace zvec { +namespace ailego { + +// ISA-specific forward declarations (implementations in +// fht_scalar/sse/avx2/avx512/neon) +void fht_flip_sign_scalar(const uint8_t *flip, float *data, size_t dim); +void fht_kacs_walk_scalar(float *data, size_t len); +void fht_inv_kacs_walk_scalar(float *data, size_t len); +void fht_inplace_scalar(float *data, size_t n); +#if defined(__SSE2__) +void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim); +void fht_kacs_walk_sse(float *data, size_t len); +void fht_inv_kacs_walk_sse(float *data, size_t len); +#endif +#if defined(__AVX2__) +void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim); +void fht_kacs_walk_avx2(float *data, size_t len); +void fht_inv_kacs_walk_avx2(float *data, size_t len); +void fht_inplace_avx2(float *data, size_t n); +#endif +#if defined(__AVX512F__) +void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim); +void fht_kacs_walk_avx512(float *data, size_t len); +void fht_inv_kacs_walk_avx512(float *data, size_t len); +void fht_inplace_avx512(float *data, size_t n); +#endif +#if defined(__ARM_NEON) && defined(__aarch64__) +void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim); +void fht_kacs_walk_neon(float *data, size_t len); +void fht_inv_kacs_walk_neon(float *data, size_t len); +#endif + +// ============================================================================ +// Runtime dispatch entry points +// ============================================================================ + +void fht_flip_sign(const uint8_t *flip, float *data, size_t dim) { +#if defined(__ARM_NEON) && defined(__aarch64__) + fht_flip_sign_neon(flip, data, dim); +#else +#if defined(__AVX512F__) + if (internal::CpuFeatures::static_flags_.AVX512F && + internal::CpuFeatures::static_flags_.AVX512DQ) { + fht_flip_sign_avx512(flip, data, dim); + return; + } +#endif +#if defined(__AVX2__) + if (internal::CpuFeatures::static_flags_.AVX2) { + fht_flip_sign_avx2(flip, data, dim); + return; + } +#endif +#if defined(__SSE2__) + if (internal::CpuFeatures::static_flags_.SSE2) { + fht_flip_sign_sse(flip, data, dim); + return; + } +#endif + fht_flip_sign_scalar(flip, data, dim); +#endif // __ARM_NEON +} + +void fht_kacs_walk(float *data, size_t len) { +#if defined(__ARM_NEON) && defined(__aarch64__) + fht_kacs_walk_neon(data, len); +#else +#if defined(__AVX512F__) + if (internal::CpuFeatures::static_flags_.AVX512F) { + fht_kacs_walk_avx512(data, len); + return; + } +#endif +#if defined(__AVX2__) + if (internal::CpuFeatures::static_flags_.AVX2) { + fht_kacs_walk_avx2(data, len); + return; + } +#endif +#if defined(__SSE2__) + if (internal::CpuFeatures::static_flags_.SSE2) { + fht_kacs_walk_sse(data, len); + return; + } +#endif + fht_kacs_walk_scalar(data, len); +#endif // __ARM_NEON +} + +void fht_inv_kacs_walk(float *data, size_t len) { +#if defined(__ARM_NEON) && defined(__aarch64__) + fht_inv_kacs_walk_neon(data, len); +#else +#if defined(__AVX512F__) + if (internal::CpuFeatures::static_flags_.AVX512F) { + fht_inv_kacs_walk_avx512(data, len); + return; + } +#endif +#if defined(__AVX2__) + if (internal::CpuFeatures::static_flags_.AVX2) { + fht_inv_kacs_walk_avx2(data, len); + return; + } +#endif +#if defined(__SSE2__) + if (internal::CpuFeatures::static_flags_.SSE2) { + fht_inv_kacs_walk_sse(data, len); + return; + } +#endif + fht_inv_kacs_walk_scalar(data, len); +#endif // __ARM_NEON +} + +void fht_inplace(float *data, size_t n) { +#if defined(__AVX512F__) + if (internal::CpuFeatures::static_flags_.AVX512F) { + fht_inplace_avx512(data, n); + return; + } +#endif +#if defined(__AVX2__) + if (internal::CpuFeatures::static_flags_.AVX2) { + fht_inplace_avx2(data, n); + return; + } +#endif + fht_inplace_scalar(data, n); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/fht_neon.cc b/src/ailego/math/fht_neon.cc new file mode 100644 index 0000000..6e70644 --- /dev/null +++ b/src/ailego/math/fht_neon.cc @@ -0,0 +1,107 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__ARM_NEON) && defined(__aarch64__) + +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim) { + const uint32x4_t sign_bit = vdupq_n_u32(0x80000000u); + size_t simd_end = dim & ~3u; + size_t flip_bytes = (dim + 7) / 8; + for (size_t i = 0; i < simd_end; i += 4) { + uint16_t bits16; + size_t byte_pos = i / 8; + if (byte_pos + 1 < flip_bytes) { + std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16)); + } else { + bits16 = flip[byte_pos]; + } + bits16 >>= (i % 8); + uint32_t b0 = bits16 & 1u; + uint32_t b1 = (bits16 >> 1) & 1u; + uint32_t b2 = (bits16 >> 2) & 1u; + uint32_t b3 = (bits16 >> 3) & 1u; + uint32x4_t bit_mask = {b0, b1, b2, b3}; + uint32x4_t sign_mask = vmulq_u32(bit_mask, sign_bit); + float32x4_t v = vld1q_f32(&data[i]); + v = vreinterpretq_f32_u32(veorq_u32(vreinterpretq_u32_f32(v), sign_mask)); + vst1q_f32(&data[i], v); + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_neon(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~3u; + for (size_t i = 0; i < half_end; i += 4) { + float32x4_t x = vld1q_f32(&data[i]); + float32x4_t y = vld1q_f32(&data[i + offset]); + vst1q_f32(&data[i], vaddq_f32(x, y)); + vst1q_f32(&data[i + offset], vsubq_f32(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_neon(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~3u; + const float32x4_t half_fac = vdupq_n_f32(0.5f); + for (size_t i = 0; i < half_end; i += 4) { + float32x4_t a = vld1q_f32(&data[i]); + float32x4_t b = vld1q_f32(&data[i + offset]); + vst1q_f32(&data[i], vmulq_f32(vaddq_f32(a, b), half_fac)); + vst1q_f32(&data[i + offset], vmulq_f32(vsubq_f32(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +} // namespace ailego +} // namespace zvec + +#endif // __ARM_NEON && __aarch64__ diff --git a/src/ailego/math/fht_scalar.cc b/src/ailego/math/fht_scalar.cc new file mode 100644 index 0000000..643d3c9 --- /dev/null +++ b/src/ailego/math/fht_scalar.cc @@ -0,0 +1,80 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +namespace zvec { +namespace ailego { + +void fht_flip_sign_scalar(const uint8_t *flip, float *data, size_t dim) { + for (size_t i = 0; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_scalar(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + for (size_t i = 0; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_scalar(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + for (size_t i = 0; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace_scalar(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + for (size_t i = 0; i < n; i += len << 1) { + for (size_t j = i; j < i + len; ++j) { + float u = data[j]; + float v = data[j + len]; + data[j] = u + v; + data[j + len] = u - v; + } + } + } +} + +void fht_vec_rescale(float *data, size_t n, float factor) { + for (size_t i = 0; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/fht_sse.cc b/src/ailego/math/fht_sse.cc new file mode 100644 index 0000000..1c03b16 --- /dev/null +++ b/src/ailego/math/fht_sse.cc @@ -0,0 +1,106 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__SSE2__) + +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~3u; + size_t flip_bytes = (dim + 7) / 8; + for (size_t i = 0; i < simd_end; i += 4) { + uint16_t bits16; + size_t byte_pos = i / 8; + if (byte_pos + 1 < flip_bytes) { + std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16)); + } else { + bits16 = flip[byte_pos]; + } + bits16 >>= (i % 8); + uint32_t b0 = bits16 & 1u; + uint32_t b1 = (bits16 >> 1) & 1u; + uint32_t b2 = (bits16 >> 2) & 1u; + uint32_t b3 = (bits16 >> 3) & 1u; + __m128i bit_mask = _mm_set_epi32(b3, b2, b1, b0); + __m128i sign_mask = _mm_slli_epi32(bit_mask, 31); + __m128 v = _mm_loadu_ps(&data[i]); + v = _mm_xor_ps(v, _mm_castsi128_ps(sign_mask)); + _mm_storeu_ps(&data[i], v); + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_sse(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~3u; + for (size_t i = 0; i < half_end; i += 4) { + __m128 x = _mm_loadu_ps(&data[i]); + __m128 y = _mm_loadu_ps(&data[i + offset]); + _mm_storeu_ps(&data[i], _mm_add_ps(x, y)); + _mm_storeu_ps(&data[i + offset], _mm_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_sse(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~3u; + const __m128 half_fac = _mm_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 4) { + __m128 a = _mm_loadu_ps(&data[i]); + __m128 b = _mm_loadu_ps(&data[i + offset]); + _mm_storeu_ps(&data[i], _mm_mul_ps(_mm_add_ps(a, b), half_fac)); + _mm_storeu_ps(&data[i + offset], _mm_mul_ps(_mm_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +} // namespace ailego +} // namespace zvec + +#endif // __SSE2__ diff --git a/src/ailego/math/inner_product_matrix.h b/src/ailego/math/inner_product_matrix.h new file mode 100644 index 0000000..6e799ed --- /dev/null +++ b/src/ailego/math/inner_product_matrix.h @@ -0,0 +1,995 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include "distance_utility.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +/*! Inner Product Matrix + */ +template +struct InnerProductMatrix; + +/*! Inner Product Matrix + */ +template +struct MinusInnerProductMatrix; + +/*! Inner Product Matrix (M=1, N=1) + */ +template +struct InnerProductMatrix< + T, 1, 1, typename std::enable_if::value>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(m[i] * q[i]); + } + *out = sum; + } +}; + +/*! Minus Inner Product Matrix (M=1, N=1) + */ +template +struct MinusInnerProductMatrix< + T, 1, 1, typename std::enable_if::value>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(m[i] * q[i]); + } + *out = -sum; + } +}; + +template <> +struct InnerProductMatrix { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct InnerProductMatrix { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct InnerProductMatrix { + //! Type of value + using ValueType = Float16; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct InnerProductMatrix { + //! Type of value + using ValueType = float; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct MinusInnerProductMatrix { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct MinusInnerProductMatrix { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct MinusInnerProductMatrix { + //! Type of value + using ValueType = Float16; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +template <> +struct MinusInnerProductMatrix { + //! Type of value + using ValueType = float; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Inner Product Matrix + */ +template +struct InnerProductMatrix< + T, M, N, + typename std::enable_if::value && sizeof(T) >= 2 && + M >= 2 && N >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + ValueType m_val = m[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = static_cast(m_val * q[j]); + r += M; + } + } + m += M; + q += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + ValueType m_val = m[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r += m_val * q[j]; + r += M; + } + } + m += M; + q += N; + } + } +}; + +/*! Inner Product Matrix (N=1) + */ +template +struct InnerProductMatrix< + T, M, 1, + typename std::enable_if::value && sizeof(T) >= 2 && + M >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + const ValueType *q_end = q + dim; + if (q != q_end) { + ValueType q_val = *q++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = static_cast(m[i] * q_val); + } + m += M; + } + + while (q != q_end) { + ValueType q_val = *q++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) += m[i] * q_val; + } + m += M; + } + } +}; + +/*! Inner Product Matrix (INT8) + */ +template +struct InnerProductMatrix= 2 && N >= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + + dim >>= 2; + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = FusedMultiplyAdd(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r += FusedMultiplyAdd(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + volatile int32_t sum = ((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) + + (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) + + (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) + + (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24)); + + return static_cast(sum); + } +}; + +/*! Inner Product Matrix (INT8, N=1) + */ +template +struct InnerProductMatrix= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 2); + + if (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = FusedMultiplyAdd(m_it[i], q_val); + } + m_it += M; + } + + while (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) += FusedMultiplyAdd(m_it[i], q_val); + } + m_it += M; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + volatile int32_t sum = ((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) + + (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) + + (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) + + (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24)); + + return static_cast(sum); + } +}; + +/*! Inner Product Matrix (INT4) + */ +template +struct InnerProductMatrix= 2 && N >= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 7) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + + dim >>= 3; + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = FusedMultiplyAdd(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r += FusedMultiplyAdd(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } +}; + +/*! Inner Product Matrix (INT4, N=1) + */ +template +struct InnerProductMatrix= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 7) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 3); + + if (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = FusedMultiplyAdd(m_it[i], q_val); + } + m_it += M; + } + + while (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) += FusedMultiplyAdd(m_it[i], q_val); + } + m_it += M; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } +}; + + +/*! Minus Inner Product Matrix + */ +template +struct MinusInnerProductMatrix< + T, M, N, + typename std::enable_if::value && sizeof(T) >= 2 && + M >= 2 && N >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + ValueType m_val = m[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = -static_cast(m_val * q[j]); + r += M; + } + } + m += M; + q += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + ValueType m_val = m[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r -= m_val * q[j]; + r += M; + } + } + m += M; + q += N; + } + } +}; + +/*! Minus Inner Product Matrix (N=1) + */ +template +struct MinusInnerProductMatrix< + T, M, 1, + typename std::enable_if::value && sizeof(T) >= 2 && + M >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && out); + + const ValueType *q_end = q + dim; + if (q != q_end) { + ValueType q_val = *q++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = -static_cast(m[i] * q_val); + } + m += M; + } + + while (q != q_end) { + ValueType q_val = *q++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) -= m[i] * q_val; + } + m += M; + } + } +}; + +/*! Minus Inner Product Matrix (INT8) + */ +template +struct MinusInnerProductMatrix< + int8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + + dim >>= 2; + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = -FusedMultiplyAdd(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r -= FusedMultiplyAdd(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + volatile int32_t sum = ((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) + + (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) + + (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) + + (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24)); + + return static_cast(sum); + } +}; + +/*! Minus Inner Product Matrix (INT8, N=1) + */ +template +struct MinusInnerProductMatrix= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 2); + + if (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = -FusedMultiplyAdd(m_it[i], q_val); + } + m_it += M; + } + + while (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) -= FusedMultiplyAdd(m_it[i], q_val); + } + m_it += M; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + volatile int32_t sum = ((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) + + (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) + + (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) + + (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24)); + + return static_cast(sum); + } +}; + +/*! Minus Inner Product Matrix (INT4) + */ +template +struct MinusInnerProductMatrix< + uint8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 7) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + + dim >>= 3; + if (dim > 0) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r = -FusedMultiplyAdd(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + uint32_t m_val = m_it[i]; + float *r = out + i; + + for (size_t j = 0; j < N; ++j) { + *r -= FusedMultiplyAdd(m_val, q_it[j]); + r += M; + } + } + m_it += M; + q_it += N; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } +}; + +/*! Minus Inner Product Matrix (INT4, N=1) + */ +template +struct MinusInnerProductMatrix= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + //! Compute the distance between matrix and query + static inline void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out) { + ailego_assert(m && q && dim && !(dim & 7) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 3); + + if (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) = -FusedMultiplyAdd(m_it[i], q_val); + } + m_it += M; + } + + while (q_it != q_end) { + uint32_t q_val = *q_it++; + + for (size_t i = 0; i < M; ++i) { + *(out + i) -= FusedMultiplyAdd(m_it[i], q_val); + } + m_it += M; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } +}; + +//-------------------------------------------------- +// Sparse +//-------------------------------------------------- +struct SparseSegmentInfo { + public: + uint32_t seg_id_{-1U}; + uint32_t vec_cnt_{0}; + + public: + SparseSegmentInfo() : seg_id_{-1U}, vec_cnt_{0} {} + + SparseSegmentInfo(uint32_t seg_id, uint32_t vec_cnt) + : seg_id_{seg_id}, vec_cnt_{vec_cnt} {} +}; + +constexpr static uint32_t SEGMENT_ID_BITS = 16; +constexpr static uint32_t SEGMENT_ID_MASK = 0xFFFF; + +template +struct MinusInnerProductSparseMatrix { + //! Type of value + using ValueType = typename std::remove_cv::type; + + static float ComputeInnerProductSparseInSegment( + uint32_t m_sparse_count, const uint16_t *m_sparse_index, + const ValueType *m_sparse_value, uint32_t q_sparse_count, + const uint16_t *q_sparse_index, const ValueType *q_sparse_value); + + //! Compute the distance between matrix and query + static inline void Compute(const void *m_sparse_data_in, + const void *q_sparse_data_in, float *out); + + static inline void transform_sparse_format(uint32_t sparse_count, + const uint32_t *sparse_index, + const void *sparse_value, + std::string &buffer); +}; + +template <> +struct MinusInnerProductSparseMatrix { + //! Type of value + using ValueType = Float16; + + static float ComputeInnerProductSparseInSegment( + uint32_t m_sparse_count, const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, uint32_t q_sparse_count, + const uint16_t *q_sparse_index, const Float16 *q_sparse_value); + + //! Compute the distance between matrix and query + static void Compute(const void *m_sparse_data_in, + const void *q_sparse_data_in, float *out); + + static void transform_sparse_format(uint32_t sparse_count, + const uint32_t *sparse_index, + const void *sparse_value, + std::string &buffer) { + uint32_t unit_size = sizeof(ValueType); + + uint32_t seg_count = 0; + if (sparse_count == 0) { + buffer.reserve(sizeof(uint32_t) + sizeof(uint32_t)); + + buffer.append(reinterpret_cast(&sparse_count), + sizeof(uint32_t)); + + buffer.append(reinterpret_cast(&seg_count), + sizeof(uint32_t)); + + return; + } + + std::vector seg_infos; + + uint32_t cur_seg_id = -1U; + uint32_t cur_vec_cnt = 0; + + for (size_t i = 0; i < sparse_count; ++i) { + uint32_t seg_id = sparse_index[i] >> SEGMENT_ID_BITS; + if (cur_seg_id == -1U) { + cur_seg_id = seg_id; + cur_vec_cnt++; + } else { + if (seg_id == cur_seg_id) { + cur_vec_cnt++; + } else if (seg_id > cur_seg_id) { + seg_infos.emplace_back(cur_seg_id, cur_vec_cnt); + + cur_seg_id = seg_id; + cur_vec_cnt = 1; + } else { + // std::abort(); + } + } + } + + if (cur_vec_cnt > 0) { + seg_infos.emplace_back(cur_seg_id, cur_vec_cnt); + } + + uint32_t buffer_len = 2 * sizeof(uint32_t) + + seg_infos.size() * 2 * sizeof(uint32_t) + + sparse_count * (sizeof(uint16_t) + sizeof(ValueType)); + + buffer.reserve(buffer_len); + + buffer.append(reinterpret_cast(&sparse_count), + sizeof(uint32_t)); + + seg_count = seg_infos.size(); + buffer.append(reinterpret_cast(&seg_count), sizeof(uint32_t)); + + for (size_t i = 0; i < seg_count; ++i) { + uint32_t seg_id = seg_infos[i].seg_id_; + buffer.append(reinterpret_cast(&seg_id), sizeof(uint32_t)); + } + + for (size_t i = 0; i < seg_count; ++i) { + uint32_t vec_cnt = seg_infos[i].vec_cnt_; + buffer.append(reinterpret_cast(&vec_cnt), sizeof(uint32_t)); + } + + for (size_t i = 0; i < sparse_count; ++i) { + uint16_t temp_dim = sparse_index[i] & SEGMENT_ID_MASK; + buffer.append(reinterpret_cast(&temp_dim), + sizeof(uint16_t)); + } + + const char *sparse_value_ptr = reinterpret_cast(sparse_value); + for (size_t i = 0; i < sparse_count; ++i) { + buffer.append(sparse_value_ptr, unit_size); + sparse_value_ptr += unit_size; + } + } +}; + +template <> +struct MinusInnerProductSparseMatrix { + //! Type of value + using ValueType = float; + + static float ComputeInnerProductSparseInSegment( + uint32_t m_sparse_count, const uint16_t *m_sparse_index, + const float *m_sparse_value, uint32_t q_sparse_count, + const uint16_t *q_sparse_index, const float *q_sparse_value); + + //! Compute the distance between matrix and query + static void Compute(const void *m_sparse_data_in, + const void *q_sparse_data_in, float *out); + + static void transform_sparse_format(uint32_t sparse_count, + const uint32_t *sparse_index, + const void *sparse_value, + std::string &buffer) { + uint32_t unit_size = sizeof(ValueType); + + uint32_t seg_count = 0; + if (sparse_count == 0) { + buffer.reserve(sizeof(uint32_t) + sizeof(uint32_t)); + + buffer.append(reinterpret_cast(&sparse_count), + sizeof(uint32_t)); + + buffer.append(reinterpret_cast(&seg_count), + sizeof(uint32_t)); + + return; + } + + std::vector seg_infos; + + uint32_t cur_seg_id = -1U; + uint32_t cur_vec_cnt = 0; + + for (size_t i = 0; i < sparse_count; ++i) { + uint32_t seg_id = sparse_index[i] >> SEGMENT_ID_BITS; + if (cur_seg_id == -1U) { + cur_seg_id = seg_id; + cur_vec_cnt++; + } else { + if (seg_id == cur_seg_id) { + cur_vec_cnt++; + } else if (seg_id > cur_seg_id) { + seg_infos.emplace_back(cur_seg_id, cur_vec_cnt); + + cur_seg_id = seg_id; + cur_vec_cnt = 1; + } else { + // std::abort(); + } + } + } + + if (cur_vec_cnt > 0) { + seg_infos.emplace_back(cur_seg_id, cur_vec_cnt); + } + + uint32_t buffer_len = 2 * sizeof(uint32_t) + + seg_infos.size() * 2 * sizeof(uint32_t) + + sparse_count * (sizeof(uint16_t) + sizeof(ValueType)); + + buffer.reserve(buffer_len); + + buffer.append(reinterpret_cast(&sparse_count), + sizeof(uint32_t)); + + seg_count = seg_infos.size(); + buffer.append(reinterpret_cast(&seg_count), sizeof(uint32_t)); + + for (size_t i = 0; i < seg_count; ++i) { + uint32_t seg_id = seg_infos[i].seg_id_; + buffer.append(reinterpret_cast(&seg_id), sizeof(uint32_t)); + } + + for (size_t i = 0; i < seg_count; ++i) { + uint32_t vec_cnt = seg_infos[i].vec_cnt_; + buffer.append(reinterpret_cast(&vec_cnt), sizeof(uint32_t)); + } + + for (size_t i = 0; i < sparse_count; ++i) { + uint16_t temp_dim = sparse_index[i] & SEGMENT_ID_MASK; + buffer.append(reinterpret_cast(&temp_dim), + sizeof(uint16_t)); + } + + const char *sparse_value_ptr = reinterpret_cast(sparse_value); + for (size_t i = 0; i < sparse_count; ++i) { + buffer.append(sparse_value_ptr, unit_size); + sparse_value_ptr += unit_size; + } + } +}; + + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_fp16.cc b/src/ailego/math/inner_product_matrix_fp16.cc new file mode 100644 index 0000000..e69de29 diff --git a/src/ailego/math/inner_product_matrix_fp16_avx.cc b/src/ailego/math/inner_product_matrix_fp16_avx.cc new file mode 100644 index 0000000..3415aa6 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp16_avx.cc @@ -0,0 +1,718 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__AVX__) +float InnerProductFp16AVX(const Float16 *lhs, const Float16 *rhs, size_t size) { + float score{0.0f}; + + ACCUM_FP16_1X1_AVX(lhs, rhs, size, &score, 0ull, ) + + return score; +} + +float MinusInnerProductFp16AVX(const Float16 *lhs, const Float16 *rhs, + size_t size) { + float score{0.0f}; + + ACCUM_FP16_1X1_AVX(lhs, rhs, size, &score, 0ull, NEGATE_FP32_GENERAL) + + return score; +} +#endif + +//-------------------------------------------------- +// Sparse +//-------------------------------------------------- +#if defined(__AVX__) +const static __m128i SHUFFLE_MASK256[256] = { + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, -127, -127), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 5, + 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 7, + 6, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 7, + 6, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 7, + 6, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 7, 6, 5, 4, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9, + 8, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9, + 8, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9, + 8, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 5, 4, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9, + 8, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9, + 8, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9, + 8, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 5, + 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 5, + 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 11, 10), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 5, 4, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 7, 6, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11, + 10, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 13, 12), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 5, 4, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 7, 6, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 11, 10), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 11, 10, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 11, 10, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 11, 10, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 11, 10, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13, + 12, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10, + 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, + 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, + 2), + _mm_set_epi8(-127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 15, 14), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 5, 4, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 7, 6, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 11, 10), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 11, 10, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 11, 10, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 11, 10, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 11, 10, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10, + 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, + 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 3, + 2), + _mm_set_epi8(-127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 13, 12), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 13, 12, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 13, 12, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 13, 12, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 13, 12, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 13, 12, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, + 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 3, + 2), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15, + 14, 13, 12, 11, 10), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 11, 10, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 11, 10, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 11, 10, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 5, + 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 5, + 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 11, 10, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, + 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, + 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, + 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, 3, + 2), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, + 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, + 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, + 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, 3, + 2), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, + 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 3, + 2), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, + 4), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2), + _mm_set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), +}; + +constexpr uint32_t MAX_SPARSE_BUFFER_LENGTH = 65536; + +float InnerProductSparseInSegmentFp16AVX(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const Float16 *q_sparse_value) { + float sum = 0.0f; + + // handle if the first dim is zero + bool m_zero = false; + Float16 m_zero_value{0.0f}; + if (m_sparse_count > 0 && m_sparse_index[0] == 0) { + m_sparse_count--; + m_sparse_index++; + m_zero_value = *m_sparse_value++; + m_zero = true; + } + + bool q_zero = false; + Float16 q_zero_value{0.0f}; + if (q_sparse_count > 0 && q_sparse_index[0] == 0) { + q_sparse_count--; + q_sparse_index++; + q_zero_value = *q_sparse_value++; + q_zero = true; + } + + if (m_zero && q_zero) { + sum = m_zero_value * q_zero_value; + } + + size_t i1 = 0, i2 = 0; + size_t end1 = m_sparse_count / 8 * 8; + size_t end2 = q_sparse_count / 8 * 8; + + uint16_t fixed_buffer_1[MAX_SPARSE_BUFFER_LENGTH]; + uint16_t fixed_buffer_2[MAX_SPARSE_BUFFER_LENGTH]; + + Float16 *val_start_1 = reinterpret_cast(fixed_buffer_1); + Float16 *val_start_2 = reinterpret_cast(fixed_buffer_2); + + Float16 *val_1 = val_start_1; + Float16 *val_2 = val_start_2; + + if (i1 < end1 && i2 < end2) { + while (m_sparse_index[i1 + 7] < q_sparse_index[i2]) { + i1 += 8; + if (i1 >= end1) goto do_scalar; + } + + while (q_sparse_index[i2 + 7] < m_sparse_index[i1]) { + i2 += 8; + if (i2 >= end2) goto do_scalar; + } + + __m128i mm_index_m = + _mm_loadu_si128(reinterpret_cast(&m_sparse_index[i1])); + __m128i mm_index_q = + _mm_loadu_si128(reinterpret_cast(&q_sparse_index[i2])); + + while (true) { +#ifdef DEBUG_PRINT + std::cout << "index 1: " << std::endl; + print_data16(&mm_index_m); + + std::cout << "index 2: " << std::endl; + print_data16(&mm_index_q); +#endif + + __m128i mm_cmp_res = + _mm_cmpistrm(mm_index_q, mm_index_m, + _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK); + +#ifdef DEBUG_PRINT + std::cout << "cmp res: " << std::endl; + print_data16(&mm_cmp_res); +#endif + + int r = _mm_extract_epi32(mm_cmp_res, 0); + + if (r) { + int r1 = r; + + __m128i v = _mm_loadu_si128( + reinterpret_cast(&m_sparse_value[i1])); + __m128i vs = _mm_shuffle_epi8(v, SHUFFLE_MASK256[r1]); + + _mm_storeu_si128(reinterpret_cast<__m128i *>(val_1), vs); + val_1 += _mm_popcnt_u32(r1); + + mm_cmp_res = _mm_cmpistrm( + mm_index_m, mm_index_q, + _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK); + r = _mm_extract_epi32(mm_cmp_res, 0); + + r1 = r; + + v = _mm_loadu_si128( + reinterpret_cast(&q_sparse_value[i2])); + vs = _mm_shuffle_epi8(v, SHUFFLE_MASK256[r1]); + + _mm_storeu_si128(reinterpret_cast<__m128i *>(val_2), vs); + val_2 += _mm_popcnt_u32(r1); + } + + const uint16_t id1_max = m_sparse_index[i1 + 7]; + + if (id1_max <= q_sparse_index[i2 + 7]) { + i1 += 8; + if (i1 >= end1) goto do_scalar; + mm_index_m = _mm_loadu_si128( + reinterpret_cast(&m_sparse_index[i1])); + } + + if (id1_max >= q_sparse_index[i2 + 7]) { + i2 += 8; + if (i2 >= end2) goto do_scalar; + mm_index_q = _mm_loadu_si128( + reinterpret_cast(&q_sparse_index[i2])); + } + } + } + +do_scalar: + while (i1 < m_sparse_count && i2 < q_sparse_count) { + if (m_sparse_index[i1] == q_sparse_index[i2]) { + *val_1++ = m_sparse_value[i1]; + *val_2++ = q_sparse_value[i2]; + + ++i1; + ++i2; + } else if (m_sparse_index[i1] < q_sparse_index[i2]) { + ++i1; + } else { + ++i2; + } + } + + size_t res_num = val_1 - val_start_1; + + size_t res_num8 = res_num / 8 * 8; + + if (res_num8) { + __m256 sum256 = _mm256_setzero_ps(); + + for (size_t k = 0; k < res_num8; k += 8) { + __m256 ymm_1 = + _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(val_start_1 + k))); + __m256 ymm_2 = + _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(val_start_2 + k))); + ACCUM_FP32_STEP_AVX(ymm_1, ymm_2, sum256); + } + + sum += HorizontalAdd_FP32_V256(sum256); + } + + for (size_t k = res_num8; k < res_num; ++k) + sum += val_start_1[k] * val_start_2[k]; + + return sum; +} + +#endif // __AVX__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_fp16_avx512.cc b/src/ailego/math/inner_product_matrix_fp16_avx512.cc new file mode 100644 index 0000000..388976c --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp16_avx512.cc @@ -0,0 +1,43 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX512F__) +float InnerProductFp16AVX512(const Float16 *lhs, const Float16 *rhs, + size_t size) { + float score{0.0f}; + + ACCUM_FP16_1X1_AVX512(lhs, rhs, size, &score, 0ull, ) + + return score; +} + +float MinusInnerProductFp16AVX512(const Float16 *lhs, const Float16 *rhs, + size_t size) { + float score{0.0f}; + + ACCUM_FP16_1X1_AVX512(lhs, rhs, size, &score, 0ull, NEGATE_FP32_GENERAL) + + return score; +} +#endif //__AVX512F__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_fp16_avx512fp16.cc b/src/ailego/math/inner_product_matrix_fp16_avx512fp16.cc new file mode 100644 index 0000000..5a10d9a --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp16_avx512fp16.cc @@ -0,0 +1,757 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX512FP16__) +//! Inner Product +float InnerProductFp16AVX512FP16(const Float16 *lhs, const Float16 *rhs, + size_t size) { + const Float16 *last = lhs + size; + const Float16 *last_aligned = lhs + ((size >> 6) << 6); + + __m512h zmm_sum_0 = _mm512_setzero_ph(); + __m512h zmm_sum_1 = _mm512_setzero_ph(); + + if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + FMA_FP16_AVX512FP16(_mm512_load_ph(lhs + 0), _mm512_load_ph(rhs + 0), + zmm_sum_0) + + FMA_FP16_AVX512FP16(_mm512_load_ph(lhs + 32), _mm512_load_ph(rhs + 32), + zmm_sum_1) + } + + if (last >= last_aligned + 32) { + FMA_FP16_AVX512FP16(_mm512_load_ph(lhs), _mm512_load_ph(rhs), zmm_sum_0) + lhs += 32; + rhs += 32; + } + } else { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + FMA_FP16_AVX512FP16(_mm512_loadu_ph(lhs + 0), _mm512_loadu_ph(rhs + 0), + zmm_sum_0) + + FMA_FP16_AVX512FP16(_mm512_loadu_ph(lhs + 32), _mm512_loadu_ph(rhs + 32), + zmm_sum_1) + } + + if (last >= last_aligned + 32) { + FMA_FP16_AVX512FP16(_mm512_loadu_ph(lhs), _mm512_loadu_ph(rhs), zmm_sum_0) + lhs += 32; + rhs += 32; + } + } + + zmm_sum_0 = _mm512_add_ph(zmm_sum_0, zmm_sum_1); + + if (lhs != last) { + __mmask32 mask = (__mmask32)((1 << (last - lhs)) - 1); + __m512i zmm_undefined = _mm512_undefined_epi32(); + zmm_sum_0 = _mm512_mask3_fmadd_ph( + _mm512_castsi512_ph(_mm512_mask_loadu_epi16(zmm_undefined, mask, lhs)), + _mm512_castsi512_ph(_mm512_mask_loadu_epi16(zmm_undefined, mask, rhs)), + zmm_sum_0, mask); + } + + return HorizontalAdd_FP16_V512(zmm_sum_0); +} + +float MinusInnerProductFp16AVX512FP16(const Float16 *lhs, const Float16 *rhs, + size_t size) { + return -1 * InnerProductFp16AVX512FP16(lhs, rhs, size); +} +#endif + +// sparse +#if defined(__AVX512FP16__) +constexpr uint32_t MAX_SPARSE_BUFFER_LENGTH = 65536; + +float InnerProductSparseInSegmentFp16AVX512FP16(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const Float16 *q_sparse_value) { + const static __m128i SHUFFLE_MASK256[256] = { + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, -127, -127), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 7, 6, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 7, 6, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 7, 6, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 7, 6, 5, 4, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 9, 8, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 9, 8, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 9, 8, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 5, 4, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 9, 8, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 9, 8, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 9, 8, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 11, 10), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 5, 4, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 7, 6, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 11, 10, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 13, 12), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 5, 4, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 7, 6, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 13, 12, 11, 10), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 11, 10, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 11, 10, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 11, 10, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 5, 4, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 11, 10, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, + 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, + 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 13, 12, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, + 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, + 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, + 10, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, + 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, + 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, + 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, + 2), + _mm_set_epi8(-127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, 15, 14), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 5, 4, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 7, 6, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 11, 10), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 11, 10, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 11, 10, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 11, 10, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 5, 4, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 11, 10, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, + 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, + 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, + 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, + 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, + 10, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, + 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, + 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, + 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 3, + 2), + _mm_set_epi8(-127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 13, 12), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 13, 12, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 13, 12, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 13, 12, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 5, 4, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 13, 12, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, + 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, + 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 13, 12, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, + 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, + 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, + 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, + 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, + 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 3, + 2), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + 15, 14, 13, 12, 11, 10), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 11, 10, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 11, 10, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 11, 10, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 5, 4, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 5, 4, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 5, 4, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 11, 10, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 7, 6, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 7, 6, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, + 12, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 9, 8, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 9, 8, 3, 2), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 3, 2, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 9, 8, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, + 3, 2), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, + 9, 8, 7, 6), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, + 1, 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, + 3, 2), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, + 5, 4), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 1, 0), + _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2), + _mm_set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + }; + + float sum = 0.0f; + + // handle if the first dim is zero + bool m_zero = false; + Float16 m_zero_value{0.0f}; + if (m_sparse_count > 0 && m_sparse_index[0] == 0) { + m_sparse_count--; + m_sparse_index++; + m_zero_value = *m_sparse_value++; + m_zero = true; + } + + bool q_zero = false; + Float16 q_zero_value{0.0f}; + if (q_sparse_count > 0 && q_sparse_index[0] == 0) { + q_sparse_count--; + q_sparse_index++; + q_zero_value = *q_sparse_value++; + q_zero = true; + } + + if (m_zero && q_zero) { + sum = m_zero_value * q_zero_value; + } + + size_t i1 = 0, i2 = 0; + size_t end1 = m_sparse_count / 8 * 8; + size_t end2 = q_sparse_count / 8 * 8; + + uint16_t fixed_buffer_1[MAX_SPARSE_BUFFER_LENGTH]; + uint16_t fixed_buffer_2[MAX_SPARSE_BUFFER_LENGTH]; + + Float16 *val_start_1 = reinterpret_cast(fixed_buffer_1); + Float16 *val_start_2 = reinterpret_cast(fixed_buffer_2); + + Float16 *val_1 = val_start_1; + Float16 *val_2 = val_start_2; + + if (i1 < end1 && i2 < end2) { + while (m_sparse_index[i1 + 7] < q_sparse_index[i2]) { + i1 += 8; + if (i1 >= end1) goto do_scalar; + } + + while (q_sparse_index[i2 + 7] < m_sparse_index[i1]) { + i2 += 8; + if (i2 >= end2) goto do_scalar; + } + + __m128i mm_index_m = + _mm_loadu_si128(reinterpret_cast(&m_sparse_index[i1])); + __m128i mm_index_q = + _mm_loadu_si128(reinterpret_cast(&q_sparse_index[i2])); + + while (true) { +#ifdef DEBUG_PRINT + std::cout << "index 1: " << std::endl; + print_data16(&mm_index_m); + + std::cout << "index 2: " << std::endl; + print_data16(&mm_index_q); +#endif + + __m128i mm_cmp_res = + _mm_cmpistrm(mm_index_q, mm_index_m, + _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK); + +#ifdef DEBUG_PRINT + std::cout << "cmp res: " << std::endl; + print_data16(&mm_cmp_res); +#endif + + int r = _mm_extract_epi32(mm_cmp_res, 0); + + if (r) { + int r1 = r; + + __m128i v = _mm_loadu_si128( + reinterpret_cast(&m_sparse_value[i1])); + __m128h vs = _mm_castsi128_ph(_mm_shuffle_epi8(v, SHUFFLE_MASK256[r1])); + + _mm_storeu_ph(val_1, vs); + val_1 += _mm_popcnt_u32(r1); + + mm_cmp_res = _mm_cmpistrm( + mm_index_m, mm_index_q, + _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK); + r = _mm_extract_epi32(mm_cmp_res, 0); + + r1 = r; + + v = _mm_loadu_si128( + reinterpret_cast(&q_sparse_value[i2])); + vs = _mm_castsi128_ph(_mm_shuffle_epi8(v, SHUFFLE_MASK256[r1])); + + _mm_storeu_ph(val_2, vs); + val_2 += _mm_popcnt_u32(r1); + } + + const uint16_t id1_max = m_sparse_index[i1 + 7]; + + if (id1_max <= q_sparse_index[i2 + 7]) { + i1 += 8; + if (i1 >= end1) goto do_scalar; + mm_index_m = _mm_loadu_si128( + reinterpret_cast(&m_sparse_index[i1])); + } + + if (id1_max >= q_sparse_index[i2 + 7]) { + i2 += 8; + if (i2 >= end2) goto do_scalar; + mm_index_q = _mm_loadu_si128( + reinterpret_cast(&q_sparse_index[i2])); + } + } + } + +do_scalar: + while (i1 < m_sparse_count && i2 < q_sparse_count) { + if (m_sparse_index[i1] == q_sparse_index[i2]) { + *val_1++ = m_sparse_value[i1]; + *val_2++ = q_sparse_value[i2]; + + ++i1; + ++i2; + } else if (m_sparse_index[i1] < q_sparse_index[i2]) { + ++i1; + } else { + ++i2; + } + } + + size_t res_num = val_1 - val_start_1; + + size_t res_num8 = res_num / 8 * 8; + + if (res_num8) { + __m128h sum128 = _mm_set1_ph(0); + + for (size_t k = 0; k < res_num8; k += 8) { + sum128 = _mm_add_ph(sum128, _mm_mul_ph(_mm_loadu_ph(val_start_1 + k), + _mm_loadu_ph(val_start_2 + k))); + } + + Float16 __attribute__((aligned(16))) tmp_res[8]; + _mm_store_ph(tmp_res, sum128); + sum += (tmp_res[0] + tmp_res[1] + tmp_res[2] + tmp_res[3] + tmp_res[4] + + tmp_res[5] + tmp_res[6] + tmp_res[7]); + } + + for (size_t k = res_num8; k < res_num; ++k) + sum += val_start_1[k] * val_start_2[k]; + + return sum; +} + +#endif // __AVX512FP16__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_fp16_dispatch.cc b/src/ailego/math/inner_product_matrix_fp16_dispatch.cc new file mode 100644 index 0000000..3c46bc3 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp16_dispatch.cc @@ -0,0 +1,179 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__ARM_NEON) +float InnerProductFp16NEON(const Float16 *lhs, const Float16 *rhs, size_t size); +float MinusInnerProductFp16NEON(const Float16 *lhs, const Float16 *rhs, + size_t size); +#endif + +#if defined(__AVX__) +float InnerProductFp16AVX(const Float16 *lhs, const Float16 *rhs, size_t size); +float MinusInnerProductFp16AVX(const Float16 *lhs, const Float16 *rhs, + size_t size); +#endif + +#if defined(__AVX512F__) +float InnerProductFp16AVX512(const Float16 *lhs, const Float16 *rhs, + size_t size); +float MinusInnerProductFp16AVX512(const Float16 *lhs, const Float16 *rhs, + size_t size); +#endif + +#if defined(__AVX512FP16__) +float InnerProductFp16AVX512FP16(const Float16 *lhs, const Float16 *rhs, + size_t size); +float MinusInnerProductFp16AVX512FP16(const Float16 *lhs, const Float16 *rhs, + size_t size); +#endif + +float InnerProductFp16Scalar(const Float16 *lhs, const Float16 *rhs, + size_t size); +float MinusInnerProductFp16Scalar(const Float16 *lhs, const Float16 *rhs, + size_t size); + +//! Compute the distance between matrix and query (FP16, M=1, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + *out = InnerProductFp16NEON(m, q, dim); +#else +#if defined(__AVX512FP16__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) { + *out = InnerProductFp16AVX512FP16(m, q, dim); + return; + } +#endif //__AVX512FP16__ +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = InnerProductFp16AVX512(m, q, dim); + return; + } +#endif //__AVX512F__ +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = InnerProductFp16AVX(m, q, dim); + return; + } +#endif //__AVX__ + *out = InnerProductFp16Scalar(m, q, dim); + +#endif //__ARM_NEON +} + +//! Compute the distance between matrix and query (FP16, M=1, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + *out = MinusInnerProductFp16NEON(m, q, dim); +#else +#if defined(__AVX512FP16__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) { + *out = MinusInnerProductFp16AVX512FP16(m, q, dim); + return; + } +#endif //__AVX512FP16__ +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = MinusInnerProductFp16AVX512(m, q, dim); + return; + } +#endif //__AVX512F__ +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = MinusInnerProductFp16AVX(m, q, dim); + return; + } +#endif //__AVX__ + + *out = MinusInnerProductFp16Scalar(m, q, dim); + +#endif //__ARM_NEON +} + +//-------------------------------------------------- +// Sparse +//-------------------------------------------------- +#if defined(__AVX512FP16__) +float InnerProductSparseInSegmentFp16AVX512FP16(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const Float16 *q_sparse_value); +#endif //__AVX512FP16__ + +#if defined(__AVX__) +float InnerProductSparseInSegmentFp16AVX(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const Float16 *q_sparse_value); +#endif //__AVX__ + +float InnerProductSparseInSegmentFp16Scalar(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const Float16 *q_sparse_value); + +float MinusInnerProductSparseFp16Scalar(const void *m_sparse_data_in, + const void *q_sparse_data_in); + +//! Compute the distance between matrix and query +void MinusInnerProductSparseMatrix::Compute( + const void *m_sparse_data_in, const void *q_sparse_data_in, float *out) { + *out = MinusInnerProductSparseFp16Scalar(m_sparse_data_in, q_sparse_data_in); +} + +float ComputeInnerProductSparseInSegmentFp16(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const Float16 *q_sparse_value) { +#if defined(__AVX512FP16__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) { + return InnerProductSparseInSegmentFp16AVX512FP16( + m_sparse_count, m_sparse_index, m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); + } +#endif +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + return InnerProductSparseInSegmentFp16AVX(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); + } +#endif + return InnerProductSparseInSegmentFp16Scalar(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_fp16_neon.cc b/src/ailego/math/inner_product_matrix_fp16_neon.cc new file mode 100644 index 0000000..3d6c0d6 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp16_neon.cc @@ -0,0 +1,43 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) +float InnerProductFp16NEON(const Float16 *lhs, const Float16 *rhs, + size_t size) { + float score; + + ACCUM_FP16_1X1_NEON(lhs, rhs, size, &score, 0ull, ) + + return score; +} + +float MinusInnerProductFp16NEON(const Float16 *lhs, const Float16 *rhs, + size_t size) { + float score; + + ACCUM_FP16_1X1_NEON(lhs, rhs, size, &score, 0ull, NEGATE_FP32_GENERAL) + + return score; +} +#endif + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_fp32.cc b/src/ailego/math/inner_product_matrix_fp32.cc new file mode 100644 index 0000000..e69de29 diff --git a/src/ailego/math/inner_product_matrix_fp32_avx.cc b/src/ailego/math/inner_product_matrix_fp32_avx.cc new file mode 100644 index 0000000..2d65f46 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp32_avx.cc @@ -0,0 +1,114 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__AVX__) +float InnerProductFp32SSEInternal(const float *lhs, const float *rhs, + size_t size); + +//! Inner Product +float InnerProductFp32AVXInternal(const float *lhs, const float *rhs, + size_t size) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 4) << 4); + + __m256 ymm_sum_0 = _mm256_setzero_ps(); + __m256 ymm_sum_1 = _mm256_setzero_ps(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m256 ymm_lhs_0 = _mm256_load_ps(lhs + 0); + __m256 ymm_lhs_1 = _mm256_load_ps(lhs + 8); + __m256 ymm_rhs_0 = _mm256_load_ps(rhs + 0); + __m256 ymm_rhs_1 = _mm256_load_ps(rhs + 8); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); + } + + if (last >= last_aligned + 8) { + ymm_sum_0 = + _mm256_fmadd_ps(_mm256_load_ps(lhs), _mm256_load_ps(rhs), ymm_sum_0); + lhs += 8; + rhs += 8; + } + } else { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m256 ymm_lhs_0 = _mm256_loadu_ps(lhs + 0); + __m256 ymm_lhs_1 = _mm256_loadu_ps(lhs + 8); + __m256 ymm_rhs_0 = _mm256_loadu_ps(rhs + 0); + __m256 ymm_rhs_1 = _mm256_loadu_ps(rhs + 8); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); + } + + if (last >= last_aligned + 8) { + ymm_sum_0 = _mm256_fmadd_ps(_mm256_loadu_ps(lhs), _mm256_loadu_ps(rhs), + ymm_sum_0); + lhs += 8; + rhs += 8; + } + } + float result = HorizontalAdd_FP32_V256(_mm256_add_ps(ymm_sum_0, ymm_sum_1)); + + switch (last - lhs) { + case 7: + FMA_FP32_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + FMA_FP32_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + FMA_FP32_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + FMA_FP32_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + FMA_FP32_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + FMA_FP32_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + FMA_FP32_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float InnerProductFp32AVX(const float *lhs, const float *rhs, size_t size) { + if (size > 7) { + return InnerProductFp32AVXInternal(lhs, rhs, size); + } + + return InnerProductFp32SSEInternal(lhs, rhs, size); +} + +float MinusInnerProductFp32AVX(const float *lhs, const float *rhs, + size_t size) { + return -1 * InnerProductFp32AVX(lhs, rhs, size); +} + +#endif // __AVX__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_fp32_avx512.cc b/src/ailego/math/inner_product_matrix_fp32_avx512.cc new file mode 100644 index 0000000..8b2b008 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp32_avx512.cc @@ -0,0 +1,102 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__AVX512F__) +float InnerProductFp32AVXInternal(const float *lhs, const float *rhs, + size_t size); + +float InnerProductFp32SSEInternal(const float *lhs, const float *rhs, + size_t size); + +//! Inner Product +float InnerProductFp32AVX512Internal(const float *lhs, const float *rhs, + size_t size) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 5) << 5); + + __m512 zmm_sum_0 = _mm512_setzero_ps(); + __m512 zmm_sum_1 = _mm512_setzero_ps(); + + if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + FMA_FP32_AVX512(_mm512_load_ps(lhs + 0), _mm512_load_ps(rhs + 0), + zmm_sum_0) + + FMA_FP32_AVX512(_mm512_load_ps(lhs + 16), _mm512_load_ps(rhs + 16), + zmm_sum_1) + } + + if (last >= last_aligned + 16) { + FMA_FP32_AVX512(_mm512_load_ps(lhs), _mm512_load_ps(rhs), zmm_sum_0) + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + FMA_FP32_AVX512(_mm512_loadu_ps(lhs + 0), _mm512_loadu_ps(rhs + 0), + zmm_sum_0) + + FMA_FP32_AVX512(_mm512_loadu_ps(lhs + 16), _mm512_loadu_ps(rhs + 16), + zmm_sum_1) + } + + if (last >= last_aligned + 16) { + FMA_FP32_AVX512(_mm512_loadu_ps(lhs), _mm512_loadu_ps(rhs), zmm_sum_0) + lhs += 16; + rhs += 16; + } + } + + zmm_sum_0 = _mm512_add_ps(zmm_sum_0, zmm_sum_1); + if (lhs != last) { + __mmask16 mask = (__mmask16)((1 << (last - lhs)) - 1); + __m512 zmm_undefined = _mm512_undefined_ps(); + zmm_sum_0 = _mm512_mask3_fmadd_ps( + _mm512_mask_loadu_ps(zmm_undefined, mask, lhs), + _mm512_mask_loadu_ps(zmm_undefined, mask, rhs), zmm_sum_0, mask); + } + return HorizontalAdd_FP32_V512(zmm_sum_0); +} + +float InnerProductFp32AVX512(const float *lhs, const float *rhs, size_t size) { + if (size > 15) { + return InnerProductFp32AVX512Internal(lhs, rhs, size); + } + + if (size > 7) { + return InnerProductFp32AVXInternal(lhs, rhs, size); + } + + return InnerProductFp32SSEInternal(lhs, rhs, size); +} + +float MinusInnerProductFp32AVX512(const float *lhs, const float *rhs, + size_t size) { + return -1 * InnerProductFp32AVX512(lhs, rhs, size); +} + +#endif + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_fp32_dispatch.cc b/src/ailego/math/inner_product_matrix_fp32_dispatch.cc new file mode 100644 index 0000000..8b289b6 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp32_dispatch.cc @@ -0,0 +1,155 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__ARM_NEON) +float InnerProductFp32NEON(const float *lhs, const float *rhs, size_t size); +float MinusInnerProductFp32NEON(const float *lhs, const float *rhs, + size_t size); +#endif + +#if defined(__AVX512F__) +float InnerProductFp32AVX512(const float *lhs, const float *rhs, size_t size); +float MinusInnerProductFp32AVX512(const float *lhs, const float *rhs, + size_t size); +#endif + +#if defined(__AVX__) +float InnerProductFp32AVX(const float *lhs, const float *rhs, size_t size); +float MinusInnerProductFp32AVX(const float *lhs, const float *rhs, size_t size); +#endif + +#if defined(__SSE__) +float InnerProductFp32SSE(const float *lhs, const float *rhs, size_t size); +float MinusInnerProductFp32SSE(const float *lhs, const float *rhs, size_t size); +#endif + +float InnerProductFp32Scalar(const float *lhs, const float *rhs, size_t size); +float MinusInnerProductFp32Scalar(const float *lhs, const float *rhs, + size_t size); + +//! Compute the distance between matrix and query (FP32, M=1, N=1) +void InnerProductMatrix::Compute(const float *m, const float *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + *out = InnerProductFp32NEON(m, q, dim); +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = InnerProductFp32AVX512(m, q, dim); + return; + } +#endif // __AVX512F__ + +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = InnerProductFp32AVX(m, q, dim); + return; + } +#endif // __AVX__ + +#if defined(__SSE__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) { + *out = InnerProductFp32SSE(m, q, dim); + return; + } +#endif // __SSE__ + *out = InnerProductFp32Scalar(m, q, dim); +#endif // __ARM_NEON +} + +//! Compute the distance between matrix and query (FP32, M=1, N=1) +void MinusInnerProductMatrix::Compute(const float *m, + const float *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + *out = MinusInnerProductFp32NEON(m, q, dim); +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = MinusInnerProductFp32AVX512(m, q, dim); + return; + } +#endif // __AVX512F__ + +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = MinusInnerProductFp32AVX(m, q, dim); + return; + } +#endif // __AVX__ + +#if defined(__SSE__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) { + *out = MinusInnerProductFp32SSE(m, q, dim); + return; + } +#endif // __SSE__ + *out = MinusInnerProductFp32Scalar(m, q, dim); +#endif // __ARM_NEON +} + +//-------------------------------------------------- +// Sparse +//-------------------------------------------------- +#if defined(__SSE4_1__) +float InnerProductSparseInSegmentFp32SSE(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value); +#endif +float InnerProductSparseInSegmentFp32Scalar(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value); + +float MinusInnerProductSparseFp32Scalar(const void *m_sparse_data_in, + const void *q_sparse_data_in); + +void MinusInnerProductSparseMatrix::Compute(const void *m_sparse_data_in, + const void *q_sparse_data_in, + float *out) { + *out = MinusInnerProductSparseFp32Scalar(m_sparse_data_in, q_sparse_data_in); +} + +float ComputeInnerProductSparseInSegmentFp32(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value) { +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + return InnerProductSparseInSegmentFp32SSE(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); + } +#endif + return InnerProductSparseInSegmentFp32Scalar(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +} +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_fp32_neon.cc b/src/ailego/math/inner_product_matrix_fp32_neon.cc new file mode 100644 index 0000000..c457b3e --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp32_neon.cc @@ -0,0 +1,65 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__ARM_NEON) +float InnerProductFp32NEON(const float *lhs, const float *rhs, size_t size) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 3) << 3); + + float32x4_t v_sum_0 = vdupq_n_f32(0); + float32x4_t v_sum_1 = vdupq_n_f32(0); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + v_sum_0 = vfmaq_f32(v_sum_0, vld1q_f32(lhs + 0), vld1q_f32(rhs + 0)); + v_sum_1 = vfmaq_f32(v_sum_1, vld1q_f32(lhs + 4), vld1q_f32(rhs + 4)); + } + if (last >= last_aligned + 4) { + v_sum_0 = vfmaq_f32(v_sum_0, vld1q_f32(lhs), vld1q_f32(rhs)); + lhs += 4; + rhs += 4; + } + + float result = vaddvq_f32(vaddq_f32(v_sum_0, v_sum_1)); + switch (last - lhs) { + case 3: + FMA_FP32_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + FMA_FP32_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + FMA_FP32_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float MinusInnerProductFp32NEON(const float *lhs, const float *rhs, + size_t size) { + return -1 * InnerProductFp32NEON(lhs, rhs, size); +} + +#endif // __ARM_NEON + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_fp32_sse.cc b/src/ailego/math/inner_product_matrix_fp32_sse.cc new file mode 100644 index 0000000..6bbca76 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp32_sse.cc @@ -0,0 +1,323 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__SSE__) +float InnerProductFp32SSEInternal(const float *lhs, const float *rhs, + size_t size) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 3) << 3); + + __m128 xmm_sum_0 = _mm_setzero_ps(); + __m128 xmm_sum_1 = _mm_setzero_ps(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m128 xmm_lhs_0 = _mm_load_ps(lhs + 0); + __m128 xmm_lhs_1 = _mm_load_ps(lhs + 4); + __m128 xmm_rhs_0 = _mm_load_ps(rhs + 0); + __m128 xmm_rhs_1 = _mm_load_ps(rhs + 4); + xmm_sum_0 = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum_0); + xmm_sum_1 = _mm_fmadd_ps(xmm_lhs_1, xmm_rhs_1, xmm_sum_1); + } + + if (last >= last_aligned + 4) { + xmm_sum_0 = _mm_fmadd_ps(_mm_load_ps(lhs), _mm_load_ps(rhs), xmm_sum_0); + lhs += 4; + rhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m128 xmm_lhs_0 = _mm_loadu_ps(lhs + 0); + __m128 xmm_lhs_1 = _mm_loadu_ps(lhs + 4); + __m128 xmm_rhs_0 = _mm_loadu_ps(rhs + 0); + __m128 xmm_rhs_1 = _mm_loadu_ps(rhs + 4); + xmm_sum_0 = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum_0); + xmm_sum_1 = _mm_fmadd_ps(xmm_lhs_1, xmm_rhs_1, xmm_sum_1); + } + + if (last >= last_aligned + 4) { + xmm_sum_0 = _mm_fmadd_ps(_mm_loadu_ps(lhs), _mm_loadu_ps(rhs), xmm_sum_0); + lhs += 4; + rhs += 4; + } + } + float result = HorizontalAdd_FP32_V128(_mm_add_ps(xmm_sum_0, xmm_sum_1)); + + switch (last - lhs) { + case 3: + FMA_FP32_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + FMA_FP32_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + FMA_FP32_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float InnerProductFp32SSE(const float *lhs, const float *rhs, size_t size) { + return InnerProductFp32SSEInternal(lhs, rhs, size); +} + +float MinusInnerProductFp32SSE(const float *lhs, const float *rhs, + size_t size) { + return -1 * InnerProductFp32SSE(lhs, rhs, size); +} + +#endif // __SSE__ + +//-------------------------------------------------- +// Sparse +//-------------------------------------------------- +#if defined(__SSE4_1__) +const static __m128i SHUFFLE_MASK16[16] = { + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, -127, -127), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 7, 6, 5, 4, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 13, 12), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, + 4), + _mm_set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), +}; + +constexpr uint32_t MAX_SPARSE_BUFFER_LENGTH = 65536; + +float InnerProductSparseInSegmentFp32SSE(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value) { + float sum = 0.0f; + + // handle if the first dim is zero + bool m_zero = false; + float m_zero_value = 0.0f; + if (m_sparse_count > 0 && m_sparse_index[0] == 0) { + m_sparse_count--; + m_sparse_index++; + m_zero_value = *m_sparse_value++; + m_zero = true; + } + + bool q_zero = false; + float q_zero_value = 0.0f; + if (q_sparse_count > 0 && q_sparse_index[0] == 0) { + q_sparse_count--; + q_sparse_index++; + q_zero_value = *q_sparse_value++; + q_zero = true; + } + + if (m_zero && q_zero) { + sum = m_zero_value * q_zero_value; + } + + size_t i1 = 0, i2 = 0; + size_t end1 = m_sparse_count / 8 * 8; + size_t end2 = q_sparse_count / 8 * 8; + + // std::vector mem1; + // std::vector mem2; + + float fixed_buffer_1[MAX_SPARSE_BUFFER_LENGTH]; + float fixed_buffer_2[MAX_SPARSE_BUFFER_LENGTH]; + + float *val_start_1 = fixed_buffer_1; + float *val_start_2 = fixed_buffer_2; + + // uint32_t max_count = std::max(m_sparse_count, q_sparse_count); + + // if (MAX_SPARSE_BUFFER_LENGTH < max_count) { + // mem1.reserve(max_count); + // mem2.reserve(max_count); + + // val_start_1 = mem1.data(); + // val_start_2 = mem2.data(); + // } + + float *val_1 = val_start_1; + float *val_2 = val_start_2; + + if (i1 < end1 && i2 < end2) { + while (m_sparse_index[i1 + 7] < q_sparse_index[i2]) { + i1 += 8; + if (i1 >= end1) goto do_scalar; + } + + while (q_sparse_index[i2 + 7] < m_sparse_index[i1]) { + i2 += 8; + if (i2 >= end2) goto do_scalar; + } + + __m128i mm_index_m = + _mm_loadu_si128(reinterpret_cast(&m_sparse_index[i1])); + __m128i mm_index_q = + _mm_loadu_si128(reinterpret_cast(&q_sparse_index[i2])); + + while (true) { +#ifdef DEBUG_PRINT + std::cout << "index 1: " << std::endl; + print_data16(&mm_index_m); + + std::cout << "index 2: " << std::endl; + print_data16(&mm_index_q); +#endif + + __m128i mm_cmp_res = + _mm_cmpistrm(mm_index_q, mm_index_m, + _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK); + +#ifdef DEBUG_PRINT + std::cout << "cmp res: " << std::endl; + print_data16(&mm_cmp_res); +#endif + + int r = _mm_extract_epi32(mm_cmp_res, 0); + + if (r) { + int r1 = r & 15; + + __m128i v = _mm_loadu_si128( + reinterpret_cast(&m_sparse_value[i1])); + __m128 vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r1])); + + _mm_storeu_ps(val_1, vs); + val_1 += _mm_popcnt_u32(r1); + + int r2 = (r >> 4) & 15; + v = _mm_loadu_si128( + reinterpret_cast(&m_sparse_value[i1 + 4])); + vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r2])); + _mm_storeu_ps(val_1, vs); + val_1 += _mm_popcnt_u32(r2); + + mm_cmp_res = _mm_cmpistrm( + mm_index_m, mm_index_q, + _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK); + r = _mm_extract_epi32(mm_cmp_res, 0); + + r1 = r & 15; + + v = _mm_loadu_si128( + reinterpret_cast(&q_sparse_value[i2])); + vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r1])); + _mm_storeu_ps(val_2, vs); + val_2 += _mm_popcnt_u32(r1); + + r2 = (r >> 4) & 15; + v = _mm_loadu_si128( + reinterpret_cast(&q_sparse_value[i2 + 4])); + vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r2])); + _mm_storeu_ps(val_2, vs); + val_2 += _mm_popcnt_u32(r2); + } + + const uint16_t id1_max = m_sparse_index[i1 + 7]; + + if (id1_max <= q_sparse_index[i2 + 7]) { + i1 += 8; + if (i1 >= end1) goto do_scalar; + mm_index_m = _mm_loadu_si128( + reinterpret_cast(&m_sparse_index[i1])); + } + + if (id1_max >= q_sparse_index[i2 + 7]) { + i2 += 8; + if (i2 >= end2) goto do_scalar; + mm_index_q = _mm_loadu_si128( + reinterpret_cast(&q_sparse_index[i2])); + } + } + } + +do_scalar: + while (i1 < m_sparse_count && i2 < q_sparse_count) { + if (m_sparse_index[i1] == q_sparse_index[i2]) { + *val_1++ = m_sparse_value[i1]; + *val_2++ = q_sparse_value[i2]; + + ++i1; + ++i2; + } else if (m_sparse_index[i1] < q_sparse_index[i2]) { + ++i1; + } else { + ++i2; + } + } + + size_t res_num = val_1 - val_start_1; + + // if (res_num != val_2 - val_start_2) { + // std::cerr << "size mismatch!" << std::endl; + // } + + size_t res_num4 = res_num / 4 * 4; + + if (res_num4) { + __m128 sum128 = _mm_set1_ps(0); + + for (size_t k = 0; k < res_num4; k += 4) { + sum128 = _mm_add_ps(sum128, _mm_mul_ps(_mm_loadu_ps(val_start_1 + k), + _mm_loadu_ps(val_start_2 + k))); + } + + alignas(16) float tmp_res[4]; + _mm_store_ps(tmp_res, sum128); + sum += (tmp_res[0] + tmp_res[1] + tmp_res[2] + tmp_res[3]); + } + + for (size_t k = res_num4; k < res_num; ++k) + sum += val_start_1[k] * val_start_2[k]; + + return sum; +} +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_int4_avx2.cc b/src/ailego/math/inner_product_matrix_int4_avx2.cc new file mode 100644 index 0000000..3fcc9f0 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_int4_avx2.cc @@ -0,0 +1,138 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int4.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__AVX2__) +float InnerProductInt4SSEInternal(const uint8_t *lhs, const uint8_t *rhs, + size_t size); + +//! Inner Product +float InnerProductInt4AVX2Internal(const uint8_t *lhs, const uint8_t *rhs, + size_t size) { + const uint8_t *last = lhs + size; + const uint8_t *last_aligned = lhs + ((size >> 5) << 5); + __m256i ymm_sum = _mm256_setzero_si256(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m256i ymm_lhs = _mm256_load_si256((const __m256i *)(lhs)); + __m256i ymm_rhs = _mm256_load_si256((const __m256i *)(rhs)); + FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum) + } + + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs); + __m128i xmm_sum = _mm_setzero_si128(); + FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) + ymm_sum = _mm256_add_epi32(_mm256_set_m128i(_mm_setzero_si128(), xmm_sum), + ymm_sum); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)(lhs)); + __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)(rhs)); + FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum) + } + + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + __m128i xmm_sum = _mm_setzero_si128(); + FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) + ymm_sum = _mm256_add_epi32(_mm256_set_m128i(_mm_setzero_si128(), xmm_sum), + ymm_sum); + lhs += 16; + rhs += 16; + } + } + float result = static_cast(HorizontalAdd_INT32_V256(ymm_sum)); + + switch (last - lhs) { + case 15: + FMA_INT4_GENERAL(lhs[14], rhs[14], result) + /* FALLTHRU */ + case 14: + FMA_INT4_GENERAL(lhs[13], rhs[13], result) + /* FALLTHRU */ + case 13: + FMA_INT4_GENERAL(lhs[12], rhs[12], result) + /* FALLTHRU */ + case 12: + FMA_INT4_GENERAL(lhs[11], rhs[11], result) + /* FALLTHRU */ + case 11: + FMA_INT4_GENERAL(lhs[10], rhs[10], result) + /* FALLTHRU */ + case 10: + FMA_INT4_GENERAL(lhs[9], rhs[9], result) + /* FALLTHRU */ + case 9: + FMA_INT4_GENERAL(lhs[8], rhs[8], result) + /* FALLTHRU */ + case 8: + FMA_INT4_GENERAL(lhs[7], rhs[7], result) + /* FALLTHRU */ + case 7: + FMA_INT4_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + FMA_INT4_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + FMA_INT4_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + FMA_INT4_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + FMA_INT4_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + FMA_INT4_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + FMA_INT4_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float InnerProductInt4AVX2(const uint8_t *lhs, const uint8_t *rhs, + size_t size) { + if (size > 63) { + return InnerProductInt4AVX2Internal(lhs, rhs, size >> 1); + } + + return InnerProductInt4SSEInternal(lhs, rhs, size >> 1); +} + +float MinusInnerProductInt4AVX2(const uint8_t *lhs, const uint8_t *rhs, + size_t size) { + return -InnerProductInt4AVX2(lhs, rhs, size); +} + +#endif // __AVX2__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_int4_dispatch.cc b/src/ailego/math/inner_product_matrix_int4_dispatch.cc new file mode 100644 index 0000000..83bfd5e --- /dev/null +++ b/src/ailego/math/inner_product_matrix_int4_dispatch.cc @@ -0,0 +1,80 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__AVX2__) +float InnerProductInt4AVX2(const uint8_t *lhs, const uint8_t *rhs, size_t size); +float MinusInnerProductInt4AVX2(const uint8_t *lhs, const uint8_t *rhs, + size_t size); +#endif + +#if defined(__SSE4_1__) +float InnerProductInt4SSE(const uint8_t *lhs, const uint8_t *rhs, size_t size); +float MinusInnerProductInt4SSE(const uint8_t *lhs, const uint8_t *rhs, + size_t size); +#endif + +float InnerProductInt4Scalar(const uint8_t *m, const uint8_t *q, size_t dim); +float MinusInnerProductInt4Scalar(const uint8_t *m, const uint8_t *q, + size_t dim); + +//! Compute the distance between matrix and query (INT4, M=1, N=1) +void InnerProductMatrix::Compute(const uint8_t *m, + const uint8_t *q, size_t dim, + float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = InnerProductInt4AVX2(m, q, dim); + return; + } +#endif // __AVX2__ + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = InnerProductInt4SSE(m, q, dim); + return; + } +#endif //__SSE4_1__ + *out = InnerProductInt4Scalar(m, q, dim); +} + +//! Compute the distance between matrix and query (INT4, M=1, N=1) +void MinusInnerProductMatrix::Compute(const uint8_t *m, + const uint8_t *q, + size_t dim, float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = MinusInnerProductInt4AVX2(m, q, dim); + return; + } +#endif // __AVX2__ + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = MinusInnerProductInt4SSE(m, q, dim); + return; + } +#endif //__SSE4_1__ + *out = MinusInnerProductInt4Scalar(m, q, dim); +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_int4_sse.cc b/src/ailego/math/inner_product_matrix_int4_sse.cc new file mode 100644 index 0000000..39f9d29 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_int4_sse.cc @@ -0,0 +1,107 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int4.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__SSE4_1__) +float InnerProductInt4SSEInternal(const uint8_t *lhs, const uint8_t *rhs, + size_t size) { + const uint8_t *last = lhs + size; + const uint8_t *last_aligned = lhs + ((size >> 4) << 4); + __m128i xmm_sum = _mm_setzero_si128(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)(lhs)); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)(rhs)); + FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) + } + } else { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)(lhs)); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)(rhs)); + FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) + } + } + float result = static_cast(HorizontalAdd_INT32_V128(xmm_sum)); + + switch (last - lhs) { + case 15: + FMA_INT4_GENERAL(lhs[14], rhs[14], result) + /* FALLTHRU */ + case 14: + FMA_INT4_GENERAL(lhs[13], rhs[13], result) + /* FALLTHRU */ + case 13: + FMA_INT4_GENERAL(lhs[12], rhs[12], result) + /* FALLTHRU */ + case 12: + FMA_INT4_GENERAL(lhs[11], rhs[11], result) + /* FALLTHRU */ + case 11: + FMA_INT4_GENERAL(lhs[10], rhs[10], result) + /* FALLTHRU */ + case 10: + FMA_INT4_GENERAL(lhs[9], rhs[9], result) + /* FALLTHRU */ + case 9: + FMA_INT4_GENERAL(lhs[8], rhs[8], result) + /* FALLTHRU */ + case 8: + FMA_INT4_GENERAL(lhs[7], rhs[7], result) + /* FALLTHRU */ + case 7: + FMA_INT4_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + FMA_INT4_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + FMA_INT4_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + FMA_INT4_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + FMA_INT4_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + FMA_INT4_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + FMA_INT4_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float InnerProductInt4SSE(const uint8_t *lhs, const uint8_t *rhs, size_t size) { + return InnerProductInt4SSEInternal(lhs, rhs, size >> 1); +} + +float MinusInnerProductInt4SSE(const uint8_t *lhs, const uint8_t *rhs, + size_t size) { + return -InnerProductInt4SSE(lhs, rhs, size); +} + +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_int8_avx2.cc b/src/ailego/math/inner_product_matrix_int8_avx2.cc new file mode 100644 index 0000000..0b9b6d6 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_int8_avx2.cc @@ -0,0 +1,204 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int8.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__AVX2__) +float InnerProductInt8SSEInternal(const int8_t *lhs, const int8_t *rhs, + size_t size); + +inline float InnerProductInt8AVX2Internal(const int8_t *lhs, const int8_t *rhs, + size_t size) { + const int8_t *last = lhs + size; + const int8_t *last_aligned = lhs + ((size >> 6) << 6); + float result = 0.0; + + __m256i ymm_sum_0 = _mm256_setzero_si256(); + __m256i ymm_sum_1 = _mm256_setzero_si256(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + __m256i ymm_lhs_0 = _mm256_load_si256((const __m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_load_si256((const __m256i *)(lhs + 32)); + __m256i ymm_rhs_0 = _mm256_load_si256((const __m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_load_si256((const __m256i *)(rhs + 32)); + + ymm_lhs_0 = _mm256_sign_epi8(ymm_lhs_0, ymm_rhs_0); + ymm_lhs_1 = _mm256_sign_epi8(ymm_lhs_1, ymm_rhs_1); + ymm_rhs_0 = _mm256_abs_epi8(ymm_rhs_0); + ymm_rhs_1 = _mm256_abs_epi8(ymm_rhs_1); + + ymm_sum_0 = _mm256_add_epi32( + _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_0, ymm_lhs_0), + ONES_INT16_AVX), + ymm_sum_0); + ymm_sum_1 = _mm256_add_epi32( + _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_1, ymm_lhs_1), + ONES_INT16_AVX), + ymm_sum_1); + } + + if (last >= last_aligned + 32) { + __m256i ymm_lhs = _mm256_load_si256((const __m256i *)lhs); + __m256i ymm_rhs = _mm256_load_si256((const __m256i *)rhs); + ymm_lhs = _mm256_sign_epi8(ymm_lhs, ymm_rhs); + ymm_rhs = _mm256_abs_epi8(ymm_rhs); + ymm_sum_0 = _mm256_add_epi32( + _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs, ymm_lhs), + ONES_INT16_AVX), + ymm_sum_0); + lhs += 32; + rhs += 32; + } + + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs); + xmm_lhs = _mm_sign_epi8(xmm_lhs, xmm_rhs); + xmm_rhs = _mm_abs_epi8(xmm_rhs); + ymm_sum_0 = _mm256_add_epi32( + _mm256_set_m128i(_mm_setzero_si128(), + _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs, xmm_lhs), + ONES_INT16_SSE)), + ymm_sum_0); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + __m256i ymm_lhs_0 = _mm256_loadu_si256((const __m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_loadu_si256((const __m256i *)(lhs + 32)); + __m256i ymm_rhs_0 = _mm256_loadu_si256((const __m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_loadu_si256((const __m256i *)(rhs + 32)); + + ymm_lhs_0 = _mm256_sign_epi8(ymm_lhs_0, ymm_rhs_0); + ymm_lhs_1 = _mm256_sign_epi8(ymm_lhs_1, ymm_rhs_1); + ymm_rhs_0 = _mm256_abs_epi8(ymm_rhs_0); + ymm_rhs_1 = _mm256_abs_epi8(ymm_rhs_1); + + ymm_sum_0 = _mm256_add_epi32( + _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_0, ymm_lhs_0), + ONES_INT16_AVX), + ymm_sum_0); + ymm_sum_1 = _mm256_add_epi32( + _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_1, ymm_lhs_1), + ONES_INT16_AVX), + ymm_sum_1); + } + + if (last >= last_aligned + 32) { + __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)lhs); + __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)rhs); + ymm_lhs = _mm256_sign_epi8(ymm_lhs, ymm_rhs); + ymm_rhs = _mm256_abs_epi8(ymm_rhs); + ymm_sum_0 = _mm256_add_epi32( + _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs, ymm_lhs), + ONES_INT16_AVX), + ymm_sum_0); + lhs += 32; + rhs += 32; + } + + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + xmm_lhs = _mm_sign_epi8(xmm_lhs, xmm_rhs); + xmm_rhs = _mm_abs_epi8(xmm_rhs); + ymm_sum_0 = _mm256_add_epi32( + _mm256_set_m128i(_mm_setzero_si128(), + _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs, xmm_lhs), + ONES_INT16_SSE)), + ymm_sum_0); + lhs += 16; + rhs += 16; + } + } + result = static_cast( + HorizontalAdd_INT32_V256(_mm256_add_epi32(ymm_sum_0, ymm_sum_1))); + + switch (last - lhs) { + case 15: + FMA_INT8_GENERAL(lhs[14], rhs[14], result) + /* FALLTHRU */ + case 14: + FMA_INT8_GENERAL(lhs[13], rhs[13], result) + /* FALLTHRU */ + case 13: + FMA_INT8_GENERAL(lhs[12], rhs[12], result) + /* FALLTHRU */ + case 12: + FMA_INT8_GENERAL(lhs[11], rhs[11], result) + /* FALLTHRU */ + case 11: + FMA_INT8_GENERAL(lhs[10], rhs[10], result) + /* FALLTHRU */ + case 10: + FMA_INT8_GENERAL(lhs[9], rhs[9], result) + /* FALLTHRU */ + case 9: + FMA_INT8_GENERAL(lhs[8], rhs[8], result) + /* FALLTHRU */ + case 8: + FMA_INT8_GENERAL(lhs[7], rhs[7], result) + /* FALLTHRU */ + case 7: + FMA_INT8_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + FMA_INT8_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + FMA_INT8_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + FMA_INT8_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + FMA_INT8_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + FMA_INT8_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + FMA_INT8_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float InnerProductInt8AVX2(const int8_t *lhs, const int8_t *rhs, size_t size) { + if (size > 31) { + return InnerProductInt8AVX2Internal(lhs, rhs, size); + } + + return InnerProductInt8SSEInternal(lhs, rhs, size); +} + +float MinusInnerProductInt8AVX2(const int8_t *lhs, const int8_t *rhs, + size_t size) { + return -InnerProductInt8AVX2(lhs, rhs, size); +} + +#endif // __AVX2__ + + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_int8_dispatch.cc b/src/ailego/math/inner_product_matrix_int8_dispatch.cc new file mode 100644 index 0000000..d2faac2 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_int8_dispatch.cc @@ -0,0 +1,82 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__AVX2__) +float InnerProductInt8AVX2(const int8_t *lhs, const int8_t *rhs, size_t size); +float MinusInnerProductInt8AVX2(const int8_t *lhs, const int8_t *rhs, + size_t size); +#endif + +#if defined(__SSE4_1__) +float InnerProductInt8SSE(const int8_t *lhs, const int8_t *rhs, size_t size); +float MinusInnerProductInt8SSE(const int8_t *lhs, const int8_t *rhs, + size_t size); +#endif + +float InnerProductInt8Scalar(const int8_t *m, const int8_t *q, size_t dim); +float MinusInnerProductInt8Scalar(const int8_t *m, const int8_t *q, size_t dim); + +//! Compute the distance between matrix and query (INT8, M=1, N=1) +void InnerProductMatrix::Compute(const int8_t *m, const int8_t *q, + size_t dim, float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = InnerProductInt8AVX2(m, q, dim); + return; + } +#endif // __AVX2__ + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = InnerProductInt8SSE(m, q, dim); + return; + } + +#endif //__SSE4_1__ + + *out = InnerProductInt8Scalar(m, q, dim); +} + +//! Compute the distance between matrix and query (INT8, M=1, N=1) +void MinusInnerProductMatrix::Compute(const int8_t *m, + const int8_t *q, size_t dim, + float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = MinusInnerProductInt8AVX2(m, q, dim); + return; + } +#endif // __AVX2__ + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = MinusInnerProductInt8SSE(m, q, dim); + return; + } +#endif //__SSE4_1__ + + *out = MinusInnerProductInt8Scalar(m, q, dim); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix_int8_sse.cc b/src/ailego/math/inner_product_matrix_int8_sse.cc new file mode 100644 index 0000000..dd84bd5 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_int8_sse.cc @@ -0,0 +1,166 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int8.i" +#include "distance_matrix_inner_product_utility.i" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +#if defined(__SSE4_1__) +//! Inner Product +float InnerProductInt8SSEInternal(const int8_t *lhs, const int8_t *rhs, + size_t size) { + const int8_t *last = lhs + size; + const int8_t *last_aligned = lhs + ((size >> 5) << 5); + + __m128i xmm_sum_0 = _mm_setzero_si128(); + __m128i xmm_sum_1 = _mm_setzero_si128(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m128i xmm_lhs_0 = _mm_load_si128((const __m128i *)(lhs + 0)); + __m128i xmm_lhs_1 = _mm_load_si128((const __m128i *)(lhs + 16)); + __m128i xmm_rhs_0 = _mm_load_si128((const __m128i *)(rhs + 0)); + __m128i xmm_rhs_1 = _mm_load_si128((const __m128i *)(rhs + 16)); + + xmm_lhs_0 = _mm_sign_epi8(xmm_lhs_0, xmm_rhs_0); + xmm_lhs_1 = _mm_sign_epi8(xmm_lhs_1, xmm_rhs_1); + xmm_rhs_0 = _mm_abs_epi8(xmm_rhs_0); + xmm_rhs_1 = _mm_abs_epi8(xmm_rhs_1); + xmm_sum_0 = + _mm_add_epi32(_mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_0, xmm_lhs_0), + ONES_INT16_SSE), + xmm_sum_0); + xmm_sum_1 = + _mm_add_epi32(_mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_1, xmm_lhs_1), + ONES_INT16_SSE), + xmm_sum_1); + } + + if (last >= last_aligned + 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs); + + xmm_lhs = _mm_sign_epi8(xmm_lhs, xmm_rhs); + xmm_rhs = _mm_abs_epi8(xmm_rhs); + xmm_sum_0 = _mm_add_epi32( + _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs, xmm_lhs), ONES_INT16_SSE), + xmm_sum_0); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m128i xmm_lhs_0 = _mm_loadu_si128((const __m128i *)(lhs + 0)); + __m128i xmm_lhs_1 = _mm_loadu_si128((const __m128i *)(lhs + 16)); + __m128i xmm_rhs_0 = _mm_loadu_si128((const __m128i *)(rhs + 0)); + __m128i xmm_rhs_1 = _mm_loadu_si128((const __m128i *)(rhs + 16)); + + xmm_lhs_0 = _mm_sign_epi8(xmm_lhs_0, xmm_rhs_0); + xmm_lhs_1 = _mm_sign_epi8(xmm_lhs_1, xmm_rhs_1); + xmm_rhs_0 = _mm_abs_epi8(xmm_rhs_0); + xmm_rhs_1 = _mm_abs_epi8(xmm_rhs_1); + xmm_sum_0 = + _mm_add_epi32(_mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_0, xmm_lhs_0), + ONES_INT16_SSE), + xmm_sum_0); + xmm_sum_1 = + _mm_add_epi32(_mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_1, xmm_lhs_1), + ONES_INT16_SSE), + xmm_sum_1); + } + + if (last >= last_aligned + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + + xmm_lhs = _mm_sign_epi8(xmm_lhs, xmm_rhs); + xmm_rhs = _mm_abs_epi8(xmm_rhs); + xmm_sum_0 = _mm_add_epi32( + _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs, xmm_lhs), ONES_INT16_SSE), + xmm_sum_0); + lhs += 16; + rhs += 16; + } + } + float result = static_cast( + HorizontalAdd_INT32_V128(_mm_add_epi32(xmm_sum_0, xmm_sum_1))); + + switch (last - lhs) { + case 15: + FMA_INT8_GENERAL(lhs[14], rhs[14], result) + /* FALLTHRU */ + case 14: + FMA_INT8_GENERAL(lhs[13], rhs[13], result) + /* FALLTHRU */ + case 13: + FMA_INT8_GENERAL(lhs[12], rhs[12], result) + /* FALLTHRU */ + case 12: + FMA_INT8_GENERAL(lhs[11], rhs[11], result) + /* FALLTHRU */ + case 11: + FMA_INT8_GENERAL(lhs[10], rhs[10], result) + /* FALLTHRU */ + case 10: + FMA_INT8_GENERAL(lhs[9], rhs[9], result) + /* FALLTHRU */ + case 9: + FMA_INT8_GENERAL(lhs[8], rhs[8], result) + /* FALLTHRU */ + case 8: + FMA_INT8_GENERAL(lhs[7], rhs[7], result) + /* FALLTHRU */ + case 7: + FMA_INT8_GENERAL(lhs[6], rhs[6], result) + /* FALLTHRU */ + case 6: + FMA_INT8_GENERAL(lhs[5], rhs[5], result) + /* FALLTHRU */ + case 5: + FMA_INT8_GENERAL(lhs[4], rhs[4], result) + /* FALLTHRU */ + case 4: + FMA_INT8_GENERAL(lhs[3], rhs[3], result) + /* FALLTHRU */ + case 3: + FMA_INT8_GENERAL(lhs[2], rhs[2], result) + /* FALLTHRU */ + case 2: + FMA_INT8_GENERAL(lhs[1], rhs[1], result) + /* FALLTHRU */ + case 1: + FMA_INT8_GENERAL(lhs[0], rhs[0], result) + } + return result; +} + +float InnerProductInt8SSE(const int8_t *lhs, const int8_t *rhs, size_t size) { + return InnerProductInt8SSEInternal(lhs, rhs, size); +} + +float MinusInnerProductInt8SSE(const int8_t *lhs, const int8_t *rhs, + size_t size) { + return -InnerProductInt8SSEInternal(lhs, rhs, size); +} + +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_scalar.cc b/src/ailego/math/inner_product_matrix_scalar.cc new file mode 100644 index 0000000..4205f6a --- /dev/null +++ b/src/ailego/math/inner_product_matrix_scalar.cc @@ -0,0 +1,299 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include "distance_utility.h" +#include "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +template +inline float InnerProductScalar(const T *m, const T *q, size_t dim) { + ailego_assert(m && q && dim); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(m[i] * q[i]); + } + return sum; +} + +template +inline float MinusInnerProductScalar(const T *m, const T *q, size_t dim) { + ailego_assert(m && q && dim); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(m[i] * q[i]); + } + return -sum; +} + +float InnerProductInt4Scalar(const uint8_t *m, const uint8_t *q, size_t dim) { + ailego_assert(m && q && dim && !(dim & 1)); + + float sum = 0.0; + for (size_t i = 0; i < (dim >> 1); ++i) { + uint8_t m_val = m[i]; + uint8_t q_val = q[i]; + sum += Int4MulTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] + + Int4MulTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)]; + } + + return sum; +} + +float MinusInnerProductInt4Scalar(const uint8_t *m, const uint8_t *q, + size_t dim) { + ailego_assert(m && q && dim && !(dim & 1)); + + float sum = 0.0; + for (size_t i = 0; i < (dim >> 1); ++i) { + uint8_t m_val = m[i]; + uint8_t q_val = q[i]; + sum -= Int4MulTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] + + Int4MulTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)]; + } + return sum; +} + +float InnerProductInt8Scalar(const int8_t *m, const int8_t *q, size_t dim) { + return InnerProductScalar(m, q, dim); +} + +float MinusInnerProductInt8Scalar(const int8_t *m, const int8_t *q, + size_t dim) { + return MinusInnerProductScalar(m, q, dim); +} + +float InnerProductFp16Scalar(const ailego::Float16 *m, const ailego::Float16 *q, + size_t dim) { + return InnerProductScalar(m, q, dim); +} + +float MinusInnerProductFp16Scalar(const ailego::Float16 *m, + const ailego::Float16 *q, size_t dim) { + return MinusInnerProductScalar(m, q, dim); +} + +float InnerProductFp32Scalar(const float *m, const float *q, size_t dim) { + return InnerProductScalar(m, q, dim); +} + +float MinusInnerProductFp32Scalar(const float *m, const float *q, size_t dim) { + return MinusInnerProductScalar(m, q, dim); +} + +//-------------------------------------------------- +// Sparse +//-------------------------------------------------- +float ComputeInnerProductSparseInSegmentFp32(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value); + +float ComputeInnerProductSparseInSegmentFp16(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const Float16 *q_sparse_value); + +template +float ComputeInnerProductSparseInSegment(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const T *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const T *q_sparse_value); + +template <> +float ComputeInnerProductSparseInSegment(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value) { + return ComputeInnerProductSparseInSegmentFp32(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +} + +template <> +float ComputeInnerProductSparseInSegment( + uint32_t m_sparse_count, const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, uint32_t q_sparse_count, + const uint16_t *q_sparse_index, const Float16 *q_sparse_value) { + return ComputeInnerProductSparseInSegmentFp16(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +} + +template +float ComputeSegments(const void *m_sparse_data_in, + const void *q_sparse_data_in) { + ailego_assert(m_sparse_data_in && q_sparse_data_in); + + float sum{0.0f}; + + const uint8_t *m_sparse_data = + reinterpret_cast(m_sparse_data_in); + const uint8_t *q_sparse_data = + reinterpret_cast(q_sparse_data_in); + + const uint32_t m_sparse_count = + *reinterpret_cast(m_sparse_data); + const uint32_t q_sparse_count = + *reinterpret_cast(q_sparse_data); + + if (m_sparse_count == 0 || q_sparse_count == 0) { + return 0.0f; + } + + const uint32_t m_seg_count = + *reinterpret_cast(m_sparse_data + sizeof(uint32_t)); + const uint32_t q_seg_count = + *reinterpret_cast(q_sparse_data + sizeof(uint32_t)); + + const uint32_t *m_seg_id = + reinterpret_cast(m_sparse_data + 2 * sizeof(uint32_t)); + const uint32_t *q_seg_id = + reinterpret_cast(q_sparse_data + 2 * sizeof(uint32_t)); + + const uint32_t *m_seg_vec_cnt = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t) + m_seg_count * sizeof(uint32_t)); + const uint32_t *q_seg_vec_cnt = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t) + q_seg_count * sizeof(uint32_t)); + + const uint16_t *m_sparse_index = + reinterpret_cast(m_sparse_data + 2 * sizeof(uint32_t) + + m_seg_count * 2 * sizeof(uint32_t)); + const uint16_t *q_sparse_index = + reinterpret_cast(q_sparse_data + 2 * sizeof(uint32_t) + + q_seg_count * 2 * sizeof(uint32_t)); + + const T *m_sparse_value = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t) + + m_seg_count * 2 * sizeof(uint32_t) + m_sparse_count * sizeof(uint16_t)); + const T *q_sparse_value = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t) + + q_seg_count * 2 * sizeof(uint32_t) + q_sparse_count * sizeof(uint16_t)); + + size_t m_s = 0; + size_t q_s = 0; + + size_t m_count = 0; + size_t q_count = 0; + + while (m_s < m_seg_count && q_s < q_seg_count) { + if (m_seg_id[m_s] == q_seg_id[q_s]) { + sum += ComputeInnerProductSparseInSegment( + m_seg_vec_cnt[m_s], m_sparse_index + m_count, + m_sparse_value + m_count, q_seg_vec_cnt[q_s], + q_sparse_index + q_count, q_sparse_value + q_count); + + m_count += m_seg_vec_cnt[m_s]; + q_count += q_seg_vec_cnt[q_s]; + + ++m_s; + ++q_s; + } else if (m_seg_id[m_s] < q_seg_id[q_s]) { + m_count += m_seg_vec_cnt[m_s]; + + ++m_s; + } else { + q_count += q_seg_vec_cnt[q_s]; + + ++q_s; + } + } + + return -sum; +} + +float MinusInnerProductSparseFp16Scalar(const void *m_sparse_data_in, + const void *q_sparse_data_in) { + return ComputeSegments(m_sparse_data_in, q_sparse_data_in); +} + +float MinusInnerProductSparseFp32Scalar(const void *m_sparse_data_in, + const void *q_sparse_data_in) { + return ComputeSegments(m_sparse_data_in, q_sparse_data_in); +} + +float InnerProductSparseInSegmentFp16Scalar(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const Float16 *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const Float16 *q_sparse_value) { + float sum = 0.0f; + + size_t m_i = 0; + size_t q_i = 0; + while (m_i < m_sparse_count && q_i < q_sparse_count) { + if (m_sparse_index[m_i] == q_sparse_index[q_i]) { + sum += m_sparse_value[m_i] * q_sparse_value[q_i]; + + ++m_i; + ++q_i; + } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) { + ++m_i; + } else { + ++q_i; + } + } + + return sum; +} + +float InnerProductSparseInSegmentFp32Scalar(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value) { + float sum = 0.0f; + + size_t m_i = 0; + size_t q_i = 0; + while (m_i < m_sparse_count && q_i < q_sparse_count) { + if (m_sparse_index[m_i] == q_sparse_index[q_i]) { + sum += m_sparse_value[m_i] * q_sparse_value[q_i]; + + ++m_i; + ++q_i; + } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) { + ++m_i; + } else { + ++q_i; + } + } + + return sum; +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/matrix_define.i b/src/ailego/math/matrix_define.i new file mode 100644 index 0000000..f796087 --- /dev/null +++ b/src/ailego/math/matrix_define.i @@ -0,0 +1,1850 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#define MATRIX_VAR_INIT_1X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_0 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_1X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_1X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_1X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_1X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_1X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_1X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_7 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_1X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_1X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_15 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_2X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_1X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_1_0 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_2X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_2X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_1 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_2X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_2X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_3 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_2X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_2X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_7 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_2X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_2X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_15 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_2X32(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_2X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_31 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_4X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_2X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_2_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_0 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_4X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_4X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_1 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_4X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_4X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_3 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_4X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_4X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_7 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_4X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_4X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_15 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_4X32(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_4X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_31 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_8X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_4X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_4_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_0 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_8X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_8X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_1 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_8X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_8X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_3 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_8X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_8X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_4 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_5 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_6 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_7 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_7 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_8X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_8X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_8 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_9 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_10 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_11 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_12 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_13 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_14 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_15 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_15 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_8X32(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_8X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_16 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_17 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_18 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_19 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_20 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_21 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_22 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_23 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_24 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_25 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_26 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_27 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_28 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_29 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_30 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_31 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_31 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_16X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_8X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_8_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_9_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_10_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_11_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_12_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_13_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_14_0 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_15_0 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_16X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_16X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_8_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_9_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_10_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_11_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_12_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_13_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_14_1 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_15_1 = (_VAR_INIT); + +#define MATRIX_VAR_INIT_16X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_16X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_8_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_9_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_10_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_11_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_12_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_13_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_14_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_15_2 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_1_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_2_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_3_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_4_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_5_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_6_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_7_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_8_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_9_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_10_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_11_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_12_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_13_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_14_3 = (_VAR_INIT); \ + _VAR_TYPE _VAR_NAME##_15_3 = (_VAR_INIT); + +#define MATRIX_VAR_STORE_1X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + _STORE((_ARRAY) + (_STEP) * (0), _NORM((_VAR##_0_0), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_1X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_1X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (1), _NORM((_VAR##_0_1), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_1X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_1X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (2), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (3), _NORM((_VAR##_0_3), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_1X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_1X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (4), _NORM((_VAR##_0_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (5), _NORM((_VAR##_0_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (6), _NORM((_VAR##_0_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (7), _NORM((_VAR##_0_7), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_1X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_1X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_0_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_0_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_0_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_0_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_0_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_0_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_0_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_0_15), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_2X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_1X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (1), _NORM((_VAR##_1_0), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_2X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_2X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (2), _NORM((_VAR##_0_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (3), _NORM((_VAR##_1_1), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_2X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_2X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (4), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (5), _NORM((_VAR##_1_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (6), _NORM((_VAR##_0_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (7), _NORM((_VAR##_1_3), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_2X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_2X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_0_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_1_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_0_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_1_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_0_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_1_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_0_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_1_7), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_2X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_2X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (16), _NORM((_VAR##_0_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (17), _NORM((_VAR##_1_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (18), _NORM((_VAR##_0_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (19), _NORM((_VAR##_1_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (20), _NORM((_VAR##_0_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (21), _NORM((_VAR##_1_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (22), _NORM((_VAR##_0_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (23), _NORM((_VAR##_1_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (24), _NORM((_VAR##_0_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (25), _NORM((_VAR##_1_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (26), _NORM((_VAR##_0_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (27), _NORM((_VAR##_1_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (28), _NORM((_VAR##_0_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (29), _NORM((_VAR##_1_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (30), _NORM((_VAR##_0_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (31), _NORM((_VAR##_1_15), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_2X32(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_2X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (32), _NORM((_VAR##_0_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (33), _NORM((_VAR##_1_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (34), _NORM((_VAR##_0_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (35), _NORM((_VAR##_1_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (36), _NORM((_VAR##_0_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (37), _NORM((_VAR##_1_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (38), _NORM((_VAR##_0_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (39), _NORM((_VAR##_1_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (40), _NORM((_VAR##_0_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (41), _NORM((_VAR##_1_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (42), _NORM((_VAR##_0_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (43), _NORM((_VAR##_1_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (44), _NORM((_VAR##_0_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (45), _NORM((_VAR##_1_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (46), _NORM((_VAR##_0_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (47), _NORM((_VAR##_1_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (48), _NORM((_VAR##_0_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (49), _NORM((_VAR##_1_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (50), _NORM((_VAR##_0_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (51), _NORM((_VAR##_1_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (52), _NORM((_VAR##_0_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (53), _NORM((_VAR##_1_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (54), _NORM((_VAR##_0_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (55), _NORM((_VAR##_1_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (56), _NORM((_VAR##_0_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (57), _NORM((_VAR##_1_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (58), _NORM((_VAR##_0_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (59), _NORM((_VAR##_1_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (60), _NORM((_VAR##_0_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (61), _NORM((_VAR##_1_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (62), _NORM((_VAR##_0_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (63), _NORM((_VAR##_1_31), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_4X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_2X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (2), _NORM((_VAR##_2_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (3), _NORM((_VAR##_3_0), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_4X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_4X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (4), _NORM((_VAR##_0_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (5), _NORM((_VAR##_1_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (6), _NORM((_VAR##_2_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (7), _NORM((_VAR##_3_1), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_4X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_4X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_1_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_2_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_3_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_0_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_1_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_2_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_3_3), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_4X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_4X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (16), _NORM((_VAR##_0_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (17), _NORM((_VAR##_1_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (18), _NORM((_VAR##_2_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (19), _NORM((_VAR##_3_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (20), _NORM((_VAR##_0_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (21), _NORM((_VAR##_1_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (22), _NORM((_VAR##_2_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (23), _NORM((_VAR##_3_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (24), _NORM((_VAR##_0_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (25), _NORM((_VAR##_1_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (26), _NORM((_VAR##_2_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (27), _NORM((_VAR##_3_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (28), _NORM((_VAR##_0_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (29), _NORM((_VAR##_1_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (30), _NORM((_VAR##_2_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (31), _NORM((_VAR##_3_7), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_4X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_4X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (32), _NORM((_VAR##_0_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (33), _NORM((_VAR##_1_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (34), _NORM((_VAR##_2_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (35), _NORM((_VAR##_3_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (36), _NORM((_VAR##_0_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (37), _NORM((_VAR##_1_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (38), _NORM((_VAR##_2_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (39), _NORM((_VAR##_3_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (40), _NORM((_VAR##_0_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (41), _NORM((_VAR##_1_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (42), _NORM((_VAR##_2_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (43), _NORM((_VAR##_3_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (44), _NORM((_VAR##_0_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (45), _NORM((_VAR##_1_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (46), _NORM((_VAR##_2_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (47), _NORM((_VAR##_3_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (48), _NORM((_VAR##_0_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (49), _NORM((_VAR##_1_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (50), _NORM((_VAR##_2_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (51), _NORM((_VAR##_3_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (52), _NORM((_VAR##_0_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (53), _NORM((_VAR##_1_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (54), _NORM((_VAR##_2_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (55), _NORM((_VAR##_3_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (56), _NORM((_VAR##_0_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (57), _NORM((_VAR##_1_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (58), _NORM((_VAR##_2_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (59), _NORM((_VAR##_3_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (60), _NORM((_VAR##_0_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (61), _NORM((_VAR##_1_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (62), _NORM((_VAR##_2_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (63), _NORM((_VAR##_3_15), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_4X32(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_4X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (64), _NORM((_VAR##_0_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (65), _NORM((_VAR##_1_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (66), _NORM((_VAR##_2_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (67), _NORM((_VAR##_3_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (68), _NORM((_VAR##_0_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (69), _NORM((_VAR##_1_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (70), _NORM((_VAR##_2_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (71), _NORM((_VAR##_3_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (72), _NORM((_VAR##_0_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (73), _NORM((_VAR##_1_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (74), _NORM((_VAR##_2_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (75), _NORM((_VAR##_3_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (76), _NORM((_VAR##_0_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (77), _NORM((_VAR##_1_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (78), _NORM((_VAR##_2_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (79), _NORM((_VAR##_3_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (80), _NORM((_VAR##_0_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (81), _NORM((_VAR##_1_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (82), _NORM((_VAR##_2_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (83), _NORM((_VAR##_3_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (84), _NORM((_VAR##_0_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (85), _NORM((_VAR##_1_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (86), _NORM((_VAR##_2_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (87), _NORM((_VAR##_3_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (88), _NORM((_VAR##_0_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (89), _NORM((_VAR##_1_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (90), _NORM((_VAR##_2_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (91), _NORM((_VAR##_3_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (92), _NORM((_VAR##_0_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (93), _NORM((_VAR##_1_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (94), _NORM((_VAR##_2_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (95), _NORM((_VAR##_3_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (96), _NORM((_VAR##_0_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (97), _NORM((_VAR##_1_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (98), _NORM((_VAR##_2_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (99), _NORM((_VAR##_3_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (100), _NORM((_VAR##_0_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (101), _NORM((_VAR##_1_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (102), _NORM((_VAR##_2_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (103), _NORM((_VAR##_3_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (104), _NORM((_VAR##_0_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (105), _NORM((_VAR##_1_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (106), _NORM((_VAR##_2_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (107), _NORM((_VAR##_3_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (108), _NORM((_VAR##_0_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (109), _NORM((_VAR##_1_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (110), _NORM((_VAR##_2_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (111), _NORM((_VAR##_3_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (112), _NORM((_VAR##_0_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (113), _NORM((_VAR##_1_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (114), _NORM((_VAR##_2_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (115), _NORM((_VAR##_3_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (116), _NORM((_VAR##_0_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (117), _NORM((_VAR##_1_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (118), _NORM((_VAR##_2_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (119), _NORM((_VAR##_3_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (120), _NORM((_VAR##_0_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (121), _NORM((_VAR##_1_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (122), _NORM((_VAR##_2_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (123), _NORM((_VAR##_3_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (124), _NORM((_VAR##_0_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (125), _NORM((_VAR##_1_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (126), _NORM((_VAR##_2_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (127), _NORM((_VAR##_3_31), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_8X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_4X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (4), _NORM((_VAR##_4_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (5), _NORM((_VAR##_5_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (6), _NORM((_VAR##_6_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (7), _NORM((_VAR##_7_0), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_8X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_8X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_0_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_1_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_2_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_3_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_4_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_5_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_6_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_7_1), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_8X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_8X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (16), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (17), _NORM((_VAR##_1_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (18), _NORM((_VAR##_2_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (19), _NORM((_VAR##_3_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (20), _NORM((_VAR##_4_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (21), _NORM((_VAR##_5_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (22), _NORM((_VAR##_6_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (23), _NORM((_VAR##_7_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (24), _NORM((_VAR##_0_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (25), _NORM((_VAR##_1_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (26), _NORM((_VAR##_2_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (27), _NORM((_VAR##_3_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (28), _NORM((_VAR##_4_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (29), _NORM((_VAR##_5_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (30), _NORM((_VAR##_6_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (31), _NORM((_VAR##_7_3), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_8X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_8X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (32), _NORM((_VAR##_0_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (33), _NORM((_VAR##_1_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (34), _NORM((_VAR##_2_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (35), _NORM((_VAR##_3_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (36), _NORM((_VAR##_4_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (37), _NORM((_VAR##_5_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (38), _NORM((_VAR##_6_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (39), _NORM((_VAR##_7_4), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (40), _NORM((_VAR##_0_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (41), _NORM((_VAR##_1_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (42), _NORM((_VAR##_2_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (43), _NORM((_VAR##_3_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (44), _NORM((_VAR##_4_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (45), _NORM((_VAR##_5_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (46), _NORM((_VAR##_6_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (47), _NORM((_VAR##_7_5), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (48), _NORM((_VAR##_0_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (49), _NORM((_VAR##_1_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (50), _NORM((_VAR##_2_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (51), _NORM((_VAR##_3_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (52), _NORM((_VAR##_4_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (53), _NORM((_VAR##_5_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (54), _NORM((_VAR##_6_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (55), _NORM((_VAR##_7_6), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (56), _NORM((_VAR##_0_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (57), _NORM((_VAR##_1_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (58), _NORM((_VAR##_2_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (59), _NORM((_VAR##_3_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (60), _NORM((_VAR##_4_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (61), _NORM((_VAR##_5_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (62), _NORM((_VAR##_6_7), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (63), _NORM((_VAR##_7_7), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_8X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_8X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (64), _NORM((_VAR##_0_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (65), _NORM((_VAR##_1_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (66), _NORM((_VAR##_2_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (67), _NORM((_VAR##_3_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (68), _NORM((_VAR##_4_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (69), _NORM((_VAR##_5_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (70), _NORM((_VAR##_6_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (71), _NORM((_VAR##_7_8), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (72), _NORM((_VAR##_0_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (73), _NORM((_VAR##_1_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (74), _NORM((_VAR##_2_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (75), _NORM((_VAR##_3_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (76), _NORM((_VAR##_4_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (77), _NORM((_VAR##_5_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (78), _NORM((_VAR##_6_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (79), _NORM((_VAR##_7_9), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (80), _NORM((_VAR##_0_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (81), _NORM((_VAR##_1_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (82), _NORM((_VAR##_2_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (83), _NORM((_VAR##_3_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (84), _NORM((_VAR##_4_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (85), _NORM((_VAR##_5_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (86), _NORM((_VAR##_6_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (87), _NORM((_VAR##_7_10), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (88), _NORM((_VAR##_0_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (89), _NORM((_VAR##_1_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (90), _NORM((_VAR##_2_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (91), _NORM((_VAR##_3_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (92), _NORM((_VAR##_4_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (93), _NORM((_VAR##_5_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (94), _NORM((_VAR##_6_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (95), _NORM((_VAR##_7_11), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (96), _NORM((_VAR##_0_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (97), _NORM((_VAR##_1_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (98), _NORM((_VAR##_2_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (99), _NORM((_VAR##_3_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (100), _NORM((_VAR##_4_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (101), _NORM((_VAR##_5_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (102), _NORM((_VAR##_6_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (103), _NORM((_VAR##_7_12), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (104), _NORM((_VAR##_0_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (105), _NORM((_VAR##_1_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (106), _NORM((_VAR##_2_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (107), _NORM((_VAR##_3_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (108), _NORM((_VAR##_4_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (109), _NORM((_VAR##_5_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (110), _NORM((_VAR##_6_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (111), _NORM((_VAR##_7_13), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (112), _NORM((_VAR##_0_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (113), _NORM((_VAR##_1_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (114), _NORM((_VAR##_2_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (115), _NORM((_VAR##_3_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (116), _NORM((_VAR##_4_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (117), _NORM((_VAR##_5_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (118), _NORM((_VAR##_6_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (119), _NORM((_VAR##_7_14), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (120), _NORM((_VAR##_0_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (121), _NORM((_VAR##_1_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (122), _NORM((_VAR##_2_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (123), _NORM((_VAR##_3_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (124), _NORM((_VAR##_4_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (125), _NORM((_VAR##_5_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (126), _NORM((_VAR##_6_15), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (127), _NORM((_VAR##_7_15), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_8X32(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_8X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (128), _NORM((_VAR##_0_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (129), _NORM((_VAR##_1_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (130), _NORM((_VAR##_2_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (131), _NORM((_VAR##_3_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (132), _NORM((_VAR##_4_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (133), _NORM((_VAR##_5_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (134), _NORM((_VAR##_6_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (135), _NORM((_VAR##_7_16), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (136), _NORM((_VAR##_0_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (137), _NORM((_VAR##_1_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (138), _NORM((_VAR##_2_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (139), _NORM((_VAR##_3_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (140), _NORM((_VAR##_4_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (141), _NORM((_VAR##_5_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (142), _NORM((_VAR##_6_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (143), _NORM((_VAR##_7_17), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (144), _NORM((_VAR##_0_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (145), _NORM((_VAR##_1_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (146), _NORM((_VAR##_2_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (147), _NORM((_VAR##_3_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (148), _NORM((_VAR##_4_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (149), _NORM((_VAR##_5_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (150), _NORM((_VAR##_6_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (151), _NORM((_VAR##_7_18), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (152), _NORM((_VAR##_0_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (153), _NORM((_VAR##_1_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (154), _NORM((_VAR##_2_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (155), _NORM((_VAR##_3_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (156), _NORM((_VAR##_4_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (157), _NORM((_VAR##_5_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (158), _NORM((_VAR##_6_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (159), _NORM((_VAR##_7_19), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (160), _NORM((_VAR##_0_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (161), _NORM((_VAR##_1_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (162), _NORM((_VAR##_2_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (163), _NORM((_VAR##_3_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (164), _NORM((_VAR##_4_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (165), _NORM((_VAR##_5_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (166), _NORM((_VAR##_6_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (167), _NORM((_VAR##_7_20), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (168), _NORM((_VAR##_0_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (169), _NORM((_VAR##_1_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (170), _NORM((_VAR##_2_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (171), _NORM((_VAR##_3_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (172), _NORM((_VAR##_4_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (173), _NORM((_VAR##_5_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (174), _NORM((_VAR##_6_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (175), _NORM((_VAR##_7_21), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (176), _NORM((_VAR##_0_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (177), _NORM((_VAR##_1_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (178), _NORM((_VAR##_2_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (179), _NORM((_VAR##_3_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (180), _NORM((_VAR##_4_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (181), _NORM((_VAR##_5_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (182), _NORM((_VAR##_6_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (183), _NORM((_VAR##_7_22), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (184), _NORM((_VAR##_0_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (185), _NORM((_VAR##_1_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (186), _NORM((_VAR##_2_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (187), _NORM((_VAR##_3_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (188), _NORM((_VAR##_4_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (189), _NORM((_VAR##_5_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (190), _NORM((_VAR##_6_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (191), _NORM((_VAR##_7_23), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (192), _NORM((_VAR##_0_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (193), _NORM((_VAR##_1_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (194), _NORM((_VAR##_2_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (195), _NORM((_VAR##_3_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (196), _NORM((_VAR##_4_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (197), _NORM((_VAR##_5_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (198), _NORM((_VAR##_6_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (199), _NORM((_VAR##_7_24), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (200), _NORM((_VAR##_0_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (201), _NORM((_VAR##_1_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (202), _NORM((_VAR##_2_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (203), _NORM((_VAR##_3_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (204), _NORM((_VAR##_4_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (205), _NORM((_VAR##_5_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (206), _NORM((_VAR##_6_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (207), _NORM((_VAR##_7_25), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (208), _NORM((_VAR##_0_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (209), _NORM((_VAR##_1_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (210), _NORM((_VAR##_2_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (211), _NORM((_VAR##_3_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (212), _NORM((_VAR##_4_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (213), _NORM((_VAR##_5_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (214), _NORM((_VAR##_6_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (215), _NORM((_VAR##_7_26), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (216), _NORM((_VAR##_0_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (217), _NORM((_VAR##_1_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (218), _NORM((_VAR##_2_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (219), _NORM((_VAR##_3_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (220), _NORM((_VAR##_4_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (221), _NORM((_VAR##_5_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (222), _NORM((_VAR##_6_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (223), _NORM((_VAR##_7_27), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (224), _NORM((_VAR##_0_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (225), _NORM((_VAR##_1_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (226), _NORM((_VAR##_2_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (227), _NORM((_VAR##_3_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (228), _NORM((_VAR##_4_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (229), _NORM((_VAR##_5_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (230), _NORM((_VAR##_6_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (231), _NORM((_VAR##_7_28), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (232), _NORM((_VAR##_0_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (233), _NORM((_VAR##_1_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (234), _NORM((_VAR##_2_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (235), _NORM((_VAR##_3_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (236), _NORM((_VAR##_4_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (237), _NORM((_VAR##_5_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (238), _NORM((_VAR##_6_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (239), _NORM((_VAR##_7_29), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (240), _NORM((_VAR##_0_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (241), _NORM((_VAR##_1_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (242), _NORM((_VAR##_2_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (243), _NORM((_VAR##_3_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (244), _NORM((_VAR##_4_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (245), _NORM((_VAR##_5_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (246), _NORM((_VAR##_6_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (247), _NORM((_VAR##_7_30), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (248), _NORM((_VAR##_0_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (249), _NORM((_VAR##_1_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (250), _NORM((_VAR##_2_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (251), _NORM((_VAR##_3_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (252), _NORM((_VAR##_4_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (253), _NORM((_VAR##_5_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (254), _NORM((_VAR##_6_31), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (255), _NORM((_VAR##_7_31), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_16X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_8X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_8_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_9_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_10_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_11_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_12_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_13_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_14_0), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_15_0), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_16X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_16X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (16), _NORM((_VAR##_0_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (17), _NORM((_VAR##_1_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (18), _NORM((_VAR##_2_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (19), _NORM((_VAR##_3_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (20), _NORM((_VAR##_4_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (21), _NORM((_VAR##_5_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (22), _NORM((_VAR##_6_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (23), _NORM((_VAR##_7_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (24), _NORM((_VAR##_8_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (25), _NORM((_VAR##_9_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (26), _NORM((_VAR##_10_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (27), _NORM((_VAR##_11_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (28), _NORM((_VAR##_12_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (29), _NORM((_VAR##_13_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (30), _NORM((_VAR##_14_1), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (31), _NORM((_VAR##_15_1), ##__VA_ARGS__)); + +#define MATRIX_VAR_STORE_16X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_16X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \ + _STORE((_ARRAY) + (_STEP) * (32), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (33), _NORM((_VAR##_1_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (34), _NORM((_VAR##_2_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (35), _NORM((_VAR##_3_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (36), _NORM((_VAR##_4_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (37), _NORM((_VAR##_5_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (38), _NORM((_VAR##_6_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (39), _NORM((_VAR##_7_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (40), _NORM((_VAR##_8_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (41), _NORM((_VAR##_9_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (42), _NORM((_VAR##_10_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (43), _NORM((_VAR##_11_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (44), _NORM((_VAR##_12_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (45), _NORM((_VAR##_13_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (46), _NORM((_VAR##_14_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (47), _NORM((_VAR##_15_2), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (48), _NORM((_VAR##_0_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (49), _NORM((_VAR##_1_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (50), _NORM((_VAR##_2_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (51), _NORM((_VAR##_3_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (52), _NORM((_VAR##_4_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (53), _NORM((_VAR##_5_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (54), _NORM((_VAR##_6_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (55), _NORM((_VAR##_7_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (56), _NORM((_VAR##_8_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (57), _NORM((_VAR##_9_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (58), _NORM((_VAR##_10_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (59), _NORM((_VAR##_11_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (60), _NORM((_VAR##_12_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (61), _NORM((_VAR##_13_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (62), _NORM((_VAR##_14_3), ##__VA_ARGS__)); \ + _STORE((_ARRAY) + (_STEP) * (63), _NORM((_VAR##_15_3), ##__VA_ARGS__)); + +#define MATRIX_VAR_PERMUTE_1X1(_VAR, _PERMUTE, ...) \ + (_VAR##_0_0) = _PERMUTE((_VAR##_0_0), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_1X2(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_1X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_1X4(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_1X2(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \ + (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_1X8(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_1X4(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_4) = _PERMUTE((_VAR##_0_4), ##__VA_ARGS__); \ + (_VAR##_0_5) = _PERMUTE((_VAR##_0_5), ##__VA_ARGS__); \ + (_VAR##_0_6) = _PERMUTE((_VAR##_0_6), ##__VA_ARGS__); \ + (_VAR##_0_7) = _PERMUTE((_VAR##_0_7), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_1X16(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_1X8(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_8) = _PERMUTE((_VAR##_0_8), ##__VA_ARGS__); \ + (_VAR##_0_9) = _PERMUTE((_VAR##_0_9), ##__VA_ARGS__); \ + (_VAR##_0_10) = _PERMUTE((_VAR##_0_10), ##__VA_ARGS__); \ + (_VAR##_0_11) = _PERMUTE((_VAR##_0_11), ##__VA_ARGS__); \ + (_VAR##_0_12) = _PERMUTE((_VAR##_0_12), ##__VA_ARGS__); \ + (_VAR##_0_13) = _PERMUTE((_VAR##_0_13), ##__VA_ARGS__); \ + (_VAR##_0_14) = _PERMUTE((_VAR##_0_14), ##__VA_ARGS__); \ + (_VAR##_0_15) = _PERMUTE((_VAR##_0_15), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_2X1(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_1X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_1_0) = _PERMUTE((_VAR##_1_0), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_2X2(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_2X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); \ + (_VAR##_1_1) = _PERMUTE((_VAR##_1_1), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_2X4(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_2X2(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \ + (_VAR##_1_2) = _PERMUTE((_VAR##_1_2), ##__VA_ARGS__); \ + (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); \ + (_VAR##_1_3) = _PERMUTE((_VAR##_1_3), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_2X8(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_2X4(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_4) = _PERMUTE((_VAR##_0_4), ##__VA_ARGS__); \ + (_VAR##_1_4) = _PERMUTE((_VAR##_1_4), ##__VA_ARGS__); \ + (_VAR##_0_5) = _PERMUTE((_VAR##_0_5), ##__VA_ARGS__); \ + (_VAR##_1_5) = _PERMUTE((_VAR##_1_5), ##__VA_ARGS__); \ + (_VAR##_0_6) = _PERMUTE((_VAR##_0_6), ##__VA_ARGS__); \ + (_VAR##_1_6) = _PERMUTE((_VAR##_1_6), ##__VA_ARGS__); \ + (_VAR##_0_7) = _PERMUTE((_VAR##_0_7), ##__VA_ARGS__); \ + (_VAR##_1_7) = _PERMUTE((_VAR##_1_7), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_2X16(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_2X8(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_8) = _PERMUTE((_VAR##_0_8), ##__VA_ARGS__); \ + (_VAR##_1_8) = _PERMUTE((_VAR##_1_8), ##__VA_ARGS__); \ + (_VAR##_0_9) = _PERMUTE((_VAR##_0_9), ##__VA_ARGS__); \ + (_VAR##_1_9) = _PERMUTE((_VAR##_1_9), ##__VA_ARGS__); \ + (_VAR##_0_10) = _PERMUTE((_VAR##_0_10), ##__VA_ARGS__); \ + (_VAR##_1_10) = _PERMUTE((_VAR##_1_10), ##__VA_ARGS__); \ + (_VAR##_0_11) = _PERMUTE((_VAR##_0_11), ##__VA_ARGS__); \ + (_VAR##_1_11) = _PERMUTE((_VAR##_1_11), ##__VA_ARGS__); \ + (_VAR##_0_12) = _PERMUTE((_VAR##_0_12), ##__VA_ARGS__); \ + (_VAR##_1_12) = _PERMUTE((_VAR##_1_12), ##__VA_ARGS__); \ + (_VAR##_0_13) = _PERMUTE((_VAR##_0_13), ##__VA_ARGS__); \ + (_VAR##_1_13) = _PERMUTE((_VAR##_1_13), ##__VA_ARGS__); \ + (_VAR##_0_14) = _PERMUTE((_VAR##_0_14), ##__VA_ARGS__); \ + (_VAR##_1_14) = _PERMUTE((_VAR##_1_14), ##__VA_ARGS__); \ + (_VAR##_0_15) = _PERMUTE((_VAR##_0_15), ##__VA_ARGS__); \ + (_VAR##_1_15) = _PERMUTE((_VAR##_1_15), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_2X32(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_2X16(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_16) = _PERMUTE((_VAR##_0_16), ##__VA_ARGS__); \ + (_VAR##_1_16) = _PERMUTE((_VAR##_1_16), ##__VA_ARGS__); \ + (_VAR##_0_17) = _PERMUTE((_VAR##_0_17), ##__VA_ARGS__); \ + (_VAR##_1_17) = _PERMUTE((_VAR##_1_17), ##__VA_ARGS__); \ + (_VAR##_0_18) = _PERMUTE((_VAR##_0_18), ##__VA_ARGS__); \ + (_VAR##_1_18) = _PERMUTE((_VAR##_1_18), ##__VA_ARGS__); \ + (_VAR##_0_19) = _PERMUTE((_VAR##_0_19), ##__VA_ARGS__); \ + (_VAR##_1_19) = _PERMUTE((_VAR##_1_19), ##__VA_ARGS__); \ + (_VAR##_0_20) = _PERMUTE((_VAR##_0_20), ##__VA_ARGS__); \ + (_VAR##_1_20) = _PERMUTE((_VAR##_1_20), ##__VA_ARGS__); \ + (_VAR##_0_21) = _PERMUTE((_VAR##_0_21), ##__VA_ARGS__); \ + (_VAR##_1_21) = _PERMUTE((_VAR##_1_21), ##__VA_ARGS__); \ + (_VAR##_0_22) = _PERMUTE((_VAR##_0_22), ##__VA_ARGS__); \ + (_VAR##_1_22) = _PERMUTE((_VAR##_1_22), ##__VA_ARGS__); \ + (_VAR##_0_23) = _PERMUTE((_VAR##_0_23), ##__VA_ARGS__); \ + (_VAR##_1_23) = _PERMUTE((_VAR##_1_23), ##__VA_ARGS__); \ + (_VAR##_0_24) = _PERMUTE((_VAR##_0_24), ##__VA_ARGS__); \ + (_VAR##_1_24) = _PERMUTE((_VAR##_1_24), ##__VA_ARGS__); \ + (_VAR##_0_25) = _PERMUTE((_VAR##_0_25), ##__VA_ARGS__); \ + (_VAR##_1_25) = _PERMUTE((_VAR##_1_25), ##__VA_ARGS__); \ + (_VAR##_0_26) = _PERMUTE((_VAR##_0_26), ##__VA_ARGS__); \ + (_VAR##_1_26) = _PERMUTE((_VAR##_1_26), ##__VA_ARGS__); \ + (_VAR##_0_27) = _PERMUTE((_VAR##_0_27), ##__VA_ARGS__); \ + (_VAR##_1_27) = _PERMUTE((_VAR##_1_27), ##__VA_ARGS__); \ + (_VAR##_0_28) = _PERMUTE((_VAR##_0_28), ##__VA_ARGS__); \ + (_VAR##_1_28) = _PERMUTE((_VAR##_1_28), ##__VA_ARGS__); \ + (_VAR##_0_29) = _PERMUTE((_VAR##_0_29), ##__VA_ARGS__); \ + (_VAR##_1_29) = _PERMUTE((_VAR##_1_29), ##__VA_ARGS__); \ + (_VAR##_0_30) = _PERMUTE((_VAR##_0_30), ##__VA_ARGS__); \ + (_VAR##_1_30) = _PERMUTE((_VAR##_1_30), ##__VA_ARGS__); \ + (_VAR##_0_31) = _PERMUTE((_VAR##_0_31), ##__VA_ARGS__); \ + (_VAR##_1_31) = _PERMUTE((_VAR##_1_31), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_4X1(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_2X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_2_0) = _PERMUTE((_VAR##_2_0), ##__VA_ARGS__); \ + (_VAR##_3_0) = _PERMUTE((_VAR##_3_0), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_4X2(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_4X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); \ + (_VAR##_1_1) = _PERMUTE((_VAR##_1_1), ##__VA_ARGS__); \ + (_VAR##_2_1) = _PERMUTE((_VAR##_2_1), ##__VA_ARGS__); \ + (_VAR##_3_1) = _PERMUTE((_VAR##_3_1), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_4X4(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_4X2(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \ + (_VAR##_1_2) = _PERMUTE((_VAR##_1_2), ##__VA_ARGS__); \ + (_VAR##_2_2) = _PERMUTE((_VAR##_2_2), ##__VA_ARGS__); \ + (_VAR##_3_2) = _PERMUTE((_VAR##_3_2), ##__VA_ARGS__); \ + (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); \ + (_VAR##_1_3) = _PERMUTE((_VAR##_1_3), ##__VA_ARGS__); \ + (_VAR##_2_3) = _PERMUTE((_VAR##_2_3), ##__VA_ARGS__); \ + (_VAR##_3_3) = _PERMUTE((_VAR##_3_3), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_4X8(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_4X4(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_4) = _PERMUTE((_VAR##_0_4), ##__VA_ARGS__); \ + (_VAR##_1_4) = _PERMUTE((_VAR##_1_4), ##__VA_ARGS__); \ + (_VAR##_2_4) = _PERMUTE((_VAR##_2_4), ##__VA_ARGS__); \ + (_VAR##_3_4) = _PERMUTE((_VAR##_3_4), ##__VA_ARGS__); \ + (_VAR##_0_5) = _PERMUTE((_VAR##_0_5), ##__VA_ARGS__); \ + (_VAR##_1_5) = _PERMUTE((_VAR##_1_5), ##__VA_ARGS__); \ + (_VAR##_2_5) = _PERMUTE((_VAR##_2_5), ##__VA_ARGS__); \ + (_VAR##_3_5) = _PERMUTE((_VAR##_3_5), ##__VA_ARGS__); \ + (_VAR##_0_6) = _PERMUTE((_VAR##_0_6), ##__VA_ARGS__); \ + (_VAR##_1_6) = _PERMUTE((_VAR##_1_6), ##__VA_ARGS__); \ + (_VAR##_2_6) = _PERMUTE((_VAR##_2_6), ##__VA_ARGS__); \ + (_VAR##_3_6) = _PERMUTE((_VAR##_3_6), ##__VA_ARGS__); \ + (_VAR##_0_7) = _PERMUTE((_VAR##_0_7), ##__VA_ARGS__); \ + (_VAR##_1_7) = _PERMUTE((_VAR##_1_7), ##__VA_ARGS__); \ + (_VAR##_2_7) = _PERMUTE((_VAR##_2_7), ##__VA_ARGS__); \ + (_VAR##_3_7) = _PERMUTE((_VAR##_3_7), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_4X16(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_4X8(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_8) = _PERMUTE((_VAR##_0_8), ##__VA_ARGS__); \ + (_VAR##_1_8) = _PERMUTE((_VAR##_1_8), ##__VA_ARGS__); \ + (_VAR##_2_8) = _PERMUTE((_VAR##_2_8), ##__VA_ARGS__); \ + (_VAR##_3_8) = _PERMUTE((_VAR##_3_8), ##__VA_ARGS__); \ + (_VAR##_0_9) = _PERMUTE((_VAR##_0_9), ##__VA_ARGS__); \ + (_VAR##_1_9) = _PERMUTE((_VAR##_1_9), ##__VA_ARGS__); \ + (_VAR##_2_9) = _PERMUTE((_VAR##_2_9), ##__VA_ARGS__); \ + (_VAR##_3_9) = _PERMUTE((_VAR##_3_9), ##__VA_ARGS__); \ + (_VAR##_0_10) = _PERMUTE((_VAR##_0_10), ##__VA_ARGS__); \ + (_VAR##_1_10) = _PERMUTE((_VAR##_1_10), ##__VA_ARGS__); \ + (_VAR##_2_10) = _PERMUTE((_VAR##_2_10), ##__VA_ARGS__); \ + (_VAR##_3_10) = _PERMUTE((_VAR##_3_10), ##__VA_ARGS__); \ + (_VAR##_0_11) = _PERMUTE((_VAR##_0_11), ##__VA_ARGS__); \ + (_VAR##_1_11) = _PERMUTE((_VAR##_1_11), ##__VA_ARGS__); \ + (_VAR##_2_11) = _PERMUTE((_VAR##_2_11), ##__VA_ARGS__); \ + (_VAR##_3_11) = _PERMUTE((_VAR##_3_11), ##__VA_ARGS__); \ + (_VAR##_0_12) = _PERMUTE((_VAR##_0_12), ##__VA_ARGS__); \ + (_VAR##_1_12) = _PERMUTE((_VAR##_1_12), ##__VA_ARGS__); \ + (_VAR##_2_12) = _PERMUTE((_VAR##_2_12), ##__VA_ARGS__); \ + (_VAR##_3_12) = _PERMUTE((_VAR##_3_12), ##__VA_ARGS__); \ + (_VAR##_0_13) = _PERMUTE((_VAR##_0_13), ##__VA_ARGS__); \ + (_VAR##_1_13) = _PERMUTE((_VAR##_1_13), ##__VA_ARGS__); \ + (_VAR##_2_13) = _PERMUTE((_VAR##_2_13), ##__VA_ARGS__); \ + (_VAR##_3_13) = _PERMUTE((_VAR##_3_13), ##__VA_ARGS__); \ + (_VAR##_0_14) = _PERMUTE((_VAR##_0_14), ##__VA_ARGS__); \ + (_VAR##_1_14) = _PERMUTE((_VAR##_1_14), ##__VA_ARGS__); \ + (_VAR##_2_14) = _PERMUTE((_VAR##_2_14), ##__VA_ARGS__); \ + (_VAR##_3_14) = _PERMUTE((_VAR##_3_14), ##__VA_ARGS__); \ + (_VAR##_0_15) = _PERMUTE((_VAR##_0_15), ##__VA_ARGS__); \ + (_VAR##_1_15) = _PERMUTE((_VAR##_1_15), ##__VA_ARGS__); \ + (_VAR##_2_15) = _PERMUTE((_VAR##_2_15), ##__VA_ARGS__); \ + (_VAR##_3_15) = _PERMUTE((_VAR##_3_15), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_4X32(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_4X16(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_16) = _PERMUTE((_VAR##_0_16), ##__VA_ARGS__); \ + (_VAR##_1_16) = _PERMUTE((_VAR##_1_16), ##__VA_ARGS__); \ + (_VAR##_2_16) = _PERMUTE((_VAR##_2_16), ##__VA_ARGS__); \ + (_VAR##_3_16) = _PERMUTE((_VAR##_3_16), ##__VA_ARGS__); \ + (_VAR##_0_17) = _PERMUTE((_VAR##_0_17), ##__VA_ARGS__); \ + (_VAR##_1_17) = _PERMUTE((_VAR##_1_17), ##__VA_ARGS__); \ + (_VAR##_2_17) = _PERMUTE((_VAR##_2_17), ##__VA_ARGS__); \ + (_VAR##_3_17) = _PERMUTE((_VAR##_3_17), ##__VA_ARGS__); \ + (_VAR##_0_18) = _PERMUTE((_VAR##_0_18), ##__VA_ARGS__); \ + (_VAR##_1_18) = _PERMUTE((_VAR##_1_18), ##__VA_ARGS__); \ + (_VAR##_2_18) = _PERMUTE((_VAR##_2_18), ##__VA_ARGS__); \ + (_VAR##_3_18) = _PERMUTE((_VAR##_3_18), ##__VA_ARGS__); \ + (_VAR##_0_19) = _PERMUTE((_VAR##_0_19), ##__VA_ARGS__); \ + (_VAR##_1_19) = _PERMUTE((_VAR##_1_19), ##__VA_ARGS__); \ + (_VAR##_2_19) = _PERMUTE((_VAR##_2_19), ##__VA_ARGS__); \ + (_VAR##_3_19) = _PERMUTE((_VAR##_3_19), ##__VA_ARGS__); \ + (_VAR##_0_20) = _PERMUTE((_VAR##_0_20), ##__VA_ARGS__); \ + (_VAR##_1_20) = _PERMUTE((_VAR##_1_20), ##__VA_ARGS__); \ + (_VAR##_2_20) = _PERMUTE((_VAR##_2_20), ##__VA_ARGS__); \ + (_VAR##_3_20) = _PERMUTE((_VAR##_3_20), ##__VA_ARGS__); \ + (_VAR##_0_21) = _PERMUTE((_VAR##_0_21), ##__VA_ARGS__); \ + (_VAR##_1_21) = _PERMUTE((_VAR##_1_21), ##__VA_ARGS__); \ + (_VAR##_2_21) = _PERMUTE((_VAR##_2_21), ##__VA_ARGS__); \ + (_VAR##_3_21) = _PERMUTE((_VAR##_3_21), ##__VA_ARGS__); \ + (_VAR##_0_22) = _PERMUTE((_VAR##_0_22), ##__VA_ARGS__); \ + (_VAR##_1_22) = _PERMUTE((_VAR##_1_22), ##__VA_ARGS__); \ + (_VAR##_2_22) = _PERMUTE((_VAR##_2_22), ##__VA_ARGS__); \ + (_VAR##_3_22) = _PERMUTE((_VAR##_3_22), ##__VA_ARGS__); \ + (_VAR##_0_23) = _PERMUTE((_VAR##_0_23), ##__VA_ARGS__); \ + (_VAR##_1_23) = _PERMUTE((_VAR##_1_23), ##__VA_ARGS__); \ + (_VAR##_2_23) = _PERMUTE((_VAR##_2_23), ##__VA_ARGS__); \ + (_VAR##_3_23) = _PERMUTE((_VAR##_3_23), ##__VA_ARGS__); \ + (_VAR##_0_24) = _PERMUTE((_VAR##_0_24), ##__VA_ARGS__); \ + (_VAR##_1_24) = _PERMUTE((_VAR##_1_24), ##__VA_ARGS__); \ + (_VAR##_2_24) = _PERMUTE((_VAR##_2_24), ##__VA_ARGS__); \ + (_VAR##_3_24) = _PERMUTE((_VAR##_3_24), ##__VA_ARGS__); \ + (_VAR##_0_25) = _PERMUTE((_VAR##_0_25), ##__VA_ARGS__); \ + (_VAR##_1_25) = _PERMUTE((_VAR##_1_25), ##__VA_ARGS__); \ + (_VAR##_2_25) = _PERMUTE((_VAR##_2_25), ##__VA_ARGS__); \ + (_VAR##_3_25) = _PERMUTE((_VAR##_3_25), ##__VA_ARGS__); \ + (_VAR##_0_26) = _PERMUTE((_VAR##_0_26), ##__VA_ARGS__); \ + (_VAR##_1_26) = _PERMUTE((_VAR##_1_26), ##__VA_ARGS__); \ + (_VAR##_2_26) = _PERMUTE((_VAR##_2_26), ##__VA_ARGS__); \ + (_VAR##_3_26) = _PERMUTE((_VAR##_3_26), ##__VA_ARGS__); \ + (_VAR##_0_27) = _PERMUTE((_VAR##_0_27), ##__VA_ARGS__); \ + (_VAR##_1_27) = _PERMUTE((_VAR##_1_27), ##__VA_ARGS__); \ + (_VAR##_2_27) = _PERMUTE((_VAR##_2_27), ##__VA_ARGS__); \ + (_VAR##_3_27) = _PERMUTE((_VAR##_3_27), ##__VA_ARGS__); \ + (_VAR##_0_28) = _PERMUTE((_VAR##_0_28), ##__VA_ARGS__); \ + (_VAR##_1_28) = _PERMUTE((_VAR##_1_28), ##__VA_ARGS__); \ + (_VAR##_2_28) = _PERMUTE((_VAR##_2_28), ##__VA_ARGS__); \ + (_VAR##_3_28) = _PERMUTE((_VAR##_3_28), ##__VA_ARGS__); \ + (_VAR##_0_29) = _PERMUTE((_VAR##_0_29), ##__VA_ARGS__); \ + (_VAR##_1_29) = _PERMUTE((_VAR##_1_29), ##__VA_ARGS__); \ + (_VAR##_2_29) = _PERMUTE((_VAR##_2_29), ##__VA_ARGS__); \ + (_VAR##_3_29) = _PERMUTE((_VAR##_3_29), ##__VA_ARGS__); \ + (_VAR##_0_30) = _PERMUTE((_VAR##_0_30), ##__VA_ARGS__); \ + (_VAR##_1_30) = _PERMUTE((_VAR##_1_30), ##__VA_ARGS__); \ + (_VAR##_2_30) = _PERMUTE((_VAR##_2_30), ##__VA_ARGS__); \ + (_VAR##_3_30) = _PERMUTE((_VAR##_3_30), ##__VA_ARGS__); \ + (_VAR##_0_31) = _PERMUTE((_VAR##_0_31), ##__VA_ARGS__); \ + (_VAR##_1_31) = _PERMUTE((_VAR##_1_31), ##__VA_ARGS__); \ + (_VAR##_2_31) = _PERMUTE((_VAR##_2_31), ##__VA_ARGS__); \ + (_VAR##_3_31) = _PERMUTE((_VAR##_3_31), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_8X1(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_4X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_4_0) = _PERMUTE((_VAR##_4_0), ##__VA_ARGS__); \ + (_VAR##_5_0) = _PERMUTE((_VAR##_5_0), ##__VA_ARGS__); \ + (_VAR##_6_0) = _PERMUTE((_VAR##_6_0), ##__VA_ARGS__); \ + (_VAR##_7_0) = _PERMUTE((_VAR##_7_0), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_8X2(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_8X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); \ + (_VAR##_1_1) = _PERMUTE((_VAR##_1_1), ##__VA_ARGS__); \ + (_VAR##_2_1) = _PERMUTE((_VAR##_2_1), ##__VA_ARGS__); \ + (_VAR##_3_1) = _PERMUTE((_VAR##_3_1), ##__VA_ARGS__); \ + (_VAR##_4_1) = _PERMUTE((_VAR##_4_1), ##__VA_ARGS__); \ + (_VAR##_5_1) = _PERMUTE((_VAR##_5_1), ##__VA_ARGS__); \ + (_VAR##_6_1) = _PERMUTE((_VAR##_6_1), ##__VA_ARGS__); \ + (_VAR##_7_1) = _PERMUTE((_VAR##_7_1), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_8X4(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_8X2(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \ + (_VAR##_1_2) = _PERMUTE((_VAR##_1_2), ##__VA_ARGS__); \ + (_VAR##_2_2) = _PERMUTE((_VAR##_2_2), ##__VA_ARGS__); \ + (_VAR##_3_2) = _PERMUTE((_VAR##_3_2), ##__VA_ARGS__); \ + (_VAR##_4_2) = _PERMUTE((_VAR##_4_2), ##__VA_ARGS__); \ + (_VAR##_5_2) = _PERMUTE((_VAR##_5_2), ##__VA_ARGS__); \ + (_VAR##_6_2) = _PERMUTE((_VAR##_6_2), ##__VA_ARGS__); \ + (_VAR##_7_2) = _PERMUTE((_VAR##_7_2), ##__VA_ARGS__); \ + (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); \ + (_VAR##_1_3) = _PERMUTE((_VAR##_1_3), ##__VA_ARGS__); \ + (_VAR##_2_3) = _PERMUTE((_VAR##_2_3), ##__VA_ARGS__); \ + (_VAR##_3_3) = _PERMUTE((_VAR##_3_3), ##__VA_ARGS__); \ + (_VAR##_4_3) = _PERMUTE((_VAR##_4_3), ##__VA_ARGS__); \ + (_VAR##_5_3) = _PERMUTE((_VAR##_5_3), ##__VA_ARGS__); \ + (_VAR##_6_3) = _PERMUTE((_VAR##_6_3), ##__VA_ARGS__); \ + (_VAR##_7_3) = _PERMUTE((_VAR##_7_3), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_8X8(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_8X4(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_4) = _PERMUTE((_VAR##_0_4), ##__VA_ARGS__); \ + (_VAR##_1_4) = _PERMUTE((_VAR##_1_4), ##__VA_ARGS__); \ + (_VAR##_2_4) = _PERMUTE((_VAR##_2_4), ##__VA_ARGS__); \ + (_VAR##_3_4) = _PERMUTE((_VAR##_3_4), ##__VA_ARGS__); \ + (_VAR##_4_4) = _PERMUTE((_VAR##_4_4), ##__VA_ARGS__); \ + (_VAR##_5_4) = _PERMUTE((_VAR##_5_4), ##__VA_ARGS__); \ + (_VAR##_6_4) = _PERMUTE((_VAR##_6_4), ##__VA_ARGS__); \ + (_VAR##_7_4) = _PERMUTE((_VAR##_7_4), ##__VA_ARGS__); \ + (_VAR##_0_5) = _PERMUTE((_VAR##_0_5), ##__VA_ARGS__); \ + (_VAR##_1_5) = _PERMUTE((_VAR##_1_5), ##__VA_ARGS__); \ + (_VAR##_2_5) = _PERMUTE((_VAR##_2_5), ##__VA_ARGS__); \ + (_VAR##_3_5) = _PERMUTE((_VAR##_3_5), ##__VA_ARGS__); \ + (_VAR##_4_5) = _PERMUTE((_VAR##_4_5), ##__VA_ARGS__); \ + (_VAR##_5_5) = _PERMUTE((_VAR##_5_5), ##__VA_ARGS__); \ + (_VAR##_6_5) = _PERMUTE((_VAR##_6_5), ##__VA_ARGS__); \ + (_VAR##_7_5) = _PERMUTE((_VAR##_7_5), ##__VA_ARGS__); \ + (_VAR##_0_6) = _PERMUTE((_VAR##_0_6), ##__VA_ARGS__); \ + (_VAR##_1_6) = _PERMUTE((_VAR##_1_6), ##__VA_ARGS__); \ + (_VAR##_2_6) = _PERMUTE((_VAR##_2_6), ##__VA_ARGS__); \ + (_VAR##_3_6) = _PERMUTE((_VAR##_3_6), ##__VA_ARGS__); \ + (_VAR##_4_6) = _PERMUTE((_VAR##_4_6), ##__VA_ARGS__); \ + (_VAR##_5_6) = _PERMUTE((_VAR##_5_6), ##__VA_ARGS__); \ + (_VAR##_6_6) = _PERMUTE((_VAR##_6_6), ##__VA_ARGS__); \ + (_VAR##_7_6) = _PERMUTE((_VAR##_7_6), ##__VA_ARGS__); \ + (_VAR##_0_7) = _PERMUTE((_VAR##_0_7), ##__VA_ARGS__); \ + (_VAR##_1_7) = _PERMUTE((_VAR##_1_7), ##__VA_ARGS__); \ + (_VAR##_2_7) = _PERMUTE((_VAR##_2_7), ##__VA_ARGS__); \ + (_VAR##_3_7) = _PERMUTE((_VAR##_3_7), ##__VA_ARGS__); \ + (_VAR##_4_7) = _PERMUTE((_VAR##_4_7), ##__VA_ARGS__); \ + (_VAR##_5_7) = _PERMUTE((_VAR##_5_7), ##__VA_ARGS__); \ + (_VAR##_6_7) = _PERMUTE((_VAR##_6_7), ##__VA_ARGS__); \ + (_VAR##_7_7) = _PERMUTE((_VAR##_7_7), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_8X16(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_8X8(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_8) = _PERMUTE((_VAR##_0_8), ##__VA_ARGS__); \ + (_VAR##_1_8) = _PERMUTE((_VAR##_1_8), ##__VA_ARGS__); \ + (_VAR##_2_8) = _PERMUTE((_VAR##_2_8), ##__VA_ARGS__); \ + (_VAR##_3_8) = _PERMUTE((_VAR##_3_8), ##__VA_ARGS__); \ + (_VAR##_4_8) = _PERMUTE((_VAR##_4_8), ##__VA_ARGS__); \ + (_VAR##_5_8) = _PERMUTE((_VAR##_5_8), ##__VA_ARGS__); \ + (_VAR##_6_8) = _PERMUTE((_VAR##_6_8), ##__VA_ARGS__); \ + (_VAR##_7_8) = _PERMUTE((_VAR##_7_8), ##__VA_ARGS__); \ + (_VAR##_0_9) = _PERMUTE((_VAR##_0_9), ##__VA_ARGS__); \ + (_VAR##_1_9) = _PERMUTE((_VAR##_1_9), ##__VA_ARGS__); \ + (_VAR##_2_9) = _PERMUTE((_VAR##_2_9), ##__VA_ARGS__); \ + (_VAR##_3_9) = _PERMUTE((_VAR##_3_9), ##__VA_ARGS__); \ + (_VAR##_4_9) = _PERMUTE((_VAR##_4_9), ##__VA_ARGS__); \ + (_VAR##_5_9) = _PERMUTE((_VAR##_5_9), ##__VA_ARGS__); \ + (_VAR##_6_9) = _PERMUTE((_VAR##_6_9), ##__VA_ARGS__); \ + (_VAR##_7_9) = _PERMUTE((_VAR##_7_9), ##__VA_ARGS__); \ + (_VAR##_0_10) = _PERMUTE((_VAR##_0_10), ##__VA_ARGS__); \ + (_VAR##_1_10) = _PERMUTE((_VAR##_1_10), ##__VA_ARGS__); \ + (_VAR##_2_10) = _PERMUTE((_VAR##_2_10), ##__VA_ARGS__); \ + (_VAR##_3_10) = _PERMUTE((_VAR##_3_10), ##__VA_ARGS__); \ + (_VAR##_4_10) = _PERMUTE((_VAR##_4_10), ##__VA_ARGS__); \ + (_VAR##_5_10) = _PERMUTE((_VAR##_5_10), ##__VA_ARGS__); \ + (_VAR##_6_10) = _PERMUTE((_VAR##_6_10), ##__VA_ARGS__); \ + (_VAR##_7_10) = _PERMUTE((_VAR##_7_10), ##__VA_ARGS__); \ + (_VAR##_0_11) = _PERMUTE((_VAR##_0_11), ##__VA_ARGS__); \ + (_VAR##_1_11) = _PERMUTE((_VAR##_1_11), ##__VA_ARGS__); \ + (_VAR##_2_11) = _PERMUTE((_VAR##_2_11), ##__VA_ARGS__); \ + (_VAR##_3_11) = _PERMUTE((_VAR##_3_11), ##__VA_ARGS__); \ + (_VAR##_4_11) = _PERMUTE((_VAR##_4_11), ##__VA_ARGS__); \ + (_VAR##_5_11) = _PERMUTE((_VAR##_5_11), ##__VA_ARGS__); \ + (_VAR##_6_11) = _PERMUTE((_VAR##_6_11), ##__VA_ARGS__); \ + (_VAR##_7_11) = _PERMUTE((_VAR##_7_11), ##__VA_ARGS__); \ + (_VAR##_0_12) = _PERMUTE((_VAR##_0_12), ##__VA_ARGS__); \ + (_VAR##_1_12) = _PERMUTE((_VAR##_1_12), ##__VA_ARGS__); \ + (_VAR##_2_12) = _PERMUTE((_VAR##_2_12), ##__VA_ARGS__); \ + (_VAR##_3_12) = _PERMUTE((_VAR##_3_12), ##__VA_ARGS__); \ + (_VAR##_4_12) = _PERMUTE((_VAR##_4_12), ##__VA_ARGS__); \ + (_VAR##_5_12) = _PERMUTE((_VAR##_5_12), ##__VA_ARGS__); \ + (_VAR##_6_12) = _PERMUTE((_VAR##_6_12), ##__VA_ARGS__); \ + (_VAR##_7_12) = _PERMUTE((_VAR##_7_12), ##__VA_ARGS__); \ + (_VAR##_0_13) = _PERMUTE((_VAR##_0_13), ##__VA_ARGS__); \ + (_VAR##_1_13) = _PERMUTE((_VAR##_1_13), ##__VA_ARGS__); \ + (_VAR##_2_13) = _PERMUTE((_VAR##_2_13), ##__VA_ARGS__); \ + (_VAR##_3_13) = _PERMUTE((_VAR##_3_13), ##__VA_ARGS__); \ + (_VAR##_4_13) = _PERMUTE((_VAR##_4_13), ##__VA_ARGS__); \ + (_VAR##_5_13) = _PERMUTE((_VAR##_5_13), ##__VA_ARGS__); \ + (_VAR##_6_13) = _PERMUTE((_VAR##_6_13), ##__VA_ARGS__); \ + (_VAR##_7_13) = _PERMUTE((_VAR##_7_13), ##__VA_ARGS__); \ + (_VAR##_0_14) = _PERMUTE((_VAR##_0_14), ##__VA_ARGS__); \ + (_VAR##_1_14) = _PERMUTE((_VAR##_1_14), ##__VA_ARGS__); \ + (_VAR##_2_14) = _PERMUTE((_VAR##_2_14), ##__VA_ARGS__); \ + (_VAR##_3_14) = _PERMUTE((_VAR##_3_14), ##__VA_ARGS__); \ + (_VAR##_4_14) = _PERMUTE((_VAR##_4_14), ##__VA_ARGS__); \ + (_VAR##_5_14) = _PERMUTE((_VAR##_5_14), ##__VA_ARGS__); \ + (_VAR##_6_14) = _PERMUTE((_VAR##_6_14), ##__VA_ARGS__); \ + (_VAR##_7_14) = _PERMUTE((_VAR##_7_14), ##__VA_ARGS__); \ + (_VAR##_0_15) = _PERMUTE((_VAR##_0_15), ##__VA_ARGS__); \ + (_VAR##_1_15) = _PERMUTE((_VAR##_1_15), ##__VA_ARGS__); \ + (_VAR##_2_15) = _PERMUTE((_VAR##_2_15), ##__VA_ARGS__); \ + (_VAR##_3_15) = _PERMUTE((_VAR##_3_15), ##__VA_ARGS__); \ + (_VAR##_4_15) = _PERMUTE((_VAR##_4_15), ##__VA_ARGS__); \ + (_VAR##_5_15) = _PERMUTE((_VAR##_5_15), ##__VA_ARGS__); \ + (_VAR##_6_15) = _PERMUTE((_VAR##_6_15), ##__VA_ARGS__); \ + (_VAR##_7_15) = _PERMUTE((_VAR##_7_15), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_8X32(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_8X16(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_16) = _PERMUTE((_VAR##_0_16), ##__VA_ARGS__); \ + (_VAR##_1_16) = _PERMUTE((_VAR##_1_16), ##__VA_ARGS__); \ + (_VAR##_2_16) = _PERMUTE((_VAR##_2_16), ##__VA_ARGS__); \ + (_VAR##_3_16) = _PERMUTE((_VAR##_3_16), ##__VA_ARGS__); \ + (_VAR##_4_16) = _PERMUTE((_VAR##_4_16), ##__VA_ARGS__); \ + (_VAR##_5_16) = _PERMUTE((_VAR##_5_16), ##__VA_ARGS__); \ + (_VAR##_6_16) = _PERMUTE((_VAR##_6_16), ##__VA_ARGS__); \ + (_VAR##_7_16) = _PERMUTE((_VAR##_7_16), ##__VA_ARGS__); \ + (_VAR##_0_17) = _PERMUTE((_VAR##_0_17), ##__VA_ARGS__); \ + (_VAR##_1_17) = _PERMUTE((_VAR##_1_17), ##__VA_ARGS__); \ + (_VAR##_2_17) = _PERMUTE((_VAR##_2_17), ##__VA_ARGS__); \ + (_VAR##_3_17) = _PERMUTE((_VAR##_3_17), ##__VA_ARGS__); \ + (_VAR##_4_17) = _PERMUTE((_VAR##_4_17), ##__VA_ARGS__); \ + (_VAR##_5_17) = _PERMUTE((_VAR##_5_17), ##__VA_ARGS__); \ + (_VAR##_6_17) = _PERMUTE((_VAR##_6_17), ##__VA_ARGS__); \ + (_VAR##_7_17) = _PERMUTE((_VAR##_7_17), ##__VA_ARGS__); \ + (_VAR##_0_18) = _PERMUTE((_VAR##_0_18), ##__VA_ARGS__); \ + (_VAR##_1_18) = _PERMUTE((_VAR##_1_18), ##__VA_ARGS__); \ + (_VAR##_2_18) = _PERMUTE((_VAR##_2_18), ##__VA_ARGS__); \ + (_VAR##_3_18) = _PERMUTE((_VAR##_3_18), ##__VA_ARGS__); \ + (_VAR##_4_18) = _PERMUTE((_VAR##_4_18), ##__VA_ARGS__); \ + (_VAR##_5_18) = _PERMUTE((_VAR##_5_18), ##__VA_ARGS__); \ + (_VAR##_6_18) = _PERMUTE((_VAR##_6_18), ##__VA_ARGS__); \ + (_VAR##_7_18) = _PERMUTE((_VAR##_7_18), ##__VA_ARGS__); \ + (_VAR##_0_19) = _PERMUTE((_VAR##_0_19), ##__VA_ARGS__); \ + (_VAR##_1_19) = _PERMUTE((_VAR##_1_19), ##__VA_ARGS__); \ + (_VAR##_2_19) = _PERMUTE((_VAR##_2_19), ##__VA_ARGS__); \ + (_VAR##_3_19) = _PERMUTE((_VAR##_3_19), ##__VA_ARGS__); \ + (_VAR##_4_19) = _PERMUTE((_VAR##_4_19), ##__VA_ARGS__); \ + (_VAR##_5_19) = _PERMUTE((_VAR##_5_19), ##__VA_ARGS__); \ + (_VAR##_6_19) = _PERMUTE((_VAR##_6_19), ##__VA_ARGS__); \ + (_VAR##_7_19) = _PERMUTE((_VAR##_7_19), ##__VA_ARGS__); \ + (_VAR##_0_20) = _PERMUTE((_VAR##_0_20), ##__VA_ARGS__); \ + (_VAR##_1_20) = _PERMUTE((_VAR##_1_20), ##__VA_ARGS__); \ + (_VAR##_2_20) = _PERMUTE((_VAR##_2_20), ##__VA_ARGS__); \ + (_VAR##_3_20) = _PERMUTE((_VAR##_3_20), ##__VA_ARGS__); \ + (_VAR##_4_20) = _PERMUTE((_VAR##_4_20), ##__VA_ARGS__); \ + (_VAR##_5_20) = _PERMUTE((_VAR##_5_20), ##__VA_ARGS__); \ + (_VAR##_6_20) = _PERMUTE((_VAR##_6_20), ##__VA_ARGS__); \ + (_VAR##_7_20) = _PERMUTE((_VAR##_7_20), ##__VA_ARGS__); \ + (_VAR##_0_21) = _PERMUTE((_VAR##_0_21), ##__VA_ARGS__); \ + (_VAR##_1_21) = _PERMUTE((_VAR##_1_21), ##__VA_ARGS__); \ + (_VAR##_2_21) = _PERMUTE((_VAR##_2_21), ##__VA_ARGS__); \ + (_VAR##_3_21) = _PERMUTE((_VAR##_3_21), ##__VA_ARGS__); \ + (_VAR##_4_21) = _PERMUTE((_VAR##_4_21), ##__VA_ARGS__); \ + (_VAR##_5_21) = _PERMUTE((_VAR##_5_21), ##__VA_ARGS__); \ + (_VAR##_6_21) = _PERMUTE((_VAR##_6_21), ##__VA_ARGS__); \ + (_VAR##_7_21) = _PERMUTE((_VAR##_7_21), ##__VA_ARGS__); \ + (_VAR##_0_22) = _PERMUTE((_VAR##_0_22), ##__VA_ARGS__); \ + (_VAR##_1_22) = _PERMUTE((_VAR##_1_22), ##__VA_ARGS__); \ + (_VAR##_2_22) = _PERMUTE((_VAR##_2_22), ##__VA_ARGS__); \ + (_VAR##_3_22) = _PERMUTE((_VAR##_3_22), ##__VA_ARGS__); \ + (_VAR##_4_22) = _PERMUTE((_VAR##_4_22), ##__VA_ARGS__); \ + (_VAR##_5_22) = _PERMUTE((_VAR##_5_22), ##__VA_ARGS__); \ + (_VAR##_6_22) = _PERMUTE((_VAR##_6_22), ##__VA_ARGS__); \ + (_VAR##_7_22) = _PERMUTE((_VAR##_7_22), ##__VA_ARGS__); \ + (_VAR##_0_23) = _PERMUTE((_VAR##_0_23), ##__VA_ARGS__); \ + (_VAR##_1_23) = _PERMUTE((_VAR##_1_23), ##__VA_ARGS__); \ + (_VAR##_2_23) = _PERMUTE((_VAR##_2_23), ##__VA_ARGS__); \ + (_VAR##_3_23) = _PERMUTE((_VAR##_3_23), ##__VA_ARGS__); \ + (_VAR##_4_23) = _PERMUTE((_VAR##_4_23), ##__VA_ARGS__); \ + (_VAR##_5_23) = _PERMUTE((_VAR##_5_23), ##__VA_ARGS__); \ + (_VAR##_6_23) = _PERMUTE((_VAR##_6_23), ##__VA_ARGS__); \ + (_VAR##_7_23) = _PERMUTE((_VAR##_7_23), ##__VA_ARGS__); \ + (_VAR##_0_24) = _PERMUTE((_VAR##_0_24), ##__VA_ARGS__); \ + (_VAR##_1_24) = _PERMUTE((_VAR##_1_24), ##__VA_ARGS__); \ + (_VAR##_2_24) = _PERMUTE((_VAR##_2_24), ##__VA_ARGS__); \ + (_VAR##_3_24) = _PERMUTE((_VAR##_3_24), ##__VA_ARGS__); \ + (_VAR##_4_24) = _PERMUTE((_VAR##_4_24), ##__VA_ARGS__); \ + (_VAR##_5_24) = _PERMUTE((_VAR##_5_24), ##__VA_ARGS__); \ + (_VAR##_6_24) = _PERMUTE((_VAR##_6_24), ##__VA_ARGS__); \ + (_VAR##_7_24) = _PERMUTE((_VAR##_7_24), ##__VA_ARGS__); \ + (_VAR##_0_25) = _PERMUTE((_VAR##_0_25), ##__VA_ARGS__); \ + (_VAR##_1_25) = _PERMUTE((_VAR##_1_25), ##__VA_ARGS__); \ + (_VAR##_2_25) = _PERMUTE((_VAR##_2_25), ##__VA_ARGS__); \ + (_VAR##_3_25) = _PERMUTE((_VAR##_3_25), ##__VA_ARGS__); \ + (_VAR##_4_25) = _PERMUTE((_VAR##_4_25), ##__VA_ARGS__); \ + (_VAR##_5_25) = _PERMUTE((_VAR##_5_25), ##__VA_ARGS__); \ + (_VAR##_6_25) = _PERMUTE((_VAR##_6_25), ##__VA_ARGS__); \ + (_VAR##_7_25) = _PERMUTE((_VAR##_7_25), ##__VA_ARGS__); \ + (_VAR##_0_26) = _PERMUTE((_VAR##_0_26), ##__VA_ARGS__); \ + (_VAR##_1_26) = _PERMUTE((_VAR##_1_26), ##__VA_ARGS__); \ + (_VAR##_2_26) = _PERMUTE((_VAR##_2_26), ##__VA_ARGS__); \ + (_VAR##_3_26) = _PERMUTE((_VAR##_3_26), ##__VA_ARGS__); \ + (_VAR##_4_26) = _PERMUTE((_VAR##_4_26), ##__VA_ARGS__); \ + (_VAR##_5_26) = _PERMUTE((_VAR##_5_26), ##__VA_ARGS__); \ + (_VAR##_6_26) = _PERMUTE((_VAR##_6_26), ##__VA_ARGS__); \ + (_VAR##_7_26) = _PERMUTE((_VAR##_7_26), ##__VA_ARGS__); \ + (_VAR##_0_27) = _PERMUTE((_VAR##_0_27), ##__VA_ARGS__); \ + (_VAR##_1_27) = _PERMUTE((_VAR##_1_27), ##__VA_ARGS__); \ + (_VAR##_2_27) = _PERMUTE((_VAR##_2_27), ##__VA_ARGS__); \ + (_VAR##_3_27) = _PERMUTE((_VAR##_3_27), ##__VA_ARGS__); \ + (_VAR##_4_27) = _PERMUTE((_VAR##_4_27), ##__VA_ARGS__); \ + (_VAR##_5_27) = _PERMUTE((_VAR##_5_27), ##__VA_ARGS__); \ + (_VAR##_6_27) = _PERMUTE((_VAR##_6_27), ##__VA_ARGS__); \ + (_VAR##_7_27) = _PERMUTE((_VAR##_7_27), ##__VA_ARGS__); \ + (_VAR##_0_28) = _PERMUTE((_VAR##_0_28), ##__VA_ARGS__); \ + (_VAR##_1_28) = _PERMUTE((_VAR##_1_28), ##__VA_ARGS__); \ + (_VAR##_2_28) = _PERMUTE((_VAR##_2_28), ##__VA_ARGS__); \ + (_VAR##_3_28) = _PERMUTE((_VAR##_3_28), ##__VA_ARGS__); \ + (_VAR##_4_28) = _PERMUTE((_VAR##_4_28), ##__VA_ARGS__); \ + (_VAR##_5_28) = _PERMUTE((_VAR##_5_28), ##__VA_ARGS__); \ + (_VAR##_6_28) = _PERMUTE((_VAR##_6_28), ##__VA_ARGS__); \ + (_VAR##_7_28) = _PERMUTE((_VAR##_7_28), ##__VA_ARGS__); \ + (_VAR##_0_29) = _PERMUTE((_VAR##_0_29), ##__VA_ARGS__); \ + (_VAR##_1_29) = _PERMUTE((_VAR##_1_29), ##__VA_ARGS__); \ + (_VAR##_2_29) = _PERMUTE((_VAR##_2_29), ##__VA_ARGS__); \ + (_VAR##_3_29) = _PERMUTE((_VAR##_3_29), ##__VA_ARGS__); \ + (_VAR##_4_29) = _PERMUTE((_VAR##_4_29), ##__VA_ARGS__); \ + (_VAR##_5_29) = _PERMUTE((_VAR##_5_29), ##__VA_ARGS__); \ + (_VAR##_6_29) = _PERMUTE((_VAR##_6_29), ##__VA_ARGS__); \ + (_VAR##_7_29) = _PERMUTE((_VAR##_7_29), ##__VA_ARGS__); \ + (_VAR##_0_30) = _PERMUTE((_VAR##_0_30), ##__VA_ARGS__); \ + (_VAR##_1_30) = _PERMUTE((_VAR##_1_30), ##__VA_ARGS__); \ + (_VAR##_2_30) = _PERMUTE((_VAR##_2_30), ##__VA_ARGS__); \ + (_VAR##_3_30) = _PERMUTE((_VAR##_3_30), ##__VA_ARGS__); \ + (_VAR##_4_30) = _PERMUTE((_VAR##_4_30), ##__VA_ARGS__); \ + (_VAR##_5_30) = _PERMUTE((_VAR##_5_30), ##__VA_ARGS__); \ + (_VAR##_6_30) = _PERMUTE((_VAR##_6_30), ##__VA_ARGS__); \ + (_VAR##_7_30) = _PERMUTE((_VAR##_7_30), ##__VA_ARGS__); \ + (_VAR##_0_31) = _PERMUTE((_VAR##_0_31), ##__VA_ARGS__); \ + (_VAR##_1_31) = _PERMUTE((_VAR##_1_31), ##__VA_ARGS__); \ + (_VAR##_2_31) = _PERMUTE((_VAR##_2_31), ##__VA_ARGS__); \ + (_VAR##_3_31) = _PERMUTE((_VAR##_3_31), ##__VA_ARGS__); \ + (_VAR##_4_31) = _PERMUTE((_VAR##_4_31), ##__VA_ARGS__); \ + (_VAR##_5_31) = _PERMUTE((_VAR##_5_31), ##__VA_ARGS__); \ + (_VAR##_6_31) = _PERMUTE((_VAR##_6_31), ##__VA_ARGS__); \ + (_VAR##_7_31) = _PERMUTE((_VAR##_7_31), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_16X1(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_8X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_8_0) = _PERMUTE((_VAR##_8_0), ##__VA_ARGS__); \ + (_VAR##_9_0) = _PERMUTE((_VAR##_9_0), ##__VA_ARGS__); \ + (_VAR##_10_0) = _PERMUTE((_VAR##_10_0), ##__VA_ARGS__); \ + (_VAR##_11_0) = _PERMUTE((_VAR##_11_0), ##__VA_ARGS__); \ + (_VAR##_12_0) = _PERMUTE((_VAR##_12_0), ##__VA_ARGS__); \ + (_VAR##_13_0) = _PERMUTE((_VAR##_13_0), ##__VA_ARGS__); \ + (_VAR##_14_0) = _PERMUTE((_VAR##_14_0), ##__VA_ARGS__); \ + (_VAR##_15_0) = _PERMUTE((_VAR##_15_0), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_16X2(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_16X1(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); \ + (_VAR##_1_1) = _PERMUTE((_VAR##_1_1), ##__VA_ARGS__); \ + (_VAR##_2_1) = _PERMUTE((_VAR##_2_1), ##__VA_ARGS__); \ + (_VAR##_3_1) = _PERMUTE((_VAR##_3_1), ##__VA_ARGS__); \ + (_VAR##_4_1) = _PERMUTE((_VAR##_4_1), ##__VA_ARGS__); \ + (_VAR##_5_1) = _PERMUTE((_VAR##_5_1), ##__VA_ARGS__); \ + (_VAR##_6_1) = _PERMUTE((_VAR##_6_1), ##__VA_ARGS__); \ + (_VAR##_7_1) = _PERMUTE((_VAR##_7_1), ##__VA_ARGS__); \ + (_VAR##_8_1) = _PERMUTE((_VAR##_8_1), ##__VA_ARGS__); \ + (_VAR##_9_1) = _PERMUTE((_VAR##_9_1), ##__VA_ARGS__); \ + (_VAR##_10_1) = _PERMUTE((_VAR##_10_1), ##__VA_ARGS__); \ + (_VAR##_11_1) = _PERMUTE((_VAR##_11_1), ##__VA_ARGS__); \ + (_VAR##_12_1) = _PERMUTE((_VAR##_12_1), ##__VA_ARGS__); \ + (_VAR##_13_1) = _PERMUTE((_VAR##_13_1), ##__VA_ARGS__); \ + (_VAR##_14_1) = _PERMUTE((_VAR##_14_1), ##__VA_ARGS__); \ + (_VAR##_15_1) = _PERMUTE((_VAR##_15_1), ##__VA_ARGS__); + +#define MATRIX_VAR_PERMUTE_16X4(_VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_16X2(_VAR, _PERMUTE, ##__VA_ARGS__) \ + (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \ + (_VAR##_1_2) = _PERMUTE((_VAR##_1_2), ##__VA_ARGS__); \ + (_VAR##_2_2) = _PERMUTE((_VAR##_2_2), ##__VA_ARGS__); \ + (_VAR##_3_2) = _PERMUTE((_VAR##_3_2), ##__VA_ARGS__); \ + (_VAR##_4_2) = _PERMUTE((_VAR##_4_2), ##__VA_ARGS__); \ + (_VAR##_5_2) = _PERMUTE((_VAR##_5_2), ##__VA_ARGS__); \ + (_VAR##_6_2) = _PERMUTE((_VAR##_6_2), ##__VA_ARGS__); \ + (_VAR##_7_2) = _PERMUTE((_VAR##_7_2), ##__VA_ARGS__); \ + (_VAR##_8_2) = _PERMUTE((_VAR##_8_2), ##__VA_ARGS__); \ + (_VAR##_9_2) = _PERMUTE((_VAR##_9_2), ##__VA_ARGS__); \ + (_VAR##_10_2) = _PERMUTE((_VAR##_10_2), ##__VA_ARGS__); \ + (_VAR##_11_2) = _PERMUTE((_VAR##_11_2), ##__VA_ARGS__); \ + (_VAR##_12_2) = _PERMUTE((_VAR##_12_2), ##__VA_ARGS__); \ + (_VAR##_13_2) = _PERMUTE((_VAR##_13_2), ##__VA_ARGS__); \ + (_VAR##_14_2) = _PERMUTE((_VAR##_14_2), ##__VA_ARGS__); \ + (_VAR##_15_2) = _PERMUTE((_VAR##_15_2), ##__VA_ARGS__); \ + (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); \ + (_VAR##_1_3) = _PERMUTE((_VAR##_1_3), ##__VA_ARGS__); \ + (_VAR##_2_3) = _PERMUTE((_VAR##_2_3), ##__VA_ARGS__); \ + (_VAR##_3_3) = _PERMUTE((_VAR##_3_3), ##__VA_ARGS__); \ + (_VAR##_4_3) = _PERMUTE((_VAR##_4_3), ##__VA_ARGS__); \ + (_VAR##_5_3) = _PERMUTE((_VAR##_5_3), ##__VA_ARGS__); \ + (_VAR##_6_3) = _PERMUTE((_VAR##_6_3), ##__VA_ARGS__); \ + (_VAR##_7_3) = _PERMUTE((_VAR##_7_3), ##__VA_ARGS__); \ + (_VAR##_8_3) = _PERMUTE((_VAR##_8_3), ##__VA_ARGS__); \ + (_VAR##_9_3) = _PERMUTE((_VAR##_9_3), ##__VA_ARGS__); \ + (_VAR##_10_3) = _PERMUTE((_VAR##_10_3), ##__VA_ARGS__); \ + (_VAR##_11_3) = _PERMUTE((_VAR##_11_3), ##__VA_ARGS__); \ + (_VAR##_12_3) = _PERMUTE((_VAR##_12_3), ##__VA_ARGS__); \ + (_VAR##_13_3) = _PERMUTE((_VAR##_13_3), ##__VA_ARGS__); \ + (_VAR##_14_3) = _PERMUTE((_VAR##_14_3), ##__VA_ARGS__); \ + (_VAR##_15_3) = _PERMUTE((_VAR##_15_3), ##__VA_ARGS__); + +#define MATRIX_VAR_PROC_2X1(_K, _LHS, _RHS, _RES, _PROCESS) \ + _PROCESS((_LHS##_0), (_RHS), (_RES##_0_##_K)) \ + _PROCESS((_LHS##_1), (_RHS), (_RES##_1_##_K)) + +#define MATRIX_VAR_PROC_4X1(_K, _LHS, _RHS, _RES, _PROCESS) \ + MATRIX_VAR_PROC_2X1(_K, _LHS, _RHS, _RES, _PROCESS) \ + _PROCESS((_LHS##_2), (_RHS), (_RES##_2_##_K)) \ + _PROCESS((_LHS##_3), (_RHS), (_RES##_3_##_K)) + +#define MATRIX_VAR_PROC_8X1(_K, _LHS, _RHS, _RES, _PROCESS) \ + MATRIX_VAR_PROC_4X1(_K, _LHS, _RHS, _RES, _PROCESS) \ + _PROCESS((_LHS##_4), (_RHS), (_RES##_4_##_K)) \ + _PROCESS((_LHS##_5), (_RHS), (_RES##_5_##_K)) \ + _PROCESS((_LHS##_6), (_RHS), (_RES##_6_##_K)) \ + _PROCESS((_LHS##_7), (_RHS), (_RES##_7_##_K)) + +#define MATRIX_VAR_PROC_16X1(_K, _LHS, _RHS, _RES, _PROCESS) \ + MATRIX_VAR_PROC_8X1(_K, _LHS, _RHS, _RES, _PROCESS) \ + _PROCESS((_LHS##_8), (_RHS), (_RES##_8_##_K)) \ + _PROCESS((_LHS##_9), (_RHS), (_RES##_9_##_K)) \ + _PROCESS((_LHS##_10), (_RHS), (_RES##_10_##_K)) \ + _PROCESS((_LHS##_11), (_RHS), (_RES##_11_##_K)) \ + _PROCESS((_LHS##_12), (_RHS), (_RES##_12_##_K)) \ + _PROCESS((_LHS##_13), (_RHS), (_RES##_13_##_K)) \ + _PROCESS((_LHS##_14), (_RHS), (_RES##_14_##_K)) \ + _PROCESS((_LHS##_15), (_RHS), (_RES##_15_##_K)) + +#define MATRIX_VAR_PROC_1X2(_K, _LHS, _RHS, _RES, _PROCESS) \ + _PROCESS((_LHS), (_RHS##_0), (_RES##_##_K##_0)) \ + _PROCESS((_LHS), (_RHS##_1), (_RES##_##_K##_1)) + +#define MATRIX_VAR_PROC_1X4(_K, _LHS, _RHS, _RES, _PROCESS) \ + MATRIX_VAR_PROC_1X2(_K, _LHS, _RHS, _RES, _PROCESS) \ + _PROCESS((_LHS), (_RHS##_2), (_RES##_##_K##_2)) \ + _PROCESS((_LHS), (_RHS##_3), (_RES##_##_K##_3)) + +#define MATRIX_VAR_PROC_1X8(_K, _LHS, _RHS, _RES, _PROCESS) \ + MATRIX_VAR_PROC_1X4(_K, _LHS, _RHS, _RES, _PROCESS) \ + _PROCESS((_LHS), (_RHS##_4), (_RES##_##_K##_4)) \ + _PROCESS((_LHS), (_RHS##_5), (_RES##_##_K##_5)) \ + _PROCESS((_LHS), (_RHS##_6), (_RES##_##_K##_6)) \ + _PROCESS((_LHS), (_RHS##_7), (_RES##_##_K##_7)) + +#define MATRIX_VAR_PROC_1X16(_K, _LHS, _RHS, _RES, _PROCESS) \ + MATRIX_VAR_PROC_1X8(_K, _LHS, _RHS, _RES, _PROCESS) \ + _PROCESS((_LHS), (_RHS##_8), (_RES##_##_K##_8)) \ + _PROCESS((_LHS), (_RHS##_9), (_RES##_##_K##_9)) \ + _PROCESS((_LHS), (_RHS##_10), (_RES##_##_K##_10)) \ + _PROCESS((_LHS), (_RHS##_11), (_RES##_##_K##_11)) \ + _PROCESS((_LHS), (_RHS##_12), (_RES##_##_K##_12)) \ + _PROCESS((_LHS), (_RHS##_13), (_RES##_##_K##_13)) \ + _PROCESS((_LHS), (_RHS##_14), (_RES##_##_K##_14)) \ + _PROCESS((_LHS), (_RHS##_15), (_RES##_##_K##_15)) + +#define MATRIX_VAR_INIT(_M, _N, _VAR_TYPE, _VAR_NAME, _VAR_INIT) \ + MATRIX_VAR_INIT_##_M##X##_N(_VAR_TYPE, _VAR_NAME, _VAR_INIT) + +#define MATRIX_VAR_STORE(_M, _N, _STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \ + MATRIX_VAR_STORE_##_M##X##_N(_STEP, _VAR, _ARRAY, _STORE, _NORM, \ + ##__VA_ARGS__) + +#define MATRIX_VAR_PERMUTE(_M, _N, _VAR, _PERMUTE, ...) \ + MATRIX_VAR_PERMUTE_##_M##X##_N(_VAR, _PERMUTE, ##__VA_ARGS__) + +#define MATRIX_VAR_PROC(_M, _N, _K, _LHS, _RHS, _RES, _PROCESS) \ + MATRIX_VAR_PROC_##_M##X##_N(_K, _LHS, _RHS, _RES, _PROCESS) diff --git a/src/ailego/math/matrix_utility.i b/src/ailego/math/matrix_utility.i new file mode 100644 index 0000000..7716a49 --- /dev/null +++ b/src/ailego/math/matrix_utility.i @@ -0,0 +1,265 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +//! Absolute value of a float +static inline float FastAbs(float x) { + uint32_t *p = reinterpret_cast(&x); + *p &= 0x7fffffffu; + return *reinterpret_cast(p); +} + +#if defined(__SSE__) +static inline float HorizontalMax_FP32_V128(__m128 v) { + __m128 x1 = _mm_shuffle_ps(v, v, _MM_SHUFFLE(0, 0, 3, 2)); + __m128 x2 = _mm_max_ps(v, x1); + __m128 x3 = _mm_shuffle_ps(x2, x2, _MM_SHUFFLE(0, 0, 0, 1)); + __m128 x4 = _mm_max_ps(x2, x3); + return _mm_cvtss_f32(x4); +} + +static inline float HorizontalAdd_FP32_V128(__m128 v) { +#ifdef __SSE3__ + __m128 x1 = _mm_hadd_ps(v, v); + __m128 x2 = _mm_hadd_ps(x1, x1); + return _mm_cvtss_f32(x2); +#else + __m128 x1 = _mm_movehl_ps(v, v); + __m128 x2 = _mm_add_ps(v, x1); + __m128 x3 = _mm_shuffle_ps(x2, x2, 1); + __m128 x4 = _mm_add_ss(x2, x3); + return _mm_cvtss_f32(x4); +#endif +} +#endif // __SSE__ + +#if defined(__SSE2__) +static inline int32_t HorizontalAdd_INT32_V128(__m128i v) { +#ifdef __SSE3__ + __m128i x1 = _mm_hadd_epi32(v, v); + __m128i x2 = _mm_hadd_epi32(x1, x1); + return _mm_cvtsi128_si32(x2); +#else + __m128i x1 = _mm_shuffle_epi32(v, _MM_SHUFFLE(0, 0, 3, 2)); + __m128i x2 = _mm_add_epi32(v, x1); + __m128i x3 = _mm_shuffle_epi32(x2, _MM_SHUFFLE(0, 0, 0, 1)); + __m128i x4 = _mm_add_epi32(x2, x3); + return _mm_cvtsi128_si32(x4); +#endif +} + +static inline int64_t HorizontalAdd_INT64_V128(__m128i v) { +#ifdef __SSE4_1__ + return (_mm_extract_epi64(v, 0) + _mm_extract_epi64(v, 1)); +#else + return _mm_cvtsi128_si64( + _mm_add_epi64(_mm_shuffle_epi32(v, _MM_SHUFFLE(0, 0, 3, 2)), v)); +#endif +} +#endif // __SSE2__ + +#if defined(__SSSE3__) +static const __m128i POPCNT_LOOKUP_SSE = + _mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4); + +static inline __m128i VerticalPopCount_INT8_V128(__m128i v) { +#if defined(__AVX512VL__) && defined(__AVX512BITALG__) + return _mm_popcnt_epi8(v); +#else + const __m128i low_mask = _mm_set1_epi8(0x0f); + __m128i lo = _mm_shuffle_epi8(POPCNT_LOOKUP_SSE, _mm_and_si128(v, low_mask)); + __m128i hi = _mm_shuffle_epi8(POPCNT_LOOKUP_SSE, + _mm_and_si128(_mm_srli_epi32(v, 4), low_mask)); + return _mm_add_epi8(lo, hi); +#endif // __AVX512VL__ && __AVX512BITALG__ +} + +static inline __m128i VerticalPopCount_INT16_V128(__m128i v) { +#if defined(__AVX512VL__) && defined(__AVX512BITALG__) + return _mm_popcnt_epi16(v); +#else + __m128i total = VerticalPopCount_INT8_V128(v); + return _mm_add_epi16(_mm_srli_epi16(total, 8), + _mm_and_si128(total, _mm_set1_epi16(0xff))); +#endif // __AVX512VL__ && __AVX512BITALG__ +} + +static inline __m128i VerticalPopCount_INT32_V128(__m128i v) { +#if defined(__AVX512VL__) && defined(__AVX512VPOPCNTDQ__) + return _mm_popcnt_epi32(v); +#else + __m128i total = + _mm_madd_epi16(VerticalPopCount_INT8_V128(v), _mm_set1_epi16(1)); + return _mm_add_epi32(_mm_srli_epi32(total, 8), + _mm_and_si128(total, _mm_set1_epi32(0xff))); +#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__ +} + +static inline __m128i VerticalPopCount_INT64_V128(__m128i v) { +#if defined(__AVX512VL__) && defined(__AVX512VPOPCNTDQ__) + return _mm_popcnt_epi64(v); +#else + return _mm_sad_epu8(VerticalPopCount_INT8_V128(v), _mm_setzero_si128()); +#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__ +} +#endif // __SSSE3__ + +#if defined(__SSE4_1__) +static inline int16_t HorizontalMax_UINT8_V128(__m128i v) { + v = _mm_max_epu8(v, _mm_shuffle_epi32(v, _MM_SHUFFLE(3, 2, 3, 2))); + v = _mm_max_epu8(v, _mm_shuffle_epi32(v, _MM_SHUFFLE(1, 1, 1, 1))); + v = _mm_max_epu8(v, _mm_shufflelo_epi16(v, _MM_SHUFFLE(1, 1, 1, 1))); + v = _mm_max_epu8(v, _mm_srli_epi16(v, 8)); + return static_cast(_mm_cvtsi128_si32(v)); +} +#endif // __SSE4_1__ + +#if defined(__AVX__) +static inline float HorizontalMax_FP32_V256(__m256 v) { + __m256 x1 = _mm256_permute_ps(v, _MM_SHUFFLE(0, 0, 3, 2)); + __m256 x2 = _mm256_max_ps(v, x1); + __m256 x3 = _mm256_permute_ps(x2, _MM_SHUFFLE(0, 0, 0, 1)); + __m256 x4 = _mm256_max_ps(x2, x3); + __m128 x5 = _mm256_extractf128_ps(x4, 1); + __m128 x6 = _mm_max_ss(_mm256_castps256_ps128(x4), x5); + return _mm_cvtss_f32(x6); +} + +static inline float HorizontalAdd_FP32_V256(__m256 v) { + __m256 x1 = _mm256_hadd_ps(v, v); + __m256 x2 = _mm256_hadd_ps(x1, x1); + __m128 x3 = _mm256_extractf128_ps(x2, 1); + __m128 x4 = _mm_add_ss(_mm256_castps256_ps128(x2), x3); + return _mm_cvtss_f32(x4); +} +#endif // __AVX__ + +#if defined(__AVX2__) +#define POPCNT_MASK1_INT8_AVX _mm256_set1_epi8(0x0f) +#define POPCNT_MASK1_INT16_AVX _mm256_set1_epi16(1) +#define POPCNT_MASK2_INT16_AVX _mm256_set1_epi16(0xff) +#define POPCNT_MASK1_INT32_AVX _mm256_set1_epi32(0xff) +#define POPCNT_ZERO_AVX _mm256_setzero_si256() +#define POPCNT_LOOKUP_AVX _mm256_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4) + +static inline __m256i VerticalPopCount_INT8_V256(__m256i v) { +#if defined(__AVX512VL__) && defined(__AVX512BITALG__) + return _mm256_popcnt_epi8(v); +#else + __m256i lo = _mm256_shuffle_epi8(POPCNT_LOOKUP_AVX, + _mm256_and_si256(v, POPCNT_MASK1_INT8_AVX)); + __m256i hi = _mm256_shuffle_epi8( + POPCNT_LOOKUP_AVX, + _mm256_and_si256(_mm256_srli_epi32(v, 4), POPCNT_MASK1_INT8_AVX)); + return _mm256_add_epi8(lo, hi); +#endif // __AVX512VL__ && __AVX512BITALG__ +} + +static inline __m256i VerticalPopCount_INT16_V256(__m256i v) { +#if defined(__AVX512VL__) && defined(__AVX512BITALG__) + return _mm256_popcnt_epi16(v); +#else + __m256i total = VerticalPopCount_INT8_V256(v); + return _mm256_add_epi16(_mm256_srli_epi16(total, 8), + _mm256_and_si256(total, POPCNT_MASK2_INT16_AVX)); +#endif // __AVX512VL__ && __AVX512BITALG__ +} + +static inline __m256i VerticalPopCount_INT32_V256(__m256i v) { +#if defined(__AVX512VL__) && defined(__AVX512VPOPCNTDQ__) + return _mm256_popcnt_epi32(v); +#else + __m256i total = + _mm256_madd_epi16(VerticalPopCount_INT8_V256(v), POPCNT_MASK1_INT16_AVX); + return _mm256_add_epi32(_mm256_srli_epi32(total, 8), + _mm256_and_si256(total, POPCNT_MASK1_INT32_AVX)); +#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__ +} + +static inline __m256i VerticalPopCount_INT64_V256(__m256i v) { +#if defined(__AVX512VL__) && defined(__AVX512VPOPCNTDQ__) + return _mm256_popcnt_epi64(v); +#else + return _mm256_sad_epu8(VerticalPopCount_INT8_V256(v), POPCNT_ZERO_AVX); +#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__ +} + +static inline int16_t HorizontalMax_UINT8_V256(__m256i v) { + v = _mm256_max_epu8(v, _mm256_shuffle_epi32(v, _MM_SHUFFLE(3, 2, 3, 2))); + v = _mm256_max_epu8(v, _mm256_shuffle_epi32(v, _MM_SHUFFLE(1, 1, 1, 1))); + v = _mm256_max_epu8(v, _mm256_shufflelo_epi16(v, _MM_SHUFFLE(1, 1, 1, 1))); + __m128i x = + _mm_max_epu8(_mm256_castsi256_si128(v), _mm256_extractf128_si256(v, 1)); + x = _mm_max_epu8(x, _mm_srli_epi16(x, 8)); + return static_cast(_mm_cvtsi128_si32(x)); +} + +static inline int32_t HorizontalAdd_INT32_V256(__m256i v) { + __m256i x1 = _mm256_hadd_epi32(v, v); + __m256i x2 = _mm256_hadd_epi32(x1, x1); + __m128i x3 = _mm256_extractf128_si256(x2, 1); + __m128i x4 = _mm_add_epi32(_mm256_castsi256_si128(x2), x3); + return _mm_cvtsi128_si32(x4); +} + +static inline int64_t HorizontalAdd_INT64_V256(__m256i v) { + __m256i x1 = _mm256_shuffle_epi32(v, _MM_SHUFFLE(1, 0, 3, 2)); + __m256i x2 = _mm256_add_epi64(v, x1); + __m128i x3 = _mm256_extractf128_si256(x2, 1); + __m128i x4 = _mm_add_epi64(_mm256_extractf128_si256(x2, 0), x3); + return _mm_cvtsi128_si64(x4); +} +#endif // __AVX2__ + +#if defined(__AVX512F__) +static inline float HorizontalMax_FP32_V512(__m512 v) { + __m256 low = _mm512_castps512_ps256(v); + __m256 high = + _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(v), 1)); + return HorizontalMax_FP32_V256(_mm256_max_ps(low, high)); +} + +static inline float HorizontalAdd_FP32_V512(__m512 v) { + __m256 low = _mm512_castps512_ps256(v); + __m256 high = + _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(v), 1)); + return HorizontalAdd_FP32_V256(_mm256_add_ps(low, high)); +} +#endif // __AVX512F__ + +#if defined(__AVX512FP16__) +static inline float HorizontalMax_FP16_V512(__m512h v) { + __m512 low = _mm512_cvtxph_ps(_mm512_castph512_ph256(v)); + __m512 high = _mm512_cvtxph_ps( + _mm256_castpd_ph(_mm512_extractf64x4_pd(_mm512_castph_pd(v), 1))); + return HorizontalMax_FP32_V512(_mm512_max_ps(low, high)); +} + +static inline float HorizontalAdd_FP16_V512(__m512h v) { + __m512 low = _mm512_cvtxph_ps(_mm512_castph512_ph256(v)); + __m512 high = _mm512_cvtxph_ps( + _mm256_castpd_ph(_mm512_extractf64x4_pd(_mm512_castph_pd(v), 1))); + + return HorizontalAdd_FP32_V512(_mm512_add_ps(low, high)); +} +#endif // __AVX512FP16__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix.h b/src/ailego/math/mips_euclidean_distance_matrix.h new file mode 100644 index 0000000..1fdd380 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix.h @@ -0,0 +1,1115 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include "distance_utility.h" + +namespace zvec { +namespace ailego { + +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +/*! Compute the Mips SphericalInjection Squared Euclidean Distance with the two + * vectors's InnerProduct and each squared l2-normlized value, and the e2 is + * 1.0 / max_squared_l2_norm + */ +static float inline ComputeSphericalInjection(double ip, double u2, double v2, + double e2) { + if (e2 == 0.0) { + // Implies *localized* spherical injection. + return static_cast(2.0 - 2.0 * ip / std::max(u2, v2)); + } + auto v = (1.0 - e2 * u2) * (1.0 - e2 * v2); + auto score = v > 0.0 ? (1.0 - e2 * ip - std::sqrt(v)) : (1.0 - e2 * ip); + return static_cast(score * 2.0); +} + +/*! Mips Squared Euclidean Distance Matrix + */ +template +struct MipsSquaredEuclideanDistanceMatrix; + +/*! Mips Squared Euclidean Distance Matrix (M=1, N=1) + */ +template +struct MipsSquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Compute the distance between matrix and query by SphericalInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out) { + ailego_assert(p && q && dim && out); + + float sum = 0.0; + float u2 = 0.0; + float v2 = 0.0; + for (size_t i = 0; i < dim; ++i) { + u2 += p[i] * p[i]; + v2 += q[i] * q[i]; + sum += static_cast(p[i] * q[i]); + } + *out = ComputeSphericalInjection(sum, u2, v2, e2); + } + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out) { + ailego_assert(p && q && dim && out); + + float sum = 0.0; + float u2 = 0.0; + float v2 = 0.0; + for (size_t i = 0; i < dim; ++i) { + u2 += p[i] * p[i]; + v2 += q[i] * q[i]; + sum += MathHelper::SquaredDifference(p[i], q[i]); + } + + sum *= e2; + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + *out = sum; + } +}; + +template <> +struct MipsSquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = uint8_t; + + // Compute the distance between matrix and query by SphericalInjection + static void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out); + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out); +}; + +template <> +struct MipsSquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = int8_t; + + // Compute the distance between matrix and query by SphericalInjection + static void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out); + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out); +}; + +template <> +struct MipsSquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = Float16; + + // Compute the distance between matrix and query by SphericalInjection + static void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out); + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out); +}; + +template <> +struct MipsSquaredEuclideanDistanceMatrix { + //! Type of value + using ValueType = float; + + // Compute the distance between matrix and query by SphericalInjection + static void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out); + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out); +}; + +/*! Mips Squared Euclidean Distance Matrix (M >= 2, N >= 2) + */ +template +struct MipsSquaredEuclideanDistanceMatrix< + T, M, N, typename std::enable_if= 2 && N >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Compute the distance between matrix and query by SphericalInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out) { + ailego_assert(p && q && dim && out); + if (dim == 0) { + return; + } + + std::array u2; + std::array v2; + for (size_t i = 0; i < M; ++i) { + const ValueType p_val = p[i]; + u2[i] = static_cast(p_val * p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r = static_cast(p_val * q[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = static_cast(q[i] * q[i]); + } + p += M; + q += N; + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + const ValueType p_val = p[i]; + u2[i] += static_cast(p_val * p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r += static_cast(p_val * q[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] += static_cast(q[i] * q[i]); + } + p += M; + q += N; + } + + // Compute the injection + for (size_t i = 0; i < M; ++i) { + float *r = out + i; + const float u2_val = u2[i]; + for (size_t j = 0; j < N; ++j) { + *r = ComputeSphericalInjection(*r, u2_val, v2[j], e2); + r += M; + } + } + } + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out) { + ailego_assert(p && q && dim && out); + if (dim == 0) { + return; + } + + std::array u2; + std::array v2; + for (size_t i = 0; i < M; ++i) { + const ValueType p_val = p[i]; + u2[i] = static_cast(p_val * p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r = MathHelper::SquaredDifference(p_val, q[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = static_cast(q[i] * q[i]); + } + p += M; + q += N; + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + const ValueType p_val = p[i]; + u2[i] += static_cast(p_val * p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r += MathHelper::SquaredDifference(p_val, q[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] += static_cast(q[i] * q[i]); + } + p += M; + q += N; + } + + // Compute the injections + float *r = out; + for (size_t i = 0; i < M; ++i) { + u2[i] *= e2; + for (size_t j = 0; j < N; ++j) { + (*r++) *= e2; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] *= e2; + } + for (size_t k = 0; k < m; ++k) { + for (size_t i = 0; i < M; ++i) { + r = out + i; + float u2_val = u2[i]; + u2[i] = u2_val * u2_val; + for (size_t j = 0; j < N; ++j) { + *r += (u2_val - v2[j]) * (u2_val - v2[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = v2[i] * v2[i]; + } + } + } +}; + +/*! Mips Squared Euclidean Distance Matrix (N=1) + */ +template +struct MipsSquaredEuclideanDistanceMatrix< + T, M, 1, typename std::enable_if= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Compute the distance between matrix and query by SphericalInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out) { + ailego_assert(p && q && dim && out); + const ValueType *q_end = q + dim; + if (q == q_end) { + return; + } + + std::array u2; + ValueType q_val = *q++; + float v2 = static_cast(q_val * q_val); + for (size_t i = 0; i < M; ++i) { + u2[i] = static_cast(p[i] * p[i]); + out[i] = static_cast(p[i] * q_val); + } + p += M; + + while (q != q_end) { + q_val = *q++; + v2 += static_cast(q_val * q_val); + for (size_t i = 0; i < M; ++i) { + u2[i] += static_cast(p[i] * p[i]); + out[i] += static_cast(p[i] * q_val); + } + p += M; + } + + // Compute the injection + for (size_t i = 0; i < M; ++i) { + out[i] = ComputeSphericalInjection(out[i], u2[i], v2, e2); + } + } + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out) { + ailego_assert(p && q && dim && out); + const ValueType *q_end = q + dim; + if (q == q_end) { + return; + } + + std::array u2; + ValueType q_val = *q++; + float v2 = static_cast(q_val * q_val); + for (size_t i = 0; i < M; ++i) { + u2[i] = static_cast(p[i] * p[i]); + out[i] = MathHelper::SquaredDifference(p[i], q_val); + } + p += M; + + while (q != q_end) { + q_val = *q++; + v2 += static_cast(q_val * q_val); + for (size_t i = 0; i < M; ++i) { + u2[i] += static_cast(p[i] * p[i]); + out[i] += MathHelper::SquaredDifference(p[i], q_val); + } + p += M; + } + + // Compute the injections + for (size_t i = 0; i < M; ++i) { + out[i] *= e2; + u2[i] *= e2; + } + v2 *= e2; + for (size_t k = 0; k < m; ++k) { + for (size_t i = 0; i < M; ++i) { + const float u_val = u2[i]; + u2[i] = u_val * u_val; + out[i] += (u_val - v2) * (u_val - v2); + } + v2 = v2 * v2; + } + } +}; + +/*! Mips Squared Euclidean Distance Matrix (INT8, M >=2, N >= 2) + */ +template +struct MipsSquaredEuclideanDistanceMatrix< + int8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> { + //! Type of value + using ValueType = int8_t; + + // Compute the distance between matrix and query by SphericalInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out) { + ailego_assert(p && q && dim && !(dim & 3) && out); + dim >>= 2; + if (dim == 0) { + return; + } + + std::array u2; + std::array v2; + const uint32_t *p_it = reinterpret_cast(p); + const uint32_t *q_it = reinterpret_cast(q); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = p_it[i]; + u2[i] = Squared(p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r = FusedMultiplyAdd(p_val, q_it[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = Squared(q_it[i]); + } + p_it += M; + q_it += N; + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = p_it[i]; + u2[i] += Squared(p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r += FusedMultiplyAdd(p_val, q_it[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] += Squared(q_it[i]); + } + p_it += M; + q_it += N; + } + + // Compute the injection + for (size_t i = 0; i < M; ++i) { + float *r = out + i; + const float u2_val = u2[i]; + for (size_t j = 0; j < N; ++j) { + *r = ComputeSphericalInjection(*r, u2_val, v2[j], e2); + r += M; + } + } + } + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out) { + ailego_assert(p && q && dim && !(dim & 3) && out); + dim >>= 2; + if (dim == 0) { + return; + } + + std::array u2; + std::array v2; + const uint32_t *p_it = reinterpret_cast(p); + const uint32_t *q_it = reinterpret_cast(q); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = p_it[i]; + u2[i] = Squared(p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r = SquaredDifference(p_val, q_it[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = Squared(q_it[i]); + } + p_it += M; + q_it += N; + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = p_it[i]; + u2[i] += Squared(p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r += SquaredDifference(p_val, q_it[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] += Squared(q_it[i]); + } + p_it += M; + q_it += N; + } + + // Compute the injections + float *r = out; + for (size_t i = 0; i < M; ++i) { + u2[i] *= e2; + for (size_t j = 0; j < N; ++j) { + (*r++) *= e2; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] *= e2; + } + for (size_t k = 0; k < m; ++k) { + for (size_t i = 0; i < M; ++i) { + r = out + i; + float u2_val = u2[i]; + u2[i] = u2_val * u2_val; + for (size_t j = 0; j < N; ++j) { + *r += (u2_val - v2[j]) * (u2_val - v2[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = v2[i] * v2[i]; + } + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + return static_cast((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) + + (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) + + (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) + + (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24)); + } + + //! Calculate the squared difference + static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) { + return static_cast(MathHelper::SquaredDifference( + (int8_t)(lhs >> 0), (int8_t)(rhs >> 0)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 8), (int8_t)(rhs >> 8)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 16), (int8_t)(rhs >> 16)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 24), (int8_t)(rhs >> 24))); + } + + //! Calculate sum of squared values + static inline float Squared(uint32_t v) { + return static_cast((int8_t)(v >> 0) * (int8_t)(v >> 0) + + (int8_t)(v >> 8) * (int8_t)(v >> 8) + + (int8_t)(v >> 16) * (int8_t)(v >> 16) + + (int8_t)(v >> 24) * (int8_t)(v >> 24)); + } +}; + +/*! Mips Squared Euclidean Distance Matrix (INT8, N=1) + */ +template +struct MipsSquaredEuclideanDistanceMatrix< + int8_t, M, 1, typename std::enable_if= 2>::type> { + //! Type of value + using ValueType = int8_t; + + // Compute the distance between matrix and query by SphericalInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out) { + ailego_assert(p && q && dim && !(dim & 3) && out); + const uint32_t *p_it = reinterpret_cast(p); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 2); + if (q_it == q_end) { + return; + } + + std::array u2; + uint32_t q_val = *q_it++; + float v2 = Squared(q_val); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = *p_it++; + u2[i] = Squared(p_val); + out[i] = FusedMultiplyAdd(p_val, q_val); + } + + while (q_it != q_end) { + q_val = *q_it++; + v2 += Squared(q_val); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = *p_it++; + u2[i] += Squared(p_val); + out[i] += FusedMultiplyAdd(p_val, q_val); + } + } + + // Compute the injection + for (size_t i = 0; i < M; ++i) { + out[i] = ComputeSphericalInjection(out[i], u2[i], v2, e2); + } + } + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out) { + ailego_assert(p && q && dim && !(dim & 3) && out); + const uint32_t *p_it = reinterpret_cast(p); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 2); + if (q_it == q_end) { + return; + } + + std::array u2; + uint32_t q_val = *q_it++; + float v2 = Squared(q_val); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = *p_it++; + u2[i] = Squared(p_val); + out[i] = SquaredDifference(p_val, q_val); + } + + while (q_it != q_end) { + q_val = *q_it++; + v2 += Squared(q_val); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = *p_it++; + u2[i] += Squared(p_val); + out[i] += SquaredDifference(p_val, q_val); + } + } + + // Compute the injections + for (size_t i = 0; i < M; ++i) { + out[i] *= e2; + u2[i] *= e2; + } + v2 *= e2; + for (size_t k = 0; k < m; ++k) { + for (size_t i = 0; i < M; ++i) { + const float u_val = u2[i]; + u2[i] = u_val * u_val; + out[i] += (u_val - v2) * (u_val - v2); + } + v2 = v2 * v2; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + return static_cast((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) + + (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) + + (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) + + (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24)); + } + + //! Calculate the squared difference + static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) { + return static_cast(MathHelper::SquaredDifference( + (int8_t)(lhs >> 0), (int8_t)(rhs >> 0)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 8), (int8_t)(rhs >> 8)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 16), (int8_t)(rhs >> 16)) + + MathHelper::SquaredDifference( + (int8_t)(lhs >> 24), (int8_t)(rhs >> 24))); + } + + //! Calculate sum of squared values + static inline float Squared(uint32_t v) { + return static_cast((int8_t)(v >> 0) * (int8_t)(v >> 0) + + (int8_t)(v >> 8) * (int8_t)(v >> 8) + + (int8_t)(v >> 16) * (int8_t)(v >> 16) + + (int8_t)(v >> 24) * (int8_t)(v >> 24)); + } +}; + +/*! Mips Squared Euclidean Distance Matrix (INT4, M >=2, N >= 2) + */ +template +struct MipsSquaredEuclideanDistanceMatrix< + uint8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + // Compute the distance between matrix and query by SphericalInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out) { + ailego_assert(p && q && dim && !(dim & 7) && out); + dim >>= 3; + if (dim == 0) { + return; + } + + std::array u2; + std::array v2; + const uint32_t *p_it = reinterpret_cast(p); + const uint32_t *q_it = reinterpret_cast(q); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = p_it[i]; + u2[i] = Squared(p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r = FusedMultiplyAdd(p_val, q_it[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = Squared(q_it[i]); + } + p_it += M; + q_it += N; + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = p_it[i]; + u2[i] += Squared(p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r += FusedMultiplyAdd(p_val, q_it[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] += Squared(q_it[i]); + } + p_it += M; + q_it += N; + } + + // Compute the injection + for (size_t i = 0; i < M; ++i) { + float *r = out + i; + const float u2_val = u2[i]; + for (size_t j = 0; j < N; ++j) { + *r = ComputeSphericalInjection(*r, u2_val, v2[j], e2); + r += M; + } + } + } + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out) { + ailego_assert(p && q && dim && !(dim & 7) && out); + dim >>= 3; + if (dim == 0) { + return; + } + + std::array u2; + std::array v2; + const uint32_t *p_it = reinterpret_cast(p); + const uint32_t *q_it = reinterpret_cast(q); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = p_it[i]; + u2[i] = Squared(p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r = SquaredDifference(p_val, q_it[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = Squared(q_it[i]); + } + p_it += M; + q_it += N; + + for (size_t k = 1; k < dim; ++k) { + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = p_it[i]; + u2[i] += Squared(p_val); + float *r = out + i; + for (size_t j = 0; j < N; ++j) { + *r += SquaredDifference(p_val, q_it[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] += Squared(q_it[i]); + } + p_it += M; + q_it += N; + } + + // Compute the injections + float *r = out; + for (size_t i = 0; i < M; ++i) { + u2[i] *= e2; + for (size_t j = 0; j < N; ++j) { + (*r++) *= e2; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] *= e2; + } + for (size_t k = 0; k < m; ++k) { + for (size_t i = 0; i < M; ++i) { + r = out + i; + float u2_val = u2[i]; + u2[i] = u2_val * u2_val; + for (size_t j = 0; j < N; ++j) { + *r += (u2_val - v2[j]) * (u2_val - v2[j]); + r += M; + } + } + for (size_t i = 0; i < N; ++i) { + v2[i] = v2[i] * v2[i]; + } + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } + + //! Calculate the squared difference + static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4SquaredDiffTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } + + //! Calculate sum of squared values + static inline float Squared(uint32_t u) { + float sum = 0.0f; + for (size_t i = 0; i < 32; i += 8) { + uint8_t v = (uint8_t)(u >> i); + int8_t lo = (int8_t)(v << 4) >> 4; + int8_t hi = (int8_t)(v & 0xf0) >> 4; + sum += hi * hi + lo * lo; + } + return sum; + } +}; + +/*! Mips Squared Euclidean Distance Matrix (INT4, N=1) + */ +template +struct MipsSquaredEuclideanDistanceMatrix< + uint8_t, M, 1, typename std::enable_if= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + // Compute the distance between matrix and query by SphericalInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + float e2, float *out) { + ailego_assert(p && q && dim && !(dim & 7) && out); + const uint32_t *p_it = reinterpret_cast(p); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 3); + if (q_it == q_end) { + return; + } + + std::array u2; + uint32_t q_val = *q_it++; + float v2 = Squared(q_val); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = *p_it++; + u2[i] = Squared(p_val); + out[i] = FusedMultiplyAdd(p_val, q_val); + } + + while (q_it != q_end) { + q_val = *q_it++; + v2 += Squared(q_val); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = *p_it++; + u2[i] += Squared(p_val); + out[i] += FusedMultiplyAdd(p_val, q_val); + } + } + + // Compute the injection + for (size_t i = 0; i < M; ++i) { + out[i] = ComputeSphericalInjection(out[i], u2[i], v2, e2); + } + } + + // Compute the distance between matrix and query by RepeatedQuadraticInjection + static inline void Compute(const ValueType *p, const ValueType *q, size_t dim, + size_t m, float e2, float *out) { + ailego_assert(p && q && dim && !(dim & 7) && out); + const uint32_t *p_it = reinterpret_cast(p); + const uint32_t *q_it = reinterpret_cast(q); + const uint32_t *q_end = q_it + (dim >> 3); + if (q_it == q_end) { + return; + } + + std::array u2; + uint32_t q_val = *q_it++; + float v2 = Squared(q_val); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = *p_it++; + u2[i] = Squared(p_val); + out[i] = SquaredDifference(p_val, q_val); + } + + while (q_it != q_end) { + q_val = *q_it++; + v2 += Squared(q_val); + for (size_t i = 0; i < M; ++i) { + const uint32_t p_val = *p_it++; + u2[i] += Squared(p_val); + out[i] += SquaredDifference(p_val, q_val); + } + } + + // Compute the injections + for (size_t i = 0; i < M; ++i) { + out[i] *= e2; + u2[i] *= e2; + } + v2 *= e2; + for (size_t k = 0; k < m; ++k) { + for (size_t i = 0; i < M; ++i) { + const float u_val = u2[i]; + u2[i] = u_val * u_val; + out[i] += (u_val - v2) * (u_val - v2); + } + v2 = v2 * v2; + } + } + + protected: + //! Calculate Fused-Multiply-Add + static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } + + //! Calculate the squared difference + static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) { + return static_cast( + Int4SquaredDiffTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] + + Int4SquaredDiffTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]); + } + + //! Calculate sum of squared values + static inline float Squared(uint32_t u) { + float sum = 0.0f; + for (size_t i = 0; i < 32; i += 8) { + uint8_t v = (uint8_t)(u >> i); + int8_t lo = (int8_t)(v << 4) >> 4; + int8_t hi = (int8_t)(v & 0xf0) >> 4; + sum += hi * hi + lo * lo; + } + return sum; + } +}; + +//-------------------------------------------------- +// Sparse +//-------------------------------------------------- +/*! Mips Squared Euclidean Sparse Distance Matrix + */ +template +struct MipsSquaredEuclideanSparseDistanceMatrix { + //! Type of value + using ValueType = typename std::remove_cv::type; + + static float ComputeInnerProductSparseInSegment( + uint32_t m_sparse_count, const uint16_t *m_sparse_index, + const ValueType *m_sparse_value, uint32_t q_sparse_count, + const uint16_t *q_sparse_index, const ValueType *q_sparse_value); + + // Compute the distance between matrix and query by SphericalInjection + static inline void Compute(const void *m_sparse_data_in, + const void *q_sparse_data_in, float *out) { + ailego_assert(m_sparse_data_in && q_sparse_data_in && out); + + const uint8_t *m_sparse_data = + reinterpret_cast(m_sparse_data_in); + const uint8_t *q_sparse_data = + reinterpret_cast(q_sparse_data_in); + + const uint32_t m_sparse_count = + *reinterpret_cast(m_sparse_data); + const uint32_t q_sparse_count = + *reinterpret_cast(q_sparse_data); + + if (m_sparse_count == 0 && q_sparse_count == 0) { + *out = 0; + return; + } + + if (m_sparse_count == 0 || q_sparse_count == 0) { + *out = 2; + return; + } + + const uint32_t m_seg_count = + *reinterpret_cast(m_sparse_data + sizeof(uint32_t)); + const uint32_t q_seg_count = + *reinterpret_cast(q_sparse_data + sizeof(uint32_t)); + + const uint32_t *m_seg_id = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t)); + const uint32_t *q_seg_id = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t)); + + const uint32_t *m_seg_vec_cnt = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t) + m_seg_count * sizeof(uint32_t)); + const uint32_t *q_seg_vec_cnt = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t) + q_seg_count * sizeof(uint32_t)); + + const uint16_t *m_sparse_index = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t) + + m_seg_count * 2 * sizeof(uint32_t)); + const uint16_t *q_sparse_index = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t) + + q_seg_count * 2 * sizeof(uint32_t)); + + const ValueType *m_sparse_value = reinterpret_cast( + m_sparse_data + 2 * sizeof(uint32_t) + + m_seg_count * 2 * sizeof(uint32_t) + m_sparse_count * sizeof(uint16_t)); + const ValueType *q_sparse_value = reinterpret_cast( + q_sparse_data + 2 * sizeof(uint32_t) + + q_seg_count * 2 * sizeof(uint32_t) + q_sparse_count * sizeof(uint16_t)); + + float ip = 0.0f; + + size_t m_s = 0; + size_t q_s = 0; + + size_t m_count = 0; + size_t q_count = 0; + + while (m_s < m_seg_count && q_s < q_seg_count) { + if (m_seg_id[m_s] == q_seg_id[q_s]) { + ip += ComputeInnerProductSparseInSegment( + m_seg_vec_cnt[m_s], m_sparse_index + m_count, + m_sparse_value + m_count, q_seg_vec_cnt[q_s], + q_sparse_index + q_count, q_sparse_value + q_count); + + m_count += m_seg_vec_cnt[m_s]; + q_count += q_seg_vec_cnt[q_s]; + + ++m_s; + ++q_s; + } else if (m_seg_id[m_s] < q_seg_id[q_s]) { + m_count += m_seg_vec_cnt[m_s]; + + ++m_s; + } else { + q_count += q_seg_vec_cnt[q_s]; + + ++q_s; + } + } + + float l2_m{0.0f}; + SquaredNorm2Matrix::Compute(m_sparse_value, m_sparse_count, + &l2_m); + + float l2_q{0.0f}; + SquaredNorm2Matrix::Compute(q_sparse_value, q_sparse_count, + &l2_q); + + *out = ComputeSphericalInjection(ip, l2_m, l2_q, 0.0f); + } +}; + +template +float MipsSquaredEuclideanSparseDistanceMatrix< + T>::ComputeInnerProductSparseInSegment(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const ValueType *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const ValueType *q_sparse_value) { + float sum = 0.0f; + + size_t m_i = 0; + size_t q_i = 0; + while (m_i < m_sparse_count && q_i < q_sparse_count) { + if (m_sparse_index[m_i] == q_sparse_index[q_i]) { + sum += m_sparse_value[m_i] * q_sparse_value[q_i]; + + ++m_i; + ++q_i; + } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) { + ++m_i; + } else { + ++q_i; + } + } + + return sum; +} + +template <> +float MipsSquaredEuclideanSparseDistanceMatrix< + float>::ComputeInnerProductSparseInSegment(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const ValueType *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const ValueType *q_sparse_value); + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc new file mode 100644 index 0000000..2ffe638 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc @@ -0,0 +1,149 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormFp16AVX(const Float16 *lhs, const Float16 *rhs, + size_t size, float *sql, float *sqr) { + __m256 ymm_sum_0 = _mm256_setzero_ps(); + __m256 ymm_sum_1 = _mm256_setzero_ps(); + __m256 ymm_sum_norm1 = _mm256_setzero_ps(); + __m256 ymm_sum_norm2 = _mm256_setzero_ps(); + + const Float16 *last = lhs + size; + const Float16 *last_aligned = lhs + ((size >> 4) << 4); + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m256i ymm_lhs = _mm256_load_si256((const __m256i *)lhs); + __m256i ymm_rhs = _mm256_load_si256((const __m256i *)rhs); + __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_lhs)); + __m256 ymm_lhs_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_lhs, 1)); + __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_rhs)); + __m256 ymm_rhs_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_rhs, 1)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2); + } + if (last >= last_aligned + 8) { + __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)lhs)); + __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)rhs)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + lhs += 8; + rhs += 8; + } + } else { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)lhs); + __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)rhs); + __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_lhs)); + __m256 ymm_lhs_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_lhs, 1)); + __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_rhs)); + __m256 ymm_rhs_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_rhs, 1)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2); + } + if (last >= last_aligned + 8) { + __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)lhs)); + __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)rhs)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + lhs += 8; + rhs += 8; + } + } + + float result = HorizontalAdd_FP32_V256(_mm256_add_ps(ymm_sum_0, ymm_sum_1)); + float norm1 = HorizontalAdd_FP32_V256(ymm_sum_norm1); + float norm2 = HorizontalAdd_FP32_V256(ymm_sum_norm2); + switch (last - lhs) { + case 7: + FMA_FP16_GENERAL(lhs[6], rhs[6], result, norm1, norm2); + /* FALLTHRU */ + case 6: + FMA_FP16_GENERAL(lhs[5], rhs[5], result, norm1, norm2); + /* FALLTHRU */ + case 5: + FMA_FP16_GENERAL(lhs[4], rhs[4], result, norm1, norm2); + /* FALLTHRU */ + case 4: + FMA_FP16_GENERAL(lhs[3], rhs[3], result, norm1, norm2); + /* FALLTHRU */ + case 3: + FMA_FP16_GENERAL(lhs[2], rhs[2], result, norm1, norm2); + /* FALLTHRU */ + case 2: + FMA_FP16_GENERAL(lhs[1], rhs[1], result, norm1, norm2); + /* FALLTHRU */ + case 1: + FMA_FP16_GENERAL(lhs[0], rhs[0], result, norm1, norm2); + } + + *sql = norm1; + *sqr = norm2; + return result; +} + +float MipsEuclideanDistanceSphericalInjectionFp16AVX(const Float16 *lhs, + const Float16 *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormFp16AVX(lhs, rhs, size, &u2, &v2); + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX( + const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormFp16AVX(lhs, rhs, size, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} + +#endif // __AVX__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc new file mode 100644 index 0000000..f5e86ba --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc @@ -0,0 +1,168 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX512F__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormFp16AVX512(const Float16 *lhs, + const Float16 *rhs, size_t size, + float *sql, float *sqr) { + __m512 zmm_sum_0 = _mm512_setzero_ps(); + __m512 zmm_sum_1 = _mm512_setzero_ps(); + __m512 zmm_sum_norm1 = _mm512_setzero_ps(); + __m512 zmm_sum_norm2 = _mm512_setzero_ps(); + + const Float16 *last = lhs + size; + const Float16 *last_aligned = lhs + ((size >> 5) << 5); + if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m512i zmm_lhs = _mm512_load_si512((const __m512i *)lhs); + __m512i zmm_rhs = _mm512_load_si512((const __m512i *)rhs); + __m512 zmm_lhs_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_lhs)); + __m512 zmm_lhs_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_lhs, 1)); + __m512 zmm_rhs_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_rhs)); + __m512 zmm_rhs_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_rhs, 1)); + FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0) + FMA_FP32_AVX512(zmm_lhs_1, zmm_rhs_1, zmm_sum_1) + FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_lhs_1, zmm_lhs_1, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2) + FMA_FP32_AVX512(zmm_rhs_1, zmm_rhs_1, zmm_sum_norm2) + } + if (last >= last_aligned + 16) { + __m512 zmm_lhs_0 = + _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)lhs)); + __m512 zmm_rhs_0 = + _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)rhs)); + FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0) + FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2) + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m512i zmm_lhs = _mm512_loadu_si512((const __m512i *)lhs); + __m512i zmm_rhs = _mm512_loadu_si512((const __m512i *)rhs); + __m512 zmm_lhs_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_lhs)); + __m512 zmm_lhs_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_lhs, 1)); + __m512 zmm_rhs_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_rhs)); + __m512 zmm_rhs_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_rhs, 1)); + FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0) + FMA_FP32_AVX512(zmm_lhs_1, zmm_rhs_1, zmm_sum_1) + FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_lhs_1, zmm_lhs_1, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2) + FMA_FP32_AVX512(zmm_rhs_1, zmm_rhs_1, zmm_sum_norm2) + } + if (last >= last_aligned + 16) { + __m512 zmm_lhs_0 = + _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)lhs)); + __m512 zmm_rhs_0 = + _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)rhs)); + FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0) + FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2) + lhs += 16; + rhs += 16; + } + } + + __m256 ymm_sum_0 = + HorizontalAdd_FP32_V512_TO_V256(_mm512_add_ps(zmm_sum_0, zmm_sum_1)); + __m256 ymm_sum_norm1 = HorizontalAdd_FP32_V512_TO_V256(zmm_sum_norm1); + __m256 ymm_sum_norm2 = HorizontalAdd_FP32_V512_TO_V256(zmm_sum_norm2); + if (last >= lhs + 8) { + __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)lhs)); + __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)rhs)); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + lhs += 8; + rhs += 8; + } + + float result = HorizontalAdd_FP32_V256(ymm_sum_0); + float norm1 = HorizontalAdd_FP32_V256(ymm_sum_norm1); + float norm2 = HorizontalAdd_FP32_V256(ymm_sum_norm2); + switch (last - lhs) { + case 7: + FMA_FP16_GENERAL(lhs[6], rhs[6], result, norm1, norm2); + /* FALLTHRU */ + case 6: + FMA_FP16_GENERAL(lhs[5], rhs[5], result, norm1, norm2); + /* FALLTHRU */ + case 5: + FMA_FP16_GENERAL(lhs[4], rhs[4], result, norm1, norm2); + /* FALLTHRU */ + case 4: + FMA_FP16_GENERAL(lhs[3], rhs[3], result, norm1, norm2); + /* FALLTHRU */ + case 3: + FMA_FP16_GENERAL(lhs[2], rhs[2], result, norm1, norm2); + /* FALLTHRU */ + case 2: + FMA_FP16_GENERAL(lhs[1], rhs[1], result, norm1, norm2); + /* FALLTHRU */ + case 1: + FMA_FP16_GENERAL(lhs[0], rhs[0], result, norm1, norm2); + } + + *sql = norm1; + *sqr = norm2; + return result; +} + +float MipsEuclideanDistanceSphericalInjectionFp16AVX512(const Float16 *lhs, + const Float16 *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormFp16AVX512(lhs, rhs, size, &u2, &v2); + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX512( + const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormFp16AVX512(lhs, rhs, size, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} + +#endif // __AVX512F__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc new file mode 100644 index 0000000..8e40563 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc @@ -0,0 +1,103 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp16NEON( + const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionFp16NEON(const Float16 *lhs, + const Float16 *rhs, + size_t size, float e2); +#endif + +#if defined(__AVX512F__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX512( + const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionFp16AVX512(const Float16 *lhs, + const Float16 *rhs, + size_t size, float e2); +#endif + +#if defined(__AVX__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX( + const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionFp16AVX(const Float16 *lhs, + const Float16 *rhs, + size_t size, float e2); +#endif + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp16Scalar( + const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionFp16Scalar( + const ailego::Float16 *p, const ailego::Float16 *q, size_t dim, float e2); + + +//! Compute the distance between matrix and query by SphericalInjection +void MipsSquaredEuclideanDistanceMatrix::Compute( + const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) { +#if defined(__ARM_NEON) + *out = MipsEuclideanDistanceSphericalInjectionFp16NEON(p, q, dim, e2); +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = MipsEuclideanDistanceSphericalInjectionFp16AVX512(p, q, dim, e2); + return; + } +#endif +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = MipsEuclideanDistanceSphericalInjectionFp16AVX(p, q, dim, e2); + return; + } +#endif //__AVX__ + *out = MipsEuclideanDistanceSphericalInjectionFp16Scalar(p, q, dim, e2); + return; +#endif //__ARM_NEON +} + +//! Compute the distance between matrix and query by RepeatedQuadraticInjection +void MipsSquaredEuclideanDistanceMatrix::Compute( + const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2, + float *out) { +#if defined(__ARM_NEON) + *out = + MipsEuclideanDistanceRepeatedQuadraticInjectionFp16NEON(p, q, dim, m, e2); +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX512(p, q, dim, + m, e2); + return; + } +#endif +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX(p, q, dim, m, + e2); + return; + } +#endif //__AVX__ + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp16Scalar(p, q, dim, m, + e2); + return; +#endif //__ARM_NEON +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc new file mode 100644 index 0000000..b4f4c97 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc @@ -0,0 +1,159 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp16.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) && defined(__aarch64__) +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormFp16NEON(const Float16 *lhs, const Float16 *rhs, + size_t size, float *sql, float *sqr) { + const Float16 *last = lhs + size; + const Float16 *last_aligned = lhs + ((size >> 3) << 3); + float16x8_t v_sum = vdupq_n_f16(0); + float16x8_t v_sum_norm1 = vdupq_n_f16(0); + float16x8_t v_sum_norm2 = vdupq_n_f16(0); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + float16x8_t v_lhs = vld1q_f16((const float16_t *)lhs); + float16x8_t v_rhs = vld1q_f16((const float16_t *)rhs); + v_sum = vfmaq_f16(v_sum, v_lhs, v_rhs); + v_sum_norm1 = vfmaq_f16(v_sum_norm1, v_lhs, v_lhs); + v_sum_norm2 = vfmaq_f16(v_sum_norm2, v_rhs, v_rhs); + } + if (last >= last_aligned + 4) { + float16x8_t v_lhs = vcombine_f16(vld1_f16((const float16_t *)lhs), + vreinterpret_f16_u64(vdup_n_u64(0ul))); + float16x8_t v_rhs = vcombine_f16(vld1_f16((const float16_t *)rhs), + vreinterpret_f16_u64(vdup_n_u64(0ul))); + v_sum = vfmaq_f16(v_sum, v_lhs, v_rhs); + v_sum_norm1 = vfmaq_f16(v_sum_norm1, v_lhs, v_lhs); + v_sum_norm2 = vfmaq_f16(v_sum_norm2, v_rhs, v_rhs); + lhs += 4; + rhs += 4; + } + + float result = HorizontalAdd_FP16_NEON(v_sum); + float norm1 = HorizontalAdd_FP16_NEON(v_sum_norm1); + float norm2 = HorizontalAdd_FP16_NEON(v_sum_norm2); + + switch (last - lhs) { + case 3: + FMA_FP16_GENERAL(lhs[2], rhs[2], result, norm1, norm2); + /* FALLTHRU */ + case 2: + FMA_FP16_GENERAL(lhs[1], rhs[1], result, norm1, norm2); + /* FALLTHRU */ + case 1: + FMA_FP16_GENERAL(lhs[0], rhs[0], result, norm1, norm2); + } + *sql = norm1; + *sqr = norm2; + return result; +} +#else +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormFp16NEON(const Float16 *lhs, const Float16 *rhs, + size_t size, float *sql, float *sqr) { + const Float16 *last = lhs + size; + const Float16 *last_aligned = lhs + ((size >> 3) << 3); + float32x4_t v_sum_0 = vdupq_n_f32(0); + float32x4_t v_sum_1 = vdupq_n_f32(0); + float32x4_t v_sum_norm1 = vdupq_n_f32(0); + float32x4_t v_sum_norm2 = vdupq_n_f32(0); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + float16x8_t v_lhs = vld1q_f16((const float16_t *)lhs); + float16x8_t v_rhs = vld1q_f16((const float16_t *)rhs); + float32x4_t v_lhs_0 = vcvt_f32_f16(vget_low_f16(v_lhs)); + float32x4_t v_rhs_0 = vcvt_f32_f16(vget_low_f16(v_rhs)); + float32x4_t v_lhs_1 = vcvt_high_f32_f16(v_lhs); + float32x4_t v_rhs_1 = vcvt_high_f32_f16(v_rhs); + v_sum_0 = vfmaq_f32(v_sum_0, v_lhs_0, v_rhs_0); + v_sum_1 = vfmaq_f32(v_sum_1, v_lhs_1, v_rhs_1); + v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_0, v_lhs_0); + v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_1, v_lhs_1); + v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_0, v_rhs_0); + v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_1, v_rhs_1); + } + if (last >= last_aligned + 4) { + float32x4_t v_lhs_0 = vcvt_f32_f16(vld1_f16((const float16_t *)lhs)); + float32x4_t v_rhs_0 = vcvt_f32_f16(vld1_f16((const float16_t *)rhs)); + v_sum_0 = vfmaq_f32(v_sum_0, v_lhs_0, v_rhs_0); + v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_0, v_lhs_0); + v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_0, v_rhs_0); + lhs += 4; + rhs += 4; + } + + float result = vaddvq_f32(vaddq_f32(v_sum_0, v_sum_1)); + float norm1 = vaddvq_f32(v_sum_norm1); + float norm2 = vaddvq_f32(v_sum_norm2); + switch (last - lhs) { + case 3: + FMA_FP16_GENERAL(lhs[2], rhs[2], result, norm1, norm2); + /* FALLTHRU */ + case 2: + FMA_FP16_GENERAL(lhs[1], rhs[1], result, norm1, norm2); + /* FALLTHRU */ + case 1: + FMA_FP16_GENERAL(lhs[0], rhs[0], result, norm1, norm2); + } + *sql = norm1; + *sqr = norm2; + return result; +} + +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC + +float MipsEuclideanDistanceSphericalInjectionFp16NEON(const Float16 *lhs, + const Float16 *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormFp16NEON(lhs, rhs, size, &u2, &v2); + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp16NEON( + const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormFp16NEON(lhs, rhs, size, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} +#endif // __ARM_NEON && __aarch64__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32.cc new file mode 100644 index 0000000..e69de29 diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc new file mode 100644 index 0000000..331e342 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc @@ -0,0 +1,159 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE__) +float InnerProductAndSquaredNormFp32SSE(const float *lhs, const float *rhs, + size_t size, float *sql, float *sqr); +#endif + +#if defined(__AVX__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormFp32AVX(const float *lhs, const float *rhs, + size_t size, float *sql, float *sqr) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 4) << 4); + + __m256 ymm_sum_0 = _mm256_setzero_ps(); + __m256 ymm_sum_1 = _mm256_setzero_ps(); + __m256 ymm_sum_norm1 = _mm256_setzero_ps(); + __m256 ymm_sum_norm2 = _mm256_setzero_ps(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m256 ymm_lhs_0 = _mm256_load_ps(lhs + 0); + __m256 ymm_lhs_1 = _mm256_load_ps(lhs + 8); + __m256 ymm_rhs_0 = _mm256_load_ps(rhs + 0); + __m256 ymm_rhs_1 = _mm256_load_ps(rhs + 8); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2); + } + + if (last >= last_aligned + 8) { + __m256 ymm_lhs_0 = _mm256_load_ps(lhs); + __m256 ymm_rhs_0 = _mm256_load_ps(rhs); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + lhs += 8; + rhs += 8; + } + } else { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m256 ymm_lhs_0 = _mm256_loadu_ps(lhs + 0); + __m256 ymm_lhs_1 = _mm256_loadu_ps(lhs + 8); + __m256 ymm_rhs_0 = _mm256_loadu_ps(rhs + 0); + __m256 ymm_rhs_1 = _mm256_loadu_ps(rhs + 8); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2); + } + + if (last >= last_aligned + 8) { + __m256 ymm_lhs_0 = _mm256_loadu_ps(lhs); + __m256 ymm_rhs_0 = _mm256_loadu_ps(rhs); + ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + lhs += 8; + rhs += 8; + } + } + float result = HorizontalAdd_FP32_V256(_mm256_add_ps(ymm_sum_0, ymm_sum_1)); + float norm1 = HorizontalAdd_FP32_V256(ymm_sum_norm1); + float norm2 = HorizontalAdd_FP32_V256(ymm_sum_norm2); + + switch (last - lhs) { + case 7: + FMA_FP32_GENERAL(lhs[6], rhs[6], result, norm1, norm2) + /* FALLTHRU */ + case 6: + FMA_FP32_GENERAL(lhs[5], rhs[5], result, norm1, norm2) + /* FALLTHRU */ + case 5: + FMA_FP32_GENERAL(lhs[4], rhs[4], result, norm1, norm2) + /* FALLTHRU */ + case 4: + FMA_FP32_GENERAL(lhs[3], rhs[3], result, norm1, norm2) + /* FALLTHRU */ + case 3: + FMA_FP32_GENERAL(lhs[2], rhs[2], result, norm1, norm2) + /* FALLTHRU */ + case 2: + FMA_FP32_GENERAL(lhs[1], rhs[1], result, norm1, norm2) + /* FALLTHRU */ + case 1: + FMA_FP32_GENERAL(lhs[0], rhs[0], result, norm1, norm2) + } + *sql = norm1; + *sqr = norm2; + return result; +} + +float MipsEuclideanDistanceSphericalInjectionFp32AVX(const float *lhs, + const float *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + if (size > 7) { + sum = InnerProductAndSquaredNormFp32AVX(lhs, rhs, size, &u2, &v2); + } else { + sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2); + } + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX( + const float *lhs, const float *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + if (size > 7) { + sum = InnerProductAndSquaredNormFp32AVX(lhs, rhs, size, &u2, &v2); + } else { + sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2); + } + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} +#endif // __AVX__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc new file mode 100644 index 0000000..b5fffd9 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc @@ -0,0 +1,155 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE__) +float InnerProductAndSquaredNormFp32SSE(const float *lhs, const float *rhs, + size_t size, float *sql, float *sqr); +#endif + +#if defined(__AVX__) +float InnerProductAndSquaredNormFp32AVX(const float *lhs, const float *rhs, + size_t size, float *sql, float *sqr); +#endif + +#if defined(__AVX512F__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormFp32AVX512(const float *lhs, const float *rhs, + size_t size, float *sql, + float *sqr) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 5) << 5); + + __m512 zmm_sum_0 = _mm512_setzero_ps(); + __m512 zmm_sum_1 = _mm512_setzero_ps(); + __m512 zmm_sum_norm1 = _mm512_setzero_ps(); + __m512 zmm_sum_norm2 = _mm512_setzero_ps(); + + if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m512 zmm_lhs_0 = _mm512_load_ps(lhs + 0); + __m512 zmm_lhs_1 = _mm512_load_ps(lhs + 16); + __m512 zmm_rhs_0 = _mm512_load_ps(rhs + 0); + __m512 zmm_rhs_1 = _mm512_load_ps(rhs + 16); + FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0) + FMA_FP32_AVX512(zmm_lhs_1, zmm_rhs_1, zmm_sum_1) + FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_lhs_1, zmm_lhs_1, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2) + FMA_FP32_AVX512(zmm_rhs_1, zmm_rhs_1, zmm_sum_norm2) + } + + if (last >= last_aligned + 16) { + __m512 zmm_lhs_0 = _mm512_load_ps(lhs); + __m512 zmm_rhs_0 = _mm512_load_ps(rhs); + FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0) + FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2) + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m512 zmm_lhs_0 = _mm512_loadu_ps(lhs + 0); + __m512 zmm_lhs_1 = _mm512_loadu_ps(lhs + 16); + __m512 zmm_rhs_0 = _mm512_loadu_ps(rhs + 0); + __m512 zmm_rhs_1 = _mm512_loadu_ps(rhs + 16); + FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0) + FMA_FP32_AVX512(zmm_lhs_1, zmm_rhs_1, zmm_sum_1) + FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_lhs_1, zmm_lhs_1, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2) + FMA_FP32_AVX512(zmm_rhs_1, zmm_rhs_1, zmm_sum_norm2) + } + + if (last >= last_aligned + 16) { + __m512 zmm_lhs_0 = _mm512_loadu_ps(lhs); + __m512 zmm_rhs_0 = _mm512_loadu_ps(rhs); + FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0) + FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1) + FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2) + lhs += 16; + rhs += 16; + } + } + + zmm_sum_0 = _mm512_add_ps(zmm_sum_0, zmm_sum_1); + if (lhs != last) { + __mmask16 mask = (__mmask16)((1 << (last - lhs)) - 1); + __m512 zmm_undefined = _mm512_undefined_ps(); + __m512 zmm_lhs_0 = _mm512_mask_loadu_ps(zmm_undefined, mask, lhs); + __m512 zmm_rhs_0 = _mm512_mask_loadu_ps(zmm_undefined, mask, rhs); + FMA_MASK_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0, mask); + FMA_MASK_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1, mask); + FMA_MASK_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2, mask); + } + + *sql = HorizontalAdd_FP32_V512(zmm_sum_norm1); + *sqr = HorizontalAdd_FP32_V512(zmm_sum_norm2); + return HorizontalAdd_FP32_V512(zmm_sum_0); +} + +float MipsEuclideanDistanceSphericalInjectionFp32AVX512(const float *lhs, + const float *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + if (size > 15) { + sum = InnerProductAndSquaredNormFp32AVX512(lhs, rhs, size, &u2, &v2); + } else if (size > 7) { + sum = InnerProductAndSquaredNormFp32AVX(lhs, rhs, size, &u2, &v2); + } else { + sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2); + } + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX512( + const float *lhs, const float *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + if (size > 15) { + sum = InnerProductAndSquaredNormFp32AVX512(lhs, rhs, size, &u2, &v2); + } else if (size > 7) { + sum = InnerProductAndSquaredNormFp32AVX(lhs, rhs, size, &u2, &v2); + } else { + sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2); + } + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} +#endif // __AVX512F__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc new file mode 100644 index 0000000..f48626a --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc @@ -0,0 +1,176 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) +float InnerProductAndSquaredNormFp32NEON(const float *lhs, const float *rhs, + size_t size, float *sql, float *sqr); +#endif + +#if defined(__AVX512F__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX512( + const float *lhs, const float *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionFp32AVX512(const float *lhs, + const float *rhs, + size_t size, float e2); +#endif + +#if defined(__AVX__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX( + const float *lhs, const float *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionFp32AVX(const float *lhs, + const float *rhs, + size_t size, float e2); +#endif + +#if defined(__SSE__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp32SSE( + const float *lhs, const float *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionFp32SSE(const float *lhs, + const float *rhs, + size_t size, float e2); +#endif + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp32Scalar( + const float *p, const float *q, size_t dim, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionFp32Scalar(const float *p, + const float *q, + size_t dim, float e2); + +float MipsInnerProductSparseInSegment(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value); + +//! Compute the distance between matrix and query by SphericalInjection +void MipsSquaredEuclideanDistanceMatrix::Compute( + const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) { +#if __ARM_NEON + float u2{0.0f}; + float v2{0.0f}; + float sum = InnerProductAndSquaredNormFp32NEON(p, q, dim, &u2, &v2); + + *out = ComputeSphericalInjection(sum, u2, v2, e2); + return; +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = MipsEuclideanDistanceSphericalInjectionFp32AVX512(p, q, dim, e2); + return; + } +#endif //__AVX512F__ +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = MipsEuclideanDistanceSphericalInjectionFp32AVX(p, q, dim, e2); + return; + } +#endif // __AVX__ +#if defined(__SSE__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) { + *out = MipsEuclideanDistanceSphericalInjectionFp32SSE(p, q, dim, e2); + return; + } +#endif // __SSE__ + *out = MipsEuclideanDistanceSphericalInjectionFp32Scalar(p, q, dim, e2); + return; +#endif //__ARM_NEON +} + +//! Compute the distance between matrix and query by RepeatedQuadraticInjection +void MipsSquaredEuclideanDistanceMatrix::Compute( + const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2, + float *out) { +#if defined(__ARM_NEON) + float u2{0.0f}; + float v2{0.0f}; + float sum = InnerProductAndSquaredNormFp32NEON(p, q, dim, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + *out = sum; + return; +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX512(p, q, dim, + m, e2); + return; + } +#endif //__AVX512F__ +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX(p, q, dim, m, + e2); + return; + } +#endif // __AVX__ + +#if defined(__SSE__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp32SSE(p, q, dim, m, + e2); + return; + } +#endif //__SSE__ + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp32Scalar(p, q, dim, m, + e2); + + return; +#endif //__ARM_NEON +} + +// Sparse +#if defined(__SSE4_1__) +float MipsInnerProductSparseInSegmentSSE(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value); +#endif + +template <> +float MipsSquaredEuclideanSparseDistanceMatrix:: + ComputeInnerProductSparseInSegment(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const ValueType *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const ValueType *q_sparse_value) { +#if defined(__SSE4_1__) + return MipsInnerProductSparseInSegmentSSE(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +#else + return MipsInnerProductSparseInSegment(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +#endif +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc new file mode 100644 index 0000000..6491f22 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc @@ -0,0 +1,77 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__ARM_NEON) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormFp32NEON(const float *lhs, const float *rhs, + size_t size, float *sql, float *sqr) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 3) << 3); + + float32x4_t v_sum_0 = vdupq_n_f32(0); + float32x4_t v_sum_1 = vdupq_n_f32(0); + float32x4_t v_sum_norm1 = vdupq_n_f32(0); + float32x4_t v_sum_norm2 = vdupq_n_f32(0); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + float32x4_t v_lhs_0 = vld1q_f32(lhs + 0); + float32x4_t v_lhs_1 = vld1q_f32(lhs + 4); + float32x4_t v_rhs_0 = vld1q_f32(rhs + 0); + float32x4_t v_rhs_1 = vld1q_f32(rhs + 4); + v_sum_0 = vfmaq_f32(v_sum_0, v_lhs_0, v_rhs_0); + v_sum_1 = vfmaq_f32(v_sum_1, v_lhs_1, v_rhs_1); + v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_0, v_lhs_0); + v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_1, v_lhs_1); + v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_0, v_rhs_0); + v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_1, v_rhs_1); + } + if (last >= last_aligned + 4) { + float32x4_t v_lhs_0 = vld1q_f32(lhs); + float32x4_t v_rhs_0 = vld1q_f32(rhs); + v_sum_0 = vfmaq_f32(v_sum_0, v_lhs_0, v_rhs_0); + v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_0, v_lhs_0); + v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_0, v_rhs_0); + lhs += 4; + rhs += 4; + } + + float result = vaddvq_f32(vaddq_f32(v_sum_0, v_sum_1)); + float norm1 = vaddvq_f32(v_sum_norm1); + float norm2 = vaddvq_f32(v_sum_norm2); + switch (last - lhs) { + case 3: + FMA_FP32_GENERAL(lhs[2], rhs[2], result, norm1, norm2) + /* FALLTHRU */ + case 2: + FMA_FP32_GENERAL(lhs[1], rhs[1], result, norm1, norm2) + /* FALLTHRU */ + case 1: + FMA_FP32_GENERAL(lhs[0], rhs[0], result, norm1, norm2) + } + *sql = norm1; + *sqr = norm2; + return result; +} + +#endif //__ARM_NEON + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc new file mode 100644 index 0000000..8f1650d --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc @@ -0,0 +1,368 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_fp32.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormFp32SSE(const float *lhs, const float *rhs, + size_t size, float *sql, float *sqr) { + const float *last = lhs + size; + const float *last_aligned = lhs + ((size >> 3) << 3); + + __m128 xmm_sum = _mm_setzero_ps(); + __m128 xmm_sum_norm1 = _mm_setzero_ps(); + __m128 xmm_sum_norm2 = _mm_setzero_ps(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m128 xmm_lhs_0 = _mm_load_ps(lhs + 0); + __m128 xmm_lhs_1 = _mm_load_ps(lhs + 4); + __m128 xmm_rhs_0 = _mm_load_ps(rhs + 0); + __m128 xmm_rhs_1 = _mm_load_ps(rhs + 4); + xmm_sum = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum); + xmm_sum = _mm_fmadd_ps(xmm_lhs_1, xmm_rhs_1, xmm_sum); + xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1); + xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1); + xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2); + xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2); + } + + if (last >= last_aligned + 4) { + __m128 xmm_lhs_0 = _mm_load_ps(lhs); + __m128 xmm_rhs_0 = _mm_load_ps(rhs); + xmm_sum = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum); + xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1); + xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2); + lhs += 4; + rhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m128 xmm_lhs_0 = _mm_loadu_ps(lhs + 0); + __m128 xmm_lhs_1 = _mm_loadu_ps(lhs + 4); + __m128 xmm_rhs_0 = _mm_loadu_ps(rhs + 0); + __m128 xmm_rhs_1 = _mm_loadu_ps(rhs + 4); + xmm_sum = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum); + xmm_sum = _mm_fmadd_ps(xmm_lhs_1, xmm_rhs_1, xmm_sum); + xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1); + xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1); + xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2); + xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2); + } + + if (last >= last_aligned + 4) { + __m128 xmm_lhs_0 = _mm_loadu_ps(lhs); + __m128 xmm_rhs_0 = _mm_loadu_ps(rhs); + xmm_sum = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum); + xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1); + xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2); + lhs += 4; + rhs += 4; + } + } + float result = HorizontalAdd_FP32_V128(xmm_sum); + float norm1 = HorizontalAdd_FP32_V128(xmm_sum_norm1); + float norm2 = HorizontalAdd_FP32_V128(xmm_sum_norm2); + + switch (last - lhs) { + case 3: + FMA_FP32_GENERAL(lhs[2], rhs[2], result, norm1, norm2) + /* FALLTHRU */ + case 2: + FMA_FP32_GENERAL(lhs[1], rhs[1], result, norm1, norm2) + /* FALLTHRU */ + case 1: + FMA_FP32_GENERAL(lhs[0], rhs[0], result, norm1, norm2) + } + *sql = norm1; + *sqr = norm2; + return result; +} + +float MipsEuclideanDistanceSphericalInjectionFp32SSE(const float *lhs, + const float *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2); + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp32SSE( + const float *lhs, const float *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} + +#endif // __SSE__ + +// #if 1 +#if defined(__SSE4_1__) +const static __m128i SHUFFLE_MASK16[16] = { + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, -127, -127, -127, -127), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 7, 6, 5, 4, 3, + 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, 15, 14, 13, 12), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 3, 2, 1, 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 7, 6, 5, 4), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12, + 11, 10, 9, 8), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 3, 2, 1, + 0), + _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, + 4), + _mm_set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0), +}; + +constexpr uint32_t MAX_SPARSE_BUFFER_LENGTH = 65536; + +float MipsInnerProductSparseInSegmentSSE(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value) { + float sum = 0.0f; + + // size_t alloc_size = 0; + + size_t i1 = 0, i2 = 0; + size_t end1 = m_sparse_count / 8 * 8; + size_t end2 = q_sparse_count / 8 * 8; + + // std::vector mem1; + // std::vector mem2; + + float fixed_buffer_1[MAX_SPARSE_BUFFER_LENGTH]; + float fixed_buffer_2[MAX_SPARSE_BUFFER_LENGTH]; + + float *val_start_1 = fixed_buffer_1; + float *val_start_2 = fixed_buffer_2; + + // uint32_t max_count = std::max(m_sparse_count, q_sparse_count); + + // if (MAX_SPARSE_BUFFER_LENGTH < max_count) { + // mem1.reserve(max_count); + // mem2.reserve(max_count); + + // val_start_1 = mem1.data(); + // val_start_2 = mem2.data(); + // } + + float *val_1 = val_start_1; + float *val_2 = val_start_2; + + if (i1 < end1 && i2 < end2) { + while (m_sparse_index[i1 + 7] < q_sparse_index[i2]) { + i1 += 8; + if (i1 >= end1) goto do_scalar; + } + + while (q_sparse_index[i2 + 7] < m_sparse_index[i1]) { + i2 += 8; + if (i2 >= end2) goto do_scalar; + } + + __m128i mm_index_m = + _mm_loadu_si128(reinterpret_cast(&m_sparse_index[i1])); + __m128i mm_index_q = + _mm_loadu_si128(reinterpret_cast(&q_sparse_index[i2])); + + while (true) { +#ifdef DEBUG_PRINT + std::cout << "index 1: " << std::endl; + print_data16(&mm_index_m); + + std::cout << "index 2: " << std::endl; + print_data16(&mm_index_q); +#endif + + __m128i mm_cmp_res = + _mm_cmpistrm(mm_index_q, mm_index_m, + _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK); + +#ifdef DEBUG_PRINT + std::cout << "cmp res: " << std::endl; + print_data16(&mm_cmp_res); +#endif + + int r = _mm_extract_epi32(mm_cmp_res, 0); + + if (r) { + int r1 = r & 15; + + __m128i v = _mm_loadu_si128( + reinterpret_cast(&m_sparse_value[i1])); + __m128 vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r1])); + + _mm_storeu_ps(val_1, vs); + val_1 += _mm_popcnt_u32(r1); + + int r2 = (r >> 4) & 15; + v = _mm_loadu_si128( + reinterpret_cast(&m_sparse_value[i1 + 4])); + vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r2])); + _mm_storeu_ps(val_1, vs); + val_1 += _mm_popcnt_u32(r2); + + mm_cmp_res = _mm_cmpistrm( + mm_index_m, mm_index_q, + _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK); + r = _mm_extract_epi32(mm_cmp_res, 0); + + r1 = r & 15; + + v = _mm_loadu_si128( + reinterpret_cast(&q_sparse_value[i2])); + vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r1])); + _mm_storeu_ps(val_2, vs); + val_2 += _mm_popcnt_u32(r1); + + r2 = (r >> 4) & 15; + v = _mm_loadu_si128( + reinterpret_cast(&q_sparse_value[i2 + 4])); + vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r2])); + _mm_storeu_ps(val_2, vs); + val_2 += _mm_popcnt_u32(r2); + } + + const uint16_t id1_max = m_sparse_index[i1 + 7]; + + if (id1_max <= q_sparse_index[i2 + 7]) { + i1 += 8; + if (i1 >= end1) goto do_scalar; + mm_index_m = _mm_loadu_si128( + reinterpret_cast(&m_sparse_index[i1])); + } + + if (id1_max >= q_sparse_index[i2 + 7]) { + i2 += 8; + if (i2 >= end2) goto do_scalar; + mm_index_q = _mm_loadu_si128( + reinterpret_cast(&q_sparse_index[i2])); + } + } + } + +do_scalar: + while (i1 < m_sparse_count && i2 < q_sparse_count) { + if (m_sparse_index[i1] == q_sparse_index[i2]) { + *val_1++ = m_sparse_value[i1]; + *val_2++ = q_sparse_value[i2]; + + ++i1; + ++i2; + } else if (m_sparse_index[i1] < q_sparse_index[i2]) { + ++i1; + } else { + ++i2; + } + } + + size_t res_num = val_1 - val_start_1; + + // if (res_num != val_2 - val_start_2) { + // std::cerr << "size mismatch!" << std::endl; + // } + + size_t res_num4 = res_num / 4 * 4; + + if (res_num4) { + __m128 sum128 = _mm_set1_ps(0); + + for (size_t k = 0; k < res_num4; k += 4) { + sum128 = _mm_add_ps(sum128, _mm_mul_ps(_mm_loadu_ps(val_start_1 + k), + _mm_loadu_ps(val_start_2 + k))); + } + + alignas(16) float tmp_res[4]; + _mm_store_ps(tmp_res, sum128); + sum += (tmp_res[0] + tmp_res[1] + tmp_res[2] + tmp_res[3]); + } + + for (size_t k = res_num4; k < res_num; ++k) + sum += val_start_1[k] * val_start_2[k]; + + return sum; +} +#else +float MipsInnerProductSparseInSegment(uint32_t m_sparse_count, + const uint16_t *m_sparse_index, + const float *m_sparse_value, + uint32_t q_sparse_count, + const uint16_t *q_sparse_index, + const float *q_sparse_value) { + float sum = 0.0f; + + size_t m_i = 0; + size_t q_i = 0; + while (m_i < m_sparse_count && q_i < q_sparse_count) { + if (m_sparse_index[m_i] == q_sparse_index[q_i]) { + sum += m_sparse_value[m_i] * q_sparse_value[q_i]; + + ++m_i; + ++q_i; + } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) { + ++m_i; + } else { + ++q_i; + } + } + + return sum; +} +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc b/src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc new file mode 100644 index 0000000..ba50c21 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc @@ -0,0 +1,172 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int8.i" +#include "distance_matrix_mips_utility.i" +#include "inner_product_matrix.h" +#include "mips_euclidean_distance_matrix.h" +#include "norm_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX2__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormInt4AVX2(const uint8_t *lhs, const uint8_t *rhs, + size_t size, float *sql, float *sqr) { + const uint8_t *last = lhs + size; + const uint8_t *last_aligned = lhs + ((size >> 5) << 5); + __m256i ymm_sum_0 = _mm256_setzero_si256(); + __m256i ymm_sum_1 = _mm256_setzero_si256(); + __m256i ymm_sum_norm1 = _mm256_setzero_si256(); + __m256i ymm_sum_norm2 = _mm256_setzero_si256(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m256i ymm_lhs = _mm256_load_si256((const __m256i *)(lhs)); + __m256i ymm_rhs = _mm256_load_si256((const __m256i *)(rhs)); + FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum_0, ymm_sum1, ymm_sum_norm1, + ymm_sum_norm2) + } + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs); + __m128i xmm_sum = _mm_setzero_si128(); + __m128i xmm_sum_norm1 = _mm_setzero_si128(); + __m128i xmm_sum_norm2 = _mm_setzero_si128(); + FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum, xmm_sum_norm1, xmm_sum_norm2) + ymm_sum_0 = _mm256_add_epi32( + _mm256_set_m128i(_mm_setzero_si128(), xmm_sum), ymm_sum_0); + ymm_sum_norm1 = _mm256_add_epi32( + _mm256_set_m128i(_mm_setzero_si128(), xmm_sum_norm1), ymm_sum_norm1); + ymm_sum_norm2 = _mm256_add_epi32( + _mm256_set_m128i(_mm_setzero_si128(), xmm_sum_norm2), ymm_sum_norm2); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)(lhs)); + __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)(rhs)); + FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum_0, ymm_sum1, ymm_sum_norm1, + ymm_sum_norm2) + } + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + __m128i xmm_sum = _mm_setzero_si128(); + __m128i xmm_sum_norm1 = _mm_setzero_si128(); + __m128i xmm_sum_norm2 = _mm_setzero_si128(); + FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum, xmm_sum_norm1, xmm_sum_norm2) + ymm_sum_0 = _mm256_add_epi32( + _mm256_set_m128i(_mm_setzero_si128(), xmm_sum), ymm_sum_0); + ymm_sum_norm1 = _mm256_add_epi32( + _mm256_set_m128i(_mm_setzero_si128(), xmm_sum_norm1), ymm_sum_norm1); + ymm_sum_norm2 = _mm256_add_epi32( + _mm256_set_m128i(_mm_setzero_si128(), xmm_sum_norm2), ymm_sum_norm2); + lhs += 16; + rhs += 16; + } + } + float result = static_cast( + HorizontalAdd_INT32_V256(_mm256_add_epi32(ymm_sum_0, ymm_sum_1))); + float norm1 = static_cast(HorizontalAdd_INT32_V256(ymm_sum_norm1)); + float norm2 = static_cast(HorizontalAdd_INT32_V256(ymm_sum_norm2)); + + switch (last - lhs) { + case 15: + FMA_INT4_GENERAL(lhs[14], rhs[14], result, norm1, norm2) + /* FALLTHRU */ + case 14: + FMA_INT4_GENERAL(lhs[13], rhs[13], result, norm1, norm2) + /* FALLTHRU */ + case 13: + FMA_INT4_GENERAL(lhs[12], rhs[12], result, norm1, norm2) + /* FALLTHRU */ + case 12: + FMA_INT4_GENERAL(lhs[11], rhs[11], result, norm1, norm2) + /* FALLTHRU */ + case 11: + FMA_INT4_GENERAL(lhs[10], rhs[10], result, norm1, norm2) + /* FALLTHRU */ + case 10: + FMA_INT4_GENERAL(lhs[9], rhs[9], result, norm1, norm2) + /* FALLTHRU */ + case 9: + FMA_INT4_GENERAL(lhs[8], rhs[8], result, norm1, norm2) + /* FALLTHRU */ + case 8: + FMA_INT4_GENERAL(lhs[7], rhs[7], result, norm1, norm2) + /* FALLTHRU */ + case 7: + FMA_INT4_GENERAL(lhs[6], rhs[6], result, norm1, norm2) + /* FALLTHRU */ + case 6: + FMA_INT4_GENERAL(lhs[5], rhs[5], result, norm1, norm2) + /* FALLTHRU */ + case 5: + FMA_INT4_GENERAL(lhs[4], rhs[4], result, norm1, norm2) + /* FALLTHRU */ + case 4: + FMA_INT4_GENERAL(lhs[3], rhs[3], result, norm1, norm2) + /* FALLTHRU */ + case 3: + FMA_INT4_GENERAL(lhs[2], rhs[2], result, norm1, norm2) + /* FALLTHRU */ + case 2: + FMA_INT4_GENERAL(lhs[1], rhs[1], result, norm1, norm2) + /* FALLTHRU */ + case 1: + FMA_INT4_GENERAL(lhs[0], rhs[0], result, norm1, norm2) + } + *sql = norm1; + *sqr = norm2; + return result; +} + +float MipsEuclideanDistanceSphericalInjectionInt4AVX2(const uint8_t *lhs, + const uint8_t *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormInt4AVX2(lhs, rhs, size >> 1, &u2, &v2); + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt4AVX2( + const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormInt4AVX2(lhs, rhs, size >> 1, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} +#endif // __AVX2__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc new file mode 100644 index 0000000..86b6183 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc @@ -0,0 +1,90 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "inner_product_matrix.h" +#include "mips_euclidean_distance_matrix.h" +#include "norm_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX2__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt4AVX2( + const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionInt4AVX2(const uint8_t *lhs, + const uint8_t *rhs, + size_t size, float e2); +#endif + +#if defined(__SSE4_1__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt4SSE( + const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionInt4SSE(const uint8_t *lhs, + const uint8_t *rhs, + size_t size, float e2); +#endif + +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt4Scalar( + const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionInt4Scalar(const uint8_t *lhs, + const uint8_t *rhs, + size_t size, float e2); + +//! Compute the distance between matrix and query by SphericalInjection +void MipsSquaredEuclideanDistanceMatrix::Compute( + const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = MipsEuclideanDistanceSphericalInjectionInt4AVX2(p, q, dim, e2); + return; + } +#endif + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = MipsEuclideanDistanceSphericalInjectionInt4SSE(p, q, dim, e2); + return; + } +#endif + + *out = MipsEuclideanDistanceSphericalInjectionInt4Scalar(p, q, dim, e2); +} + +//! Compute the distance between matrix and query by RepeatedQuadraticInjection +void MipsSquaredEuclideanDistanceMatrix::Compute( + const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2, + float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt4AVX2(p, q, dim, m, + e2); + return; + } +#endif + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt4SSE(p, q, dim, m, + e2); + return; + } +#endif + + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt4Scalar(p, q, dim, m, + e2); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc b/src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc new file mode 100644 index 0000000..464071a --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc @@ -0,0 +1,137 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int8.i" +#include "distance_matrix_mips_utility.i" +#include "inner_product_matrix.h" +#include "mips_euclidean_distance_matrix.h" +#include "norm_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE4_1__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormInt4SSE(const uint8_t *lhs, const uint8_t *rhs, + size_t size, float *sql, float *sqr) { + const uint8_t *last = lhs + size; + const uint8_t *last_aligned = lhs + ((size >> 4) << 4); + __m128i xmm_sum = _mm_setzero_si128(); + __m128i xmm_sum_norm1 = _mm_setzero_si128(); + __m128i xmm_sum_norm2 = _mm_setzero_si128(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)(lhs)); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)(rhs)); + FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum, xmm_sum_norm1, xmm_sum_norm2) + } + } else { + for (; lhs != last_aligned; lhs += 16, rhs += 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)(lhs)); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)(rhs)); + FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum, xmm_sum_norm1, xmm_sum_norm2) + } + } + float result = static_cast(HorizontalAdd_INT32_V128(xmm_sum)); + float norm1 = static_cast(HorizontalAdd_INT32_V128(xmm_sum_norm1)); + float norm2 = static_cast(HorizontalAdd_INT32_V128(xmm_sum_norm2)); + + switch (last - lhs) { + case 15: + FMA_INT4_GENERAL(lhs[14], rhs[14], result, norm1, norm2) + /* FALLTHRU */ + case 14: + FMA_INT4_GENERAL(lhs[13], rhs[13], result, norm1, norm2) + /* FALLTHRU */ + case 13: + FMA_INT4_GENERAL(lhs[12], rhs[12], result, norm1, norm2) + /* FALLTHRU */ + case 12: + FMA_INT4_GENERAL(lhs[11], rhs[11], result, norm1, norm2) + /* FALLTHRU */ + case 11: + FMA_INT4_GENERAL(lhs[10], rhs[10], result, norm1, norm2) + /* FALLTHRU */ + case 10: + FMA_INT4_GENERAL(lhs[9], rhs[9], result, norm1, norm2) + /* FALLTHRU */ + case 9: + FMA_INT4_GENERAL(lhs[8], rhs[8], result, norm1, norm2) + /* FALLTHRU */ + case 8: + FMA_INT4_GENERAL(lhs[7], rhs[7], result, norm1, norm2) + /* FALLTHRU */ + case 7: + FMA_INT4_GENERAL(lhs[6], rhs[6], result, norm1, norm2) + /* FALLTHRU */ + case 6: + FMA_INT4_GENERAL(lhs[5], rhs[5], result, norm1, norm2) + /* FALLTHRU */ + case 5: + FMA_INT4_GENERAL(lhs[4], rhs[4], result, norm1, norm2) + /* FALLTHRU */ + case 4: + FMA_INT4_GENERAL(lhs[3], rhs[3], result, norm1, norm2) + /* FALLTHRU */ + case 3: + FMA_INT4_GENERAL(lhs[2], rhs[2], result, norm1, norm2) + /* FALLTHRU */ + case 2: + FMA_INT4_GENERAL(lhs[1], rhs[1], result, norm1, norm2) + /* FALLTHRU */ + case 1: + FMA_INT4_GENERAL(lhs[0], rhs[0], result, norm1, norm2) + } + *sql = norm1; + *sqr = norm2; + return result; +} + +float MipsEuclideanDistanceSphericalInjectionInt4SSE(const uint8_t *lhs, + const uint8_t *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormInt4SSE(lhs, rhs, size >> 1, &u2, &v2); + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt4SSE( + const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormInt4SSE(lhs, rhs, size >> 1, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} + +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc b/src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc new file mode 100644 index 0000000..0f95cd2 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc @@ -0,0 +1,191 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int8.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX2__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormInt8AVX2(const int8_t *lhs, const int8_t *rhs, + size_t size, float *sql, float *sqr) { + const int8_t *last = lhs + size; + const int8_t *last_aligned = lhs + ((size >> 6) << 6); + + __m256i ymm_sum_0 = _mm256_setzero_si256(); + __m256i ymm_sum_1 = _mm256_setzero_si256(); + __m256i ymm_sum_norm1 = _mm256_setzero_si256(); + __m256i ymm_sum_norm2 = _mm256_setzero_si256(); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + __m256i ymm_lhs_0 = _mm256_load_si256((const __m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_load_si256((const __m256i *)(lhs + 32)); + __m256i ymm_rhs_0 = _mm256_load_si256((const __m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_load_si256((const __m256i *)(rhs + 32)); + FMA_INT8_AVX(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + FMA_INT8_AVX(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); + FMA_INT8_AVX(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + FMA_INT8_AVX(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1); + FMA_INT8_AVX(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + FMA_INT8_AVX(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2); + } + + if (last >= last_aligned + 32) { + __m256i ymm_lhs = _mm256_load_si256((const __m256i *)lhs); + __m256i ymm_rhs = _mm256_load_si256((const __m256i *)rhs); + FMA_INT8_AVX(ymm_lhs, ymm_rhs, ymm_sum_0); + FMA_INT8_AVX(ymm_lhs, ymm_lhs, ymm_sum_norm1); + FMA_INT8_AVX(ymm_rhs, ymm_rhs, ymm_sum_norm2); + lhs += 32; + rhs += 32; + } + + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs); + FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_rhs, ymm_sum_0); + FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_lhs, ymm_sum_norm1); + FMA_INT8_AVX_SSE_HYBRID(xmm_rhs, xmm_rhs, ymm_sum_norm2); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 64, rhs += 64) { + __m256i ymm_lhs_0 = _mm256_loadu_si256((const __m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_loadu_si256((const __m256i *)(lhs + 32)); + __m256i ymm_rhs_0 = _mm256_loadu_si256((const __m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_loadu_si256((const __m256i *)(rhs + 32)); + FMA_INT8_AVX(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); + FMA_INT8_AVX(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); + FMA_INT8_AVX(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); + FMA_INT8_AVX(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1); + FMA_INT8_AVX(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); + FMA_INT8_AVX(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2); + } + + if (last >= last_aligned + 32) { + __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)lhs); + __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)rhs); + FMA_INT8_AVX(ymm_lhs, ymm_rhs, ymm_sum_0); + FMA_INT8_AVX(ymm_lhs, ymm_lhs, ymm_sum_norm1); + FMA_INT8_AVX(ymm_rhs, ymm_rhs, ymm_sum_norm2); + lhs += 32; + rhs += 32; + } + + if (last >= lhs + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_rhs, ymm_sum_0); + FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_lhs, ymm_sum_norm1); + FMA_INT8_AVX_SSE_HYBRID(xmm_rhs, xmm_rhs, ymm_sum_norm2); + lhs += 16; + rhs += 16; + } + } + float result = static_cast( + HorizontalAdd_INT32_V256(_mm256_add_epi32(ymm_sum_0, ymm_sum_1))); + float norm1 = static_cast(HorizontalAdd_INT32_V256(ymm_sum_norm1)); + float norm2 = static_cast(HorizontalAdd_INT32_V256(ymm_sum_norm2)); + + switch (last - lhs) { + case 15: + FMA_INT8_GENERAL(lhs[14], rhs[14], result, norm1, norm2) + /* FALLTHRU */ + case 14: + FMA_INT8_GENERAL(lhs[13], rhs[13], result, norm1, norm2) + /* FALLTHRU */ + case 13: + FMA_INT8_GENERAL(lhs[12], rhs[12], result, norm1, norm2) + /* FALLTHRU */ + case 12: + FMA_INT8_GENERAL(lhs[11], rhs[11], result, norm1, norm2) + /* FALLTHRU */ + case 11: + FMA_INT8_GENERAL(lhs[10], rhs[10], result, norm1, norm2) + /* FALLTHRU */ + case 10: + FMA_INT8_GENERAL(lhs[9], rhs[9], result, norm1, norm2) + /* FALLTHRU */ + case 9: + FMA_INT8_GENERAL(lhs[8], rhs[8], result, norm1, norm2) + /* FALLTHRU */ + case 8: + FMA_INT8_GENERAL(lhs[7], rhs[7], result, norm1, norm2) + /* FALLTHRU */ + case 7: + FMA_INT8_GENERAL(lhs[6], rhs[6], result, norm1, norm2) + /* FALLTHRU */ + case 6: + FMA_INT8_GENERAL(lhs[5], rhs[5], result, norm1, norm2) + /* FALLTHRU */ + case 5: + FMA_INT8_GENERAL(lhs[4], rhs[4], result, norm1, norm2) + /* FALLTHRU */ + case 4: + FMA_INT8_GENERAL(lhs[3], rhs[3], result, norm1, norm2) + /* FALLTHRU */ + case 3: + FMA_INT8_GENERAL(lhs[2], rhs[2], result, norm1, norm2) + /* FALLTHRU */ + case 2: + FMA_INT8_GENERAL(lhs[1], rhs[1], result, norm1, norm2) + /* FALLTHRU */ + case 1: + FMA_INT8_GENERAL(lhs[0], rhs[0], result, norm1, norm2) + } + *sql = norm1; + *sqr = norm2; + return result; +} + +float MipsEuclideanDistanceSphericalInjectionInt8AVX2(const int8_t *lhs, + const int8_t *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormInt8AVX2(lhs, rhs, size, &u2, &v2); + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt8AVX2( + const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormInt8AVX2(lhs, rhs, size, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} +#endif // __AVX2__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc new file mode 100644 index 0000000..f0f7449 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc @@ -0,0 +1,87 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__AVX2__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt8AVX2( + const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionInt8AVX2(const int8_t *lhs, + const int8_t *rhs, + size_t size, float e2); +#endif + +#if defined(__SSE4_1__) +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt8SSE( + const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionInt8SSE(const int8_t *lhs, + const int8_t *rhs, + size_t size, float e2); +#endif + +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt8Scalar( + const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2); +float MipsEuclideanDistanceSphericalInjectionInt8Scalar(const int8_t *lhs, + const int8_t *rhs, + size_t size, float e2); + +//! Compute the distance between matrix and query by SphericalInjection +void MipsSquaredEuclideanDistanceMatrix::Compute( + const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = MipsEuclideanDistanceSphericalInjectionInt8AVX2(p, q, dim, e2); + return; + } +#endif + +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = MipsEuclideanDistanceSphericalInjectionInt8SSE(p, q, dim, e2); + return; + } +#endif //__SSE4_1__ + + *out = MipsEuclideanDistanceSphericalInjectionInt8Scalar(p, q, dim, e2); +} + +//! Compute the distance between matrix and query by RepeatedQuadraticInjection +void MipsSquaredEuclideanDistanceMatrix::Compute( + const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2, + float *out) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt8AVX2(p, q, dim, m, + e2); + return; + } +#endif +#if defined(__SSE4_1__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) { + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt8SSE(p, q, dim, m, + e2); + return; + } +#endif //__SSE4_1__ + + *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt8Scalar(p, q, dim, m, + e2); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc b/src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc new file mode 100644 index 0000000..86a19ea --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc @@ -0,0 +1,170 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "distance_matrix_accum_int8.i" +#include "distance_matrix_mips_utility.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE4_1__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +float InnerProductAndSquaredNormInt8SSE(const int8_t *lhs, const int8_t *rhs, + size_t size, float *sql, float *sqr) { + const int8_t *last = lhs + size; + const int8_t *last_aligned = lhs + ((size >> 5) << 5); + + __m128i xmm_sum = _mm_setzero_si128(); + __m128i xmm_sum_norm1 = _mm_setzero_si128(); + __m128i xmm_sum_norm2 = _mm_setzero_si128(); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m128i xmm_lhs_0 = _mm_load_si128((const __m128i *)(lhs + 0)); + __m128i xmm_lhs_1 = _mm_load_si128((const __m128i *)(lhs + 16)); + __m128i xmm_rhs_0 = _mm_load_si128((const __m128i *)(rhs + 0)); + __m128i xmm_rhs_1 = _mm_load_si128((const __m128i *)(rhs + 16)); + FMA_INT8_SSE(xmm_lhs_0, xmm_rhs_0, xmm_sum); + FMA_INT8_SSE(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1); + FMA_INT8_SSE(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2); + FMA_INT8_SSE(xmm_lhs_1, xmm_rhs_1, xmm_sum); + FMA_INT8_SSE(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1); + FMA_INT8_SSE(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2); + } + + if (last >= last_aligned + 16) { + __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs); + FMA_INT8_SSE(xmm_lhs, xmm_rhs, xmm_sum); + FMA_INT8_SSE(xmm_lhs, xmm_lhs, xmm_sum_norm1); + FMA_INT8_SSE(xmm_rhs, xmm_rhs, xmm_sum_norm2); + lhs += 16; + rhs += 16; + } + } else { + for (; lhs != last_aligned; lhs += 32, rhs += 32) { + __m128i xmm_lhs_0 = _mm_loadu_si128((const __m128i *)(lhs + 0)); + __m128i xmm_lhs_1 = _mm_loadu_si128((const __m128i *)(lhs + 16)); + __m128i xmm_rhs_0 = _mm_loadu_si128((const __m128i *)(rhs + 0)); + __m128i xmm_rhs_1 = _mm_loadu_si128((const __m128i *)(rhs + 16)); + FMA_INT8_SSE(xmm_lhs_0, xmm_rhs_0, xmm_sum); + FMA_INT8_SSE(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1); + FMA_INT8_SSE(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2); + FMA_INT8_SSE(xmm_lhs_1, xmm_rhs_1, xmm_sum); + FMA_INT8_SSE(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1); + FMA_INT8_SSE(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2); + } + + if (last >= last_aligned + 16) { + __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs); + __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs); + FMA_INT8_SSE(xmm_lhs, xmm_rhs, xmm_sum); + FMA_INT8_SSE(xmm_lhs, xmm_lhs, xmm_sum_norm1); + FMA_INT8_SSE(xmm_rhs, xmm_rhs, xmm_sum_norm2); + lhs += 16; + rhs += 16; + } + } + float result = static_cast(HorizontalAdd_INT32_V128(xmm_sum)); + float norm1 = static_cast(HorizontalAdd_INT32_V128(xmm_sum_norm1)); + float norm2 = static_cast(HorizontalAdd_INT32_V128(xmm_sum_norm2)); + + switch (last - lhs) { + case 15: + FMA_INT8_GENERAL(lhs[14], rhs[14], result, norm1, norm2) + /* FALLTHRU */ + case 14: + FMA_INT8_GENERAL(lhs[13], rhs[13], result, norm1, norm2) + /* FALLTHRU */ + case 13: + FMA_INT8_GENERAL(lhs[12], rhs[12], result, norm1, norm2) + /* FALLTHRU */ + case 12: + FMA_INT8_GENERAL(lhs[11], rhs[11], result, norm1, norm2) + /* FALLTHRU */ + case 11: + FMA_INT8_GENERAL(lhs[10], rhs[10], result, norm1, norm2) + /* FALLTHRU */ + case 10: + FMA_INT8_GENERAL(lhs[9], rhs[9], result, norm1, norm2) + /* FALLTHRU */ + case 9: + FMA_INT8_GENERAL(lhs[8], rhs[8], result, norm1, norm2) + /* FALLTHRU */ + case 8: + FMA_INT8_GENERAL(lhs[7], rhs[7], result, norm1, norm2) + /* FALLTHRU */ + case 7: + FMA_INT8_GENERAL(lhs[6], rhs[6], result, norm1, norm2) + /* FALLTHRU */ + case 6: + FMA_INT8_GENERAL(lhs[5], rhs[5], result, norm1, norm2) + /* FALLTHRU */ + case 5: + FMA_INT8_GENERAL(lhs[4], rhs[4], result, norm1, norm2) + /* FALLTHRU */ + case 4: + FMA_INT8_GENERAL(lhs[3], rhs[3], result, norm1, norm2) + /* FALLTHRU */ + case 3: + FMA_INT8_GENERAL(lhs[2], rhs[2], result, norm1, norm2) + /* FALLTHRU */ + case 2: + FMA_INT8_GENERAL(lhs[1], rhs[1], result, norm1, norm2) + /* FALLTHRU */ + case 1: + FMA_INT8_GENERAL(lhs[0], rhs[0], result, norm1, norm2) + } + *sql = norm1; + *sqr = norm2; + return result; +} + +float MipsEuclideanDistanceSphericalInjectionInt8SSE(const int8_t *lhs, + const int8_t *rhs, + size_t size, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormInt8SSE(lhs, rhs, size, &u2, &v2); + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt8SSE( + const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2) { + float u2{0.0f}; + float v2{0.0f}; + float sum{0.0f}; + + sum = InnerProductAndSquaredNormInt8SSE(lhs, rhs, size, &u2, &v2); + + sum = e2 * (u2 + v2 - 2 * sum); + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} + +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_scalar.cc b/src/ailego/math/mips_euclidean_distance_matrix_scalar.cc new file mode 100644 index 0000000..06f39da --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_scalar.cc @@ -0,0 +1,172 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "distance_utility.h" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { +//-------------------------------------------------- +// Dense +//-------------------------------------------------- +// Compute the distance between matrix and query by SphericalInjection +template +inline float MipsEuclideanDistanceSphericalInjectionScalar(const T *p, + const T *q, + size_t dim, + float e2) { + ailego_assert(p && q && dim); + + float sum = 0.0; + float u2 = 0.0; + float v2 = 0.0; + for (size_t i = 0; i < dim; ++i) { + u2 += p[i] * p[i]; + v2 += q[i] * q[i]; + sum += static_cast(p[i] * q[i]); + } + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +// Compute the distance between matrix and query by RepeatedQuadraticInjection +template +inline float MipsEuclideanDistanceRepeatedQuadraticInjectionScalar( + const T *p, const T *q, size_t dim, size_t m, float e2) { + ailego_assert(p && q && dim); + + float sum = 0.0; + float u2 = 0.0; + float v2 = 0.0; + for (size_t i = 0; i < dim; ++i) { + u2 += p[i] * p[i]; + v2 += q[i] * q[i]; + sum += MathHelper::SquaredDifference(p[i], q[i]); + } + + sum *= e2; + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} + +/*! Mips Squared Euclidean Distance Matrix (INT4, M=1, N=1) + */ +//! Calculate sum of squared values +static inline float Squared(uint8_t v) { + return static_cast(((int8_t)(v << 4) >> 4) * ((int8_t)(v << 4) >> 4) + + ((int8_t)(v & 0xf0) >> 4) * + ((int8_t)(v & 0xf0) >> 4)); +} + +// Compute the distance between matrix and query by SphericalInjection +float MipsEuclideanDistanceSphericalInjectionInt4Scalar(const uint8_t *p, + const uint8_t *q, + size_t dim, float e2) { + ailego_assert(p && q && dim && !(dim & 1)); + + float sum = 0.0; + float u2 = 0.0; + float v2 = 0.0; + for (size_t i = 0; i < (dim >> 1); ++i) { + const uint8_t p_val = p[i]; + const uint8_t q_val = q[i]; + u2 += Squared(p_val); + v2 += Squared(q_val); + sum += Int4MulTable[((p_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] + + Int4MulTable[((p_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)]; + } + + return ComputeSphericalInjection(sum, u2, v2, e2); +} + +// Compute the distance between matrix and query by RepeatedQuadraticInjection +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt4Scalar( + const uint8_t *p, const uint8_t *q, size_t dim, size_t m, float e2) { + ailego_assert(p && q && dim && !(dim & 1)); + + float sum = 0.0; + float u2 = 0.0; + float v2 = 0.0; + for (size_t i = 0; i < (dim >> 1); ++i) { + const uint8_t p_val = p[i]; + const uint8_t q_val = q[i]; + u2 += Squared(p_val); + v2 += Squared(q_val); + sum += Int4SquaredDiffTable[((p_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] + + Int4SquaredDiffTable[((p_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)]; + } + sum *= e2; + u2 *= e2; + v2 *= e2; + for (size_t i = 0; i < m; ++i) { + sum += (u2 - v2) * (u2 - v2); + u2 = u2 * u2; + v2 = v2 * v2; + } + + return sum; +} + +float MipsEuclideanDistanceSphericalInjectionInt8Scalar(const int8_t *p, + const int8_t *q, + size_t dim, float e2) { + return MipsEuclideanDistanceSphericalInjectionScalar(p, q, dim, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionInt8Scalar( + const int8_t *p, const int8_t *q, size_t dim, size_t m, float e2) { + return MipsEuclideanDistanceRepeatedQuadraticInjectionScalar( + p, q, dim, m, e2); +} + +float MipsEuclideanDistanceSphericalInjectionFp16Scalar( + const ailego::Float16 *p, const ailego::Float16 *q, size_t dim, float e2) { + return MipsEuclideanDistanceSphericalInjectionScalar( + p, q, dim, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp16Scalar( + const ailego::Float16 *p, const ailego::Float16 *q, size_t dim, size_t m, + float e2) { + return MipsEuclideanDistanceRepeatedQuadraticInjectionScalar( + p, q, dim, m, e2); +} + +float MipsEuclideanDistanceSphericalInjectionFp32Scalar(const float *p, + const float *q, + size_t dim, float e2) { + return MipsEuclideanDistanceSphericalInjectionScalar(p, q, dim, e2); +} + +float MipsEuclideanDistanceRepeatedQuadraticInjectionFp32Scalar( + const float *p, const float *q, size_t dim, size_t m, float e2) { + return MipsEuclideanDistanceRepeatedQuadraticInjectionScalar(p, q, dim, + m, e2); +} + + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/norm1_matrix.h b/src/ailego/math/norm1_matrix.h new file mode 100644 index 0000000..7e8d9cb --- /dev/null +++ b/src/ailego/math/norm1_matrix.h @@ -0,0 +1,147 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! L1-Norm Matrix + */ +template +struct Norm1Matrix; + +/*! L1-Norm Matrix + */ +template +struct Norm1Matrix::value && + sizeof(T) >= 2 && M >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && out); + + const ValueType *m_end = m + dim * M; + if (m != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) = MathHelper::Absolute(m[i]); + } + m += M; + } + while (m != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) += MathHelper::Absolute(m[i]); + } + m += M; + } + } +}; + +/*! L1-Norm Matrix (INT8) + */ +template +struct Norm1Matrix= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *m_end = m_it + (dim >> 2) * M; + + if (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) = Absolute(m_it[i]); + } + m_it += M; + } + while (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) += Absolute(m_it[i]); + } + m_it += M; + } + } + + protected: + //! Calculate sum of absolute values + static inline float Absolute(uint32_t v) { + return static_cast( + MathHelper::Absolute((int8_t)(v >> 0)) + + MathHelper::Absolute((int8_t)(v >> 8)) + + MathHelper::Absolute((int8_t)(v >> 16)) + + MathHelper::Absolute((int8_t)(v >> 24))); + } +}; + +/*! L1-Norm Matrix (M=1) + */ +template +struct Norm1Matrix< + T, 1, typename std::enable_if::value>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && out); + + const ValueType *m_end = m + dim; + if (m != m_end) { + *out = MathHelper::Absolute(*m++); + } + while (m != m_end) { + *out += MathHelper::Absolute(*m++); + } + } +}; + +#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__)) +/*! L1-Norm Matrix (FP32, M=1) + */ +template <> +struct Norm1Matrix { + //! Type of value + using ValueType = float; + + //! Compute the L1-norm of vectors + static void Compute(const ValueType *m, size_t dim, float *out); +}; +#endif // __SSE__ || (__ARM_NEON && __aarch64__) + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +/*! L1-Norm Matrix (FP16, M=1) + */ +template <> +struct Norm1Matrix { + //! Type of value + using ValueType = Float16; + + //! Compute the L1-norm of vectors + static void Compute(const ValueType *m, size_t dim, float *out); +}; +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/norm1_matrix_fp16.cc b/src/ailego/math/norm1_matrix_fp16.cc new file mode 100644 index 0000000..e75b3e0 --- /dev/null +++ b/src/ailego/math/norm1_matrix_fp16.cc @@ -0,0 +1,90 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "ailego/internal/cpu_features.h" +#include "norm1_matrix.h" +#include "norm_matrix_fp16.i" + +namespace zvec { +namespace ailego { + +#define NORM_FP32_STEP_GENERAL SA_FP32_GENERAL +#define NORM_FP32_STEP_SSE SA_FP32_SSE +#define NORM_FP32_STEP_AVX SA_FP32_AVX +#define NORM_FP32_STEP_AVX512 SA_FP32_AVX512 +#define NORM_FP32_STEP_NEON SA_FP32_NEON +#define NORM_FP16_STEP_GENERAL SA_FP16_GENERAL +#define NORM_FP16_STEP_NEON SA_FP16_NEON + +#if defined(__SSE__) +static const __m128 ABS_MASK_FP32_SSE = + _mm_castsi128_ps(_mm_set1_epi32(0x7fffffffu)); +#endif // __SSE__ + +#if defined(__AVX__) +static const __m256 ABS_MASK_FP32_AVX = + _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffffu)); +#endif // __AVX__ + +#if defined(__AVX512F__) +static const __m512 ABS_MASK_FP32_AVX512 = + _mm512_castsi512_ps(_mm512_set1_epi32(0x7fffffffu)); +#endif // __AVX512F__ + +//! Calculate sum of absolute (GENERAL) +#define SA_FP32_GENERAL(m, sum) sum += FastAbs(m); + +//! Calculate sum of absolute (SSE) +#define SA_FP32_SSE(xmm_m, xmm_sum) \ + xmm_sum = _mm_add_ps(_mm_and_ps(xmm_m, ABS_MASK_FP32_SSE), xmm_sum); + +//! Calculate sum of absolute (AVX) +#define SA_FP32_AVX(ymm_m, ymm_sum) \ + ymm_sum = _mm256_add_ps(_mm256_and_ps(ymm_m, ABS_MASK_FP32_AVX), ymm_sum); + +//! Calculate sum of absolute (AVX512) +#define SA_FP32_AVX512(zmm_m, zmm_sum) \ + zmm_sum = _mm512_add_ps(_mm512_and_ps(zmm_m, ABS_MASK_FP32_AVX512), zmm_sum); + +//! Calculate sum of absolute (NEON) +#define SA_FP32_NEON(v_m, v_sum) v_sum = vaddq_f32(vabsq_f32(v_m), v_sum); + +//! Calculate sum of absolute (GENERAL) +#define SA_FP16_GENERAL(m, sum) sum += Float16::Absolute(m); + +//! Calculate sum of absolute (NEON) +#define SA_FP16_NEON(v_m, v_sum) v_sum = vaddq_f16(vabsq_f16(v_m), v_sum); + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +//! Compute the L1-norm of vectors (FP16, M=1) +void Norm1Matrix::Compute(const ValueType *m, size_t dim, + float *out) { +#if defined(__ARM_NEON) + NORM_FP16_1_NEON(m, dim, out, ) +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + NORM_FP16_1_AVX512(m, dim, out, ) + return; + } +#endif + NORM_FP16_1_AVX(m, dim, out, ) +#endif +} +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/norm1_matrix_fp32.cc b/src/ailego/math/norm1_matrix_fp32.cc new file mode 100644 index 0000000..2e72791 --- /dev/null +++ b/src/ailego/math/norm1_matrix_fp32.cc @@ -0,0 +1,84 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "ailego/internal/cpu_features.h" +#include "norm1_matrix.h" +#include "norm_matrix_fp32.i" + +namespace zvec { +namespace ailego { + +#define NORM_FP32_STEP_GENERAL SA_FP32_GENERAL +#define NORM_FP32_STEP_SSE SA_FP32_SSE +#define NORM_FP32_STEP_AVX SA_FP32_AVX +#define NORM_FP32_STEP_AVX512 SA_FP32_AVX512 +#define NORM_FP32_STEP_NEON SA_FP32_NEON + +#if defined(__SSE__) +#define ABS_MASK_FP32_SSE _mm_castsi128_ps(_mm_set1_epi32(0x7fffffffu)) +#endif // __SSE__ + +#if defined(__AVX__) +#define ABS_MASK_FP32_AVX _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffffu)) +#endif // __AVX__ + +#if defined(__AVX512F__) +#define ABS_MASK_FP32_AVX512 _mm512_castsi512_ps(_mm512_set1_epi32(0x7fffffffu)) +#endif // __AVX512F__ + +//! Calculate sum of absolute (GENERAL) +#define SA_FP32_GENERAL(m, sum) sum += FastAbs(m); + +//! Calculate sum of absolute (SSE) +#define SA_FP32_SSE(xmm_m, xmm_sum) \ + xmm_sum = _mm_add_ps(_mm_and_ps(xmm_m, ABS_MASK_FP32_SSE), xmm_sum); + +//! Calculate sum of absolute (AVX) +#define SA_FP32_AVX(ymm_m, ymm_sum) \ + ymm_sum = _mm256_add_ps(_mm256_and_ps(ymm_m, ABS_MASK_FP32_AVX), ymm_sum); + +//! Calculate sum of absolute (AVX512) +#define SA_FP32_AVX512(zmm_m, zmm_sum) \ + zmm_sum = _mm512_add_ps(_mm512_and_ps(zmm_m, ABS_MASK_FP32_AVX512), zmm_sum); + +//! Calculate sum of absolute (NEON) +#define SA_FP32_NEON(v_m, v_sum) v_sum = vaddq_f32(vabsq_f32(v_m), v_sum); + +#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__)) +//! Compute the L1-norm of vectors (FP32, M=1) +void Norm1Matrix::Compute(const ValueType *m, size_t dim, + float *out) { +#if defined(__ARM_NEON) + NORM_FP32_1_NEON(m, dim, out, ) +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + NORM_FP32_1_AVX512(m, dim, out, ) + return; + } +#endif +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + NORM_FP32_1_AVX(m, dim, out, ) + return; + } +#endif + NORM_FP32_1_SSE(m, dim, out, ) +#endif +} +#endif // __SSE__ || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/norm2_matrix.h b/src/ailego/math/norm2_matrix.h new file mode 100644 index 0000000..3c90514 --- /dev/null +++ b/src/ailego/math/norm2_matrix.h @@ -0,0 +1,424 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! L2-Norm Matrix + */ +template +struct Norm2Matrix; + +/*! L2-Norm Matrix + */ +template +struct Norm2Matrix::value && + sizeof(T) >= 2 && M >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && out); + + const ValueType *m_end = m + dim * M; + if (m != m_end) { + for (size_t i = 0; i < M; ++i) { + ValueType v = m[i]; + *(out + i) = static_cast(v * v); + } + m += M; + } + while (m != m_end) { + for (size_t i = 0; i < M; ++i) { + ValueType v = m[i]; + *(out + i) += static_cast(v * v); + } + m += M; + } + for (size_t i = 0; i < M; ++i) { + float v = *out; + *out++ = std::sqrt(v); + } + } +}; + +/*! L2-Norm Matrix (INT8) + */ +template +struct Norm2Matrix= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *m_end = m_it + (dim >> 2) * M; + + if (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) = Squared(m_it[i]); + } + m_it += M; + } + while (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) += Squared(m_it[i]); + } + m_it += M; + } + for (size_t i = 0; i < M; ++i) { + float v = *out; + *out++ = std::sqrt(v); + } + } + + protected: + //! Calculate sum of squared values + static inline float Squared(uint32_t v) { + return static_cast((int8_t)(v >> 0) * (int8_t)(v >> 0) + + (int8_t)(v >> 8) * (int8_t)(v >> 8) + + (int8_t)(v >> 16) * (int8_t)(v >> 16) + + (int8_t)(v >> 24) * (int8_t)(v >> 24)); + } +}; + +/*! L2-Norm Matrix (M=1) + */ +template +struct Norm2Matrix< + T, 1, typename std::enable_if::value>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && out); + + const ValueType *m_end = m + dim; + if (m != m_end) { + ValueType v = *m++; + *out = static_cast(v * v); + } + while (m != m_end) { + ValueType v = *m++; + *out += static_cast(v * v); + } + *out = std::sqrt(*out); + } +}; + +/*! L2-Norm Matrix (M=1, INT4) + */ +template <> +struct Norm2Matrix { + //! Type of value + using ValueType = uint8_t; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && !(dim & 1) && dim && out); + + const uint8_t *m_end = m + (dim >> 1); + float square = 0.0f; + while (m != m_end) { + square += Squared(*m++); + } + *out = std::sqrt(square); + } + + protected: + //! Calculate sum of squared values + static inline float Squared(uint8_t v) { + return static_cast( + ((int8_t)(v << 4) >> 4) * ((int8_t)(v << 4) >> 4) + + ((int8_t)(v & 0xf0) >> 4) * ((int8_t)(v & 0xf0) >> 4)); + } +}; + +/*! L2-Norm Matrix (INT4) + */ +template +struct Norm2Matrix= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && !(dim & 7) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *m_end = m_it + (dim >> 3) * M; + + if (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) = Squared(m_it[i]); + } + m_it += M; + } + while (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) += Squared(m_it[i]); + } + m_it += M; + } + for (size_t i = 0; i < M; ++i) { + float v = *out; + *out++ = std::sqrt(v); + } + } + + protected: + //! Calculate sum of squared values + static inline float Squared(uint32_t u) { + float sum = 0.0f; + for (size_t i = 0; i < 32; i += 8) { + uint8_t v = (uint8_t)(u >> i); + int8_t lo = (int8_t)(v << 4) >> 4; + int8_t hi = (int8_t)(v & 0xf0) >> 4; + sum += hi * hi + lo * lo; + } + return sum; + } +}; + +/*! Squared L2-Norm Matrix + */ +template +struct SquaredNorm2Matrix; + +/*! Squared L2-Norm Matrix + */ +template +struct SquaredNorm2Matrix< + T, M, + typename std::enable_if::value && sizeof(T) >= 2 && + M >= 2>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && out); + + const ValueType *m_end = m + dim * M; + if (m != m_end) { + for (size_t i = 0; i < M; ++i) { + ValueType v = m[i]; + *(out + i) = static_cast(v * v); + } + m += M; + } + while (m != m_end) { + for (size_t i = 0; i < M; ++i) { + ValueType v = m[i]; + *(out + i) += static_cast(v * v); + } + m += M; + } + } +}; + +/*! Squared L2-Norm Matrix (INT8) + */ +template +struct SquaredNorm2Matrix= 2>::type> { + //! Type of value + using ValueType = int8_t; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && !(dim & 3) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *m_end = m_it + (dim >> 2) * M; + + if (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) = Squared(m_it[i]); + } + m_it += M; + } + while (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) += Squared(m_it[i]); + } + m_it += M; + } + } + + protected: + //! Calculate sum of squared values + static inline float Squared(uint32_t v) { + return static_cast((int8_t)(v >> 0) * (int8_t)(v >> 0) + + (int8_t)(v >> 8) * (int8_t)(v >> 8) + + (int8_t)(v >> 16) * (int8_t)(v >> 16) + + (int8_t)(v >> 24) * (int8_t)(v >> 24)); + } +}; + +/*! Squared L2-Norm Matrix (M=1) + */ +template +struct SquaredNorm2Matrix< + T, 1, typename std::enable_if::value>::type> { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && out); + + const ValueType *m_end = m + dim; + if (m != m_end) { + ValueType v = *m++; + *out = static_cast(v * v); + } + while (m != m_end) { + ValueType v = *m++; + *out += static_cast(v * v); + } + } +}; + +/*! L2-Norm Matrix (M=1, INT4) + */ +template <> +struct SquaredNorm2Matrix { + //! Type of value + using ValueType = uint8_t; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && !(dim & 1) && out); + + const uint8_t *m_end = m + (dim >> 1); + *out = 0.0f; + while (m != m_end) { + *out += Squared(*m++); + } + } + + protected: + //! Calculate sum of squared values + static inline float Squared(uint8_t v) { + return static_cast( + ((int8_t)(v << 4) >> 4) * ((int8_t)(v << 4) >> 4) + + ((int8_t)(v & 0xf0) >> 4) * ((int8_t)(v & 0xf0) >> 4)); + } +}; + +/*! Squared L2-Norm Matrix (INT4) + */ +template +struct SquaredNorm2Matrix= 2>::type> { + //! Type of value + using ValueType = uint8_t; + + //! Compute the norm of vectors + static inline void Compute(const ValueType *m, size_t dim, float *out) { + ailego_assert(m && dim && !(dim & 7) && out); + + const uint32_t *m_it = reinterpret_cast(m); + const uint32_t *m_end = m_it + (dim >> 3) * M; + + if (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) = Squared(m_it[i]); + } + m_it += M; + } + while (m_it != m_end) { + for (size_t i = 0; i < M; ++i) { + *(out + i) += Squared(m_it[i]); + } + m_it += M; + } + } + + protected: + //! Calculate sum of squared values + static inline float Squared(uint32_t u) { + float sum = 0.0f; + for (size_t i = 0; i < 32; i += 8) { + uint8_t v = (uint8_t)(u >> i); + int8_t lo = (int8_t)(v << 4) >> 4; + int8_t hi = (int8_t)(v & 0xf0) >> 4; + sum += hi * hi + lo * lo; + } + return sum; + } +}; + +#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__)) +/*! L2-Norm Matrix (FP32, M=1) + */ +template <> +struct Norm2Matrix { + //! Type of value + using ValueType = float; + + //! Compute the L2-norm of vectors + static void Compute(const ValueType *m, size_t dim, float *out); +}; + +/*! Squared L2-Norm Matrix (FP32, M=1) + */ +template <> +struct SquaredNorm2Matrix { + //! Type of value + using ValueType = float; + + //! Compute the squared L2-norm of vectors + static void Compute(const ValueType *m, size_t dim, float *out); +}; +#endif // __SSE__ || (__ARM_NEON && __aarch64__) + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +/*! L2-Norm Matrix (FP16, M=1) + */ +template <> +struct Norm2Matrix { + //! Type of value + using ValueType = Float16; + + //! Compute the L2-norm of vectors + static void Compute(const ValueType *m, size_t dim, float *out); +}; + +/*! Squared L2-Norm Matrix (FP16, M=1) + */ +template <> +struct SquaredNorm2Matrix { + //! Type of value + using ValueType = Float16; + + //! Compute the squared L2-norm of vectors + static void Compute(const ValueType *m, size_t dim, float *out); +}; +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/norm2_matrix_fp16.cc b/src/ailego/math/norm2_matrix_fp16.cc new file mode 100644 index 0000000..6bb8dd0 --- /dev/null +++ b/src/ailego/math/norm2_matrix_fp16.cc @@ -0,0 +1,91 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "ailego/internal/cpu_features.h" +#include "norm2_matrix.h" +#include "norm_matrix_fp16.i" + +namespace zvec { +namespace ailego { + +#define NORM_FP32_STEP_GENERAL SS_FP32_GENERAL +#define NORM_FP32_STEP_SSE SS_FP32_SSE +#define NORM_FP32_STEP_AVX SS_FP32_AVX +#define NORM_FP32_STEP_AVX512 SS_FP32_AVX512 +#define NORM_FP32_STEP_NEON SS_FP32_NEON +#define NORM_FP16_STEP_GENERAL SS_FP16_GENERAL +#define NORM_FP16_STEP_NEON SS_FP16_NEON + +//! Calculate sum of squared (GENERAL) +#define SS_FP32_GENERAL(m, sum) sum += (m) * (m); + +//! Calculate sum of squared (SSE) +#define SS_FP32_SSE(xmm_m, xmm_sum) \ + xmm_sum = _mm_fmadd_ps(xmm_m, xmm_m, xmm_sum); + +//! Calculate sum of squared (AVX) +#define SS_FP32_AVX(ymm_m, ymm_sum) \ + ymm_sum = _mm256_fmadd_ps(ymm_m, ymm_m, ymm_sum); + +//! Calculate sum of squared (AVX512) +#define SS_FP32_AVX512(zmm_m, zmm_sum) \ + zmm_sum = _mm512_fmadd_ps(zmm_m, zmm_m, zmm_sum); + +//! Calculate sum of squared (NEON) +#define SS_FP32_NEON(v_m, v_sum) v_sum = vfmaq_f32(v_sum, v_m, v_m); + +//! Calculate sum of squared (GENERAL) +#define SS_FP16_GENERAL(m, sum) sum += (m) * (m); + +//! Calculate sum of squared (NEON) +#define SS_FP16_NEON(v_m, v_sum) v_sum = vfmaq_f16(v_sum, v_m, v_m); + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +//! Compute the L2-norm of vectors (FP16, M=1) +void Norm2Matrix::Compute(const ValueType *m, size_t dim, + float *out) { +#if defined(__ARM_NEON) + NORM_FP16_1_NEON(m, dim, out, std::sqrt) +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + NORM_FP16_1_AVX512(m, dim, out, std::sqrt) + return; + } +#endif + NORM_FP16_1_AVX(m, dim, out, std::sqrt) +#endif +} + +//! Compute the L2-norm of vectors (FP16, M=1) +void SquaredNorm2Matrix::Compute(const ValueType *m, size_t dim, + float *out) { +#if defined(__ARM_NEON) + NORM_FP16_1_NEON(m, dim, out, ) +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + NORM_FP16_1_AVX512(m, dim, out, ) + return; + } +#endif + NORM_FP16_1_AVX(m, dim, out, ) +#endif +} +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/norm2_matrix_fp32.cc b/src/ailego/math/norm2_matrix_fp32.cc new file mode 100644 index 0000000..8cc76c1 --- /dev/null +++ b/src/ailego/math/norm2_matrix_fp32.cc @@ -0,0 +1,93 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "norm2_matrix.h" +#include "norm_matrix_fp32.i" + +namespace zvec { +namespace ailego { + +#define NORM_FP32_STEP_GENERAL SS_FP32_GENERAL +#define NORM_FP32_STEP_SSE SS_FP32_SSE +#define NORM_FP32_STEP_AVX SS_FP32_AVX +#define NORM_FP32_STEP_AVX512 SS_FP32_AVX512 +#define NORM_FP32_STEP_NEON SS_FP32_NEON + +//! Calculate sum of squared (GENERAL) +#define SS_FP32_GENERAL(m, sum) sum += (m) * (m); + +//! Calculate sum of squared (SSE) +#define SS_FP32_SSE(xmm_m, xmm_sum) \ + xmm_sum = _mm_fmadd_ps(xmm_m, xmm_m, xmm_sum); + +//! Calculate sum of squared (AVX) +#define SS_FP32_AVX(ymm_m, ymm_sum) \ + ymm_sum = _mm256_fmadd_ps(ymm_m, ymm_m, ymm_sum); + +//! Calculate sum of squared (AVX512) +#define SS_FP32_AVX512(zmm_m, zmm_sum) \ + zmm_sum = _mm512_fmadd_ps(zmm_m, zmm_m, zmm_sum); + +//! Calculate sum of squared (NEON) +#define SS_FP32_NEON(v_m, v_sum) v_sum = vfmaq_f32(v_sum, v_m, v_m); + +#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__)) +//! Compute the L2-norm of vectors (FP32, M=1) +void Norm2Matrix::Compute(const ValueType *m, size_t dim, + float *out) { +#if defined(__ARM_NEON) + NORM_FP32_1_NEON(m, dim, out, std::sqrt) +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + NORM_FP32_1_AVX512(m, dim, out, std::sqrt) + return; + } +#endif +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + NORM_FP32_1_AVX(m, dim, out, std::sqrt) + return; + } +#endif + NORM_FP32_1_SSE(m, dim, out, std::sqrt) +#endif +} + +//! Compute the squared L2-norm of vectors (FP32, M=1) +void SquaredNorm2Matrix::Compute(const ValueType *m, size_t dim, + float *out) { +#if defined(__ARM_NEON) + NORM_FP32_1_NEON(m, dim, out, ) +#else +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + NORM_FP32_1_AVX512(m, dim, out, ) + return; + } +#endif +#if defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + NORM_FP32_1_AVX(m, dim, out, ) + return; + } +#endif + NORM_FP32_1_SSE(m, dim, out, ) +#endif +} +#endif // __SSE__ || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/norm_matrix.h b/src/ailego/math/norm_matrix.h new file mode 100644 index 0000000..87642a2 --- /dev/null +++ b/src/ailego/math/norm_matrix.h @@ -0,0 +1,18 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "norm1_matrix.h" +#include "norm2_matrix.h" diff --git a/src/ailego/math/norm_matrix_fp16.i b/src/ailego/math/norm_matrix_fp16.i new file mode 100644 index 0000000..c1d3bd2 --- /dev/null +++ b/src/ailego/math/norm_matrix_fp16.i @@ -0,0 +1,219 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "matrix_define.i" +#include "matrix_utility.i" + +#if !defined(__FMA__) +#define _mm_fmadd_ps(a, b, c) _mm_add_ps(_mm_mul_ps((a), (b)), (c)) +#define _mm256_fmadd_ps(a, b, c) _mm256_add_ps(_mm256_mul_ps((a), (b)), (c)) +#endif // !__FMA__ + +//! Mask process of computing norm (FP16) +#define NORM_FP16_MASK_AVX(m, cnt, _RES) \ + switch (cnt) { \ + case 7: { \ + __m256 ymm_m = _mm256_cvtph_ps( \ + _mm_set_epi16(0, *((const short *)(m) + 6), \ + *((const short *)(m) + 5), *((const short *)(m) + 4), \ + *((const short *)(m) + 3), *((const short *)(m) + 2), \ + *((const short *)(m) + 1), *((const short *)(m)))); \ + NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \ + break; \ + } \ + case 6: { \ + __m256 ymm_m = _mm256_cvtph_ps(_mm_set_epi32(0, *((const int *)(m) + 2), \ + *((const int *)(m) + 1), \ + *((const int *)(m)))); \ + NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \ + break; \ + } \ + case 5: { \ + __m256 ymm_m = _mm256_cvtph_ps( \ + _mm_set_epi16(0, 0, 0, *((const short *)(m) + 4), \ + *((const short *)(m) + 3), *((const short *)(m) + 2), \ + *((const short *)(m) + 1), *((const short *)(m)))); \ + NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \ + break; \ + } \ + case 4: { \ + __m256 ymm_m = _mm256_cvtph_ps( \ + _mm_set_epi64x(0LL, *(const long long *)(m))); \ + NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \ + break; \ + } \ + case 3: { \ + __m256 ymm_m = _mm256_cvtph_ps( \ + _mm_set_epi16(0, 0, 0, 0, 0, *((const short *)(m) + 2), \ + *((const short *)(m) + 1), *((const short *)(m)))); \ + NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \ + break; \ + } \ + case 2: { \ + __m256 ymm_m = \ + _mm256_cvtph_ps(_mm_set_epi32(0, 0, 0, *((const int *)(m)))); \ + NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \ + break; \ + } \ + case 1: { \ + __m256 ymm_m = _mm256_cvtph_ps( \ + _mm_set_epi16(0, 0, 0, 0, 0, 0, 0, *((const short *)(m)))); \ + NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \ + break; \ + } \ + } + +//! Compute the norm of vectors (FP16, M=1) +#define NORM_FP16_1_AVX(m, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + const Float16 *last = m + dim; \ + const Float16 *last_aligned = m + ((dim >> 4) << 4); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; m != last_aligned; m += 16) { \ + __m256i ymm_mi = _mm256_load_si256((const __m256i *)m); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + NORM_FP32_STEP_AVX(ymm_m_0, ymm_sum_0_0) \ + NORM_FP32_STEP_AVX(ymm_m_1, ymm_sum_0_0) \ + } \ + if (last >= last_aligned + 8) { \ + __m256 ymm_m = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)m)); \ + NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \ + m += 8; \ + } \ + } else { \ + for (; m != last_aligned; m += 16) { \ + __m256i ymm_mi = _mm256_loadu_si256((const __m256i *)m); \ + __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \ + __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \ + NORM_FP32_STEP_AVX(ymm_m_0, ymm_sum_0_0) \ + NORM_FP32_STEP_AVX(ymm_m_1, ymm_sum_0_0) \ + } \ + if (last >= last_aligned + 8) { \ + __m256 ymm_m = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)m)); \ + NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \ + m += 8; \ + } \ + } \ + NORM_FP16_MASK_AVX(m, (last - m), ymm_sum) \ + *out = _NORM(HorizontalAdd_FP32_V256(ymm_sum_0_0)); + +//! Compute the norm of vectors (FP16, M=1) +#define NORM_FP16_1_AVX512(m, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m512, zmm_sum, _mm512_setzero_ps()) \ + const Float16 *last = m + dim; \ + const Float16 *last_aligned = m + ((dim >> 5) << 5); \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (; m != last_aligned; m += 32) { \ + __m512i zmm_mi = _mm512_load_si512((const __m512i *)m); \ + __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + NORM_FP32_STEP_AVX512(zmm_m_0, zmm_sum_0_0) \ + NORM_FP32_STEP_AVX512(zmm_m_1, zmm_sum_0_1) \ + } \ + if (last >= last_aligned + 16) { \ + __m512 zmm_m = _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)m)); \ + NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \ + m += 16; \ + } \ + } else { \ + for (; m != last_aligned; m += 32) { \ + __m512i zmm_mi = _mm512_loadu_si512((const __m512i *)m); \ + __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \ + __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \ + NORM_FP32_STEP_AVX512(zmm_m_0, zmm_sum_0_0) \ + NORM_FP32_STEP_AVX512(zmm_m_1, zmm_sum_0_1) \ + } \ + if (last >= last_aligned + 16) { \ + __m512 zmm_m = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)m)); \ + NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \ + m += 16; \ + } \ + } \ + float result = \ + HorizontalAdd_FP32_V512(_mm512_add_ps(zmm_sum_0_0, zmm_sum_0_1)); \ + if (m != last) { \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + if (last >= m + 8) { \ + __m256 ymm_m = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)m)); \ + NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \ + m += 8; \ + } \ + NORM_FP16_MASK_AVX(m, (last - m), ymm_sum) \ + result += HorizontalAdd_FP32_V256(ymm_sum_0_0); \ + } \ + *out = _NORM(result); + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) +//! Compute the norm of vectors (FP16, M=1) +#define NORM_FP16_1_NEON(m, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, float16x8_t, v_sum, vdupq_n_f16(0)) \ + const Float16 *last = m + dim; \ + const Float16 *last_aligned = m + ((dim >> 3) << 3); \ + for (; m != last_aligned; m += 8) { \ + float16x8_t v_m = vld1q_f16((const float16_t *)m); \ + NORM_FP16_STEP_NEON(v_m, v_sum_0_0) \ + } \ + if (last >= m + 4) { \ + float16x8_t v_m = vreinterpretq_f16_u64( \ + vld1q_lane_u64((const uint64_t *)m, vdupq_n_u64(0), 0)); \ + NORM_FP16_STEP_NEON(v_m, v_sum_0_0) \ + m += 4; \ + } \ + float result = vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(v_sum_0_0)), \ + vcvt_high_f32_f16(v_sum_0_0))); \ + switch (last - m) { \ + case 3: \ + NORM_FP16_STEP_GENERAL(m[2], result) \ + /* FALLTHRU */ \ + case 2: \ + NORM_FP16_STEP_GENERAL(m[1], result) \ + /* FALLTHRU */ \ + case 1: \ + NORM_FP16_STEP_GENERAL(m[0], result) \ + } \ + *out = _NORM(result); + +#else +//! Compute the norm of vectors (FP16, M=1) +#define NORM_FP16_1_NEON(m, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \ + const Float16 *last = m + dim; \ + const Float16 *last_aligned = m + ((dim >> 3) << 3); \ + for (; m != last_aligned; m += 8) { \ + float16x8_t v_m = vld1q_f16((const float16_t *)m); \ + float32x4_t v_n_0 = vcvt_f32_f16(vget_low_f16(v_m)); \ + float32x4_t v_n_1 = vcvt_high_f32_f16(v_m); \ + NORM_FP32_STEP_NEON(v_n_0, v_sum_0_0) \ + NORM_FP32_STEP_NEON(v_n_1, v_sum_0_1) \ + } \ + if (last >= m + 4) { \ + float32x4_t v_m = vcvt_f32_f16(vld1_f16((const float16_t *)m)); \ + NORM_FP32_STEP_NEON(v_m, v_sum_0_0) \ + m += 4; \ + } \ + float result = vaddvq_f32(vaddq_f32(v_sum_0_0, v_sum_0_1)); \ + switch (last - m) { \ + case 3: \ + NORM_FP16_STEP_GENERAL(m[2], result) \ + /* FALLTHRU */ \ + case 2: \ + NORM_FP16_STEP_GENERAL(m[1], result) \ + /* FALLTHRU */ \ + case 1: \ + NORM_FP16_STEP_GENERAL(m[0], result) \ + } \ + *out = _NORM(result); + +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC \ No newline at end of file diff --git a/src/ailego/math/norm_matrix_fp32.i b/src/ailego/math/norm_matrix_fp32.i new file mode 100644 index 0000000..ba2e4a6 --- /dev/null +++ b/src/ailego/math/norm_matrix_fp32.i @@ -0,0 +1,185 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "matrix_define.i" +#include "matrix_utility.i" + +#if !defined(__FMA__) +#define _mm_fmadd_ps(a, b, c) _mm_add_ps(_mm_mul_ps((a), (b)), (c)) +#define _mm256_fmadd_ps(a, b, c) _mm256_add_ps(_mm256_mul_ps((a), (b)), (c)) +#endif // !__FMA__ + +//! Mask process of computing norm (FP32) +#define NORM_FP32_MASK_SSE(m, cnt, _RES) \ + switch (cnt) { \ + case 3: { \ + __m128 xmm_m = _mm_set_ps(0.0f, m[2], m[1], m[0]); \ + NORM_FP32_STEP_SSE(xmm_m, _RES##_0_0) \ + break; \ + } \ + case 2: { \ + __m128 xmm_m = _mm_set_ps(0.0f, 0.0f, m[1], m[0]); \ + NORM_FP32_STEP_SSE(xmm_m, _RES##_0_0) \ + break; \ + } \ + case 1: { \ + __m128 xmm_m = _mm_set_ps(0.0f, 0.0f, 0.0f, m[0]); \ + NORM_FP32_STEP_SSE(xmm_m, _RES##_0_0) \ + break; \ + } \ + } + +//! Compute the norm of vectors (FP32, M=1) +#define NORM_FP32_1_SSE(m, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m128, xmm_sum, _mm_setzero_ps()) \ + const float *last = m + dim; \ + const float *last_aligned = m + ((dim >> 3) << 3); \ + if (((uintptr_t)m & 0xf) == 0) { \ + for (; m != last_aligned; m += 8) { \ + __m128 xmm_m_0 = _mm_load_ps(m + 0); \ + __m128 xmm_m_1 = _mm_load_ps(m + 4); \ + NORM_FP32_STEP_SSE(xmm_m_0, xmm_sum_0_0) \ + NORM_FP32_STEP_SSE(xmm_m_1, xmm_sum_0_0) \ + } \ + if (last >= last_aligned + 4) { \ + __m128 xmm_m = _mm_load_ps(m); \ + NORM_FP32_STEP_SSE(xmm_m, xmm_sum_0_0) \ + m += 4; \ + } \ + } else { \ + for (; m != last_aligned; m += 8) { \ + __m128 xmm_m_0 = _mm_loadu_ps(m + 0); \ + __m128 xmm_m_1 = _mm_loadu_ps(m + 4); \ + NORM_FP32_STEP_SSE(xmm_m_0, xmm_sum_0_0) \ + NORM_FP32_STEP_SSE(xmm_m_1, xmm_sum_0_0) \ + } \ + if (last >= last_aligned + 4) { \ + __m128 xmm_m = _mm_loadu_ps(m); \ + NORM_FP32_STEP_SSE(xmm_m, xmm_sum_0_0) \ + m += 4; \ + } \ + } \ + NORM_FP32_MASK_SSE(m, (last - m), xmm_sum) \ + *out = _NORM(HorizontalAdd_FP32_V128(xmm_sum_0_0)); + +//! Compute the norm of vectors (FP32, M=1) +#define NORM_FP32_1_AVX(m, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \ + const float *last = m + dim; \ + const float *last_aligned = m + ((dim >> 4) << 4); \ + if (((uintptr_t)m & 0x1f) == 0) { \ + for (; m != last_aligned; m += 16) { \ + __m256 ymm_m_0 = _mm256_load_ps(m + 0); \ + __m256 ymm_m_1 = _mm256_load_ps(m + 8); \ + NORM_FP32_STEP_AVX(ymm_m_0, ymm_sum_0_0) \ + NORM_FP32_STEP_AVX(ymm_m_1, ymm_sum_0_0) \ + } \ + if (last >= last_aligned + 8) { \ + __m256 ymm_m = _mm256_load_ps(m); \ + NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \ + m += 8; \ + } \ + } else { \ + for (; m != last_aligned; m += 16) { \ + __m256 ymm_m_0 = _mm256_loadu_ps(m + 0); \ + __m256 ymm_m_1 = _mm256_loadu_ps(m + 8); \ + NORM_FP32_STEP_AVX(ymm_m_0, ymm_sum_0_0) \ + NORM_FP32_STEP_AVX(ymm_m_1, ymm_sum_0_0) \ + } \ + if (last >= last_aligned + 8) { \ + __m256 ymm_m = _mm256_loadu_ps(m); \ + NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \ + m += 8; \ + } \ + } \ + float result = HorizontalAdd_FP32_V256(ymm_sum_0_0); \ + if (m != last) { \ + __m128 xmm_sum_0_0 = _mm_setzero_ps(); \ + if (last >= m + 4) { \ + __m128 xmm_m = _mm_loadu_ps(m); \ + NORM_FP32_STEP_SSE(xmm_m, xmm_sum_0_0) \ + m += 4; \ + } \ + NORM_FP32_MASK_SSE(m, (last - m), xmm_sum) \ + result += HorizontalAdd_FP32_V128(xmm_sum_0_0); \ + } \ + *out = _NORM(result); + +//! Compute the norm of vectors (FP32, M=1) +#define NORM_FP32_1_AVX512(m, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, __m512, zmm_sum, _mm512_setzero_ps()) \ + const float *last = m + dim; \ + const float *last_aligned = m + ((dim >> 5) << 5); \ + if (((uintptr_t)m & 0x3f) == 0) { \ + for (; m != last_aligned; m += 32) { \ + __m512 zmm_m_0 = _mm512_load_ps(m + 0); \ + NORM_FP32_STEP_AVX512(zmm_m_0, zmm_sum_0_0) \ + __m512 zmm_m_1 = _mm512_load_ps(m + 16); \ + NORM_FP32_STEP_AVX512(zmm_m_1, zmm_sum_0_1) \ + } \ + if (last >= last_aligned + 16) { \ + __m512 zmm_m = _mm512_load_ps(m); \ + NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \ + m += 16; \ + } \ + } else { \ + for (; m != last_aligned; m += 32) { \ + __m512 zmm_m_0 = _mm512_loadu_ps(m + 0); \ + NORM_FP32_STEP_AVX512(zmm_m_0, zmm_sum_0_0) \ + __m512 zmm_m_1 = _mm512_loadu_ps(m + 16); \ + NORM_FP32_STEP_AVX512(zmm_m_1, zmm_sum_0_1) \ + } \ + if (last >= last_aligned + 16) { \ + __m512 zmm_m = _mm512_loadu_ps(m); \ + NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \ + m += 16; \ + } \ + } \ + if (m != last) { \ + __mmask16 mask = (__mmask16)((1 << (last - m)) - 1); \ + __m512 zmm_m = _mm512_mask_loadu_ps(_mm512_setzero_ps(), mask, m); \ + NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \ + } \ + float result = \ + HorizontalAdd_FP32_V512(_mm512_add_ps(zmm_sum_0_0, zmm_sum_0_1)); \ + *out = _NORM(result); + +//! Compute the norm of vectors (FP32, M=1) +#define NORM_FP32_1_NEON(m, dim, out, _NORM) \ + MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \ + const float *last = m + dim; \ + const float *last_aligned = m + ((dim >> 3) << 3); \ + for (; m != last_aligned; m += 8) { \ + float32x4_t v_m_0 = vld1q_f32(m + 0); \ + float32x4_t v_m_1 = vld1q_f32(m + 4); \ + NORM_FP32_STEP_NEON(v_m_0, v_sum_0_0) \ + NORM_FP32_STEP_NEON(v_m_1, v_sum_0_1) \ + } \ + if (last >= last_aligned + 4) { \ + float32x4_t v_m = vld1q_f32(m); \ + NORM_FP32_STEP_NEON(v_m, v_sum_0_0) \ + m += 4; \ + } \ + float result = vaddvq_f32(vaddq_f32(v_sum_0_0, v_sum_0_1)); \ + switch (last - m) { \ + case 3: \ + NORM_FP32_STEP_GENERAL(m[2], result) \ + /* FALLTHRU */ \ + case 2: \ + NORM_FP32_STEP_GENERAL(m[1], result) \ + /* FALLTHRU */ \ + case 1: \ + NORM_FP32_STEP_GENERAL(m[0], result) \ + } \ + *out = _NORM(result); diff --git a/src/ailego/math/normalizer.cc b/src/ailego/math/normalizer.cc new file mode 100644 index 0000000..a31a9f3 --- /dev/null +++ b/src/ailego/math/normalizer.cc @@ -0,0 +1,437 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "normalizer.h" + +namespace zvec { +namespace ailego { + +#if (defined(__ARM_NEON) && defined(__aarch64__)) +static inline void NormalizeNEON(float *arr, size_t dim, float norm) { + float *last = arr + dim; + float *last_aligned = arr + ((dim >> 3) << 3); + + float32x4_t v_norm = vdupq_n_f32(norm); + for (; arr != last_aligned; arr += 8) { + vst1q_f32(arr + 0, vdivq_f32(vld1q_f32(arr + 0), v_norm)); + vst1q_f32(arr + 4, vdivq_f32(vld1q_f32(arr + 4), v_norm)); + } + if (last >= last_aligned + 4) { + vst1q_f32(arr, vdivq_f32(vld1q_f32(arr), v_norm)); + arr += 4; + } + switch (last - arr) { + case 3: + arr[2] /= norm; + /* FALLTHRU */ + case 2: + arr[1] /= norm; + /* FALLTHRU */ + case 1: + arr[0] /= norm; + } +} + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) +static inline void NormalizeNEON(float16_t *arr, size_t dim, float norm) { + float16_t *last = arr + dim; + float16_t *last_aligned = arr + ((dim >> 4) << 4); + + float16x8_t v_norm = vdupq_n_f16((float16_t)norm); + for (; arr != last_aligned; arr += 16) { + vst1q_f16(arr + 0, vdivq_f16(vld1q_f16(arr + 0), v_norm)); + vst1q_f16(arr + 8, vdivq_f16(vld1q_f16(arr + 8), v_norm)); + } + if (last >= arr + 8) { + vst1q_f16(arr, vdivq_f16(vld1q_f16(arr), v_norm)); + arr += 8; + } + if (last >= arr + 4) { + vst1_f16(arr, vdiv_f16(vld1_f16(arr), vget_low_f16(v_norm))); + arr += 4; + } + switch (last - arr) { + case 3: + arr[2] /= norm; + /* FALLTHRU */ + case 2: + arr[1] /= norm; + /* FALLTHRU */ + case 1: + arr[0] /= norm; + } +} +#else +static inline void NormalizeNEON(float16_t *arr, size_t dim, float norm) { + float16_t *last = arr + dim; + float16_t *last_aligned = arr + ((dim >> 4) << 4); + + float32x4_t v_norm = vdupq_n_f32(norm); + for (; arr != last_aligned; arr += 16) { + float16x8_t vf16_0 = vld1q_f16(arr + 0); + float16x8_t vf16_1 = vld1q_f16(arr + 8); + vf16_0 = vcombine_f16( + vcvt_f16_f32(vdivq_f32(vcvt_f32_f16(vget_low_f16(vf16_0)), v_norm)), + vcvt_f16_f32(vdivq_f32(vcvt_high_f32_f16(vf16_0), v_norm))); + vf16_1 = vcombine_f16( + vcvt_f16_f32(vdivq_f32(vcvt_f32_f16(vget_low_f16(vf16_1)), v_norm)), + vcvt_f16_f32(vdivq_f32(vcvt_high_f32_f16(vf16_1), v_norm))); + vst1q_f16(arr + 0, vf16_0); + vst1q_f16(arr + 8, vf16_1); + } + if (last >= arr + 8) { + float16x8_t vf16 = vld1q_f16(arr); + vf16 = vcombine_f16( + vcvt_f16_f32(vdivq_f32(vcvt_f32_f16(vget_low_f16(vf16)), v_norm)), + vcvt_f16_f32(vdivq_f32(vcvt_high_f32_f16(vf16), v_norm))); + vst1q_f16(arr, vf16); + arr += 8; + } + if (last >= arr + 4) { + vst1_f16(arr, vcvt_f16_f32(vdivq_f32(vcvt_f32_f16(vld1_f16(arr)), v_norm))); + arr += 4; + } + switch (last - arr) { + case 3: + arr[2] /= norm; + /* FALLTHRU */ + case 2: + arr[1] /= norm; + /* FALLTHRU */ + case 1: + arr[0] /= norm; + } +} +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC +#endif // __ARM_NEON && __aarch64__ + +#if defined(__AVX__) +#if defined(__AVX512F__) +static inline void NormalizeAVX512(float *arr, size_t dim, float norm) { + float *last = arr + dim; + float *last_aligned = arr + ((dim >> 4) << 4); + + __m512 zmm_norm = _mm512_set1_ps(norm); + if (((uintptr_t)arr & 0x3f) == 0) { + for (; arr != last_aligned; arr += 16) { + _mm512_store_ps(arr, _mm512_div_ps(_mm512_load_ps(arr), zmm_norm)); + } + if (last >= arr + 8) { + __m256 ymm_norm = _mm256_set1_ps(norm); + _mm256_store_ps(arr, _mm256_div_ps(_mm256_load_ps(arr), ymm_norm)); + arr += 8; + } + if (last >= arr + 4) { + __m128 xmm_norm = _mm_set1_ps(norm); + _mm_store_ps(arr, _mm_div_ps(_mm_load_ps(arr), xmm_norm)); + arr += 4; + } + } else { + for (; arr != last_aligned; arr += 16) { + _mm512_storeu_ps(arr, _mm512_div_ps(_mm512_loadu_ps(arr), zmm_norm)); + } + if (last >= arr + 8) { + __m256 ymm_norm = _mm256_set1_ps(norm); + _mm256_storeu_ps(arr, _mm256_div_ps(_mm256_loadu_ps(arr), ymm_norm)); + arr += 8; + } + if (last >= arr + 4) { + __m128 xmm_norm = _mm_set1_ps(norm); + _mm_storeu_ps(arr, _mm_div_ps(_mm_loadu_ps(arr), xmm_norm)); + arr += 4; + } + } + switch (last - arr) { + case 3: + arr[2] /= norm; + /* FALLTHRU */ + case 2: + arr[1] /= norm; + /* FALLTHRU */ + case 1: + arr[0] /= norm; + } +} +#endif // __AVX512F__ + +static inline void NormalizeAVX(float *arr, size_t dim, float norm) { + float *last = arr + dim; + float *last_aligned = arr + ((dim >> 4) << 4); + + __m256 ymm_norm = _mm256_set1_ps(norm); + if (((uintptr_t)arr & 0x1f) == 0) { + for (; arr != last_aligned; arr += 16) { + _mm256_store_ps(arr + 0, + _mm256_div_ps(_mm256_load_ps(arr + 0), ymm_norm)); + _mm256_store_ps(arr + 8, + _mm256_div_ps(_mm256_load_ps(arr + 8), ymm_norm)); + } + if (last >= arr + 8) { + _mm256_store_ps(arr, _mm256_div_ps(_mm256_load_ps(arr), ymm_norm)); + arr += 8; + } + if (last >= arr + 4) { + __m128 xmm_norm = _mm_set1_ps(norm); + _mm_store_ps(arr, _mm_div_ps(_mm_load_ps(arr), xmm_norm)); + arr += 4; + } + } else { + for (; arr != last_aligned; arr += 16) { + _mm256_storeu_ps(arr + 0, + _mm256_div_ps(_mm256_loadu_ps(arr + 0), ymm_norm)); + _mm256_storeu_ps(arr + 8, + _mm256_div_ps(_mm256_loadu_ps(arr + 8), ymm_norm)); + } + if (last >= arr + 8) { + _mm256_storeu_ps(arr, _mm256_div_ps(_mm256_loadu_ps(arr), ymm_norm)); + arr += 8; + } + if (last >= arr + 4) { + __m128 xmm_norm = _mm_set1_ps(norm); + _mm_storeu_ps(arr, _mm_div_ps(_mm_loadu_ps(arr), xmm_norm)); + arr += 4; + } + } + switch (last - arr) { + case 3: + arr[2] /= norm; + /* FALLTHRU */ + case 2: + arr[1] /= norm; + /* FALLTHRU */ + case 1: + arr[0] /= norm; + } +} +#endif // __AVX__ + +#if defined(__AVX__) && defined(__F16C__) +#if defined(__AVX512F__) +static inline void NormalizeAVX512(uint16_t *arr, size_t dim, float norm) { + uint16_t *last = arr + dim; + uint16_t *last_aligned = arr + ((dim >> 4) << 4); + + __m512 zmm_norm = _mm512_set1_ps(norm); + if (((uintptr_t)arr & 0x1f) == 0) { + for (; arr != last_aligned; arr += 16) { + _mm256_store_si256( + (__m256i *)arr, + _mm512_cvtps_ph(_mm512_div_ps(_mm512_cvtph_ps(_mm256_load_si256( + (const __m256i *)arr)), + zmm_norm), + _MM_FROUND_NO_EXC)); + } + if (last >= arr + 8) { + __m256 ymm_norm = _mm256_set1_ps(norm); + _mm_store_si128( + (__m128i *)arr, + _mm256_cvtps_ph(_mm256_div_ps(_mm256_cvtph_ps(_mm_load_si128( + (const __m128i *)arr)), + ymm_norm), + _MM_FROUND_NO_EXC)); + arr += 8; + } + } else { + for (; arr != last_aligned; arr += 16) { + _mm256_storeu_si256( + (__m256i *)arr, + _mm512_cvtps_ph(_mm512_div_ps(_mm512_cvtph_ps(_mm256_loadu_si256( + (const __m256i *)arr)), + zmm_norm), + _MM_FROUND_NO_EXC)); + } + if (last >= arr + 8) { + __m256 ymm_norm = _mm256_set1_ps(norm); + _mm_storeu_si128( + (__m128i *)arr, + _mm256_cvtps_ph(_mm256_div_ps(_mm256_cvtph_ps(_mm_loadu_si128( + (const __m128i *)arr)), + ymm_norm), + _MM_FROUND_NO_EXC)); + arr += 8; + } + } + if (last >= arr + 4) { + __m128 xmm_norm = _mm_set1_ps(norm); + _mm_storel_epi64( + (__m128i *)arr, + _mm_cvtps_ph( + _mm_div_ps(_mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)arr)), + xmm_norm), + _MM_FROUND_NO_EXC)); + arr += 8; + } + switch (last - arr) { + case 3: + arr[2] = _cvtss_sh(_cvtsh_ss(arr[2]) / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 2: + arr[1] = _cvtss_sh(_cvtsh_ss(arr[1]) / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 1: + arr[0] = _cvtss_sh(_cvtsh_ss(arr[0]) / norm, _MM_FROUND_NO_EXC); + } +} +#endif // __AVX512F__ + +static inline void NormalizeAVX(uint16_t *arr, size_t dim, float norm) { + uint16_t *last = arr + dim; + uint16_t *last_aligned = arr + ((dim >> 4) << 4); + + __m256 ymm_norm = _mm256_set1_ps(norm); + if (((uintptr_t)arr & 0xf) == 0) { + for (; arr != last_aligned; arr += 16) { + __m128i xmm_0 = _mm_load_si128((const __m128i *)(arr + 0)); + __m128i xmm_1 = _mm_load_si128((const __m128i *)(arr + 8)); + __m256 ymm_0 = _mm256_div_ps(_mm256_cvtph_ps(xmm_0), ymm_norm); + __m256 ymm_1 = _mm256_div_ps(_mm256_cvtph_ps(xmm_1), ymm_norm); + _mm_store_si128((__m128i *)(arr + 0), + _mm256_cvtps_ph(ymm_0, _MM_FROUND_NO_EXC)); + _mm_store_si128((__m128i *)(arr + 8), + _mm256_cvtps_ph(ymm_1, _MM_FROUND_NO_EXC)); + } + if (last >= arr + 8) { + _mm_store_si128( + (__m128i *)arr, + _mm256_cvtps_ph(_mm256_div_ps(_mm256_cvtph_ps(_mm_load_si128( + (const __m128i *)arr)), + ymm_norm), + _MM_FROUND_NO_EXC)); + arr += 8; + } + } else { + for (; arr != last_aligned; arr += 16) { + __m128i xmm_0 = _mm_loadu_si128((const __m128i *)(arr + 0)); + __m128i xmm_1 = _mm_loadu_si128((const __m128i *)(arr + 8)); + __m256 ymm_0 = _mm256_div_ps(_mm256_cvtph_ps(xmm_0), ymm_norm); + __m256 ymm_1 = _mm256_div_ps(_mm256_cvtph_ps(xmm_1), ymm_norm); + _mm_storeu_si128((__m128i *)(arr + 0), + _mm256_cvtps_ph(ymm_0, _MM_FROUND_NO_EXC)); + _mm_storeu_si128((__m128i *)(arr + 8), + _mm256_cvtps_ph(ymm_1, _MM_FROUND_NO_EXC)); + } + if (last >= arr + 8) { + _mm_storeu_si128( + (__m128i *)arr, + _mm256_cvtps_ph(_mm256_div_ps(_mm256_cvtph_ps(_mm_loadu_si128( + (const __m128i *)arr)), + ymm_norm), + _MM_FROUND_NO_EXC)); + arr += 8; + } + } + if (last >= arr + 4) { + __m128 xmm_norm = _mm_set1_ps(norm); + _mm_storel_epi64( + (__m128i *)arr, + _mm_cvtps_ph( + _mm_div_ps(_mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)arr)), + xmm_norm), + _MM_FROUND_NO_EXC)); + arr += 8; + } + switch (last - arr) { + case 3: + arr[2] = _cvtss_sh(_cvtsh_ss(arr[2]) / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 2: + arr[1] = _cvtss_sh(_cvtsh_ss(arr[1]) / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 1: + arr[0] = _cvtss_sh(_cvtsh_ss(arr[0]) / norm, _MM_FROUND_NO_EXC); + } +} +#endif // __AVX__ && __F16C__ + +#if defined(__SSE__) +static inline void NormalizeSSE(float *arr, size_t dim, float norm) { + float *last = arr + dim; + float *last_aligned = arr + ((dim >> 3) << 3); + + __m128 xmm_norm = _mm_set1_ps(norm); + if (((uintptr_t)arr & 0xf) == 0) { + for (; arr != last_aligned; arr += 8) { + _mm_store_ps(arr + 0, _mm_div_ps(_mm_load_ps(arr + 0), xmm_norm)); + _mm_store_ps(arr + 4, _mm_div_ps(_mm_load_ps(arr + 4), xmm_norm)); + } + if (last >= last_aligned + 4) { + _mm_store_ps(arr, _mm_div_ps(_mm_load_ps(arr), xmm_norm)); + arr += 4; + } + } else { + for (; arr != last_aligned; arr += 8) { + _mm_storeu_ps(arr + 0, _mm_div_ps(_mm_loadu_ps(arr + 0), xmm_norm)); + _mm_storeu_ps(arr + 4, _mm_div_ps(_mm_loadu_ps(arr + 4), xmm_norm)); + } + if (last >= last_aligned + 4) { + _mm_storeu_ps(arr, _mm_div_ps(_mm_loadu_ps(arr), xmm_norm)); + arr += 4; + } + } + switch (last - arr) { + case 3: + arr[2] /= norm; + /* FALLTHRU */ + case 2: + arr[1] /= norm; + /* FALLTHRU */ + case 1: + arr[0] /= norm; + } +} +#endif // __SSE__ + +#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__)) +//! Compute the norm of vector +void Normalizer::Compute(ValueType *arr, size_t dim, float norm) { +#if defined(__ARM_NEON) + NormalizeNEON(arr, dim, norm); +#else +#if defined(__AVX512F__) + if (dim > 15) { + NormalizeAVX512(arr, dim, norm); + return; + } +#endif // __AVX512F__ +#if defined(__AVX__) + if (dim > 7) { + NormalizeAVX(arr, dim, norm); + return; + } +#endif // __AVX__ + NormalizeSSE(arr, dim, norm); +#endif // __ARM_NEON +} +#endif // __SSE__ || (__ARM_NEON && __aarch64__) + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +//! Compute the norm of vector +void Normalizer::Compute(ValueType *arr, size_t dim, float norm) { +#if defined(__ARM_NEON) + NormalizeNEON(reinterpret_cast(arr), dim, norm); +#else +#if defined(__AVX512F__) + if (dim > 31) { + NormalizeAVX512(reinterpret_cast(arr), dim, norm); + return; + } +#endif // __AVX512F__ + NormalizeAVX(reinterpret_cast(arr), dim, norm); +#endif // __ARM_NEON +} +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/normalizer.h b/src/ailego/math/normalizer.h new file mode 100644 index 0000000..2c191b0 --- /dev/null +++ b/src/ailego/math/normalizer.h @@ -0,0 +1,114 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "norm_matrix.h" + +namespace zvec { +namespace ailego { + +/*! Normalizer + */ +template ::value>::type> +struct Normalizer { + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! Compute the norm of vector + static inline void Compute(ValueType *arr, size_t dim, float norm) { + for (size_t i = 0; i < dim; ++i) { + arr[i] /= norm; + } + } + + //! Normalize a vector (L1) + static inline void L1(ValueType *arr, size_t dim, float *norm) { + Norm1Matrix::Compute(arr, dim, norm); + if (*norm > 0.0f) { + Compute(arr, dim, *norm); + } + } + + //! Normalize a vector (L2) + static inline void L2(ValueType *arr, size_t dim, float *norm) { + Norm2Matrix::Compute(arr, dim, norm); + if (*norm > 0.0f) { + Compute(arr, dim, *norm); + } + } +}; + +#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__)) +/*! Normalizer (FP32) + */ +template <> +struct Normalizer { + //! Type of value + using ValueType = float; + + //! Compute the norm of vector + static void Compute(ValueType *arr, size_t dim, float norm); + + //! Normalize a vector (L1) + static inline void L1(ValueType *arr, size_t dim, float *norm) { + Norm1Matrix::Compute(arr, dim, norm); + if (*norm > 0.0f) { + Compute(arr, dim, *norm); + } + } + + //! Normalize a vector (L2) + static inline void L2(ValueType *arr, size_t dim, float *norm) { + Norm2Matrix::Compute(arr, dim, norm); + if (*norm > 0.0f) { + Compute(arr, dim, *norm); + } + } +}; +#endif // __SSE__ || (__ARM_NEON && __aarch64__) + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +/*! Normalizer (FP16) + */ +template <> +struct Normalizer { + //! Type of value + using ValueType = Float16; + + //! Compute the norm of vector + static void Compute(ValueType *arr, size_t dim, float norm); + + //! Normalize a vector (L1) + static inline void L1(ValueType *arr, size_t dim, float *norm) { + Norm1Matrix::Compute(arr, dim, norm); + if (*norm > 0.0f) { + Compute(arr, dim, *norm); + } + } + + //! Normalize a vector (L2) + static inline void L2(ValueType *arr, size_t dim, float *norm) { + Norm2Matrix::Compute(arr, dim, norm); + if (*norm > 0.0f) { + Compute(arr, dim, *norm); + } + } +}; +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math_batch/cosine_distance_batch.h b/src/ailego/math_batch/cosine_distance_batch.h new file mode 100644 index 0000000..b8a8309 --- /dev/null +++ b/src/ailego/math_batch/cosine_distance_batch.h @@ -0,0 +1,58 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include "inner_product_distance_batch.h" + +namespace zvec::ailego::DistanceBatch { + + +template +struct CosineDistanceBatch; + +template +struct CosineDistanceBatch { + using ValueType = typename std::remove_cv::type; + + static inline void ComputeBatch(const ValueType **vecs, + const ValueType *query, size_t num_vecs, + size_t dim, float *results) { + constexpr size_t extra_dim = sizeof(float) / sizeof(ValueType); + size_t _dim = dim - extra_dim; + + InnerProductDistanceBatch::ComputeBatch( + vecs, query, num_vecs, _dim, results); + + for (size_t i = 0; i < num_vecs; ++i) { + results[i] = 1 - results[i]; + } + } + + using IPImplType = + InnerProductDistanceBatch; + + static void QueryPreprocess(void *query, size_t dim) { + return IPImplType::QueryPreprocess(query, + dim - sizeof(float) / sizeof(ValueType)); + } +}; + + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/math_batch/distance_batch.h b/src/ailego/math_batch/distance_batch.h new file mode 100644 index 0000000..9494be8 --- /dev/null +++ b/src/ailego/math_batch/distance_batch.h @@ -0,0 +1,65 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include "ailego/math/distance_matrix.h" +#include "cosine_distance_batch.h" +#include "euclidean_distance_batch.h" +#include "inner_product_distance_batch.h" + +namespace zvec::ailego { + +template < + template class DistanceType, + typename ValueType, size_t BatchSize, size_t PrefetchStep, typename = void> +struct BaseDistance { + static inline void _ComputeBatch(const ValueType **m, const ValueType *q, + size_t num, size_t dim, float *out) { + for (size_t i = 0; i < num; ++i) { + DistanceType::Compute(m[i], q, dim, out + i); + } + } + + // If Distance has ComputeBatch, use it; otherwise fall back to _ComputeBatch. + static inline void ComputeBatch(const ValueType **m, const ValueType *q, + size_t num, size_t dim, float *out) { + if constexpr (std::is_same_v, + CosineDistanceMatrix>) { + return DistanceBatch::CosineDistanceBatch< + ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim, + out); + } + + if constexpr (std::is_same_v, + EuclideanDistanceMatrix>) { + return DistanceBatch::EuclideanDistanceBatch< + ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim, + out); + } + + if constexpr (std::is_same_v< + DistanceType, + SquaredEuclideanDistanceMatrix>) { + return DistanceBatch::SquaredEuclideanDistanceBatch< + ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim, + out); + } + + _ComputeBatch(m, q, num, dim, out); + } +}; + +} // namespace zvec::ailego \ No newline at end of file diff --git a/src/ailego/math_batch/distance_batch_math.h b/src/ailego/math_batch/distance_batch_math.h new file mode 100644 index 0000000..8ab4014 --- /dev/null +++ b/src/ailego/math_batch/distance_batch_math.h @@ -0,0 +1,30 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#if defined(__AVX2__) + +inline float sum4(__m128 v) { + v = _mm_add_ps(v, _mm_castsi128_ps(_mm_srli_si128(_mm_castps_si128(v), 8))); + return _mm_cvtss_f32(v) + _mm_cvtss_f32(_mm_shuffle_ps(v, v, 1)); +} + +inline __m128 sum_top_bottom_avx(__m256 v) { + const __m128 high = _mm256_extractf128_ps(v, 1); + const __m128 low = _mm256_castps256_ps128(v); + return _mm_add_ps(high, low); +} + +#endif diff --git a/src/ailego/math_batch/euclidean_distance_batch.h b/src/ailego/math_batch/euclidean_distance_batch.h new file mode 100644 index 0000000..84d0a8f --- /dev/null +++ b/src/ailego/math_batch/euclidean_distance_batch.h @@ -0,0 +1,143 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +// SquaredEuclideanDistanceBatch +template +struct SquaredEuclideanDistanceBatch; + +template +static void compute_one_to_many_squared_euclidean_fallback( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + for (size_t j = 0; j < BatchSize; ++j) { + sums[j] = 0.0; + SquaredEuclideanDistanceMatrix::Compute(ptrs[j], query, + dim, sums + j); + ailego_prefetch(&prefetch_ptrs[j]); + } +} + +// Function template partial specialization is not allowed, +// therefore the wrapper struct is required. +template +struct SquaredEuclideanDistanceBatchImpl { + using ValueType = typename std::remove_cv::type; + static void compute_one_to_many( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_squared_euclidean_fallback( + query, ptrs, prefetch_ptrs, dim, sums); + } +}; + +template +struct SquaredEuclideanDistanceBatch { + using ValueType = typename std::remove_cv::type; + + static inline void ComputeBatch(const ValueType **vecs, + const ValueType *query, size_t num_vecs, + size_t dim, float *results) { + size_t i = 0; + for (; i + BatchSize <= num_vecs; i += BatchSize) { + std::array prefetch_ptrs; + for (size_t j = 0; j < BatchSize; ++j) { + if (i + j + BatchSize * PrefetchStep < num_vecs) { + prefetch_ptrs[j] = vecs[i + j + BatchSize * PrefetchStep]; + } else { + prefetch_ptrs[j] = nullptr; + } + } + SquaredEuclideanDistanceBatchImpl< + ValueType, BatchSize>::compute_one_to_many(query, &vecs[i], + prefetch_ptrs, dim, + &results[i]); + } + for (; i < num_vecs; ++i) { // TODO: unroll by 1, 2, 4, 8, etc. + std::array prefetch_ptrs{nullptr}; + SquaredEuclideanDistanceBatchImpl::compute_one_to_many( + query, &vecs[i], prefetch_ptrs, dim, &results[i]); + } + } +}; + +// EuclideanDistanceBatch +template +struct EuclideanDistanceBatch; + +template +struct EuclideanDistanceBatch { + using ValueType = typename std::remove_cv::type; + + static inline void ComputeBatch(const ValueType **vecs, + const ValueType *query, size_t num_vecs, + size_t dim, float *results) { + SquaredEuclideanDistanceBatch::ComputeBatch( + vecs, query, num_vecs, dim, results); + + for (size_t i = 0; i < num_vecs; ++i) { + results[i] = std::sqrt(results[i]); + } + } +}; + + +template <> +struct SquaredEuclideanDistanceBatchImpl { + using ValueType = ailego::Float16; + static void compute_one_to_many( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums); +}; + +template <> +struct SquaredEuclideanDistanceBatchImpl { + using ValueType = float; + static void compute_one_to_many(const float *query, const float **ptrs, + std::array &prefetch_ptrs, + size_t dim, float *sums); +}; + +template <> +struct SquaredEuclideanDistanceBatchImpl { + using ValueType = ailego::Float16; + static void compute_one_to_many( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums); +}; + +template <> +struct SquaredEuclideanDistanceBatchImpl { + using ValueType = float; + static void compute_one_to_many(const float *query, const float **ptrs, + std::array &prefetch_ptrs, + size_t dim, float *sums); +}; + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/euclidean_distance_batch_dispatch.cc b/src/ailego/math_batch/euclidean_distance_batch_dispatch.cc new file mode 100644 index 0000000..5c8ffb2 --- /dev/null +++ b/src/ailego/math_batch/euclidean_distance_batch_dispatch.cc @@ -0,0 +1,191 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "euclidean_distance_batch.h" + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512FP16__) +void compute_one_to_many_squared_euclidean_avx512fp16_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +void compute_one_to_many_squared_euclidean_avx512fp16_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); +#endif //__AVX512FP16__ + +#if defined(__AVX512F__) +void compute_one_to_many_squared_euclidean_avx512f_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +void compute_one_to_many_squared_euclidean_avx512f_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +void compute_one_to_many_squared_euclidean_avx512f_fp32_1( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); + +void compute_one_to_many_squared_euclidean_avx512f_fp32_12( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); +#endif //__AVX512F__ + +#if defined(__AVX2__) +void compute_one_to_many_squared_euclidean_avx2_fp32_1( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); + +void compute_one_to_many_squared_euclidean_avx2_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +// void compute_one_to_many_squared_euclidean_avx2_int8_1( +// const int8_t *query, const int8_t **ptrs, +// std::array &prefetch_ptrs, size_t dimensionality, +// float *results); + +void compute_one_to_many_squared_euclidean_avx2_fp32_12( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); + +void compute_one_to_many_squared_euclidean_avx2_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +// void compute_one_to_many_squared_euclidean_avx2_int8_12( +// const int8_t *query, const int8_t **ptrs, +// std::array &prefetch_ptrs, size_t dimensionality, +// float *results); +#endif + +void SquaredEuclideanDistanceBatchImpl::compute_one_to_many( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_squared_euclidean_avx2_fp32_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_squared_euclidean_fallback( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void SquaredEuclideanDistanceBatchImpl::compute_one_to_many( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { +#if defined(__AVX512FP16__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) { + return compute_one_to_many_squared_euclidean_avx512fp16_fp16_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return compute_one_to_many_squared_euclidean_avx512f_fp16_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_squared_euclidean_avx2_fp16_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_squared_euclidean_fallback( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void SquaredEuclideanDistanceBatchImpl::compute_one_to_many( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return compute_one_to_many_squared_euclidean_avx512f_fp32_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_squared_euclidean_avx2_fp32_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_squared_euclidean_fallback( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void SquaredEuclideanDistanceBatchImpl:: + compute_one_to_many(const ailego::Float16 *query, + const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dim, float *sums) { +#if defined(__AVX512FP16__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) { + return compute_one_to_many_squared_euclidean_avx512fp16_fp16_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return compute_one_to_many_squared_euclidean_avx512f_fp16_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_squared_euclidean_avx2_fp16_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_squared_euclidean_fallback( + query, ptrs, prefetch_ptrs, dim, sums); +} + +// void SquaredEuclideanDistanceBatchImpl::compute_one_to_many( +// const int8_t *query, const int8_t **ptrs, +// std::array &prefetch_ptrs, size_t dim, float *sums) { + +// #if defined(__AVX2__) +// if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { +// return compute_one_to_many_squared_euclidean_avx2_int8_12( +// query, ptrs, prefetch_ptrs, dim, sums); +// } +// #endif +// return compute_one_to_many_squared_euclidean_fallback(query, ptrs, +// prefetch_ptrs, +// dim, sums); +// } + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx2.cc b/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx2.cc new file mode 100644 index 0000000..e6d793d --- /dev/null +++ b/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx2.cc @@ -0,0 +1,114 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX2__) + +template +static std::enable_if_t, void> +compute_one_to_many_squared_euclidean_avx2_fp16( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m256 accs[dp_batch]; + + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm256_setzero_ps(); + } + + size_t dim = 0; + for (; dim + 16 <= dimensionality; dim += 16) { + __m256i q = + _mm256_loadu_si256(reinterpret_cast(query + dim)); + + __m256 q1 = _mm256_cvtph_ps(_mm256_castsi256_si128(q)); + __m256 q2 = _mm256_cvtph_ps(_mm256_extractf128_si256(q, 1)); + + __m256 data_regs_1[dp_batch]; + __m256 data_regs_2[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + __m256i m = + _mm256_loadu_si256(reinterpret_cast(ptrs[i] + dim)); + + data_regs_1[i] = _mm256_cvtph_ps(_mm256_castsi256_si128(m)); + data_regs_2[i] = _mm256_cvtph_ps(_mm256_extractf128_si256(m, 1)); + } + + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + __m256 diff1 = _mm256_sub_ps(q1, data_regs_1[i]); + accs[i] = _mm256_fmadd_ps(diff1, diff1, accs[i]); + + __m256 diff2 = _mm256_sub_ps(q2, data_regs_2[i]); + accs[i] = _mm256_fmadd_ps(diff2, diff2, accs[i]); + } + } + + if (dim + 8 <= dimensionality) { + __m256 q = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(query + dim))); + + __m256 data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(ptrs[i] + dim))); + + __m256 diff = _mm256_sub_ps(q, data_regs[i]); + accs[i] = _mm256_fmadd_ps(diff, diff, accs[i]); + } + + dim += 8; + } + + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = HorizontalAdd_FP32_V256(accs[i]); + } + + for (; dim < dimensionality; ++dim) { + for (size_t i = 0; i < dp_batch; ++i) { + float diff = (*(query + dim)) - (*(ptrs[i] + dim)); + results[i] += diff * diff; + } + } +} + +void compute_one_to_many_squared_euclidean_avx2_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_squared_euclidean_avx2_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_squared_euclidean_avx2_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_squared_euclidean_avx2_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} +#endif + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512.cc b/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512.cc new file mode 100644 index 0000000..1d4750e --- /dev/null +++ b/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512.cc @@ -0,0 +1,137 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512F__) + +template +static std::enable_if_t, void> +compute_one_to_many_squared_euclidean_avx512f_fp16( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m512 accs[dp_batch]; + + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_setzero_ps(); + } + + size_t dim = 0; + for (; dim + 32 <= dimensionality; dim += 32) { + __m512i q = + _mm512_loadu_si512(reinterpret_cast(query + dim)); + + __m512 q1 = _mm512_cvtph_ps(_mm512_castsi512_si256(q)); + __m512 q2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(q, 1)); + + __m512 data_regs_1[dp_batch]; + __m512 data_regs_2[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + __m512i m = + _mm512_loadu_si512(reinterpret_cast(ptrs[i] + dim)); + + data_regs_1[i] = _mm512_cvtph_ps(_mm512_castsi512_si256(m)); + data_regs_2[i] = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(m, 1)); + } + + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + __m512 diff1 = _mm512_sub_ps(q1, data_regs_1[i]); + accs[i] = _mm512_fmadd_ps(diff1, diff1, accs[i]); + + __m512 diff2 = _mm512_sub_ps(q2, data_regs_2[i]); + accs[i] = _mm512_fmadd_ps(diff2, diff2, accs[i]); + } + } + + if (dim + 16 <= dimensionality) { + __m512 q = _mm512_cvtph_ps( + _mm256_loadu_si256(reinterpret_cast(query + dim))); + + for (size_t i = 0; i < dp_batch; ++i) { + __m512 m = _mm512_cvtph_ps( + _mm256_loadu_si256(reinterpret_cast(ptrs[i] + dim))); + + __m512 diff = _mm512_sub_ps(m, q); + accs[i] = _mm512_fmadd_ps(diff, diff, accs[i]); + } + + dim += 16; + } + + __m256 acc_new[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + acc_new[i] = _mm256_add_ps( + _mm512_castps512_ps256(accs[i]), + _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(accs[i]), 1))); + } + + if (dim + 8 < dimensionality) { + __m256 q = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(query + dim))); + + for (size_t i = 0; i < dp_batch; ++i) { + __m256 m = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(ptrs[i] + dim))); + + __m256 diff = _mm256_sub_ps(m, q); + acc_new[i] = _mm256_fmadd_ps(diff, diff, acc_new[i]); + } + + dim += 8; + } + + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = HorizontalAdd_FP32_V256(acc_new[i]); + } + + for (; dim < dimensionality; ++dim) { + for (size_t i = 0; i < dp_batch; ++i) { + float diff = (*(query + dim)) - (*(ptrs[i] + dim)); + results[i] += diff * diff; + } + } +} + +void compute_one_to_many_squared_euclidean_avx512f_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_squared_euclidean_avx512f_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_squared_euclidean_avx512f_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_squared_euclidean_avx512f_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +#endif + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512fp16.cc b/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512fp16.cc new file mode 100644 index 0000000..70e0240 --- /dev/null +++ b/src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512fp16.cc @@ -0,0 +1,99 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512FP16__) + +template +static std::enable_if_t, void> +compute_one_to_many_squared_euclidean_avx512fp16_fp16( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m512h accs[dp_batch]; + + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_setzero_ph(); + } + + size_t dim = 0; + for (; dim + 32 <= dimensionality; dim += 32) { + __m512h q = _mm512_loadu_ph(query + dim); + + __m512h data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm512_loadu_ph(ptrs[i] + dim); + } + + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + __m512h diff = _mm512_sub_ph(data_regs[i], q); + accs[i] = _mm512_fmadd_ph(diff, diff, accs[i]); + } + } + + if (dim < dimensionality) { + __mmask32 mask = (__mmask32)(((uint32_t)1 << (dimensionality - dim)) - 1); + + for (size_t i = 0; i < dp_batch; ++i) { + __m512i zmm_undefined = _mm512_undefined_epi32(); + __m512h zmm_undefined_ph = _mm512_undefined_ph(); + __m512h zmm_d = + _mm512_mask_sub_ph(zmm_undefined_ph, mask, + _mm512_castsi512_ph(_mm512_mask_loadu_epi16( + zmm_undefined, mask, query + dim)), + _mm512_castsi512_ph(_mm512_mask_loadu_epi16( + zmm_undefined, mask, ptrs[i] + dim))); + + accs[i] = _mm512_mask3_fmadd_ph(zmm_d, zmm_d, accs[i], mask); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = HorizontalAdd_FP16_V512(accs[i]); + } +} + +void compute_one_to_many_squared_euclidean_avx512fp16_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_squared_euclidean_avx512fp16_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_squared_euclidean_avx512fp16_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_squared_euclidean_avx512fp16_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +#endif + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx2.cc b/src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx2.cc new file mode 100644 index 0000000..ce27e92 --- /dev/null +++ b/src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx2.cc @@ -0,0 +1,104 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "distance_batch_math.h" + +#define SSD_FP32_GENERAL(m, q, sum) \ + { \ + float x = m - q; \ + sum += (x * x); \ + } + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX2__) + +template +static std::enable_if_t, void> +compute_one_to_many_squared_euclidean_avx2_fp32( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m256 accs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm256_setzero_ps(); + } + size_t dim = 0; + for (; dim + 8 <= dimensionality; dim += 8) { + __m256 q = _mm256_loadu_ps(query + dim); + __m256 data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm256_loadu_ps(ptrs[i] + dim); + } + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + for (size_t i = 0; i < dp_batch; ++i) { + __m256 diff = _mm256_sub_ps(q, data_regs[i]); + accs[i] = _mm256_fmadd_ps(diff, diff, accs[i]); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = HorizontalAdd_FP32_V256(accs[i]); + + switch (dimensionality - dim) { + case 7: + SSD_FP32_GENERAL(query[dim + 6], ptrs[i][dim + 6], results[i]); + /* FALLTHRU */ + case 6: + SSD_FP32_GENERAL(query[dim + 5], ptrs[i][dim + 5], results[i]); + /* FALLTHRU */ + case 5: + SSD_FP32_GENERAL(query[dim + 4], ptrs[i][dim + 4], results[i]); + /* FALLTHRU */ + case 4: + SSD_FP32_GENERAL(query[dim + 3], ptrs[i][dim + 3], results[i]); + /* FALLTHRU */ + case 3: + SSD_FP32_GENERAL(query[dim + 2], ptrs[i][dim + 2], results[i]); + /* FALLTHRU */ + case 2: + SSD_FP32_GENERAL(query[dim + 1], ptrs[i][dim + 1], results[i]); + /* FALLTHRU */ + case 1: + SSD_FP32_GENERAL(query[dim + 0], ptrs[i][dim + 0], results[i]); + } + } +} + +void compute_one_to_many_squared_euclidean_avx2_fp32_1( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_squared_euclidean_avx2_fp32( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_squared_euclidean_avx2_fp32_12( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_squared_euclidean_avx2_fp32( + query, ptrs, prefetch_ptrs, dim, sums); +} + +#endif + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx512.cc b/src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx512.cc new file mode 100644 index 0000000..d7e5847 --- /dev/null +++ b/src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx512.cc @@ -0,0 +1,95 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "distance_batch_math.h" + +#define SSD_FP32_GENERAL(m, q, sum) \ + { \ + float x = m - q; \ + sum += (x * x); \ + } + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512F__) + +template +static std::enable_if_t, void> +compute_one_to_many_squared_euclidean_avx512f_fp32( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m512 accs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_setzero_ps(); + } + size_t dim = 0; + for (; dim + 16 <= dimensionality; dim += 16) { + __m512 q = _mm512_loadu_ps(query + dim); + __m512 data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm512_loadu_ps(ptrs[i] + dim); + } + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + for (size_t i = 0; i < dp_batch; ++i) { + __m512 diff = _mm512_sub_ps(q, data_regs[i]); + accs[i] = _mm512_fmadd_ps(diff, diff, accs[i]); + } + } + + if (dim < dimensionality) { + __mmask16 mask = (__mmask16)((1 << (dimensionality - dim)) - 1); + + for (size_t i = 0; i < dp_batch; ++i) { + __m512 zmm_undefined = _mm512_undefined_ps(); + + __m512 q = _mm512_mask_loadu_ps(zmm_undefined, mask, query + dim); + __m512 m = _mm512_mask_loadu_ps(zmm_undefined, mask, ptrs[i] + dim); + __m512 diff = _mm512_mask_sub_ps(zmm_undefined, mask, q, m); + + accs[i] = _mm512_mask3_fmadd_ps(diff, diff, accs[i], mask); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = HorizontalAdd_FP32_V512(accs[i]); + } +} + +void compute_one_to_many_squared_euclidean_avx512f_fp32_1( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_squared_euclidean_avx512f_fp32( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_squared_euclidean_avx512f_fp32_12( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_squared_euclidean_avx512f_fp32( + query, ptrs, prefetch_ptrs, dim, sums); +} + +#endif + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/inner_product_distance_batch.h b/src/ailego/math_batch/inner_product_distance_batch.h new file mode 100644 index 0000000..6375482 --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch.h @@ -0,0 +1,143 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +template +struct InnerProductDistanceBatch; + +template +static void compute_one_to_many_inner_product_fallback( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + for (size_t j = 0; j < BatchSize; ++j) { + sums[j] = 0.0; + InnerProductMatrix::Compute(ptrs[j], query, dim, sums + j); + ailego_prefetch(&prefetch_ptrs[j]); + } +} + +// Function template partial specialization is not allowed, +// therefore the wrapper struct is required. +template +struct InnerProductDistanceBatchImpl { + using ValueType = typename std::remove_cv::type; + static void compute_one_to_many( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_inner_product_fallback(query, ptrs, + prefetch_ptrs, dim, sums); + } + static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc() { + return nullptr; + } +}; + +template +struct InnerProductDistanceBatch { + using ValueType = typename std::remove_cv::type; + + static inline void ComputeBatch(const ValueType **vecs, + const ValueType *query, size_t num_vecs, + size_t dim, float *results) { + size_t i = 0; + for (; i + BatchSize <= num_vecs; i += BatchSize) { + std::array prefetch_ptrs; + for (size_t j = 0; j < BatchSize; ++j) { + if (i + j + BatchSize * PrefetchStep < num_vecs) { + prefetch_ptrs[j] = vecs[i + j + BatchSize * PrefetchStep]; + } else { + prefetch_ptrs[j] = nullptr; + } + } + InnerProductDistanceBatchImpl::compute_one_to_many( + query, &vecs[i], prefetch_ptrs, dim, &results[i]); + } + for (; i < num_vecs; ++i) { // TODO: unroll by 1, 2, 4, 8, etc. + std::array prefetch_ptrs{nullptr}; + InnerProductDistanceBatchImpl::compute_one_to_many( + query, &vecs[i], prefetch_ptrs, dim, &results[i]); + } + } + + static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc() { + return InnerProductDistanceBatchImpl::GetQueryPreprocessFunc(); + } +}; + +template <> +struct InnerProductDistanceBatchImpl { + using ValueType = ailego::Float16; + static void compute_one_to_many( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums); +}; + +template <> +struct InnerProductDistanceBatchImpl { + using ValueType = float; + static void compute_one_to_many(const float *query, const float **ptrs, + std::array &prefetch_ptrs, + size_t dim, float *sums); +}; + +template <> +struct InnerProductDistanceBatchImpl { + using ValueType = int8_t; + static void compute_one_to_many(const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, + size_t dim, float *sums); + + static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc(); +}; + +template <> +struct InnerProductDistanceBatchImpl { + using ValueType = ailego::Float16; + static void compute_one_to_many( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums); +}; + +template <> +struct InnerProductDistanceBatchImpl { + using ValueType = float; + static void compute_one_to_many(const float *query, const float **ptrs, + std::array &prefetch_ptrs, + size_t dim, float *sums); +}; + +template <> +struct InnerProductDistanceBatchImpl { + using ValueType = int8_t; + static void compute_one_to_many(const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, + size_t dim, float *sums); +}; + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/inner_product_distance_batch_dispatch.cc b/src/ailego/math_batch/inner_product_distance_batch_dispatch.cc new file mode 100644 index 0000000..7837662 --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_dispatch.cc @@ -0,0 +1,228 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include "inner_product_distance_batch.h" + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512VNNI__) +void compute_one_to_many_inner_product_avx512_vnni_int8_query_preprocess( + void *query, size_t dim); + +void compute_one_to_many_inner_product_avx512_vnni_int8_1( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); + +void compute_one_to_many_inner_product_avx512_vnni_int8_12( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); +#endif + +#if defined(__AVX512FP16__) +void compute_one_to_many_inner_product_avx512fp16_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +void compute_one_to_many_inner_product_avx512fp16_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); +#endif //__AVX512FP16__ + +#if defined(__AVX512F__) +void compute_one_to_many_inner_product_avx512f_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +void compute_one_to_many_inner_product_avx512f_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); +#endif //__AVX512F__ + +#if defined(__AVX2__) +void compute_one_to_many_inner_product_avx2_fp32_1( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); + +void compute_one_to_many_inner_product_avx2_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +void compute_one_to_many_inner_product_avx2_int8_1( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); + +void compute_one_to_many_inner_product_avx2_fp32_12( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); + +void compute_one_to_many_inner_product_avx2_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results); + +void compute_one_to_many_inner_product_avx2_int8_12( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results); +#endif + +void InnerProductDistanceBatchImpl::compute_one_to_many( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_inner_product_avx2_fp32_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs, + dim, sums); +} + +void InnerProductDistanceBatchImpl::compute_one_to_many( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { +#if defined(__AVX512FP16__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) { + return compute_one_to_many_inner_product_avx512fp16_fp16_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return compute_one_to_many_inner_product_avx512f_fp16_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_inner_product_avx2_fp16_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs, + dim, sums); +} + +void InnerProductDistanceBatchImpl::compute_one_to_many( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { +// #if defined(__AVX512BW__) // TODO: this version is problematic +// return compute_one_to_many_avx512_int8( +// query, ptrs, prefetch_ptrs, dim, sums); +#if defined(__AVX512VNNI__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + return compute_one_to_many_inner_product_avx512_vnni_int8_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_inner_product_avx2_int8_1( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs, + dim, sums); +} + +DistanceBatchQueryPreprocessFunc +InnerProductDistanceBatchImpl::GetQueryPreprocessFunc() { +#if defined(__AVX512VNNI__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + return compute_one_to_many_inner_product_avx512_vnni_int8_query_preprocess; + } +#endif + return nullptr; +} + +void InnerProductDistanceBatchImpl::compute_one_to_many( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_inner_product_avx2_fp32_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs, + dim, sums); +} + +void InnerProductDistanceBatchImpl::compute_one_to_many( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { +#if defined(__AVX512FP16__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) { + return compute_one_to_many_inner_product_avx512fp16_fp16_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return compute_one_to_many_inner_product_avx512f_fp16_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_inner_product_avx2_fp16_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs, + dim, sums); +} + +void InnerProductDistanceBatchImpl::compute_one_to_many( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { +// #if defined(__AVX512BW__) // TODO: this version is problematic +// return compute_one_to_many_avx512_int8( +// query, ptrs, prefetch_ptrs, dim, sums); +#if defined(__AVX512VNNI__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + return compute_one_to_many_inner_product_avx512_vnni_int8_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + return compute_one_to_many_inner_product_avx2_int8_12( + query, ptrs, prefetch_ptrs, dim, sums); + } +#endif + return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs, + dim, sums); +} + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc new file mode 100644 index 0000000..d6fe475 --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc @@ -0,0 +1,109 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX2__) + +template +static std::enable_if_t, void> +compute_one_to_many_inner_product_avx2_fp16( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m256 accs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm256_setzero_ps(); + } + + size_t dim = 0; + for (; dim + 16 <= dimensionality; dim += 16) { + __m256i q = + _mm256_loadu_si256(reinterpret_cast(query + dim)); + + __m256 q1 = _mm256_cvtph_ps(_mm256_castsi256_si128(q)); + __m256 q2 = _mm256_cvtph_ps(_mm256_extractf128_si256(q, 1)); + + __m256 data_regs_1[dp_batch]; + __m256 data_regs_2[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + __m256i m = + _mm256_loadu_si256(reinterpret_cast(ptrs[i] + dim)); + + data_regs_1[i] = _mm256_cvtph_ps(_mm256_castsi256_si128(m)); + data_regs_2[i] = _mm256_cvtph_ps(_mm256_extractf128_si256(m, 1)); + } + + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm256_fmadd_ps(q1, data_regs_1[i], accs[i]); + accs[i] = _mm256_fmadd_ps(q2, data_regs_2[i], accs[i]); + } + } + + if (dim + 8 <= dimensionality) { + __m256 q = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(query + dim))); + + __m256 data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(ptrs[i] + dim))); + accs[i] = _mm256_fmadd_ps(q, data_regs[i], accs[i]); + } + + dim += 8; + } + + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = HorizontalAdd_FP32_V256(accs[i]); + } + + for (; dim < dimensionality; ++dim) { + for (size_t i = 0; i < dp_batch; ++i) { + results[i] += (*(query + dim)) * (*(ptrs[i] + dim)); + } + } +} + +void compute_one_to_many_inner_product_avx2_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_inner_product_avx2_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_inner_product_avx2_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_inner_product_avx2_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +#endif + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc new file mode 100644 index 0000000..805da8d --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc @@ -0,0 +1,128 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512F__) + +template +static std::enable_if_t, void> +compute_one_to_many_inner_product_avx512f_fp16( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m512 accs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_setzero_ps(); + } + + size_t dim = 0; + for (; dim + 32 <= dimensionality; dim += 32) { + __m512i q = + _mm512_loadu_si512(reinterpret_cast(query + dim)); + + __m512 q1 = _mm512_cvtph_ps(_mm512_castsi512_si256(q)); + __m512 q2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(q, 1)); + + __m512 data_regs_1[dp_batch]; + __m512 data_regs_2[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + __m512i m = + _mm512_loadu_si512(reinterpret_cast(ptrs[i] + dim)); + + data_regs_1[i] = _mm512_cvtph_ps(_mm512_castsi512_si256(m)); + data_regs_2[i] = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(m, 1)); + } + + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_fmadd_ps(q1, data_regs_1[i], accs[i]); + accs[i] = _mm512_fmadd_ps(q2, data_regs_2[i], accs[i]); + } + } + + if (dim + 16 <= dimensionality) { + __m512 q = _mm512_cvtph_ps( + _mm256_loadu_si256(reinterpret_cast(query + dim))); + + __m512 data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm512_cvtph_ps( + _mm256_loadu_si256(reinterpret_cast(ptrs[i] + dim))); + accs[i] = _mm512_fmadd_ps(q, data_regs[i], accs[i]); + } + + dim += 16; + } + + __m256 acc_new[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + acc_new[i] = _mm256_add_ps( + _mm512_castps512_ps256(accs[i]), + _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(accs[i]), 1))); + } + + if (dim + 8 <= dimensionality) { + __m256 q = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(query + dim))); + + for (size_t i = 0; i < dp_batch; ++i) { + __m256 m = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(ptrs[i] + dim))); + acc_new[i] = _mm256_fmadd_ps(m, q, acc_new[i]); + } + + dim += 8; + } + + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = HorizontalAdd_FP32_V256(acc_new[i]); + } + + for (; dim < dimensionality; ++dim) { + for (size_t i = 0; i < dp_batch; ++i) { + results[i] += (*(query + dim)) * (*(ptrs[i] + dim)); + } + } +} + +void compute_one_to_many_inner_product_avx512f_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_inner_product_avx512f_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_inner_product_avx512f_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_inner_product_avx512f_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} +#endif + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512fp16.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512fp16.cc new file mode 100644 index 0000000..b69e60b --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512fp16.cc @@ -0,0 +1,92 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512FP16__) +template +static std::enable_if_t, void> +compute_one_to_many_inner_product_avx512fp16_fp16( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m512h accs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_setzero_ph(); + } + + size_t dim = 0; + for (; dim + 32 <= dimensionality; dim += 32) { + __m512h q = _mm512_loadu_ph(query + dim); + + __m512h data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm512_loadu_ph(ptrs[i] + dim); + } + + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_fmadd_ph(data_regs[i], q, accs[i]); + } + } + + if (dim < dimensionality) { + __mmask32 mask = (__mmask32)((1 << (dimensionality - dim)) - 1); + + for (size_t i = 0; i < dp_batch; ++i) { + __m512i zmm_undefined = _mm512_undefined_epi32(); + + accs[i] = + _mm512_mask3_fmadd_ph(_mm512_castsi512_ph(_mm512_mask_loadu_epi16( + zmm_undefined, mask, query + dim)), + _mm512_castsi512_ph(_mm512_mask_loadu_epi16( + zmm_undefined, mask, ptrs[i] + dim)), + accs[i], mask); + } + } + + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = HorizontalAdd_FP16_V512(accs[i]); + } +} + +void compute_one_to_many_inner_product_avx512fp16_fp16_1( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_inner_product_avx512fp16_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_inner_product_avx512fp16_fp16_12( + const ailego::Float16 *query, const ailego::Float16 **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { + return compute_one_to_many_inner_product_avx512fp16_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +} +#endif + +} // namespace zvec::ailego::DistanceBatch diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc new file mode 100644 index 0000000..bead507 --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc @@ -0,0 +1,134 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX2__) + +inline float sum4(__m128 v) { + v = _mm_add_ps(v, _mm_castsi128_ps(_mm_srli_si128(_mm_castps_si128(v), 8))); + return _mm_cvtss_f32(v) + _mm_cvtss_f32(_mm_shuffle_ps(v, v, 1)); +} + +inline __m128 sum_top_bottom_avx(__m256 v) { + const __m128 high = _mm256_extractf128_ps(v, 1); + const __m128 low = _mm256_castps256_ps128(v); + return _mm_add_ps(high, low); +} + +template +static std::enable_if_t, void> +compute_one_to_many_inner_product_avx2_fp32( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, + size_t dimensionality, float *results) { + __m256 accs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm256_setzero_ps(); + } + size_t dim = 0; + for (; dim + 8 <= dimensionality; dim += 8) { + __m256 q = _mm256_loadu_ps(query + dim); + + __m256 data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm256_loadu_ps(ptrs[i] + dim); + } + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm256_fnmadd_ps(q, data_regs[i], accs[i]); + } + } + + __m128 sum128_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + sum128_regs[i] = sum_top_bottom_avx(accs[i]); + } + if (dim + 4 <= dimensionality) { + __m128 q = _mm_loadu_ps(query + dim); + + __m128 data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm_loadu_ps(ptrs[i] + dim); + } + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + for (size_t i = 0; i < dp_batch; ++i) { + sum128_regs[i] = _mm_fnmadd_ps(q, data_regs[i], sum128_regs[i]); + } + dim += 4; + } + if (dim + 2 <= dimensionality) { + __m128 q = _mm_setzero_ps(); + + __m128 data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm_setzero_ps(); + } + + q = _mm_loadh_pi(q, (const __m64 *)(query + dim)); + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm_loadh_pi(data_regs[i], (const __m64 *)(ptrs[i] + dim)); + } + for (size_t i = 0; i < dp_batch; ++i) { + sum128_regs[i] = _mm_fnmadd_ps(q, data_regs[i], sum128_regs[i]); + } + dim += 2; + } + + float res[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + res[i] = sum4(sum128_regs[i]); + } + if (dim < dimensionality) { + float q = query[dim]; + for (size_t i = 0; i < dp_batch; ++i) { + res[i] -= q * ptrs[i][dim]; + } + } + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = -res[i]; + } +} + +void compute_one_to_many_inner_product_avx2_fp32_1( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_inner_product_avx2_fp32( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_inner_product_avx2_fp32_12( + const float *query, const float **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_inner_product_avx2_fp32( + query, ptrs, prefetch_ptrs, dim, sums); +} + +#endif + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc new file mode 100644 index 0000000..66d7e15 --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc @@ -0,0 +1,104 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX2__) + +template +static std::enable_if_t, void> +compute_one_to_many_inner_product_avx2_int8( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results) { + __m256i accs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm256_setzero_si256(); + } + size_t dim = 0; + for (; dim + 32 <= dimensionality; dim += 32) { + __m256i q = _mm256_loadu_si256((const __m256i *)(query + dim)); + + __m256i data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = _mm256_loadu_si256((const __m256i *)(ptrs[i] + dim)); + } + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + __m256i q_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q)); + __m256i q_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q, 1)); + __m256i data_lo[dp_batch]; + __m256i data_hi[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_lo[i] = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(data_regs[i])); + data_hi[i] = + _mm256_cvtepi8_epi16(_mm256_extracti128_si256(data_regs[i], 1)); + } + __m256i prod_lo[dp_batch]; + __m256i prod_hi[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + prod_lo[i] = _mm256_madd_epi16(q_lo, data_lo[i]); + prod_hi[i] = _mm256_madd_epi16(q_hi, data_hi[i]); + } + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = + _mm256_add_epi32(accs[i], _mm256_add_epi32(prod_lo[i], prod_hi[i])); + } + } + + int temp_results[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + __m128i lo = _mm256_castsi256_si128(accs[i]); + __m128i hi = _mm256_extracti128_si256(accs[i], 1); + __m128i sum128 = _mm_add_epi32(lo, hi); + sum128 = _mm_hadd_epi32(sum128, sum128); + sum128 = _mm_hadd_epi32(sum128, sum128); + temp_results[i] = _mm_cvtsi128_si32(sum128); + } + for (; dim < dimensionality; ++dim) { + int8_t q = query[dim]; + for (size_t i = 0; i < dp_batch; ++i) { + temp_results[i] += q * static_cast(ptrs[i][dim]); + } + } + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = static_cast(temp_results[i]); + } +} + +void compute_one_to_many_inner_product_avx2_int8_1( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_inner_product_avx2_int8( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_inner_product_avx2_int8_12( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_inner_product_avx2_int8( + query, ptrs, prefetch_ptrs, dim, sums); +} + +#endif + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512fp16.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512fp16.cc new file mode 100644 index 0000000..96c68dc --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512fp16.cc @@ -0,0 +1,173 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512VNNI__) + +void compute_one_to_many_inner_product_avx512_vnni_int8_query_preprocess( + void *query, size_t dim) { + const int8_t *input = reinterpret_cast(query); + uint8_t *output = reinterpret_cast(query); + + // // AVX512 constant: 128 in each byte (cast to int8_t, which becomes -128 + // // in signed representation, but addition works correctly due to two's + // // complement arithmetic) + const __m512i offset = _mm512_set1_epi8(static_cast(128)); + // + size_t i = 0; + // // Process 64 bytes at a time using AVX512 + for (; i + 64 <= dim; i += 64) { + __m512i data = + _mm512_loadu_si512(reinterpret_cast(input + i)); + __m512i result = _mm512_add_epi8(data, offset); + _mm512_storeu_si512(reinterpret_cast<__m512i *>(output + i), result); + } + + // Handle remaining elements with scalar loop + for (; i < dim; ++i) { + output[i] = static_cast(static_cast(input[i]) + 128); + } +} + +// query is unsigned +template +static void compute_one_to_many_inner_product_avx512_vnni_int8( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dimensionality, + float *results) { + __m512i accs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_setzero_si512(); + } + size_t dim = 0; + for (; dim + 64 <= dimensionality; dim += 64) { + __m512i q = + _mm512_loadu_si512(reinterpret_cast(query + dim)); + + __m512i data_regs[dp_batch]; + for (size_t i = 0; i < dp_batch; ++i) { + data_regs[i] = + _mm512_loadu_si512(reinterpret_cast(ptrs[i] + dim)); + } + if (prefetch_ptrs[0]) { + for (size_t i = 0; i < dp_batch; ++i) { + ailego_prefetch(prefetch_ptrs[i] + dim); + } + } + for (size_t i = 0; i < dp_batch; ++i) { + accs[i] = _mm512_dpbusd_epi32(accs[i], q, data_regs[i]); + } + } + + int temp_results[dp_batch]{}; + for (size_t i = 0; i < dp_batch; ++i) { + temp_results[i] = _mm512_reduce_add_epi32(accs[i]); + } + for (; dim < dimensionality; ++dim) { + uint q = reinterpret_cast(query)[dim]; + for (size_t i = 0; i < dp_batch; ++i) { + temp_results[i] += q * static_cast(ptrs[i][dim]); + } + } + for (size_t i = 0; i < dp_batch; ++i) { + results[i] = static_cast(temp_results[i]); + } +} + +// +// #elif defined(__AVX512BW__) +// // TODO: this version is problematic +// template +// static std::enable_if_t, void> +// compute_one_to_many_avx512_int8( +// const int8_t *query, const int8_t **ptrs, +// std::array &prefetch_ptrs, size_t +// dimensionality, float *results) { +// std::array<__m512i, dp_batch> accs; +// size_t dim = 0; +// for (; dim + 64 <= dimensionality; dim += 64) { +// __m512i q = +// _mm512_loadu_si512(reinterpret_cast(query + dim)); +// std::array<__m512i, dp_batch> data_regs; +// for (size_t i = 0; i < dp_batch; ++i) { +// data_regs[i] = +// _mm512_loadu_si512(reinterpret_cast(ptrs[i] + +// dim)); +// } +// if (prefetch_ptrs[0]) { +// for (size_t i = 0; i < dp_batch; ++i) { +// ailego_prefetch(prefetch_ptrs[i] + dim); +// } +// } +// __m512i q_lo = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(q, 0)); +// __m512i q_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(q, 1)); +// std::array<__m512i, dp_batch> data_lo; +// std::array<__m512i, dp_batch> data_hi; +// for (size_t i = 0; i < dp_batch; ++i) { +// data_lo[i] = +// _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(data_regs[i], 0)); +// data_hi[i] = +// _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(data_regs[i], 1)); +// } +// std::array<__m512i, dp_batch> prod_lo; +// std::array<__m512i, dp_batch> prod_hi; +// for (size_t i = 0; i < dp_batch; ++i) { +// prod_lo[i] = _mm512_madd_epi16(q_lo, data_lo[i]); +// prod_hi[i] = _mm512_madd_epi16(q_hi, data_hi[i]); +// } +// for (size_t i = 0; i < dp_batch; ++i) { +// accs[i] = _mm512_add_epi32( +// accs[i], _mm512_add_epi32( +// _mm512_madd_epi16(prod_lo[i], _mm512_set1_epi16(1)), +// _mm512_madd_epi16(prod_hi[i], _mm512_set1_epi16(1)))); +// } +// } +// std::array temp_results; +// for (size_t i = 0; i < dp_batch; ++i) { +// temp_results[i] = _mm512_reduce_add_epi32(accs[i]); +// } +// for (; dim < dimensionality; ++dim) { +// int8_t q = query[dim]; +// for (size_t i = 0; i < dp_batch; ++i) { +// temp_results[i] += q * static_cast(ptrs[i][dim]); +// } +// } +// for (size_t i = 0; i < dp_batch; ++i) { +// results[i] = static_cast(temp_results[i]); +// } +// } + +void compute_one_to_many_inner_product_avx512_vnni_int8_1( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_inner_product_avx512_vnni_int8<1>( + query, ptrs, prefetch_ptrs, dim, sums); +} + +void compute_one_to_many_inner_product_avx512_vnni_int8_12( + const int8_t *query, const int8_t **ptrs, + std::array &prefetch_ptrs, size_t dim, float *sums) { + return compute_one_to_many_inner_product_avx512_vnni_int8<12>( + query, ptrs, prefetch_ptrs, dim, sums); +} + +#endif + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/parallel/lock.h b/src/ailego/parallel/lock.h new file mode 100644 index 0000000..0d19f1a --- /dev/null +++ b/src/ailego/parallel/lock.h @@ -0,0 +1,294 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#if __cplusplus >= 201703L +#include +#endif +#include +#include + +namespace zvec { +namespace ailego { + +// Test if atomic_bool is always lock free. +// Arm may be always lock free using some compiler flags, +// see https://stackoverflow.com/a/64253858/486350. +#if ATOMIC_BOOL_LOCK_FREE == 2 + +/*! Spin Mutex (The atomic type is always lock-free) + */ +class SpinMutex { + public: + //! Constructor + SpinMutex(void) {} + + //! Locking + void lock(void) { + bool expected = false; + while (!flag_.compare_exchange_weak( + expected, true, std::memory_order_acquire, std::memory_order_relaxed)) { + expected = false; + // Provide a hint to the processor that the code sequence is a spin-wait + // loop. This can help improve the performance and power consumption of + // spin-wait loops. + ailego_yield(); + } + } + + //! Try locking + bool try_lock(void) { + bool expected = false; + return flag_.compare_exchange_strong( + expected, true, std::memory_order_acquire, std::memory_order_relaxed); + } + + //! Unlocking + void unlock(void) { + flag_.store(false, std::memory_order_release); + } + + private: + //! Disable them + SpinMutex(const SpinMutex &) = delete; + SpinMutex(SpinMutex &&) = delete; + SpinMutex &operator=(const SpinMutex &) = delete; + SpinMutex &operator=(SpinMutex &&) = delete; + + //! Members + std::atomic_bool flag_{false}; +}; +#else + +/*! Spin Mutex (General) + */ +class SpinMutex { + public: + //! Constructor + SpinMutex(void) {} + + //! Locking + void lock(void) { + while (flag_.test_and_set(std::memory_order_acquire)); + } + + //! Try locking + bool try_lock(void) { + return (!flag_.test_and_set(std::memory_order_acquire)); + } + + //! Unlocking + void unlock(void) { + flag_.clear(std::memory_order_release); + } + + private: + //! Disable them + SpinMutex(const SpinMutex &) = delete; + SpinMutex(SpinMutex &&) = delete; + SpinMutex &operator=(const SpinMutex &) = delete; + SpinMutex &operator=(SpinMutex &&) = delete; + + //! Members + std::atomic_flag flag_{}; +}; +#endif // ATOMIC_BOOL_LOCK_FREE == 2 + +#if __cplusplus >= 201703L + +using SharedMutex = std::shared_mutex; + +#else + +/*! Shared Mutex + */ +class SharedMutex { + public: + //! Constructor + SharedMutex(void) {} + + //! Locking + void lock(void) { + std::unique_lock q(mutex_); + ++write_count_; + write_cond_.wait(q, [this]() { return (pending_count_ == 0); }); + --write_count_; + --pending_count_; + } + + //! Try locking + bool try_lock(void) { + std::unique_lock q(mutex_, std::defer_lock); + if (q.try_lock()) { + if (pending_count_ == 0) { + --pending_count_; + return true; + } + } + return false; + } + + //! Unlocking + void unlock(void) { + std::lock_guard q(mutex_); + ++pending_count_; + + if (write_count_ != 0) { + write_cond_.notify_one(); + } else { + read_cond_.notify_all(); + } + } + + //! Locking (shared) + void lock_shared(void) { + std::unique_lock q(mutex_); + ++read_count_; + read_cond_.wait( + q, [this]() { return (write_count_ == 0 && pending_count_ >= 0); }); + --read_count_; + ++pending_count_; + } + + //! Try locking (shared) + bool try_lock_shared(void) { + std::lock_guard q(mutex_); + if (write_count_ == 0 && pending_count_ >= 0) { + ++pending_count_; + return true; + } + return false; + } + + //! Unlocking (shared) + void unlock_shared(void) { + std::lock_guard q(mutex_); + --pending_count_; + + if (write_count_ != 0 && pending_count_ == 0) { + write_cond_.notify_one(); + } else { + read_cond_.notify_all(); + } + } + + private: + //! Disable them + SharedMutex(const SharedMutex &) = delete; + SharedMutex(SharedMutex &&) = delete; + SharedMutex &operator=(const SharedMutex &) = delete; + SharedMutex &operator=(SharedMutex &&) = delete; + + //! Members + int32_t pending_count_{0}; + int32_t read_count_{0}; + int32_t write_count_{0}; + std::mutex mutex_{}; + std::condition_variable read_cond_{}; + std::condition_variable write_cond_{}; +}; + +#endif // __cplusplus >= 201703L + +/*! Write Lock + */ +class WriteLock { + public: + //! Constructor + WriteLock(SharedMutex &mutex) : mutex_(mutex) {} + + //! Locking + void lock(void) { + mutex_.lock(); + } + + //! Try locking + bool try_lock(void) { + return mutex_.try_lock(); + } + + //! Unlocking + void unlock(void) { + mutex_.unlock(); + } + + private: + //! Disable them + WriteLock(void) = delete; + WriteLock(const WriteLock &) = delete; + WriteLock(WriteLock &&) = delete; + WriteLock &operator=(const WriteLock &) = delete; + WriteLock &operator=(WriteLock &&) = delete; + + //! Members + SharedMutex &mutex_; +}; + +/*! Read Lock + */ +class ReadLock { + public: + //! Constructor + ReadLock(SharedMutex &mutex) : mutex_(mutex) {} + + //! Locking + void lock(void) { + mutex_.lock_shared(); + } + + //! Try locking + bool try_lock(void) { + return mutex_.try_lock_shared(); + } + + //! Unlocking + void unlock(void) { + mutex_.unlock_shared(); + } + + private: + //! Disable them + ReadLock(void) = delete; + ReadLock(const ReadLock &) = delete; + ReadLock(ReadLock &&) = delete; + ReadLock &operator=(const ReadLock &) = delete; + ReadLock &operator=(ReadLock &&) = delete; + + //! Members + SharedMutex &mutex_; +}; + +/* + Atomic Close Lock + */ + +#define AILEGO_SAFE_ACCESS(CLOSE_ERR) \ + counter_.fetch_add(1); \ + AILEGO_DEFER([this] { counter_.fetch_sub(1); }); \ + if (!opened_.load()) { \ + return CLOSE_ERR; \ + } + +#define AILEGO_SAFE_CLOSE \ + opened_.store(false); \ + while (counter_.load() > 0) { \ + std::this_thread::sleep_for(std::chrono::milliseconds(1)); \ + } + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/parallel/multi_thread_list.h b/src/ailego/parallel/multi_thread_list.h new file mode 100644 index 0000000..6c2bf0b --- /dev/null +++ b/src/ailego/parallel/multi_thread_list.h @@ -0,0 +1,102 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! Multi-Thread list + */ +template +class MultiThreadList { + public: + MultiThreadList(size_t size_limit = 1000) : size_limit_(size_limit) {} + bool produce(const T &item) { + std::unique_lock lk(lock_); + not_full_.wait( + lk, [&]() { return (list_.size() < size_limit_) || done_.load(); }); + if (done_.load()) { + return false; + } + list_.emplace_back(item); + not_empty_.notify_one(); + return true; + } + + bool produce(T &&item) { + std::unique_lock lk(lock_); + not_full_.wait( + lk, [&]() { return (list_.size() < size_limit_) || done_.load(); }); + if (done_.load()) { + return false; + } + list_.emplace_back(std::move(item)); + not_empty_.notify_one(); + return true; + } + + bool consume(T *item) { + std::unique_lock lk(lock_); + not_empty_.wait(lk, [&]() { + return !list_.empty() || done_.load() || consume_stopped_.load(); + }); + if ((list_.empty() && done_.load()) || consume_stopped_.load()) { + return false; + } + *item = std::move(list_.front()); + list_.pop_front(); + not_full_.notify_one(); + return true; + } + + void done() { + std::unique_lock lk(lock_); + done_.store(true); + not_empty_.notify_all(); + not_full_.notify_all(); + } + + void reset() { + done_.store(false); + list_.clear(); + } + + void stop_consume() { + std::unique_lock lk(lock_); + consume_stopped_.store(true); + not_empty_.notify_all(); + } + + void resume_consume() { + consume_stopped_.store(false); + } + + private: + std::deque list_; + size_t size_limit_{0}; + std::mutex lock_; + std::condition_variable not_empty_, not_full_; + + std::atomic done_{false}; + std::atomic consume_stopped_{false}; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/parallel/semaphore.h b/src/ailego/parallel/semaphore.h new file mode 100644 index 0000000..89fdbef --- /dev/null +++ b/src/ailego/parallel/semaphore.h @@ -0,0 +1,181 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! Semaphore + */ +class Semaphore { + public: + //! Constructor + Semaphore(void) : Semaphore{1} {} + + //! Constructor + Semaphore(uint32_t count) : count_(count) {} + + //! Acquire a permit from this semaphore, suspending until one is available + void lock(void) { + while (!this->try_lock()) { + std::unique_lock latch(mutex_); + cond_.wait(latch, [this]() { return (count_ > 0); }); + } + } + + //! Try to acquire a permit from this semaphore without suspension + bool try_lock(void) { + uint32_t count = count_.load(std::memory_order_acquire); + return (count > 0 ? count_.compare_exchange_strong( + count, count - 1, std::memory_order_release, + std::memory_order_relaxed) + : false); + } + + //! Release a permit, returning it into this semaphore + void unlock(void) { + ++count_; + std::lock_guard latch(mutex_); + cond_.notify_one(); + } + + private: + //! Disable them + Semaphore(const Semaphore &) = delete; + Semaphore(Semaphore &&) = delete; + Semaphore &operator=(const Semaphore &) = delete; + Semaphore &operator=(Semaphore &&) = delete; + + //! Members + std::atomic count_{0}; + std::mutex mutex_{}; + std::condition_variable cond_{}; +}; + +/*! Binary Semaphores + */ +template ::type> +class BinarySemaphores { + public: + using BitwiseType = typename std::conditional< + N <= 32u, + typename std::conditional< + N <= 16u, typename std::conditional::type, + uint32_t>::type, + uint64_t>::type; + + //! Constructor + BinarySemaphores(void) : BinarySemaphores{1} {} + + //! Constructor + BinarySemaphores(uint32_t count) { + if (count == 0 || count > N) { + count = N; + } + count_ = count; + mask_ = static_cast(BitwiseType(1) << (count - 1)); + mask_ |= static_cast(mask_ - 1); + flags_.store(mask_); + } + + //! Acquire a permit from this semaphore, suspending until one is available + int acquire(void) { + int index = -1; + while ((index = this->try_acquire()) < 0) { + std::unique_lock latch(mutex_); + cond_.wait(latch, [this]() { return (flags_ > 0); }); + } + return index; + } + + //! Try to acquire a permit from this semaphore without suspension + int try_acquire(void) { + BitwiseType flags = flags_.load(std::memory_order_relaxed); + while (flags > 0) { + int index = CountTrailingZeros(flags); + if (flags_.compare_exchange_weak( + flags, flags & (~(BitwiseType(1) << index)), + std::memory_order_release, std::memory_order_relaxed)) { + return index; + } + flags = flags_.load(std::memory_order_relaxed); + } + return -1; + } + + //! Acquire a specified permit from this semaphore, suspending until index is + //! available + int acquire(int index) { + if (index < 0 || (uint32_t)index >= count_) { + return -1; + } + BitwiseType flags = flags_.load(std::memory_order_relaxed); + BitwiseType mask = BitwiseType(1) << index; + while (true) { + if ((flags & mask) && + flags_.compare_exchange_weak(flags, flags & (~mask), + std::memory_order_release, + std::memory_order_relaxed)) { + return index; + } + flags = flags_.load(std::memory_order_relaxed); + } + } + + //! Release a permit, returning it into this semaphore + void release(int index) { + flags_.fetch_or((BitwiseType(1) << index) & mask_); + std::lock_guard latch(mutex_); + cond_.notify_one(); + } + + protected: + //! Count the trailing zeros (32 bits) + template + static inline auto CountTrailingZeros(T val) -> + typename std::enable_if::type { + return ailego_ctz32(val); + } + + //! Count the trailing zeros (64 bits) + template + static inline auto CountTrailingZeros(T val) -> + typename std::enable_if::type { + return ailego_ctz64(val); + } + + private: + //! Disable them + BinarySemaphores(const BinarySemaphores &) = delete; + BinarySemaphores(BinarySemaphores &&) = delete; + BinarySemaphores &operator=(const BinarySemaphores &) = delete; + BinarySemaphores &operator=(BinarySemaphores &&) = delete; + + //! Members + uint32_t count_{0}; + BitwiseType mask_{0}; + std::atomic flags_{0}; + std::mutex mutex_{}; + std::condition_variable cond_{}; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/parallel/thread_pool.cc b/src/ailego/parallel/thread_pool.cc new file mode 100644 index 0000000..54a8970 --- /dev/null +++ b/src/ailego/parallel/thread_pool.cc @@ -0,0 +1,133 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#if (defined(__linux) || defined(__linux__)) && !defined(__ANDROID__) +#include + +static inline void BindThreads(std::vector &pool) { + uint32_t hc = std::thread::hardware_concurrency(); + if (hc > 1) { + cpu_set_t mask; + + for (size_t i = 0u; i < pool.size(); ++i) { + CPU_ZERO(&mask); + CPU_SET(i % hc, &mask); + pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask); + } + } +} + +static inline void UnbindThreads(std::vector &pool) { + cpu_set_t mask; + CPU_ZERO(&mask); + + for (size_t i = 0u; i < CPU_SETSIZE; ++i) { + CPU_SET(i, &mask); + } + for (size_t i = 0u; i < pool.size(); ++i) { + pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask); + } +} +#else +static inline void BindThreads(std::vector &) {} +static inline void UnbindThreads(std::vector &) {} +#endif + +namespace zvec { +namespace ailego { + +ThreadPool::ThreadPool(uint32_t size, bool binding) { + for (uint32_t i = 0u; i < size; ++i) { + pool_.emplace_back(&ThreadPool::worker, this); + } + if (binding) { + this->bind(); + } +} + +void ThreadPool::bind(void) { + BindThreads(pool_); +} + +void ThreadPool::unbind(void) { + UnbindThreads(pool_); +} + +void ThreadPool::worker(void) { + // Counter of workers + ++worker_count_; + + ThreadPool::Task task; + while (this->picking(&task)) { + // Run the task + task.handle->run(); + task.handle = nullptr; + + // Notify task finished + if (task.control) { + task.control->notify(); + } + + // Notify task group + if (task.group) { + task.group->notify(); + task.group = nullptr; + } + + // Decrease count of active works + std::lock_guard lock(wait_mutex_); + if (--active_count_ == 0 && pending_count_ == 0) { + finished_cond_.notify_all(); + } + } + + // Decrease count of workers + std::lock_guard lock(wait_mutex_); + if (--worker_count_ == 0) { + stopped_cond_.notify_all(); + } +} + +bool ThreadPool::picking(ThreadPool::Task *task) { + std::unique_lock latch(queue_mutex_); + work_cond_.wait(latch, + [this]() { return (pending_count_ > 0 || stopping_); }); + if (stopping_) { + return false; + } + + // Pop a task + auto &head = queue_.front(); + task->control = head.control; + task->group = std::move(head.group); + task->handle = std::move(head.handle); + queue_.pop(); + + // Update group control + if (task->group) { + task->group->mark_task_actived(); + } + + // Counter of active tasks + std::unique_lock lock(wait_mutex_); + ++active_count_; + --pending_count_; + + return true; +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/pattern/defer.h b/src/ailego/pattern/defer.h new file mode 100644 index 0000000..9156852 --- /dev/null +++ b/src/ailego/pattern/defer.h @@ -0,0 +1,24 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "scope_guard.h" + +#define AILEGO_DEFER_NAME_(x, y) x##y +#define AILEGO_DEFER_NAME(x) AILEGO_DEFER_NAME_(__ailegoDefer_, x) + +//! Defer operator +#define AILEGO_DEFER(...) \ + auto AILEGO_DEFER_NAME(__LINE__) = ailego::ScopeGuard::Make(__VA_ARGS__) diff --git a/src/ailego/pattern/scope_guard.h b/src/ailego/pattern/scope_guard.h new file mode 100644 index 0000000..075fd39 --- /dev/null +++ b/src/ailego/pattern/scope_guard.h @@ -0,0 +1,177 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +/*! Scope Guard Implementation + */ +template +class ScopeGuardImpl { + public: + using Object = CallbackObject; + using Functor = CallbackFunctor; + + //! Constructor + ScopeGuardImpl(ScopeGuardImpl &&rhs) + : obj_(rhs.obj_), + impl_(std::move(rhs.impl_)), + tuple_(std::move(rhs.tuple_)) { + rhs.obj_ = nullptr; + } + + //! Constructor + template + ScopeGuardImpl(typename Object::Type *obj, const typename Functor::Type &impl, + TArgs &&...args) + : obj_(obj), impl_(impl), tuple_(std::forward(args)...) {} + + //! Constructor + template + ScopeGuardImpl(typename Object::Type *obj, typename Functor::Type &&impl, + TArgs &&...args) + : obj_(obj), + impl_(std::move(impl)), + tuple_(std::forward(args)...) {} + + // Destructor + ~ScopeGuardImpl(void) { + if (obj_) { + Functor::Run(obj_, impl_, tuple_); + } + } + + protected: + //! Disable them + ScopeGuardImpl(void) = delete; + ScopeGuardImpl(const ScopeGuardImpl &) = delete; + ScopeGuardImpl &operator=(const ScopeGuardImpl &) = delete; + + private: + //! Members + typename Object::Type *obj_; + typename Functor::Type impl_; + typename Functor::TupleType tuple_; +}; + +/*! Scope Guard Implementation (void, TFunc) + */ +template +class ScopeGuardImpl { + public: + //! Callback Functor Type + using Functor = CallbackFunctor; + + //! Constructor + ScopeGuardImpl(ScopeGuardImpl &&rhs) + : impl_(std::move(rhs.impl_)), + tuple_(std::move(rhs.tuple_)), + valid_(rhs.valid_) { + rhs.valid_ = false; + } + + //! Constructor + template + ScopeGuardImpl(const typename Functor::Type &impl, TArgs &&...args) + : impl_(impl), tuple_(std::forward(args)...), valid_(true) {} + + //! Constructor + template + ScopeGuardImpl(typename Functor::Type &&impl, TArgs &&...args) + : impl_(std::move(impl)), + tuple_(std::forward(args)...), + valid_(true) {} + + // Destructor + ~ScopeGuardImpl(void) { + if (valid_) { + Functor::Run(impl_, tuple_); + } + } + + protected: + //! Disable them + ScopeGuardImpl(void) = delete; + ScopeGuardImpl(const ScopeGuardImpl &) = delete; + ScopeGuardImpl &operator=(const ScopeGuardImpl &) = delete; + + private: + //! Members + typename Functor::Type impl_; + typename Functor::TupleType tuple_; + bool valid_; +}; + +/*! Scope Guard + */ +struct ScopeGuard { + //! Make a scope guard object (member function pointer) + template + static inline auto Make(T *obj, R (T::*impl)(TParams...), TArgs &&...args) + -> ScopeGuardImpl::Type> { + return ScopeGuardImpl::Type>( + obj, impl, std::forward(args)...); + } + + //! Make a scope guard object (constable member function pointer) + template + static inline auto Make(const T *obj, R (T::*impl)(TParams...) const, + TArgs &&...args) + -> ScopeGuardImpl::Type> { + return ScopeGuardImpl::Type>( + obj, impl, std::forward(args)...); + } + + //! Make a scope guard object (volatile member function pointer) + template + static inline auto Make(volatile T *obj, R (T::*impl)(TParams...) volatile, + TArgs &&...args) + -> ScopeGuardImpl::Type> { + return ScopeGuardImpl::Type>( + obj, impl, std::forward(args)...); + } + + //! Make a scope guard object (constable volatile member function pointer) + template + static inline auto Make(const volatile T *obj, + R (T::*impl)(TParams...) const volatile, + TArgs &&...args) + -> ScopeGuardImpl::Type> { + return ScopeGuardImpl::Type>( + obj, impl, std::forward(args)...); + } + + //! Make a scope guard object (function) + template < + typename TFunc, typename... TArgs, + typename = typename std::enable_if::Value>::type> + static inline auto Make(TFunc &&impl, TArgs &&...args) + -> ScopeGuardImpl::Type> { + return ScopeGuardImpl::Type>( + std::forward(impl), std::forward(args)...); + } +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/bit_string_helper.h b/src/ailego/utility/bit_string_helper.h new file mode 100644 index 0000000..a50e1ee --- /dev/null +++ b/src/ailego/utility/bit_string_helper.h @@ -0,0 +1,118 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec { + +namespace ailego { + +class BitStringWriter { + public: + BitStringWriter(uint8_t *buffer, size_t buffer_size) + : buffer_(buffer), buffer_size_(buffer_size), offset_(0) { + ::memset(buffer_, 0, buffer_size_); + } + + bool write(uint64_t data, int nbit) { + if (buffer_size_ * 8 < nbit + offset_) { + return false; + } + + int bits_remain = 8 - (offset_ & 7); + + if (nbit <= bits_remain) { + buffer_[offset_ >> 3] |= data << (offset_ & 7); + offset_ += nbit; + } else { + size_t j = offset_ >> 3; + buffer_[j++] |= data << (offset_ & 7); + offset_ += nbit; + data >>= bits_remain; + while (data != 0) { + buffer_[j++] |= data; + data >>= 8; + } + } + + return true; + } + + size_t offset() { + return offset_; + } + + private: + uint8_t *buffer_; + size_t buffer_size_; + size_t offset_; +}; + +class BitStringReader { + public: + BitStringReader(const uint8_t *buffer, size_t buffer_size) + : buffer_(buffer), buffer_size_(buffer_size), offset_(0) {} + + bool read(uint64_t &data, int nbit) { + if (buffer_size_ * 8 < nbit + offset_) { + return false; + } + + int bits_remain = 8 - (offset_ & 7); + + uint64_t result = buffer_[offset_ >> 3] >> (offset_ & 7); + if (nbit <= bits_remain) { + result &= (1 << nbit) - 1; + offset_ += nbit; + + data = result; + } else { + int temp = bits_remain; + size_t i = (offset_ >> 3) + 1; + offset_ += nbit; + nbit -= bits_remain; + + while (nbit > 8) { + result |= ((uint64_t)buffer_[i++]) << temp; + temp += 8; + nbit -= 8; + } + + uint64_t last_byte = buffer_[i]; + + last_byte &= (1 << nbit) - 1; + result |= last_byte << temp; + + data = result; + } + + return true; + } + + size_t offset() { + return offset_; + } + + private: + const uint8_t *buffer_; + size_t buffer_size_; + size_t offset_; +}; + +} // namespace ailego + +} // namespace zvec diff --git a/src/ailego/utility/bitset_helper.cc b/src/ailego/utility/bitset_helper.cc new file mode 100644 index 0000000..19be348 --- /dev/null +++ b/src/ailego/utility/bitset_helper.cc @@ -0,0 +1,2049 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "bitset_helper.h" +#include + +#ifndef __SSE4_2__ +#define bitset_popcount32 ailego_popcount32 +#define bitset_popcount64 ailego_popcount64 +#else +#define bitset_popcount32 _mm_popcnt_u32 +#define bitset_popcount64 _mm_popcnt_u64 +#endif // !__SSE4_2__ + +#if defined(__ARM_NEON) +static inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + vst1q_u32(lhs, vandq_u32(vld1q_u32(lhs), vld1q_u32(rhs))); + } + switch (last - last_aligned) { + case 3: + lhs[2] &= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= rhs[0]; + } +} + +static inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs, + size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + vst1q_u32(lhs, vbicq_u32(vld1q_u32(lhs), vld1q_u32(rhs))); + } + switch (last - last_aligned) { + case 3: + lhs[2] &= ~rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= ~rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= ~rhs[0]; + } +} + +static inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + vst1q_u32(lhs, vorrq_u32(vld1q_u32(lhs), vld1q_u32(rhs))); + } + switch (last - last_aligned) { + case 3: + lhs[2] |= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] |= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] |= rhs[0]; + } +} + +static inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + vst1q_u32(lhs, veorq_u32(vld1q_u32(lhs), vld1q_u32(rhs))); + } + switch (last - last_aligned) { + case 3: + lhs[2] ^= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] ^= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] ^= rhs[0]; + } +} + +static inline void bitset_not(uint32_t *lhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + static const uint32x4_t v_zero = vdupq_n_u32(0); + + for (; lhs != last_aligned; lhs += 4) { + vst1q_u32(lhs, vornq_u32(v_zero, vld1q_u32(lhs))); + } + switch (last - last_aligned) { + case 3: + lhs[2] = ~lhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] = ~lhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] = ~lhs[0]; + } +} + +static inline bool bitset_test_all(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4) { + uint64x2_t vu64 = vld1q_u64((const uint64_t *)lhs); + if ((vgetq_lane_u64(vu64, 0) & vgetq_lane_u64(vu64, 1)) != (uint64_t)-1) { + return false; + } + } + switch (last - last_aligned) { + case 3: + if (lhs[2] != 0xffffffffu) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0xffffffffu) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0xffffffffu) { + return false; + } + } + return true; +} + +static inline bool bitset_test_any(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4) { + uint64x2_t vu64 = vld1q_u64((const uint64_t *)lhs); + if (vgetq_lane_u64(vu64, 0) | vgetq_lane_u64(vu64, 1)) { + return true; + } + } + switch (last - last_aligned) { + case 3: + if (lhs[2] != 0u) { + return true; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return true; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return true; + } + } + return false; +} + +static inline bool bitset_test_none(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4) { + uint64x2_t vu64 = vld1q_u64((const uint64_t *)lhs); + if (vgetq_lane_u64(vu64, 0) | vgetq_lane_u64(vu64, 1)) { + return false; + } + } + switch (last - last_aligned) { + case 3: + if (lhs[2] != 0u) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return false; + } + } + return true; +} + +#elif defined(__AVX2__) +static inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m256i ymm0 = _mm256_load_si256((__m256i *)lhs); + __m256i ymm1 = _mm256_load_si256((__m256i *)rhs); + _mm256_store_si256((__m256i *)lhs, _mm256_and_si256(ymm1, ymm0)); + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + __m128i xmm1 = _mm_load_si128((__m128i *)rhs); + _mm_store_si128((__m128i *)lhs, _mm_and_si128(xmm1, xmm0)); + lhs += 4; + rhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs); + __m256i ymm1 = _mm256_loadu_si256((__m256i *)rhs); + _mm256_storeu_si256((__m256i *)lhs, _mm256_and_si256(ymm1, ymm0)); + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs); + _mm_storeu_si128((__m128i *)lhs, _mm_and_si128(xmm1, xmm0)); + lhs += 4; + rhs += 4; + } + } + switch (last - lhs) { + case 3: + lhs[2] &= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= rhs[0]; + } +} + +static inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs, + size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m256i ymm0 = _mm256_load_si256((__m256i *)lhs); + __m256i ymm1 = _mm256_load_si256((__m256i *)rhs); + _mm256_store_si256((__m256i *)lhs, _mm256_andnot_si256(ymm1, ymm0)); + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + __m128i xmm1 = _mm_load_si128((__m128i *)rhs); + _mm_store_si128((__m128i *)lhs, _mm_andnot_si128(xmm1, xmm0)); + lhs += 4; + rhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs); + __m256i ymm1 = _mm256_loadu_si256((__m256i *)rhs); + _mm256_storeu_si256((__m256i *)lhs, _mm256_andnot_si256(ymm1, ymm0)); + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs); + _mm_storeu_si128((__m128i *)lhs, _mm_andnot_si128(xmm1, xmm0)); + lhs += 4; + rhs += 4; + } + } + switch (last - lhs) { + case 3: + lhs[2] &= ~rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= ~rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= ~rhs[0]; + } +} + +static inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m256i ymm0 = _mm256_load_si256((__m256i *)lhs); + __m256i ymm1 = _mm256_load_si256((__m256i *)rhs); + _mm256_store_si256((__m256i *)lhs, _mm256_or_si256(ymm1, ymm0)); + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + __m128i xmm1 = _mm_load_si128((__m128i *)rhs); + _mm_store_si128((__m128i *)lhs, _mm_or_si128(xmm1, xmm0)); + lhs += 4; + rhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs); + __m256i ymm1 = _mm256_loadu_si256((__m256i *)rhs); + _mm256_storeu_si256((__m256i *)lhs, _mm256_or_si256(ymm1, ymm0)); + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs); + _mm_storeu_si128((__m128i *)lhs, _mm_or_si128(xmm1, xmm0)); + lhs += 4; + rhs += 4; + } + } + switch (last - lhs) { + case 3: + lhs[2] |= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] |= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] |= rhs[0]; + } +} + +static inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m256i ymm0 = _mm256_load_si256((__m256i *)lhs); + __m256i ymm1 = _mm256_load_si256((__m256i *)rhs); + _mm256_store_si256((__m256i *)lhs, _mm256_xor_si256(ymm1, ymm0)); + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + __m128i xmm1 = _mm_load_si128((__m128i *)rhs); + _mm_store_si128((__m128i *)lhs, _mm_xor_si128(xmm1, xmm0)); + lhs += 4; + rhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs); + __m256i ymm1 = _mm256_loadu_si256((__m256i *)rhs); + _mm256_storeu_si256((__m256i *)lhs, _mm256_xor_si256(ymm1, ymm0)); + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs); + _mm_storeu_si128((__m128i *)lhs, _mm_xor_si128(xmm1, xmm0)); + lhs += 4; + rhs += 4; + } + } + switch (last - lhs) { + case 3: + lhs[2] ^= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] ^= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] ^= rhs[0]; + } +} + +static inline void bitset_not(uint32_t *lhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + static const __m256i mask_256 = _mm256_set1_epi32(0xffffffffu); + static const __m128i mask_128 = _mm_set1_epi32(0xffffffffu); + + if (((uintptr_t)lhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 8) { + _mm256_store_si256( + (__m256i *)lhs, + _mm256_andnot_si256(_mm256_load_si256((__m256i *)lhs), mask_256)); + } + if (last >= last_aligned + 4) { + _mm_store_si128( + (__m128i *)lhs, + _mm_andnot_si128(_mm_load_si128((__m128i *)lhs), mask_128)); + lhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8) { + _mm256_storeu_si256( + (__m256i *)lhs, + _mm256_andnot_si256(_mm256_loadu_si256((__m256i *)lhs), mask_256)); + } + if (last >= last_aligned + 4) { + _mm_storeu_si128( + (__m128i *)lhs, + _mm_andnot_si128(_mm_lddqu_si128((__m128i *)lhs), mask_128)); + lhs += 4; + } + } + switch (last - lhs) { + case 3: + lhs[2] = ~lhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] = ~lhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] = ~lhs[0]; + } +} + +static inline bool bitset_test_all(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + static const __m256i mask_256 = _mm256_set1_epi32(0xffffffffu); + static const __m128i mask_128 = _mm_set1_epi32(0xffffffffu); + + if (((uintptr_t)lhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 8) { + __m256i neq = + _mm256_xor_si256(_mm256_load_si256((__m256i *)lhs), mask_256); + if (!_mm256_testz_si256(neq, neq)) { + return false; + } + } + if (last >= last_aligned + 4) { + __m128i neq = _mm_xor_si128(_mm_load_si128((__m128i *)lhs), mask_128); + if (!_mm_testz_si128(neq, neq)) { + return false; + } + lhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8) { + __m256i neq = + _mm256_xor_si256(_mm256_loadu_si256((__m256i *)lhs), mask_256); + if (!_mm256_testz_si256(neq, neq)) { + return false; + } + } + if (last >= last_aligned + 4) { + __m128i neq = _mm_xor_si128(_mm_lddqu_si128((__m128i *)lhs), mask_128); + if (!_mm_testz_si128(neq, neq)) { + return false; + } + lhs += 4; + } + } + switch (last - lhs) { + case 3: + if (lhs[2] != 0xffffffffu) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0xffffffffu) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0xffffffffu) { + return false; + } + } + return true; +} + +static inline bool bitset_test_any(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + if (((uintptr_t)lhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 8) { + __m256i ymm0 = _mm256_load_si256((__m256i *)lhs); + if (!_mm256_testz_si256(ymm0, ymm0)) { + return true; + } + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + if (!_mm_testz_si128(xmm0, xmm0)) { + return true; + } + lhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8) { + __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs); + if (!_mm256_testz_si256(ymm0, ymm0)) { + return true; + } + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + if (!_mm_testz_si128(xmm0, xmm0)) { + return true; + } + lhs += 4; + } + } + switch (last - lhs) { + case 3: + if (lhs[2] != 0u) { + return true; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return true; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return true; + } + } + return false; +} + +static inline bool bitset_test_none(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + if (((uintptr_t)lhs & 0x1f) == 0) { + for (; lhs != last_aligned; lhs += 8) { + __m256i ymm0 = _mm256_load_si256((__m256i *)lhs); + if (!_mm256_testz_si256(ymm0, ymm0)) { + return false; + } + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + if (!_mm_testz_si128(xmm0, xmm0)) { + return false; + } + lhs += 4; + } + } else { + for (; lhs != last_aligned; lhs += 8) { + __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs); + if (!_mm256_testz_si256(ymm0, ymm0)) { + return false; + } + } + if (last >= last_aligned + 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + if (!_mm_testz_si128(xmm0, xmm0)) { + return false; + } + lhs += 4; + } + } + switch (last - lhs) { + case 3: + if (lhs[2] != 0u) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return false; + } + } + return true; +} + +#elif defined(__SSE2__) +#ifndef __SSE3__ +#define _mm_lddqu_si128 _mm_loadu_si128 +#endif // !__SSE3__ + +static inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + __m128i xmm1 = _mm_load_si128((__m128i *)rhs); + _mm_store_si128((__m128i *)lhs, _mm_and_si128(xmm1, xmm0)); + } + } else { + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs); + _mm_storeu_si128((__m128i *)lhs, _mm_and_si128(xmm1, xmm0)); + } + } + switch (last - last_aligned) { + case 3: + lhs[2] &= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= rhs[0]; + } +} + +static inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs, + size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + __m128i xmm1 = _mm_load_si128((__m128i *)rhs); + _mm_store_si128((__m128i *)lhs, _mm_andnot_si128(xmm1, xmm0)); + } + } else { + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs); + _mm_storeu_si128((__m128i *)lhs, _mm_andnot_si128(xmm1, xmm0)); + } + } + switch (last - last_aligned) { + case 3: + lhs[2] &= ~rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= ~rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= ~rhs[0]; + } +} + +static inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + __m128i xmm1 = _mm_load_si128((__m128i *)rhs); + _mm_store_si128((__m128i *)lhs, _mm_or_si128(xmm1, xmm0)); + } + } else { + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs); + _mm_storeu_si128((__m128i *)lhs, _mm_or_si128(xmm1, xmm0)); + } + } + switch (last - last_aligned) { + case 3: + lhs[2] |= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] |= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] |= rhs[0]; + } +} + +static inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + __m128i xmm1 = _mm_load_si128((__m128i *)rhs); + _mm_store_si128((__m128i *)lhs, _mm_xor_si128(xmm1, xmm0)); + } + } else { + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs); + _mm_storeu_si128((__m128i *)lhs, _mm_xor_si128(xmm1, xmm0)); + } + } + switch (last - last_aligned) { + case 3: + lhs[2] ^= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] ^= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] ^= rhs[0]; + } +} + +static inline void bitset_not(uint32_t *lhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + static const __m128i mask = _mm_set1_epi32(0xffffffffu); + + if (((uintptr_t)lhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4) { + _mm_store_si128((__m128i *)lhs, + _mm_andnot_si128(_mm_load_si128((__m128i *)lhs), mask)); + } + } else { + for (; lhs != last_aligned; lhs += 4) { + _mm_storeu_si128((__m128i *)lhs, + _mm_andnot_si128(_mm_lddqu_si128((__m128i *)lhs), mask)); + } + } + switch (last - last_aligned) { + case 3: + lhs[2] = ~lhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] = ~lhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] = ~lhs[0]; + } +} + +static inline bool bitset_test_all(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + static const __m128i mask = _mm_set1_epi32(0xffffffffu); + +#ifndef __SSE4_1__ + if (((uintptr_t)lhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4) { + __m128i eq = _mm_cmpeq_epi32(_mm_load_si128((__m128i *)lhs), mask); + if (_mm_movemask_epi8(eq) != 0xffffu) { + return false; + } + } + } else { + for (; lhs != last_aligned; lhs += 4) { + __m128i eq = _mm_cmpeq_epi32(_mm_lddqu_si128((__m128i *)lhs), mask); + if (_mm_movemask_epi8(eq) != 0xffffu) { + return false; + } + } + } +#else + if (((uintptr_t)lhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4) { + __m128i neq = _mm_xor_si128(_mm_load_si128((__m128i *)lhs), mask); + if (!_mm_testz_si128(neq, neq)) { + return false; + } + } + } else { + for (; lhs != last_aligned; lhs += 4) { + __m128i neq = _mm_xor_si128(_mm_lddqu_si128((__m128i *)lhs), mask); + if (!_mm_testz_si128(neq, neq)) { + return false; + } + } + } +#endif // !__SSE4_1__ + + switch (last - last_aligned) { + case 3: + if (lhs[2] != 0xffffffffu) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0xffffffffu) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0xffffffffu) { + return false; + } + } + return true; +} + +static inline bool bitset_test_any(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + +#ifndef __SSE4_1__ + static const __m128i zero = _mm_setzero_si128(); + + if (((uintptr_t)lhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4) { + __m128i eq = _mm_cmpeq_epi32(_mm_load_si128((__m128i *)lhs), zero); + if (_mm_movemask_epi8(eq) != 0xffffu) { + return true; + } + } + } else { + for (; lhs != last_aligned; lhs += 4) { + __m128i eq = _mm_cmpeq_epi32(_mm_lddqu_si128((__m128i *)lhs), zero); + if (_mm_movemask_epi8(eq) != 0xffffu) { + return true; + } + } + } +#else + if (((uintptr_t)lhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + if (!_mm_testz_si128(xmm0, xmm0)) { + return true; + } + } + } else { + for (; lhs != last_aligned; lhs += 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + if (!_mm_testz_si128(xmm0, xmm0)) { + return true; + } + } + } +#endif // !__SSE4_1__ + + switch (last - last_aligned) { + case 3: + if (lhs[2] != 0u) { + return true; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return true; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return true; + } + } + return false; +} + +static inline bool bitset_test_none(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + +#ifndef __SSE4_1__ + static __m128i zero = _mm_setzero_si128(); + + if (((uintptr_t)lhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4) { + __m128i eq = _mm_cmpeq_epi32(_mm_load_si128((__m128i *)lhs), zero); + if (_mm_movemask_epi8(eq) != 0xffffu) { + return false; + } + } + } else { + for (; lhs != last_aligned; lhs += 4) { + __m128i eq = _mm_cmpeq_epi32(_mm_lddqu_si128((__m128i *)lhs), zero); + if (_mm_movemask_epi8(eq) != 0xffffu) { + return false; + } + } + } +#else + if (((uintptr_t)lhs & 0xf) == 0) { + for (; lhs != last_aligned; lhs += 4) { + __m128i xmm0 = _mm_load_si128((__m128i *)lhs); + if (!_mm_testz_si128(xmm0, xmm0)) { + return false; + } + } + } else { + for (; lhs != last_aligned; lhs += 4) { + __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs); + if (!_mm_testz_si128(xmm0, xmm0)) { + return false; + } + } + } +#endif // !__SSE4_1__ + + switch (last - last_aligned) { + case 3: + if (lhs[2] != 0u) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return false; + } + } + return true; +} + +#else +#if defined(AILEGO_M64) +static inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + *(uint64_t *)(&lhs[6]) &= *(uint64_t *)(&rhs[6]); + *(uint64_t *)(&lhs[4]) &= *(uint64_t *)(&rhs[4]); + *(uint64_t *)(&lhs[2]) &= *(uint64_t *)(&rhs[2]); + *(uint64_t *)(&lhs[0]) &= *(uint64_t *)(&rhs[0]); + } + switch (last - last_aligned) { + case 7: + lhs[6] &= rhs[6]; + /* FALLTHRU */ + case 6: + lhs[5] &= rhs[5]; + /* FALLTHRU */ + case 5: + lhs[4] &= rhs[4]; + /* FALLTHRU */ + case 4: + lhs[3] &= rhs[3]; + /* FALLTHRU */ + case 3: + lhs[2] &= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= rhs[0]; + } +} + +static inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs, + size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + *(uint64_t *)(&lhs[6]) &= ~(*(uint64_t *)(&rhs[6])); + *(uint64_t *)(&lhs[4]) &= ~(*(uint64_t *)(&rhs[4])); + *(uint64_t *)(&lhs[2]) &= ~(*(uint64_t *)(&rhs[2])); + *(uint64_t *)(&lhs[0]) &= ~(*(uint64_t *)(&rhs[0])); + } + switch (last - last_aligned) { + case 7: + lhs[6] &= ~rhs[6]; + /* FALLTHRU */ + case 6: + lhs[5] &= ~rhs[5]; + /* FALLTHRU */ + case 5: + lhs[4] &= ~rhs[4]; + /* FALLTHRU */ + case 4: + lhs[3] &= ~rhs[3]; + /* FALLTHRU */ + case 3: + lhs[2] &= ~rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= ~rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= ~rhs[0]; + } +} + +static inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + *(uint64_t *)(&lhs[6]) |= *(uint64_t *)(&rhs[6]); + *(uint64_t *)(&lhs[4]) |= *(uint64_t *)(&rhs[4]); + *(uint64_t *)(&lhs[2]) |= *(uint64_t *)(&rhs[2]); + *(uint64_t *)(&lhs[0]) |= *(uint64_t *)(&rhs[0]); + } + switch (last - last_aligned) { + case 7: + lhs[6] |= rhs[6]; + /* FALLTHRU */ + case 6: + lhs[5] |= rhs[5]; + /* FALLTHRU */ + case 5: + lhs[4] |= rhs[4]; + /* FALLTHRU */ + case 4: + lhs[3] |= rhs[3]; + /* FALLTHRU */ + case 3: + lhs[2] |= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] |= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] |= rhs[0]; + } +} + +static inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + *(uint64_t *)(&lhs[6]) ^= *(uint64_t *)(&rhs[6]); + *(uint64_t *)(&lhs[4]) ^= *(uint64_t *)(&rhs[4]); + *(uint64_t *)(&lhs[2]) ^= *(uint64_t *)(&rhs[2]); + *(uint64_t *)(&lhs[0]) ^= *(uint64_t *)(&rhs[0]); + } + switch (last - last_aligned) { + case 7: + lhs[6] ^= rhs[6]; + /* FALLTHRU */ + case 6: + lhs[5] ^= rhs[5]; + /* FALLTHRU */ + case 5: + lhs[4] ^= rhs[4]; + /* FALLTHRU */ + case 4: + lhs[3] ^= rhs[3]; + /* FALLTHRU */ + case 3: + lhs[2] ^= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] ^= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] ^= rhs[0]; + } +} + +static inline void bitset_not(uint32_t *lhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + for (; lhs != last_aligned; lhs += 8) { + *(uint64_t *)(&lhs[6]) = ~(*(uint64_t *)(&lhs[6])); + *(uint64_t *)(&lhs[4]) = ~(*(uint64_t *)(&lhs[4])); + *(uint64_t *)(&lhs[2]) = ~(*(uint64_t *)(&lhs[2])); + *(uint64_t *)(&lhs[0]) = ~(*(uint64_t *)(&lhs[0])); + } + switch (last - last_aligned) { + case 7: + lhs[6] = ~lhs[6]; + /* FALLTHRU */ + case 6: + lhs[5] = ~lhs[5]; + /* FALLTHRU */ + case 5: + lhs[4] = ~lhs[4]; + /* FALLTHRU */ + case 4: + lhs[3] = ~lhs[3]; + /* FALLTHRU */ + case 3: + lhs[2] = ~lhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] = ~lhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] = ~lhs[0]; + } +} + +static inline bool bitset_test_all(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + for (; lhs != last_aligned; lhs += 8) { + if (*(uint64_t *)(&lhs[6]) != (uint64_t)-1) { + return false; + } + if (*(uint64_t *)(&lhs[4]) != (uint64_t)-1) { + return false; + } + if (*(uint64_t *)(&lhs[2]) != (uint64_t)-1) { + return false; + } + if (*(uint64_t *)(&lhs[0]) != (uint64_t)-1) { + return false; + } + } + switch (last - last_aligned) { + case 7: + if (lhs[6] != (uint32_t)-1) { + return false; + } + /* FALLTHRU */ + case 6: + if (lhs[5] != (uint32_t)-1) { + return false; + } + /* FALLTHRU */ + case 5: + if (lhs[4] != (uint32_t)-1) { + return false; + } + /* FALLTHRU */ + case 4: + if (lhs[3] != (uint32_t)-1) { + return false; + } + /* FALLTHRU */ + case 3: + if (lhs[2] != (uint32_t)-1) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != (uint32_t)-1) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != (uint32_t)-1) { + return false; + } + } + return true; +} + +static inline bool bitset_test_any(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + for (; lhs != last_aligned; lhs += 8) { + if (*(uint64_t *)(&lhs[6]) != 0u) { + return true; + } + if (*(uint64_t *)(&lhs[4]) != 0u) { + return true; + } + if (*(uint64_t *)(&lhs[2]) != 0u) { + return true; + } + if (*(uint64_t *)(&lhs[0]) != 0u) { + return true; + } + } + switch (last - last_aligned) { + case 7: + if (lhs[6] != 0u) { + return true; + } + /* FALLTHRU */ + case 6: + if (lhs[5] != 0u) { + return true; + } + /* FALLTHRU */ + case 5: + if (lhs[4] != 0u) { + return true; + } + /* FALLTHRU */ + case 4: + if (lhs[3] != 0u) { + return true; + } + /* FALLTHRU */ + case 3: + if (lhs[2] != 0u) { + return true; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return true; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return true; + } + } + return false; +} + +static inline bool bitset_test_none(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + + for (; lhs != last_aligned; lhs += 8) { + if (*(uint64_t *)(&lhs[6]) != 0u) { + return false; + } + if (*(uint64_t *)(&lhs[4]) != 0u) { + return false; + } + if (*(uint64_t *)(&lhs[2]) != 0u) { + return false; + } + if (*(uint64_t *)(&lhs[0]) != 0u) { + return false; + } + } + switch (last - last_aligned) { + case 7: + if (lhs[6] != 0u) { + return false; + } + /* FALLTHRU */ + case 6: + if (lhs[5] != 0u) { + return false; + } + /* FALLTHRU */ + case 5: + if (lhs[4] != 0u) { + return false; + } + /* FALLTHRU */ + case 4: + if (lhs[3] != 0u) { + return false; + } + /* FALLTHRU */ + case 3: + if (lhs[2] != 0u) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return false; + } + } + return true; +} + +#else // AILEGO_M64 +static inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + lhs[3] &= rhs[3]; + lhs[2] &= rhs[2]; + lhs[1] &= rhs[1]; + lhs[0] &= rhs[0]; + } + switch (last - last_aligned) { + case 3: + lhs[2] &= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= rhs[0]; + } +} + +static inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs, + size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + lhs[3] &= ~rhs[3]; + lhs[2] &= ~rhs[2]; + lhs[1] &= ~rhs[1]; + lhs[0] &= ~rhs[0]; + } + switch (last - last_aligned) { + case 3: + lhs[2] &= ~rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] &= ~rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] &= ~rhs[0]; + } +} + +static inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + lhs[3] |= rhs[3]; + lhs[2] |= rhs[2]; + lhs[1] |= rhs[1]; + lhs[0] |= rhs[0]; + } + switch (last - last_aligned) { + case 3: + lhs[2] |= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] |= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] |= rhs[0]; + } +} + +static inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + lhs[3] ^= rhs[3]; + lhs[2] ^= rhs[2]; + lhs[1] ^= rhs[1]; + lhs[0] ^= rhs[0]; + } + switch (last - last_aligned) { + case 3: + lhs[2] ^= rhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] ^= rhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] ^= rhs[0]; + } +} + +static inline void bitset_not(uint32_t *lhs, size_t size) { + uint32_t *last = lhs + size; + uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4) { + lhs[3] = ~lhs[3]; + lhs[2] = ~lhs[2]; + lhs[1] = ~lhs[1]; + lhs[0] = ~lhs[0]; + } + switch (last - last_aligned) { + case 3: + lhs[2] = ~lhs[2]; + /* FALLTHRU */ + case 2: + lhs[1] = ~lhs[1]; + /* FALLTHRU */ + case 1: + lhs[0] = ~lhs[0]; + } +} + +static inline bool bitset_test_all(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4) { + if (lhs[3] != (uint32_t)-1) { + return false; + } + if (lhs[2] != (uint32_t)-1) { + return false; + } + if (lhs[1] != (uint32_t)-1) { + return false; + } + if (lhs[0] != (uint32_t)-1) { + return false; + } + } + switch (last - last_aligned) { + case 3: + if (lhs[2] != (uint32_t)-1) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != (uint32_t)-1) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != (uint32_t)-1) { + return false; + } + } + return true; +} + +static inline bool bitset_test_any(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4) { + if (lhs[3] != 0u) { + return true; + } + if (lhs[2] != 0u) { + return true; + } + if (lhs[1] != 0u) { + return true; + } + if (lhs[0] != 0u) { + return true; + } + } + switch (last - last_aligned) { + case 3: + if (lhs[2] != 0u) { + return true; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return true; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return true; + } + } + return false; +} + +static inline bool bitset_test_none(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + + for (; lhs != last_aligned; lhs += 4) { + if (lhs[3] != 0u) { + return false; + } + if (lhs[2] != 0u) { + return false; + } + if (lhs[1] != 0u) { + return false; + } + if (lhs[0] != 0u) { + return false; + } + } + switch (last - last_aligned) { + case 3: + if (lhs[2] != 0u) { + return false; + } + /* FALLTHRU */ + case 2: + if (lhs[1] != 0u) { + return false; + } + /* FALLTHRU */ + case 1: + if (lhs[0] != 0u) { + return false; + } + } + return true; +} +#endif // AILEGO_M64 +#endif // __AVX2__ + +#if (defined(__ARM_NEON) && defined(__aarch64__)) +static inline size_t bitset_cardinality(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + while (lhs != last_aligned) { + const uint32_t *last_stage = + (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u; + + uint8x16_t v_count = vdupq_n_u8(0); + for (; lhs != last_stage; lhs += 4) { + v_count = vaddq_u8(vcntq_u8(vld1q_u8((const uint8_t *)lhs)), v_count); + } + + v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count)); + count += vaddvq_u16(vreinterpretq_u16_u8(v_count)); + } + + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0]); + } + return count; +} + +static inline size_t bitset_xor_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + while (lhs != last_aligned) { + const uint32_t *last_stage = + (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u; + + uint8x16_t v_count = vdupq_n_u8(0); + for (; lhs != last_stage; lhs += 4, rhs += 4) { + v_count = vaddq_u8(vcntq_u8(veorq_u8(vld1q_u8((const uint8_t *)lhs), + vld1q_u8((const uint8_t *)rhs))), + v_count); + } + + v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count)); + count += vaddvq_u16(vreinterpretq_u16_u8(v_count)); + } + + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2] ^ rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] ^ rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] ^ rhs[0]); + } + return count; +} + +static inline size_t bitset_and_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + while (lhs != last_aligned) { + const uint32_t *last_stage = + (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u; + + uint8x16_t v_count = vdupq_n_u8(0); + for (; lhs != last_stage; lhs += 4, rhs += 4) { + v_count = vaddq_u8(vcntq_u8(vandq_u8(vld1q_u8((const uint8_t *)lhs), + vld1q_u8((const uint8_t *)rhs))), + v_count); + } + + v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count)); + count += vaddvq_u16(vreinterpretq_u16_u8(v_count)); + } + + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2] & rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] & rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] & rhs[0]); + } + return count; +} + +static inline size_t bitset_andnot_cardinality(const uint32_t *lhs, + const uint32_t *rhs, + size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + while (lhs != last_aligned) { + const uint32_t *last_stage = + (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u; + + uint8x16_t v_count = vdupq_n_u8(0); + for (; lhs != last_stage; lhs += 4, rhs += 4) { + v_count = vaddq_u8(vcntq_u8(vbicq_u8(vld1q_u8((const uint8_t *)lhs), + vld1q_u8((const uint8_t *)rhs))), + v_count); + } + + v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count)); + count += vaddvq_u16(vreinterpretq_u16_u8(v_count)); + } + + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2] & ~rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] & ~rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] & ~rhs[0]); + } + return count; +} + +static inline size_t bitset_or_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + while (lhs != last_aligned) { + const uint32_t *last_stage = + (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u; + + uint8x16_t v_count = vdupq_n_u8(0); + for (; lhs != last_stage; lhs += 4, rhs += 4) { + v_count = vaddq_u8(vcntq_u8(vorrq_u8(vld1q_u8((const uint8_t *)lhs), + vld1q_u8((const uint8_t *)rhs))), + v_count); + } + + v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count)); + count += vaddvq_u16(vreinterpretq_u16_u8(v_count)); + } + + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2] | rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] | rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] | rhs[0]); + } + return count; +} + +#elif defined(AILEGO_M64) +static inline size_t bitset_cardinality(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 8) { + count += bitset_popcount64(*(uint64_t *)(&lhs[6])); + count += bitset_popcount64(*(uint64_t *)(&lhs[4])); + count += bitset_popcount64(*(uint64_t *)(&lhs[2])); + count += bitset_popcount64(*(uint64_t *)(&lhs[0])); + } + switch (last - last_aligned) { + case 7: + count += bitset_popcount32(lhs[6]); + /* FALLTHRU */ + case 6: + count += bitset_popcount32(lhs[5]); + /* FALLTHRU */ + case 5: + count += bitset_popcount32(lhs[4]); + /* FALLTHRU */ + case 4: + count += bitset_popcount32(lhs[3]); + /* FALLTHRU */ + case 3: + count += bitset_popcount32(lhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0]); + } + return count; +} + +static inline size_t bitset_xor_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + count += bitset_popcount64(*(uint64_t *)(&lhs[6]) ^ *(uint64_t *)(&rhs[6])); + count += bitset_popcount64(*(uint64_t *)(&lhs[4]) ^ *(uint64_t *)(&rhs[4])); + count += bitset_popcount64(*(uint64_t *)(&lhs[2]) ^ *(uint64_t *)(&rhs[2])); + count += bitset_popcount64(*(uint64_t *)(&lhs[0]) ^ *(uint64_t *)(&rhs[0])); + } + switch (last - last_aligned) { + case 7: + count += bitset_popcount32(lhs[6] ^ rhs[6]); + /* FALLTHRU */ + case 6: + count += bitset_popcount32(lhs[5] ^ rhs[5]); + /* FALLTHRU */ + case 5: + count += bitset_popcount32(lhs[4] ^ rhs[4]); + /* FALLTHRU */ + case 4: + count += bitset_popcount32(lhs[3] ^ rhs[3]); + /* FALLTHRU */ + case 3: + count += bitset_popcount32(lhs[2] ^ rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] ^ rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] ^ rhs[0]); + } + return count; +} + +static inline size_t bitset_and_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + count += bitset_popcount64(*(uint64_t *)(&lhs[6]) & *(uint64_t *)(&rhs[6])); + count += bitset_popcount64(*(uint64_t *)(&lhs[4]) & *(uint64_t *)(&rhs[4])); + count += bitset_popcount64(*(uint64_t *)(&lhs[2]) & *(uint64_t *)(&rhs[2])); + count += bitset_popcount64(*(uint64_t *)(&lhs[0]) & *(uint64_t *)(&rhs[0])); + } + switch (last - last_aligned) { + case 7: + count += bitset_popcount32(lhs[6] & rhs[6]); + /* FALLTHRU */ + case 6: + count += bitset_popcount32(lhs[5] & rhs[5]); + /* FALLTHRU */ + case 5: + count += bitset_popcount32(lhs[4] & rhs[4]); + /* FALLTHRU */ + case 4: + count += bitset_popcount32(lhs[3] & rhs[3]); + /* FALLTHRU */ + case 3: + count += bitset_popcount32(lhs[2] & rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] & rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] & rhs[0]); + } + return count; +} + +static inline size_t bitset_andnot_cardinality(const uint32_t *lhs, + const uint32_t *rhs, + size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + count += + bitset_popcount64(*(uint64_t *)(&lhs[6]) & ~(*(uint64_t *)(&rhs[6]))); + count += + bitset_popcount64(*(uint64_t *)(&lhs[4]) & ~(*(uint64_t *)(&rhs[4]))); + count += + bitset_popcount64(*(uint64_t *)(&lhs[2]) & ~(*(uint64_t *)(&rhs[2]))); + count += + bitset_popcount64(*(uint64_t *)(&lhs[0]) & ~(*(uint64_t *)(&rhs[0]))); + } + switch (last - last_aligned) { + case 7: + count += bitset_popcount32(lhs[6] & ~rhs[6]); + /* FALLTHRU */ + case 6: + count += bitset_popcount32(lhs[5] & ~rhs[5]); + /* FALLTHRU */ + case 5: + count += bitset_popcount32(lhs[4] & ~rhs[4]); + /* FALLTHRU */ + case 4: + count += bitset_popcount32(lhs[3] & ~rhs[3]); + /* FALLTHRU */ + case 3: + count += bitset_popcount32(lhs[2] & ~rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] & ~rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] & ~rhs[0]); + } + return count; +} + +static inline size_t bitset_or_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 3) << 3); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 8, rhs += 8) { + count += bitset_popcount64(*(uint64_t *)(&lhs[6]) | *(uint64_t *)(&rhs[6])); + count += bitset_popcount64(*(uint64_t *)(&lhs[4]) | *(uint64_t *)(&rhs[4])); + count += bitset_popcount64(*(uint64_t *)(&lhs[2]) | *(uint64_t *)(&rhs[2])); + count += bitset_popcount64(*(uint64_t *)(&lhs[0]) | *(uint64_t *)(&rhs[0])); + } + switch (last - last_aligned) { + case 7: + count += bitset_popcount32(lhs[6] | rhs[6]); + /* FALLTHRU */ + case 6: + count += bitset_popcount32(lhs[5] | rhs[5]); + /* FALLTHRU */ + case 5: + count += bitset_popcount32(lhs[4] | rhs[4]); + /* FALLTHRU */ + case 4: + count += bitset_popcount32(lhs[3] | rhs[3]); + /* FALLTHRU */ + case 3: + count += bitset_popcount32(lhs[2] | rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] | rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] | rhs[0]); + } + return count; +} + +#else // !__ARM_NEON && !AILEGO_M64 +static inline size_t bitset_cardinality(const uint32_t *lhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 4) { + count += bitset_popcount32(lhs[3]); + count += bitset_popcount32(lhs[2]); + count += bitset_popcount32(lhs[1]); + count += bitset_popcount32(lhs[0]); + } + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0]); + } + return count; +} + +static inline size_t bitset_xor_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + count += bitset_popcount32(lhs[3] ^ rhs[3]); + count += bitset_popcount32(lhs[2] ^ rhs[2]); + count += bitset_popcount32(lhs[1] ^ rhs[1]); + count += bitset_popcount32(lhs[0] ^ rhs[0]); + } + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2] ^ rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] ^ rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] ^ rhs[0]); + } + return count; +} + +static inline size_t bitset_and_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + count += bitset_popcount32(lhs[3] & rhs[3]); + count += bitset_popcount32(lhs[2] & rhs[2]); + count += bitset_popcount32(lhs[1] & rhs[1]); + count += bitset_popcount32(lhs[0] & rhs[0]); + } + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2] & rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] & rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] & rhs[0]); + } + return count; +} + +static inline size_t bitset_andnot_cardinality(const uint32_t *lhs, + const uint32_t *rhs, + size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + count += bitset_popcount32(lhs[3] & ~rhs[3]); + count += bitset_popcount32(lhs[2] & ~rhs[2]); + count += bitset_popcount32(lhs[1] & ~rhs[1]); + count += bitset_popcount32(lhs[0] & ~rhs[0]); + } + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2] & ~rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] & ~rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] & ~rhs[0]); + } + return count; +} + +static inline size_t bitset_or_cardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + const uint32_t *last = lhs + size; + const uint32_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + count += bitset_popcount32(lhs[3] | rhs[3]); + count += bitset_popcount32(lhs[2] | rhs[2]); + count += bitset_popcount32(lhs[1] | rhs[1]); + count += bitset_popcount32(lhs[0] | rhs[0]); + } + switch (last - last_aligned) { + case 3: + count += bitset_popcount32(lhs[2] | rhs[2]); + /* FALLTHRU */ + case 2: + count += bitset_popcount32(lhs[1] | rhs[1]); + /* FALLTHRU */ + case 1: + count += bitset_popcount32(lhs[0] | rhs[0]); + } + return count; +} +#endif // __ARM_NEON && __aarch64__ + +namespace zvec { + +namespace ailego { + +void BitsetHelper::BitwiseAnd(uint32_t *lhs, const uint32_t *rhs, size_t size) { + bitset_and(lhs, rhs, size); +} + +void BitsetHelper::BitwiseAndnot(uint32_t *lhs, const uint32_t *rhs, + size_t size) { + bitset_andnot(lhs, rhs, size); +} + +void BitsetHelper::BitwiseOr(uint32_t *lhs, const uint32_t *rhs, size_t size) { + bitset_or(lhs, rhs, size); +} + +void BitsetHelper::BitwiseXor(uint32_t *lhs, const uint32_t *rhs, size_t size) { + bitset_xor(lhs, rhs, size); +} + +void BitsetHelper::BitwiseNot(uint32_t *arr, size_t size) { + bitset_not(arr, size); +} + +bool BitsetHelper::TestAll(const uint32_t *arr, size_t size) { + return bitset_test_all(arr, size); +} + +bool BitsetHelper::TestAny(const uint32_t *arr, size_t size) { + return bitset_test_any(arr, size); +} + +bool BitsetHelper::TestNone(const uint32_t *arr, size_t size) { + return bitset_test_none(arr, size); +} + +size_t BitsetHelper::BitwiseAndCardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + return bitset_and_cardinality(lhs, rhs, size); +} + +size_t BitsetHelper::BitwiseOrCardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + return bitset_or_cardinality(lhs, rhs, size); +} + +size_t BitsetHelper::BitwiseAndnotCardinality(const uint32_t *lhs, + const uint32_t *rhs, + size_t size) { + return bitset_andnot_cardinality(lhs, rhs, size); +} + +size_t BitsetHelper::BitwiseXorCardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + return bitset_xor_cardinality(lhs, rhs, size); +} + +size_t BitsetHelper::Cardinality(const uint32_t *arr, size_t size) { + return bitset_cardinality(arr, size); +} + +bool BitsetHelper::test_all(void) const { + return bitset_test_all(array_, size_); +} + +bool BitsetHelper::test_any(void) const { + return bitset_test_any(array_, size_); +} + +bool BitsetHelper::test_none(void) const { + return bitset_test_none(array_, size_); +} + +size_t BitsetHelper::cardinality(void) const { + return bitset_cardinality(array_, size_); +} + +} // namespace ailego + +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/utility/bitset_helper.h b/src/ailego/utility/bitset_helper.h new file mode 100644 index 0000000..56ac7e9 --- /dev/null +++ b/src/ailego/utility/bitset_helper.h @@ -0,0 +1,174 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec { + +namespace ailego { + +/*! Bitset Helper + */ +class BitsetHelper { + public: + //! Constructor + BitsetHelper(void) {} + + //! Constructor + BitsetHelper(void *buf, size_t len) + : array_(reinterpret_cast(buf)), + size_(len / sizeof(uint32_t)) {} + + //! Mount a buffer as bitset + void mount(void *buf, size_t len) { + array_ = reinterpret_cast(buf); + size_ = len / sizeof(uint32_t); + } + + //! Umount the buffer + void umount(void) { + array_ = nullptr; + size_ = 0u; + } + + // ！Clear the bitset + void clear(void) { + memset(array_, 0, sizeof(uint32_t) * size_); + } + + //! Test a bit in bitset + bool test(size_t num) const { + ailego_assert_with((size_ << 5) > num, "overflow argument"); + return ((array_[num >> 5] & (1u << (num & 0x1f))) != 0); + } + + //! Set a bit in bitset + void set(size_t num) { + ailego_assert_with((size_ << 5) > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] |= mask; + } + + //! Reset a bit in bitset + void reset(size_t num) { + ailego_assert_with((size_ << 5) > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] &= ~mask; + } + + //! Toggle a bit in bitset + void flip(size_t num) { + ailego_assert_with((size_ << 5) > num, "overflow argument"); + uint32_t mask = (1u << (num & 0x1f)); + array_[num >> 5] ^= mask; + } + + //! Extract the bitset to an array + void extract(size_t base, std::vector *out) const { + const uint32_t *iter = array_; + const uint32_t *last = array_ + size_; + + for (; iter != last; ++iter) { + uint32_t w = *iter; + + while (w != 0) { + uint32_t c = ailego_ctz32(w); + w &= ~(1u << c); + out->push_back(base + c); + } + base += 32u; + } + } + + //! Extract the bitset to an array + void extract(std::vector *out) const { + this->extract(0, out); + } + + //! Check if all bits are set to true + bool test_all(void) const; + + //! Check if any bits are set to true + bool test_any(void) const; + + //! Check if none of the bits are set to true + bool test_none(void) const; + + //! Compute the cardinality of a bitset + size_t cardinality(void) const; + + //! Calculate the size of buffer if it contains N bits + static size_t BufferSize(size_t N) { + return (((N + 0x1f) >> 5) << 2); + } + + //! Calculate the count of bits can be contained + static size_t BitsCount(size_t len) { + return ((len >> 2) << 2); + } + + //! Check if all bits are set to true + static bool TestAll(const uint32_t *arr, size_t size); + + //! Check if cube bits are set to true + static bool TestAny(const uint32_t *arr, size_t size); + + //! Check if none of the bits are set to true + static bool TestNone(const uint32_t *arr, size_t size); + + //! Compute the AND cardinality between two bitsets + static size_t BitwiseAndCardinality(const uint32_t *lhs, const uint32_t *rhs, + size_t size); + + //! Compute the OR cardinality between two bitsets + static size_t BitwiseOrCardinality(const uint32_t *lhs, const uint32_t *rhs, + size_t size); + + //! Compute the ANDNOT cardinality between two bitsets + static size_t BitwiseAndnotCardinality(const uint32_t *lhs, + const uint32_t *rhs, size_t size); + + //! Compute the XOR cardinality between two bitsets + static size_t BitwiseXorCardinality(const uint32_t *lhs, const uint32_t *rhs, + size_t size); + + //! Compute the cardinality of a bitset + static size_t Cardinality(const uint32_t *arr, size_t size); + + //! Perform binary AND + static void BitwiseAnd(uint32_t *lhs, const uint32_t *rhs, size_t size); + + //! Perform binary AND_NOT + static void BitwiseAndnot(uint32_t *lhs, const uint32_t *rhs, size_t size); + + //! Perform binary OR + static void BitwiseOr(uint32_t *lhs, const uint32_t *rhs, size_t size); + + //! Perform binary XOR + static void BitwiseXor(uint32_t *lhs, const uint32_t *rhs, size_t size); + + //! Perform binary NOT + static void BitwiseNot(uint32_t *arr, size_t size); + + private: + uint32_t *array_{nullptr}; + size_t size_{0u}; +}; + +} // namespace ailego + +} // namespace zvec diff --git a/src/ailego/utility/concurrency_helper.cc b/src/ailego/utility/concurrency_helper.cc new file mode 100644 index 0000000..3f79c9c --- /dev/null +++ b/src/ailego/utility/concurrency_helper.cc @@ -0,0 +1,70 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "concurrency_helper.h" +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +// Refer to: +// https://stackoverflow.com/questions/65551215/get-docker-cpu-memory-limit-inside-container +ConcurrencyHelper::ConcurrencyHelper() { + std::string cfs_quota_us = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"; + std::string cfs_period_us = "/sys/fs/cgroup/cpu/cpu.cfs_period_us"; + + concurrency_ = std::thread::hardware_concurrency(); + if (FileHelper::IsExist(cfs_quota_us.c_str()) && + FileHelper::IsExist(cfs_period_us.c_str())) { + std::ifstream quota_ifs; + std::string quota_str{""}; + uint32_t quota_val = 0; + quota_ifs.open(cfs_quota_us, std::ios::in); + if (quota_ifs.is_open()) { + quota_ifs >> quota_str; + if (quota_str != "-1") { + StringHelper::ToUint32(quota_str, "a_val); + } + quota_ifs.close(); + } + + if (quota_val > 0) { + std::ifstream period_ifs; + std::string period_str{""}; + uint32_t period_val = 0; + period_ifs.open(cfs_period_us, std::ios::in); + if (period_ifs.is_open()) { + period_ifs >> period_str; + StringHelper::ToUint32(period_str, &period_val); + period_ifs.close(); + } + + if (period_val > 0) { + concurrency_ = (quota_val + period_val - 1) / period_val; + } + } + } +} + +uint32_t ConcurrencyHelper::container_aware_concurrency() { + static ConcurrencyHelper concurrency_helper; + return concurrency_helper.concurrency_; +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/concurrency_helper.h b/src/ailego/utility/concurrency_helper.h new file mode 100644 index 0000000..993d753 --- /dev/null +++ b/src/ailego/utility/concurrency_helper.h @@ -0,0 +1,34 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +class ConcurrencyHelper { + public: + ConcurrencyHelper(); + + //! get hardware concurrency from either vm or container + static uint32_t container_aware_concurrency(); + + private: + uint32_t concurrency_{0}; +}; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/utility/dl_helper.cc b/src/ailego/utility/dl_helper.cc new file mode 100644 index 0000000..545c875 --- /dev/null +++ b/src/ailego/utility/dl_helper.cc @@ -0,0 +1,74 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "dl_helper.h" +#if !defined(_WIN64) && !defined(_WIN32) +#include +#else +#include +#endif + +namespace zvec { +namespace ailego { + +#if !defined(_WIN64) && !defined(_WIN32) +void *DLHelper::Load(const char *path, std::string *err) { + void *handle = dlopen(path, RTLD_NOW); + if (!handle && err) { + *err = dlerror(); + } + return handle; +} + +void DLHelper::Unload(void *handle) { + ailego_return_if_false(handle); + dlclose(handle); +} + +void *DLHelper::Symbol(void *handle, const char *symbol) { + ailego_null_if_false(handle && symbol); + return dlsym(handle, symbol); +} + +#else +void *DLHelper::Load(const char *path, std::string *err) { + HMODULE handle = LoadLibraryA(path); + if (!handle && err) { + DWORD error_code = GetLastError(); + LPSTR error_msg = nullptr; + + DWORD len = FormatMessageA( + FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | + FORMAT_MESSAGE_IGNORE_INSERTS, + nullptr, error_code, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + (LPSTR)&error_msg, 0, nullptr); + err->assign(error_msg, len); + LocalFree(error_msg); + } + return handle; +} + +void DLHelper::Unload(void *handle) { + ailego_return_if_false(handle); + FreeLibrary((HMODULE)handle); +} + +void *DLHelper::Symbol(void *handle, const char *symbol) { + ailego_null_if_false(handle && symbol); + return GetProcAddress((HMODULE)handle, symbol); +} +#endif // !_WIN64 && !_WIN32 + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/utility/dl_helper.h b/src/ailego/utility/dl_helper.h new file mode 100644 index 0000000..d72b0dc --- /dev/null +++ b/src/ailego/utility/dl_helper.h @@ -0,0 +1,47 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec { +namespace ailego { + +/*! Dynamic Library Helper + */ +struct DLHelper { + //! Load library from path + static void *Load(const char *path, std::string *err); + + //! Unload a library + static void Unload(void *handle); + + //! Retrieve a symbol from a library handle + static void *Symbol(void *handle, const char *symbol); + + //! Load library from path + static void *Load(const std::string &path, std::string *err) { + return DLHelper::Load(path.c_str(), err); + } + + //! Retrieve a symbol from a library handle + static void *Symbol(void *handle, const std::string &symbol) { + return DLHelper::Symbol(handle, symbol.c_str()); + } +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/file_helper.cc b/src/ailego/utility/file_helper.cc new file mode 100644 index 0000000..5026f40 --- /dev/null +++ b/src/ailego/utility/file_helper.cc @@ -0,0 +1,404 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) || defined(_WIN64) +#include +#ifdef RemoveDirectory +#undef RemoveDirectory +#endif +#ifdef DeleteFile +#undef DeleteFile +#endif +#ifdef GetFileSize +#undef GetFileSize +#endif +#else +#if defined(__APPLE__) || defined(__MACH__) +#include +#endif +#if defined(__FreeBSD__) +#include +#endif +#include +#include +#endif + +namespace fs = std::filesystem; + +namespace zvec { +namespace ailego { + +namespace { + +thread_local std::error_code g_last_fs_error; + +void ClearFsError() { + g_last_fs_error.clear(); +} + +void SetFsError(std::error_code ec) { + g_last_fs_error = ec; +} + +} // namespace + +// ---------- public UTF-8 / wide helpers ---------- + +fs::path FileHelper::PathFromUtf8(const char *s) { +#if defined(_WIN32) || defined(_WIN64) + if (!s || !*s) { + return fs::path(); + } + return fs::u8path(s); +#else + return fs::path(s ? s : ""); +#endif +} + +fs::path FileHelper::PathFromUtf8(const std::string &s) { + return PathFromUtf8(s.c_str()); +} + +std::string FileHelper::PathToUtf8(const fs::path &p) { + return p.u8string(); +} + +#if defined(_WIN32) || defined(_WIN64) +std::wstring FileHelper::Utf8ToWide(const std::string &src) { + if (src.empty()) { + return {}; + } + int src_len = static_cast(src.size()); + int dst_len = + MultiByteToWideChar(CP_UTF8, 0, src.data(), src_len, nullptr, 0); + if (dst_len <= 0) { + return {}; + } + std::wstring dst(static_cast(dst_len), L'\0'); + if (MultiByteToWideChar(CP_UTF8, 0, src.data(), src_len, dst.data(), + dst_len) != dst_len) { + return {}; + } + return dst; +} + +std::string FileHelper::WideToUtf8(const std::wstring &src) { + if (src.empty()) { + return {}; + } + int src_len = static_cast(src.size()); + int dst_len = WideCharToMultiByte(CP_UTF8, 0, src.data(), src_len, nullptr, 0, + nullptr, nullptr); + if (dst_len <= 0) { + return {}; + } + std::string dst(static_cast(dst_len), '\0'); + if (WideCharToMultiByte(CP_UTF8, 0, src.data(), src_len, dst.data(), dst_len, + nullptr, nullptr) != dst_len) { + return {}; + } + return dst; +} + +#endif + +// ---------- internal helpers ---------- + +namespace { + +static bool GetFileSizeImpl(const fs::path &p, size_t *psz) { + ClearFsError(); + std::error_code ec; + auto sz = fs::file_size(p, ec); + if (ec) { + SetFsError(ec); + return false; + } + *psz = static_cast(sz); + return true; +} + +static bool DeleteFileImpl(const fs::path &p) { + ClearFsError(); + std::error_code ec; + fs::file_status st = fs::symlink_status(p, ec); + if (ec) { + SetFsError(ec); + return false; + } + if (fs::is_directory(st) && !fs::is_symlink(st)) { + ec = std::make_error_code(std::errc::is_a_directory); + SetFsError(ec); + return false; + } + if (!fs::remove(p, ec)) { + SetFsError(ec ? ec + : std::make_error_code(std::errc::no_such_file_or_directory)); + return false; + } + return true; +} + +static bool RenameFileImpl(const fs::path &from, const fs::path &to) { + ClearFsError(); + std::error_code ec; + fs::rename(from, to, ec); + if (ec) { + SetFsError(ec); + return false; + } + return true; +} + +static bool MakePathImpl(const fs::path &p) { + ClearFsError(); + std::error_code ec; + fs::create_directories(p, ec); + if (ec) { + SetFsError(ec); + return false; + } + return true; +} + +static bool RemoveDirectoryImpl(const fs::path &p) { + ClearFsError(); + std::error_code ec; + if (!fs::is_directory(p, ec)) { + if (ec) { + SetFsError(ec); + } + return false; + } + std::uintmax_t n = fs::remove_all(p, ec); + if (ec) { + SetFsError(ec); + return false; + } + (void)n; + return true; +} + +static bool IsExistImpl(const fs::path &p) { + std::error_code ec; + return fs::exists(p, ec); +} + +static bool IsRegularImpl(const fs::path &p) { + std::error_code ec; + return fs::is_regular_file(p, ec); +} + +static bool IsDirectoryImpl(const fs::path &p) { + std::error_code ec; + return fs::is_directory(p, ec); +} + +static bool IsSymbolicLinkImpl(const fs::path &p) { + std::error_code ec; + return fs::is_symlink(p, ec); +} + +static bool IsSameImpl(const fs::path &a, const fs::path &b) { + std::error_code ec; + return fs::equivalent(a, b, ec); +} + +} // namespace + +bool FileHelper::GetSelfPath(std::string *path) { +#if defined(_WIN32) || defined(_WIN64) + std::wstring wbuf(4096, L'\0'); + DWORD n = + GetModuleFileNameW(nullptr, wbuf.data(), static_cast(wbuf.size())); + while (n >= wbuf.size() - 1) { + if (wbuf.size() > 65536) { + return false; + } + wbuf.resize(wbuf.size() * 2); + n = GetModuleFileNameW(nullptr, wbuf.data(), + static_cast(wbuf.size())); + } + if (n == 0) { + return false; + } + wbuf.resize(n); + *path = WideToUtf8(wbuf); + return !path->empty(); +#elif defined(__APPLE__) || defined(__MACH__) + char buf[PATH_MAX]; + size_t len = 0; + + char dirty_buf[PATH_MAX]; + uint32_t size = sizeof(dirty_buf); + if (_NSGetExecutablePath(dirty_buf, &size) == 0) { + realpath(dirty_buf, buf); + len = strlen(buf); + } +#elif defined(__FreeBSD__) + char buf[PATH_MAX]; + size_t len = PATH_MAX; + int mib[4] = {CTL_KERN, KERN_PROC, KERN_PROC_PATHNAME, -1}; + if (sysctl(mib, 4, &buf, &len, NULL, 0) != 0) { + len = 0; + } +#else + char buf[PATH_MAX]; + ssize_t len = readlink("/proc/self/exe", buf, PATH_MAX); +#endif + +#if !defined(_WIN32) && !defined(_WIN64) + if (len <= 0) { + return false; + } + path->assign(buf, len); + return true; +#endif +} + +bool FileHelper::GetFilePath(NativeHandle handle, std::string *path) { +#if defined(_WIN32) || defined(_WIN64) + DWORD need = GetFinalPathNameByHandleW(static_cast(handle), nullptr, + 0, FILE_NAME_OPENED); + if (need == 0) { + return false; + } + std::wstring wbuf(static_cast(need) + 1, L'\0'); + DWORD got = GetFinalPathNameByHandleW( + static_cast(handle), wbuf.data(), need + 1, FILE_NAME_OPENED); + if (got == 0 || got > need) { + return false; + } + wbuf.resize(got); + *path = WideToUtf8(wbuf); + return !path->empty(); +#elif defined(__linux) || defined(__linux__) + char buf[PATH_MAX]; + char src[32]; + snprintf(src, sizeof(src), "/proc/self/fd/%d", handle); + ssize_t len = readlink(src, buf, PATH_MAX); +#else + char buf[PATH_MAX]; + size_t len = 0; + if (fcntl(handle, F_GETPATH, buf) != -1) { + len = strlen(buf); + } +#endif + +#if !defined(_WIN32) && !defined(_WIN64) + if (len <= 0) { + return false; + } + path->assign(buf, len); + return true; +#endif +} + +bool FileHelper::GetWorkingDirectory(std::string *path) { + ClearFsError(); + std::error_code ec; + fs::path cwd = fs::current_path(ec); + if (ec) { + SetFsError(ec); + return false; + } + *path = PathToUtf8(cwd); + return !path->empty(); +} + +bool FileHelper::GetFileSize(const char *path, size_t *psz) { + return GetFileSizeImpl(PathFromUtf8(path), psz); +} + +bool FileHelper::DeleteFile(const char *path) { + return DeleteFileImpl(PathFromUtf8(path)); +} + +bool FileHelper::RenameFile(const char *oldpath, const char *newpath) { + return RenameFileImpl(PathFromUtf8(oldpath), PathFromUtf8(newpath)); +} + +bool FileHelper::MakePath(const char *path) { + return MakePathImpl(PathFromUtf8(path)); +} + +bool FileHelper::RemoveDirectory(const char *path) { + if (path == nullptr || *path == '\0') { + return false; + } + return RemoveDirectoryImpl(PathFromUtf8(path)); +} + +bool FileHelper::IsExist(const char *path) { + return IsExistImpl(PathFromUtf8(path)); +} + +bool FileHelper::IsRegular(const char *path) { + return IsRegularImpl(PathFromUtf8(path)); +} + +bool FileHelper::IsDirectory(const char *path) { + return IsDirectoryImpl(PathFromUtf8(path)); +} + +bool FileHelper::IsSymbolicLink(const char *path) { + return IsSymbolicLinkImpl(PathFromUtf8(path)); +} + +bool FileHelper::IsSame(const char *path1, const char *path2) { + return IsSameImpl(PathFromUtf8(path1), PathFromUtf8(path2)); +} + +std::string FileHelper::GetLastErrorString() { + if (g_last_fs_error) { + return g_last_fs_error.message(); + } +#if defined(_WIN32) || defined(_WIN64) + DWORD err = GetLastError(); + if (err == 0) { + return "No error"; + } + char buf[256]; + DWORD len = FormatMessageA( + FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, nullptr, err, + MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), buf, sizeof(buf), nullptr); + if (len > 0) { + while (len > 0 && (buf[len - 1] == '\r' || buf[len - 1] == '\n')) { + buf[--len] = '\0'; + } + return std::string(buf, len); + } + return "Unknown error " + std::to_string(err); +#else + return strerror(errno); +#endif +} + +bool FileHelper::RemovePath(const char *path) { + if (FileHelper::IsDirectory(path)) { + return FileHelper::RemoveDirectory(path); + } + return FileHelper::DeleteFile(path); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/float_helper.cc b/src/ailego/utility/float_helper.cc new file mode 100644 index 0000000..bc07eec --- /dev/null +++ b/src/ailego/utility/float_helper.cc @@ -0,0 +1,1246 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +// #if defined(__F16C__) && defined(__AVX__) +// #define float16(x) _cvtss_sh((x), _MM_FROUND_NO_EXC) +// #define float32(x) _cvtsh_ss(x) +// #endif // __F16C__ && __AVX__ + +#if defined(__aarch64__) +static inline float float32(uint16_t val) { + __fp16 f; + memcpy(&f, &val, sizeof(val)); + return static_cast(f); +} + +static inline uint16_t float16(float val) { + __fp16 f = static_cast<__fp16>(val); + uint16_t result; + memcpy(&result, &f, sizeof(result)); + return result; +} + +static inline void convert_fp16_to_fp32(const uint16_t *arr, size_t size, + float *out) { + for (size_t i = 0; i != size; ++i) { + out[i] = float32(arr[i]); + } +} + +static inline void convert_fp16_to_fp32(const uint16_t *arr, size_t size, + float norm, float *out) { + for (size_t i = 0; i != size; ++i) { + out[i] = float32(arr[i]) / norm; + } +} + +static inline void convert_fp32_to_fp16(const float *arr, size_t size, + uint16_t *out) { + for (size_t i = 0; i != size; ++i) { + out[i] = float16(arr[i]); + } +} + +static inline void convert_fp32_to_fp16(const float *arr, size_t size, + float norm, uint16_t *out) { + for (size_t i = 0; i != size; ++i) { + out[i] = float16(arr[i] / norm); + } +} +#else +// Refer: https://github.com/Maratyszcza/FP16/blob/master/third-party/half.hpp +static inline float float32(uint16_t val) { + static const uint32_t mantissa_table[2048] = { + 0x00000000, 0x33800000, 0x34000000, 0x34400000, 0x34800000, 0x34A00000, + 0x34C00000, 0x34E00000, 0x35000000, 0x35100000, 0x35200000, 0x35300000, + 0x35400000, 0x35500000, 0x35600000, 0x35700000, 0x35800000, 0x35880000, + 0x35900000, 0x35980000, 0x35A00000, 0x35A80000, 0x35B00000, 0x35B80000, + 0x35C00000, 0x35C80000, 0x35D00000, 0x35D80000, 0x35E00000, 0x35E80000, + 0x35F00000, 0x35F80000, 0x36000000, 0x36040000, 0x36080000, 0x360C0000, + 0x36100000, 0x36140000, 0x36180000, 0x361C0000, 0x36200000, 0x36240000, + 0x36280000, 0x362C0000, 0x36300000, 0x36340000, 0x36380000, 0x363C0000, + 0x36400000, 0x36440000, 0x36480000, 0x364C0000, 0x36500000, 0x36540000, + 0x36580000, 0x365C0000, 0x36600000, 0x36640000, 0x36680000, 0x366C0000, + 0x36700000, 0x36740000, 0x36780000, 0x367C0000, 0x36800000, 0x36820000, + 0x36840000, 0x36860000, 0x36880000, 0x368A0000, 0x368C0000, 0x368E0000, + 0x36900000, 0x36920000, 0x36940000, 0x36960000, 0x36980000, 0x369A0000, + 0x369C0000, 0x369E0000, 0x36A00000, 0x36A20000, 0x36A40000, 0x36A60000, + 0x36A80000, 0x36AA0000, 0x36AC0000, 0x36AE0000, 0x36B00000, 0x36B20000, + 0x36B40000, 0x36B60000, 0x36B80000, 0x36BA0000, 0x36BC0000, 0x36BE0000, + 0x36C00000, 0x36C20000, 0x36C40000, 0x36C60000, 0x36C80000, 0x36CA0000, + 0x36CC0000, 0x36CE0000, 0x36D00000, 0x36D20000, 0x36D40000, 0x36D60000, + 0x36D80000, 0x36DA0000, 0x36DC0000, 0x36DE0000, 0x36E00000, 0x36E20000, + 0x36E40000, 0x36E60000, 0x36E80000, 0x36EA0000, 0x36EC0000, 0x36EE0000, + 0x36F00000, 0x36F20000, 0x36F40000, 0x36F60000, 0x36F80000, 0x36FA0000, + 0x36FC0000, 0x36FE0000, 0x37000000, 0x37010000, 0x37020000, 0x37030000, + 0x37040000, 0x37050000, 0x37060000, 0x37070000, 0x37080000, 0x37090000, + 0x370A0000, 0x370B0000, 0x370C0000, 0x370D0000, 0x370E0000, 0x370F0000, + 0x37100000, 0x37110000, 0x37120000, 0x37130000, 0x37140000, 0x37150000, + 0x37160000, 0x37170000, 0x37180000, 0x37190000, 0x371A0000, 0x371B0000, + 0x371C0000, 0x371D0000, 0x371E0000, 0x371F0000, 0x37200000, 0x37210000, + 0x37220000, 0x37230000, 0x37240000, 0x37250000, 0x37260000, 0x37270000, + 0x37280000, 0x37290000, 0x372A0000, 0x372B0000, 0x372C0000, 0x372D0000, + 0x372E0000, 0x372F0000, 0x37300000, 0x37310000, 0x37320000, 0x37330000, + 0x37340000, 0x37350000, 0x37360000, 0x37370000, 0x37380000, 0x37390000, + 0x373A0000, 0x373B0000, 0x373C0000, 0x373D0000, 0x373E0000, 0x373F0000, + 0x37400000, 0x37410000, 0x37420000, 0x37430000, 0x37440000, 0x37450000, + 0x37460000, 0x37470000, 0x37480000, 0x37490000, 0x374A0000, 0x374B0000, + 0x374C0000, 0x374D0000, 0x374E0000, 0x374F0000, 0x37500000, 0x37510000, + 0x37520000, 0x37530000, 0x37540000, 0x37550000, 0x37560000, 0x37570000, + 0x37580000, 0x37590000, 0x375A0000, 0x375B0000, 0x375C0000, 0x375D0000, + 0x375E0000, 0x375F0000, 0x37600000, 0x37610000, 0x37620000, 0x37630000, + 0x37640000, 0x37650000, 0x37660000, 0x37670000, 0x37680000, 0x37690000, + 0x376A0000, 0x376B0000, 0x376C0000, 0x376D0000, 0x376E0000, 0x376F0000, + 0x37700000, 0x37710000, 0x37720000, 0x37730000, 0x37740000, 0x37750000, + 0x37760000, 0x37770000, 0x37780000, 0x37790000, 0x377A0000, 0x377B0000, + 0x377C0000, 0x377D0000, 0x377E0000, 0x377F0000, 0x37800000, 0x37808000, + 0x37810000, 0x37818000, 0x37820000, 0x37828000, 0x37830000, 0x37838000, + 0x37840000, 0x37848000, 0x37850000, 0x37858000, 0x37860000, 0x37868000, + 0x37870000, 0x37878000, 0x37880000, 0x37888000, 0x37890000, 0x37898000, + 0x378A0000, 0x378A8000, 0x378B0000, 0x378B8000, 0x378C0000, 0x378C8000, + 0x378D0000, 0x378D8000, 0x378E0000, 0x378E8000, 0x378F0000, 0x378F8000, + 0x37900000, 0x37908000, 0x37910000, 0x37918000, 0x37920000, 0x37928000, + 0x37930000, 0x37938000, 0x37940000, 0x37948000, 0x37950000, 0x37958000, + 0x37960000, 0x37968000, 0x37970000, 0x37978000, 0x37980000, 0x37988000, + 0x37990000, 0x37998000, 0x379A0000, 0x379A8000, 0x379B0000, 0x379B8000, + 0x379C0000, 0x379C8000, 0x379D0000, 0x379D8000, 0x379E0000, 0x379E8000, + 0x379F0000, 0x379F8000, 0x37A00000, 0x37A08000, 0x37A10000, 0x37A18000, + 0x37A20000, 0x37A28000, 0x37A30000, 0x37A38000, 0x37A40000, 0x37A48000, + 0x37A50000, 0x37A58000, 0x37A60000, 0x37A68000, 0x37A70000, 0x37A78000, + 0x37A80000, 0x37A88000, 0x37A90000, 0x37A98000, 0x37AA0000, 0x37AA8000, + 0x37AB0000, 0x37AB8000, 0x37AC0000, 0x37AC8000, 0x37AD0000, 0x37AD8000, + 0x37AE0000, 0x37AE8000, 0x37AF0000, 0x37AF8000, 0x37B00000, 0x37B08000, + 0x37B10000, 0x37B18000, 0x37B20000, 0x37B28000, 0x37B30000, 0x37B38000, + 0x37B40000, 0x37B48000, 0x37B50000, 0x37B58000, 0x37B60000, 0x37B68000, + 0x37B70000, 0x37B78000, 0x37B80000, 0x37B88000, 0x37B90000, 0x37B98000, + 0x37BA0000, 0x37BA8000, 0x37BB0000, 0x37BB8000, 0x37BC0000, 0x37BC8000, + 0x37BD0000, 0x37BD8000, 0x37BE0000, 0x37BE8000, 0x37BF0000, 0x37BF8000, + 0x37C00000, 0x37C08000, 0x37C10000, 0x37C18000, 0x37C20000, 0x37C28000, + 0x37C30000, 0x37C38000, 0x37C40000, 0x37C48000, 0x37C50000, 0x37C58000, + 0x37C60000, 0x37C68000, 0x37C70000, 0x37C78000, 0x37C80000, 0x37C88000, + 0x37C90000, 0x37C98000, 0x37CA0000, 0x37CA8000, 0x37CB0000, 0x37CB8000, + 0x37CC0000, 0x37CC8000, 0x37CD0000, 0x37CD8000, 0x37CE0000, 0x37CE8000, + 0x37CF0000, 0x37CF8000, 0x37D00000, 0x37D08000, 0x37D10000, 0x37D18000, + 0x37D20000, 0x37D28000, 0x37D30000, 0x37D38000, 0x37D40000, 0x37D48000, + 0x37D50000, 0x37D58000, 0x37D60000, 0x37D68000, 0x37D70000, 0x37D78000, + 0x37D80000, 0x37D88000, 0x37D90000, 0x37D98000, 0x37DA0000, 0x37DA8000, + 0x37DB0000, 0x37DB8000, 0x37DC0000, 0x37DC8000, 0x37DD0000, 0x37DD8000, + 0x37DE0000, 0x37DE8000, 0x37DF0000, 0x37DF8000, 0x37E00000, 0x37E08000, + 0x37E10000, 0x37E18000, 0x37E20000, 0x37E28000, 0x37E30000, 0x37E38000, + 0x37E40000, 0x37E48000, 0x37E50000, 0x37E58000, 0x37E60000, 0x37E68000, + 0x37E70000, 0x37E78000, 0x37E80000, 0x37E88000, 0x37E90000, 0x37E98000, + 0x37EA0000, 0x37EA8000, 0x37EB0000, 0x37EB8000, 0x37EC0000, 0x37EC8000, + 0x37ED0000, 0x37ED8000, 0x37EE0000, 0x37EE8000, 0x37EF0000, 0x37EF8000, + 0x37F00000, 0x37F08000, 0x37F10000, 0x37F18000, 0x37F20000, 0x37F28000, + 0x37F30000, 0x37F38000, 0x37F40000, 0x37F48000, 0x37F50000, 0x37F58000, + 0x37F60000, 0x37F68000, 0x37F70000, 0x37F78000, 0x37F80000, 0x37F88000, + 0x37F90000, 0x37F98000, 0x37FA0000, 0x37FA8000, 0x37FB0000, 0x37FB8000, + 0x37FC0000, 0x37FC8000, 0x37FD0000, 0x37FD8000, 0x37FE0000, 0x37FE8000, + 0x37FF0000, 0x37FF8000, 0x38000000, 0x38004000, 0x38008000, 0x3800C000, + 0x38010000, 0x38014000, 0x38018000, 0x3801C000, 0x38020000, 0x38024000, + 0x38028000, 0x3802C000, 0x38030000, 0x38034000, 0x38038000, 0x3803C000, + 0x38040000, 0x38044000, 0x38048000, 0x3804C000, 0x38050000, 0x38054000, + 0x38058000, 0x3805C000, 0x38060000, 0x38064000, 0x38068000, 0x3806C000, + 0x38070000, 0x38074000, 0x38078000, 0x3807C000, 0x38080000, 0x38084000, + 0x38088000, 0x3808C000, 0x38090000, 0x38094000, 0x38098000, 0x3809C000, + 0x380A0000, 0x380A4000, 0x380A8000, 0x380AC000, 0x380B0000, 0x380B4000, + 0x380B8000, 0x380BC000, 0x380C0000, 0x380C4000, 0x380C8000, 0x380CC000, + 0x380D0000, 0x380D4000, 0x380D8000, 0x380DC000, 0x380E0000, 0x380E4000, + 0x380E8000, 0x380EC000, 0x380F0000, 0x380F4000, 0x380F8000, 0x380FC000, + 0x38100000, 0x38104000, 0x38108000, 0x3810C000, 0x38110000, 0x38114000, + 0x38118000, 0x3811C000, 0x38120000, 0x38124000, 0x38128000, 0x3812C000, + 0x38130000, 0x38134000, 0x38138000, 0x3813C000, 0x38140000, 0x38144000, + 0x38148000, 0x3814C000, 0x38150000, 0x38154000, 0x38158000, 0x3815C000, + 0x38160000, 0x38164000, 0x38168000, 0x3816C000, 0x38170000, 0x38174000, + 0x38178000, 0x3817C000, 0x38180000, 0x38184000, 0x38188000, 0x3818C000, + 0x38190000, 0x38194000, 0x38198000, 0x3819C000, 0x381A0000, 0x381A4000, + 0x381A8000, 0x381AC000, 0x381B0000, 0x381B4000, 0x381B8000, 0x381BC000, + 0x381C0000, 0x381C4000, 0x381C8000, 0x381CC000, 0x381D0000, 0x381D4000, + 0x381D8000, 0x381DC000, 0x381E0000, 0x381E4000, 0x381E8000, 0x381EC000, + 0x381F0000, 0x381F4000, 0x381F8000, 0x381FC000, 0x38200000, 0x38204000, + 0x38208000, 0x3820C000, 0x38210000, 0x38214000, 0x38218000, 0x3821C000, + 0x38220000, 0x38224000, 0x38228000, 0x3822C000, 0x38230000, 0x38234000, + 0x38238000, 0x3823C000, 0x38240000, 0x38244000, 0x38248000, 0x3824C000, + 0x38250000, 0x38254000, 0x38258000, 0x3825C000, 0x38260000, 0x38264000, + 0x38268000, 0x3826C000, 0x38270000, 0x38274000, 0x38278000, 0x3827C000, + 0x38280000, 0x38284000, 0x38288000, 0x3828C000, 0x38290000, 0x38294000, + 0x38298000, 0x3829C000, 0x382A0000, 0x382A4000, 0x382A8000, 0x382AC000, + 0x382B0000, 0x382B4000, 0x382B8000, 0x382BC000, 0x382C0000, 0x382C4000, + 0x382C8000, 0x382CC000, 0x382D0000, 0x382D4000, 0x382D8000, 0x382DC000, + 0x382E0000, 0x382E4000, 0x382E8000, 0x382EC000, 0x382F0000, 0x382F4000, + 0x382F8000, 0x382FC000, 0x38300000, 0x38304000, 0x38308000, 0x3830C000, + 0x38310000, 0x38314000, 0x38318000, 0x3831C000, 0x38320000, 0x38324000, + 0x38328000, 0x3832C000, 0x38330000, 0x38334000, 0x38338000, 0x3833C000, + 0x38340000, 0x38344000, 0x38348000, 0x3834C000, 0x38350000, 0x38354000, + 0x38358000, 0x3835C000, 0x38360000, 0x38364000, 0x38368000, 0x3836C000, + 0x38370000, 0x38374000, 0x38378000, 0x3837C000, 0x38380000, 0x38384000, + 0x38388000, 0x3838C000, 0x38390000, 0x38394000, 0x38398000, 0x3839C000, + 0x383A0000, 0x383A4000, 0x383A8000, 0x383AC000, 0x383B0000, 0x383B4000, + 0x383B8000, 0x383BC000, 0x383C0000, 0x383C4000, 0x383C8000, 0x383CC000, + 0x383D0000, 0x383D4000, 0x383D8000, 0x383DC000, 0x383E0000, 0x383E4000, + 0x383E8000, 0x383EC000, 0x383F0000, 0x383F4000, 0x383F8000, 0x383FC000, + 0x38400000, 0x38404000, 0x38408000, 0x3840C000, 0x38410000, 0x38414000, + 0x38418000, 0x3841C000, 0x38420000, 0x38424000, 0x38428000, 0x3842C000, + 0x38430000, 0x38434000, 0x38438000, 0x3843C000, 0x38440000, 0x38444000, + 0x38448000, 0x3844C000, 0x38450000, 0x38454000, 0x38458000, 0x3845C000, + 0x38460000, 0x38464000, 0x38468000, 0x3846C000, 0x38470000, 0x38474000, + 0x38478000, 0x3847C000, 0x38480000, 0x38484000, 0x38488000, 0x3848C000, + 0x38490000, 0x38494000, 0x38498000, 0x3849C000, 0x384A0000, 0x384A4000, + 0x384A8000, 0x384AC000, 0x384B0000, 0x384B4000, 0x384B8000, 0x384BC000, + 0x384C0000, 0x384C4000, 0x384C8000, 0x384CC000, 0x384D0000, 0x384D4000, + 0x384D8000, 0x384DC000, 0x384E0000, 0x384E4000, 0x384E8000, 0x384EC000, + 0x384F0000, 0x384F4000, 0x384F8000, 0x384FC000, 0x38500000, 0x38504000, + 0x38508000, 0x3850C000, 0x38510000, 0x38514000, 0x38518000, 0x3851C000, + 0x38520000, 0x38524000, 0x38528000, 0x3852C000, 0x38530000, 0x38534000, + 0x38538000, 0x3853C000, 0x38540000, 0x38544000, 0x38548000, 0x3854C000, + 0x38550000, 0x38554000, 0x38558000, 0x3855C000, 0x38560000, 0x38564000, + 0x38568000, 0x3856C000, 0x38570000, 0x38574000, 0x38578000, 0x3857C000, + 0x38580000, 0x38584000, 0x38588000, 0x3858C000, 0x38590000, 0x38594000, + 0x38598000, 0x3859C000, 0x385A0000, 0x385A4000, 0x385A8000, 0x385AC000, + 0x385B0000, 0x385B4000, 0x385B8000, 0x385BC000, 0x385C0000, 0x385C4000, + 0x385C8000, 0x385CC000, 0x385D0000, 0x385D4000, 0x385D8000, 0x385DC000, + 0x385E0000, 0x385E4000, 0x385E8000, 0x385EC000, 0x385F0000, 0x385F4000, + 0x385F8000, 0x385FC000, 0x38600000, 0x38604000, 0x38608000, 0x3860C000, + 0x38610000, 0x38614000, 0x38618000, 0x3861C000, 0x38620000, 0x38624000, + 0x38628000, 0x3862C000, 0x38630000, 0x38634000, 0x38638000, 0x3863C000, + 0x38640000, 0x38644000, 0x38648000, 0x3864C000, 0x38650000, 0x38654000, + 0x38658000, 0x3865C000, 0x38660000, 0x38664000, 0x38668000, 0x3866C000, + 0x38670000, 0x38674000, 0x38678000, 0x3867C000, 0x38680000, 0x38684000, + 0x38688000, 0x3868C000, 0x38690000, 0x38694000, 0x38698000, 0x3869C000, + 0x386A0000, 0x386A4000, 0x386A8000, 0x386AC000, 0x386B0000, 0x386B4000, + 0x386B8000, 0x386BC000, 0x386C0000, 0x386C4000, 0x386C8000, 0x386CC000, + 0x386D0000, 0x386D4000, 0x386D8000, 0x386DC000, 0x386E0000, 0x386E4000, + 0x386E8000, 0x386EC000, 0x386F0000, 0x386F4000, 0x386F8000, 0x386FC000, + 0x38700000, 0x38704000, 0x38708000, 0x3870C000, 0x38710000, 0x38714000, + 0x38718000, 0x3871C000, 0x38720000, 0x38724000, 0x38728000, 0x3872C000, + 0x38730000, 0x38734000, 0x38738000, 0x3873C000, 0x38740000, 0x38744000, + 0x38748000, 0x3874C000, 0x38750000, 0x38754000, 0x38758000, 0x3875C000, + 0x38760000, 0x38764000, 0x38768000, 0x3876C000, 0x38770000, 0x38774000, + 0x38778000, 0x3877C000, 0x38780000, 0x38784000, 0x38788000, 0x3878C000, + 0x38790000, 0x38794000, 0x38798000, 0x3879C000, 0x387A0000, 0x387A4000, + 0x387A8000, 0x387AC000, 0x387B0000, 0x387B4000, 0x387B8000, 0x387BC000, + 0x387C0000, 0x387C4000, 0x387C8000, 0x387CC000, 0x387D0000, 0x387D4000, + 0x387D8000, 0x387DC000, 0x387E0000, 0x387E4000, 0x387E8000, 0x387EC000, + 0x387F0000, 0x387F4000, 0x387F8000, 0x387FC000, 0x38000000, 0x38002000, + 0x38004000, 0x38006000, 0x38008000, 0x3800A000, 0x3800C000, 0x3800E000, + 0x38010000, 0x38012000, 0x38014000, 0x38016000, 0x38018000, 0x3801A000, + 0x3801C000, 0x3801E000, 0x38020000, 0x38022000, 0x38024000, 0x38026000, + 0x38028000, 0x3802A000, 0x3802C000, 0x3802E000, 0x38030000, 0x38032000, + 0x38034000, 0x38036000, 0x38038000, 0x3803A000, 0x3803C000, 0x3803E000, + 0x38040000, 0x38042000, 0x38044000, 0x38046000, 0x38048000, 0x3804A000, + 0x3804C000, 0x3804E000, 0x38050000, 0x38052000, 0x38054000, 0x38056000, + 0x38058000, 0x3805A000, 0x3805C000, 0x3805E000, 0x38060000, 0x38062000, + 0x38064000, 0x38066000, 0x38068000, 0x3806A000, 0x3806C000, 0x3806E000, + 0x38070000, 0x38072000, 0x38074000, 0x38076000, 0x38078000, 0x3807A000, + 0x3807C000, 0x3807E000, 0x38080000, 0x38082000, 0x38084000, 0x38086000, + 0x38088000, 0x3808A000, 0x3808C000, 0x3808E000, 0x38090000, 0x38092000, + 0x38094000, 0x38096000, 0x38098000, 0x3809A000, 0x3809C000, 0x3809E000, + 0x380A0000, 0x380A2000, 0x380A4000, 0x380A6000, 0x380A8000, 0x380AA000, + 0x380AC000, 0x380AE000, 0x380B0000, 0x380B2000, 0x380B4000, 0x380B6000, + 0x380B8000, 0x380BA000, 0x380BC000, 0x380BE000, 0x380C0000, 0x380C2000, + 0x380C4000, 0x380C6000, 0x380C8000, 0x380CA000, 0x380CC000, 0x380CE000, + 0x380D0000, 0x380D2000, 0x380D4000, 0x380D6000, 0x380D8000, 0x380DA000, + 0x380DC000, 0x380DE000, 0x380E0000, 0x380E2000, 0x380E4000, 0x380E6000, + 0x380E8000, 0x380EA000, 0x380EC000, 0x380EE000, 0x380F0000, 0x380F2000, + 0x380F4000, 0x380F6000, 0x380F8000, 0x380FA000, 0x380FC000, 0x380FE000, + 0x38100000, 0x38102000, 0x38104000, 0x38106000, 0x38108000, 0x3810A000, + 0x3810C000, 0x3810E000, 0x38110000, 0x38112000, 0x38114000, 0x38116000, + 0x38118000, 0x3811A000, 0x3811C000, 0x3811E000, 0x38120000, 0x38122000, + 0x38124000, 0x38126000, 0x38128000, 0x3812A000, 0x3812C000, 0x3812E000, + 0x38130000, 0x38132000, 0x38134000, 0x38136000, 0x38138000, 0x3813A000, + 0x3813C000, 0x3813E000, 0x38140000, 0x38142000, 0x38144000, 0x38146000, + 0x38148000, 0x3814A000, 0x3814C000, 0x3814E000, 0x38150000, 0x38152000, + 0x38154000, 0x38156000, 0x38158000, 0x3815A000, 0x3815C000, 0x3815E000, + 0x38160000, 0x38162000, 0x38164000, 0x38166000, 0x38168000, 0x3816A000, + 0x3816C000, 0x3816E000, 0x38170000, 0x38172000, 0x38174000, 0x38176000, + 0x38178000, 0x3817A000, 0x3817C000, 0x3817E000, 0x38180000, 0x38182000, + 0x38184000, 0x38186000, 0x38188000, 0x3818A000, 0x3818C000, 0x3818E000, + 0x38190000, 0x38192000, 0x38194000, 0x38196000, 0x38198000, 0x3819A000, + 0x3819C000, 0x3819E000, 0x381A0000, 0x381A2000, 0x381A4000, 0x381A6000, + 0x381A8000, 0x381AA000, 0x381AC000, 0x381AE000, 0x381B0000, 0x381B2000, + 0x381B4000, 0x381B6000, 0x381B8000, 0x381BA000, 0x381BC000, 0x381BE000, + 0x381C0000, 0x381C2000, 0x381C4000, 0x381C6000, 0x381C8000, 0x381CA000, + 0x381CC000, 0x381CE000, 0x381D0000, 0x381D2000, 0x381D4000, 0x381D6000, + 0x381D8000, 0x381DA000, 0x381DC000, 0x381DE000, 0x381E0000, 0x381E2000, + 0x381E4000, 0x381E6000, 0x381E8000, 0x381EA000, 0x381EC000, 0x381EE000, + 0x381F0000, 0x381F2000, 0x381F4000, 0x381F6000, 0x381F8000, 0x381FA000, + 0x381FC000, 0x381FE000, 0x38200000, 0x38202000, 0x38204000, 0x38206000, + 0x38208000, 0x3820A000, 0x3820C000, 0x3820E000, 0x38210000, 0x38212000, + 0x38214000, 0x38216000, 0x38218000, 0x3821A000, 0x3821C000, 0x3821E000, + 0x38220000, 0x38222000, 0x38224000, 0x38226000, 0x38228000, 0x3822A000, + 0x3822C000, 0x3822E000, 0x38230000, 0x38232000, 0x38234000, 0x38236000, + 0x38238000, 0x3823A000, 0x3823C000, 0x3823E000, 0x38240000, 0x38242000, + 0x38244000, 0x38246000, 0x38248000, 0x3824A000, 0x3824C000, 0x3824E000, + 0x38250000, 0x38252000, 0x38254000, 0x38256000, 0x38258000, 0x3825A000, + 0x3825C000, 0x3825E000, 0x38260000, 0x38262000, 0x38264000, 0x38266000, + 0x38268000, 0x3826A000, 0x3826C000, 0x3826E000, 0x38270000, 0x38272000, + 0x38274000, 0x38276000, 0x38278000, 0x3827A000, 0x3827C000, 0x3827E000, + 0x38280000, 0x38282000, 0x38284000, 0x38286000, 0x38288000, 0x3828A000, + 0x3828C000, 0x3828E000, 0x38290000, 0x38292000, 0x38294000, 0x38296000, + 0x38298000, 0x3829A000, 0x3829C000, 0x3829E000, 0x382A0000, 0x382A2000, + 0x382A4000, 0x382A6000, 0x382A8000, 0x382AA000, 0x382AC000, 0x382AE000, + 0x382B0000, 0x382B2000, 0x382B4000, 0x382B6000, 0x382B8000, 0x382BA000, + 0x382BC000, 0x382BE000, 0x382C0000, 0x382C2000, 0x382C4000, 0x382C6000, + 0x382C8000, 0x382CA000, 0x382CC000, 0x382CE000, 0x382D0000, 0x382D2000, + 0x382D4000, 0x382D6000, 0x382D8000, 0x382DA000, 0x382DC000, 0x382DE000, + 0x382E0000, 0x382E2000, 0x382E4000, 0x382E6000, 0x382E8000, 0x382EA000, + 0x382EC000, 0x382EE000, 0x382F0000, 0x382F2000, 0x382F4000, 0x382F6000, + 0x382F8000, 0x382FA000, 0x382FC000, 0x382FE000, 0x38300000, 0x38302000, + 0x38304000, 0x38306000, 0x38308000, 0x3830A000, 0x3830C000, 0x3830E000, + 0x38310000, 0x38312000, 0x38314000, 0x38316000, 0x38318000, 0x3831A000, + 0x3831C000, 0x3831E000, 0x38320000, 0x38322000, 0x38324000, 0x38326000, + 0x38328000, 0x3832A000, 0x3832C000, 0x3832E000, 0x38330000, 0x38332000, + 0x38334000, 0x38336000, 0x38338000, 0x3833A000, 0x3833C000, 0x3833E000, + 0x38340000, 0x38342000, 0x38344000, 0x38346000, 0x38348000, 0x3834A000, + 0x3834C000, 0x3834E000, 0x38350000, 0x38352000, 0x38354000, 0x38356000, + 0x38358000, 0x3835A000, 0x3835C000, 0x3835E000, 0x38360000, 0x38362000, + 0x38364000, 0x38366000, 0x38368000, 0x3836A000, 0x3836C000, 0x3836E000, + 0x38370000, 0x38372000, 0x38374000, 0x38376000, 0x38378000, 0x3837A000, + 0x3837C000, 0x3837E000, 0x38380000, 0x38382000, 0x38384000, 0x38386000, + 0x38388000, 0x3838A000, 0x3838C000, 0x3838E000, 0x38390000, 0x38392000, + 0x38394000, 0x38396000, 0x38398000, 0x3839A000, 0x3839C000, 0x3839E000, + 0x383A0000, 0x383A2000, 0x383A4000, 0x383A6000, 0x383A8000, 0x383AA000, + 0x383AC000, 0x383AE000, 0x383B0000, 0x383B2000, 0x383B4000, 0x383B6000, + 0x383B8000, 0x383BA000, 0x383BC000, 0x383BE000, 0x383C0000, 0x383C2000, + 0x383C4000, 0x383C6000, 0x383C8000, 0x383CA000, 0x383CC000, 0x383CE000, + 0x383D0000, 0x383D2000, 0x383D4000, 0x383D6000, 0x383D8000, 0x383DA000, + 0x383DC000, 0x383DE000, 0x383E0000, 0x383E2000, 0x383E4000, 0x383E6000, + 0x383E8000, 0x383EA000, 0x383EC000, 0x383EE000, 0x383F0000, 0x383F2000, + 0x383F4000, 0x383F6000, 0x383F8000, 0x383FA000, 0x383FC000, 0x383FE000, + 0x38400000, 0x38402000, 0x38404000, 0x38406000, 0x38408000, 0x3840A000, + 0x3840C000, 0x3840E000, 0x38410000, 0x38412000, 0x38414000, 0x38416000, + 0x38418000, 0x3841A000, 0x3841C000, 0x3841E000, 0x38420000, 0x38422000, + 0x38424000, 0x38426000, 0x38428000, 0x3842A000, 0x3842C000, 0x3842E000, + 0x38430000, 0x38432000, 0x38434000, 0x38436000, 0x38438000, 0x3843A000, + 0x3843C000, 0x3843E000, 0x38440000, 0x38442000, 0x38444000, 0x38446000, + 0x38448000, 0x3844A000, 0x3844C000, 0x3844E000, 0x38450000, 0x38452000, + 0x38454000, 0x38456000, 0x38458000, 0x3845A000, 0x3845C000, 0x3845E000, + 0x38460000, 0x38462000, 0x38464000, 0x38466000, 0x38468000, 0x3846A000, + 0x3846C000, 0x3846E000, 0x38470000, 0x38472000, 0x38474000, 0x38476000, + 0x38478000, 0x3847A000, 0x3847C000, 0x3847E000, 0x38480000, 0x38482000, + 0x38484000, 0x38486000, 0x38488000, 0x3848A000, 0x3848C000, 0x3848E000, + 0x38490000, 0x38492000, 0x38494000, 0x38496000, 0x38498000, 0x3849A000, + 0x3849C000, 0x3849E000, 0x384A0000, 0x384A2000, 0x384A4000, 0x384A6000, + 0x384A8000, 0x384AA000, 0x384AC000, 0x384AE000, 0x384B0000, 0x384B2000, + 0x384B4000, 0x384B6000, 0x384B8000, 0x384BA000, 0x384BC000, 0x384BE000, + 0x384C0000, 0x384C2000, 0x384C4000, 0x384C6000, 0x384C8000, 0x384CA000, + 0x384CC000, 0x384CE000, 0x384D0000, 0x384D2000, 0x384D4000, 0x384D6000, + 0x384D8000, 0x384DA000, 0x384DC000, 0x384DE000, 0x384E0000, 0x384E2000, + 0x384E4000, 0x384E6000, 0x384E8000, 0x384EA000, 0x384EC000, 0x384EE000, + 0x384F0000, 0x384F2000, 0x384F4000, 0x384F6000, 0x384F8000, 0x384FA000, + 0x384FC000, 0x384FE000, 0x38500000, 0x38502000, 0x38504000, 0x38506000, + 0x38508000, 0x3850A000, 0x3850C000, 0x3850E000, 0x38510000, 0x38512000, + 0x38514000, 0x38516000, 0x38518000, 0x3851A000, 0x3851C000, 0x3851E000, + 0x38520000, 0x38522000, 0x38524000, 0x38526000, 0x38528000, 0x3852A000, + 0x3852C000, 0x3852E000, 0x38530000, 0x38532000, 0x38534000, 0x38536000, + 0x38538000, 0x3853A000, 0x3853C000, 0x3853E000, 0x38540000, 0x38542000, + 0x38544000, 0x38546000, 0x38548000, 0x3854A000, 0x3854C000, 0x3854E000, + 0x38550000, 0x38552000, 0x38554000, 0x38556000, 0x38558000, 0x3855A000, + 0x3855C000, 0x3855E000, 0x38560000, 0x38562000, 0x38564000, 0x38566000, + 0x38568000, 0x3856A000, 0x3856C000, 0x3856E000, 0x38570000, 0x38572000, + 0x38574000, 0x38576000, 0x38578000, 0x3857A000, 0x3857C000, 0x3857E000, + 0x38580000, 0x38582000, 0x38584000, 0x38586000, 0x38588000, 0x3858A000, + 0x3858C000, 0x3858E000, 0x38590000, 0x38592000, 0x38594000, 0x38596000, + 0x38598000, 0x3859A000, 0x3859C000, 0x3859E000, 0x385A0000, 0x385A2000, + 0x385A4000, 0x385A6000, 0x385A8000, 0x385AA000, 0x385AC000, 0x385AE000, + 0x385B0000, 0x385B2000, 0x385B4000, 0x385B6000, 0x385B8000, 0x385BA000, + 0x385BC000, 0x385BE000, 0x385C0000, 0x385C2000, 0x385C4000, 0x385C6000, + 0x385C8000, 0x385CA000, 0x385CC000, 0x385CE000, 0x385D0000, 0x385D2000, + 0x385D4000, 0x385D6000, 0x385D8000, 0x385DA000, 0x385DC000, 0x385DE000, + 0x385E0000, 0x385E2000, 0x385E4000, 0x385E6000, 0x385E8000, 0x385EA000, + 0x385EC000, 0x385EE000, 0x385F0000, 0x385F2000, 0x385F4000, 0x385F6000, + 0x385F8000, 0x385FA000, 0x385FC000, 0x385FE000, 0x38600000, 0x38602000, + 0x38604000, 0x38606000, 0x38608000, 0x3860A000, 0x3860C000, 0x3860E000, + 0x38610000, 0x38612000, 0x38614000, 0x38616000, 0x38618000, 0x3861A000, + 0x3861C000, 0x3861E000, 0x38620000, 0x38622000, 0x38624000, 0x38626000, + 0x38628000, 0x3862A000, 0x3862C000, 0x3862E000, 0x38630000, 0x38632000, + 0x38634000, 0x38636000, 0x38638000, 0x3863A000, 0x3863C000, 0x3863E000, + 0x38640000, 0x38642000, 0x38644000, 0x38646000, 0x38648000, 0x3864A000, + 0x3864C000, 0x3864E000, 0x38650000, 0x38652000, 0x38654000, 0x38656000, + 0x38658000, 0x3865A000, 0x3865C000, 0x3865E000, 0x38660000, 0x38662000, + 0x38664000, 0x38666000, 0x38668000, 0x3866A000, 0x3866C000, 0x3866E000, + 0x38670000, 0x38672000, 0x38674000, 0x38676000, 0x38678000, 0x3867A000, + 0x3867C000, 0x3867E000, 0x38680000, 0x38682000, 0x38684000, 0x38686000, + 0x38688000, 0x3868A000, 0x3868C000, 0x3868E000, 0x38690000, 0x38692000, + 0x38694000, 0x38696000, 0x38698000, 0x3869A000, 0x3869C000, 0x3869E000, + 0x386A0000, 0x386A2000, 0x386A4000, 0x386A6000, 0x386A8000, 0x386AA000, + 0x386AC000, 0x386AE000, 0x386B0000, 0x386B2000, 0x386B4000, 0x386B6000, + 0x386B8000, 0x386BA000, 0x386BC000, 0x386BE000, 0x386C0000, 0x386C2000, + 0x386C4000, 0x386C6000, 0x386C8000, 0x386CA000, 0x386CC000, 0x386CE000, + 0x386D0000, 0x386D2000, 0x386D4000, 0x386D6000, 0x386D8000, 0x386DA000, + 0x386DC000, 0x386DE000, 0x386E0000, 0x386E2000, 0x386E4000, 0x386E6000, + 0x386E8000, 0x386EA000, 0x386EC000, 0x386EE000, 0x386F0000, 0x386F2000, + 0x386F4000, 0x386F6000, 0x386F8000, 0x386FA000, 0x386FC000, 0x386FE000, + 0x38700000, 0x38702000, 0x38704000, 0x38706000, 0x38708000, 0x3870A000, + 0x3870C000, 0x3870E000, 0x38710000, 0x38712000, 0x38714000, 0x38716000, + 0x38718000, 0x3871A000, 0x3871C000, 0x3871E000, 0x38720000, 0x38722000, + 0x38724000, 0x38726000, 0x38728000, 0x3872A000, 0x3872C000, 0x3872E000, + 0x38730000, 0x38732000, 0x38734000, 0x38736000, 0x38738000, 0x3873A000, + 0x3873C000, 0x3873E000, 0x38740000, 0x38742000, 0x38744000, 0x38746000, + 0x38748000, 0x3874A000, 0x3874C000, 0x3874E000, 0x38750000, 0x38752000, + 0x38754000, 0x38756000, 0x38758000, 0x3875A000, 0x3875C000, 0x3875E000, + 0x38760000, 0x38762000, 0x38764000, 0x38766000, 0x38768000, 0x3876A000, + 0x3876C000, 0x3876E000, 0x38770000, 0x38772000, 0x38774000, 0x38776000, + 0x38778000, 0x3877A000, 0x3877C000, 0x3877E000, 0x38780000, 0x38782000, + 0x38784000, 0x38786000, 0x38788000, 0x3878A000, 0x3878C000, 0x3878E000, + 0x38790000, 0x38792000, 0x38794000, 0x38796000, 0x38798000, 0x3879A000, + 0x3879C000, 0x3879E000, 0x387A0000, 0x387A2000, 0x387A4000, 0x387A6000, + 0x387A8000, 0x387AA000, 0x387AC000, 0x387AE000, 0x387B0000, 0x387B2000, + 0x387B4000, 0x387B6000, 0x387B8000, 0x387BA000, 0x387BC000, 0x387BE000, + 0x387C0000, 0x387C2000, 0x387C4000, 0x387C6000, 0x387C8000, 0x387CA000, + 0x387CC000, 0x387CE000, 0x387D0000, 0x387D2000, 0x387D4000, 0x387D6000, + 0x387D8000, 0x387DA000, 0x387DC000, 0x387DE000, 0x387E0000, 0x387E2000, + 0x387E4000, 0x387E6000, 0x387E8000, 0x387EA000, 0x387EC000, 0x387EE000, + 0x387F0000, 0x387F2000, 0x387F4000, 0x387F6000, 0x387F8000, 0x387FA000, + 0x387FC000, 0x387FE000}; + static const uint32_t exponent_table[64] = { + 0x00000000, 0x00800000, 0x01000000, 0x01800000, 0x02000000, 0x02800000, + 0x03000000, 0x03800000, 0x04000000, 0x04800000, 0x05000000, 0x05800000, + 0x06000000, 0x06800000, 0x07000000, 0x07800000, 0x08000000, 0x08800000, + 0x09000000, 0x09800000, 0x0A000000, 0x0A800000, 0x0B000000, 0x0B800000, + 0x0C000000, 0x0C800000, 0x0D000000, 0x0D800000, 0x0E000000, 0x0E800000, + 0x0F000000, 0x47800000, 0x80000000, 0x80800000, 0x81000000, 0x81800000, + 0x82000000, 0x82800000, 0x83000000, 0x83800000, 0x84000000, 0x84800000, + 0x85000000, 0x85800000, 0x86000000, 0x86800000, 0x87000000, 0x87800000, + 0x88000000, 0x88800000, 0x89000000, 0x89800000, 0x8A000000, 0x8A800000, + 0x8B000000, 0x8B800000, 0x8C000000, 0x8C800000, 0x8D000000, 0x8D800000, + 0x8E000000, 0x8E800000, 0x8F000000, 0xC7800000}; + static const uint16_t offset_table[64] = { + 0, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, + 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, + 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 0, + 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, + 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, + 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024}; + uint16_t hval = static_cast(val >> 10); + uint32_t bits = + mantissa_table[offset_table[hval] + (val & 0x3FF)] + exponent_table[hval]; + float result; + std::memcpy(&result, &bits, sizeof(result)); + return result; +} + +// Refer: https://github.com/Maratyszcza/FP16/blob/master/third-party/half.hpp +static inline uint16_t float16(float val) { + static const uint16_t base_table[512] = { + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, + 0x0000, 0x0000, 0x0000, 0x0000, 0x0001, 0x0002, 0x0004, 0x0008, 0x0010, + 0x0020, 0x0040, 0x0080, 0x0100, 0x0200, 0x0400, 0x0800, 0x0C00, 0x1000, + 0x1400, 0x1800, 0x1C00, 0x2000, 0x2400, 0x2800, 0x2C00, 0x3000, 0x3400, + 0x3800, 0x3C00, 0x4000, 0x4400, 0x4800, 0x4C00, 0x5000, 0x5400, 0x5800, + 0x5C00, 0x6000, 0x6400, 0x6800, 0x6C00, 0x7000, 0x7400, 0x7800, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, + 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, + 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8001, + 0x8002, 0x8004, 0x8008, 0x8010, 0x8020, 0x8040, 0x8080, 0x8100, 0x8200, + 0x8400, 0x8800, 0x8C00, 0x9000, 0x9400, 0x9800, 0x9C00, 0xA000, 0xA400, + 0xA800, 0xAC00, 0xB000, 0xB400, 0xB800, 0xBC00, 0xC000, 0xC400, 0xC800, + 0xCC00, 0xD000, 0xD400, 0xD800, 0xDC00, 0xE000, 0xE400, 0xE800, 0xEC00, + 0xF000, 0xF400, 0xF800, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, + 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00}; + static const uint8_t shift_table[512] = { + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 23, 22, 21, 20, 19, + 18, 17, 16, 15, 14, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, + 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 13, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 23, + 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 13, 13, 13, 13, 13, 13, 13, 13, + 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, + 13, 13, 13, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 13}; + uint32_t *p = reinterpret_cast(&val); + uint32_t hbits = + base_table[*p >> 23] + + static_cast((*p & 0x7FFFFF) >> shift_table[*p >> 23]); + hbits += (((*p & 0x7FFFFF) >> (shift_table[*p >> 23] - 1)) | + (((*p >> 23) & 0xFF) == 102)) & + ((hbits & 0x7C00) != 0x7C00); + return static_cast(hbits); +} +#if defined(__F16C__) && defined(__AVX512F__) +static inline void convert_fp16_to_fp32_avx512f(const uint16_t *arr, + size_t size, float *out) { + const uint16_t *last = arr + size; + const uint16_t *last_aligned = arr + ((size >> 5) << 5); + + if (((uintptr_t)arr & 0x1f) == 0 && ((uintptr_t)out & 0x3f) == 0) { + for (; arr != last_aligned; arr += 32, out += 32) { + _mm512_store_ps(out + 0, + _mm512_cvtph_ps(_mm256_load_si256((__m256i *)(arr + 0)))); + _mm512_store_ps( + out + 16, _mm512_cvtph_ps(_mm256_load_si256((__m256i *)(arr + 16)))); + } + + if (last >= last_aligned + 16) { + _mm512_store_ps(out, _mm512_cvtph_ps(_mm256_load_si256((__m256i *)arr))); + arr += 16; + out += 16; + } + if (last >= arr + 8) { + _mm256_store_ps(out, _mm256_cvtph_ps(_mm_load_si128((__m128i *)arr))); + arr += 8; + out += 8; + } + } else { + for (; arr != last_aligned; arr += 32, out += 32) { + _mm512_storeu_ps( + out + 0, _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)(arr + 0)))); + _mm512_storeu_ps( + out + 16, _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)(arr + 16)))); + } + + if (last >= last_aligned + 16) { + _mm512_storeu_ps(out, + _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)arr))); + arr += 16; + out += 16; + } + if (last >= arr + 8) { + _mm256_storeu_ps(out, _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)arr))); + arr += 8; + out += 8; + } + } + switch (last - arr) { + case 7: + out[6] = float32(arr[6]); + /* FALLTHRU */ + case 6: + out[5] = float32(arr[5]); + /* FALLTHRU */ + case 5: + out[4] = float32(arr[4]); + /* FALLTHRU */ + case 4: + out[3] = float32(arr[3]); + /* FALLTHRU */ + case 3: + out[2] = float32(arr[2]); + /* FALLTHRU */ + case 2: + out[1] = float32(arr[1]); + /* FALLTHRU */ + case 1: + out[0] = float32(arr[0]); + } +} + +static inline void convert_fp16_to_fp32_avx512f(const uint16_t *arr, + size_t size, float norm, + float *out) { + const uint16_t *last = arr + size; + const uint16_t *last_aligned = arr + ((size >> 5) << 5); + __m512 zmm_norm = _mm512_set1_ps(norm); + + if (((uintptr_t)arr & 0x1f) == 0 && ((uintptr_t)out & 0x3f) == 0) { + for (; arr != last_aligned; arr += 32, out += 32) { + __m512 zmm_0 = _mm512_div_ps( + _mm512_cvtph_ps(_mm256_load_si256((__m256i *)(arr + 0))), zmm_norm); + __m512 zmm_1 = _mm512_div_ps( + _mm512_cvtph_ps(_mm256_load_si256((__m256i *)(arr + 16))), zmm_norm); + _mm512_store_ps(out + 0, zmm_0); + _mm512_store_ps(out + 16, zmm_1); + } + + if (last >= last_aligned + 16) { + _mm512_store_ps( + out, _mm512_div_ps(_mm512_cvtph_ps(_mm256_load_si256((__m256i *)arr)), + zmm_norm)); + arr += 16; + out += 16; + } + if (last >= arr + 8) { + _mm256_store_ps( + out, _mm256_div_ps(_mm256_cvtph_ps(_mm_load_si128((__m128i *)arr)), + _mm256_set1_ps(norm))); + arr += 8; + out += 8; + } + } else { + for (; arr != last_aligned; arr += 32, out += 32) { + __m512 zmm_0 = _mm512_div_ps( + _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)(arr + 0))), zmm_norm); + __m512 zmm_1 = _mm512_div_ps( + _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)(arr + 16))), zmm_norm); + _mm512_storeu_ps(out + 0, zmm_0); + _mm512_storeu_ps(out + 16, zmm_1); + } + + if (last >= last_aligned + 16) { + _mm512_storeu_ps( + out, + _mm512_div_ps(_mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)arr)), + zmm_norm)); + arr += 16; + out += 16; + } + if (last >= arr + 8) { + _mm256_storeu_ps( + out, _mm256_div_ps(_mm256_cvtph_ps(_mm_loadu_si128((__m128i *)arr)), + _mm256_set1_ps(norm))); + arr += 8; + out += 8; + } + } + switch (last - arr) { + case 7: + out[6] = float32(arr[6]) / norm; + /* FALLTHRU */ + case 6: + out[5] = float32(arr[5]) / norm; + /* FALLTHRU */ + case 5: + out[4] = float32(arr[4]) / norm; + /* FALLTHRU */ + case 4: + out[3] = float32(arr[3]) / norm; + /* FALLTHRU */ + case 3: + out[2] = float32(arr[2]) / norm; + /* FALLTHRU */ + case 2: + out[1] = float32(arr[1]) / norm; + /* FALLTHRU */ + case 1: + out[0] = float32(arr[0]) / norm; + } +} + +static inline void convert_fp32_to_fp16_avx512f(const float *arr, size_t size, + uint16_t *out) { + const float *last = arr + size; + const float *last_aligned = arr + ((size >> 5) << 5); + + if (((uintptr_t)arr & 0x3f) == 0 && ((uintptr_t)out & 0x1f) == 0) { + for (; arr != last_aligned; arr += 32, out += 32) { + _mm256_store_si256( + (__m256i *)(out + 0), + _mm512_cvtps_ph(_mm512_load_ps(arr + 0), _MM_FROUND_NO_EXC)); + _mm256_store_si256( + (__m256i *)(out + 16), + _mm512_cvtps_ph(_mm512_load_ps(arr + 16), _MM_FROUND_NO_EXC)); + } + + if (last >= last_aligned + 16) { + _mm256_store_si256( + (__m256i *)(out + 0), + _mm512_cvtps_ph(_mm512_load_ps(arr + 0), _MM_FROUND_NO_EXC)); + arr += 16; + out += 16; + } + if (last >= arr + 8) { + _mm_store_si128( + (__m128i *)(out + 0), + _mm256_cvtps_ph(_mm256_load_ps(arr + 0), _MM_FROUND_NO_EXC)); + arr += 8; + out += 8; + } + } else { + for (; arr != last_aligned; arr += 32, out += 32) { + _mm256_storeu_si256( + (__m256i *)(out + 0), + _mm512_cvtps_ph(_mm512_loadu_ps(arr + 0), _MM_FROUND_NO_EXC)); + _mm256_storeu_si256( + (__m256i *)(out + 16), + _mm512_cvtps_ph(_mm512_loadu_ps(arr + 16), _MM_FROUND_NO_EXC)); + } + + if (last >= last_aligned + 16) { + _mm256_storeu_si256( + (__m256i *)(out + 0), + _mm512_cvtps_ph(_mm512_loadu_ps(arr + 0), _MM_FROUND_NO_EXC)); + arr += 16; + out += 16; + } + if (last >= arr + 8) { + _mm_storeu_si128( + (__m128i *)(out + 0), + _mm256_cvtps_ph(_mm256_loadu_ps(arr + 0), _MM_FROUND_NO_EXC)); + arr += 8; + out += 8; + } + } + switch (last - arr) { + case 7: + out[6] = float16(arr[6]); + /* FALLTHRU */ + case 6: + out[5] = float16(arr[5]); + /* FALLTHRU */ + case 5: + out[4] = float16(arr[4]); + /* FALLTHRU */ + case 4: + out[3] = float16(arr[3]); + /* FALLTHRU */ + case 3: + out[2] = float16(arr[2]); + /* FALLTHRU */ + case 2: + out[1] = float16(arr[1]); + /* FALLTHRU */ + case 1: + out[0] = float16(arr[0]); + } +} + +static inline void convert_fp32_to_fp16_avx512f(const float *arr, size_t size, + float norm, uint16_t *out) { + const float *last = arr + size; + const float *last_aligned = arr + ((size >> 5) << 5); + __m512 zmm_norm = _mm512_set1_ps(norm); + + if (((uintptr_t)arr & 0x3f) == 0 && ((uintptr_t)out & 0x1f) == 0) { + for (; arr != last_aligned; arr += 32, out += 32) { + __m512 zmm_0 = _mm512_div_ps(_mm512_load_ps(arr + 0), zmm_norm); + __m512 zmm_1 = _mm512_div_ps(_mm512_load_ps(arr + 16), zmm_norm); + _mm256_store_si256((__m256i *)(out + 0), + _mm512_cvtps_ph(zmm_0, _MM_FROUND_NO_EXC)); + _mm256_store_si256((__m256i *)(out + 16), + _mm512_cvtps_ph(zmm_1, _MM_FROUND_NO_EXC)); + } + + if (last >= last_aligned + 16) { + _mm256_store_si256( + (__m256i *)out, + _mm512_cvtps_ph(_mm512_div_ps(_mm512_load_ps(arr), zmm_norm), + _MM_FROUND_NO_EXC)); + arr += 16; + out += 16; + } + if (last >= arr + 8) { + _mm_store_si128((__m128i *)out, + _mm256_cvtps_ph(_mm256_div_ps(_mm256_load_ps(arr), + _mm256_set1_ps(norm)), + _MM_FROUND_NO_EXC)); + arr += 8; + out += 8; + } + } else { + for (; arr != last_aligned; arr += 32, out += 32) { + __m512 zmm_0 = _mm512_div_ps(_mm512_loadu_ps(arr + 0), zmm_norm); + __m512 zmm_1 = _mm512_div_ps(_mm512_loadu_ps(arr + 16), zmm_norm); + _mm256_storeu_si256((__m256i *)(out + 0), + _mm512_cvtps_ph(zmm_0, _MM_FROUND_NO_EXC)); + _mm256_storeu_si256((__m256i *)(out + 16), + _mm512_cvtps_ph(zmm_1, _MM_FROUND_NO_EXC)); + } + + if (last >= last_aligned + 16) { + _mm256_storeu_si256( + (__m256i *)out, + _mm512_cvtps_ph(_mm512_div_ps(_mm512_loadu_ps(arr), zmm_norm), + _MM_FROUND_NO_EXC)); + arr += 16; + out += 16; + } + if (last >= arr + 8) { + _mm_storeu_si128((__m128i *)out, + _mm256_cvtps_ph(_mm256_div_ps(_mm256_loadu_ps(arr), + _mm256_set1_ps(norm)), + _MM_FROUND_NO_EXC)); + arr += 8; + out += 8; + } + } + switch (last - arr) { + case 7: + out[6] = float16(arr[6] / norm); + /* FALLTHRU */ + case 6: + out[5] = float16(arr[5] / norm); + /* FALLTHRU */ + case 5: + out[4] = float16(arr[4] / norm); + /* FALLTHRU */ + case 4: + out[3] = float16(arr[3] / norm); + /* FALLTHRU */ + case 3: + out[2] = float16(arr[2] / norm); + /* FALLTHRU */ + case 2: + out[1] = float16(arr[1] / norm); + /* FALLTHRU */ + case 1: + out[0] = float16(arr[0] / norm); + } +} +#endif //__F16C__ && __AVX512F__ + +#if defined(__F16C__) && defined(__AVX__) +static inline void convert_fp16_to_fp32_avx(const uint16_t *arr, size_t size, + float *out) { + const uint16_t *last = arr + size; + const uint16_t *last_aligned = arr + ((size >> 4) << 4); + + if (((uintptr_t)arr & 0xf) == 0 && ((uintptr_t)out & 0x1f) == 0) { + for (; arr != last_aligned; arr += 16, out += 16) { + _mm256_store_ps(out + 0, + _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 0)))); + _mm256_store_ps(out + 8, + _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 8)))); + } + + if (last >= last_aligned + 8) { + _mm256_store_ps(out + 0, + _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 0)))); + arr += 8; + out += 8; + } + } else { + for (; arr != last_aligned; arr += 16, out += 16) { + _mm256_storeu_ps(out + 0, + _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 0)))); + _mm256_storeu_ps(out + 8, + _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 8)))); + } + + if (last >= last_aligned + 8) { + _mm256_storeu_ps(out + 0, + _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 0)))); + arr += 8; + out += 8; + } + } + switch (last - arr) { + case 7: + out[6] = _cvtsh_ss(arr[6]); + /* FALLTHRU */ + case 6: + out[5] = _cvtsh_ss(arr[5]); + /* FALLTHRU */ + case 5: + out[4] = _cvtsh_ss(arr[4]); + /* FALLTHRU */ + case 4: + out[3] = _cvtsh_ss(arr[3]); + /* FALLTHRU */ + case 3: + out[2] = _cvtsh_ss(arr[2]); + /* FALLTHRU */ + case 2: + out[1] = _cvtsh_ss(arr[1]); + /* FALLTHRU */ + case 1: + out[0] = _cvtsh_ss(arr[0]); + } +} + +static inline void convert_fp16_to_fp32_avx(const uint16_t *arr, size_t size, + float norm, float *out) { + const uint16_t *last = arr + size; + const uint16_t *last_aligned = arr + ((size >> 4) << 4); + __m256 ymm_norm = _mm256_set1_ps(norm); + + if (((uintptr_t)arr & 0xf) == 0 && ((uintptr_t)out & 0x1f) == 0) { + for (; arr != last_aligned; arr += 16, out += 16) { + __m256 ymm_0 = _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 0))); + __m256 ymm_1 = _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 8))); + ymm_0 = _mm256_div_ps(ymm_0, ymm_norm); + ymm_1 = _mm256_div_ps(ymm_1, ymm_norm); + _mm256_store_ps(out + 0, ymm_0); + _mm256_store_ps(out + 8, ymm_1); + } + + if (last >= last_aligned + 8) { + _mm256_store_ps( + out, _mm256_div_ps(_mm256_cvtph_ps(_mm_load_si128((__m128i *)arr)), + ymm_norm)); + arr += 8; + out += 8; + } + } else { + for (; arr != last_aligned; arr += 16, out += 16) { + __m256 ymm_0 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 0))); + __m256 ymm_1 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 8))); + ymm_0 = _mm256_div_ps(ymm_0, ymm_norm); + ymm_1 = _mm256_div_ps(ymm_1, ymm_norm); + _mm256_storeu_ps(out + 0, ymm_0); + _mm256_storeu_ps(out + 8, ymm_1); + } + + if (last >= last_aligned + 8) { + _mm256_storeu_ps( + out, _mm256_div_ps(_mm256_cvtph_ps(_mm_loadu_si128((__m128i *)arr)), + ymm_norm)); + arr += 8; + out += 8; + } + } + switch (last - arr) { + case 7: + out[6] = _cvtsh_ss(arr[6]) / norm; + /* FALLTHRU */ + case 6: + out[5] = _cvtsh_ss(arr[5]) / norm; + /* FALLTHRU */ + case 5: + out[4] = _cvtsh_ss(arr[4]) / norm; + /* FALLTHRU */ + case 4: + out[3] = _cvtsh_ss(arr[3]) / norm; + /* FALLTHRU */ + case 3: + out[2] = _cvtsh_ss(arr[2]) / norm; + /* FALLTHRU */ + case 2: + out[1] = _cvtsh_ss(arr[1]) / norm; + /* FALLTHRU */ + case 1: + out[0] = _cvtsh_ss(arr[0]) / norm; + } +} + +static inline void convert_fp32_to_fp16_avx(const float *arr, size_t size, + uint16_t *out) { + const float *last = arr + size; + const float *last_aligned = arr + ((size >> 4) << 4); + + if (((uintptr_t)arr & 0x1f) == 0 && ((uintptr_t)out & 0xf) == 0) { + for (; arr != last_aligned; arr += 16, out += 16) { + _mm_store_si128( + (__m128i *)(out + 0), + _mm256_cvtps_ph(_mm256_load_ps(arr + 0), _MM_FROUND_NO_EXC)); + _mm_store_si128( + (__m128i *)(out + 8), + _mm256_cvtps_ph(_mm256_load_ps(arr + 8), _MM_FROUND_NO_EXC)); + } + + if (last >= last_aligned + 8) { + _mm_store_si128( + (__m128i *)(out + 0), + _mm256_cvtps_ph(_mm256_load_ps(arr + 0), _MM_FROUND_NO_EXC)); + arr += 8; + out += 8; + } + } else { + for (; arr != last_aligned; arr += 16, out += 16) { + _mm_storeu_si128( + (__m128i *)(out + 0), + _mm256_cvtps_ph(_mm256_loadu_ps(arr + 0), _MM_FROUND_NO_EXC)); + _mm_storeu_si128( + (__m128i *)(out + 8), + _mm256_cvtps_ph(_mm256_loadu_ps(arr + 8), _MM_FROUND_NO_EXC)); + } + + if (last >= last_aligned + 8) { + _mm_storeu_si128( + (__m128i *)(out + 0), + _mm256_cvtps_ph(_mm256_loadu_ps(arr + 0), _MM_FROUND_NO_EXC)); + arr += 8; + out += 8; + } + } + switch (last - arr) { + case 7: + out[6] = _cvtss_sh(arr[6], _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 6: + out[5] = _cvtss_sh(arr[5], _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 5: + out[4] = _cvtss_sh(arr[4], _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 4: + out[3] = _cvtss_sh(arr[3], _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 3: + out[2] = _cvtss_sh(arr[2], _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 2: + out[1] = _cvtss_sh(arr[1], _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 1: + out[0] = _cvtss_sh(arr[0], _MM_FROUND_NO_EXC); + } +} + +static inline void convert_fp32_to_fp16_avx(const float *arr, size_t size, + float norm, uint16_t *out) { + const float *last = arr + size; + const float *last_aligned = arr + ((size >> 4) << 4); + __m256 ymm_norm = _mm256_set1_ps(norm); + + if (((uintptr_t)arr & 0x1f) == 0 && ((uintptr_t)out & 0xf) == 0) { + for (; arr != last_aligned; arr += 16, out += 16) { + __m256 ymm_0 = _mm256_load_ps(arr + 0); + __m256 ymm_1 = _mm256_load_ps(arr + 8); + ymm_0 = _mm256_div_ps(ymm_0, ymm_norm); + ymm_1 = _mm256_div_ps(ymm_1, ymm_norm); + _mm_store_si128((__m128i *)(out + 0), + _mm256_cvtps_ph(ymm_0, _MM_FROUND_NO_EXC)); + _mm_store_si128((__m128i *)(out + 8), + _mm256_cvtps_ph(ymm_1, _MM_FROUND_NO_EXC)); + } + + if (last >= last_aligned + 8) { + _mm_store_si128( + (__m128i *)out, + _mm256_cvtps_ph(_mm256_div_ps(_mm256_load_ps(arr), ymm_norm), + _MM_FROUND_NO_EXC)); + arr += 8; + out += 8; + } + } else { + for (; arr != last_aligned; arr += 16, out += 16) { + __m256 ymm_0 = _mm256_loadu_ps(arr + 0); + __m256 ymm_1 = _mm256_loadu_ps(arr + 8); + ymm_0 = _mm256_div_ps(ymm_0, ymm_norm); + ymm_1 = _mm256_div_ps(ymm_1, ymm_norm); + _mm_storeu_si128((__m128i *)(out + 0), + _mm256_cvtps_ph(ymm_0, _MM_FROUND_NO_EXC)); + _mm_storeu_si128((__m128i *)(out + 8), + _mm256_cvtps_ph(ymm_1, _MM_FROUND_NO_EXC)); + } + + if (last >= last_aligned + 8) { + _mm_storeu_si128( + (__m128i *)out, + _mm256_cvtps_ph(_mm256_div_ps(_mm256_loadu_ps(arr), ymm_norm), + _MM_FROUND_NO_EXC)); + arr += 8; + out += 8; + } + } + switch (last - arr) { + case 7: + out[6] = _cvtss_sh(arr[6] / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 6: + out[5] = _cvtss_sh(arr[5] / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 5: + out[4] = _cvtss_sh(arr[4] / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 4: + out[3] = _cvtss_sh(arr[3] / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 3: + out[2] = _cvtss_sh(arr[2] / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 2: + out[1] = _cvtss_sh(arr[1] / norm, _MM_FROUND_NO_EXC); + /* FALLTHRU */ + case 1: + out[0] = _cvtss_sh(arr[0] / norm, _MM_FROUND_NO_EXC); + } +} +#endif // __F16C__ && __AVX__ + +static inline void convert_fp16_to_fp32_fallback(const uint16_t *arr, + size_t size, float *out) { + for (size_t i = 0; i != size; ++i) { + out[i] = float32(arr[i]); + } +} + +static inline void convert_fp16_to_fp32_fallback(const uint16_t *arr, + size_t size, float norm, + float *out) { + for (size_t i = 0; i != size; ++i) { + out[i] = float32(arr[i]) / norm; + } +} + +static inline void convert_fp32_to_fp16_fallback(const float *arr, size_t size, + uint16_t *out) { + for (size_t i = 0; i != size; ++i) { + out[i] = float16(arr[i]); + } +} + +static inline void convert_fp32_to_fp16_fallback(const float *arr, size_t size, + float norm, uint16_t *out) { + for (size_t i = 0; i != size; ++i) { + out[i] = float16(arr[i] / norm); + } +} + +static inline void convert_fp16_to_fp32(const uint16_t *arr, size_t size, + float *out) { +#if defined(__F16C__) && defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return convert_fp16_to_fp32_avx512f(arr, size, out); + } +#endif + +#if defined(__F16C__) && defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + return convert_fp16_to_fp32_avx(arr, size, out); + } +#endif + + return convert_fp16_to_fp32_fallback(arr, size, out); +} + +static inline void convert_fp16_to_fp32(const uint16_t *arr, size_t size, + float norm, float *out) { +#if defined(__F16C__) && defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return convert_fp16_to_fp32_avx512f(arr, size, norm, out); + } +#endif + +#if defined(__F16C__) && defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + return convert_fp16_to_fp32_avx(arr, size, norm, out); + } +#endif + + return convert_fp16_to_fp32_fallback(arr, size, norm, out); +} + +static inline void convert_fp32_to_fp16(const float *arr, size_t size, + uint16_t *out) { +#if defined(__F16C__) && defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return convert_fp32_to_fp16_avx512f(arr, size, out); + } +#endif + +#if defined(__F16C__) && defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + return convert_fp32_to_fp16_avx(arr, size, out); + } +#endif + + return convert_fp32_to_fp16_fallback(arr, size, out); +} + +static inline void convert_fp32_to_fp16(const float *arr, size_t size, + float norm, uint16_t *out) { +#if defined(__F16C__) && defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + return convert_fp32_to_fp16_avx512f(arr, size, norm, out); + } +#endif + +#if defined(__F16C__) && defined(__AVX__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX) { + return convert_fp32_to_fp16_avx(arr, size, norm, out); + } +#endif + + return convert_fp32_to_fp16_fallback(arr, size, norm, out); +} + +#endif // + +namespace zvec { +namespace ailego { + +float FloatHelper::ToFP32(uint16_t val) { + return float32(val); +} + +void FloatHelper::ToFP32(const uint16_t *arr, size_t size, float *out) { + return convert_fp16_to_fp32(arr, size, out); +} + +void FloatHelper::ToFP32(const uint16_t *arr, size_t size, float norm, + float *out) { + return convert_fp16_to_fp32(arr, size, norm, out); +} + +uint16_t FloatHelper::ToFP16(float val) { + return float16(val); +} + +void FloatHelper::ToFP16(const float *arr, size_t size, uint16_t *out) { + return convert_fp32_to_fp16(arr, size, out); +} + +void FloatHelper::ToFP16(const float *arr, size_t size, float norm, + uint16_t *out) { + return convert_fp32_to_fp16(arr, size, norm, out); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/math_helper.h b/src/ailego/utility/math_helper.h new file mode 100644 index 0000000..6c83e24 --- /dev/null +++ b/src/ailego/utility/math_helper.h @@ -0,0 +1,109 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +namespace zvec { +namespace ailego { + +/*! Math Helper + */ +struct MathHelper { + //! Calculate the absolute value + template + static inline auto Absolute(const T &x) -> + typename std::enable_if::value, R>::type { + return static_cast(std::abs(x)); + } + + //! Calculate the absolute value + template + static inline R Absolute(const Float16 &x) { + return static_cast(Float16::Absolute(x)); + } + + //! Calculate the absolute difference + template + static inline auto AbsoluteDifference(const T &x, const T &y) -> + typename std::enable_if::value, R>::type { + auto m = ((x ^ y) & -(x < y)); + auto d = + static_cast::type>((x ^ m) - (y ^ m)); + return static_cast(d); + } + + //! Calculate the absolute difference + template + static inline auto AbsoluteDifference(const T &x, const T &y) -> + typename std::enable_if::value, R>::type { + return static_cast(std::abs(x - y)); + } + + //! Calculate the absolute difference + template + static inline R AbsoluteDifference(const Float16 &x, const Float16 &y) { + return static_cast(std::abs(x - y)); + } + + //! Calculate the squared difference + template + static inline auto SquaredDifference(const T &x, const T &y) -> + typename std::enable_if::value, R>::type { + auto m = ((x ^ y) & -(x < y)); + auto d = + static_cast::type>((x ^ m) - (y ^ m)); + return static_cast(d * d); + } + + //! Calculate the squared difference + template + static inline auto SquaredDifference(const T &x, const T &y) -> + typename std::enable_if::value, R>::type { + auto d = x - y; + return static_cast(d * d); + } + + //! Calculate the squared difference + template + static inline R SquaredDifference(const Float16 &x, const Float16 &y) { + auto d = x - y; + return static_cast(d * d); + } + + //! Test whether two integral numbers are equal + template + static inline auto IsAlmostEqual(const T &x, const T &y, int) -> + typename std::enable_if::value, bool>::type { + return (x == y); + } + + //! Test whether two floating point numbers are equal + template + static inline auto IsAlmostEqual(const T &x, const T &y, int ulp) -> + typename std::enable_if::value, bool>::type { + // the machine epsilon has to be scaled to the magnitude of the values used + // and multiplied by the desired precision in ULPs (units in the last place) + return ((std::fabs(x - y) <= + std::numeric_limits::epsilon() * std::fabs(x + y) * ulp) || + (std::fabs(x - y) < std::numeric_limits::min())); + } +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/matrix_helper.h b/src/ailego/utility/matrix_helper.h new file mode 100644 index 0000000..1a05ad5 --- /dev/null +++ b/src/ailego/utility/matrix_helper.h @@ -0,0 +1,70 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +struct MatrixHelper { + //! Transpose a matrix + template + static inline void Transpose(const void *src, size_t N, void *dst) { + for (size_t i = 0; i < M; ++i) { + for (size_t j = 0; j < N; ++j) { + *(reinterpret_cast(dst) + (j * M + i)) = + *(reinterpret_cast(src) + (i * N + j)); + } + } + } + + //! Reverse transpose a matrix + template + static inline void ReverseTranspose(const void *src, size_t N, void *dst) { + for (size_t i = 0; i < N; ++i) { + for (size_t j = 0; j < M; ++j) { + *(reinterpret_cast(dst) + (j * N + i)) = + *(reinterpret_cast(src) + (i * M + j)); + } + } + } + + //! Transpose a matrix + template + static inline void Transpose(const void *src, size_t M, size_t N, void *dst) { + for (size_t i = 0; i < M; ++i) { + for (size_t j = 0; j < N; ++j) { + *(reinterpret_cast(dst) + (j * M + i)) = + *(reinterpret_cast(src) + (i * N + j)); + } + } + } + + //! Reverse transpose a matrix + template + static inline void ReverseTranspose(const void *src, size_t M, size_t N, + void *dst) { + for (size_t i = 0; i < N; ++i) { + for (size_t j = 0; j < M; ++j) { + *(reinterpret_cast(dst) + (j * N + i)) = + *(reinterpret_cast(src) + (i * M + j)); + } + } + } +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/memory_helper.cc b/src/ailego/utility/memory_helper.cc new file mode 100644 index 0000000..168c38b --- /dev/null +++ b/src/ailego/utility/memory_helper.cc @@ -0,0 +1,492 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "memory_helper.h" +#include +#include +#include +#include +#include +#include + +#if defined(_WIN64) || defined(_WIN32) +#include +#include +#else +#if defined(__linux__) || defined(__linux) +#include +#elif defined(__APPLE__) && defined(__MACH__) +#include +#include +#endif +#include +#include +#endif + +namespace zvec { +namespace ailego { + +#if defined(__linux__) || defined(__linux) +bool MemoryHelper::SelfUsage(size_t *vsz, size_t *rss) { + FILE *fp = fopen("/proc/self/statm", "r"); + if (!fp) { + return false; + } + + if (fscanf(fp, "%zd %zd", vsz, rss) == EOF) { + fclose(fp); + return false; + } + fclose(fp); + + long pagesz = sysconf(_SC_PAGESIZE); + *vsz *= (size_t)pagesz; + *rss *= (size_t)pagesz; + return true; +} + +size_t MemoryHelper::SelfRSS(void) { + FILE *fp = fopen("/proc/self/statm", "r"); + if (!fp) { + return 0; + } + + size_t rss = 0; + if (fscanf(fp, "%*d %zd %*d", &rss) == EOF) { + fclose(fp); + return 0; + } + fclose(fp); + return (rss * sysconf(_SC_PAGESIZE)); +} + +size_t MemoryHelper::SelfPeakRSS(void) { + struct rusage rusage; + getrusage(RUSAGE_SELF, &rusage); + return (size_t)(rusage.ru_maxrss * 1024); +} + +size_t MemoryHelper::TotalRamSize(void) { + return (sysconf(_SC_PHYS_PAGES) * sysconf(_SC_PAGESIZE)); +} + +size_t MemoryHelper::AvailableRamSize(void) { + FILE *fp = fopen("/proc/meminfo", "r"); + if (!fp) { + return 0; + } + + size_t avail = 0; + char buf[128]; + while (fgets(buf, sizeof(buf), fp)) { + if (strncmp(buf, "MemAvailable:", 13) == 0) { + avail = (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + break; + } + } + + // No found 'MemAvailable' + if (avail == 0) { + fseek(fp, 0L, SEEK_SET); + + size_t count = 0; + while (fgets(buf, sizeof(buf), fp)) { + switch (buf[0]) { + case 'M': + if (strncmp(buf, "MemFree:", 8) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + ++count; + } + break; + + case 'B': + if (strncmp(buf, "Buffers:", 8) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + ++count; + } + break; + + case 'C': + if (strncmp(buf, "Cached:", 7) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + ++count; + } + break; + } + // All read + if (count == 3) { + break; + } + } + } + fclose(fp); + return (avail * 1024); +} + +size_t MemoryHelper::UsedRamSize(void) { + FILE *fp = fopen("/proc/meminfo", "r"); + if (!fp) { + return 0; + } + + size_t total = 0, avail = 0, count = 0; + char buf[128]; + + while (fgets(buf, sizeof(buf), fp)) { + switch (buf[0]) { + case 'M': + if (strncmp(buf, "MemTotal:", 9) == 0) { + total = (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + ++count; + } else if (strncmp(buf, "MemFree:", 8) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + ++count; + } + break; + + case 'B': + if (strncmp(buf, "Buffers:", 8) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + ++count; + } + break; + + case 'C': + if (strncmp(buf, "Cached:", 7) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + ++count; + } + break; + + case 'S': + if (strncmp(buf, "Slab:", 5) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, nullptr, 10); + ++count; + } + break; + } + // All read + if (count == 5) { + break; + } + } + fclose(fp); + + if (total == 0) { + total = (sysconf(_SC_PHYS_PAGES) * sysconf(_SC_PAGESIZE)) / 1024; + } + return ((total - avail) * 1024); +} + +size_t MemoryHelper::ContainerAwareTotalRamSize(void) { + size_t total_ram_size = TotalRamSize(); + std::string limit_in_bytes = "/sys/fs/cgroup/memory/memory.limit_in_bytes"; + if (FileHelper::IsExist(limit_in_bytes.c_str())) { + std::ifstream memory_limit_ifs; + std::string memory_limit_str{""}; + memory_limit_ifs.open(limit_in_bytes, std::ios::in); + if (memory_limit_ifs.is_open()) { + uint64_t limit = 0; + memory_limit_ifs >> memory_limit_str; + if (memory_limit_str != "-1") { + // Refer to: + // https://access.redhat.com/documentation/zh-cn/red_hat_enterprise_linux/7/html/resource_management_guide/sec-memory + StringHelper::ToUint64(memory_limit_str, &limit); + if (limit != 0x7FFFFFFFFFFFF000) { + // Refer to: + // https://stackoverflow.com/questions/70332396/why-cgroups-file-memory-limit-in-bytes-use-9223372036854771712-as-a-default-valu + total_ram_size = static_cast(limit); + } + } + memory_limit_ifs.close(); + } + } + return total_ram_size; +} + +#elif defined(__APPLE__) && defined(__MACH__) +bool MemoryHelper::SelfUsage(size_t *vsz, size_t *rss) { + struct mach_task_basic_info info; + mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT; + + if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&info, + &count) != KERN_SUCCESS) { + return false; + } + *vsz = info.virtual_size; + *rss = info.resident_size; + return true; +} + +size_t MemoryHelper::SelfRSS(void) { + struct mach_task_basic_info info; + mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT; + + if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&info, + &count) != KERN_SUCCESS) { + return 0; + } + return info.resident_size; +} + +size_t MemoryHelper::SelfPeakRSS(void) { + struct mach_task_basic_info info; + mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT; + + if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&info, + &count) != KERN_SUCCESS) { + return 0; + } + return info.resident_size_max; +} + +size_t MemoryHelper::TotalRamSize(void) { + int mib[2] = {CTL_HW, HW_MEMSIZE}; + uint64_t size = 0; + size_t len = sizeof(size); + if (sysctl(mib, 2, &size, &len, nullptr, 0) != 0) { + return 0; + } + return (size_t)size; +} + +size_t MemoryHelper::AvailableRamSize(void) { + struct vm_statistics stat; + mach_msg_type_number_t count = HOST_VM_INFO_COUNT; + vm_size_t pagesize = 0; + + if (host_page_size(mach_host_self(), &pagesize) != KERN_SUCCESS) { + return 0; + } + if (host_statistics(mach_host_self(), HOST_VM_INFO, (host_info_t)&stat, + &count) != KERN_SUCCESS) { + return 0; + } + return ((stat.free_count + stat.inactive_count) * pagesize); +} + +size_t MemoryHelper::UsedRamSize(void) { + struct vm_statistics stat; + mach_msg_type_number_t count = HOST_VM_INFO_COUNT; + vm_size_t pagesize = 0; + + if (host_page_size(mach_host_self(), &pagesize) != KERN_SUCCESS) { + return 0; + } + if (host_statistics(mach_host_self(), HOST_VM_INFO, (host_info_t)&stat, + &count) != KERN_SUCCESS) { + return 0; + } + return ((stat.active_count + stat.wire_count) * pagesize); +} + +size_t MemoryHelper::ContainerAwareTotalRamSize(void) { + return 0u; +} + +#elif defined(_WIN64) || defined(_WIN32) +static inline int getpagesize(void) { + SYSTEM_INFO info; + GetSystemInfo(&info); + return info.dwPageSize; +} + +bool MemoryHelper::SelfUsage(size_t *vsz, size_t *rss) { + PROCESS_MEMORY_COUNTERS info; + if (!GetProcessMemoryInfo(GetCurrentProcess(), &info, sizeof(info))) { + return false; + } + *vsz = (size_t)info.PagefileUsage; + *rss = (size_t)info.WorkingSetSize; + return true; +} + +size_t MemoryHelper::SelfRSS(void) { + PROCESS_MEMORY_COUNTERS info; + if (!GetProcessMemoryInfo(GetCurrentProcess(), &info, sizeof(info))) { + return 0u; + } + return (size_t)info.WorkingSetSize; +} + +size_t MemoryHelper::SelfPeakRSS(void) { + PROCESS_MEMORY_COUNTERS info; + GetProcessMemoryInfo(GetCurrentProcess(), &info, sizeof(info)); + return (size_t)info.PeakWorkingSetSize; +} + +size_t MemoryHelper::TotalRamSize(void) { + MEMORYSTATUSEX status; + status.dwLength = sizeof(status); + GlobalMemoryStatusEx(&status); + return (size_t)status.ullTotalPhys; +} + +size_t MemoryHelper::AvailableRamSize(void) { + MEMORYSTATUSEX status; + status.dwLength = sizeof(status); + GlobalMemoryStatusEx(&status); + return (size_t)status.ullAvailPhys; +} + +size_t MemoryHelper::UsedRamSize(void) { + MEMORYSTATUSEX status; + status.dwLength = sizeof(status); + GlobalMemoryStatusEx(&status); + return (size_t)(status.ullTotalPhys - status.ullAvailPhys); +} + +size_t MemoryHelper::ContainerAwareTotalRamSize(void) { + return 0u; +} + +#else +bool MemoryHelper::SelfUsage(size_t *vsz, size_t *rss) { + *vsz = 0u; + *rss = 0u; + return false; +} + +size_t MemoryHelper::SelfRSS(void) { + return 0u; +} + +size_t MemoryHelper::SelfPeakRSS(void) { + return 0u; +} + +size_t MemoryHelper::TotalRamSize(void) { + return 0u; +} + +size_t MemoryHelper::AvailableRamSize(void) { + return 0u; +} + +size_t MemoryHelper::UsedRamSize(void) { + return 0u; +} + +size_t MemoryHelper::ContainerAwareTotalRamSize(void) { + return 0u; +} +#endif + +size_t MemoryHelper::PageSize(void) { + static size_t page_size = static_cast(getpagesize()); + return page_size; +} + +size_t MemoryHelper::HugePageSize(void) { + static size_t page_size = static_cast(2 * 1024 * 1024); + return page_size; +} + +size_t MemoryHelper::AlignHugePageSize(size_t size) { + const size_t page_mask = HugePageSize() - 1; + return (size + page_mask) & (~page_mask); +} + +void *MemoryHelper::AllocateHugePage(size_t size, bool zero_fill) { + if (size == 0) { + return nullptr; + } + const size_t aligned_size = AlignHugePageSize(size); + +#if defined(_WIN64) || defined(_WIN32) + void *ptr = ::_aligned_malloc(aligned_size, PageSize()); + if (ptr == nullptr) { + return nullptr; + } + if (zero_fill) { + std::memset(ptr, 0, aligned_size); + } + return ptr; +#else + void *ptr = ::mmap(nullptr, aligned_size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (ptr == MAP_FAILED) { + return nullptr; + } + // MADV_HUGEPAGE is a Linux-only hint for transparent huge pages. On + // macOS/BSD (which manage superpages differently) it is intentionally + // absent; skipping it only forgoes a performance hint, not correctness. +#if defined(MADV_HUGEPAGE) + ::madvise(ptr, aligned_size, MADV_HUGEPAGE); +#endif + // mmap with MAP_ANONYMOUS already returns zero-filled pages, so an explicit + // memset is only needed when the caller relies on it for a non-anonymous + // fallback; here it is redundant and skipped to avoid touching every page. + (void)zero_fill; + return ptr; +#endif +} + +void MemoryHelper::FreeHugePage(void *ptr, size_t size) { + if (ptr == nullptr) { + return; + } +#if defined(_WIN64) || defined(_WIN32) + (void)size; + ::_aligned_free(ptr); +#else + ::munmap(ptr, AlignHugePageSize(size)); +#endif +} + +void *MemoryHelper::AllocateAligned(size_t size, size_t alignment, + bool zero_fill) { + assert(alignment != 0 && (alignment & (alignment - 1)) == 0 && + "alignment must be a power of two"); + if (size == 0) { + return nullptr; + } + if (size >= HugePageSize()) { + return AllocateHugePage(size, zero_fill); + } + + // Small block: a regular aligned allocation avoids reserving a whole huge + // page. std::aligned_alloc requires the size to be a multiple of alignment. + const size_t aligned_size = (size + alignment - 1) / alignment * alignment; +#if defined(_WIN64) || defined(_WIN32) + void *ptr = ::_aligned_malloc(aligned_size, alignment); +#else + void *ptr = std::aligned_alloc(alignment, aligned_size); +#endif + if (ptr == nullptr) { + return nullptr; + } + if (zero_fill) { + std::memset(ptr, 0, aligned_size); + } + return ptr; +} + +void MemoryHelper::FreeAligned(void *ptr, size_t size) { + if (ptr == nullptr) { + return; + } + if (size >= HugePageSize()) { + FreeHugePage(ptr, size); + return; + } +#if defined(_WIN64) || defined(_WIN32) + ::_aligned_free(ptr); +#else + std::free(ptr); +#endif +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/memory_helper.h b/src/ailego/utility/memory_helper.h new file mode 100644 index 0000000..ee68629 --- /dev/null +++ b/src/ailego/utility/memory_helper.h @@ -0,0 +1,113 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +/*! Memory Helper + */ +struct MemoryHelper { + //! Retrieve the page size of memory + static size_t PageSize(void); + + //! Retrieve the huge page size of memory + static size_t HugePageSize(void); + + //! Round `size` up to a multiple of the huge page size. + static size_t AlignHugePageSize(size_t size); + + //! Allocate a large, page-aligned block that prefers transparent huge pages. + //! + //! On Linux the block is obtained via anonymous mmap and hinted with + //! MADV_HUGEPAGE; on other platforms it falls back to a page-aligned + //! allocation without the huge-page hint (which is a performance hint, not a + //! correctness requirement). Returns nullptr on failure. + //! + //! `size` is rounded up to the huge page size internally, and the same + //! rounded value is what the corresponding FreeHugePage call expects, so + //! callers should treat the returned block as exactly AlignHugePageSize(size) + //! bytes. + //! + //! `zero_fill` requests zeroed memory: when true the returned block is + //! guaranteed to be zero-initialized. When false the caller does not require + //! zeroing, but the implementation is still free to return zeroed memory and + //! does so on the anonymous-mmap path (MAP_ANONYMOUS pages are always zero), + //! where an explicit memset is skipped to preserve lazy paging. In other + //! words, true => always zeroed; false => zeroing is not guaranteed either + //! way. Never assume non-zero contents. + //! + //! Blocks returned here MUST be released with FreeHugePage (never free()), + //! because the underlying allocator differs per platform. + static void *AllocateHugePage(size_t size, bool zero_fill = true); + + //! Release a block previously returned by AllocateHugePage. + //! + //! `size` must be the same value originally passed to AllocateHugePage; it is + //! required because the Linux mmap path needs the length for munmap. + static void FreeHugePage(void *ptr, size_t size); + + //! Allocate an aligned block, choosing the backing allocator by size. + //! + //! When `size` is at least the huge page size, the block is obtained via + //! AllocateHugePage (huge-page-backed, page-aligned). Otherwise a regular + //! `alignment`-aligned allocation is used, which avoids wasting a full huge + //! page on small buffers. Returns nullptr on failure. + //! + //! `alignment` must be a power of two. + //! + //! `zero_fill` follows the same contract as AllocateHugePage: true guarantees + //! zeroed memory; false does not require zeroing but the implementation may + //! still return zeroed memory (it does on the huge-page mmap path). Never + //! assume non-zero contents. + //! + //! Blocks returned here MUST be released with FreeAligned, passing the same + //! `size`, because the chosen allocator (and therefore the matching free) is + //! derived from `size`. + static void *AllocateAligned(size_t size, size_t alignment = 64, + bool zero_fill = true); + + //! Release a block previously returned by AllocateAligned. + //! + //! `size` must be the same value originally passed to AllocateAligned so the + //! same allocator path is selected for releasing the block. + static void FreeAligned(void *ptr, size_t size); + + //! Retrieve the VSZ and RSS of self process in bytes + static bool SelfUsage(size_t *vsz, size_t *rss); + + //! Retrieve the RSS of self process in bytes + static size_t SelfRSS(void); + + //! Retrieve the peak RSS of self process in bytes + static size_t SelfPeakRSS(void); + + //! Retrieve the total size of physical memory (RAM) in bytes + static size_t TotalRamSize(void); + + //! Retrieve the available size of physical memory (RAM) in bytes + static size_t AvailableRamSize(void); + + //! Retrieve the used size of physical memory (RAM) in bytes + static size_t UsedRamSize(void); + + //! Retrieve the total size of physical memory (RAM) in bytes in container + static size_t ContainerAwareTotalRamSize(void); +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/string_helper.cc b/src/ailego/utility/string_helper.cc new file mode 100644 index 0000000..74d3aff --- /dev/null +++ b/src/ailego/utility/string_helper.cc @@ -0,0 +1,121 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +bool StringHelper::StartsWith(const std::string &ref, + const std::string &prefix) { + return (ref.size() >= prefix.size()) && + (ref.compare(0, prefix.size(), prefix) == 0); +} + +bool StringHelper::EndsWith(const std::string &ref, const std::string &suffix) { + size_t s1 = ref.size(); + size_t s2 = suffix.size(); + return (s1 >= s2) && (ref.compare(s1 - s2, s2, suffix) == 0); +} + +void StringHelper::LeftTrim(std::string &str) { + str.erase(str.begin(), std::find_if(str.begin(), str.end(), [](int ch) { + return !std::isspace(ch); + })); +} + +void StringHelper::RightTrim(std::string &str) { + str.erase(std::find_if(str.rbegin(), str.rend(), + [](int ch) { return !std::isspace(ch); }) + .base(), + str.end()); +} + +void StringHelper::Trim(std::string &str) { + StringHelper::RightTrim(str); + StringHelper::LeftTrim(str); +} + +std::string StringHelper::CopyLeftTrim(std::string str) { + StringHelper::LeftTrim(str); + return str; +} + +std::string StringHelper::CopyRightTrim(std::string str) { + StringHelper::RightTrim(str); + return str; +} + +std::string StringHelper::CopyTrim(std::string str) { + StringHelper::Trim(str); + return str; +} + +bool StringHelper::CompareIgnoreCase(const std::string &a, + const std::string &b) { + if (a.size() != b.size()) { + return false; + } + return (strncasecmp(a.data(), b.data(), a.size()) == 0); +} + +void StringHelper::Append(std::string *str, const internal::Alphameric &a) { + str->reserve(str->size() + a.size()); + str->append(a.data(), a.size()); +} + +void StringHelper::Append(std::string *str, const internal::Alphameric &a, + const internal::Alphameric &b) { + str->reserve(str->size() + a.size() + b.size()); + str->append(a.data(), a.size()); + str->append(b.data(), b.size()); +} + +void StringHelper::Append(std::string *str, const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c) { + str->reserve(str->size() + a.size() + b.size() + c.size()); + str->append(a.data(), a.size()); + str->append(b.data(), b.size()); + str->append(c.data(), c.size()); +} + +void StringHelper::Append(std::string *str, const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c, + const internal::Alphameric &d) { + str->reserve(str->size() + a.size() + b.size() + c.size() + d.size()); + str->append(a.data(), a.size()); + str->append(b.data(), b.size()); + str->append(c.data(), c.size()); + str->append(d.data(), d.size()); +} + +void StringHelper::AppendViews(std::string *str, + std::initializer_list views) { + size_t new_size = str->size(); + for (auto &v : views) { + new_size += v.size(); + } + str->reserve(new_size); + for (auto &v : views) { + str->append(v.data(), v.size()); + } +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/utility/time_helper.cc b/src/ailego/utility/time_helper.cc new file mode 100644 index 0000000..7f0231e --- /dev/null +++ b/src/ailego/utility/time_helper.cc @@ -0,0 +1,215 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#if defined(_WIN64) || defined(_WIN32) +#include +#endif + +namespace zvec { +namespace ailego { + +#if defined(_WIN64) || defined(_WIN32) +uint64_t Monotime::NanoSeconds(void) { + LARGE_INTEGER stamp, freq; + QueryPerformanceFrequency(&freq); + QueryPerformanceCounter(&stamp); + return (uint64_t)((double)stamp.QuadPart * + (1000000000.0 / (double)freq.QuadPart)); +} + +uint64_t Monotime::MicroSeconds(void) { + LARGE_INTEGER stamp, freq; + QueryPerformanceFrequency(&freq); + QueryPerformanceCounter(&stamp); + return (stamp.QuadPart * 1000000u / freq.QuadPart); +} + +uint64_t Monotime::MilliSeconds(void) { + LARGE_INTEGER stamp, freq; + QueryPerformanceFrequency(&freq); + QueryPerformanceCounter(&stamp); + return (stamp.QuadPart * 1000u / freq.QuadPart); +} + +uint64_t Monotime::Seconds(void) { + LARGE_INTEGER stamp, freq; + QueryPerformanceFrequency(&freq); + QueryPerformanceCounter(&stamp); + return (stamp.QuadPart / freq.QuadPart); +} + +// January 1, 1970 (start of Unix epoch) in "ticks" +#define UNIX_TIME_START 0x019DB1DED53E8000ull + +uint64_t Realtime::NanoSeconds(void) { + LARGE_INTEGER stamp; + FILETIME file; + GetSystemTimeAsFileTime(&file); + stamp.HighPart = file.dwHighDateTime; + stamp.LowPart = file.dwLowDateTime; + return (stamp.QuadPart - UNIX_TIME_START) * 100u; +} + +uint64_t Realtime::MicroSeconds(void) { + LARGE_INTEGER stamp; + FILETIME file; + GetSystemTimeAsFileTime(&file); + stamp.HighPart = file.dwHighDateTime; + stamp.LowPart = file.dwLowDateTime; + return (stamp.QuadPart - UNIX_TIME_START) / 10u; +} + +uint64_t Realtime::MilliSeconds(void) { + LARGE_INTEGER stamp; + FILETIME file; + GetSystemTimeAsFileTime(&file); + stamp.HighPart = file.dwHighDateTime; + stamp.LowPart = file.dwLowDateTime; + return (stamp.QuadPart - UNIX_TIME_START) / 10000u; +} + +uint64_t Realtime::Seconds(void) { + LARGE_INTEGER stamp; + FILETIME file; + GetSystemTimeAsFileTime(&file); + stamp.HighPart = file.dwHighDateTime; + stamp.LowPart = file.dwLowDateTime; + return (stamp.QuadPart - UNIX_TIME_START) / 10000000u; +} + +size_t Realtime::Localtime(uint64_t stamp, const char *format, char *buf, + size_t len) { + time_t val = static_cast(stamp); + return strftime(buf, len, format, localtime(&val)); +} + +size_t Realtime::Gmtime(uint64_t stamp, const char *format, char *buf, + size_t len) { + time_t val = static_cast(stamp); + return strftime(buf, len, format, gmtime(&val)); +} + +size_t Realtime::Localtime(const char *format, char *buf, size_t len) { + time_t now = time(0); + return strftime(buf, len, format, localtime(&now)); +} + +size_t Realtime::Gmtime(const char *format, char *buf, size_t len) { + time_t now = time(0); + return strftime(buf, len, format, gmtime(&now)); +} +#else +uint64_t Monotime::NanoSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_MONOTONIC, &tspec); + return (tspec.tv_sec * 1000000000u + tspec.tv_nsec); +} + +uint64_t Monotime::MicroSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_MONOTONIC, &tspec); + return (tspec.tv_sec * 1000000u + tspec.tv_nsec / 1000u); +} + +uint64_t Monotime::MilliSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_MONOTONIC, &tspec); + return (tspec.tv_sec * 1000u + tspec.tv_nsec / 1000000u); +} + +uint64_t Monotime::Seconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_MONOTONIC, &tspec); + return (tspec.tv_sec); +} + +uint64_t Realtime::NanoSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_REALTIME, &tspec); + return (tspec.tv_sec * 1000000000u + tspec.tv_nsec); +} + +uint64_t Realtime::MicroSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_REALTIME, &tspec); + return (tspec.tv_sec * 1000000u + tspec.tv_nsec / 1000u); +} + +uint64_t Realtime::MilliSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_REALTIME, &tspec); + return (tspec.tv_sec * 1000u + tspec.tv_nsec / 1000000u); +} + +uint64_t Realtime::Seconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_REALTIME, &tspec); + return (tspec.tv_sec); +} + +size_t Realtime::Localtime(uint64_t stamp, const char *format, char *buf, + size_t len) { + struct tm tmbuf; + time_t val = static_cast(stamp); + return strftime(buf, len, format, localtime_r(&val, &tmbuf)); +} + +size_t Realtime::Gmtime(uint64_t stamp, const char *format, char *buf, + size_t len) { + struct tm tmbuf; + time_t val = static_cast(stamp); + return strftime(buf, len, format, gmtime_r(&val, &tmbuf)); +} + +size_t Realtime::Localtime(const char *format, char *buf, size_t len) { + struct tm tmbuf; + time_t now = time(nullptr); + return strftime(buf, len, format, localtime_r(&now, &tmbuf)); +} + +size_t Realtime::Gmtime(const char *format, char *buf, size_t len) { + struct tm tmbuf; + time_t now = time(nullptr); + return strftime(buf, len, format, gmtime_r(&now, &tmbuf)); +} + +uint64_t CPUtime::NanoSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &tspec); + return (tspec.tv_sec * 1000000000u + tspec.tv_nsec); +} + +uint64_t CPUtime::MicroSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &tspec); + return (tspec.tv_sec * 1000000u + tspec.tv_nsec / 1000u); +} + +uint64_t CPUtime::MilliSeconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &tspec); + return (tspec.tv_sec * 1000u + tspec.tv_nsec / 1000000u); +} + +uint64_t CPUtime::Seconds(void) { + struct timespec tspec; + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &tspec); + return (tspec.tv_sec); +} +#endif // _WIN64 || _WIN32 + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/version.cc b/src/ailego/version.cc new file mode 100644 index 0000000..ac457a3 --- /dev/null +++ b/src/ailego/version.cc @@ -0,0 +1,48 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "version.h" +#include "version.i" + +#ifdef ailego_VERSION +#define AILEGO_VERSION_STRING ailego_VERSION +#else +#define AILEGO_VERSION_STRING "unknown" +#endif + +namespace zvec { + +namespace ailego { + +static const char AILEGO_VERSION_DETAILS[] = AILEGO_VERSION_COMPILE_DETAILS( + "AiLego Library Version " AILEGO_VERSION_STRING + ".\nCopyright (C) The Software Authors. All rights reserved.\n"); + +const char *Version::String(void) { + return AILEGO_VERSION_STRING; +} + +const char *Version::Details(void) { + return AILEGO_VERSION_DETAILS; +} + +} // namespace ailego +} // namespace zvec + +// extern "C" int __wrap_main(int, char *[]) { +// fwrite(ailego::AILEGO_VERSION_DETAILS, 1, +// strlen(ailego::AILEGO_VERSION_DETAILS), stdout); +// fflush(stdout); +// _Exit(0); +// } diff --git a/src/ailego/version.h b/src/ailego/version.h new file mode 100644 index 0000000..bf7b4ae --- /dev/null +++ b/src/ailego/version.h @@ -0,0 +1,33 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +namespace zvec { + +namespace ailego { + +/*! AiLego Version + */ +struct Version { + //! Retrieve the version number in string + static const char *String(void); + + //! Retrieve the detailed version information + static const char *Details(void); +}; + +} // namespace ailego + +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/version.i b/src/ailego/version.i new file mode 100644 index 0000000..c1b14be --- /dev/null +++ b/src/ailego/version.i @@ -0,0 +1,293 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#ifndef AILEGO_VERSION_TO_STRING_ +#define AILEGO_VERSION_TO_STRING_(x) #x +#endif + +#ifndef AILEGO_VERSION_TO_STRING +#define AILEGO_VERSION_TO_STRING(x) AILEGO_VERSION_TO_STRING_(x) +#endif + +/*! http://nadeausoftware.com/articles/2012/01/ + * c_c_tip_how_use_compiler_predefined_macros_detect_operating_system + */ +#if defined(__linux) || defined(__linux__) +#define AILEGO_VERSION_PLATFORM "Linux" +#elif defined(__FreeBSD__) +#define AILEGO_VERSION_PLATFORM "FreeBSD" +#elif defined(__NetBSD__) +#define AILEGO_VERSION_PLATFORM "NetBSD" +#elif defined(__OpenBSD__) +#define AILEGO_VERSION_PLATFORM "OpenBSD" +#elif defined(__APPLE__) || defined(__MACH__) +#define AILEGO_VERSION_PLATFORM "Darwin" +#elif defined(__CYGWIN__) && !defined(_WIN32) +#define AILEGO_VERSION_PLATFORM "Cygwin" +#elif defined(_WIN64) +#define AILEGO_VERSION_PLATFORM "Microsoft Windows (64-bit)" +#elif defined(_WIN32) +#define AILEGO_VERSION_PLATFORM "Microsoft Windows (32-bit)" +#elif defined(__sun) && defined(__SVR4) +#define AILEGO_VERSION_PLATFORM "Solaris" +#elif defined(_AIX) +#define AILEGO_VERSION_PLATFORM "AIX" +#elif defined(__hpux) +#define AILEGO_VERSION_PLATFORM "HP-UX" +#elif defined(__unix) || defined(__unix__) +#define AILEGO_VERSION_PLATFORM "Unix" +#else +#define AILEGO_VERSION_PLATFORM "Unknown Platform" +#endif + +/*! http://nadeausoftware.com/articles/2012/10/ + * c_c_tip_how_detect_compiler_name_and_version_using_compiler_predefined_macros + */ +#if defined(__NVCC__) +#define AILEGO_VERSION_COMPILER_NAME "Nvidia CUDA Compiler" +#elif defined(__clang__) +#define AILEGO_VERSION_COMPILER_NAME "Clang/LLVM" +#elif defined(__ICC) || defined(__INTEL_COMPILER) +#define AILEGO_VERSION_COMPILER_NAME "Intel ICC/ICPC" +#elif defined(__GNUC__) || defined(__GNUG__) +#define AILEGO_VERSION_COMPILER_NAME "GNU GCC/G++" +#elif defined(__HP_cc) || defined(__HP_aCC) +#define AILEGO_VERSION_COMPILER_NAME "Hewlett-Packard C/aC++" +#elif defined(__IBMC__) || defined(__IBMCPP__) +#define AILEGO_VERSION_COMPILER_NAME "IBM XL C/C++" +#elif defined(_MSC_VER) +#define AILEGO_VERSION_COMPILER_NAME "Microsoft Visual C++" +#elif defined(__PGI) +#define AILEGO_VERSION_COMPILER_NAME "Portland Group PGCC/PGCPP" +#elif defined(__SUNPRO_C) || defined(__SUNPRO_CC) +#define AILEGO_VERSION_COMPILER_NAME "Oracle Solaris Studio" +#else +#define AILEGO_VERSION_COMPILER_NAME "Unknown Compiler" +#endif + +#if defined(__CUDACC_VER_MAJOR__) +#define AILEGO_VERSION_COMPILER \ + AILEGO_VERSION_COMPILER_NAME \ + " (" AILEGO_VERSION_TO_STRING(__CUDACC_VER_MAJOR__) \ + "." AILEGO_VERSION_TO_STRING(__CUDACC_VER_MINOR__) \ + "." AILEGO_VERSION_TO_STRING(__CUDACC_VER_BUILD__) ")" +#elif defined(__VERSION__) +#define AILEGO_VERSION_COMPILER \ + AILEGO_VERSION_COMPILER_NAME " (" __VERSION__ ")" +#elif defined(_MSC_FULL_VER) +#define AILEGO_VERSION_COMPILER \ + AILEGO_VERSION_COMPILER_NAME " (" AILEGO_VERSION_TO_STRING(_MSC_FULL_VER) ")" +#elif defined(_MSC_VER) +#define AILEGO_VERSION_COMPILER \ + AILEGO_VERSION_COMPILER_NAME " (" AILEGO_VERSION_TO_STRING(_MSC_VER) ")" +#elif defined(__PGIC__) +#define AILEGO_VERSION_COMPILER \ + AILEGO_VERSION_COMPILER_NAME \ + " (" AILEGO_VERSION_TO_STRING(__PGIC__) "." AILEGO_VERSION_TO_STRING( \ + __PGIC_MINOR__) "." AILEGO_VERSION_TO_STRING(__PGIC_PATCHLEVEL__) ")" +#elif defined(__xlc__) +#define AILEGO_VERSION_COMPILER AILEGO_VERSION_COMPILER_NAME " (" __xlc__ ")" +#elif defined(__SUNPRO_C) +#define AILEGO_VERSION_COMPILER \ + AILEGO_VERSION_COMPILER_NAME " (" AILEGO_VERSION_TO_STRING(__SUNPRO_C) ")" +#elif defined(__HP_cc) +#define AILEGO_VERSION_COMPILER \ + AILEGO_VERSION_COMPILER_NAME " (" AILEGO_VERSION_TO_STRING(__HP_cc) ")" +#else +#define AILEGO_VERSION_COMPILER AILEGO_VERSION_COMPILER_NAME +#endif + +#if defined(__x86_64__) || defined(_M_X64) +#define AILEGO_VERSION_PROCESSOR "x86 64-bit Processor" +#elif defined(__i386) || defined(_M_IX86) +#define AILEGO_VERSION_PROCESSOR "x86 32-bit Processor" +#elif defined(__ARM_ARCH) +#if defined(__ARM_64BIT_STATE) +#define AILEGO_VERSION_PROCESSOR "ARM 64-bit Processor" +#else +#define AILEGO_VERSION_PROCESSOR "ARM 32-bit Processor" +#endif +#elif defined(__ia64) || defined(__itanium__) || defined(_M_IA64) +#define AILEGO_VERSION_PROCESSOR "Itanium Processor" +#elif defined(__powerpc64__) || defined(__ppc64__) || defined(__PPC64__) +#define AILEGO_VERSION_PROCESSOR "PowerPC 64-bit Processor" +#elif defined(__powerpc__) || defined(__ppc__) || defined(__PPC__) +#define AILEGO_VERSION_PROCESSOR "PowerPC 32-bit Processor" +#elif defined(__sparc) +#define AILEGO_VERSION_PROCESSOR "SPARC Processor" +#else +#define AILEGO_VERSION_PROCESSOR "Unknown Processor" +#endif + +#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ +#define AILEGO_VERSION_BYTE_ORDER " Little-endian Byte Order\n" +#elif __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ +#define AILEGO_VERSION_BYTE_ORDER " Big-endian Byte Order\n" +#elif __BYTE_ORDER__ == __ORDER_PDP_ENDIAN__ +#define AILEGO_VERSION_BYTE_ORDER " PDP-endian Byte Order\n" +#else +#define AILEGO_VERSION_BYTE_ORDER "" +#endif + +#if defined(_DEBUG) || (!defined(__OPTIMIZE__) && !defined(NDEBUG)) +#define AILEGO_VERSION_DEBUG_INFO " Debug Information\n" +#else +#define AILEGO_VERSION_DEBUG_INFO "" +#endif + +#if defined(__SANITIZE_ADDRESS__) +#define AILEGO_VERSION_ASAN " Address Sanitizer\n" +#else +#define AILEGO_VERSION_ASAN "" +#endif + +#if defined(__STDC_VERSION__) +#define AILEGO_VERSION_STDC \ + " C Standard " AILEGO_VERSION_TO_STRING(__STDC_VERSION__) "\n" +#else +#define AILEGO_VERSION_STDC "" +#endif + +#if defined(__cplusplus) +#define AILEGO_VERSION_CPLUSPLUS \ + " C++ Standard " AILEGO_VERSION_TO_STRING(__cplusplus) "\n" +#else +#define AILEGO_VERSION_CPLUSPLUS "" +#endif + +#if defined(__GXX_ABI_VERSION) +#define AILEGO_VERSION_GXX_ABI \ + " GNU C++ ABI " AILEGO_VERSION_TO_STRING(__GXX_ABI_VERSION) "\n" +#else +#define AILEGO_VERSION_GXX_ABI "" +#endif + +#if defined(__GLIBC__) +#define AILEGO_VERSION_GLIBC \ + " GNU glibc " AILEGO_VERSION_TO_STRING( \ + __GLIBC__) "." AILEGO_VERSION_TO_STRING(__GLIBC_MINOR__) "\n" +#else +#define AILEGO_VERSION_GLIBC "" +#endif + +#if defined(WINVER) +#define AILEGO_VERSION_WINSDK \ + " Microsoft Windows SDK " AILEGO_VERSION_TO_STRING(WINVER) "\n" +#else +#define AILEGO_VERSION_WINSDK "" +#endif + +#if defined(__CLR_VER) +#define AILEGO_VERSION_CLR \ + " Microsoft CLR " AILEGO_VERSION_TO_STRING(__CLR_VER) "\n" +#else +#define AILEGO_VERSION_CLR "" +#endif + +#if defined(__LSB_VERSION__) +#define AILEGO_VERSION_LSB \ + " Linux Standards Base " AILEGO_VERSION_TO_STRING(__LSB_VERSION__) "\n" +#else +#define AILEGO_VERSION_LSB "" +#endif + +#if defined(_POSIX_VERSION) +#define AILEGO_VERSION_POSIX \ + " POSIX Specification " AILEGO_VERSION_TO_STRING(_POSIX_VERSION) "\n" +#else +#define AILEGO_VERSION_POSIX "" +#endif + +#if defined(_XOPEN_VERSION) +#define AILEGO_VERSION_XOPEN \ + " X/Open Specification " AILEGO_VERSION_TO_STRING(_XOPEN_VERSION) "\n" +#else +#define AILEGO_VERSION_XOPEN "" +#endif + +#if defined(_OPENMP) +#define AILEGO_VERSION_OPENMP \ + " OpenMP API " AILEGO_VERSION_TO_STRING(_OPENMP) "\n" +#else +#define AILEGO_VERSION_OPENMP "" +#endif + +#if defined(__ARM_NEON) +#define AILEGO_VERSION_SIMD " Arm Neon Instruction Set\n" +#elif defined(__AVX512FP16__) +#define AILEGO_VERSION_SIMD " AVX-512FP16 Instruction Set\n" +#elif defined(__AVX512F__) +#define AILEGO_VERSION_SIMD " AVX-512F Instruction Set\n" +#elif defined(__AVX2__) +#define AILEGO_VERSION_SIMD " AVX-2 Instruction Set\n" +#elif defined(__AVX__) +#define AILEGO_VERSION_SIMD " AVX Instruction Set\n" +#elif defined(__SSE4_2__) +#define AILEGO_VERSION_SIMD " SSE-4.2 Instruction Set\n" +#elif defined(__SSE4_1__) +#define AILEGO_VERSION_SIMD " SSE-4.1 Instruction Set\n" +#elif defined(__SSSE3__) +#define AILEGO_VERSION_SIMD " SSSE-3 Instruction Set\n" +#elif defined(__SSE3__) +#define AILEGO_VERSION_SIMD " SSE-3 Instruction Set\n" +#elif defined(__SSE2__) +#define AILEGO_VERSION_SIMD " SSE-2 Instruction Set\n" +#elif defined(__SSE__) +#define AILEGO_VERSION_SIMD " SSE Instruction Set\n" +#elif defined(__MMX__) +#define AILEGO_VERSION_SIMD " MMX Instruction Set\n" +#else +#define AILEGO_VERSION_SIMD "" +#endif + +#if defined(PY_VERSION) +#if PY_RELEASE_LEVEL == PY_RELEASE_LEVEL_ALPHA +#define AILEGO_VERSION_PYTHON \ + " Python API " PY_VERSION \ + " Alpha " AILEGO_VERSION_TO_STRING(PY_RELEASE_SERIAL) "\n" +#elif PY_RELEASE_LEVEL == PY_RELEASE_LEVEL_BETA +#define AILEGO_VERSION_PYTHON \ + " Python API " PY_VERSION \ + " Beta " AILEGO_VERSION_TO_STRING(PY_RELEASE_SERIAL) "\n" +#elif PY_RELEASE_LEVEL == PY_RELEASE_LEVEL_GAMMA +#define AILEGO_VERSION_PYTHON \ + " Python API " PY_VERSION \ + " Release Candidate " AILEGO_VERSION_TO_STRING(PY_RELEASE_SERIAL) "\n" +#elif PY_RELEASE_LEVEL == PY_RELEASE_LEVEL_FINAL +#define AILEGO_VERSION_PYTHON " Python API " PY_VERSION " Final\n" +#else +#define AILEGO_VERSION_PYTHON " Python API " PY_VERSION "\n" +#endif +#else +#define AILEGO_VERSION_PYTHON "" +#endif + +//! Gather information of compiling +#define AILEGO_VERSION_COMPILE_DETAILS(__PREFIX_INFO__) \ + __PREFIX_INFO__ \ + "Compiled by " AILEGO_VERSION_COMPILER \ + ".\n" \ + "Compiled for " AILEGO_VERSION_PROCESSOR \ + ".\n" \ + "Compiled on " AILEGO_VERSION_PLATFORM " on " __DATE__ " " __TIME__ \ + ".\n" \ + "Compiled with: \n" \ + "" AILEGO_VERSION_BYTE_ORDER "" AILEGO_VERSION_SIMD \ + "" AILEGO_VERSION_DEBUG_INFO "" AILEGO_VERSION_ASAN "" AILEGO_VERSION_STDC \ + "" AILEGO_VERSION_CPLUSPLUS "" AILEGO_VERSION_GXX_ABI \ + "" AILEGO_VERSION_POSIX "" AILEGO_VERSION_XOPEN "" AILEGO_VERSION_LSB \ + "" AILEGO_VERSION_GLIBC "" AILEGO_VERSION_WINSDK "" AILEGO_VERSION_CLR \ + "" AILEGO_VERSION_OPENMP "" AILEGO_VERSION_PYTHON "\n" diff --git a/src/binding/CMakeLists.txt b/src/binding/CMakeLists.txt new file mode 100644 index 0000000..700d081 --- /dev/null +++ b/src/binding/CMakeLists.txt @@ -0,0 +1,14 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +# Retrieve version from git repository +git_version(ZVEC_VERSION ${CMAKE_CURRENT_SOURCE_DIR}) + +# Add repositories +if(BUILD_C_BINDINGS) + cc_directory(c) +endif() + +if(BUILD_PYTHON_BINDINGS) + cc_directory(python) +endif() \ No newline at end of file diff --git a/src/binding/c/CMakeLists.txt b/src/binding/c/CMakeLists.txt new file mode 100644 index 0000000..13dfdc0 --- /dev/null +++ b/src/binding/c/CMakeLists.txt @@ -0,0 +1,256 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +include(GNUInstallDirs) + +# Retrieve version from git repository and generate version header +git_version(ZVEC_VERSION ${CMAKE_CURRENT_SOURCE_DIR}) + +# Debug: print version variables +message(STATUS "ZVEC_VERSION: ${ZVEC_VERSION}") + +# Parse version string to extract major.minor.patch +# Format: vX.Y.Z-commit-hash or vX.Y.Z +if(ZVEC_VERSION MATCHES "^v([0-9]+)\\.([0-9]+)\\.([0-9]+)") + set(ZVEC_VERSION_MAJOR "${CMAKE_MATCH_1}") + set(ZVEC_VERSION_MINOR "${CMAKE_MATCH_2}") + set(ZVEC_VERSION_PATCH "${CMAKE_MATCH_3}") + set(ZVEC_VERSION_STRING "${ZVEC_VERSION}") +else() + # Default version if parsing fails + set(ZVEC_VERSION_MAJOR 0) + set(ZVEC_VERSION_MINOR 2) + set(ZVEC_VERSION_PATCH 1) + set(ZVEC_VERSION_STRING "${ZVEC_VERSION_MAJOR}.${ZVEC_VERSION_MINOR}.${ZVEC_VERSION_PATCH}") +endif() + +message(STATUS "Parsed version: ${ZVEC_VERSION_MAJOR}.${ZVEC_VERSION_MINOR}.${ZVEC_VERSION_PATCH} (${ZVEC_VERSION_STRING})") + +# Configure version header file +configure_file( + ${PROJECT_SOURCE_DIR}/src/include/zvec/version.h.in + ${PROJECT_BINARY_DIR}/src/generated/zvec_version.h +) + +# Configure c_api.h with version information embedded +configure_file( + ${PROJECT_SOURCE_DIR}/src/include/zvec/c_api.h + ${PROJECT_BINARY_DIR}/src/generated/zvec/c_api.h + @ONLY +) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) + +# C API library source files +set(ZVEC_C_API_SOURCES + c_api.cc +) + +# C API library header files +set(ZVEC_C_API_HEADERS + ${PROJECT_SOURCE_DIR}/src/include/zvec/c_api.h +) + +# ============================================================================= +# Build FAT Shared Library (zvec_c_api.so) +# ============================================================================= +add_library(zvec_c_api SHARED + ${ZVEC_C_API_SOURCES} + ${ZVEC_C_API_HEADERS} +) + +# Set library properties +set_target_properties(zvec_c_api PROPERTIES + OUTPUT_NAME "zvec_c_api" + POSITION_INDEPENDENT_CODE ON + # Hide all symbols by default, only export C API + CXX_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON +) + +# On Windows, define ZVEC_BUILD_SHARED so that c_api.h uses __declspec(dllexport) +if(MSVC OR WIN32) + target_compile_definitions(zvec_c_api PRIVATE ZVEC_BUILD_SHARED) +endif() + +find_package(Threads REQUIRED) + +# Static linking of C++ standard library is handled in platform-specific sections + +# Determine whole-archive flag for current platform +if(MSVC) + # MSVC uses /WHOLEARCHIVE:filename syntax + set(ZVEC_WHOLE_ARCHIVE_PREFIX "/WHOLEARCHIVE:") + set(ZVEC_NO_WHOLE_ARCHIVE_SUFFIX "") +elseif(APPLE) + # macOS uses -Wl,-force_load,filename + set(ZVEC_WHOLE_ARCHIVE_PREFIX "-Wl,-force_load,") + set(ZVEC_NO_WHOLE_ARCHIVE_SUFFIX "") +else() + # Linux/Unix uses -Wl,--whole-archive and -Wl,--no-whole-archive + set(ZVEC_WHOLE_ARCHIVE_PREFIX "-Wl,--whole-archive") + set(ZVEC_NO_WHOLE_ARCHIVE_SUFFIX "-Wl,--no-whole-archive") +endif() + +# FAT mode: embed ALL libraries (including third-party) statically +# This creates a truly self-contained library with zero external dependencies +# Users only need to link libzvec_c_api.so without installing any dependencies +if(MSVC) + target_link_options(zvec_c_api PRIVATE + ${ZVEC_WHOLE_ARCHIVE_PREFIX}$ + ${ZVEC_WHOLE_ARCHIVE_PREFIX}$ + ${ZVEC_WHOLE_ARCHIVE_PREFIX}$ + ${ZVEC_WHOLE_ARCHIVE_PREFIX}$ + ) + + target_link_libraries(zvec_c_api + PRIVATE + zvec + zvec_core + zvec_ailego + zvec_turbo + ${ZVEC_CXX_STD_LIB} + Threads::Threads + ) +elseif(APPLE) + target_link_options(zvec_c_api PRIVATE + ${ZVEC_WHOLE_ARCHIVE_PREFIX}$ + ${ZVEC_WHOLE_ARCHIVE_PREFIX}$ + ${ZVEC_WHOLE_ARCHIVE_PREFIX}$ + ${ZVEC_WHOLE_ARCHIVE_PREFIX}$ + ) + + # Link the libraries themselves (without force_load flags) + target_link_libraries(zvec_c_api + PRIVATE + zvec + zvec_core + zvec_ailego + zvec_turbo + Threads::Threads + ${CMAKE_DL_LIBS} + ) + + # Optional: Static link libc++ for better portability (uncomment if needed) + # target_link_options(zvec_c_api PRIVATE -static-libstdc++) + +else() + # Linux and other Unix-like systems + target_link_libraries(zvec_c_api + PRIVATE + # Force load all zvec internal libraries (extract all objects) + ${ZVEC_WHOLE_ARCHIVE_PREFIX} + zvec + zvec_core + zvec_ailego + zvec_turbo # Required by zvec_core (metric module) + ${ZVEC_NO_WHOLE_ARCHIVE_SUFFIX} + Threads::Threads + ${CMAKE_DL_LIBS} + ) + + # Link C++ standard library appropriately for the platform + # On Linux, we need to explicitly link libstdc++ + # Try static linking first, fall back to dynamic + set(STATIC_LIBS_FOUND FALSE) + + # Try to find static libraries using compiler's built-in search + execute_process( + COMMAND ${CMAKE_CXX_COMPILER} -print-file-name=libstdc++.a + OUTPUT_VARIABLE LIBSTDCPP_A_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + execute_process( + COMMAND ${CMAKE_CXX_COMPILER} -print-file-name=libgcc.a + OUTPUT_VARIABLE LIBGCC_A_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + + # Check if both static libraries exist (not just filename returned) + # If library not found, -print-file-name returns just the filename without path + if(EXISTS "${LIBSTDCPP_A_PATH}" AND EXISTS "${LIBGCC_A_PATH}") + # Verify we got actual paths, not just filenames + if(IS_ABSOLUTE "${LIBSTDCPP_A_PATH}" AND IS_ABSOLUTE "${LIBGCC_A_PATH}") + set(STATIC_LIBS_FOUND TRUE) + message(STATUS "Found static libraries: ${LIBSTDCPP_A_PATH}") + endif() + endif() + + if(STATIC_LIBS_FOUND) + # Both static libraries found - use them + target_link_options(zvec_c_api PRIVATE + -static-libstdc++ + -static-libgcc + ) + message(STATUS "Using static libstdc++ and libgcc") + else() + # Fall back to dynamic linking + # Note: We still need to link the libraries, just not statically + message(STATUS "Static libraries not found, using dynamic linking for libstdc++ and libgcc") + endif() +endif() + +# Include directories +target_include_directories(zvec_c_api + PUBLIC + $ + $ + $ + PRIVATE + ${PROJECT_SOURCE_DIR}/src +) + +# Compile options +target_compile_options(zvec_c_api PRIVATE + $<$:-Wall -Wextra -Wpedantic> + $<$:-Wall -Wextra -Wpedantic> +) + +# Strip symbols in release builds to reduce library size +if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") + if(UNIX AND NOT APPLE) + # Linux: Use strip command to remove symbols + add_custom_command(TARGET zvec_c_api POST_BUILD + COMMAND ${CMAKE_STRIP} $ + COMMENT "Stripping symbols from libzvec_c_api.so to reduce size" + ) + elseif(APPLE) + # macOS: Use strip with platform-specific flags + add_custom_command(TARGET zvec_c_api POST_BUILD + COMMAND /usr/bin/strip -x $ + COMMENT "Stripping symbols from libzvec_c_api.dylib to reduce size" + ) + endif() +endif() + +# ============================================================================= +# Installation Rules +# ============================================================================= + +# Install shared library +install(TARGETS zvec_c_api + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} +) + +# Install headers +install(FILES ${PROJECT_SOURCE_DIR}/src/include/zvec/c_api.h + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/zvec +) diff --git a/src/binding/c/c_api.cc b/src/binding/c/c_api.cc new file mode 100644 index 0000000..0c50b32 --- /dev/null +++ b/src/binding/c/c_api.cc @@ -0,0 +1,7223 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// clang-format off +#include "zvec/c_api.h" +// Include generated version header for build-time +#if defined(__has_include) && __has_include() +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Error checking macros - these preserve __LINE__ accuracy +// Simplified macro for setting error with automatic file/line/function info +#define SET_LAST_ERROR(code, msg) \ + set_last_error_details(code, msg, __FILE__, __LINE__, __FUNCTION__) + +#define ZVEC_CHECK_NOTNULL(ptr, error_code, msg) \ + if (!(ptr)) { \ + SET_LAST_ERROR(error_code, msg); \ + return nullptr; \ + } + +#define ZVEC_CHECK_NOTNULL_ERRCODE(ptr, error_code, msg) \ + if (!(ptr)) { \ + SET_LAST_ERROR(error_code, msg); \ + return (error_code); \ + } + +#define ZVEC_CHECK_COND(cond, error_code, msg) \ + if (cond) { \ + SET_LAST_ERROR(error_code, msg); \ + return nullptr; \ + } + +#define ZVEC_CHECK_COND_ERRCODE(cond, error_code, msg) \ + if (cond) { \ + SET_LAST_ERROR(error_code, msg); \ + return (error_code); \ + } + +// For void functions (no return value): +#define ZVEC_TRY_BEGIN_VOID try { +#define ZVEC_CATCH_END_VOID \ + } \ + catch (const std::exception &e) { \ + SET_LAST_ERROR(ZVEC_ERROR_UNKNOWN, std::string("Exception: ") + e.what()); \ + } + +// For functions returning zvec_error_code_t - complete try-catch wrapper +#define ZVEC_TRY_BEGIN_CODE ZVEC_TRY_BEGIN_VOID +#define ZVEC_CATCH_END_CODE(code_on_error) \ + } \ + catch (const std::exception &e) { \ + SET_LAST_ERROR(ZVEC_ERROR_UNKNOWN, std::string("Exception: ") + e.what()); \ + return code_on_error; \ + } \ + return ZVEC_OK; + +// For functions returning pointer - complete try-catch wrapper +// Usage: ZVEC_TRY_RETURN_NULL("error msg", code...) +// Note: Use variadic macro to handle commas in template arguments +#define ZVEC_TRY_RETURN_NULL(msg, ...) \ + try { \ + { \ + __VA_ARGS__ \ + } \ + } catch (const std::bad_alloc &e) { \ + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, \ + std::string(msg) + ": " + e.what()); \ + return nullptr; \ + } catch (const std::exception &e) { \ + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, \ + std::string(msg) + ": " + e.what()); \ + return nullptr; \ + } + +// For functions returning ErrorCode +// Usage: ZVEC_TRY_RETURN_ERROR("error msg", code...) +// Note: Use variadic macro to handle commas in template arguments +#define ZVEC_TRY_RETURN_ERROR(msg, ...) \ + try { \ + { \ + __VA_ARGS__ \ + } \ + } catch (const std::bad_alloc &e) { \ + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, \ + std::string(msg) + ": " + e.what()); \ + return ZVEC_ERROR_RESOURCE_EXHAUSTED; \ + } catch (const std::exception &e) { \ + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, \ + std::string(msg) + ": " + e.what()); \ + return ZVEC_ERROR_INTERNAL_ERROR; \ + } + +// For functions returning scalar values (int, float, size_t, etc.) +// Usage: ZVEC_TRY_RETURN_SCALAR("error msg", error_value, code...) +// Note: Use variadic macro to handle commas in template arguments +#define ZVEC_TRY_RETURN_SCALAR(msg, error_val, ...) \ + try { \ + { \ + __VA_ARGS__ \ + } \ + } catch (const std::bad_alloc &e) { \ + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, \ + std::string(msg) + ": " + e.what()); \ + return (error_val); \ + } catch (const std::exception &e) { \ + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, \ + std::string(msg) + ": " + e.what()); \ + return (error_val); \ + } + +// Global status flags +static std::atomic g_initialized{false}; +static std::mutex g_init_mutex; + +// Thread-local storage for error information +static thread_local std::string last_error_message; +static thread_local zvec_error_details_t last_error_details; + +// Helper function: set error information (noexcept to avoid exceptions in error handling) +static void set_last_error(const std::string &msg) noexcept { + try { + last_error_message = msg; + last_error_details.code = ZVEC_ERROR_UNKNOWN; + last_error_details.message = last_error_message.c_str(); + last_error_details.file = nullptr; + last_error_details.line = 0; + last_error_details.function = nullptr; + } catch (...) { + // If we can't even store the error message, at least set the code + last_error_details.code = ZVEC_ERROR_RESOURCE_EXHAUSTED; + last_error_details.message = "Out of memory"; + } +} + +// Error setting function with detailed information (noexcept to avoid exceptions in error handling) +static void set_last_error_details(zvec_error_code_t code, const std::string &msg, + const char *file = nullptr, int line = 0, + const char *function = nullptr) noexcept { + try { + last_error_message = msg; + last_error_details.code = code; + last_error_details.message = last_error_message.c_str(); + last_error_details.file = file; + last_error_details.line = line; + last_error_details.function = function; + } catch (...) { + // If memory allocation fails, at least set the error code + last_error_details.code = ZVEC_ERROR_RESOURCE_EXHAUSTED; + last_error_details.message = "Out of memory"; + } +} + +// ============================================================================= +// Version information interface implementation +// ============================================================================= + +// Version string is generated at compile time by CMake +const char *zvec_get_version(void) { + // ZVEC_VERSION_STRING is a compile-time constant from zvec_version.h + // Format: "vX.Y.Z-commit-hash" or "g" + return ZVEC_VERSION_STRING; +} + +bool zvec_check_version(int major, int minor, int patch) { + if (major < 0 || minor < 0 || patch < 0) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Version numbers must be non-negative"); + return false; + } + + if (ZVEC_VERSION_MAJOR > major) return true; + if (ZVEC_VERSION_MAJOR < major) return false; + + if (ZVEC_VERSION_MINOR > minor) return true; + if (ZVEC_VERSION_MINOR < minor) return false; + + return ZVEC_VERSION_PATCH >= patch; +} + +int zvec_get_version_major(void) { + return ZVEC_VERSION_MAJOR; +} + +int zvec_get_version_minor(void) { + return ZVEC_VERSION_MINOR; +} + +int zvec_get_version_patch(void) { + return ZVEC_VERSION_PATCH; +} + +// ============================================================================= +// String management functions implementation +// ============================================================================= + +zvec_string_t *zvec_string_create(const char *str) { + if (!str) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "String pointer cannot be null"); + return nullptr; + } + + size_t len = strlen(str); + zvec_string_t *zstr = static_cast(malloc(sizeof(zvec_string_t))); + if (!zstr) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory for zvec_string_t"); + return nullptr; + } + + char *data_buffer = static_cast(malloc(len + 1)); + if (!data_buffer) { + free(zstr); + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory for string data"); + return nullptr; + } + + memcpy(data_buffer, str, len + 1); + zstr->data = data_buffer; + zstr->length = len; + zstr->capacity = len + 1; + return zstr; +} + +zvec_string_t *zvec_string_create_from_view(const zvec_string_view_t *view) { + if (!view || !view->data) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "String view or data cannot be null"); + return nullptr; + } + + zvec_string_t *zstr = static_cast(malloc(sizeof(zvec_string_t))); + if (!zstr) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory for zvec_string_t"); + return nullptr; + } + + char *data_buffer = static_cast(malloc(view->length + 1)); + if (!data_buffer) { + free(zstr); + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory for string data"); + return nullptr; + } + + memcpy(data_buffer, view->data, view->length); + data_buffer[view->length] = '\0'; + zstr->data = data_buffer; + zstr->length = view->length; + zstr->capacity = view->length + 1; + + return zstr; +} + +zvec_string_t *zvec_bin_create(const uint8_t *data, size_t length) { + if (!data) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Binary data pointer cannot be null"); + return nullptr; + } + + zvec_string_t *zstr = static_cast(malloc(sizeof(zvec_string_t))); + if (!zstr) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory for zvec_string_t"); + return nullptr; + } + + char *data_buffer = static_cast(malloc(length + 1)); + if (!data_buffer) { + free(zstr); + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory for binary data"); + return nullptr; + } + + memcpy(data_buffer, data, length); + data_buffer[length] = '\0'; + zstr->data = data_buffer; + zstr->length = length; + zstr->capacity = length + 1; + + return zstr; +} + +zvec_string_t *zvec_string_copy(const zvec_string_t *str) { + if (!str || !str->data) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Source string or data cannot be null"); + return nullptr; + } + + return zvec_string_create(str->data); +} + +const char *zvec_string_c_str(const zvec_string_t *str) { + if (!str) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "String pointer cannot be null"); + return nullptr; + } + + return str->data; +} + +size_t zvec_string_length(const zvec_string_t *str) { + if (!str) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "String pointer cannot be null"); + return 0; + } + + return str->length; +} + +int zvec_string_compare(const zvec_string_t *str1, const zvec_string_t *str2) { + if (!str1 || !str2) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "String pointers cannot be null"); + return -1; + } + + if (!str1->data || !str2->data) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "String data cannot be null"); + return -1; + } + + return strcmp(str1->data, str2->data); +} + +// ============================================================================= +// Configuration-related functions implementation +// ============================================================================= + +zvec_log_config_t *zvec_config_log_create_console(zvec_log_level_t level) { + try { + auto *config = new zvec::GlobalConfig::ConsoleLogConfig( + static_cast(level)); + return reinterpret_cast(config); + } catch (const std::exception &e) { + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, e.what()); + return nullptr; + } +} + +zvec_log_config_t *zvec_config_log_create_file(zvec_log_level_t level, const char *dir, + const char *basename, + uint32_t file_size, + uint32_t overdue_days) { + if (!dir || !basename) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Directory or basename cannot be null"); + return nullptr; + } + + try { + auto *config = new zvec::GlobalConfig::FileLogConfig( + static_cast(level), std::string(dir), + std::string(basename), file_size, overdue_days); + return reinterpret_cast(config); + } catch (const std::exception &e) { + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, e.what()); + return nullptr; + } +} + +void zvec_config_log_destroy(zvec_log_config_t *config) { + if (config) { + delete reinterpret_cast(config); + } +} + +zvec_log_level_t zvec_config_log_get_level(const zvec_log_config_t *config) { + if (!config) { + return ZVEC_LOG_LEVEL_WARN; + } + auto *cpp_config = + reinterpret_cast(config); + return static_cast(cpp_config->level); +} + +zvec_error_code_t zvec_config_log_set_level(zvec_log_config_t *config, + zvec_log_level_t level) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->level = static_cast(level); + return ZVEC_OK; +} + +bool zvec_config_log_is_file_type(const zvec_log_config_t *config) { + if (!config) { + return false; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->GetLoggerType() == zvec::FILE_LOG_TYPE_NAME; +} + +inline const zvec::GlobalConfig::FileLogConfig* file_config_from_config(const zvec_log_config_t *config) { + auto *cpp_config = + reinterpret_cast(config); + return dynamic_cast(cpp_config); +} + +inline zvec::GlobalConfig::FileLogConfig* mutable_file_config_from_config(zvec_log_config_t *config) { + auto *cpp_config = + reinterpret_cast(config); + return dynamic_cast(cpp_config); +} + +const char *zvec_config_log_get_dir(const zvec_log_config_t *config) { + auto* file_config = file_config_from_config(config); + return file_config->dir.c_str(); +} + +zvec_error_code_t zvec_config_log_set_dir(zvec_log_config_t *config, const char *dir) { + auto *file_config = mutable_file_config_from_config(config); + file_config->dir = dir; + return ZVEC_OK; +} + +const char *zvec_config_log_get_basename(const zvec_log_config_t *config) { + auto* file_config = file_config_from_config(config); + return file_config->basename.c_str(); +} + +zvec_error_code_t zvec_config_log_set_basename(zvec_log_config_t *config, + const char *basename) { + auto *file_config = mutable_file_config_from_config(config); + file_config->basename = basename; + return ZVEC_OK; +} + +uint32_t zvec_config_log_get_file_size(const zvec_log_config_t *config) { + auto* file_config = file_config_from_config(config); + return file_config->file_size; +} + +zvec_error_code_t zvec_config_log_set_file_size(zvec_log_config_t *config, + uint32_t file_size) { + auto *file_config = mutable_file_config_from_config(config); + file_config->file_size = file_size; + return ZVEC_OK; +} + +uint32_t zvec_config_log_get_overdue_days(const zvec_log_config_t *config) { + auto* file_config = file_config_from_config(config); + return file_config ? file_config->overdue_days : 0; +} + +zvec_error_code_t zvec_config_log_set_overdue_days(zvec_log_config_t *config, + uint32_t days) { + auto *file_config = mutable_file_config_from_config(config); + file_config->overdue_days = days; + return ZVEC_OK; +} + +// ============================================================================ +// Configuration Data Management Functions +// ============================================================================ + +zvec_config_data_t *zvec_config_data_create(void) { + try { + auto *config = new zvec::GlobalConfig::ConfigData(); + return reinterpret_cast(config); + } catch (const std::exception &e) { + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, e.what()); + return nullptr; + } +} + +void zvec_config_data_destroy(zvec_config_data_t *config) { + if (config) { + delete reinterpret_cast(config); + } +} + +zvec_error_code_t zvec_config_data_set_memory_limit(zvec_config_data_t *config, + uint64_t memory_limit_bytes) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->memory_limit_bytes = memory_limit_bytes; + return ZVEC_OK; +} + +uint64_t zvec_config_data_get_memory_limit(const zvec_config_data_t *config) { + if (!config) { + return 0; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->memory_limit_bytes; +} + +zvec_error_code_t zvec_config_data_set_log_config(zvec_config_data_t *config, + zvec_log_config_t *log_config) { + if (!config || !log_config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Config or log_config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *cpp_config = reinterpret_cast(config); + + // Convert raw pointer to shared_ptr for C++ internal use + auto *log_config_raw = + reinterpret_cast(log_config); + cpp_config->log_config = std::shared_ptr( + log_config_raw, [](zvec::GlobalConfig::LogConfig *ptr) { delete ptr; }); + + return ZVEC_OK; +} + +zvec_log_type_t zvec_config_data_get_log_type(const zvec_config_data_t *config) { + if (!config) { + return ZVEC_LOG_TYPE_CONSOLE; + } + + const auto *cpp_config = + reinterpret_cast(config); + if (!cpp_config->log_config) { + return ZVEC_LOG_TYPE_CONSOLE; + } + + if (cpp_config->log_config->GetLoggerType() == zvec::FILE_LOG_TYPE_NAME) { + return ZVEC_LOG_TYPE_FILE; + } + return ZVEC_LOG_TYPE_CONSOLE; +} + +zvec_error_code_t zvec_config_data_set_query_thread_count(zvec_config_data_t *config, + uint32_t thread_count) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->query_thread_count = thread_count; + return ZVEC_OK; +} + +uint32_t zvec_config_data_get_query_thread_count(const zvec_config_data_t *config) { + if (!config) { + return 1; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->query_thread_count; +} + +zvec_error_code_t zvec_config_data_set_invert_to_forward_scan_ratio( + zvec_config_data_t *config, float ratio) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->invert_to_forward_scan_ratio = ratio; + return ZVEC_OK; +} + +float zvec_config_data_get_invert_to_forward_scan_ratio( + const zvec_config_data_t *config) { + if (!config) { + return 0.0f; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->invert_to_forward_scan_ratio; +} + +zvec_error_code_t zvec_config_data_set_brute_force_by_keys_ratio( + zvec_config_data_t *config, float ratio) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->brute_force_by_keys_ratio = ratio; + return ZVEC_OK; +} + +float zvec_config_data_get_brute_force_by_keys_ratio( + const zvec_config_data_t *config) { + if (!config) { + return 0.0f; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->brute_force_by_keys_ratio; +} + +zvec_error_code_t zvec_config_data_set_fts_brute_force_by_keys_ratio( + zvec_config_data_t *config, float ratio) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->fts_brute_force_by_keys_ratio = ratio; + return ZVEC_OK; +} + +float zvec_config_data_get_fts_brute_force_by_keys_ratio( + const zvec_config_data_t *config) { + if (!config) { + return 0.0f; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->fts_brute_force_by_keys_ratio; +} + +zvec_error_code_t zvec_config_data_set_optimize_thread_count( + zvec_config_data_t *config, uint32_t thread_count) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->optimize_thread_count = thread_count; + return ZVEC_OK; +} + +uint32_t zvec_config_data_get_optimize_thread_count( + const zvec_config_data_t *config) { + if (!config) { + return 1; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->optimize_thread_count; +} + +zvec_error_code_t zvec_config_data_set_jieba_dict_dir( + zvec_config_data_t *config, const char *dir) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->jieba_dict_dir = (dir != nullptr) ? std::string(dir) : ""; + return ZVEC_OK; +} + +const char *zvec_config_data_get_jieba_dict_dir( + const zvec_config_data_t *config) { + if (!config) { + return ""; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->jieba_dict_dir.c_str(); +} + + +// ============================================================================= +// Initialization and cleanup interface implementation +// ============================================================================= + +zvec_error_code_t zvec_initialize(const zvec_config_data_t *config) { + std::lock_guard lock(g_init_mutex); + + if (g_initialized.load()) { + SET_LAST_ERROR(ZVEC_ERROR_ALREADY_EXISTS, "Library already initialized"); + return ZVEC_ERROR_ALREADY_EXISTS; + } + + ZVEC_TRY_RETURN_ERROR( + "Initialization failed", + // Convert to C++ configuration object + zvec::GlobalConfig::ConfigData cpp_config{}; + + if (config) { + auto *cpp_config_data = + reinterpret_cast(config); + cpp_config = *cpp_config_data; // Copy the C++ ConfigData + } else { + // Initialize with default configuration + cpp_config = zvec::GlobalConfig::ConfigData{}; + } + + // Initialize global configuration + auto status = zvec::GlobalConfig::Instance().Initialize(cpp_config); + if (!status.ok()) { + set_last_error(status.message()); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + g_initialized.store(true); + return ZVEC_OK;) +} + +zvec_error_code_t zvec_shutdown(void) { + std::lock_guard lock(g_init_mutex); + + if (!g_initialized.load()) { + SET_LAST_ERROR(ZVEC_ERROR_FAILED_PRECONDITION, "Library not initialized"); + return ZVEC_ERROR_FAILED_PRECONDITION; + } + // We're do nothing here for now, + // but we might add zvec::GlobalConfig::Finalize in the future. + ZVEC_TRY_RETURN_ERROR("Shutdown failed", g_initialized.store(false); + return ZVEC_OK;) +} + +bool zvec_is_initialized(void) { + return g_initialized.load(); +} + +void zvec_set_default_jieba_dict_dir(const char *dir) { + zvec::GlobalConfig::Instance().set_default_jieba_dict_dir( + (dir != nullptr) ? std::string(dir) : std::string()); +} + +const char *zvec_get_default_jieba_dict_dir(void) { + // Thread-local buffer keeps c_str() valid until the next call on this thread. + thread_local std::string cached; + cached = zvec::GlobalConfig::Instance().jieba_dict_dir(); + return cached.c_str(); +} + +// ============================================================================= +// Error handling interface implementation +// ============================================================================= + +zvec_error_code_t zvec_get_last_error_details(zvec_error_details_t *error_details) { + if (!error_details) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Error details pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + *error_details = last_error_details; + return ZVEC_OK; +} + +void zvec_clear_error(void) { + last_error_message.clear(); + last_error_details = {}; +} + +// Helper functions: convert internal status to error code +static zvec_error_code_t status_to_error_code(const zvec::Status &status) { + if (status.code() < zvec::StatusCode::OK || + status.code() > zvec::StatusCode::UNKNOWN) { + set_last_error("Unexpected status code: " + + std::to_string(static_cast(status.code()))); + return ZVEC_ERROR_UNKNOWN; + } + + return static_cast(status.code()); +} + +// Helper function: handle Expected results +template +static zvec_error_code_t handle_expected_result( + const tl::expected &result, T *out_value = nullptr) { + if (result.has_value()) { + if (out_value) { + *out_value = result.value(); + } + return ZVEC_OK; + } else { + set_last_error(result.error().message()); + return status_to_error_code(result.error()); + } +} + +/** + * @brief Copy a C++ string to C heap-allocated string + * @param str String to copy + * @return Newly allocated C string, or NULL on failure + * @note Caller must free() the returned string + */ +static char *copy_string(const std::string &str) { + if (str.empty()) return nullptr; + size_t len = str.length(); + char *copy = static_cast(malloc(len + 1)); + if (!copy) return nullptr; + strncpy(copy, str.c_str(), len); + copy[len] = '\0'; // Ensure null-termination + return copy; +} + +/** + * @brief Free write results array returned by DML APIs + * @param results Results array to free + * @param result_count Number of results + */ +static void free_write_results_internal(zvec_write_result_t *results, + size_t result_count) { + if (!results) { + return; + } + for (size_t i = 0; i < result_count; ++i) { + // pk is not stored (ordered style), only free message + if (results[i].message) { + free((void *)results[i].message); + results[i].message = nullptr; + } + } + free(results); +} + +// Helper function: convert per-doc statuses to C API write result array. +static zvec_error_code_t build_write_results( + const std::vector &statuses, zvec_write_result_t **results, + size_t *result_count) { + if (!results || !result_count) { + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + *result_count = statuses.size(); + if (*result_count == 0) { + *results = nullptr; + return ZVEC_OK; + } + + *results = static_cast( + calloc(*result_count, sizeof(zvec_write_result_t))); + if (!*results) { + set_last_error("Failed to allocate memory for write results"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // Use ordered style: result index corresponds to input index. + // No need to store pk in result, caller can access by index. + for (size_t i = 0; i < *result_count; ++i) { + const std::string message = statuses[i].message(); + (*results)[i].message = copy_string(message); + (*results)[i].code = status_to_error_code(statuses[i]); + } + + return ZVEC_OK; +} + +static std::vector collect_doc_pks(const zvec_doc_t **docs, + size_t doc_count) { + std::vector pks; + pks.reserve(doc_count); + for (size_t i = 0; i < doc_count; ++i) { + if (!docs[i]) { + pks.emplace_back(""); + continue; + } + auto *doc_ptr = reinterpret_cast(docs[i]); + pks.emplace_back(doc_ptr->pk_ref()); + } + return pks; +} + +// ============================================================================= +// Type conversion helpers +// ============================================================================= + +/** + * @brief Convert C index params to C++ shared_ptr + * @param params C index params handle + * @return Shared pointer to C++ IndexParams, or nullptr on failure + */ +static std::shared_ptr convert_c_index_params_to_cpp( + const zvec_index_params_t *params) { + if (!params) { + return nullptr; + } + return reinterpret_cast(params)->clone(); +} + +// ============================================================================= +// Memory Management interface implementation +// ============================================================================= + +/** + * @brief Allocate memory within the library + * @param size Number of bytes to allocate + * @return Pointer to allocated memory, or NULL on failure + * + * @note Use zvec_malloc instead of malloc to ensure memory is managed + * consistently within the library. All memory allocated with zvec_malloc + * should be freed with zvec_free. + */ +void* zvec_malloc(size_t size) { + return malloc(size); +} + +/** + * @brief Free memory allocated by zvec_malloc + * @param ptr Pointer to memory to free (can be NULL) + * + * @note Use zvec_free instead of free to ensure memory is managed + * consistently within the library. This function should be used to free + * any memory allocated with zvec_malloc or returned by library functions + * that document they return library-allocated memory. + */ +void zvec_free(void *ptr) { + if (ptr) { + free(ptr); + } +} + +/** + * @brief Free a zvec_string_t structure + * @param str String structure to free (can be NULL) + */ +void zvec_free_string(zvec_string_t *str) { + if (str) { + if (str->data) { + free((void *)str->data); + } + free(str); + } +} + +/** + * @brief Create a string array with given count + * @param count Number of strings in the array + * @return New string array, or NULL on failure + */ +zvec_string_array_t *zvec_string_array_create(size_t count) { + zvec_string_array_t *array = (zvec_string_array_t *)malloc(sizeof(zvec_string_array_t)); + array->count = count; + array->strings = (zvec_string_t *)malloc(sizeof(zvec_string_t) * count); + memset(array->strings, 0, sizeof(zvec_string_t) * count); + return array; +} + +/** + * @brief Create a string array from C-string array + * @param strings Array of C-strings + * @param count Number of strings + * @return New string array, or NULL on failure + */ +zvec_string_array_t *zvec_string_array_create_from_strings(const char **strings, + size_t count) { + if (!strings || count == 0) { + return nullptr; + } + zvec_string_array_t *array = zvec_string_array_create(count); + for (size_t i = 0; i < count; ++i) { + zvec_string_array_add(array, i, strings[i]); + } + return array; +} + +/** + * @brief Add a string to string array at specified index + * @param array String array to add to + * @param idx Index to add at + * @param str String to add + */ +void zvec_string_array_add(zvec_string_array_t *array, size_t idx, + const char *str) { + if (idx >= array->count) return; + size_t len = strlen(str); + array->strings[idx].data = (char *)malloc(len + 1); + memcpy(array->strings[idx].data, str, len + 1); + array->strings[idx].length = len; + array->strings[idx].capacity = len + 1; +} + +/** + * @brief Destroy a string array and free all memory + * @param array String array to destroy (can be NULL) + */ +void zvec_string_array_destroy(zvec_string_array_t *array) { + if (!array) return; + for (size_t i = 0; i < array->count; i++) { + if (array->strings[i].data) { + free((void *)array->strings[i].data); + } + } + free(array->strings); + free(array); +} + + +/** + * @brief Create a mutable byte array with given capacity + * @param capacity Initial capacity in bytes + * @return New byte array, or NULL on failure + */ +// Byte array helper functions +zvec_mutable_byte_array_t *zvec_byte_array_create(size_t capacity) { + zvec_mutable_byte_array_t *array = + (zvec_mutable_byte_array_t *)malloc(sizeof(zvec_mutable_byte_array_t)); + if (!array) return nullptr; + + array->data = (uint8_t *)malloc(capacity); + if (!array->data) { + free(array); + return nullptr; + } + + array->length = 0; + array->capacity = capacity; + memset(array->data, 0, capacity); + return array; +} + +/** + * @brief Destroy a byte array and free all memory + * @param array Byte array to destroy (can be NULL) + */ +void zvec_byte_array_destroy(zvec_mutable_byte_array_t *array) { + if (!array) return; + if (array->data) { + free(array->data); + } + free(array); +} + +/** + * @brief Create a float array with given count + * @param count Number of floats in the array + * @return New float array, or NULL on failure + */ +// Float array helper functions +zvec_float_array_t *zvec_float_array_create(size_t count) { + zvec_float_array_t *array = (zvec_float_array_t *)malloc(sizeof(zvec_float_array_t)); + if (!array) return nullptr; + + array->data = (const float *)malloc(sizeof(float) * count); + if (!array->data) { + free(array); + return nullptr; + } + + array->length = count; + memset((void *)array->data, 0, sizeof(float) * count); + return array; +} + +/** + * @brief Destroy a float array and free all memory + * @param array Float array to destroy (can be NULL) + */ +void zvec_float_array_destroy(zvec_float_array_t *array) { + if (!array) return; + if (array->data) { + free((void *)array->data); + } + free(array); +} + +// Int64 array helper functions +zvec_int64_array_t *zvec_int64_array_create(size_t count) { + zvec_int64_array_t *array = (zvec_int64_array_t *)malloc(sizeof(zvec_int64_array_t)); + if (!array) return nullptr; + + array->data = (const int64_t *)malloc(sizeof(int64_t) * count); + if (!array->data) { + free(array); + return nullptr; + } + + array->length = count; + memset((void *)array->data, 0, sizeof(int64_t) * count); + return array; +} + +void zvec_int64_array_destroy(zvec_int64_array_t *array) { + if (!array) return; + if (array->data) { + free((void *)array->data); + } + free(array); +} + +void zvec_free_str_array(char **array, size_t count) { + if (!array) return; + + // If count is 0, only free the string array itself, don't process internal + // strings + if (count == 0) { + free(array); + return; + } + + for (size_t i = 0; i < count; ++i) { + if (array[i]) { // Only free when string pointer is not null + free(array[i]); + } + } + free(array); +} + +zvec_error_code_t zvec_get_last_error(char **error_msg) { + if (!error_msg) { + set_last_error("Invalid argument: error_msg cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + *error_msg = copy_string(last_error_message); + return ZVEC_OK; +} + +void zvec_free_uint8_array(uint8_t *array) { + if (array) { + free(array); + } +} + +void zvec_free_field_schema(zvec_field_schema_t *field_schema) { + if (field_schema) { + // index_params is embedded, no need to free + free(field_schema); + } +} + +// ============================================================================= +// CollectionOptions functions implementation +// ============================================================================= + +zvec_collection_options_t *zvec_collection_options_create(void) { + ZVEC_TRY_RETURN_NULL( + "Failed to create zvec_collection_options_t", + auto *options = new zvec::CollectionOptions(); + return reinterpret_cast(options);) + return nullptr; +} + +void zvec_collection_options_destroy(zvec_collection_options_t *options) { + if (options) { + delete reinterpret_cast(options); + } +} + +zvec_error_code_t zvec_collection_options_set_enable_mmap( + zvec_collection_options_t *options, bool enable) { + if (!options) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection options pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(options); + ptr->enable_mmap_ = enable; + return ZVEC_OK; +} + +bool zvec_collection_options_get_enable_mmap( + const zvec_collection_options_t *options) { + if (!options) { + return true; // Default + } + auto *ptr = reinterpret_cast(options); + return ptr->enable_mmap_; +} + +zvec_error_code_t zvec_collection_options_set_max_buffer_size( + zvec_collection_options_t *options, size_t size) { + if (!options) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection options pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(options); + ptr->max_buffer_size_ = static_cast(size); + return ZVEC_OK; +} + +size_t zvec_collection_options_get_max_buffer_size( + const zvec_collection_options_t *options) { + if (!options) { + return zvec::DEFAULT_MAX_BUFFER_SIZE; // Default + } + auto *ptr = reinterpret_cast(options); + return ptr->max_buffer_size_; +} + +zvec_error_code_t zvec_collection_options_set_read_only( + zvec_collection_options_t *options, bool read_only) { + if (!options) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection options pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(options); + ptr->read_only_ = read_only; + return ZVEC_OK; +} + +bool zvec_collection_options_get_read_only( + const zvec_collection_options_t *options) { + if (!options) { + return false; // Default + } + auto *ptr = reinterpret_cast(options); + return ptr->read_only_; +} + +// ============================================================================= +// CollectionStats functions implementation +// ============================================================================= + +/** + * @brief Get document count from collection stats + * @param stats Collection stats handle + * @return Document count, or 0 if stats is NULL + */ +uint64_t zvec_collection_stats_get_doc_count(const zvec_collection_stats_t *stats) { + if (!stats) { + return 0; + } + auto *ptr = reinterpret_cast(stats); + return ptr->doc_count; +} + +/** + * @brief Get index count from collection stats + * @param stats Collection stats handle + * @return Number of indexes, or 0 if stats is NULL + */ +size_t zvec_collection_stats_get_index_count(const zvec_collection_stats_t *stats) { + if (!stats) { + return 0; + } + auto *ptr = reinterpret_cast(stats); + return ptr->index_completeness.size(); +} + +/** + * @brief Get index name at specified index + * @param stats Collection stats handle + * @param index Index position (0-based) + * @return Index name C-string, or NULL if invalid + * @note Returned string is owned by stats, do not free + */ +const char *zvec_collection_stats_get_index_name( + const zvec_collection_stats_t *stats, size_t index) { + if (!stats) { + return nullptr; + } + auto *ptr = reinterpret_cast(stats); + if (index >= ptr->index_completeness.size()) { + return nullptr; + } + // Return pointer to string data - caller should not free + auto it = ptr->index_completeness.begin(); + std::advance(it, index); + return it->first.c_str(); +} + +/** + * @brief Get index completeness at specified index + * @param stats Collection stats handle + * @param index Index position (0-based) + * @return Completeness value (0.0-1.0), or 0.0 if invalid + */ +float zvec_collection_stats_get_index_completeness( + const zvec_collection_stats_t *stats, size_t index) { + if (!stats) { + return 0.0f; + } + auto *ptr = reinterpret_cast(stats); + if (index >= ptr->index_completeness.size()) { + return 0.0f; + } + auto it = ptr->index_completeness.begin(); + std::advance(it, index); + return it->second; +} + +// ============================================================================= +// IndexParams functions implementation +// ============================================================================= + +/** + * @brief Create index parameters of specified type + * @param index_type Type of index to create + * @return New index params handle, or NULL on failure + * @note Caller must call zvec_index_params_destroy() to free + */ +zvec_index_params_t *zvec_index_params_create(zvec_index_type_t index_type) { + ZVEC_TRY_RETURN_NULL( + "Failed to create zvec_index_params_t", + // Create appropriate C++ IndexParams based on type with default parameters + zvec::IndexParams *cpp_params = nullptr; + + switch (index_type) { + case ZVEC_INDEX_TYPE_INVERT: + cpp_params = + new zvec::InvertIndexParams(true, // enable_range_optimization + false); // enable_extended_wildcard + break; + case ZVEC_INDEX_TYPE_FTS: + // Defaults align with FtsIndexParams default ctor: + // tokenizer="standard", filters=["lowercase"], extra="". + cpp_params = new zvec::FtsIndexParams(); + break; + case ZVEC_INDEX_TYPE_HNSW: + cpp_params = + new zvec::HnswIndexParams( + zvec::MetricType::L2, // metric_type + zvec::core_interface::kDefaultHnswNeighborCnt, // m + zvec::core_interface::kDefaultHnswEfConstruction, // ef_construction + zvec::QuantizeType::UNDEFINED); + break; + case ZVEC_INDEX_TYPE_IVF: + cpp_params = + new zvec::IVFIndexParams(zvec::MetricType::L2, // metric_type + 1024, // n_list (default) + 10, // n_iters (default) + false, // use_soar (default) + zvec::QuantizeType::UNDEFINED); + break; + case ZVEC_INDEX_TYPE_VAMANA: + cpp_params = + new zvec::VamanaIndexParams( + zvec::MetricType::L2, // metric_type + zvec::core_interface::kDefaultVamanaMaxDegree, + zvec::core_interface::kDefaultVamanaSearchListSize, + zvec::core_interface::kDefaultVamanaAlpha, + zvec::core_interface::kDefaultVamanaSaturateGraph, + false, // use_contiguous_memory + false, // use_id_map + zvec::QuantizeType::UNDEFINED); + break; + case ZVEC_INDEX_TYPE_FLAT: + default: + cpp_params = + new zvec::FlatIndexParams(zvec::MetricType::L2, // metric_type + zvec::QuantizeType::UNDEFINED); + break; + case ZVEC_INDEX_TYPE_DISKANN: + cpp_params = + new zvec::DiskAnnIndexParams( + zvec::MetricType::L2, // metric_type + 100, // max_degree (default) + 50, // list_size (default) + 0, // pq_chunk_num (default) + zvec::QuantizeType::UNDEFINED); + break; + } + + // Return as opaque pointer (raw pointer) + return reinterpret_cast(cpp_params);) + + return nullptr; +} + +/** + * @brief Destroy index parameters and free memory + * @param params Index params to destroy (can be NULL) + */ +void zvec_index_params_destroy(zvec_index_params_t *params) { + if (params) { + delete reinterpret_cast(params); + } +} + +/** + * @brief Set metric type for vector index parameters + * @param params Index parameters (must be vector index type) + * @param metric_type Metric type to set + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_set_metric_type(zvec_index_params_t *params, + zvec_metric_type_t metric_type) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Index params pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + + // Set metric type in the underlying C++ IndexParams + if (cpp_params->is_vector_index_type()) { + auto *vec_params = dynamic_cast(cpp_params); + if (vec_params) { + vec_params->set_metric_type(static_cast(metric_type)); + } + } + return ZVEC_OK; +} + +/** + * @brief Get metric type from index parameters + * @param params Index parameters + * @return Metric type, or default (L2) if NULL or not vector index + */ +zvec_metric_type_t zvec_index_params_get_metric_type( + const zvec_index_params_t *params) { + if (!params) { + return ZVEC_METRIC_TYPE_L2; // Default + } + auto *cpp_params = reinterpret_cast(params); + + if (cpp_params->is_vector_index_type()) { + auto *vec_params = + dynamic_cast(cpp_params); + if (vec_params) { + return static_cast( + static_cast(vec_params->metric_type())); + } + } + return ZVEC_METRIC_TYPE_L2; +} + +/** + * @brief Set quantization type for vector index parameters + * @param params Index parameters (must be vector index type) + * @param quantize_type Quantization type to set + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_set_quantize_type( + zvec_index_params_t *params, zvec_quantize_type_t quantize_type) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Index params pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + + // Set quantize type in the underlying C++ IndexParams + if (cpp_params->is_vector_index_type()) { + auto *vec_params = dynamic_cast(cpp_params); + if (vec_params) { + vec_params->set_quantize_type( + static_cast(quantize_type)); + } + } + return ZVEC_OK; +} + +/** + * @brief Get quantization type from index parameters + * @param params Index parameters + * @return Quantization type, or UNDEFINED if NULL or not vector index + */ +zvec_quantize_type_t zvec_index_params_get_quantize_type( + const zvec_index_params_t *params) { + if (!params) { + return ZVEC_QUANTIZE_TYPE_UNDEFINED; + } + auto *cpp_params = reinterpret_cast(params); + + if (cpp_params->is_vector_index_type()) { + auto *vec_params = + dynamic_cast(cpp_params); + if (vec_params) { + return static_cast( + static_cast(vec_params->quantize_type())); + } + } + return ZVEC_QUANTIZE_TYPE_UNDEFINED; +} + +/** + * @brief Set enable_rotate for quantizer parameters + * @param params Index parameters (must be vector index type) + * @param enable_rotate Whether to enable random rotation before quantization + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_set_quantizer_enable_rotate( + zvec_index_params_t *params, bool enable_rotate) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Index params pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + + if (!cpp_params->is_vector_index_type()) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Index params is not a vector index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *vec_params = dynamic_cast(cpp_params); + if (!vec_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Failed to cast to VectorIndexParams"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + zvec::QuantizerParam qp = vec_params->quantizer_param(); + qp.set_enable_rotate(enable_rotate); + vec_params->set_quantizer_param(qp); + return ZVEC_OK; +} + +/** + * @brief Get enable_rotate setting from quantizer parameters + * @param params Index parameters + * @return true if rotation is enabled, false otherwise + */ +bool zvec_index_params_get_quantizer_enable_rotate( + const zvec_index_params_t *params) { + if (!params) { + return false; + } + auto *cpp_params = reinterpret_cast(params); + + if (cpp_params->is_vector_index_type()) { + auto *vec_params = + dynamic_cast(cpp_params); + if (vec_params) { + return vec_params->quantizer_param().enable_rotate(); + } + } + return false; +} + +/** + * @brief Get index type from index parameters + * @param params Index parameters + * @return Index type, or FLAT as default if NULL + */ +zvec_index_type_t zvec_index_params_get_type(const zvec_index_params_t *params) { + if (!params) { + return ZVEC_INDEX_TYPE_FLAT; // Default + } + auto *cpp_params = reinterpret_cast(params); + return static_cast( + static_cast(cpp_params->type())); +} + +/** + * @brief Set HNSW-specific parameters + * @param params Index parameters (must be HNSW type) + * @param m Graph connectivity parameter + * @param ef_construction Construction exploration factor + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_set_hnsw_params(zvec_index_params_t *params, int m, + int ef_construction) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not HNSW index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *hnsw_params = dynamic_cast(cpp_params); + if (!hnsw_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not HNSW index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + hnsw_params->set_m(m); + hnsw_params->set_ef_construction(ef_construction); + return ZVEC_OK; +} + +/** + * @brief Get HNSW m parameter + * @param params Index parameters (must be HNSW type) + * @return m parameter value, or 0 on error + */ +int zvec_index_params_get_hnsw_m(const zvec_index_params_t *params) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not HNSW index type"); + return 0; + } + auto *cpp_params = reinterpret_cast(params); + auto *hnsw_params = dynamic_cast(cpp_params); + if (!hnsw_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not HNSW index type"); + return 0; + } + return hnsw_params->m(); +} + +/** + * @brief Get HNSW ef_construction parameter + * @param params Index parameters (must be HNSW type) + * @return ef_construction parameter value, or 0 on error + */ +int zvec_index_params_get_hnsw_ef_construction(const zvec_index_params_t *params) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not HNSW index type"); + return 0; + } + auto *cpp_params = reinterpret_cast(params); + auto *hnsw_params = dynamic_cast(cpp_params); + if (!hnsw_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not HNSW index type"); + return 0; + } + return hnsw_params->ef_construction(); +} + +zvec_error_code_t zvec_index_params_set_vamana_params( + zvec_index_params_t *params, int max_degree, int search_list_size, + float alpha, bool saturate_graph, bool use_contiguous_memory) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not Vamana index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *vamana_params = dynamic_cast(cpp_params); + if (!vamana_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not Vamana index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + vamana_params->set_max_degree(max_degree); + vamana_params->set_search_list_size(search_list_size); + vamana_params->set_alpha(alpha); + vamana_params->set_saturate_graph(saturate_graph); + vamana_params->set_use_contiguous_memory(use_contiguous_memory); + return ZVEC_OK; +} + +zvec_error_code_t zvec_index_params_get_vamana_params( + const zvec_index_params_t *params, int *out_max_degree, + int *out_search_list_size, float *out_alpha, bool *out_saturate_graph, + bool *out_use_contiguous_memory) { + if (!params || !out_max_degree || !out_search_list_size || !out_alpha || + !out_saturate_graph || !out_use_contiguous_memory) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or output pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *vamana_params = + dynamic_cast(cpp_params); + if (!vamana_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not Vamana index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + *out_max_degree = vamana_params->max_degree(); + *out_search_list_size = vamana_params->search_list_size(); + *out_alpha = vamana_params->alpha(); + *out_saturate_graph = vamana_params->saturate_graph(); + *out_use_contiguous_memory = vamana_params->use_contiguous_memory(); + return ZVEC_OK; +} + +/** + * @brief Set DiskANN-specific parameters + * @param params Index parameters (must be DiskANN type) + * @param max_degree Graph connectivity (max degree of Vamana graph) + * @param list_size Build-time list size + * @param pq_chunk_num PQ chunk count + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_set_diskann_params( + zvec_index_params_t *params, int max_degree, int list_size, + int pq_chunk_num) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not DiskANN index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *diskann_params = dynamic_cast(cpp_params); + if (!diskann_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not DiskANN index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + diskann_params->set_max_degree(max_degree); + diskann_params->set_list_size(list_size); + diskann_params->set_pq_chunk_num(pq_chunk_num); + return ZVEC_OK; +} + +/** + * @brief Get DiskANN max_degree parameter + * @param params Index parameters (must be DiskANN type) + * @return max_degree parameter value, or 0 on error + */ +int zvec_index_params_get_diskann_max_degree( + const zvec_index_params_t *params) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not DiskANN index type"); + return 0; + } + auto *cpp_params = reinterpret_cast(params); + auto *diskann_params = + dynamic_cast(cpp_params); + if (!diskann_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not DiskANN index type"); + return 0; + } + return diskann_params->max_degree(); +} + +/** + * @brief Get DiskANN list_size parameter + * @param params Index parameters (must be DiskANN type) + * @return list_size parameter value, or 0 on error + */ +int zvec_index_params_get_diskann_list_size( + const zvec_index_params_t *params) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not DiskANN index type"); + return 0; + } + auto *cpp_params = reinterpret_cast(params); + auto *diskann_params = + dynamic_cast(cpp_params); + if (!diskann_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not DiskANN index type"); + return 0; + } + return diskann_params->list_size(); +} + +/** + * @brief Get DiskANN pq_chunk_num parameter + * @param params Index parameters (must be DiskANN type) + * @return pq_chunk_num parameter value, or 0 on error + */ +int zvec_index_params_get_diskann_pq_chunk_num( + const zvec_index_params_t *params) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not DiskANN index type"); + return 0; + } + auto *cpp_params = reinterpret_cast(params); + auto *diskann_params = + dynamic_cast(cpp_params); + if (!diskann_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not DiskANN index type"); + return 0; + } + return diskann_params->pq_chunk_num(); +} + +/** + * @brief Set IVF-specific parameters + * @param params Index parameters (must be IVF type) + * @param n_list Number of clusters + * @param n_iters Number of k-means iterations + * @param use_soar Whether to use SOAR optimization + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_set_ivf_params(zvec_index_params_t *params, + int n_list, int n_iters, + bool use_soar) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not IVF index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *ivf_params = dynamic_cast(cpp_params); + if (!ivf_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not IVF index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + ivf_params->set_n_list(n_list); + ivf_params->set_n_iters(n_iters); + ivf_params->set_use_soar(use_soar); + return ZVEC_OK; +} + +/** + * @brief Get IVF parameters + * @param params Index parameters (must be IVF type) + * @param out_n_list Output for n_list (can be NULL) + * @param out_n_iters Output for n_iters (can be NULL) + * @param out_use_soar Output for use_soar (can be NULL) + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_get_ivf_params(const zvec_index_params_t *params, + int *out_n_list, + int *out_n_iters, + bool *out_use_soar) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not IVF index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *ivf_params = dynamic_cast(cpp_params); + if (!ivf_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not IVF index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + if (out_n_list) *out_n_list = ivf_params->n_list(); + if (out_n_iters) *out_n_iters = ivf_params->n_iters(); + if (out_use_soar) *out_use_soar = ivf_params->use_soar(); + return ZVEC_OK; +} + +/** + * @brief Set Invert index parameters + * @param params Index parameters (must be INVERT type) + * @param enable_range_opt Enable range optimization + * @param enable_wildcard Enable wildcard search + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_set_invert_params(zvec_index_params_t *params, + bool enable_range_opt, + bool enable_wildcard) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not INVERT index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *invert_params = dynamic_cast(cpp_params); + if (!invert_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not INVERT index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + invert_params->set_enable_range_optimization(enable_range_opt); + invert_params->set_enable_extended_wildcard(enable_wildcard); + return ZVEC_OK; +} + +/** + * @brief Get Invert index parameters + * @param params Index parameters (must be INVERT type) + * @param out_enable_range_opt Output for enable_range_optimization (can be NULL) + * @param out_enable_wildcard Output for enable_extended_wildcard (can be NULL) + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_index_params_get_invert_params(const zvec_index_params_t *params, + bool *out_enable_range_opt, + bool *out_enable_wildcard) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not INVERT index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *invert_params = + dynamic_cast(cpp_params); + if (!invert_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not INVERT index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + if (out_enable_range_opt) + *out_enable_range_opt = invert_params->enable_range_optimization(); + if (out_enable_wildcard) + *out_enable_wildcard = invert_params->enable_extended_wildcard(); + return ZVEC_OK; +} + +zvec_error_code_t zvec_index_params_set_fts_params( + zvec_index_params_t *params, const char *tokenizer_name, + const zvec_string_array_t *filters, const char *extra_params) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not FTS index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *fts_params = dynamic_cast(cpp_params); + if (!fts_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not FTS index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + if (tokenizer_name) { + fts_params->set_tokenizer_name(std::string(tokenizer_name)); + } + if (filters) { + std::vector filter_vec; + filter_vec.reserve(filters->count); + for (size_t i = 0; i < filters->count; ++i) { + const auto &item = filters->strings[i]; + filter_vec.emplace_back(item.data ? item.data : "", + item.data ? item.length : 0); + } + fts_params->set_filters(std::move(filter_vec)); + } + if (extra_params) { + fts_params->set_extra_params(std::string(extra_params)); + } + return ZVEC_OK; +} + +zvec_error_code_t zvec_index_params_get_fts_params( + const zvec_index_params_t *params, const char **out_tokenizer_name, + zvec_string_array_t **out_filters, const char **out_extra_params) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not FTS index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *fts_params = dynamic_cast(cpp_params); + if (!fts_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not FTS index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + if (out_tokenizer_name) { + *out_tokenizer_name = fts_params->tokenizer_name().c_str(); + } + if (out_extra_params) { + *out_extra_params = fts_params->extra_params().c_str(); + } + if (out_filters) { + const auto &filters = fts_params->filters(); + zvec_string_array_t *arr = zvec_string_array_create(filters.size()); + if (!arr) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate filters string array"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + for (size_t i = 0; i < filters.size(); ++i) { + zvec_string_array_add(arr, i, filters[i].c_str()); + } + *out_filters = arr; + } + return ZVEC_OK; +} + +// ============================================================================= +// FieldSchema management interface implementation +// ============================================================================= + +zvec_field_schema_t *zvec_field_schema_create(const char *name, + zvec_data_type_t data_type, bool nullable, + uint32_t dimension) { + if (!name) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Field name cannot be null"); + return nullptr; + } + + ZVEC_TRY_RETURN_NULL( + "Failed to create field schema", + auto cpp_schema = new zvec::FieldSchema( + std::string(name), static_cast(data_type), dimension, + nullable); + + // Return as opaque pointer (raw pointer) + return reinterpret_cast(cpp_schema);) + + return nullptr; +} + +void zvec_field_schema_destroy(zvec_field_schema_t *schema) { + if (schema) { + delete reinterpret_cast(schema); + } +} + +const char *zvec_field_schema_get_name(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return nullptr; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->name().c_str(); +} + +zvec_error_code_t zvec_field_schema_set_name(zvec_field_schema_t *schema, + const char *name) { + if (!schema || !name) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema and name cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_schema = reinterpret_cast(schema); + cpp_schema->set_name(std::string(name)); + return ZVEC_OK; +} + +zvec_data_type_t zvec_field_schema_get_data_type(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return ZVEC_DATA_TYPE_UNDEFINED; + } + auto *cpp_schema = reinterpret_cast(schema); + return static_cast( + static_cast(cpp_schema->data_type())); +} + +zvec_error_code_t zvec_field_schema_set_data_type(zvec_field_schema_t *schema, + zvec_data_type_t data_type) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_schema = reinterpret_cast(schema); + cpp_schema->set_data_type(static_cast(data_type)); + return ZVEC_OK; +} + +zvec_data_type_t zvec_field_schema_get_element_data_type( + const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return ZVEC_DATA_TYPE_UNDEFINED; + } + auto *cpp_schema = reinterpret_cast(schema); + return static_cast( + static_cast(cpp_schema->element_data_type())); +} + +size_t zvec_field_schema_get_element_data_size(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return 0; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->element_data_size(); +} + +bool zvec_field_schema_is_vector_field(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return false; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->is_vector_field(); +} + +bool zvec_field_schema_is_dense_vector(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return false; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->is_dense_vector(); +} + +bool zvec_field_schema_is_sparse_vector(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return false; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->is_sparse_vector(); +} + +bool zvec_field_schema_is_nullable(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return false; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->nullable(); +} + +zvec_error_code_t zvec_field_schema_set_nullable(zvec_field_schema_t *schema, + bool nullable) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_schema = reinterpret_cast(schema); + cpp_schema->set_nullable(nullable); + return ZVEC_OK; +} + +bool zvec_field_schema_has_invert_index(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return false; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->has_invert_index(); +} + +bool zvec_field_schema_is_array_type(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return false; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->is_array_type(); +} + +uint32_t zvec_field_schema_get_dimension(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return 0; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->dimension(); +} + +zvec_error_code_t zvec_field_schema_set_dimension(zvec_field_schema_t *schema, + uint32_t dimension) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_schema = reinterpret_cast(schema); + cpp_schema->set_dimension(dimension); + return ZVEC_OK; +} + +zvec_index_type_t zvec_field_schema_get_index_type(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return ZVEC_INDEX_TYPE_UNDEFINED; + } + auto *cpp_schema = reinterpret_cast(schema); + auto cpp_index_params = cpp_schema->index_params(); + if (!cpp_index_params) { + return ZVEC_INDEX_TYPE_UNDEFINED; + } + return static_cast( + static_cast(cpp_index_params->type())); +} + +const zvec_index_params_t *zvec_field_schema_get_index_params( + const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return nullptr; + } + const auto *cpp_schema = reinterpret_cast(schema); + auto cpp_index_params = cpp_schema->index_params(); + if (!cpp_index_params) { + return nullptr; + } + // Return internal pointer directly - caller does not own and should not free + // The pointer is valid as long as the schema is not modified or destroyed + return reinterpret_cast(cpp_index_params.get()); +} + +zvec_error_code_t zvec_field_schema_set_index_params( + zvec_field_schema_t *schema, const zvec_index_params_t *index_params) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_schema = reinterpret_cast(schema); + + if (!index_params) { + cpp_schema->set_index_params(nullptr); + return ZVEC_OK; + } + + auto cpp_index_params = convert_c_index_params_to_cpp(index_params); + cpp_schema->set_index_params(cpp_index_params); + return ZVEC_OK; +} + +zvec_error_code_t zvec_field_schema_validate(const zvec_field_schema_t *schema, + zvec_string_t **error_msg) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + if (error_msg) { + *error_msg = nullptr; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to validate field schema", + auto *cpp_schema = reinterpret_cast(schema); + auto status = cpp_schema->validate(); if (!status.ok()) { + if (error_msg) { + *error_msg = zvec_string_create(status.message().c_str()); + } + return status_to_error_code(status); + }) + + return ZVEC_OK; +} + +// Internal helper function (forward declared earlier) +bool zvec_field_schema_has_index(const zvec_field_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Field schema pointer cannot be null"); + return false; + } + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->index_params() != nullptr; +} + +// ============================================================================= +// CollectionSchema management interface implementation +// ============================================================================= + +zvec_collection_schema_t *zvec_collection_schema_create(const char *name) { + if (!name) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection name cannot be null"); + return nullptr; + } + + ZVEC_TRY_RETURN_NULL( + "Failed to create collection schema", + auto cpp_schema = new zvec::CollectionSchema(std::string(name)); + + // Return as opaque pointer (raw pointer) + return reinterpret_cast(cpp_schema);) + + return nullptr; +} + +void zvec_collection_schema_destroy(zvec_collection_schema_t *schema) { + if (schema) { + delete reinterpret_cast(schema); + } +} + +const char *zvec_collection_schema_get_name( + const zvec_collection_schema_t *schema) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema pointer cannot be null"); + return nullptr; + } + auto *cpp_schema = reinterpret_cast(schema); + // Use strdup to create a persistent copy since name() returns by value + return strdup(cpp_schema->name().c_str()); +} + +zvec_error_code_t zvec_collection_schema_set_name(zvec_collection_schema_t *schema, + const char *name) { + if (!schema || !name) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema or name cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to set collection name", + auto *cpp_schema = reinterpret_cast(schema); + cpp_schema->set_name(std::string(name)); + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_schema_add_field(zvec_collection_schema_t *schema, + const zvec_field_schema_t *field) { + if (!schema || !field) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema or field pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to add field", + auto *cpp_schema = reinterpret_cast(schema); + const auto *cpp_field = + reinterpret_cast(field); + + // Clone the field schema + auto cloned_field = std::make_shared(*cpp_field); + auto status = cpp_schema->add_field(cloned_field); + return status_to_error_code(status);) +} + +zvec_error_code_t zvec_collection_schema_alter_field( + zvec_collection_schema_t *schema, const char *field_name, + const zvec_field_schema_t *new_field) { + if (!schema || !field_name || !new_field) { + SET_LAST_ERROR( + ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema, field name, or new field cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to alter field", + auto *cpp_schema = reinterpret_cast(schema); + auto *cpp_new_field = + reinterpret_cast(new_field); + auto cloned_field = std::make_shared(*cpp_new_field); + auto status = + cpp_schema->alter_field(std::string(field_name), cloned_field); + return status_to_error_code(status);) +} + +zvec_error_code_t zvec_collection_schema_drop_field(zvec_collection_schema_t *schema, + const char *field_name) { + if (!schema || !field_name) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema or field name cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to drop field", + auto *cpp_schema = reinterpret_cast(schema); + auto status = cpp_schema->drop_field(std::string(field_name)); + return status_to_error_code(status);) +} + +bool zvec_collection_schema_has_field(const zvec_collection_schema_t *schema, + const char *field_name) { + if (!schema || !field_name) { + return false; + } + + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->has_field(std::string(field_name)); +} + +zvec_field_schema_t *zvec_collection_schema_get_field( + const zvec_collection_schema_t *schema, const char *field_name) { + if (!schema || !field_name) { + return nullptr; + } + + ZVEC_TRY_RETURN_NULL( + "Failed to get field", + auto *cpp_schema = + reinterpret_cast(schema); + const zvec::FieldSchema *cpp_field = + cpp_schema->get_field(std::string(field_name)); + if (!cpp_field) { return nullptr; } + // Return non-owning pointer - caller should NOT free this + return reinterpret_cast( + const_cast(cpp_field));) + + return nullptr; +} + +zvec_field_schema_t *zvec_collection_schema_get_forward_field( + const zvec_collection_schema_t *schema, const char *field_name) { + if (!schema || !field_name) { + return nullptr; + } + + ZVEC_TRY_RETURN_NULL( + "Failed to get forward field", + auto *cpp_schema = + reinterpret_cast(schema); + const zvec::FieldSchema *cpp_field = + cpp_schema->get_forward_field(std::string(field_name)); + if (!cpp_field) { return nullptr; } + // Return non-owning pointer - caller should NOT free this + return reinterpret_cast( + const_cast(cpp_field));) + + return nullptr; +} + +zvec_field_schema_t *zvec_collection_schema_get_vector_field( + const zvec_collection_schema_t *schema, const char *field_name) { + if (!schema || !field_name) { + return nullptr; + } + + ZVEC_TRY_RETURN_NULL( + "Failed to get vector field", + auto *cpp_schema = + reinterpret_cast(schema); + const zvec::FieldSchema *cpp_field = + cpp_schema->get_vector_field(std::string(field_name)); + if (!cpp_field) { return nullptr; } + // Return non-owning pointer - caller should NOT free this + return reinterpret_cast( + const_cast(cpp_field));) + + return nullptr; +} + +zvec_error_code_t zvec_collection_schema_get_forward_fields( + const zvec_collection_schema_t *schema, zvec_field_schema_t ***fields, + size_t *count) { + if (!schema || !fields || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Schema, fields, and count cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get forward fields", + auto *cpp_schema = + reinterpret_cast(schema); + auto forward_fields = cpp_schema->forward_fields(); + + *count = forward_fields.size(); + *fields = (zvec_field_schema_t **)zvec_malloc(*count * sizeof(zvec_field_schema_t *)); + if (!*fields) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + + for (size_t i = 0; i < *count; ++i) { + // Return non-owning pointers - caller should NOT free these + (*fields)[i] = + reinterpret_cast(forward_fields[i].get()); + } + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_schema_get_forward_fields_with_index( + const zvec_collection_schema_t *schema, zvec_field_schema_t ***fields, + size_t *count) { + if (!schema || !fields || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Schema, fields, and count cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get forward fields with index", + auto *cpp_schema = + reinterpret_cast(schema); + auto fields_with_index = cpp_schema->forward_fields_with_index(); + + *count = fields_with_index.size(); + *fields = (zvec_field_schema_t **)zvec_malloc(*count * sizeof(zvec_field_schema_t *)); + if (!*fields) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + + for (size_t i = 0; i < *count; ++i) { + // Return non-owning pointers - caller should NOT free these + (*fields)[i] = + reinterpret_cast(fields_with_index[i].get()); + } + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_schema_get_forward_field_names( + const zvec_collection_schema_t *schema, const char ***names, size_t *count) { + if (!schema || !names || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Schema, names, and count cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get forward field names", + auto *cpp_schema = + reinterpret_cast(schema); + auto forward_names = cpp_schema->forward_field_names(); + + *count = forward_names.size(); + *names = (const char **)malloc(*count * sizeof(const char *)); + if (!*names) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + + // Copy strings - caller owns the memory and should free + for (size_t i = 0; i < *count; ++i) { + (*names)[i] = strdup(forward_names[i].c_str()); + } + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_schema_get_forward_field_names_with_index( + const zvec_collection_schema_t *schema, const char ***names, size_t *count) { + if (!schema || !names || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Schema, names, and count cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get forward field names with index", + auto *cpp_schema = + reinterpret_cast(schema); + auto forward_names_with_index = + cpp_schema->forward_field_names_with_index(); + + *count = forward_names_with_index.size(); + *names = (const char **)malloc(*count * sizeof(const char *)); + if (!*names) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + + // Copy strings - caller owns the memory and should free + for (size_t i = 0; i < *count; ++i) { + (*names)[i] = strdup(forward_names_with_index[i].c_str()); + } + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_schema_get_all_field_names( + const zvec_collection_schema_t *schema, const char ***names, size_t *count) { + if (!schema || !names || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Schema, names, and count cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get all field names", + auto *cpp_schema = + reinterpret_cast(schema); + auto all_names = cpp_schema->all_field_names(); + + *count = all_names.size(); + *names = (const char **)malloc(*count * sizeof(const char *)); + if (!*names) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + + // Copy strings - caller owns the memory and should free + for (size_t i = 0; i < *count; + ++i) { (*names)[i] = strdup(all_names[i].c_str()); + } + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_schema_get_vector_fields( + const zvec_collection_schema_t *schema, zvec_field_schema_t ***fields, + size_t *count) { + if (!schema || !fields || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Schema, fields, and count cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get vector fields", + auto *cpp_schema = + reinterpret_cast(schema); + auto vector_fields = cpp_schema->vector_fields(); + + *count = vector_fields.size(); + *fields = (zvec_field_schema_t **)malloc(*count * sizeof(zvec_field_schema_t *)); + if (!*fields) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + + for (size_t i = 0; i < *count; ++i) { + // Return non-owning pointers - caller should NOT free these + (*fields)[i] = + reinterpret_cast(vector_fields[i].get()); + } + return ZVEC_OK;) +} + +uint64_t zvec_collection_schema_get_max_doc_count_per_segment( + const zvec_collection_schema_t *schema) { + if (!schema) return 0; + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->max_doc_count_per_segment(); +} + +zvec_error_code_t zvec_collection_schema_set_max_doc_count_per_segment( + zvec_collection_schema_t *schema, uint64_t max_doc_count) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to set max doc count per segment", + auto *cpp_schema = reinterpret_cast(schema); + cpp_schema->set_max_doc_count_per_segment(max_doc_count); + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_schema_validate( + const zvec_collection_schema_t *schema, zvec_string_t **error_msg) { + if (!schema) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema pointer cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + if (error_msg) { + *error_msg = nullptr; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to validate schema", + auto *cpp_schema = + reinterpret_cast(schema); + auto status = cpp_schema->validate(); if (!status.ok()) { + if (error_msg) { + *error_msg = zvec_string_create(status.message().c_str()); + } + return status_to_error_code(status); + } + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_schema_add_index( + zvec_collection_schema_t *schema, const char *field_name, + const zvec_index_params_t *index_params) { + if (!schema || !field_name || !index_params) { + SET_LAST_ERROR( + ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema, field name, or index params cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to add index", + auto *cpp_schema = reinterpret_cast(schema); + auto cpp_index_params = convert_c_index_params_to_cpp(index_params); + auto status = + cpp_schema->add_index(std::string(field_name), cpp_index_params); + return status_to_error_code(status);) +} + +zvec_error_code_t zvec_collection_schema_drop_index(zvec_collection_schema_t *schema, + const char *field_name) { + if (!schema || !field_name) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Collection schema or field name cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to drop index", + auto *cpp_schema = reinterpret_cast(schema); + // Find the field and clear its index + auto *field = cpp_schema->get_field(std::string(field_name)); + if (!field) { + SET_LAST_ERROR(ZVEC_ERROR_NOT_FOUND, "Field not found"); + return ZVEC_ERROR_NOT_FOUND; + } + const_cast(field) + ->set_index_params(nullptr); + return ZVEC_OK;) +} + +bool zvec_collection_schema_has_index(const zvec_collection_schema_t *schema, + const char *field_name) { + if (!schema || !field_name) { + return false; + } + + auto *cpp_schema = reinterpret_cast(schema); + return cpp_schema->has_index(std::string(field_name)); +} + +// ============================================================================= +// Helper functions +// ============================================================================= + +const char *zvec_error_code_to_string(zvec_error_code_t error_code) { + switch (error_code) { + case ZVEC_OK: + return "OK"; + case ZVEC_ERROR_NOT_FOUND: + return "NOT_FOUND"; + case ZVEC_ERROR_ALREADY_EXISTS: + return "ALREADY_EXISTS"; + case ZVEC_ERROR_INVALID_ARGUMENT: + return "INVALID_ARGUMENT"; + case ZVEC_ERROR_PERMISSION_DENIED: + return "PERMISSION_DENIED"; + case ZVEC_ERROR_FAILED_PRECONDITION: + return "FAILED_PRECONDITION"; + case ZVEC_ERROR_RESOURCE_EXHAUSTED: + return "RESOURCE_EXHAUSTED"; + case ZVEC_ERROR_UNAVAILABLE: + return "UNAVAILABLE"; + case ZVEC_ERROR_INTERNAL_ERROR: + return "INTERNAL_ERROR"; + case ZVEC_ERROR_NOT_SUPPORTED: + return "NOT_SUPPORTED"; + case ZVEC_ERROR_UNKNOWN: + return "UNKNOWN"; + default: + return "UNKNOWN_ERROR_CODE"; + } +} + +const char *zvec_data_type_to_string(zvec_data_type_t data_type) { + switch (data_type) { + case ZVEC_DATA_TYPE_UNDEFINED: + return "UNDEFINED"; + case ZVEC_DATA_TYPE_BINARY: + return "BINARY"; + case ZVEC_DATA_TYPE_STRING: + return "STRING"; + case ZVEC_DATA_TYPE_BOOL: + return "BOOL"; + case ZVEC_DATA_TYPE_INT32: + return "INT32"; + case ZVEC_DATA_TYPE_INT64: + return "INT64"; + case ZVEC_DATA_TYPE_UINT32: + return "UINT32"; + case ZVEC_DATA_TYPE_UINT64: + return "UINT64"; + case ZVEC_DATA_TYPE_FLOAT: + return "FLOAT"; + case ZVEC_DATA_TYPE_DOUBLE: + return "DOUBLE"; + case ZVEC_DATA_TYPE_VECTOR_BINARY32: + return "VECTOR_BINARY32"; + case ZVEC_DATA_TYPE_VECTOR_BINARY64: + return "VECTOR_BINARY64"; + case ZVEC_DATA_TYPE_VECTOR_FP16: + return "VECTOR_FP16"; + case ZVEC_DATA_TYPE_VECTOR_FP32: + return "VECTOR_FP32"; + case ZVEC_DATA_TYPE_VECTOR_FP64: + return "VECTOR_FP64"; + case ZVEC_DATA_TYPE_VECTOR_INT4: + return "VECTOR_INT4"; + case ZVEC_DATA_TYPE_VECTOR_INT8: + return "VECTOR_INT8"; + case ZVEC_DATA_TYPE_VECTOR_INT16: + return "VECTOR_INT16"; + case ZVEC_DATA_TYPE_SPARSE_VECTOR_FP16: + return "SPARSE_VECTOR_FP16"; + case ZVEC_DATA_TYPE_SPARSE_VECTOR_FP32: + return "SPARSE_VECTOR_FP32"; + case ZVEC_DATA_TYPE_ARRAY_BINARY: + return "ARRAY_BINARY"; + case ZVEC_DATA_TYPE_ARRAY_STRING: + return "ARRAY_STRING"; + case ZVEC_DATA_TYPE_ARRAY_BOOL: + return "ARRAY_BOOL"; + case ZVEC_DATA_TYPE_ARRAY_INT32: + return "ARRAY_INT32"; + case ZVEC_DATA_TYPE_ARRAY_INT64: + return "ARRAY_INT64"; + case ZVEC_DATA_TYPE_ARRAY_UINT32: + return "ARRAY_UINT32"; + case ZVEC_DATA_TYPE_ARRAY_UINT64: + return "ARRAY_UINT64"; + case ZVEC_DATA_TYPE_ARRAY_FLOAT: + return "ARRAY_FLOAT"; + case ZVEC_DATA_TYPE_ARRAY_DOUBLE: + return "ARRAY_DOUBLE"; + default: + return "UNKNOWN_DATA_TYPE"; + } +} + +const char *zvec_index_type_to_string(zvec_index_type_t index_type) { + switch (index_type) { + case ZVEC_INDEX_TYPE_UNDEFINED: + return "UNDEFINED"; + case ZVEC_INDEX_TYPE_HNSW: + return "HNSW"; + case ZVEC_INDEX_TYPE_IVF: + return "IVF"; + case ZVEC_INDEX_TYPE_FLAT: + return "FLAT"; + case ZVEC_INDEX_TYPE_INVERT: + return "INVERT"; + case ZVEC_INDEX_TYPE_FTS: + return "FTS"; + case ZVEC_INDEX_TYPE_DISKANN: + return "DiskANN"; + default: + return "UNKNOWN_INDEX_TYPE"; + } +} + +const char *zvec_metric_type_to_string(zvec_metric_type_t metric_type) { + switch (metric_type) { + case ZVEC_METRIC_TYPE_UNDEFINED: + return "UNDEFINED"; + case ZVEC_METRIC_TYPE_L2: + return "L2"; + case ZVEC_METRIC_TYPE_IP: + return "IP"; + case ZVEC_METRIC_TYPE_COSINE: + return "COSINE"; + case ZVEC_METRIC_TYPE_MIPSL2: + return "MIPSL2"; + default: + return "UNKNOWN_METRIC_TYPE"; + } +} + +// ============================================================================= +// Doc functions implementation +// ============================================================================= + +zvec_doc_t *zvec_doc_create(void) { + ZVEC_TRY_RETURN_NULL("Failed to create document", { + auto *doc_ptr = new zvec::Doc(); + return reinterpret_cast(doc_ptr); + }) +} + +void zvec_doc_destroy(zvec_doc_t *doc) { + if (doc) { + delete reinterpret_cast(doc); + } +} + +void zvec_doc_clear(zvec_doc_t *doc) { + if (!doc) return; + + ZVEC_TRY_BEGIN_VOID + auto *doc_ptr = reinterpret_cast(doc); + doc_ptr->clear(); + ZVEC_CATCH_END_VOID +} + +void zvec_docs_free(zvec_doc_t **docs, size_t count) { + if (!docs) return; + + for (size_t i = 0; i < count; ++i) { + zvec_doc_destroy(docs[i]); + } + + free(docs); +} + +void zvec_write_results_free(zvec_write_result_t *results, size_t result_count) { + free_write_results_internal(results, result_count); +} + +void zvec_doc_set_pk(zvec_doc_t *doc, const char *pk) { + if (!doc || !pk) return; + + ZVEC_TRY_BEGIN_VOID + auto *doc_ptr = reinterpret_cast(doc); + doc_ptr->set_pk(std::string(pk)); + ZVEC_CATCH_END_VOID +} + +void zvec_doc_set_doc_id(zvec_doc_t *doc, uint64_t doc_id) { + if (!doc) return; + + ZVEC_TRY_BEGIN_VOID + auto *doc_ptr = reinterpret_cast(doc); + doc_ptr->set_doc_id(doc_id); + ZVEC_CATCH_END_VOID +} + +void zvec_doc_set_score(zvec_doc_t *doc, float score) { + if (!doc) return; + + ZVEC_TRY_BEGIN_VOID + auto *doc_ptr = reinterpret_cast(doc); + doc_ptr->set_score(score); + ZVEC_CATCH_END_VOID +} + +void zvec_doc_set_operator(zvec_doc_t *doc, zvec_doc_operator_t op) { + if (!doc) return; + + ZVEC_TRY_BEGIN_VOID + auto *doc_ptr = reinterpret_cast(doc); + doc_ptr->set_operator(static_cast(op)); + ZVEC_CATCH_END_VOID +} + +zvec_error_code_t zvec_doc_set_field_null(zvec_doc_t *doc, const char *field_name) { + if (!doc || !field_name) { + set_last_error("Invalid arguments: null pointer"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR("Failed to set null field", + auto *doc_ptr = reinterpret_cast(doc); + doc_ptr->set_null(std::string(field_name)); + return ZVEC_OK;) +} + +// ============================================================================= +// Document interface implementation +// ============================================================================= + +// Helper function to extract scalar values from raw data +template +T extract_scalar_value(const void *value, size_t value_size, + zvec_error_code_t *error_code) { + if (value_size != sizeof(T)) { + if (error_code) { + *error_code = ZVEC_ERROR_INVALID_ARGUMENT; + } + return T{}; + } + return *static_cast(value); +} + +// Helper function to extract vector values from raw data +template +std::vector extract_vector_values(const void *value, size_t value_size, + zvec_error_code_t *error_code) { + if (value_size % sizeof(T) != 0) { + if (error_code) { + *error_code = ZVEC_ERROR_INVALID_ARGUMENT; + } + return std::vector(); + } + size_t count = value_size / sizeof(T); + const T *vals = static_cast(value); + return std::vector(vals, vals + count); +} + +// Helper function to extract array values from raw data +template +std::vector extract_array_values(const void *value, size_t value_size, + zvec_error_code_t *error_code) { + if (value_size % sizeof(T) != 0) { + if (error_code) { + *error_code = ZVEC_ERROR_INVALID_ARGUMENT; + } + return std::vector(); + } + size_t count = value_size / sizeof(T); + const T *vals = static_cast(value); + return std::vector(vals, vals + count); +} + +// Helper function to handle sparse vector extraction +template +std::pair, std::vector> extract_sparse_vector( + const void *value, size_t value_size, zvec_error_code_t *error_code) { + if (value_size < sizeof(uint32_t)) { + if (error_code) { + *error_code = ZVEC_ERROR_INVALID_ARGUMENT; + } + return std::make_pair(std::vector(), std::vector()); + } + + const uint32_t *data = static_cast(value); + uint32_t nnz = data[0]; + + size_t required_size = + sizeof(uint32_t) + nnz * (sizeof(uint32_t) + sizeof(T)); + if (value_size < required_size) { + if (error_code) { + *error_code = ZVEC_ERROR_INVALID_ARGUMENT; + } + return std::make_pair(std::vector(), std::vector()); + } + + const uint32_t *indices = data + 1; + const T *values = reinterpret_cast(indices + nnz); + + std::vector index_vec(indices, indices + nnz); + std::vector value_vec(values, values + nnz); + + return std::make_pair(std::move(index_vec), std::move(value_vec)); +} + +// Helper function to extract string array from raw data (C-string array) +std::vector extract_string_array(const void *value, + size_t value_size) { + std::vector string_array; + const char *data = static_cast(value); + size_t pos = 0; + + while (pos < value_size) { + size_t str_len = strlen(data + pos); + if (pos + str_len >= value_size) { + break; + } + string_array.emplace_back(data + pos, str_len); + pos += str_len + 1; + } + return string_array; +} + +// Helper function to extract string array from zvec_string_t** array +std::vector extract_string_array_from_zvec( + zvec_string_t **zvec_strings, size_t count) { + std::vector string_array; + string_array.reserve(count); + + for (size_t i = 0; i < count; ++i) { + if (zvec_strings[i] && zvec_strings[i]->data) { + string_array.emplace_back(zvec_strings[i]->data, zvec_strings[i]->length); + } else { + string_array.emplace_back("", 0); + } + } + + return string_array; +} + +// Helper function to extract binary array from raw data +std::vector extract_binary_array(const void *value, + size_t value_size) { + std::vector binary_array; + const char *data = static_cast(value); + size_t pos = 0; + + while (pos < value_size) { + if (pos + sizeof(uint32_t) > value_size) { + break; + } + uint32_t bin_len = *reinterpret_cast(data + pos); + pos += sizeof(uint32_t); + + if (pos + bin_len > value_size) { + break; + } + binary_array.emplace_back(data + pos, bin_len); + pos += bin_len; + } + return binary_array; +} + +static std::vector convert_zvec_docs_to_internal( + const zvec_doc_t **zvec_docs, size_t doc_count) { + std::vector docs; + docs.reserve(doc_count); + + for (size_t i = 0; i < doc_count; ++i) { + auto *doc_ptr = reinterpret_cast(zvec_docs[i]); + // Use copy constructor to create a deep copy + docs.push_back(*doc_ptr); + } + + return docs; +} + +static zvec::Status convert_zvec_collection_schema_to_internal( + const zvec_collection_schema_t *schema, + zvec::CollectionSchema::Ptr &collection_schema) { + // Get the underlying C++ CollectionSchema + auto *cpp_schema = reinterpret_cast(schema); + + // Create a copy of the C++ schema as shared_ptr + collection_schema = std::make_shared(*cpp_schema); + + return zvec::Status::OK(); +} + +zvec_error_code_t zvec_doc_add_field_by_value(zvec_doc_t *doc, const char *field_name, + zvec_data_type_t data_type, + const void *value, + size_t value_size) { + if (!doc || !field_name || !value) { + set_last_error("Invalid arguments: null pointer"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to add field", auto *doc_ptr = reinterpret_cast(doc); + std::string name(field_name); zvec_error_code_t error_code = ZVEC_OK; + + switch (data_type) { + // Scalar types + case ZVEC_DATA_TYPE_BINARY: + case ZVEC_DATA_TYPE_STRING: { + std::string val(static_cast(value), value_size); + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_BOOL: { + bool val = extract_scalar_value(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for bool type"); + return error_code; + } + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_INT32: { + int32_t val = + extract_scalar_value(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for int32 type"); + return error_code; + } + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_INT64: { + int64_t val = + extract_scalar_value(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for int64 type"); + return error_code; + } + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_UINT32: { + uint32_t val = + extract_scalar_value(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for uint32 type"); + return error_code; + } + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_UINT64: { + uint64_t val = + extract_scalar_value(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for uint64 type"); + return error_code; + } + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_FLOAT: { + float val = + extract_scalar_value(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for float type"); + return error_code; + } + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_DOUBLE: { + double val = + extract_scalar_value(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for double type"); + return error_code; + } + doc_ptr->set(name, val); + break; + } + + // Vector types + case ZVEC_DATA_TYPE_VECTOR_BINARY32: { + auto vec = + extract_vector_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for vector_binary32 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_BINARY64: { + auto vec = + extract_vector_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for vector_binary64 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP32: { + auto vec = + extract_vector_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for vector_fp32 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP16: { + auto vec = extract_vector_values(value, value_size, + &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for vector_fp16 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP64: { + auto vec = + extract_vector_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for vector_fp64 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT8: { + auto vec = + extract_vector_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for vector_int8 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT16: { + auto vec = + extract_vector_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for vector_int16 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT4: { + // INT4 vectors are packed - each byte contains 2 int4 values + size_t count = value_size * 2; + const int8_t *packed_vals = static_cast(value); + std::vector vec; + vec.reserve(count); + + // Unpack int4 values + for (size_t i = 0; i < value_size; ++i) { + int8_t byte_val = packed_vals[i]; + // Extract lower 4 bits + vec.push_back(byte_val & 0x0F); + // Extract upper 4 bits + vec.push_back((byte_val >> 4) & 0x0F); + } + doc_ptr->set(name, std::move(vec)); + break; + } + + // Sparse vector types + case ZVEC_DATA_TYPE_SPARSE_VECTOR_FP16: { + auto sparse_vec = extract_sparse_vector( + value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid sparse vector data size"); + return error_code; + } + doc_ptr->set(name, std::move(sparse_vec)); + break; + } + case ZVEC_DATA_TYPE_SPARSE_VECTOR_FP32: { + auto sparse_vec = + extract_sparse_vector(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid sparse vector data size"); + return error_code; + } + doc_ptr->set(name, std::move(sparse_vec)); + break; + } + + // Array types + case ZVEC_DATA_TYPE_ARRAY_BINARY: { + auto binary_array = extract_binary_array(value, value_size); + doc_ptr->set(name, std::move(binary_array)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_STRING: { + // Check if this is a zvec_string_t** array or a C-string array + // zvec_string_t** array has pointer-sized elements + constexpr size_t ptr_size = sizeof(void *); + if (value_size % ptr_size == 0) { + // Likely a zvec_string_t** array + size_t count = value_size / ptr_size; + zvec_string_t **zvec_str_array = + reinterpret_cast(const_cast(value)); + auto string_array = + extract_string_array_from_zvec(zvec_str_array, count); + doc_ptr->set(name, std::move(string_array)); + } else { + // C-string array (null-terminated strings) + auto string_array = extract_string_array(value, value_size); + doc_ptr->set(name, std::move(string_array)); + } + break; + } + case ZVEC_DATA_TYPE_ARRAY_BOOL: { + auto vec = extract_array_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for array_bool type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_INT32: { + auto vec = + extract_array_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for array_int32 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_INT64: { + auto vec = + extract_array_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for array_int64 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_UINT32: { + auto vec = + extract_array_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for array_uint32 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_UINT64: { + auto vec = + extract_array_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for array_uint64 type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_FLOAT: { + auto vec = + extract_array_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for array_float type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_DOUBLE: { + auto vec = + extract_array_values(value, value_size, &error_code); + if (error_code != ZVEC_OK) { + set_last_error("Invalid value size for array_double type"); + return error_code; + } + doc_ptr->set(name, std::move(vec)); + break; + } + + default: + set_last_error("Unsupported data type: " + std::to_string(data_type)); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + return ZVEC_OK;) +} + +zvec_error_code_t zvec_doc_add_field_by_struct(zvec_doc_t *doc, + const zvec_doc_field_t *field) { + if (!doc || !field) { + set_last_error("Invalid arguments: null pointer"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to add field", auto *doc_ptr = reinterpret_cast(doc); + + std::string name(field->name.data, field->name.length); + + switch (field->data_type) { + // Scalar types (in zvec_data_type_t order: BINARY, STRING, BOOL, INT32, + // INT64, UINT32, UINT64, FLOAT, DOUBLE) + case ZVEC_DATA_TYPE_BINARY: { + std::string val( + reinterpret_cast(field->value.binary_value.data), + field->value.binary_value.length); + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_STRING: { + std::string val(field->value.string_value.data, + field->value.string_value.length); + doc_ptr->set(name, val); + break; + } + case ZVEC_DATA_TYPE_BOOL: { + doc_ptr->set(name, field->value.bool_value); + break; + } + case ZVEC_DATA_TYPE_INT32: { + doc_ptr->set(name, field->value.int32_value); + break; + } + case ZVEC_DATA_TYPE_INT64: { + doc_ptr->set(name, field->value.int64_value); + break; + } + case ZVEC_DATA_TYPE_UINT32: { + doc_ptr->set(name, field->value.uint32_value); + break; + } + case ZVEC_DATA_TYPE_UINT64: { + doc_ptr->set(name, field->value.uint64_value); + break; + } + case ZVEC_DATA_TYPE_FLOAT: { + doc_ptr->set(name, field->value.float_value); + break; + } + case ZVEC_DATA_TYPE_DOUBLE: { + doc_ptr->set(name, field->value.double_value); + break; + } + + // Vector types (in zvec_data_type_t order: BINARY32, BINARY64, FP16, FP32, + // FP64, INT4, INT8, INT16) + case ZVEC_DATA_TYPE_VECTOR_BINARY32: { + std::vector vec(reinterpret_cast( + field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_BINARY64: { + std::vector vec(reinterpret_cast( + field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP16: { + std::vector vec( + reinterpret_cast( + field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP32: { + std::vector vec(field->value.vector_value.data, + field->value.vector_value.data + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP64: { + std::vector vec( + reinterpret_cast(field->value.vector_value.data), + reinterpret_cast(field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT4: { + size_t byte_count = (field->value.vector_value.length + 1) / 2; + const int8_t *packed_data = + reinterpret_cast(field->value.vector_value.data); + std::vector vec; + vec.reserve(field->value.vector_value.length); + + for (size_t i = 0; + i < byte_count && vec.size() < field->value.vector_value.length; + ++i) { + int8_t byte_val = packed_data[i]; + // Extract lower 4 bits + vec.push_back(byte_val & 0x0F); + // Extract upper 4 bits + if (vec.size() < field->value.vector_value.length) { + vec.push_back((byte_val >> 4) & 0x0F); + } + } + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT8: { + std::vector vec( + reinterpret_cast(field->value.vector_value.data), + reinterpret_cast(field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT16: { + std::vector vec( + reinterpret_cast(field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + + // Sparse vector types (in zvec_data_type_t order: FP16, FP32) + case ZVEC_DATA_TYPE_SPARSE_VECTOR_FP16: { + std::vector vec( + reinterpret_cast( + field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + case ZVEC_DATA_TYPE_SPARSE_VECTOR_FP32: { + std::vector vec(field->value.vector_value.data, + field->value.vector_value.data + + field->value.vector_value.length); + doc_ptr->set(name, std::move(vec)); + break; + } + + // Array types (in zvec_data_type_t order: BINARY, STRING, BOOL, INT32, + // INT64, UINT32, UINT64, FLOAT, DOUBLE) + case ZVEC_DATA_TYPE_ARRAY_BINARY: { + std::vector array_values; + const uint8_t *data_ptr = field->value.binary_value.data; + size_t total_length = field->value.binary_value.length; + size_t offset = 0; + + while (offset + sizeof(uint32_t) <= total_length) { + uint32_t elem_length = + *reinterpret_cast(data_ptr + offset); + offset += sizeof(uint32_t); + + if (offset + elem_length <= total_length) { + std::string elem( + reinterpret_cast(data_ptr + offset), + elem_length); + array_values.push_back(elem); + offset += elem_length; + } else { + break; + } + } + doc_ptr->set(name, std::move(array_values)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_STRING: { + std::vector array_values; + const char *data_ptr = field->value.string_value.data; + size_t total_length = field->value.string_value.length; + size_t offset = 0; + + while (offset < total_length) { + size_t str_len = strlen(data_ptr + offset); + if (str_len > 0 && offset + str_len <= total_length) { + array_values.emplace_back(data_ptr + offset, str_len); + offset += str_len + 1; + } else { + break; + } + } + doc_ptr->set(name, std::move(array_values)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_BOOL: { + std::vector array_values( + reinterpret_cast(field->value.binary_value.data), + reinterpret_cast(field->value.binary_value.data) + + field->value.binary_value.length); + doc_ptr->set(name, std::move(array_values)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_INT32: { + std::vector array_values( + reinterpret_cast(field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(array_values)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_INT64: { + std::vector array_values( + reinterpret_cast(field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(array_values)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_UINT32: { + std::vector array_values( + reinterpret_cast( + field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(array_values)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_UINT64: { + std::vector array_values( + reinterpret_cast( + field->value.vector_value.data), + reinterpret_cast( + field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(array_values)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_FLOAT: { + std::vector array_values(field->value.vector_value.data, + field->value.vector_value.data + + field->value.vector_value.length); + doc_ptr->set(name, std::move(array_values)); + break; + } + case ZVEC_DATA_TYPE_ARRAY_DOUBLE: { + std::vector array_values( + reinterpret_cast(field->value.vector_value.data), + reinterpret_cast(field->value.vector_value.data) + + field->value.vector_value.length); + doc_ptr->set(name, std::move(array_values)); + break; + } + + default: + set_last_error("Unsupported data type: " + + std::to_string(field->data_type)); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + return ZVEC_OK;) +} + +const char *zvec_doc_get_pk_pointer(const zvec_doc_t *doc) { + if (!doc) return nullptr; + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->pk_ref().data(); +} + +const char *zvec_doc_get_pk_copy(const zvec_doc_t *doc) { + if (!doc) return nullptr; + auto doc_ptr = reinterpret_cast(doc); + const std::string &pk = doc_ptr->pk_ref(); + if (pk.empty()) return nullptr; + + char *result = static_cast(malloc(pk.length() + 1)); + strcpy(result, pk.c_str()); + return result; +} + +uint64_t zvec_doc_get_doc_id(const zvec_doc_t *doc) { + if (!doc) return 0; + + ZVEC_TRY_RETURN_SCALAR( + "Failed to get document ID", 0, + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->doc_id();) +} + +float zvec_doc_get_score(const zvec_doc_t *doc) { + if (!doc) return 0.0f; + + ZVEC_TRY_RETURN_SCALAR( + "Failed to get document score", 0.0f, + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->score();) +} + +zvec_doc_operator_t zvec_doc_get_operator(const zvec_doc_t *doc) { + if (!doc) return ZVEC_DOC_OP_INSERT; // default + ZVEC_TRY_RETURN_SCALAR( + "Failed to get document operator", ZVEC_DOC_OP_INSERT, + auto doc_ptr = reinterpret_cast(doc); + zvec::Operator op = doc_ptr->get_operator(); + return static_cast(op);) +} + +size_t zvec_doc_get_field_count(const zvec_doc_t *doc) { + if (!doc) return 0; + + ZVEC_TRY_RETURN_SCALAR( + "Failed to get field count", 0, + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->field_names().size();) +} + +zvec_error_code_t zvec_doc_get_field_value_basic(const zvec_doc_t *doc, + const char *field_name, + zvec_data_type_t field_type, + void *value_buffer, + size_t buffer_size) { + if (!doc || !field_name || !value_buffer) { + set_last_error("Invalid arguments: null pointer"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get field value", + auto doc_ptr = reinterpret_cast(doc); + + // Check if field exists + if (!doc_ptr->has(field_name)) { + set_last_error("Field not found in document"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + // Handle basic data types that return values directly + switch (field_type) { + case ZVEC_DATA_TYPE_BOOL: { + if (buffer_size < sizeof(bool)) { + set_last_error("Buffer too small for bool value"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + const bool val = doc_ptr->get_ref(field_name); + *static_cast(value_buffer) = val; + break; + } + case ZVEC_DATA_TYPE_INT32: { + if (buffer_size < sizeof(int32_t)) { + set_last_error("Buffer too small for int32 value"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + const int32_t val = doc_ptr->get_ref(field_name); + *static_cast(value_buffer) = val; + break; + } + case ZVEC_DATA_TYPE_INT64: { + if (buffer_size < sizeof(int64_t)) { + set_last_error("Buffer too small for int64 value"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + const int64_t val = doc_ptr->get_ref(field_name); + *static_cast(value_buffer) = val; + break; + } + case ZVEC_DATA_TYPE_UINT32: { + if (buffer_size < sizeof(uint32_t)) { + set_last_error("Buffer too small for uint32 value"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + const uint32_t val = doc_ptr->get_ref(field_name); + *static_cast(value_buffer) = val; + break; + } + case ZVEC_DATA_TYPE_UINT64: { + if (buffer_size < sizeof(uint64_t)) { + set_last_error("Buffer too small for uint64 value"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + const uint64_t val = doc_ptr->get_ref(field_name); + *static_cast(value_buffer) = val; + break; + } + case ZVEC_DATA_TYPE_FLOAT: { + if (buffer_size < sizeof(float)) { + set_last_error("Buffer too small for float value"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + const float val = doc_ptr->get_ref(field_name); + *static_cast(value_buffer) = val; + break; + } + case ZVEC_DATA_TYPE_DOUBLE: { + if (buffer_size < sizeof(double)) { + set_last_error("Buffer too small for double value"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + const double val = doc_ptr->get_ref(field_name); + *static_cast(value_buffer) = val; + break; + } + default: { + set_last_error("Data type not supported for basic value return"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + } + + return ZVEC_OK;) +} + +zvec_error_code_t zvec_doc_get_field_value_copy(const zvec_doc_t *doc, + const char *field_name, + zvec_data_type_t field_type, + void **value, size_t *value_size) { + if (!doc || !field_name || !value || !value_size) { + set_last_error("Invalid arguments: null pointer"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get field value copy", + auto doc_ptr = reinterpret_cast(doc); + + // Check if field exists + if (!doc_ptr->has(field_name)) { + set_last_error("Field not found in document"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + // Handle copy-returning data types (allocate new memory) + switch (field_type) { + // Basic types - copy the actual values + case ZVEC_DATA_TYPE_BOOL: { + const bool val = doc_ptr->get_ref(field_name); + void *buffer = malloc(sizeof(bool)); + if (!buffer) { + set_last_error("Memory allocation failed for bool"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + *static_cast(buffer) = val; + *value = buffer; + *value_size = sizeof(bool); + break; + } + case ZVEC_DATA_TYPE_INT32: { + const int32_t val = doc_ptr->get_ref(field_name); + void *buffer = malloc(sizeof(int32_t)); + if (!buffer) { + set_last_error("Memory allocation failed for int32"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + *static_cast(buffer) = val; + *value = buffer; + *value_size = sizeof(int32_t); + break; + } + case ZVEC_DATA_TYPE_INT64: { + const int64_t val = doc_ptr->get_ref(field_name); + void *buffer = malloc(sizeof(int64_t)); + if (!buffer) { + set_last_error("Memory allocation failed for int64"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + *static_cast(buffer) = val; + *value = buffer; + *value_size = sizeof(int64_t); + break; + } + case ZVEC_DATA_TYPE_UINT32: { + const uint32_t val = doc_ptr->get_ref(field_name); + void *buffer = malloc(sizeof(uint32_t)); + if (!buffer) { + set_last_error("Memory allocation failed for uint32"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + *static_cast(buffer) = val; + *value = buffer; + *value_size = sizeof(uint32_t); + break; + } + case ZVEC_DATA_TYPE_UINT64: { + const uint64_t val = doc_ptr->get_ref(field_name); + void *buffer = malloc(sizeof(uint64_t)); + if (!buffer) { + set_last_error("Memory allocation failed for uint64"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + *static_cast(buffer) = val; + *value = buffer; + *value_size = sizeof(uint64_t); + break; + } + case ZVEC_DATA_TYPE_FLOAT: { + const float val = doc_ptr->get_ref(field_name); + void *buffer = malloc(sizeof(float)); + if (!buffer) { + set_last_error("Memory allocation failed for float"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + *static_cast(buffer) = val; + *value = buffer; + *value_size = sizeof(float); + break; + } + case ZVEC_DATA_TYPE_DOUBLE: { + const double val = doc_ptr->get_ref(field_name); + void *buffer = malloc(sizeof(double)); + if (!buffer) { + set_last_error("Memory allocation failed for double"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + *static_cast(buffer) = val; + *value = buffer; + *value_size = sizeof(double); + break; + } + + // String and binary types - copy the data + case ZVEC_DATA_TYPE_BINARY: + case ZVEC_DATA_TYPE_STRING: { + const std::string &val = doc_ptr->get_ref(field_name); + void *buffer = malloc(val.length()); + if (!buffer) { + set_last_error("Memory allocation failed for string/binary"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + memcpy(buffer, val.data(), val.length()); + *value = buffer; + *value_size = val.length(); + break; + } + + // Vector types - copy the data + case ZVEC_DATA_TYPE_VECTOR_BINARY32: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + size_t total_size = val.size() * sizeof(uint32_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for uint32 vector"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + memcpy(buffer, val.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_VECTOR_BINARY64: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + size_t total_size = val.size() * sizeof(uint64_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for uint64 vector"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + memcpy(buffer, val.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP16: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + size_t total_size = val.size() * sizeof(zvec::float16_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for fp16 vector"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + memcpy(buffer, val.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP32: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + size_t total_size = val.size() * sizeof(float); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for fp32 vector"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + memcpy(buffer, val.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP64: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + size_t total_size = val.size() * sizeof(double); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for fp64 vector"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + memcpy(buffer, val.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT4: + case ZVEC_DATA_TYPE_VECTOR_INT8: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + size_t total_size = val.size() * sizeof(int8_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for int8 vector"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + memcpy(buffer, val.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT16: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + size_t total_size = val.size() * sizeof(int16_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for int16 vector"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + memcpy(buffer, val.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + + // Sparse vector types - create flattened representation + case ZVEC_DATA_TYPE_SPARSE_VECTOR_FP16: { + using SparseVecFP16 = + std::pair, std::vector>; + const SparseVecFP16 &sparse_vec = + doc_ptr->get_ref(field_name); + size_t nnz = sparse_vec.first.size(); + size_t total_size = sizeof(size_t) + nnz * (sizeof(uint32_t) + + sizeof(zvec::float16_t)); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for sparse vector FP16"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + char *ptr = static_cast(buffer); + *reinterpret_cast(ptr) = nnz; + ptr += sizeof(size_t); + + for (size_t i = 0; i < nnz; ++i) { + *reinterpret_cast(ptr) = sparse_vec.first[i]; + ptr += sizeof(uint32_t); + } + for (size_t i = 0; i < nnz; ++i) { + *reinterpret_cast(ptr) = sparse_vec.second[i]; + ptr += sizeof(zvec::float16_t); + } + + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_SPARSE_VECTOR_FP32: { + using SparseVecFP32 = + std::pair, std::vector>; + const SparseVecFP32 &sparse_vec = + doc_ptr->get_ref(field_name); + size_t nnz = sparse_vec.first.size(); + size_t total_size = + sizeof(size_t) + nnz * (sizeof(uint32_t) + sizeof(float)); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for sparse vector FP32"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + char *ptr = static_cast(buffer); + *reinterpret_cast(ptr) = nnz; + ptr += sizeof(size_t); + + for (size_t i = 0; i < nnz; ++i) { + *reinterpret_cast(ptr) = sparse_vec.first[i]; + ptr += sizeof(uint32_t); + } + for (size_t i = 0; i < nnz; ++i) { + *reinterpret_cast(ptr) = sparse_vec.second[i]; + ptr += sizeof(float); + } + + *value = buffer; + *value_size = total_size; + break; + } + + // Array types - create serialized representations + case ZVEC_DATA_TYPE_ARRAY_BINARY: { + using BinaryArray = std::vector; + const BinaryArray &array_vals = + doc_ptr->get_ref(field_name); + size_t total_size = 0; + for (const auto &bin_val : array_vals) { + total_size += bin_val.length(); + } + + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for binary array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + char *ptr = static_cast(buffer); + for (const auto &bin_val : array_vals) { + memcpy(ptr, bin_val.data(), bin_val.length()); + ptr += bin_val.length(); + } + + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_ARRAY_STRING: { + using StringArray = std::vector; + const StringArray &array_vals = + doc_ptr->get_ref(field_name); + size_t total_size = 0; + for (const auto &str_val : array_vals) { + total_size += str_val.length() + 1; // +1 for null terminator + } + + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for string array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + char *ptr = static_cast(buffer); + for (const auto &str_val : array_vals) { + memcpy(ptr, str_val.c_str(), str_val.length()); + ptr += str_val.length(); + *ptr = '\0'; + ptr++; + } + + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_ARRAY_BOOL: { + using BoolArray = std::vector; + const BoolArray &array_vals = doc_ptr->get_ref(field_name); + size_t byte_count = (array_vals.size() + 7) / 8; + void *buffer = malloc(byte_count); + if (!buffer) { + set_last_error("Memory allocation failed for bool array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + uint8_t *bytes = static_cast(buffer); + memset(bytes, 0, byte_count); + + for (size_t i = 0; i < array_vals.size(); ++i) { + if (array_vals[i]) { + bytes[i / 8] |= (1 << (i % 8)); + } + } + + *value = buffer; + *value_size = byte_count; + break; + } + case ZVEC_DATA_TYPE_ARRAY_INT32: { + using Int32Array = std::vector; + const Int32Array &array_vals = + doc_ptr->get_ref(field_name); + size_t total_size = array_vals.size() * sizeof(int32_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for int32 array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + memcpy(buffer, array_vals.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_ARRAY_INT64: { + using Int64Array = std::vector; + const Int64Array &array_vals = + doc_ptr->get_ref(field_name); + size_t total_size = array_vals.size() * sizeof(int64_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for int64 array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + memcpy(buffer, array_vals.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_ARRAY_UINT32: { + using UInt32Array = std::vector; + const UInt32Array &array_vals = + doc_ptr->get_ref(field_name); + size_t total_size = array_vals.size() * sizeof(uint32_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for uint32 array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + memcpy(buffer, array_vals.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_ARRAY_UINT64: { + using UInt64Array = std::vector; + const UInt64Array &array_vals = + doc_ptr->get_ref(field_name); + size_t total_size = array_vals.size() * sizeof(uint64_t); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for uint64 array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + memcpy(buffer, array_vals.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_ARRAY_FLOAT: { + using FloatArray = std::vector; + const FloatArray &array_vals = + doc_ptr->get_ref(field_name); + size_t total_size = array_vals.size() * sizeof(float); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for float array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + memcpy(buffer, array_vals.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + case ZVEC_DATA_TYPE_ARRAY_DOUBLE: { + using DoubleArray = std::vector; + const DoubleArray &array_vals = + doc_ptr->get_ref(field_name); + size_t total_size = array_vals.size() * sizeof(double); + void *buffer = malloc(total_size); + if (!buffer) { + set_last_error("Memory allocation failed for double array"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + memcpy(buffer, array_vals.data(), total_size); + *value = buffer; + *value_size = total_size; + break; + } + default: { + set_last_error("Unknown data type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + } + + return ZVEC_OK;) +} + +zvec_error_code_t zvec_doc_get_field_value_pointer(const zvec_doc_t *doc, + const char *field_name, + zvec_data_type_t field_type, + const void **value, + size_t *value_size) { + if (!doc || !field_name || !value || !value_size) { + set_last_error("Invalid arguments: null pointer"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get field value pointer", + auto doc_ptr = reinterpret_cast(doc); + + // Check if field exists + if (!doc_ptr->has(field_name)) { + set_last_error("Field not found in document"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + // Get field value based on data type + switch (field_type) { + case ZVEC_DATA_TYPE_BINARY: { + const std::string &val = doc_ptr->get_ref(field_name); + *value = val.data(); + *value_size = val.length(); + break; + } + case ZVEC_DATA_TYPE_STRING: { + const std::string &val = doc_ptr->get_ref(field_name); + *value = val.c_str(); + *value_size = val.length(); + break; + } + case ZVEC_DATA_TYPE_BOOL: { + const bool &val = doc_ptr->get_ref(field_name); + *value = &val; + *value_size = sizeof(bool); + break; + } + case ZVEC_DATA_TYPE_INT32: { + const int32_t &val = doc_ptr->get_ref(field_name); + *value = &val; + *value_size = sizeof(int32_t); + break; + } + case ZVEC_DATA_TYPE_INT64: { + const int64_t &val = doc_ptr->get_ref(field_name); + *value = &val; + *value_size = sizeof(int64_t); + break; + } + case ZVEC_DATA_TYPE_UINT32: { + const uint32_t &val = doc_ptr->get_ref(field_name); + *value = &val; + *value_size = sizeof(uint32_t); + break; + } + case ZVEC_DATA_TYPE_UINT64: { + const uint64_t &val = doc_ptr->get_ref(field_name); + *value = &val; + *value_size = sizeof(uint64_t); + break; + } + case ZVEC_DATA_TYPE_FLOAT: { + const float &val = doc_ptr->get_ref(field_name); + *value = &val; + *value_size = sizeof(float); + break; + } + case ZVEC_DATA_TYPE_DOUBLE: { + const double &val = doc_ptr->get_ref(field_name); + *value = &val; + *value_size = sizeof(double); + break; + } + case ZVEC_DATA_TYPE_VECTOR_BINARY32: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + *value = val.data(); + *value_size = val.size() * sizeof(uint32_t); + break; + } + case ZVEC_DATA_TYPE_VECTOR_BINARY64: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + *value = val.data(); + *value_size = val.size() * sizeof(uint64_t); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP16: { + // FP16 vectors typically stored as uint16_t + const std::vector &val = + doc_ptr->get_ref>(field_name); + *value = val.data(); + *value_size = val.size() * sizeof(zvec::float16_t); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP32: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + *value = val.data(); + *value_size = val.size() * sizeof(float); + break; + } + case ZVEC_DATA_TYPE_VECTOR_FP64: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + *value = val.data(); + *value_size = val.size() * sizeof(double); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT4: { + // INT4 vectors typically stored as int8_t with 2 values per byte + const std::vector &val = + doc_ptr->get_ref>(field_name); + *value = val.data(); + *value_size = val.size() * sizeof(int8_t); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT8: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + *value = val.data(); + *value_size = val.size() * sizeof(int8_t); + break; + } + case ZVEC_DATA_TYPE_VECTOR_INT16: { + const std::vector &val = + doc_ptr->get_ref>(field_name); + *value = val.data(); + *value_size = val.size() * sizeof(int16_t); + break; + } + case ZVEC_DATA_TYPE_ARRAY_INT32: { + auto &array_vals = doc_ptr->get_ref>(field_name); + *value = array_vals.data(); + *value_size = array_vals.size() * sizeof(int32_t); + break; + } + case ZVEC_DATA_TYPE_ARRAY_INT64: { + auto &array_vals = doc_ptr->get_ref>(field_name); + *value = array_vals.data(); + *value_size = array_vals.size() * sizeof(int64_t); + break; + } + case ZVEC_DATA_TYPE_ARRAY_UINT32: { + auto &array_vals = + doc_ptr->get_ref>(field_name); + *value = array_vals.data(); + *value_size = array_vals.size() * sizeof(uint32_t); + break; + } + case ZVEC_DATA_TYPE_ARRAY_UINT64: { + auto &array_vals = + doc_ptr->get_ref>(field_name); + *value = array_vals.data(); + *value_size = array_vals.size() * sizeof(uint64_t); + break; + } + case ZVEC_DATA_TYPE_ARRAY_FLOAT: { + auto &array_vals = doc_ptr->get_ref>(field_name); + *value = array_vals.data(); + *value_size = array_vals.size() * sizeof(float); + break; + } + case ZVEC_DATA_TYPE_ARRAY_DOUBLE: { + auto &array_vals = doc_ptr->get_ref>(field_name); + *value = array_vals.data(); + *value_size = array_vals.size() * sizeof(double); + break; + } + default: { + set_last_error("Unknown data type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + } + + return ZVEC_OK;) +} + +bool zvec_doc_is_empty(const zvec_doc_t *doc) { + if (!doc) { + set_last_error("Document pointer is null"); + return true; + } + + ZVEC_TRY_RETURN_SCALAR( + "Failed to check if document is empty", true, + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->is_empty();) +} + +zvec_error_code_t zvec_doc_remove_field(zvec_doc_t *doc, const char *field_name) { + if (!doc || !field_name) { + set_last_error("Document pointer or field name is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR("Failed to remove field", + auto *doc_ptr = reinterpret_cast(doc); + doc_ptr->remove(std::string(field_name)); + return ZVEC_OK;) +} + + +bool zvec_doc_has_field(const zvec_doc_t *doc, const char *field_name) { + if (!doc || !field_name) { + set_last_error("Document pointer or field name is null"); + return false; + } + + ZVEC_TRY_RETURN_SCALAR( + "Failed to check field existence", false, + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->has(std::string(field_name));) +} + +bool zvec_doc_has_field_value(const zvec_doc_t *doc, const char *field_name) { + if (!doc || !field_name) { + set_last_error("Document pointer or field name is null"); + return false; + } + + ZVEC_TRY_RETURN_SCALAR( + "Failed to check field value existence", false, + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->has_value(std::string(field_name));) +} + +bool zvec_doc_is_field_null(const zvec_doc_t *doc, const char *field_name) { + if (!doc || !field_name) { + set_last_error("Document pointer or field name is null"); + return false; + } + + ZVEC_TRY_RETURN_SCALAR( + "Failed to check if field is null", false, + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->is_null(std::string(field_name));) +} + +zvec_error_code_t zvec_doc_get_field_names(const zvec_doc_t *doc, char ***field_names, + size_t *count) { + if (!doc || !field_names || !count) { + set_last_error("Invalid arguments"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get field names", + auto doc_ptr = reinterpret_cast(doc); + std::vector names = doc_ptr->field_names(); + + *count = names.size(); + if (*count == 0) { + *field_names = nullptr; + return ZVEC_OK; + } + + *field_names = static_cast(malloc(*count * sizeof(char *))); + if (!*field_names) { + set_last_error("Failed to allocate memory for field names"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + for (size_t i = 0; i < *count; ++i) { + (*field_names)[i] = copy_string(names[i]); + if (!(*field_names)[i]) { + for (size_t j = 0; j < i; ++j) { + free((*field_names)[j]); + } + free(*field_names); + *field_names = nullptr; + set_last_error("Failed to copy field name"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + } + + return ZVEC_OK;) +} + +zvec_error_code_t zvec_doc_serialize(const zvec_doc_t *doc, uint8_t **data, + size_t *size) { + if (!doc || !data || !size) { + set_last_error("Invalid arguments"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to serialize document", + auto doc_ptr = reinterpret_cast(doc); + std::vector serialized_data = doc_ptr->serialize(); + + *size = serialized_data.size(); + if (*size == 0) { + *data = nullptr; + return ZVEC_OK; + } + + *data = static_cast(malloc(*size)); + if (!*data) { + set_last_error("Failed to allocate memory for serialized data"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + memcpy(*data, serialized_data.data(), *size); + return ZVEC_OK;) +} + +zvec_error_code_t zvec_doc_deserialize(const uint8_t *data, size_t size, + zvec_doc_t **doc) { + if (!data || !doc || size == 0) { + set_last_error("Invalid arguments"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to deserialize document", + auto deserialized_doc = zvec::Doc::deserialize(data, size); + if (!deserialized_doc) { + set_last_error("Failed to deserialize document"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // Create a new Doc by copying the deserialized content + auto *new_doc = new zvec::Doc(*deserialized_doc); + *doc = reinterpret_cast(new_doc); + return ZVEC_OK;) +} + +void zvec_doc_merge(zvec_doc_t *doc, const zvec_doc_t *other) { + if (!doc || !other) { + set_last_error("Document pointers are null"); + return; + } + + ZVEC_TRY_BEGIN_VOID + auto *doc_ptr = reinterpret_cast(doc); + auto *other_ptr = reinterpret_cast(other); + doc_ptr->merge(*other_ptr); + ZVEC_CATCH_END_VOID +} + +size_t zvec_doc_memory_usage(const zvec_doc_t *doc) { + if (!doc) { + set_last_error("Document pointer is null"); + return 0; + } + + ZVEC_TRY_RETURN_SCALAR( + "Failed to get document memory usage", 0, + auto doc_ptr = reinterpret_cast(doc); + return doc_ptr->memory_usage();) +} + + +zvec_error_code_t zvec_doc_to_detail_string(const zvec_doc_t *doc, char **detail_str) { + if (!doc || !detail_str) { + set_last_error("Invalid arguments"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get document detail string", + auto doc_ptr = reinterpret_cast(doc); + std::string detail = doc_ptr->to_detail_string(); + *detail_str = copy_string(detail); + + if (!*detail_str && !detail.empty()) { + set_last_error("Failed to copy detail string"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + return ZVEC_OK;) +} + +// ============================================================================= +// Collection functions implementation +// ============================================================================= + +zvec_error_code_t zvec_collection_create_and_open( + const char *path, const zvec_collection_schema_t *schema, + const zvec_collection_options_t *options, zvec_collection_t **collection) { + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_create_and_open_with_schema", + if (!path || !schema || !collection) { + set_last_error("Path, schema, or collection cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + std::shared_ptr + schema_ptr = nullptr; + auto status = + convert_zvec_collection_schema_to_internal(schema, schema_ptr); + if (!status.ok()) { + set_last_error(status.message()); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + zvec::CollectionOptions collection_options; + if (options) { + auto *opts = reinterpret_cast(options); + collection_options.enable_mmap_ = opts->enable_mmap_; + collection_options.max_buffer_size_ = opts->max_buffer_size_; + collection_options.read_only_ = opts->read_only_; + } + + auto result = zvec::Collection::CreateAndOpen(path, *schema_ptr, + collection_options); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code == ZVEC_OK) { + *collection = reinterpret_cast( + new std::shared_ptr(std::move(result.value()))); + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_open(const char *path, + const zvec_collection_options_t *options, + zvec_collection_t **collection) { + if (!path || !collection) { + set_last_error("Invalid arguments: path and collection cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", zvec::CollectionOptions collection_options; + if (options) { + auto *opts = reinterpret_cast(options); + collection_options.enable_mmap_ = opts->enable_mmap_; + collection_options.max_buffer_size_ = opts->max_buffer_size_; + collection_options.read_only_ = opts->read_only_; + } + + auto result = zvec::Collection::Open(path, collection_options); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code == ZVEC_OK) { + *collection = reinterpret_cast( + new std::shared_ptr(std::move(result.value()))); + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_close(zvec_collection_t *collection) { + if (!collection) { + set_last_error("Invalid argument: collection cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + delete reinterpret_cast *>(collection); + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_destroy(zvec_collection_t *collection) { + if (!collection) { + set_last_error("Invalid argument: collection cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto &coll = + *reinterpret_cast *>(collection); + zvec::Status status = coll->Destroy(); + if (!status.ok()) { set_last_error(status.message()); } + + return status_to_error_code(status);) +} + +zvec_error_code_t zvec_collection_flush(zvec_collection_t *collection) { + if (!collection) { + set_last_error("Invalid argument: collection cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto &coll = + *reinterpret_cast *>(collection); + zvec::Status status = coll->Flush(); + + if (!status.ok()) { set_last_error(status.message()); } + + return status_to_error_code(status);) +} + +zvec_error_code_t zvec_collection_get_schema(const zvec_collection_t *collection, + zvec_collection_schema_t **schema) { + if (!collection || !schema) { + set_last_error("Invalid arguments: collection and schema cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto &coll = *reinterpret_cast *>( + collection); + auto result = coll->Schema(); + + zvec_error_code_t error_code = handle_expected_result(result); + if (error_code == ZVEC_OK) { + const auto &cpp_schema = result.value(); + + // Create a copy of the schema and return as raw pointer + auto *copied_schema = new zvec::CollectionSchema(cpp_schema); + *schema = reinterpret_cast(copied_schema); + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_get_options(const zvec_collection_t *collection, + zvec_collection_options_t **options) { + if (!collection || !options) { + set_last_error("Invalid arguments"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get collection options", + auto collection_ptr = + reinterpret_cast *>( + collection); + auto result = (*collection_ptr)->Options(); + + if (!result.has_value()) { + set_last_error("Failed to get collection option: " + + result.error().message()); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // Create and initialize options using new + *options = reinterpret_cast( + new zvec::CollectionOptions(result.value())); + + return ZVEC_OK;) +} + +zvec_error_code_t zvec_collection_get_stats(const zvec_collection_t *collection, + zvec_collection_stats_t **stats) { + if (!collection || !stats) { + set_last_error("Invalid arguments"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Failed to get detailed collection stats", + auto collection_ptr = + reinterpret_cast *>( + collection); + auto result = (*collection_ptr)->Stats(); + + if (!result.has_value()) { + set_last_error("Failed to get collection stats: " + + result.error().message()); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // Create a new CollectionStats object and return as opaque pointer + *stats = reinterpret_cast( + new zvec::CollectionStats(result.value())); + + return ZVEC_OK;) +} + +void zvec_collection_stats_destroy(zvec_collection_stats_t *stats) { + if (stats) { + delete reinterpret_cast(stats); + } +} + +// ============================================================================= +// QueryParams implementation +// ============================================================================= +// Users should create type-specific query params: +// - HnswQueryParams via zvec_query_params_hnsw_create() +// - IVFQueryParams via zvec_query_params_ivf_create() +// - FlatQueryParams via zvec_query_params_flat_create() +// +// Each type-specific instance has its own destroy function. +// Common parameters (radius, is_linear, is_using_refiner) are set via the +// type-specific create functions. + +// ============================================================================= +// HnswQueryParams implementation - wrapper around zvec::HnswQueryParams +// ============================================================================= + +zvec_hnsw_query_params_t *zvec_query_params_hnsw_create(int ef, float radius, + bool is_linear, + bool is_using_refiner) { + ZVEC_TRY_RETURN_NULL("Failed to create HnswQueryParams", + auto *params = new zvec::HnswQueryParams( + ef, radius, is_linear, is_using_refiner); + return reinterpret_cast(params);) + return nullptr; +} + +void zvec_query_params_hnsw_destroy(zvec_hnsw_query_params_t *params) { + if (params) { + delete reinterpret_cast(params); + } +} + +zvec_error_code_t zvec_query_params_hnsw_set_ef(zvec_hnsw_query_params_t *params, + int ef) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "HNSW query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_ef(ef); + return ZVEC_OK; +} + +int zvec_query_params_hnsw_get_ef(const zvec_hnsw_query_params_t *params) { + if (!params) return zvec::core_interface::kDefaultHnswEfSearch; + auto *ptr = reinterpret_cast(params); + return ptr->ef(); +} + +zvec_error_code_t zvec_query_params_hnsw_set_radius(zvec_hnsw_query_params_t *params, + float radius) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "HNSW query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_radius(radius); + return ZVEC_OK; +} + +float zvec_query_params_hnsw_get_radius(const zvec_hnsw_query_params_t *params) { + if (!params) return 0.0f; + auto *ptr = reinterpret_cast(params); + return ptr->radius(); +} + +zvec_error_code_t zvec_query_params_hnsw_set_is_linear(zvec_hnsw_query_params_t *params, + bool is_linear) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "HNSW query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_linear(is_linear); + return ZVEC_OK; +} + +bool zvec_query_params_hnsw_get_is_linear(const zvec_hnsw_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_linear(); +} + +zvec_error_code_t zvec_query_params_hnsw_set_is_using_refiner( + zvec_hnsw_query_params_t *params, bool is_using_refiner) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "HNSW query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_using_refiner(is_using_refiner); + return ZVEC_OK; +} + +bool zvec_query_params_hnsw_get_is_using_refiner( + const zvec_hnsw_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_using_refiner(); +} + +// ============================================================================= +// DiskAnnQueryParams implementation - wrapper around zvec::DiskAnnQueryParams +// ============================================================================= + +zvec_diskann_query_params_t *zvec_query_params_diskann_create(int list_size) { + ZVEC_TRY_RETURN_NULL( + "Failed to create DiskAnnQueryParams", + auto *params = new zvec::DiskAnnQueryParams(list_size); + return reinterpret_cast(params);) + return nullptr; +} + +void zvec_query_params_diskann_destroy(zvec_diskann_query_params_t *params) { + if (params) { + delete reinterpret_cast(params); + } +} + +zvec_error_code_t zvec_query_params_diskann_set_list_size( + zvec_diskann_query_params_t *params, int list_size) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "DiskANN query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_list_size(list_size); + return ZVEC_OK; +} + +int zvec_query_params_diskann_get_list_size( + const zvec_diskann_query_params_t *params) { + if (!params) return 300; // DiskAnnQueryParams default + auto *ptr = reinterpret_cast(params); + return ptr->list_size(); +} + +zvec_error_code_t zvec_query_params_diskann_set_radius( + zvec_diskann_query_params_t *params, float radius) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "DiskANN query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_radius(radius); + return ZVEC_OK; +} + +float zvec_query_params_diskann_get_radius( + const zvec_diskann_query_params_t *params) { + if (!params) return 0.0f; + auto *ptr = reinterpret_cast(params); + return ptr->radius(); +} + +zvec_error_code_t zvec_query_params_diskann_set_is_linear( + zvec_diskann_query_params_t *params, bool is_linear) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "DiskANN query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_linear(is_linear); + return ZVEC_OK; +} + +bool zvec_query_params_diskann_get_is_linear( + const zvec_diskann_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_linear(); +} + +zvec_error_code_t zvec_query_params_diskann_set_is_using_refiner( + zvec_diskann_query_params_t *params, bool is_using_refiner) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "DiskANN query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_using_refiner(is_using_refiner); + return ZVEC_OK; +} + +bool zvec_query_params_diskann_get_is_using_refiner( + const zvec_diskann_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_using_refiner(); +} + +// ============================================================================= +// IVFQueryParams implementation - wrapper around zvec::IVFQueryParams +// ============================================================================= + +zvec_ivf_query_params_t *zvec_query_params_ivf_create(int nprobe, + bool is_using_refiner, + float scale_factor) { + ZVEC_TRY_RETURN_NULL("Failed to create IVFQueryParams", + auto *params = new zvec::IVFQueryParams( + nprobe, is_using_refiner, scale_factor); + return reinterpret_cast(params);) + return nullptr; +} + +void zvec_query_params_ivf_destroy(zvec_ivf_query_params_t *params) { + if (params) { + delete reinterpret_cast(params); + } +} + +zvec_error_code_t zvec_query_params_ivf_set_nprobe(zvec_ivf_query_params_t *params, + int nprobe) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_nprobe(nprobe); + return ZVEC_OK; +} + +int zvec_query_params_ivf_get_nprobe(const zvec_ivf_query_params_t *params) { + if (!params) return 10; + auto *ptr = reinterpret_cast(params); + return ptr->nprobe(); +} + +zvec_error_code_t zvec_query_params_ivf_set_scale_factor(zvec_ivf_query_params_t *params, + float scale_factor) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_scale_factor(scale_factor); + return ZVEC_OK; +} + +float zvec_query_params_ivf_get_scale_factor(const zvec_ivf_query_params_t *params) { + if (!params) return 10.0f; + auto *ptr = reinterpret_cast(params); + return ptr->scale_factor(); +} + +zvec_error_code_t zvec_query_params_ivf_set_radius(zvec_ivf_query_params_t *params, + float radius) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_radius(radius); + return ZVEC_OK; +} + +float zvec_query_params_ivf_get_radius(const zvec_ivf_query_params_t *params) { + if (!params) return 0.0f; + auto *ptr = reinterpret_cast(params); + return ptr->radius(); +} + +zvec_error_code_t zvec_query_params_ivf_set_is_linear(zvec_ivf_query_params_t *params, + bool is_linear) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_linear(is_linear); + return ZVEC_OK; +} + +bool zvec_query_params_ivf_get_is_linear(const zvec_ivf_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_linear(); +} + +zvec_error_code_t zvec_query_params_ivf_set_is_using_refiner( + zvec_ivf_query_params_t *params, bool is_using_refiner) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_using_refiner(is_using_refiner); + return ZVEC_OK; +} + +bool zvec_query_params_ivf_get_is_using_refiner( + const zvec_ivf_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_using_refiner(); +} + +// ============================================================================= +// FlatQueryParams implementation - wrapper around zvec::FlatQueryParams +// ============================================================================= + +zvec_flat_query_params_t *zvec_query_params_flat_create(bool is_using_refiner, + float scale_factor) { + ZVEC_TRY_RETURN_NULL( + "Failed to create FlatQueryParams", + auto *params = new zvec::FlatQueryParams(is_using_refiner, scale_factor); + return reinterpret_cast(params);) + return nullptr; +} + +void zvec_query_params_flat_destroy(zvec_flat_query_params_t *params) { + if (params) { + delete reinterpret_cast(params); + } +} + +zvec_error_code_t zvec_query_params_flat_set_scale_factor( + zvec_flat_query_params_t *params, float scale_factor) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Flat query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_scale_factor(scale_factor); + return ZVEC_OK; +} + +float zvec_query_params_flat_get_scale_factor( + const zvec_flat_query_params_t *params) { + if (!params) return 10.0f; + auto *ptr = reinterpret_cast(params); + return ptr->scale_factor(); +} + +zvec_error_code_t zvec_query_params_flat_set_radius(zvec_flat_query_params_t *params, + float radius) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Flat query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_radius(radius); + return ZVEC_OK; +} + +float zvec_query_params_flat_get_radius(const zvec_flat_query_params_t *params) { + if (!params) return 0.0f; + auto *ptr = reinterpret_cast(params); + return ptr->radius(); +} + +zvec_error_code_t zvec_query_params_flat_set_is_linear(zvec_flat_query_params_t *params, + bool is_linear) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Flat query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_linear(is_linear); + return ZVEC_OK; +} + +bool zvec_query_params_flat_get_is_linear(const zvec_flat_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_linear(); +} + +zvec_error_code_t zvec_query_params_flat_set_is_using_refiner( + zvec_flat_query_params_t *params, bool is_using_refiner) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Flat query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_using_refiner(is_using_refiner); + return ZVEC_OK; +} + +bool zvec_query_params_flat_get_is_using_refiner( + const zvec_flat_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_using_refiner(); +} + +// ============================================================================= +// FtsQueryParams implementation - wrapper around zvec::FtsQueryParams +// ============================================================================= + +zvec_fts_query_params_t *zvec_query_params_fts_create( + const char *default_operator) { + ZVEC_TRY_RETURN_NULL( + "Failed to create FtsQueryParams", + auto *params = new zvec::FtsQueryParams(); + if (default_operator && *default_operator) { + params->set_default_operator(std::string(default_operator)); + } return reinterpret_cast(params);) + return nullptr; +} + +void zvec_query_params_fts_destroy(zvec_fts_query_params_t *params) { + if (params) { + delete reinterpret_cast(params); + } +} + +zvec_error_code_t zvec_query_params_fts_set_default_operator( + zvec_fts_query_params_t *params, const char *default_operator) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "FTS query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_default_operator(std::string(default_operator ? default_operator + : "")); + return ZVEC_OK; +} + +const char *zvec_query_params_fts_get_default_operator( + const zvec_fts_query_params_t *params) { + if (!params) return nullptr; + auto *ptr = reinterpret_cast(params); + return ptr->default_operator().c_str(); +} + +// ============================================================================= +// Query implementation - owns zvec::SearchQuery via raw pointer +// (external C symbol naming kept for ABI compatibility) +// ============================================================================= + +zvec_vector_query_t *zvec_vector_query_create(void) { + ZVEC_TRY_RETURN_NULL("Failed to create query object", + auto *query = new zvec::SearchQuery(); + query->topk_ = 10; query->include_doc_id_ = true; + query->include_vector_ = false; + return reinterpret_cast(query);) + return nullptr; +} + +void zvec_vector_query_destroy(zvec_vector_query_t *query) { + if (query) { + delete reinterpret_cast(query); + } +} + +zvec_error_code_t zvec_vector_query_set_topk(zvec_vector_query_t *query, int topk) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->topk_ = topk; + return ZVEC_OK; +} + +int zvec_vector_query_get_topk(const zvec_vector_query_t *query) { + if (!query) return 10; + auto *ptr = reinterpret_cast(query); + return ptr->topk_; +} + +zvec_error_code_t zvec_vector_query_set_field_name(zvec_vector_query_t *query, + const char *field_name) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->target_.field_name_ = field_name ? field_name : ""; + return ZVEC_OK; +} + +const char *zvec_vector_query_get_field_name(const zvec_vector_query_t *query) { + if (!query) return nullptr; + auto *ptr = reinterpret_cast(query); + return ptr->target_.field_name_.empty() ? nullptr + : ptr->target_.field_name_.c_str(); +} + +zvec_error_code_t zvec_vector_query_set_query_vector(zvec_vector_query_t *query, + const void *data, + size_t size) { + if (!query || !data || size == 0) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Vector query pointer or data is null/empty"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + // Copies into VectorClause (not VectorViewClause) because the C API does + // not require `data` to stay alive after this call returns. + ptr->target_.set_vector(std::string(static_cast(data), size)); + return ZVEC_OK; +} + +zvec_error_code_t zvec_vector_query_set_filter(zvec_vector_query_t *query, + const char *filter) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->filter_ = filter ? filter : ""; + return ZVEC_OK; +} + +const char *zvec_vector_query_get_filter(const zvec_vector_query_t *query) { + if (!query) return nullptr; + auto *ptr = reinterpret_cast(query); + return ptr->filter_.empty() ? nullptr : ptr->filter_.c_str(); +} + +zvec_error_code_t zvec_vector_query_set_include_vector(zvec_vector_query_t *query, + bool include) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->include_vector_ = include; + return ZVEC_OK; +} + +bool zvec_vector_query_get_include_vector(const zvec_vector_query_t *query) { + if (!query) return false; + auto *ptr = reinterpret_cast(query); + return ptr->include_vector_; +} + +zvec_error_code_t zvec_vector_query_set_include_doc_id(zvec_vector_query_t *query, + bool include) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->include_doc_id_ = include; + return ZVEC_OK; +} + +bool zvec_vector_query_get_include_doc_id(const zvec_vector_query_t *query) { + if (!query) return false; + auto *ptr = reinterpret_cast(query); + return ptr->include_doc_id_; +} + +zvec_error_code_t zvec_vector_query_set_output_fields(zvec_vector_query_t *query, + const char **fields, + size_t count) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + if (!fields || count == 0) { + ptr->output_fields_ = std::nullopt; + } else { + std::vector result; + result.reserve(count); + for (size_t i = 0; i < count; ++i) { + result.emplace_back(fields[i]); + } + ptr->output_fields_ = std::move(result); + } + return ZVEC_OK; +} + +zvec_error_code_t zvec_vector_query_get_output_fields(const zvec_vector_query_t *query, + const char ***fields, + size_t *count) { + if (!query || !fields || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query, fields, or count pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + + if (!ptr->output_fields_.has_value()) { + *fields = nullptr; + *count = 0; + } else { + const auto &output_fields = ptr->output_fields_.value(); + *count = output_fields.size(); + *fields = (const char **)malloc(*count * sizeof(const char *)); + if (!*fields) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + for (size_t i = 0; i < *count; ++i) { + (*fields)[i] = strdup(output_fields[i].c_str()); + } + } + return ZVEC_OK; +} + +// ============================================================================= +// VamanaQueryParams implementation - wrapper around zvec::VamanaQueryParams +// ============================================================================= + +zvec_vamana_query_params_t *zvec_query_params_vamana_create( + int ef_search, float radius, bool is_linear, bool is_using_refiner) { + ZVEC_TRY_RETURN_NULL( + "Failed to create VamanaQueryParams", + auto *params = new zvec::VamanaQueryParams(ef_search, radius, is_linear, + is_using_refiner); + return reinterpret_cast(params);) + return nullptr; +} + +void zvec_query_params_vamana_destroy(zvec_vamana_query_params_t *params) { + if (params) { + delete reinterpret_cast(params); + } +} + +zvec_error_code_t zvec_query_params_vamana_set_ef_search( + zvec_vamana_query_params_t *params, int ef_search) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Vamana query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_ef_search(ef_search); + return ZVEC_OK; +} + +int zvec_query_params_vamana_get_ef_search( + const zvec_vamana_query_params_t *params) { + if (!params) return zvec::core_interface::kDefaultVamanaEfSearch; + auto *ptr = reinterpret_cast(params); + return ptr->ef_search(); +} + +zvec_error_code_t zvec_query_params_vamana_set_radius( + zvec_vamana_query_params_t *params, float radius) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Vamana query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_radius(radius); + return ZVEC_OK; +} + +float zvec_query_params_vamana_get_radius( + const zvec_vamana_query_params_t *params) { + if (!params) return 0.0f; + auto *ptr = reinterpret_cast(params); + return ptr->radius(); +} + +zvec_error_code_t zvec_query_params_vamana_set_is_linear( + zvec_vamana_query_params_t *params, bool is_linear) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Vamana query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_linear(is_linear); + return ZVEC_OK; +} + +bool zvec_query_params_vamana_get_is_linear( + const zvec_vamana_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_linear(); +} + +zvec_error_code_t zvec_query_params_vamana_set_is_using_refiner( + zvec_vamana_query_params_t *params, bool is_using_refiner) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Vamana query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_using_refiner(is_using_refiner); + return ZVEC_OK; +} + +bool zvec_query_params_vamana_get_is_using_refiner( + const zvec_vamana_query_params_t *params) { + if (!params) return false; + auto *ptr = reinterpret_cast(params); + return ptr->is_using_refiner(); +} + +// ============================================================================= +// Type-safe query params attachment functions (transfer ownership to +// query object) +// ============================================================================= + +zvec_error_code_t zvec_vector_query_set_query_params(zvec_vector_query_t *query, + void *params) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + + // Cast to QueryParams* and transfer ownership via shared_ptr. + auto *params_ptr = reinterpret_cast(params); + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +// Type-specific setters for cleaner ownership transfer +zvec_error_code_t zvec_vector_query_set_hnsw_params( + zvec_vector_query_t *query, zvec_hnsw_query_params_t *hnsw_params) { + if (!query || !hnsw_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or HNSW params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(hnsw_params); + + // Transfer ownership via shared_ptr (polymorphic conversion) + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_vector_query_set_ivf_params(zvec_vector_query_t *query, + zvec_ivf_query_params_t *ivf_params) { + if (!query || !ivf_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or IVF params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(ivf_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_vector_query_set_flat_params( + zvec_vector_query_t *query, zvec_flat_query_params_t *flat_params) { + if (!query || !flat_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or Flat params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(flat_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_vector_query_set_fts_params( + zvec_vector_query_t *query, zvec_fts_query_params_t *fts_params) { + if (!query || !fts_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or FTS params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(fts_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_vector_query_set_vamana_params( + zvec_vector_query_t *query, zvec_vamana_query_params_t *vamana_params) { + if (!query || !vamana_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or Vamana params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(vamana_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_vector_query_set_diskann_params( + zvec_vector_query_t *query, zvec_diskann_query_params_t *diskann_params) { + if (!query || !diskann_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or DiskANN params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(diskann_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +// ============================================================================= +// Fts payload implementation - wrapper around zvec::FtsClause (value type) +// ============================================================================= + +zvec_fts_t *zvec_fts_create(void) { + ZVEC_TRY_RETURN_NULL("Failed to create Fts payload", + auto *fts = new zvec::FtsClause(); + return reinterpret_cast(fts);) + return nullptr; +} + +void zvec_fts_destroy(zvec_fts_t *fts) { + if (fts) { + delete reinterpret_cast(fts); + } +} + +zvec_error_code_t zvec_fts_set_query_string(zvec_fts_t *fts, + const char *query_string) { + if (!fts) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Fts pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(fts); + ptr->query_string_ = query_string ? query_string : ""; + return ZVEC_OK; +} + +zvec_error_code_t zvec_fts_set_match_string(zvec_fts_t *fts, + const char *match_string) { + if (!fts) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Fts pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(fts); + ptr->match_string_ = match_string ? match_string : ""; + return ZVEC_OK; +} + +const char *zvec_fts_get_query_string(const zvec_fts_t *fts) { + if (!fts) return nullptr; + auto *ptr = reinterpret_cast(fts); + return ptr->query_string_.c_str(); +} + +const char *zvec_fts_get_match_string(const zvec_fts_t *fts) { + if (!fts) return nullptr; + auto *ptr = reinterpret_cast(fts); + return ptr->match_string_.c_str(); +} + +zvec_error_code_t zvec_vector_query_set_fts(zvec_vector_query_t *query, + const zvec_fts_t *fts) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *query_ptr = reinterpret_cast(query); + if (!fts) { + // Clearing FTS resets the target to an empty vector clause. + query_ptr->target_.clause_ = zvec::VectorClause{}; + } else { + query_ptr->target_.clause_ = *reinterpret_cast(fts); + } + return ZVEC_OK; +} + +const zvec_fts_t *zvec_vector_query_get_fts(const zvec_vector_query_t *query) { + if (!query) return nullptr; + auto *query_ptr = reinterpret_cast(query); + const auto *fc = std::get_if(&query_ptr->target_.clause_); + if (!fc) return nullptr; + return reinterpret_cast(fc); +} + +// ============================================================================= +// GroupByVectorQuery implementation - owns zvec::GroupByVectorQuery via raw +// pointer +// ============================================================================= + +zvec_group_by_vector_query_t *zvec_group_by_vector_query_create(void) { + ZVEC_TRY_RETURN_NULL( + "Failed to create GroupByVectorQuery", + auto *query = new zvec::GroupByVectorQuery(); + query->group_count_ = 2; query->group_topk_ = 3; + return reinterpret_cast(query);) + return nullptr; +} + +void zvec_group_by_vector_query_destroy(zvec_group_by_vector_query_t *query) { + if (query) { + delete reinterpret_cast(query); + } +} + +zvec_error_code_t zvec_group_by_vector_query_set_field_name( + zvec_group_by_vector_query_t *query, const char *field_name) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "GroupByVectorQuery pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->target_.field_name_ = field_name ? field_name : ""; + return ZVEC_OK; +} + +const char *zvec_group_by_vector_query_get_field_name( + const zvec_group_by_vector_query_t *query) { + if (!query) return nullptr; + auto *ptr = reinterpret_cast(query); + return ptr->target_.field_name_.empty() ? nullptr + : ptr->target_.field_name_.c_str(); +} + +zvec_error_code_t zvec_group_by_vector_query_set_group_by_field_name( + zvec_group_by_vector_query_t *query, const char *field_name) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "GroupByVectorQuery pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->group_by_field_name_ = field_name ? field_name : ""; + return ZVEC_OK; +} + +const char *zvec_group_by_vector_query_get_group_by_field_name( + const zvec_group_by_vector_query_t *query) { + if (!query) return nullptr; + auto *ptr = reinterpret_cast(query); + return ptr->group_by_field_name_.empty() ? nullptr + : ptr->group_by_field_name_.c_str(); +} + +zvec_error_code_t zvec_group_by_vector_query_set_group_count( + zvec_group_by_vector_query_t *query, uint32_t count) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "GroupByVectorQuery pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->group_count_ = count; + return ZVEC_OK; +} + +uint32_t zvec_group_by_vector_query_get_group_count( + const zvec_group_by_vector_query_t *query) { + if (!query) return 2; + auto *ptr = reinterpret_cast(query); + return ptr->group_count_; +} + +zvec_error_code_t zvec_group_by_vector_query_set_group_topk( + zvec_group_by_vector_query_t *query, uint32_t topk) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "GroupByVectorQuery pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->group_topk_ = topk; + return ZVEC_OK; +} + +uint32_t zvec_group_by_vector_query_get_group_topk( + const zvec_group_by_vector_query_t *query) { + if (!query) return 3; + auto *ptr = reinterpret_cast(query); + return ptr->group_topk_; +} + +zvec_error_code_t zvec_group_by_vector_query_set_query_vector( + zvec_group_by_vector_query_t *query, const void *data, size_t size) { + if (!query || !data || size == 0) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "GroupByVectorQuery pointer or data is null/empty"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + // Copies into VectorClause — see comment on zvec_vector_query_set_query_vector. + ptr->target_.set_vector(std::string(static_cast(data), size)); + return ZVEC_OK; +} + +zvec_error_code_t zvec_group_by_vector_query_set_filter( + zvec_group_by_vector_query_t *query, const char *filter) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "GroupByVectorQuery pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->filter_ = filter ? filter : ""; + return ZVEC_OK; +} + +const char *zvec_group_by_vector_query_get_filter( + const zvec_group_by_vector_query_t *query) { + if (!query) return nullptr; + auto *ptr = reinterpret_cast(query); + return ptr->filter_.empty() ? nullptr : ptr->filter_.c_str(); +} + +zvec_error_code_t zvec_group_by_vector_query_set_include_vector( + zvec_group_by_vector_query_t *query, bool include) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "GroupByVectorQuery pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->include_vector_ = include; + return ZVEC_OK; +} + +bool zvec_group_by_vector_query_get_include_vector( + const zvec_group_by_vector_query_t *query) { + if (!query) return false; + auto *ptr = reinterpret_cast(query); + return ptr->include_vector_; +} + +zvec_error_code_t zvec_group_by_vector_query_set_output_fields( + zvec_group_by_vector_query_t *query, const char **fields, size_t count) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "GroupByVectorQuery pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + if (!fields || count == 0) { + ptr->output_fields_ = std::nullopt; + } else { + std::vector result; + result.reserve(count); + for (size_t i = 0; i < count; ++i) { + result.emplace_back(fields[i]); + } + ptr->output_fields_ = std::move(result); + } + return ZVEC_OK; +} + +zvec_error_code_t zvec_group_by_vector_query_get_output_fields( + zvec_group_by_vector_query_t *query, const char ***fields, size_t *count) { + if (!query || !fields || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query, fields, or count pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + + if (!ptr->output_fields_.has_value()) { + *fields = nullptr; + *count = 0; + } else { + const auto &output_fields = ptr->output_fields_.value(); + *count = output_fields.size(); + *fields = (const char **)malloc(*count * sizeof(const char *)); + if (!*fields) { + SET_LAST_ERROR(ZVEC_ERROR_RESOURCE_EXHAUSTED, + "Failed to allocate memory"); + return ZVEC_ERROR_RESOURCE_EXHAUSTED; + } + for (size_t i = 0; i < *count; ++i) { + (*fields)[i] = strdup(output_fields[i].c_str()); + } + } + return ZVEC_OK; +} + +zvec_error_code_t zvec_group_by_vector_query_set_query_params( + zvec_group_by_vector_query_t *query, void *params) { + if (!query || !params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +// Type-specific setters for GroupByVectorQuery +zvec_error_code_t zvec_group_by_vector_query_set_hnsw_params( + zvec_group_by_vector_query_t *query, zvec_hnsw_query_params_t *hnsw_params) { + if (!query || !hnsw_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or HNSW params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(hnsw_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_group_by_vector_query_set_ivf_params( + zvec_group_by_vector_query_t *query, zvec_ivf_query_params_t *ivf_params) { + if (!query || !ivf_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or IVF params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(ivf_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_group_by_vector_query_set_flat_params( + zvec_group_by_vector_query_t *query, zvec_flat_query_params_t *flat_params) { + if (!query || !flat_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or Flat params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(flat_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_group_by_vector_query_set_vamana_params( + zvec_group_by_vector_query_t *query, + zvec_vamana_query_params_t *vamana_params) { + if (!query || !vamana_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or Vamana params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(vamana_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_group_by_vector_query_set_diskann_params( + zvec_group_by_vector_query_t *query, + zvec_diskann_query_params_t *diskann_params) { + if (!query || !diskann_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or DiskANN params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(diskann_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + +// ============================================================================= +// Reranker Implementation +// ============================================================================= + +zvec_error_code_t zvec_multi_query_set_rerank_rrf( + zvec_multi_query_t *query, int rank_constant) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *mq = reinterpret_cast(query); + mq->rerank = zvec::reranker::RrfParams{rank_constant}; + return ZVEC_OK; +} + +zvec_error_code_t zvec_multi_query_set_rerank_weighted( + zvec_multi_query_t *query, const double *weights, size_t weight_count) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + if (!weights && weight_count > 0) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Weights pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *mq = reinterpret_cast(query); + mq->rerank = zvec::reranker::WeightedParams{ + std::vector(weights, weights + weight_count)}; + return ZVEC_OK; +} + +// ============================================================================= +// MultiVectorQuery Implementation +// ============================================================================= + +zvec_multi_query_t *zvec_multi_query_create(void) { + ZVEC_TRY_RETURN_NULL("Failed to create MultiVectorQuery", + auto *query = new zvec::MultiQuery(); + return reinterpret_cast( + query);) + return nullptr; +} + +void zvec_multi_query_destroy(zvec_multi_query_t *query) { + if (query) { + delete reinterpret_cast(query); + } +} + +zvec_error_code_t zvec_multi_query_add_sub_query( + zvec_multi_query_t *query, + const zvec_sub_query_t *sub_query) { + if (!query || !sub_query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or sub_query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *mvq = reinterpret_cast(query); + auto *sub = reinterpret_cast(sub_query); + mvq->queries.push_back(*sub); + + return ZVEC_OK; +} + +size_t zvec_multi_query_get_sub_query_count( + const zvec_multi_query_t *query) { + if (!query) return 0; + auto *mvq = reinterpret_cast(query); + return mvq->queries.size(); +} + +zvec_error_code_t zvec_multi_query_set_topk( + zvec_multi_query_t *query, int topk) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Multi-vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *mvq = reinterpret_cast(query); + mvq->topk = topk; + return ZVEC_OK; +} + +int zvec_multi_query_get_topk( + const zvec_multi_query_t *query) { + if (!query) return 0; + auto *mvq = reinterpret_cast(query); + return mvq->topk; +} + +zvec_error_code_t zvec_multi_query_set_filter( + zvec_multi_query_t *query, const char *filter) { + if (!query || !filter) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or filter pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *mvq = reinterpret_cast(query); + mvq->filter = std::string(filter); + return ZVEC_OK; +} + +const char *zvec_multi_query_get_filter( + const zvec_multi_query_t *query) { + if (!query) return nullptr; + auto *mvq = reinterpret_cast(query); + return mvq->filter.c_str(); +} + +zvec_error_code_t zvec_multi_query_set_include_vector( + zvec_multi_query_t *query, bool include) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Multi-vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *mvq = reinterpret_cast(query); + mvq->include_vector = include; + return ZVEC_OK; +} + +bool zvec_multi_query_get_include_vector( + const zvec_multi_query_t *query) { + if (!query) return false; + auto *mvq = reinterpret_cast(query); + return mvq->include_vector; +} + +zvec_error_code_t zvec_multi_query_set_output_fields( + zvec_multi_query_t *query, const char **fields, size_t count) { + if (!query || (!fields && count > 0)) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query pointer is null or fields is null with count > 0"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *mvq = reinterpret_cast(query); + std::vector field_vec; + field_vec.reserve(count); + for (size_t i = 0; i < count; ++i) { + if (!fields[i]) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Null field name at index " + std::to_string(i)); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + field_vec.emplace_back(fields[i]); + } + mvq->output_fields = std::move(field_vec); + + return ZVEC_OK; +} + +zvec_error_code_t zvec_multi_query_get_output_fields( + zvec_multi_query_t *query, const char ***fields, size_t *count) { + if (!query || !fields || !count) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query, fields or count pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *mvq = reinterpret_cast(query); + if (!mvq->output_fields.has_value() || mvq->output_fields->empty()) { + *fields = nullptr; + *count = 0; + return ZVEC_OK; + } + + const auto &field_vec = mvq->output_fields.value(); + *count = field_vec.size(); + *fields = static_cast(malloc(*count * sizeof(const char *))); + if (!*fields) { + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, "Failed to allocate memory"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + for (size_t i = 0; i < *count; ++i) { + (*fields)[i] = field_vec[i].c_str(); + } + + return ZVEC_OK; +} + +// ============================================================================= +// SubVectorQuery Implementation +// ============================================================================= + +zvec_sub_query_t *zvec_sub_query_create(void) { + ZVEC_TRY_RETURN_NULL("Failed to create SubVectorQuery", + auto *query = new zvec::SubQuery(); + query->num_candidates_ = 10; + return reinterpret_cast( + query);) + return nullptr; +} + +void zvec_sub_query_destroy(zvec_sub_query_t *query) { + if (query) { + delete reinterpret_cast(query); + } +} + +zvec_error_code_t zvec_sub_query_set_num_candidates( + zvec_sub_query_t *query, int num_candidates) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->num_candidates_ = num_candidates; + return ZVEC_OK; +} + +int zvec_sub_query_get_num_candidates( + const zvec_sub_query_t *query) { + if (!query) return 0; + auto *ptr = reinterpret_cast(query); + return ptr->num_candidates_; +} + +zvec_error_code_t zvec_sub_query_set_field_name( + zvec_sub_query_t *query, const char *field_name) { + if (!query || !field_name) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or field_name pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + ptr->target_.field_name_ = std::string(field_name); + return ZVEC_OK; +} + +const char *zvec_sub_query_get_field_name( + const zvec_sub_query_t *query) { + if (!query) return nullptr; + auto *ptr = reinterpret_cast(query); + return ptr->target_.field_name_.c_str(); +} + +zvec_error_code_t zvec_sub_query_set_query_vector( + zvec_sub_query_t *query, const void *data, size_t size) { + if (!query || !data || size == 0) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query pointer or data is null/empty"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + // Copies into VectorClause — see comment on zvec_vector_query_set_query_vector. + auto &payload = std::get(ptr->target_.clause_); + payload.query_vector_.assign(static_cast(data), size); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_sparse_vector( + zvec_sub_query_t *query, const uint32_t *indices, const float *values, + size_t count) { + if (!query || (!indices && count > 0) || (!values && count > 0)) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query, indices or values pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *ptr = reinterpret_cast(query); + auto &payload = std::get(ptr->target_.clause_); + if (count == 0) { + payload.sparse_indices_.clear(); + payload.sparse_values_.clear(); + return ZVEC_OK; + } + payload.sparse_indices_.assign( + reinterpret_cast(indices), count * sizeof(uint32_t)); + payload.sparse_values_.assign( + reinterpret_cast(values), count * sizeof(float)); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_sparse_indices( + zvec_sub_query_t *query, const uint32_t *indices, size_t count) { + if (!query || (!indices && count > 0)) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or indices pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto &payload = std::get(ptr->target_.clause_); + payload.sparse_indices_.assign( + reinterpret_cast(indices), count * sizeof(uint32_t)); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_sparse_values( + zvec_sub_query_t *query, const float *values, size_t count) { + if (!query || (!values && count > 0)) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or values pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto &payload = std::get(ptr->target_.clause_); + payload.sparse_values_.assign( + reinterpret_cast(values), count * sizeof(float)); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_hnsw_params( + zvec_sub_query_t *query, zvec_hnsw_query_params_t *hnsw_params) { + if (!query || !hnsw_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or HNSW params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(hnsw_params); + ptr->target_.query_params_.reset(params_ptr); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_ivf_params( + zvec_sub_query_t *query, zvec_ivf_query_params_t *ivf_params) { + if (!query || !ivf_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or IVF params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(ivf_params); + ptr->target_.query_params_.reset(params_ptr); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_flat_params( + zvec_sub_query_t *query, zvec_flat_query_params_t *flat_params) { + if (!query || !flat_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or Flat params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(flat_params); + ptr->target_.query_params_.reset(params_ptr); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_vamana_params( + zvec_sub_query_t *query, zvec_vamana_query_params_t *vamana_params) { + if (!query || !vamana_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or Vamana params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(vamana_params); + ptr->target_.query_params_.reset(params_ptr); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_fts_params( + zvec_sub_query_t *query, zvec_fts_query_params_t *fts_params) { + if (!query || !fts_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-query or FTS params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto *params_ptr = reinterpret_cast(fts_params); + ptr->target_.query_params_.reset(params_ptr); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_diskann_params( + zvec_sub_query_t *query, zvec_diskann_query_params_t *diskann_params) { + if (!query || !diskann_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or DiskANN params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(diskann_params); + ptr->target_.query_params_.reset(params_ptr); + return ZVEC_OK; +} + +zvec_error_code_t zvec_sub_query_set_fts(zvec_sub_query_t *query, + const zvec_fts_t *fts) { + if (!query) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Sub-query pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + if (!fts) { + ptr->target_.clause_ = zvec::VectorClause{}; + } else { + ptr->target_.clause_ = *reinterpret_cast(fts); + } + return ZVEC_OK; +} + +// ============================================================================= +// Index Interface Implementation +// ============================================================================= + +/** + * @brief Create index on a collection column + * @param collection Collection handle + * @param column_name Column name to create index on + * @param index_params Index parameters + * @return ZVEC_OK on success, error code on failure + * @note index_params is cloned internally, caller should still call + * zvec_index_params_destroy() to free the original + */ +zvec_error_code_t zvec_collection_create_index( + zvec_collection_t *collection, const char *column_name, + const zvec_index_params_t *index_params) { + if (!collection || !column_name || !index_params) { + set_last_error( + "Invalid arguments: collection, column_name, and index_params cannot " + "be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_create_index", + auto coll_ptr = + reinterpret_cast *>(collection); + std::string field_name_str(column_name); + + auto *cpp_params = + reinterpret_cast(index_params); + auto index_params_ptr = cpp_params->clone(); + auto status = (*coll_ptr)->CreateIndex(field_name_str, index_params_ptr); + return status_to_error_code(status);) +} + +/** + * @brief Drop index from a collection column + * @param collection Collection handle + * @param column_name Column name to drop index from + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_collection_drop_index(zvec_collection_t *collection, + const char *column_name) { + if (!collection || !column_name) { + set_last_error( + "Invalid arguments: collection and column_name cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + zvec::Status status = (*coll_ptr)->DropIndex(column_name); + if (!status.ok()) { set_last_error(status.message()); } + + return status_to_error_code(status);) +} + +/** + * @brief Optimize collection (rebuild indexes, merge segments) + * @param collection Collection handle + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_collection_optimize(zvec_collection_t *collection) { + if (!collection) { + set_last_error("Invalid argument: collection cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + zvec::Status status = (*coll_ptr)->Optimize(); + if (!status.ok()) { set_last_error(status.message()); } + + return status_to_error_code(status);) +} + +// ============================================================================= +// Column Interface Implementation +// ============================================================================= + +/** + * @brief Add a column to collection + * @param collection Collection handle + * @param field_schema Field schema (deep-copied, caller retains ownership) + * @param expression Default value expression (can be NULL for no default) + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_collection_add_column(zvec_collection_t *collection, + const zvec_field_schema_t *field_schema, + const char *expression) { + if (!collection || !field_schema) { + set_last_error( + "Invalid arguments: collection and field_schema cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + + // Deep copy the schema - caller retains ownership + auto *cpp_schema = + reinterpret_cast(field_schema); + zvec::FieldSchema::Ptr schema = + std::make_shared(*cpp_schema); + + std::string expr = expression ? expression : ""; + zvec::Status status = (*coll_ptr)->AddColumn(schema, expr); + + if (!status.ok()) { set_last_error(status.message()); } + + return status_to_error_code(status);) +} + +/** + * @brief Drop a column from collection + * @param collection Collection handle + * @param column_name Column name to drop + * @return ZVEC_OK on success, error code on failure + */ +zvec_error_code_t zvec_collection_drop_column(zvec_collection_t *collection, + const char *column_name) { + if (!collection || !column_name) { + set_last_error( + "Invalid arguments: collection and column_name cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + zvec::Status status = (*coll_ptr)->DropColumn(column_name); + + if (!status.ok()) { set_last_error(status.message()); } + + return status_to_error_code(status);) +} + +zvec_error_code_t zvec_collection_alter_column( + zvec_collection_t *collection, const char *column_name, const char *new_name, + const zvec_field_schema_t *new_schema) { + if (!collection || !column_name) { + set_last_error( + "Invalid arguments: collection and column_name cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + std::string rename = new_name ? new_name : ""; + + // Deep copy the schema - caller retains ownership and must call + // zvec_field_schema_destroy after the call + zvec::FieldSchema::Ptr schema = nullptr; + if (new_schema) { + auto *cpp_schema = + reinterpret_cast(new_schema); + // Use copy constructor to create a deep copy + schema = std::make_shared(*cpp_schema); + } + + zvec::Status status = + (*coll_ptr)->AlterColumn(column_name, rename, schema); + if (!status.ok()) { set_last_error(status.message()); } + + return status_to_error_code(status);) +} + +// ============================================================================= +// DML Interface Implementation +// ============================================================================= + +zvec_error_code_t zvec_collection_insert(zvec_collection_t *collection, + const zvec_doc_t **docs, size_t doc_count, + size_t *success_count, + size_t *error_count) { + if (!collection || !docs || doc_count == 0 || !success_count || + !error_count) { + set_last_error( + "Invalid arguments: collection, docs, doc_count, success_count and " + "error_count cannot be null/zero"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_insert_docs", + auto coll_ptr = + reinterpret_cast *>(collection); + + std::vector internal_docs = + convert_zvec_docs_to_internal(docs, doc_count); + + auto result = (*coll_ptr)->Insert(internal_docs); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code == ZVEC_OK) { + *success_count = 0; + *error_count = 0; + for (const auto &status : result.value()) { + if (status.ok()) { + (*success_count)++; + } else { + (*error_count)++; + } + } + } else { + *success_count = 0; + *error_count = doc_count; + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_insert_with_results(zvec_collection_t *collection, + const zvec_doc_t **docs, + size_t doc_count, + zvec_write_result_t **results, + size_t *result_count) { + if (!collection || !docs || doc_count == 0 || !results || !result_count) { + set_last_error( + "Invalid arguments: collection, docs, doc_count, results and " + "result_count cannot be null/zero"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + *results = nullptr; + *result_count = 0; + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_insert_with_results", + auto coll_ptr = + reinterpret_cast *>(collection); + + std::vector internal_docs = + convert_zvec_docs_to_internal(docs, doc_count); + std::vector pks = collect_doc_pks(docs, doc_count); + + auto result = (*coll_ptr)->Insert(internal_docs); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code != ZVEC_OK) { return error_code; } + + return build_write_results(result.value(), results, result_count);) +} + +zvec_error_code_t zvec_collection_update(zvec_collection_t *collection, + const zvec_doc_t **docs, size_t doc_count, + size_t *success_count, + size_t *error_count) { + if (!collection || !docs || doc_count == 0 || !success_count || + !error_count) { + set_last_error( + "Invalid arguments: collection, docs, doc_count, success_count and " + "error_count cannot be null/zero"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + + std::vector internal_docs = + convert_zvec_docs_to_internal(docs, doc_count); + + auto result = (*coll_ptr)->Update(internal_docs); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code == ZVEC_OK) { + *success_count = 0; + *error_count = 0; + for (const auto &status : result.value()) { + if (status.ok()) { + (*success_count)++; + } else { + (*error_count)++; + } + } + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_update_with_results(zvec_collection_t *collection, + const zvec_doc_t **docs, + size_t doc_count, + zvec_write_result_t **results, + size_t *result_count) { + if (!collection || !docs || doc_count == 0 || !results || !result_count) { + set_last_error( + "Invalid arguments: collection, docs, doc_count, results and " + "result_count cannot be null/zero"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + *results = nullptr; + *result_count = 0; + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_update_with_results", + auto coll_ptr = + reinterpret_cast *>(collection); + + std::vector internal_docs = + convert_zvec_docs_to_internal(docs, doc_count); + std::vector pks = collect_doc_pks(docs, doc_count); + + auto result = (*coll_ptr)->Update(internal_docs); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code != ZVEC_OK) { return error_code; } + + return build_write_results(result.value(), results, result_count);) +} + +zvec_error_code_t zvec_collection_upsert(zvec_collection_t *collection, + const zvec_doc_t **docs, size_t doc_count, + size_t *success_count, + size_t *error_count) { + if (!collection || !docs || doc_count == 0 || !success_count || + !error_count) { + set_last_error( + "Invalid arguments: collection, docs, doc_count, success_count and " + "error_count cannot be null/zero"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + + std::vector internal_docs = + convert_zvec_docs_to_internal(docs, doc_count); + + auto result = (*coll_ptr)->Upsert(internal_docs); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code == ZVEC_OK) { + *success_count = 0; + *error_count = 0; + for (const auto &status : result.value()) { + if (status.ok()) { + (*success_count)++; + } else { + (*error_count)++; + } + } + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_upsert_with_results(zvec_collection_t *collection, + const zvec_doc_t **docs, + size_t doc_count, + zvec_write_result_t **results, + size_t *result_count) { + if (!collection || !docs || doc_count == 0 || !results || !result_count) { + set_last_error( + "Invalid arguments: collection, docs, doc_count, results and " + "result_count cannot be null/zero"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + *results = nullptr; + *result_count = 0; + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_upsert_with_results", + auto coll_ptr = + reinterpret_cast *>(collection); + + std::vector internal_docs = + convert_zvec_docs_to_internal(docs, doc_count); + std::vector pks = collect_doc_pks(docs, doc_count); + + auto result = (*coll_ptr)->Upsert(internal_docs); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code != ZVEC_OK) { return error_code; } + + return build_write_results(result.value(), results, result_count);) +} + +zvec_error_code_t zvec_collection_delete(zvec_collection_t *collection, + const char *const *pks, size_t pk_count, + size_t *success_count, + size_t *error_count) { + if (!collection || !pks || pk_count == 0 || !success_count || + !error_count) { + set_last_error( + "Invalid arguments: collection, pks, pk_count, success_count and " + "error_count cannot be null/zero"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + + std::vector primary_keys; primary_keys.reserve(pk_count); + for (size_t i = 0; i < pk_count; ++i) { + if (pks[i]) { + primary_keys.emplace_back(pks[i]); + } + } + + auto result = (*coll_ptr)->Delete(primary_keys); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code == ZVEC_OK) { + *success_count = 0; + *error_count = 0; + for (const auto &status : result.value()) { + if (status.ok()) { + (*success_count)++; + } else { + (*error_count)++; + } + } + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_delete_with_results(zvec_collection_t *collection, + const char *const *pks, + size_t pk_count, + zvec_write_result_t **results, + size_t *result_count) { + if (!collection || !pks || pk_count == 0 || !results || !result_count) { + set_last_error( + "Invalid arguments: collection, pks, pk_count, results and " + "result_count cannot be null/zero"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + *results = nullptr; + *result_count = 0; + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_delete_with_results", + auto coll_ptr = + reinterpret_cast *>(collection); + + std::vector primary_keys; primary_keys.reserve(pk_count); + for (size_t i = 0; i < pk_count; ++i) { + if (pks[i]) { + primary_keys.emplace_back(pks[i]); + } else { + primary_keys.emplace_back(""); + } + } + + auto result = (*coll_ptr)->Delete(primary_keys); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code != ZVEC_OK) { return error_code; } + + return build_write_results(result.value(), results, + result_count);) +} + +zvec_error_code_t zvec_collection_delete_by_filter(zvec_collection_t *collection, + const char *filter) { + if (!collection || !filter) { + set_last_error("Invalid arguments: collection,filter cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>(collection); + + auto status = (*coll_ptr)->DeleteByFilter(filter); if (!status.ok()) { + set_last_error(status.message()); + return status_to_error_code(status); + } + return ZVEC_OK;) +} + +// ============================================================================= +// Data query interface implementation +// ============================================================================= + +// Helper function to convert document results to C API format +zvec_error_code_t convert_document_results( + const std::vector> &query_results, + zvec_doc_t ***results, size_t *result_count) { + *result_count = query_results.size(); + *results = + static_cast(malloc(*result_count * sizeof(zvec_doc_t *))); + + if (!*results) { + set_last_error("Failed to allocate memory for query results"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + for (size_t i = 0; i < *result_count; ++i) { + const auto &internal_doc = query_results[i]; + // Create new document wrapper + zvec_doc_t *c_doc = zvec_doc_create(); + if (!c_doc) { + // Clean up previously allocated documents + for (size_t j = 0; j < i; ++j) { + zvec_doc_destroy((*results)[j]); + } + free(*results); + *results = nullptr; + *result_count = 0; + set_last_error("Failed to create document wrapper"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // Copy the C++ document to our wrapper + auto *doc_ptr = reinterpret_cast(c_doc); + *doc_ptr = *internal_doc; // Copy assignment + (*results)[i] = c_doc; // Store the pointer, not dereference + } + + return ZVEC_OK; +} + +// Helper function to convert fetched document results to C API format +static void normalize_nullable_fields_for_fetch( + const zvec::CollectionSchema &schema, zvec::DocPtrMap &doc_map) { + std::vector nullable_fields; + nullable_fields.reserve(schema.fields().size()); + + for (const auto &field : schema.fields()) { + if (field && field->nullable()) { + nullable_fields.push_back(field->name()); + } + } + + if (nullable_fields.empty()) { + return; + } + + for (auto &[_, doc_ptr] : doc_map) { + if (!doc_ptr) { + continue; + } + + for (const auto &field_name : nullable_fields) { + if (!doc_ptr->has(field_name)) { + doc_ptr->set_null(field_name); + } + } + } +} + +zvec_error_code_t convert_fetched_document_results(const zvec::DocPtrMap &doc_map, + zvec_doc_t ***results, + size_t *doc_count) { + // Calculate actual document count (some PKs might not exist) + size_t actual_count = 0; + for (const auto &[pk, doc_ptr] : doc_map) { + if (doc_ptr) { + actual_count++; + } + } + + // Allocate memory for document pointers + *doc_count = actual_count; + if (*doc_count == 0) { + *results = nullptr; + return ZVEC_OK; + } + + *results = static_cast(malloc(*doc_count * sizeof(zvec_doc_t *))); + if (!*results) { + set_last_error("Failed to allocate memory for document pointers"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // Convert C++ DocPtrMap to C zvec_doc_t pointer array + size_t index = 0; + for (const auto &[pk, doc_ptr] : doc_map) { + if (doc_ptr && index < *doc_count) { + // Create new document wrapper + zvec_doc_t *c_doc = zvec_doc_create(); + if (!c_doc) { + // Clean up previously allocated documents + for (size_t j = 0; j < index; ++j) { + zvec_doc_destroy((*results)[j]); + } + free(*results); + *results = nullptr; + *doc_count = 0; + set_last_error("Failed to create document wrapper"); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // Copy the C++ document to our wrapper using copy assignment + auto *cpp_doc_ptr = reinterpret_cast(c_doc); + *cpp_doc_ptr = *doc_ptr; // Copy assignment from shared_ptr + + // Set the primary key explicitly + zvec_doc_set_pk(c_doc, pk.c_str()); + + (*results)[index] = c_doc; + ++index; + } + } + + return ZVEC_OK; +} + +zvec_error_code_t zvec_collection_query(const zvec_collection_t *collection, + const zvec_vector_query_t *query, + zvec_doc_t ***results, + size_t *result_count) { + if (!collection || !query || !results || !result_count) { + set_last_error( + "Invalid arguments: collection, query, results and result_count " + "cannot " + "be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>( + collection); + + // zvec_vector_query_t wraps zvec::SearchQuery internally. + auto *internal_query = + reinterpret_cast(query); + + auto result = (*coll_ptr)->Query(*internal_query); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code == ZVEC_OK) { + const auto &query_results = result.value(); + error_code = + convert_document_results(query_results, results, result_count); + } else { + *results = nullptr; + *result_count = 0; + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_multi_query( + const zvec_collection_t *collection, + const zvec_multi_query_t *query, + zvec_doc_t ***results, size_t *result_count) { + if (!collection || !query || !results || !result_count) { + set_last_error( + "Invalid arguments: collection, query, results and result_count " + "cannot be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception occurred", + auto coll_ptr = + reinterpret_cast *>( + collection); + + auto *internal_query = + reinterpret_cast(query); + + auto result = (*coll_ptr)->Query(*internal_query); + zvec_error_code_t error_code = handle_expected_result(result); + + if (error_code == ZVEC_OK) { + const auto &query_results = result.value(); + error_code = + convert_document_results(query_results, results, result_count); + } else { + *results = nullptr; + *result_count = 0; + } + + return error_code;) +} + +zvec_error_code_t zvec_collection_fetch(zvec_collection_t *collection, + const char *const *pks, size_t pk_count, + const char *const *output_fields, + size_t output_field_count, + bool include_vector, + zvec_doc_t ***results, size_t *doc_count) { + if (!collection || !pks || !results || !doc_count) { + set_last_error( + "Invalid arguments: collection, pks, results and doc_count cannot " + "be null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + // Handle empty case + if (pk_count == 0) { + *results = nullptr; + *doc_count = 0; + return ZVEC_OK; + } + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_fetch", + auto coll_ptr = + reinterpret_cast *>( + collection); + + // Convert C array to C++ vector + std::vector pk_vector; pk_vector.reserve(pk_count); + for (size_t i = 0; i < pk_count; ++i) { + if (pks[i]) { + pk_vector.emplace_back(pks[i]); + } else { + set_last_error("Null primary key at index " + std::to_string(i)); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + } + + // Build optional output_fields + std::optional> cpp_output_fields; + if (output_fields != nullptr && output_field_count > 0) { + std::vector fields; + fields.reserve(output_field_count); + for (size_t i = 0; i < output_field_count; ++i) { + if (output_fields[i]) { + fields.emplace_back(output_fields[i]); + } else { + set_last_error("Null output_field at index " + std::to_string(i)); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + } + cpp_output_fields = std::move(fields); + } + + // Call C++ fetch method + auto result = (*coll_ptr)->Fetch(pk_vector, cpp_output_fields, include_vector); + if (!result.has_value()) { + set_last_error("Failed to fetch documents: " + + result.error().message()); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + auto doc_map = result.value(); + auto schema_result = (*coll_ptr)->Schema(); + if (schema_result.has_value()) { + normalize_nullable_fields_for_fetch(schema_result.value(), doc_map); + } + return convert_fetched_document_results(doc_map, results, doc_count);) +} diff --git a/src/binding/python/CMakeLists.txt b/src/binding/python/CMakeLists.txt new file mode 100644 index 0000000..0db6d75 --- /dev/null +++ b/src/binding/python/CMakeLists.txt @@ -0,0 +1,133 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +find_package(pybind11 REQUIRED) + +set(SRC_LISTS + binding.cc + model/python_collection.cc + model/python_doc.cc + model/python_reranker.cc + model/param/python_param.cc + model/schema/python_schema.cc + model/common/python_config.cc + typing/python_type.cc +) + +pybind11_add_module(_zvec ${SRC_LISTS}) + +# pybind11_add_module() defaults to CXX_VISIBILITY_PRESET=hidden + +# VISIBILITY_INLINES_HIDDEN=ON, which hides the compiler-generated helper +# symbols attached to inline functions (guard variables for static locals, +# vtables, typeinfo ...). The DiskAnn runtime plugin has its own copy of +# Factory::Instance()'s guard variable; if _zvec.so's copy +# is hidden, the two guards are separate and the factory constructor runs +# twice during plugin load, wiping out the registrations that happened +# during _zvec.so import. We switch to default visibility here and rely on +# the version script (exports.map / exports.mac) to keep the dynamic +# symbol table small by exporting only zvec::* and PyInit_*. +set_target_properties(_zvec PROPERTIES + CXX_VISIBILITY_PRESET default + C_VISIBILITY_PRESET default + VISIBILITY_INLINES_HIDDEN OFF) + +# Ensure any change to the linker version script (exports.map) triggers a +# re-link of _zvec.so. target_link_options() alone is a command-line flag +# and does not register the script file as a build dependency, so stale +# wheels would silently keep the old export set otherwise. +if (CMAKE_SYSTEM_NAME STREQUAL "Linux") + set_property(TARGET _zvec APPEND PROPERTY LINK_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/exports.map") +elseif (APPLE) + set_property(TARGET _zvec APPEND PROPERTY LINK_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/exports.mac") +endif() + +if (CMAKE_SYSTEM_NAME STREQUAL "Linux") + target_link_libraries(_zvec PRIVATE + -Wl,--whole-archive + $ + $ + $ + $ + $ + $ + $ + $ + $ + $ + $ + $ + $ + -Wl,--no-whole-archive + zvec + ${CMAKE_DL_LIBS} + ) + target_link_options(_zvec PRIVATE + "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/exports.map" + ) + # DiskAnn is x86-only (it depends on libaio, unavailable on ARM64), so the + # runtime plugin (libzvec_diskann_plugin.so) is only produced on non-ARM + # builds. It is shipped as a runtime-loaded shared module and brought up + # implicitly the first time a DiskAnn index is created — users never need to + # call any load function. If libaio is missing at runtime the auto-load + # fails cleanly and the error is surfaced only when DiskAnn is actually + # used; other index types (HNSW/IVF/Flat/Vamana) remain fully functional. + # The .so must therefore be discoverable next to the extension module, + # hence the $ORIGIN rpath below. + if (NOT CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|arm") + set_target_properties(_zvec PROPERTIES + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + ) + endif() +elseif (APPLE) + target_link_libraries(_zvec PRIVATE + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + zvec + ) + target_link_libraries(_zvec PRIVATE + -Wl,-exported_symbols_list,${CMAKE_CURRENT_SOURCE_DIR}/exports.mac + ) +elseif (MSVC) + set(_zvec_whole_archive_libs + core_knn_flat_static + core_knn_flat_sparse_static + core_knn_hnsw_static + core_knn_hnsw_sparse_static + core_knn_ivf_static + core_knn_vamana_static + core_knn_cluster_static + core_mix_reducer_static + core_metric_static + core_utility_static + core_quantizer_static + core_plugin + ) + target_link_libraries(_zvec PRIVATE + ${_zvec_whole_archive_libs} + zvec + ) + foreach(_lib ${_zvec_whole_archive_libs}) + target_link_options(_zvec PRIVATE + "/WHOLEARCHIVE:$" + ) + endforeach() +endif () + +target_include_directories(_zvec PRIVATE ${PYBIND11_INCLUDE_DIR} ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/binding/python/include) \ No newline at end of file diff --git a/src/binding/python/binding.cc b/src/binding/python/binding.cc new file mode 100644 index 0000000..ecb4d91 --- /dev/null +++ b/src/binding/python/binding.cc @@ -0,0 +1,72 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include "python_collection.h" +#include "python_config.h" +#include "python_doc.h" +#include "python_param.h" +#include "python_reranker.h" +#include "python_schema.h" +#include "python_type.h" + +namespace zvec { + +namespace { + +// Expose DiskAnn plugin management to Python. The DiskAnn runtime normally +// auto-loads on first use, but tests (and diagnostic tooling) need a way to +// force a load up-front and get actionable errors when libaio is missing or +// the plugin shared object cannot be located. +void InitializeDiskAnnPluginBindings(pybind11::module_ &m) { + m.def( + "load_diskann_plugin", + [](const std::string &path) { return ::zvec::LoadDiskAnnPlugin(path); }, + pybind11::arg("path") = std::string(), + "Load the DiskAnn runtime plugin. Returns 0 on success or a negative " + "DiskAnnPluginStatus code on failure (unsupported platform, libaio " + "missing, or dlopen failure)."); + m.def("is_diskann_plugin_loaded", &::zvec::IsDiskAnnPluginLoaded, + "Return True if the DiskAnn runtime plugin is currently loaded."); + m.def("is_libaio_available", &::zvec::IsLibAioAvailable, + "Return True if libaio is resolvable on this host (required by the " + "DiskAnn runtime)."); + + // Status constants so callers can compare against well-known codes without + // hard-coding integers. + m.attr("DISKANN_PLUGIN_OK") = static_cast(::zvec::kDiskAnnPluginOk); + m.attr("DISKANN_PLUGIN_UNSUPPORTED_PLATFORM") = + static_cast(::zvec::kDiskAnnPluginUnsupportedPlatform); + m.attr("DISKANN_PLUGIN_LIBAIO_MISSING") = + static_cast(::zvec::kDiskAnnPluginLibAioMissing); + m.attr("DISKANN_PLUGIN_DLOPEN_FAILED") = + static_cast(::zvec::kDiskAnnPluginDlopenFailed); +} + +} // namespace + +PYBIND11_MODULE(_zvec, m) { + m.doc() = "Zvec core module"; + + ZVecPyTyping::Initialize(m); + ZVecPyParams::Initialize(m); + ZVecPySchemas::Initialize(m); + ZVecPyReranker::Initialize(m); + ZVecPyConfig::Initialize(m); + ZVecPyDoc::Initialize(m); + ZVecPyCollection::Initialize(m); + InitializeDiskAnnPluginBindings(m); +} +} // namespace zvec diff --git a/src/binding/python/exports.mac b/src/binding/python/exports.mac new file mode 100644 index 0000000..f605f7f --- /dev/null +++ b/src/binding/python/exports.mac @@ -0,0 +1 @@ +_PyInit__zvec diff --git a/src/binding/python/exports.map b/src/binding/python/exports.map new file mode 100644 index 0000000..553da00 --- /dev/null +++ b/src/binding/python/exports.map @@ -0,0 +1,34 @@ +{ + global: + # Python module entry point(s). + PyInit_*; + + # Expose the full zvec C++ namespace so the DiskAnn runtime plugin + # (libzvec_diskann_plugin.so), which is dlopen()ed with + # RTLD_NOW | RTLD_GLOBAL after the interpreter has loaded _zvec.so, + # can resolve its undefined references against this module. Without + # this, the plugin fails to load with errors like + # undefined symbol: _ZN4zvec6ailego6Logger10LEVEL_INFOE + # because the default version script hides every internal symbol. + extern "C++" { + "zvec::*"; + zvec::*; + # Also export the compiler-generated helper symbols that live + # alongside symbols in the zvec namespace (guard variables for + # static locals, vtables, typeinfo, VTT, construction vtables, + # thunks). Without these, the DiskAnn plugin and _zvec.so each + # get their own copy of e.g. the guard for Factory::Instance's + # static local, which causes the factory constructor to run + # twice - wiping out registrations done during _zvec.so load. + "guard variable for zvec::*"; + "vtable for zvec::*"; + "VTT for zvec::*"; + "typeinfo for zvec::*"; + "typeinfo name for zvec::*"; + "construction vtable for zvec::*"; + "non-virtual thunk to zvec::*"; + "virtual thunk to zvec::*"; + }; + local: + *; +}; diff --git a/src/binding/python/include/python_collection.h b/src/binding/python/include/python_collection.h new file mode 100644 index 0000000..7c41775 --- /dev/null +++ b/src/binding/python/include/python_collection.h @@ -0,0 +1,36 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License.#pragma once + +#include +#include + +namespace py = pybind11; + +namespace zvec { + +class ZVecPyCollection { + public: + ZVecPyCollection() = delete; + + public: + static void Initialize(py::module_ &m); + + private: + static void bind_db_methods(py::class_ &col); + static void bind_ddl_methods(py::class_ &col); + static void bind_dml_methods(py::class_ &col); + static void bind_dql_methods(py::class_ &col); +}; + +} // namespace zvec diff --git a/src/binding/python/include/python_config.h b/src/binding/python/include/python_config.h new file mode 100644 index 0000000..5cb927c --- /dev/null +++ b/src/binding/python/include/python_config.h @@ -0,0 +1,30 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License.#pragma once + +#include +#include + +namespace py = pybind11; + +namespace zvec { + +class ZVecPyConfig { + public: + ZVecPyConfig() = delete; + + public: + static void Initialize(py::module_ &m); +}; + +} // namespace zvec diff --git a/src/binding/python/include/python_doc.h b/src/binding/python/include/python_doc.h new file mode 100644 index 0000000..c386710 --- /dev/null +++ b/src/binding/python/include/python_doc.h @@ -0,0 +1,34 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License.#pragma once + +#include +#include + +namespace py = pybind11; + +namespace zvec { + +class ZVecPyDoc { + public: + ZVecPyDoc() = delete; + + public: + static void Initialize(py::module_ &m); + + private: + static void bind_doc_operator(py::module_ &m); + static void bind_doc(py::module_ &m); +}; + +} // namespace zvec diff --git a/src/binding/python/include/python_param.h b/src/binding/python/include/python_param.h new file mode 100644 index 0000000..65ac563 --- /dev/null +++ b/src/binding/python/include/python_param.h @@ -0,0 +1,37 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License.#pragma once + +#include +#include +#include + +namespace py = pybind11; + +namespace zvec { + +class ZVecPyParams { + public: + ZVecPyParams() = delete; + + public: + static void Initialize(py::module_ &m); + + private: + static void bind_index_params(py::module_ &m); + static void bind_query_params(py::module_ &m); + static void bind_options(py::module_ &m); + static void bind_vector_query(py::module_ &m); +}; + +} // namespace zvec diff --git a/src/binding/python/include/python_reranker.h b/src/binding/python/include/python_reranker.h new file mode 100644 index 0000000..4ab174a --- /dev/null +++ b/src/binding/python/include/python_reranker.h @@ -0,0 +1,31 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include + +namespace py = pybind11; + +namespace zvec { + +class ZVecPyReranker { + public: + ZVecPyReranker() = delete; + + public: + static void Initialize(py::module_ &m); +}; + +} // namespace zvec diff --git a/src/binding/python/include/python_schema.h b/src/binding/python/include/python_schema.h new file mode 100644 index 0000000..d151ef2 --- /dev/null +++ b/src/binding/python/include/python_schema.h @@ -0,0 +1,35 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License.#pragma once + +#include +#include + +namespace py = pybind11; + +namespace zvec { + +class ZVecPySchemas { + public: + ZVecPySchemas() = delete; + + public: + static void Initialize(py::module_ &m); + + private: + static void bind_field_schema(py::module_ &m); + static void bind_collection_schema(py::module_ &m); + static void bind_collection_stats(py::module_ &m); +}; + +} // namespace zvec diff --git a/src/binding/python/include/python_type.h b/src/binding/python/include/python_type.h new file mode 100644 index 0000000..90a5b67 --- /dev/null +++ b/src/binding/python/include/python_type.h @@ -0,0 +1,38 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License.#pragma once + +#include +#include +#include + +namespace py = pybind11; + +namespace zvec { + +class ZVecPyTyping { + public: + ZVecPyTyping() = delete; + + public: + static void Initialize(py::module_ &m); + + private: + static void bind_datatypes(py::module_ &m); + static void bind_index_types(py::module_ &m); + static void bind_metric_types(py::module_ &m); + static void bind_quantize_types(py::module_ &m); + static void bind_status(py::module_ &m); +}; + +} // namespace zvec diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc new file mode 100644 index 0000000..dade4bb --- /dev/null +++ b/src/binding/python/model/common/python_config.cc @@ -0,0 +1,223 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "python_config.h" +#include + +namespace zvec { + +inline bool has_key(py::dict d, const std::string &key) { + return py::bool_(d.contains(key)); +} + +template +std::optional get_if(py::dict d, const std::string &key) { + if (has_key(d, key)) { + try { + py::object obj = d[py::str(key)]; + return obj.cast(); + } catch (const py::cast_error &) { + throw py::type_error("Key '" + key + "' is not of expected type."); + } + } + return std::nullopt; +} + +inline std::string to_lower(const std::string &s) { + std::string lower; + lower.reserve(s.size()); + std::transform(s.begin(), s.end(), std::back_inserter(lower), ::tolower); + return lower; +} + +inline bool iequals(const std::string &a, const std::string &b) { + return to_lower(a) == to_lower(b); +} + +GlobalConfig::LogLevel str_to_loglevel(const std::string &s) { + if (iequals(s, "debug")) return GlobalConfig::LogLevel::kDebug; + if (iequals(s, "info")) return GlobalConfig::LogLevel::kInfo; + if (iequals(s, "warn") || iequals(s, "warning")) + return GlobalConfig::LogLevel::kWarn; + if (iequals(s, "error")) return GlobalConfig::LogLevel::kError; + if (iequals(s, "fatal")) return GlobalConfig::LogLevel::kFatal; + throw py::value_error("Invalid log level: "); +} + + +void ZVecPyConfig::Initialize(pybind11::module_ &m) { + m.def("Initialize", [](py::args args, py::kwargs kwargs) -> py::none { + py::dict config_dict; + // parse args + for (auto &arg : args) { + if (py::isinstance(arg)) { + for (auto item : arg.cast()) { + config_dict[item.first] = item.second; + } + } else { + throw py::type_error("Positional argument must be a dict if provided"); + } + } + + // parser kwargs + if (kwargs) { + for (auto item : kwargs) { + config_dict[item.first] = item.second; + } + } + + if (config_dict.empty()) { + return py::none(); + } + + GlobalConfig::ConfigData data; + // config memory_limit_mb + if (has_key(config_dict, "memory_limit_mb")) { + auto mb = get_if(config_dict, "memory_limit_mb").value(); + if (mb <= 0) throw py::value_error("memory_limit_mb must be positive"); + data.memory_limit_bytes = static_cast(mb) * 1024 * 1024; + } + + // config log + bool has_log_type = has_key(config_dict, "log_type"); + bool has_log_level = has_key(config_dict, "log_level"); + if (has_log_type || has_log_level) { + std::string log_type = "console"; + std::string log_level_str = "warn"; + + if (has_log_type) { + log_type = config_dict["log_type"].cast(); + } + if (has_log_level) { + log_level_str = config_dict["log_level"].cast(); + } + auto log_level = str_to_loglevel(log_level_str); + if (iequals(log_type, "file")) { + std::string dir = DEFAULT_LOG_DIR; + std::string basename = DEFAULT_LOG_BASENAME; + uint32_t file_size = DEFAULT_LOG_FILE_SIZE; + uint32_t overdue_days = DEFAULT_LOG_OVERDUE_DAYS; + + if (has_key(config_dict, "log_dir")) { + dir = get_if(config_dict, "log_dir").value(); + } + if (has_key(config_dict, "log_basename")) { + basename = get_if(config_dict, "log_basename").value(); + } + if (has_key(config_dict, "log_file_size")) { + auto s = get_if(config_dict, "log_file_size").value(); + if (s <= 0) { + throw py::value_error("log_file_size must be positive"); + } + file_size = static_cast(s); + } + if (has_key(config_dict, "log_overdue_days")) { + std::cout << " ** log_overdue_days: " << overdue_days << std::endl; + auto d = get_if(config_dict, "log_overdue_days").value(); + if (d <= 0) { + throw py::value_error("log_overdue_days must be positive"); + } + overdue_days = static_cast(d); + } + + data.log_config = std::make_shared( + log_level, dir, basename, file_size, overdue_days); + + } else if (iequals(log_type, "console")) { + data.log_config = + std::make_shared(log_level); + } else { + throw py::value_error("log_type must be 'console' or 'file'"); + } + } + + // set query thread count + if (has_key(config_dict, "query_threads")) { + auto q = get_if(config_dict, "query_threads").value(); + if (q <= 0) throw py::value_error("query_threads must be positive"); + data.query_thread_count = static_cast(q); + } + + // set optimize thread count + if (has_key(config_dict, "optimize_threads")) { + auto o = get_if(config_dict, "optimize_threads").value(); + if (o <= 0) throw py::value_error("optimize_threads must be positive"); + data.optimize_thread_count = static_cast(o); + } + + // set invert_to_forward_scan_ratio + if (has_key(config_dict, "invert_to_forward_scan_ratio")) { + auto v = + get_if(config_dict, "invert_to_forward_scan_ratio").value(); + if (v < 0.0 || v > 1.0) { + throw py::value_error( + "invert_to_forward_scan_ratio must be in [0.0, 1.0]"); + } + data.invert_to_forward_scan_ratio = static_cast(v); + } + + // set brute_force_by_keys_ratio + if (has_key(config_dict, "brute_force_by_keys_ratio")) { + auto v = get_if(config_dict, "brute_force_by_keys_ratio").value(); + if (v < 0.0 || v > 1.0) { + throw py::value_error( + "brute_force_by_keys_ratio must be in [0.0, 1.0]"); + } + data.brute_force_by_keys_ratio = static_cast(v); + } + + // set fts_brute_force_by_keys_ratio + if (has_key(config_dict, "fts_brute_force_by_keys_ratio")) { + auto v = + get_if(config_dict, "fts_brute_force_by_keys_ratio").value(); + if (v < 0.0 || v > 1.0) { + throw py::value_error( + "fts_brute_force_by_keys_ratio must be in [0.0, 1.0]"); + } + data.fts_brute_force_by_keys_ratio = static_cast(v); + } + + // jieba_dict_dir: optional override of the SDK-registered default. + // Empty value is a no-op (Initialize preserves the SDK default). + if (has_key(config_dict, "jieba_dict_dir")) { + data.jieba_dict_dir = + get_if(config_dict, "jieba_dict_dir").value(); + } + + // initialize (contains validate) + Status status = GlobalConfig::Instance().Initialize(data); + if (!status.ok()) { + throw std::runtime_error("Initialization failed: " + status.message()); + } + return py::none(); + }); + + // Process-wide setter, independent of Initialize(); called by __init__.py + // on import to register the wheel-bundled dict path. + m.def( + "set_default_jieba_dict_dir", + [](const std::string &dir) { + GlobalConfig::Instance().set_default_jieba_dict_dir(dir); + }, + pybind11::arg("dir"), + "Register the process-wide default jieba dict directory."); + + m.def( + "get_default_jieba_dict_dir", + []() -> std::string { return GlobalConfig::Instance().jieba_dict_dir(); }, + "Read the currently registered default jieba dict directory."); +} + + +} // namespace zvec \ No newline at end of file diff --git a/src/binding/python/model/param/python_param.cc b/src/binding/python/model/param/python_param.cc new file mode 100644 index 0000000..500455c --- /dev/null +++ b/src/binding/python/model/param/python_param.cc @@ -0,0 +1,2269 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "python_param.h" +#include +#include +#include +#include +#include +#include "python_doc.h" + +namespace zvec { + +static std::string index_type_to_string(const IndexType type) { + switch (type) { + case IndexType::INVERT: + return "INVERT"; + case IndexType::FLAT: + return "FLAT"; + case IndexType::IVF: + return "IVF"; + case IndexType::HNSW: + return "HNSW"; + case IndexType::HNSW_RABITQ: + return "HNSW_RABITQ"; + case IndexType::DISKANN: + return "DISKANN"; + case IndexType::VAMANA: + return "VAMANA"; + case IndexType::FTS: + return "FTS"; + default: + return "UNDEFINED"; + } +} + +static std::string metric_type_to_string(const MetricType type) { + switch (type) { + case MetricType::COSINE: + return "COSINE"; + case MetricType::IP: + return "IP"; + case MetricType::L2: + return "L2"; + default: + return "UNDEFINED"; + } +} + +static std::string quantize_type_to_string(const QuantizeType type) { + switch (type) { + case QuantizeType::UNDEFINED: + return "UNDEFINED"; + case QuantizeType::INT8: + return "INT8"; + case QuantizeType::INT4: + return "INT4"; + case QuantizeType::FP16: + return "FP16"; + case QuantizeType::RABITQ: + return "RABITQ"; + default: + return "UNDEFINED"; + } +} + +template +T checked_cast(const py::handle &h, const std::string &vector_field, + const std::string &expected_type) { + try { + return py::cast(h); + } catch (const py::cast_error &e) { + std::string actual_type = std::string(py::str(py::type::of(h))); + std::string msg = + vector_field + ": expected " + expected_type + ", got " + actual_type; + throw py::type_error(msg); + } +} + +template +std::pair serialize_sparse_vector( + const py::dict &sparse_dict, ValueCastFn &&value_caster) { + const size_t n = sparse_dict.size(); + if (n == 0) return {{}, {}}; + + std::string indices_buf; + indices_buf.resize(n * sizeof(uint32_t)); + auto *indices_ptr = reinterpret_cast(indices_buf.data()); + + std::string values_buf; + values_buf.resize(n * sizeof(ValueType)); + auto *values_ptr = reinterpret_cast(values_buf.data()); + + size_t i = 0; + for (const auto &[py_key, py_val] : sparse_dict) { + indices_ptr[i] = checked_cast(py_key, "Sparse indices", "UINT32"); + values_ptr[i] = value_caster(py_val, i); + ++i; + } + return {std::move(indices_buf), std::move(values_buf)}; +} + +void ZVecPyParams::Initialize(pybind11::module_ &parent) { + auto m = + parent.def_submodule("param", "This module contains the params of Zvec"); + + // binding index_params [invert/hnsw/flat/ivf] + bind_index_params(m); + + // bind query_params [hnsw/ivf] + bind_query_params(m); + + // bind options [collection/index/optimize/column] + bind_options(m); + + // bind vector query + bind_vector_query(m); +} + +void ZVecPyParams::bind_index_params(pybind11::module_ &m) { + // binding base index params + py::class_> index_params( + m, "IndexParam", R"pbdoc( +Base class for all index parameter configurations. + +This abstract base class defines the common interface for index types. +It should not be instantiated directly; use derived classes instead. + +Attributes: + type (IndexType): The type of the index (e.g., HNSW, FLAT, INVERT). +)pbdoc"); + index_params + .def_property_readonly( + "type", + [](const IndexParams &self) -> IndexType { return self.type(); }, + "IndexType: The type of the index.") + .def("clone", &IndexParams::clone, py::return_value_policy::copy) + .def( + "__eq__", + [](const IndexParams &self, const py::object &other) { + if (!py::isinstance(other)) return false; + return self == other.cast(); + }, + py::is_operator()) + .def( + "to_dict", + [](const IndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + return dict; + }, + "Convert to dictionary with all fields") + .def(py::pickle( + [](const IndexParams &self) { // __getstate__ + return py::make_tuple(self.type()); + }, + [](py::tuple t) { // __setstate__ + if (t.size() != 1) + throw std::runtime_error("Invalid state for IndexParams"); + return std::shared_ptr(); + })); + + // binding invert index params + py::class_> + invert_params(m, "InvertIndexParam", R"pbdoc( +Parameters for configuring an invert index. + +This class controls whether range query +optimization is enabled for invert index structures. + +Attributes: + type (IndexType): Always `IndexType.INVERTED`. + enable_range_optimization (bool): Whether range optimization is enabled. + enable_extended_wildcard (bool): Whether extended wildcard (suffix and infix) search is enabled. + +Examples: + >>> params = InvertIndexParam(enable_range_optimization=True, enable_extended_wildcard=False) + >>> print(params.enable_range_optimization) + True + >>> print(params.enable_extended_wildcard) + False + >>> config = params.to_dict() + >>> print(config) + {'enable_range_optimization': True, 'enable_extended_wildcard': False} +)pbdoc"); + invert_params + .def(py::init(), py::arg("enable_range_optimization") = false, + py::arg("enable_extended_wildcard") = false, + R"pbdoc( +Constructs an InvertIndexParam instance. + +Args: + enable_range_optimization (bool, optional): If True, enables range query + optimization for the invert index. Defaults to False. + enable_extended_wildcard (bool, optional): If True, enables extended wildcard + search including suffix and infix patterns. Defaults to False. +)pbdoc") + .def_property_readonly("enable_range_optimization", + &InvertIndexParams::enable_range_optimization, + R"pbdoc( +bool: Whether range optimization is enabled for this inverted index. +)pbdoc") + .def_property_readonly("enable_extended_wildcard", + &InvertIndexParams::enable_extended_wildcard, + R"pbdoc( +bool: Whether extended wildcard (suffix and infix) search is enabled. +Note: Prefix search is always enabled regardless of this setting. +)pbdoc") + .def( + "to_dict", + [](const InvertIndexParams &self) -> py::dict { + py::dict dict; + dict["enable_range_optimization"] = + self.enable_range_optimization(); + dict["enable_extended_wildcard"] = self.enable_extended_wildcard(); + return dict; + }, + "Convert to dictionary with all fields") + .def("__repr__", + [](const InvertIndexParams &self) -> std::string { + return "{" + "\"enable_range_optimization\":" + + std::to_string(self.enable_range_optimization()) + + "," + "\"enable_extended_wildcard\":" + + std::to_string(self.enable_extended_wildcard()) + "}"; + }) + .def(py::pickle( + [](const InvertIndexParams &self) { // __getstate__ + return py::make_tuple(self.enable_range_optimization(), + self.enable_extended_wildcard()); + }, + [](py::tuple t) { // __setstate__ + if (t.size() != 2) + throw std::runtime_error("Invalid state for InvertIndexParams"); + return std::make_shared(t[0].cast(), + t[1].cast()); + })); + + // binding fts index params + py::class_> + fts_index_params(m, "FtsIndexParam", R"pbdoc( +Parameters for configuring a full-text search (FTS) index. + +Controls the tokenizer pipeline used during indexing and querying. + +Attributes: + type (IndexType): Always ``IndexType.FTS``. + tokenizer_name (str): Name of the tokenizer (one of "standard", "jieba", + "whitespace"). + Default is "standard". + filters (list[str]): List of token filter names applied after tokenization. + Supported values include "lowercase", "ascii_folding", and "stemmer". + Default is ["lowercase"]. + extra_params (str): Additional tokenizer/filter parameters as an empty + string or JSON object string. Supported keys are grouped by component: + Tokenizers: + standard: + - "max_token_length" (positive integer). + jieba: + - "jieba_dict_dir" (directory containing jieba.dict.utf8 and + hmm_model.utf8). + - "user_dict_path" (user dictionary path). + - "cut_mode" ("search", "mix", "full", or "hmm"; default + "search"). + whitespace: + - no extra_params. + Filters: + lowercase: + - no extra_params. + ascii_folding: + - no extra_params. + stemmer: + - "stemmer_lang" (Snowball language/algorithm; default + "english"), for example {"stemmer_lang":"porter"} for ES + behaviour. + Default is "". + +Examples: + >>> params = FtsIndexParam( + ... tokenizer_name="jieba", filters=["lowercase", "ascii_folding"] + ... ) + >>> print(params.tokenizer_name) + jieba + >>> params = FtsIndexParam( + ... tokenizer_name="standard", + ... filters=["lowercase", "stemmer"], + ... extra_params='{"stemmer_lang":"porter"}', + ... ) +)pbdoc"); + fts_index_params + .def(py::init, std::string>(), + py::arg("tokenizer_name") = "standard", + py::arg("filters") = std::vector{"lowercase"}, + py::arg("extra_params") = "", + R"pbdoc( +Constructs an FtsIndexParam instance. + +Args: + tokenizer_name (str, optional): Tokenizer name. Defaults to "standard". + filters (list[str], optional): Token filter names. Supports "lowercase", + "ascii_folding", and "stemmer". Defaults to ["lowercase"]. + extra_params (str, optional): Extra tokenizer/filter parameters as an empty + string or JSON object string. Supported keys: + Tokenizers: + standard: + - "max_token_length" (positive integer). + jieba: + - "jieba_dict_dir". + - "user_dict_path". + - "cut_mode" ("search", "mix", "full", or "hmm"; default + "search"). + whitespace: + - no extra_params. + Filters: + lowercase: + - no extra_params. + ascii_folding: + - no extra_params. + stemmer: + - "stemmer_lang" (Snowball language/algorithm; default + "english"). + Defaults to "". +)pbdoc") + .def_property_readonly("tokenizer_name", &FtsIndexParams::tokenizer_name, + "str: Name of the tokenizer.") + .def_property_readonly("filters", &FtsIndexParams::filters, + "list[str]: Token filter names.") + .def_property_readonly("extra_params", &FtsIndexParams::extra_params, + "str: Additional tokenizer parameters.") + .def( + "to_dict", + [](const FtsIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["tokenizer_name"] = self.tokenizer_name(); + dict["filters"] = self.filters(); + dict["extra_params"] = self.extra_params(); + return dict; + }, + "Convert to dictionary with all fields") + .def("__repr__", + [](const FtsIndexParams &self) -> std::string { + std::string filters_str = "["; + for (size_t i = 0; i < self.filters().size(); ++i) { + if (i > 0) { + filters_str += ","; + } + filters_str += "\"" + self.filters()[i] + "\""; + } + filters_str += "]"; + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"tokenizer_name\":\"" + self.tokenizer_name() + + "\", \"filters\":" + filters_str + ", \"extra_params\":\"" + + self.extra_params() + "\"}"; + }) + .def(py::pickle( + [](const FtsIndexParams &self) { + return py::make_tuple(self.tokenizer_name(), self.filters(), + self.extra_params()); + }, + [](py::tuple t) { + if (t.size() != 3) { + throw std::runtime_error("Invalid state for FtsIndexParams"); + } + return std::make_shared( + t[0].cast(), t[1].cast>(), + t[2].cast()); + })); + + // binding QuantizerParam + py::class_> quantizer_param( + m, "QuantizerParam", R"pbdoc( +Parameters for quantizer configuration. + +Encapsulates quantization-related settings such as enable_rotate. +Designed for future extensibility. + +Attributes: + enable_rotate (bool): Whether to apply random rotation before INT8/INT4 + quantization to reduce quantization error. + Only effective with quantize_type=INT8 or INT4. Defaults to False. + +Examples: + >>> qp = QuantizerParam(enable_rotate=True) + >>> print(qp.enable_rotate) + True +)pbdoc"); + quantizer_param.def(py::init(), py::arg("enable_rotate") = false) + .def_property_readonly( + "enable_rotate", + [](const QuantizerParam &self) -> bool { + return self.enable_rotate(); + }, + "bool: Whether random rotation is enabled before INT8/INT4 " + "quantization.") + .def( + "to_dict", + [](const QuantizerParam &self) -> py::dict { + py::dict dict; + dict["enable_rotate"] = self.enable_rotate(); + return dict; + }, + "Convert to dictionary with all fields") + .def("__repr__", + [](const QuantizerParam &self) -> std::string { + return "{\"enable_rotate\":" + + std::string(self.enable_rotate() ? "true" : "false") + "}"; + }) + .def( + "__eq__", + [](const QuantizerParam &self, const py::object &other) { + if (!py::isinstance(other)) return false; + return self == other.cast(); + }, + py::is_operator()) + .def(py::pickle( + [](const QuantizerParam &self) { + return py::make_tuple(self.enable_rotate()); + }, + [](py::tuple t) { + if (t.size() != 1) + throw std::runtime_error("Invalid state for QuantizerParam"); + return std::make_shared(t[0].cast()); + })); + + // binding base vector index params + py::class_> + vector_params(m, "VectorIndexParam", R"pbdoc( +Base class for vector index parameter configurations. + +Encapsulates common settings for all vector index types. + +Attributes: + type (IndexType): The specific vector index type (e.g., HNSW, FLAT). + metric_type (MetricType): Distance metric used for similarity search. + quantize_type (QuantizeType): Optional vector quantization type. + quantizer_param (QuantizerParam): Quantizer configuration (e.g., enable_rotate). +)pbdoc"); + vector_params + .def_property_readonly( + "metric_type", + [](const VectorIndexParams &self) -> MetricType { + return self.metric_type(); + }, + "MetricType: Distance metric (e.g., IP, COSINE, L2).") + .def_property_readonly( + "quantize_type", + [](const VectorIndexParams &self) -> QuantizeType { + return self.quantize_type(); + }, + "QuantizeType: Vector quantization type (e.g., FP16, INT8).") + .def_property_readonly( + "quantizer_param", + [](const VectorIndexParams &self) -> QuantizerParam { + return self.quantizer_param(); + }, + "QuantizerParam: Quantizer configuration including enable_rotate.") + .def( + "to_dict", + [](const VectorIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + py::dict qp_dict; + qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); + dict["quantizer_param"] = qp_dict; + return dict; + }, + "Convert to dictionary with all fields") + .def(py::pickle( + [](const VectorIndexParams &self) { // __getstate__ + return py::make_tuple(self.type(), self.metric_type(), + self.quantize_type()); + }, + [](py::tuple t) { // __setstate__ + if (t.size() != 3) + throw std::runtime_error("Invalid state for VectorIndexParams"); + // Base class, cannot instantiate directly, used by subclasses + return std::shared_ptr(); + })); + + // binding hnsw index params + py::class_> + hnsw_params(m, "HnswIndexParam", R"pbdoc( +Parameters for configuring an HNSW (Hierarchical Navigable Small World) index. + +HNSW is a graph-based approximate nearest neighbor search index. This class +encapsulates its construction hyperparameters. + +Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + m (int): Number of bi-directional links created for every new element + during construction. Higher values improve accuracy but increase + memory usage and construction time. Default is 50. + ef_construction (int): Size of the dynamic candidate list for nearest + neighbors during index construction. Larger values yield better + graph quality at the cost of slower build time. Default is 500. + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to + disable quantization. + +Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = HnswIndexParam( + ... metric_type=MetricType.COSINE, + ... m=16, + ... ef_construction=200, + ... quantize_type=QuantizeType.INT8, + ... use_contiguous_memory=True, + ... ) + >>> print(params) + {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8', 'use_contiguous_memory': True} +)pbdoc"); + hnsw_params + .def(py::init([](MetricType metric_type, int m, int ef_construction, + QuantizeType quantize_type, bool use_contiguous_memory, + QuantizerParam quantizer_param) { + return std::make_shared( + metric_type, m, ef_construction, quantize_type, + use_contiguous_memory, quantizer_param); + }), + py::arg("metric_type") = MetricType::IP, + py::arg("m") = core_interface::kDefaultHnswNeighborCnt, + py::arg("ef_construction") = + core_interface::kDefaultHnswEfConstruction, + py::arg("quantize_type") = QuantizeType::UNDEFINED, + py::arg("use_contiguous_memory") = false, + py::arg("quantizer_param") = QuantizerParam()) + .def_property_readonly( + "m", &HnswIndexParams::m, + "int: Maximum number of neighbors per node in upper layers.") + .def_property_readonly( + "ef_construction", &HnswIndexParams::ef_construction, + "int: Candidate list size during index construction.") + .def_property_readonly( + "use_contiguous_memory", &HnswIndexParams::use_contiguous_memory, + "bool: Whether to allocate a single contiguous memory arena for " + "all HNSW graph nodes. Improves cache locality and search " + "throughput at the cost of peak memory usage. Defaults to False.") + .def( + "to_dict", + [](const HnswIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["m"] = self.m(); + dict["ef_construction"] = self.ef_construction(); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + dict["use_contiguous_memory"] = self.use_contiguous_memory(); + py::dict qp_dict; + qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); + dict["quantizer_param"] = qp_dict; + return dict; + }, + "Convert to dictionary with all fields") + .def( + "__repr__", + [](const HnswIndexParams &self) -> std::string { + return "{" + "\"metric_type\":" + + metric_type_to_string(self.metric_type()) + + ", \"m\":" + std::to_string(self.m()) + + ", \"ef_construction\":" + + std::to_string(self.ef_construction()) + + ", \"quantize_type\":" + + quantize_type_to_string(self.quantize_type()) + + ", \"use_contiguous_memory\":" + + (self.use_contiguous_memory() ? "true" : "false") + + ", \"quantizer_param\":{" + "\"enable_rotate\":" + + (self.quantizer_param().enable_rotate() ? "true" : "false") + + "}}"; + }) + .def(py::pickle( + [](const HnswIndexParams &self) { + return py::make_tuple(self.metric_type(), self.m(), + self.ef_construction(), self.quantize_type(), + self.use_contiguous_memory(), + self.quantizer_param().enable_rotate()); + }, + [](py::tuple t) { + if (t.size() != 5 && t.size() != 6) + throw std::runtime_error("Invalid state for HnswIndexParams"); + QuantizerParam qp(t.size() >= 6 ? t[5].cast() : false); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast(), qp); + })); + + // binding hnsw rabitq index params + py::class_> + hnsw_rabitq_params(m, "HnswRabitqIndexParam", R"pbdoc( +Parameters for configuring an HNSW (Hierarchical Navigable Small World) index with RabitQ quantization. + +HNSW is a graph-based approximate nearest neighbor search index. RabitQ is a +quantization method that provides high compression with minimal accuracy loss. + +Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + m (int): Number of bi-directional links created for every new element + during construction. Higher values improve accuracy but increase + memory usage and construction time. Default is 50. + ef_construction (int): Size of the dynamic candidate list for nearest + neighbors during index construction. Larger values yield better + graph quality at the cost of slower build time. Default is 500. + +Examples: + >>> from zvec.typing import MetricType + >>> params = HnswRabitqIndexParam( + ... metric_type=MetricType.COSINE, + ... m=16, + ... ef_construction=200 + ... ) + >>> print(params) + {'metric_type': 'COSINE', 'm': 16, 'ef_construction': 200} +)pbdoc"); + hnsw_rabitq_params + .def(py::init(), + py::arg("metric_type") = MetricType::IP, + py::arg("total_bits") = core_interface::kDefaultRabitqTotalBits, + py::arg("num_clusters") = core_interface::kDefaultRabitqNumClusters, + py::arg("m") = core_interface::kDefaultHnswNeighborCnt, + py::arg("ef_construction") = + core_interface::kDefaultHnswEfConstruction, + py::arg("sample_count") = 0) + .def_property_readonly("m", &HnswRabitqIndexParams::m, + "int: Maximum number of neighbors per node.") + .def_property_readonly( + "ef_construction", &HnswRabitqIndexParams::ef_construction, + "int: Candidate list size during index construction.") + .def_property_readonly("total_bits", &HnswRabitqIndexParams::total_bits, + "int: Total bits for RabitQ quantization.") + .def_property_readonly("num_clusters", + &HnswRabitqIndexParams::num_clusters, + "int: Number of clusters for RabitQ.") + .def_property_readonly("sample_count", + &HnswRabitqIndexParams::sample_count, + "int: Sample count for RabitQ training.") + .def( + "to_dict", + [](const HnswRabitqIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + dict["total_bits"] = self.total_bits(); + dict["num_clusters"] = self.num_clusters(); + dict["sample_count"] = self.sample_count(); + dict["m"] = self.m(); + dict["ef_construction"] = self.ef_construction(); + return dict; + }, + "Convert to dictionary with all fields") + .def( + "__repr__", + [](const HnswRabitqIndexParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"metric_type\":\"" + + metric_type_to_string(self.metric_type()) + + "\", \"total_bits\":" + std::to_string(self.total_bits()) + + ", \"num_clusters\":" + std::to_string(self.num_clusters()) + + ", \"sample_count\":" + std::to_string(self.sample_count()) + + ", \"m\":" + std::to_string(self.m()) + + ", \"ef_construction\":" + + std::to_string(self.ef_construction()) + + ", \"quantize_type\":\"" + + quantize_type_to_string(self.quantize_type()) + "\"}"; + }) + .def(py::pickle( + [](const HnswRabitqIndexParams &self) { + return py::make_tuple(self.metric_type(), self.total_bits(), + self.num_clusters(), self.m(), + self.ef_construction(), self.sample_count()); + }, + [](py::tuple t) { + if (t.size() != 6) + throw std::runtime_error( + "Invalid state for HnswRabitqIndexParams"); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast(), t[5].cast()); + })); + + // binding vamana index params + py::class_> + vamana_params(m, "VamanaIndexParam", R"pbdoc( +Parameters for configuring a Vamana (DiskANN) index. + +Vamana is a single-layer graph-based approximate nearest neighbor search +index originally proposed in the DiskANN paper. This class encapsulates +its construction hyperparameters. + +Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + max_degree (int): Maximum out-degree (R) of every node in the Vamana + graph. Higher values improve recall but increase memory usage and + construction time. Default is 64. + search_list_size (int): Size of the dynamic candidate list during graph + construction (analogous to HNSW's ef_construction). Larger values + yield better graph quality at the cost of slower build time. + Default is 100. + alpha (float): Pruning factor used by Vamana's RobustPrune. Values > 1.0 + keep more long-range edges and improve recall on hard datasets. + Default is 1.2. + saturate_graph (bool): If True, force every node to reach max_degree + neighbors during construction. Default is False. + use_contiguous_memory (bool): If True, allocate a single contiguous + memory arena for all graph nodes, improving cache locality and + search throughput at the cost of peak memory usage. Default is + False. + use_id_map (bool): Reserved flag for engine-level id remapping; the + db layer always supplies consecutive ids so this is currently + ignored by the engine. Default is False. + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED`` + to disable quantization. + +Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = VamanaIndexParam( + ... metric_type=MetricType.COSINE, + ... max_degree=64, + ... search_list_size=128, + ... alpha=1.2, + ... quantize_type=QuantizeType.INT8, + ... ) +)pbdoc"); + vamana_params + .def(py::init([](MetricType metric_type, int max_degree, + int search_list_size, float alpha, bool saturate_graph, + bool use_contiguous_memory, bool use_id_map, + QuantizeType quantize_type, + QuantizerParam quantizer_param) { + return std::make_shared( + metric_type, max_degree, search_list_size, alpha, + saturate_graph, use_contiguous_memory, use_id_map, + quantize_type, quantizer_param); + }), + py::arg("metric_type") = MetricType::IP, + py::arg("max_degree") = core_interface::kDefaultVamanaMaxDegree, + py::arg("search_list_size") = + core_interface::kDefaultVamanaSearchListSize, + py::arg("alpha") = core_interface::kDefaultVamanaAlpha, + py::arg("saturate_graph") = + core_interface::kDefaultVamanaSaturateGraph, + py::arg("use_contiguous_memory") = false, + py::arg("use_id_map") = false, + py::arg("quantize_type") = QuantizeType::UNDEFINED, + py::arg("quantizer_param") = QuantizerParam()) + .def_property_readonly( + "max_degree", &VamanaIndexParams::max_degree, + "int: Maximum out-degree (R) of every node in the Vamana graph.") + .def_property_readonly( + "search_list_size", &VamanaIndexParams::search_list_size, + "int: Candidate list size during Vamana graph construction.") + .def_property_readonly("alpha", &VamanaIndexParams::alpha, + "float: Vamana RobustPrune alpha factor.") + .def_property_readonly( + "saturate_graph", &VamanaIndexParams::saturate_graph, + "bool: Whether to saturate every node to max_degree neighbors.") + .def_property_readonly( + "use_contiguous_memory", &VamanaIndexParams::use_contiguous_memory, + "bool: Whether to allocate a single contiguous memory arena for " + "all Vamana graph nodes. Improves cache locality and search " + "throughput at the cost of peak memory usage. Defaults to False.") + .def_property_readonly( + "use_id_map", &VamanaIndexParams::use_id_map, + "bool: Reserved flag for engine-level id remapping. Currently " + "ignored by the engine because the db layer always supplies " + "consecutive ids.") + .def( + "to_dict", + [](const VamanaIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["max_degree"] = self.max_degree(); + dict["search_list_size"] = self.search_list_size(); + dict["alpha"] = self.alpha(); + dict["saturate_graph"] = self.saturate_graph(); + dict["use_contiguous_memory"] = self.use_contiguous_memory(); + dict["use_id_map"] = self.use_id_map(); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + py::dict qp_dict; + qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); + dict["quantizer_param"] = qp_dict; + return dict; + }, + "Convert to dictionary with all fields") + .def( + "__repr__", + [](const VamanaIndexParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"metric_type\":\"" + + metric_type_to_string(self.metric_type()) + + "\", \"max_degree\":" + std::to_string(self.max_degree()) + + ", \"search_list_size\":" + + std::to_string(self.search_list_size()) + + ", \"alpha\":" + std::to_string(self.alpha()) + + ", \"saturate_graph\":" + + std::string(self.saturate_graph() ? "true" : "false") + + ", \"use_contiguous_memory\":" + + std::string(self.use_contiguous_memory() ? "true" + : "false") + + ", \"use_id_map\":" + + std::string(self.use_id_map() ? "true" : "false") + + ", \"quantize_type\":\"" + + quantize_type_to_string(self.quantize_type()) + + "\", \"quantizer_param\":{" + "\"enable_rotate\":" + + (self.quantizer_param().enable_rotate() ? "true" : "false") + + "}}"; + }) + .def(py::pickle( + [](const VamanaIndexParams &self) { + return py::make_tuple( + self.metric_type(), self.max_degree(), self.search_list_size(), + self.alpha(), self.saturate_graph(), + self.use_contiguous_memory(), self.use_id_map(), + self.quantize_type(), self.quantizer_param().enable_rotate()); + }, + [](py::tuple t) { + if (t.size() != 8 && t.size() != 9) + throw std::runtime_error("Invalid state for VamanaIndexParams"); + QuantizerParam qp(t.size() >= 9 ? t[8].cast() : false); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast(), t[5].cast(), + t[6].cast(), t[7].cast(), qp); + })); + + // FlatIndexParams + py::class_> + flat_params(m, "FlatIndexParam", R"pbdoc( +Parameters for configuring a flat (brute-force) index. + +A flat index performs exact nearest neighbor search by comparing the query +vector against all vectors in the collection. It is simple, accurate, and +suitable for small to medium datasets or as a baseline. + +Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Use ``QuantizeType.UNDEFINED`` to + disable quantization. Default is ``QuantizeType.UNDEFINED``. + +Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = FlatIndexParam( + ... metric_type=MetricType.L2, + ... quantize_type=QuantizeType.FP16 + ... ) + >>> print(params) + {'metric_type': 'L2', 'quantize_type': 'FP16'} +)pbdoc"); + flat_params + .def(py::init([](MetricType metric_type, QuantizeType quantize_type, + QuantizerParam quantizer_param) { + return std::make_shared( + metric_type, quantize_type, quantizer_param); + }), + py::arg("metric_type") = MetricType::IP, + py::arg("quantize_type") = QuantizeType::UNDEFINED, + py::arg("quantizer_param") = QuantizerParam(), + R"pbdoc( +Constructs a FlatIndexParam instance. + +Args: + metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP. + quantize_type (QuantizeType, optional): Vector quantization type. + Defaults to QuantizeType.UNDEFINED (no quantization). + quantizer_param (QuantizerParam, optional): Quantizer configuration. + Defaults to QuantizerParam(). +)pbdoc") + .def( + "to_dict", + [](const FlatIndexParams &self) -> py::dict { + py::dict dict; + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + py::dict qp_dict; + qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); + dict["quantizer_param"] = qp_dict; + return dict; + }, + "Convert to dictionary with all fields") + .def( + "__repr__", + [](const FlatIndexParams &self) -> std::string { + return "{" + "\"metric_type\":" + + metric_type_to_string(self.metric_type()) + + ", \"quantize_type\":" + + quantize_type_to_string(self.quantize_type()) + + ", \"quantizer_param\":{" + "\"enable_rotate\":" + + (self.quantizer_param().enable_rotate() ? "true" : "false") + + "}}"; + }) + .def(py::pickle( + [](const FlatIndexParams &self) { + return py::make_tuple(self.metric_type(), self.quantize_type(), + self.quantizer_param().enable_rotate()); + }, + [](py::tuple t) { + if (t.size() != 2 && t.size() != 3) + throw std::runtime_error("Invalid state for FlatIndexParams"); + QuantizerParam qp(t.size() >= 3 ? t[2].cast() : false); + return std::make_shared( + t[0].cast(), t[1].cast(), qp); + })); + + // IVFIndexParams + py::class_> + ivf_params(m, "IVFIndexParam", R"pbdoc( +Parameters for configuring an IVF (Inverted File Index) index. + +IVF partitions the vector space into clusters (inverted lists). At query time, +only a subset of clusters is searched, providing a trade-off between speed +and accuracy. + +Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + n_list (int): Number of clusters (inverted lists) to partition the dataset into. + Default is 10. + n_iters (int): Number of iterations for k-means clustering during index training. + Higher values yield more stable centroids. Default is 10. + use_soar (bool): Whether to enable SOAR (Scalable Optimized Adaptive Routing) + for improved IVF search performance. Default is False. + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED``. + +Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = IVFIndexParam( + ... metric_type=MetricType.COSINE, + ... n_list=100, + ... n_iters=15, + ... use_soar=True, + ... quantize_type=QuantizeType.INT8 + ... ) + >>> print(params.n_list) + 100 +)pbdoc"); + ivf_params + .def(py::init([](MetricType metric_type, int n_list, int n_iters, + bool use_soar, QuantizeType quantize_type, + QuantizerParam quantizer_param) { + return std::make_shared( + metric_type, n_list, n_iters, use_soar, quantize_type, + quantizer_param); + }), + py::arg("metric_type") = MetricType::IP, py::arg("n_list") = 10, + py::arg("n_iters") = 10, py::arg("use_soar") = false, + py::arg("quantize_type") = QuantizeType::UNDEFINED, + py::arg("quantizer_param") = QuantizerParam(), + R"pbdoc( +Constructs an IVFIndexParam instance. + +Args: + metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP. + n_list (int, optional): Number of inverted lists (clusters). + Defaults to 10. + n_iters (int, optional): Number of k-means iterations during training. + Defaults to 10. + use_soar (bool, optional): Enable SOAR optimization. Defaults to False. + quantize_type (QuantizeType, optional): Vector quantization type. + Defaults to QuantizeType.UNDEFINED. + quantizer_param (QuantizerParam, optional): Quantizer configuration. + Defaults to QuantizerParam(). +)pbdoc") + .def_property_readonly("n_list", &IVFIndexParams::n_list, + "int: Number of inverted lists.") + .def_property_readonly( + "n_iters", &IVFIndexParams::n_iters, + "int: Number of k-means iterations during training.") + .def_property_readonly("use_soar", &IVFIndexParams::use_soar, + "bool: Whether SOAR optimization is enabled.") + .def( + "to_dict", + [](const IVFIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["n_list"] = self.n_list(); + dict["n_iters"] = self.n_iters(); + dict["use_soar"] = self.use_soar(); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + py::dict qp_dict; + qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); + dict["quantizer_param"] = qp_dict; + return dict; + }, + "Convert to dictionary with all fields") + .def( + "__repr__", + [](const IVFIndexParams &self) { + return "{" + "\"metric_type\":" + + metric_type_to_string(self.metric_type()) + + ", \"n_list\":" + std::to_string(self.n_list()) + + ", \"n_iters\":" + std::to_string(self.n_iters()) + + ", \"use_soar\":" + std::to_string(self.use_soar()) + + ", \"quantize_type\":" + + quantize_type_to_string(self.quantize_type()) + + ", \"quantizer_param\":{" + "\"enable_rotate\":" + + (self.quantizer_param().enable_rotate() ? "true" : "false") + + "}}"; + }) + .def(py::pickle( + [](const IVFIndexParams &self) { + return py::make_tuple(self.metric_type(), self.n_list(), + self.n_iters(), self.use_soar(), + self.quantize_type(), + self.quantizer_param().enable_rotate()); + }, + [](py::tuple t) { + if (t.size() != 5 && t.size() != 6) + throw std::runtime_error("Invalid state for IVFIndexParams"); + QuantizerParam qp(t.size() >= 6 ? t[5].cast() : false); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast(), qp); + })); + + // DiskAnnIndexParams + py::class_> + diskann_params(m, "DiskAnnIndexParam", R"pbdoc( +Parameters for configuring an DiskAnn index. + +DiskAnn stores compressed vector in memory and high-definition vector on disk. At query time, +only compressed vector will be loaded into memory. By this way, search memory at runtime is diminished. + +Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + max_degree (int): Maximum out-degree of each node in the Vamana graph. + Larger values improve recall at the cost of build time and index size. + Clamped to the range [1, 100]. Default is 100. + list_size (int): Candidate list size used during graph construction. + Larger values improve graph quality and recall at the cost of build time. + Clamped to the range [10, 100]. Default is 50. + pq_chunk_num (int): Number of PQ chunks used for product-quantizing the + in-memory compressed vectors. ``0`` means auto-pick based on dimension. + Clamped to the range [1, 1024]. Default is 0. + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED``. + +Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = DiskAnnIndexParam( + ... metric_type=MetricType.COSINE, + ... max_degree=100, + ... list_size=50, + ... pq_chunk_num=8, + ... quantize_type=QuantizeType.FP16 + ... ) + >>> print(params.max_degree) + 100 +)pbdoc"); + diskann_params + .def(py::init([](MetricType metric_type, int max_degree, int list_size, + int pq_chunk_num, QuantizeType quantize_type, + QuantizerParam quantizer_param) { + return std::make_shared( + metric_type, max_degree, list_size, pq_chunk_num, + quantize_type, quantizer_param); + }), + py::arg("metric_type") = MetricType::IP, py::arg("max_degree") = 100, + py::arg("list_size") = 50, py::arg("pq_chunk_num") = 0, + py::arg("quantize_type") = QuantizeType::UNDEFINED, + py::arg("quantizer_param") = QuantizerParam(), + R"pbdoc( +Constructs an DiskAnnIndexParams instance. + +Args: + metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP. + max_degree (int, optional): Maximum out-degree of each node in the Vamana + graph. Clamped to [1, 100]. Defaults to 100. + list_size (int, optional): Candidate list size used during graph + construction. Clamped to [10, 100]. Defaults to 50. + pq_chunk_num (int, optional): Number of PQ chunks for product + quantization. ``0`` means auto-pick based on dimension. + Clamped to [1, 1024]. Defaults to 0. + quantize_type (QuantizeType, optional): Vector quantization type. + Defaults to QuantizeType.UNDEFINED. + quantizer_param (QuantizerParam, optional): Quantizer configuration. + Defaults to QuantizerParam(). +)pbdoc") + .def_property_readonly("max_degree", &DiskAnnIndexParams::max_degree, + "int: max node degree.") + .def_property_readonly("list_size", &DiskAnnIndexParams::list_size, + "int: list size of graph construction") + .def_property_readonly( + "pq_chunk_num", + [](const DiskAnnIndexParams &self) -> int { + return self.pq_chunk_num(); + }, + "int: chunk num of production quantization.") + .def( + "to_dict", + [](const DiskAnnIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["max_degree"] = self.max_degree(); + dict["list_size"] = self.list_size(); + dict["pq_chunk_num"] = self.pq_chunk_num(); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + py::dict qp_dict; + qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); + dict["quantizer_param"] = qp_dict; + return dict; + }, + "Convert to dictionary with all fields") + .def( + "__repr__", + [](const DiskAnnIndexParams &self) { + return "{" + "\"metric_type\":" + + metric_type_to_string(self.metric_type()) + + ", \"max_degree\":" + std::to_string(self.max_degree()) + + ", \"list_size\":" + std::to_string(self.list_size()) + + ", \"pq_chunk_num\":" + std::to_string(self.pq_chunk_num()) + + ", \"quantize_type\":" + + quantize_type_to_string(self.quantize_type()) + + ", \"quantizer_param\":{" + "\"enable_rotate\":" + + (self.quantizer_param().enable_rotate() ? "true" : "false") + + "}}"; + }) + .def(py::pickle( + [](const DiskAnnIndexParams &self) { + return py::make_tuple(self.metric_type(), self.max_degree(), + self.list_size(), self.pq_chunk_num(), + self.quantize_type(), + self.quantizer_param().enable_rotate()); + }, + [](py::tuple t) { + if (t.size() != 5 && t.size() != 6) + throw std::runtime_error("Invalid state for DiskAnnIndexParams"); + QuantizerParam qp(t.size() >= 6 ? t[5].cast() : false); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast(), qp); + })); +} + +void ZVecPyParams::bind_query_params(py::module_ &m) { + // binding base query params + py::class_> query_params( + m, "QueryParam", R"pbdoc( +Base class for all query parameter configurations. + +This abstract base class defines common query settings such as search radius +and whether to force linear (brute-force) search. It should not be instantiated +directly; use derived classes like `HnswQueryParam` or `IVFQueryParam`. + +Attributes: + type (IndexType): The index type this query is configured for. + radius (float): Search radius for range queries. Used in combination with + top-k to filter results. Default is 0.0 (disabled). + is_linear (bool): If True, forces brute-force linear search instead of + using the index. Useful for debugging or small datasets. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. +)pbdoc"); + query_params + .def_property_readonly( + "type", + [](const QueryParams &self) -> IndexType { return self.type(); }, + "IndexType: The type of index this query targets.") + .def_property_readonly( + "radius", + [](const QueryParams &self) -> float { return self.radius(); }, + "IndexType: The type of index this query targets.") + .def_property_readonly( + "is_linear", + [](const QueryParams &self) -> bool { return self.is_linear(); }, + "bool: Whether to bypass the index and use brute-force linear " + "search.") + .def_property_readonly( + "is_using_refiner", + [](const QueryParams &self) -> bool { + return self.is_using_refiner(); + }, + "bool: Whether to use refiner for the query.") + .def(py::pickle( + [](const QueryParams &self) { // __getstate__ + return py::make_tuple(self.type(), self.radius(), self.is_linear()); + }, + [](py::tuple t) { // __setstate__ + if (t.size() != 3) + throw std::runtime_error("Invalid state for QueryParams"); + return std::shared_ptr(); + })); + + // binding hnsw query params + py::class_> + hnsw_params(m, "HnswQueryParam", R"pbdoc( +Query parameters for HNSW (Hierarchical Navigable Small World) index. + +Controls the trade-off between search speed and accuracy via the `ef` parameter. + +Attributes: + type (IndexType): Always ``IndexType.HNSW``. + ef (int): Size of the dynamic candidate list during search. + Larger values improve recall but slow down search. + Default is 300. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + +Examples: + >>> params = HnswQueryParam(ef=300) + >>> print(params.ef) + 300 + >>> print(params.to_dict() if hasattr(params, 'to_dict') else params) + {"type":"HNSW", "ef":300} +)pbdoc"); + hnsw_params + .def(py::init([](int ef, float radius, bool is_linear, + bool is_using_refiner, py::dict extra_params) { + auto obj = std::make_shared(ef, radius, is_linear, + is_using_refiner); + if (extra_params.contains("prefetch_offset")) { + obj->set_prefetch_offset( + extra_params["prefetch_offset"].cast()); + } + if (extra_params.contains("prefetch_lines")) { + obj->set_prefetch_lines( + extra_params["prefetch_lines"].cast()); + } + return obj; + }), + py::arg("ef") = core_interface::kDefaultHnswEfSearch, + py::arg("radius") = 0.0f, py::arg("is_linear") = false, + py::arg("is_using_refiner") = false, + py::arg("extra_params") = py::dict(), + R"pbdoc( +Constructs an HnswQueryParam instance. + +Args: + ef (int, optional): Search-time candidate list size. + Higher values improve accuracy. Defaults to 100. + radius (float, optional): Search radius for range queries. Default is 0.0. + is_linear (bool, optional): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + extra_params (dict, optional): Additional search parameters. Supported keys: + - ``prefetch_offset`` (int): Graph prefetch offset (PO). + ``0`` disables prefetching. Default is ``8``. + Values are clamped to ``256``. + - ``prefetch_lines`` (int): Number of 64B cache lines to prefetch + per neighbour vector (PL). ``0`` (default) uses the auto-derived + value ``ceil(vector_size/64)``. Values are clamped to ``256``. +)pbdoc") + .def_property_readonly( + "ef", [](const HnswQueryParams &self) -> int { return self.ef(); }, + "int: Size of the dynamic candidate list during HNSW search.") + .def_property_readonly( + "prefetch_offset", + [](const HnswQueryParams &self) -> uint32_t { + return self.prefetch_offset(); + }, + "int: Graph prefetch offset used by the HNSW fast path.") + .def_property_readonly( + "prefetch_lines", + [](const HnswQueryParams &self) -> uint32_t { + return self.prefetch_lines(); + }, + "int: Override of prefetch cache lines per vector (0=auto).") + .def("__repr__", + [](const HnswQueryParams &self) -> std::string { + return "{" + "\"type\":" + + index_type_to_string(self.type()) + + ", \"ef\":" + std::to_string(self.ef()) + + ", \"radius\":" + std::to_string(self.radius()) + + ", \"is_linear\":" + std::to_string(self.is_linear()) + + ", \"is_using_refiner\":" + + std::to_string(self.is_using_refiner()) + + ", \"prefetch_offset\":" + + std::to_string(self.prefetch_offset()) + + ", \"prefetch_lines\":" + + std::to_string(self.prefetch_lines()) + "}"; + }) + .def(py::pickle( + [](const HnswQueryParams &self) { + return py::make_tuple(self.ef(), self.radius(), self.is_linear(), + self.is_using_refiner(), + self.prefetch_offset(), + self.prefetch_lines()); + }, + [](py::tuple t) { + if (t.size() != 4 && t.size() != 5 && t.size() != 6) + throw std::runtime_error("Invalid state for HnswQueryParams"); + auto obj = std::make_shared(t[0].cast()); + obj->set_radius(t[1].cast()); + obj->set_is_linear(t[2].cast()); + obj->set_is_using_refiner(t[3].cast()); + if (t.size() >= 5) { + obj->set_prefetch_offset(t[4].cast()); + } + if (t.size() >= 6) { + obj->set_prefetch_lines(t[5].cast()); + } + return obj; + })); + + // binding ivf query params + py::class_> + ivf_params(m, "IVFQueryParam", R"pbdoc( +Query parameters for IVF (Inverted File Index) index. + +Controls how many inverted lists (`nprobe`) to visit during search. + +Attributes: + type (IndexType): Always ``IndexType.IVF``. + nprobe (int): Number of closest clusters (inverted lists) to search. + Higher values improve recall but increase latency. + Default is 10. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + +Examples: + >>> params = IVFQueryParam(nprobe=20) + >>> print(params.nprobe) + 20 +)pbdoc"); + ivf_params + .def(py::init(), py::arg("nprobe") = 10, R"pbdoc( +Constructs an IVFQueryParam instance. + +Args: + nprobe (int, optional): Number of inverted lists to probe during search. + Higher values improve accuracy. Defaults to 10. +)pbdoc") + .def_property_readonly( + "nprobe", + [](const IVFQueryParams &self) -> int { return self.nprobe(); }, + "int: Number of inverted lists to search during IVF query.") + .def("__repr__", + [](const IVFQueryParams &self) -> std::string { + return "{" + "\"type\":" + + index_type_to_string(self.type()) + + ", \"nprobe\":" + std::to_string(self.nprobe()) + "}"; + }) + .def(py::pickle( + [](const IVFQueryParams &self) { + return py::make_tuple(self.nprobe(), self.radius(), + self.is_linear()); + }, + [](py::tuple t) { + if (t.size() != 3) + throw std::runtime_error("Invalid state for IVFQueryParams"); + auto obj = std::make_shared(t[0].cast()); + obj->set_radius(t[1].cast()); + obj->set_is_linear(t[2].cast()); + return obj; + })); + + // binding hnsw rabitq query params + py::class_> + hnsw_rabitq_query_params(m, "HnswRabitqQueryParam", R"pbdoc( +Query parameters for HNSW RaBitQ (Hierarchical Navigable Small World with RaBitQ quantization) index. + +Controls the trade-off between search speed and accuracy via the `ef` parameter. +RaBitQ provides efficient quantization while maintaining high search quality. + +Attributes: + type (IndexType): Always ``IndexType.HNSW_RABITQ``. + ef (int): Size of the dynamic candidate list during search. + Larger values improve recall but slow down search. + Default is 300. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + +Examples: + >>> params = HnswRabitqQueryParam(ef=300) + >>> print(params.ef) + 300 + >>> print(params.to_dict() if hasattr(params, 'to_dict') else params) + {"type":"HNSW_RABITQ", "ef":300} +)pbdoc"); + hnsw_rabitq_query_params + .def(py::init(), + py::arg("ef") = core_interface::kDefaultHnswEfSearch, + py::arg("radius") = 0.0f, py::arg("is_linear") = false, + py::arg("is_using_refiner") = false, + R"pbdoc( +Constructs an HnswRabitqQueryParam instance. + +Args: + ef (int, optional): Search-time candidate list size. + Higher values improve accuracy. Defaults to 300. + radius (float, optional): Search radius for range queries. Default is 0.0. + is_linear (bool, optional): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. +)pbdoc") + .def_property_readonly( + "ef", + [](const HnswRabitqQueryParams &self) -> int { return self.ef(); }, + "int: Size of the dynamic candidate list during HNSW RaBitQ search.") + .def("__repr__", + [](const HnswRabitqQueryParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"ef\":" + std::to_string(self.ef()) + + ", \"radius\":" + std::to_string(self.radius()) + + ", \"is_linear\":" + std::to_string(self.is_linear()) + + ", \"is_using_refiner\":" + + std::to_string(self.is_using_refiner()) + "}"; + }) + .def(py::pickle( + [](const HnswRabitqQueryParams &self) { + return py::make_tuple(self.ef(), self.radius(), self.is_linear(), + self.is_using_refiner()); + }, + [](py::tuple t) { + if (t.size() != 4) + throw std::runtime_error( + "Invalid state for HnswRabitqQueryParams"); + auto obj = + std::make_shared(t[0].cast()); + obj->set_radius(t[1].cast()); + obj->set_is_linear(t[2].cast()); + obj->set_is_using_refiner(t[3].cast()); + return obj; + })); + + // binding diskann query params + py::class_> + diskann_params(m, "DiskAnnQueryParam", R"pbdoc( +Query parameters for DiskAnn index. + +Attributes: + type (IndexType): Always ``IndexType.DISKANN``. + list_size (int): Beam-search candidate list size used at query time. + Higher values improve recall but increase latency. Default is 10. + +Examples: + >>> params = DiskAnnQueryParam(list_size=20) + >>> print(params.list_size) + 20 +)pbdoc"); + diskann_params + .def(py::init(), py::arg("list_size") = 300, R"pbdoc( +Constructs an DiskAnnQueryParams instance. + +Args: + list_size (int, optional): Beam-search candidate list size during + graph search. Higher values improve recall at the cost of latency. + Defaults to 300. +)pbdoc") + .def_property_readonly( + "list_size", + [](const DiskAnnQueryParams &self) -> int { + return self.list_size(); + }, + "int: Beam-search candidate list size during DiskAnn query.") + .def("__repr__", + [](const DiskAnnQueryParams &self) -> std::string { + return "{" + "\"type\":" + + index_type_to_string(self.type()) + + ", \"list_size\":" + std::to_string(self.list_size()) + "}"; + }) + .def(py::pickle( + [](const DiskAnnQueryParams &self) { + return py::make_tuple(self.list_size()); + }, + [](py::tuple t) { + if (t.size() != 1) + throw std::runtime_error("Invalid state for DiskAnnQueryParams"); + return std::make_shared(t[0].cast()); + })); + + // binding vamana query params + py::class_> + vamana_query_params(m, "VamanaQueryParam", R"pbdoc( +Query parameters for the Vamana (DiskANN) index. + +Controls the trade-off between search speed and accuracy via the +``ef_search`` parameter, which sets the size of the dynamic candidate list +explored during search. + +Attributes: + type (IndexType): Always ``IndexType.VAMANA``. + ef_search (int): Size of the dynamic candidate list during Vamana + search. Larger values improve recall but slow down search. + Default is 200. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. + Default is False. + +Examples: + >>> params = VamanaQueryParam(ef_search=200) + >>> print(params.ef_search) + 200 +)pbdoc"); + vamana_query_params + .def(py::init([](int ef_search, float radius, bool is_linear, + bool is_using_refiner, py::dict extra_params) { + auto obj = std::make_shared( + ef_search, radius, is_linear, is_using_refiner); + if (extra_params.contains("prefetch_offset")) { + obj->set_prefetch_offset( + extra_params["prefetch_offset"].cast()); + } + if (extra_params.contains("prefetch_lines")) { + obj->set_prefetch_lines( + extra_params["prefetch_lines"].cast()); + } + return obj; + }), + py::arg("ef_search") = core_interface::kDefaultVamanaEfSearch, + py::arg("radius") = 0.0f, py::arg("is_linear") = false, + py::arg("is_using_refiner") = false, + py::arg("extra_params") = py::dict(), + R"pbdoc( +Constructs a VamanaQueryParam instance. + +Args: + ef_search (int, optional): Search-time candidate list size. + Higher values improve accuracy. Defaults to 200. + radius (float, optional): Search radius for range queries. Default is 0.0. + is_linear (bool, optional): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. + Default is False. + extra_params (dict, optional): Additional search parameters. Supported keys: + - ``prefetch_offset`` (int): Graph prefetch offset (PO). + ``0`` disables prefetching. Default is ``8``. + Values are clamped to ``256``. + - ``prefetch_lines`` (int): Number of 64B cache lines to prefetch + per neighbour vector (PL). ``0`` (default) uses the auto-derived + value ``ceil(dim/64)``. Values are clamped to ``256``. +)pbdoc") + .def_property_readonly( + "ef_search", + [](const VamanaQueryParams &self) -> int { return self.ef_search(); }, + "int: Size of the dynamic candidate list during Vamana search.") + .def_property_readonly( + "prefetch_offset", + [](const VamanaQueryParams &self) -> uint32_t { + return self.prefetch_offset(); + }, + "int: Graph prefetch offset used by the Vamana fast path.") + .def_property_readonly( + "prefetch_lines", + [](const VamanaQueryParams &self) -> uint32_t { + return self.prefetch_lines(); + }, + "int: Override of prefetch cache lines per vector (0=auto).") + .def("__repr__", + [](const VamanaQueryParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"ef_search\":" + std::to_string(self.ef_search()) + + ", \"radius\":" + std::to_string(self.radius()) + + ", \"is_linear\":" + std::to_string(self.is_linear()) + + ", \"is_using_refiner\":" + + std::to_string(self.is_using_refiner()) + + ", \"prefetch_offset\":" + + std::to_string(self.prefetch_offset()) + + ", \"prefetch_lines\":" + + std::to_string(self.prefetch_lines()) + "}"; + }) + .def(py::pickle( + [](const VamanaQueryParams &self) { + return py::make_tuple(self.ef_search(), self.radius(), + self.is_linear(), self.is_using_refiner(), + self.prefetch_offset(), + self.prefetch_lines()); + }, + [](py::tuple t) { + if (t.size() != 4 && t.size() != 5 && t.size() != 6) + throw std::runtime_error("Invalid state for VamanaQueryParams"); + auto obj = std::make_shared(t[0].cast()); + obj->set_radius(t[1].cast()); + obj->set_is_linear(t[2].cast()); + obj->set_is_using_refiner(t[3].cast()); + if (t.size() >= 5) { + obj->set_prefetch_offset(t[4].cast()); + } + if (t.size() >= 6) { + obj->set_prefetch_lines(t[5].cast()); + } + return obj; + })); + + // binding fts query params + py::class_> + fts_query_params(m, "FtsQueryParam", R"pbdoc( +Query parameters for full-text search (FTS) index. + +Controls the default boolean operator used to combine adjacent bare terms +in a query string. + +Attributes: + type (IndexType): Always ``IndexType.FTS``. + default_operator (str): Default boolean operator for adjacent bare terms. + Supported values (case-insensitive): "OR" (default), "AND". + +Examples: + >>> params = FtsQueryParam(default_operator="AND") + >>> print(params.default_operator) + AND +)pbdoc"); + fts_query_params + .def(py::init([](const std::string &default_operator) { + auto params = std::make_shared(); + if (!default_operator.empty()) { + params->set_default_operator(default_operator); + } + return params; + }), + py::arg("default_operator") = "", + R"pbdoc( +Constructs an FtsQueryParam instance. + +Args: + default_operator (str, optional): Default boolean operator for adjacent + bare terms. Supported: "OR", "AND". Defaults to "" (uses engine default). +)pbdoc") + .def_property_readonly("default_operator", + &FtsQueryParams::default_operator, + "str: Default boolean operator for bare terms.") + .def("__repr__", + [](const FtsQueryParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"default_operator\":\"" + self.default_operator() + + "\"}"; + }) + .def(py::pickle( + [](const FtsQueryParams &self) { + return py::make_tuple(self.default_operator()); + }, + [](py::tuple t) { + if (t.size() != 1) { + throw std::runtime_error("Invalid state for FtsQueryParams"); + } + auto obj = std::make_shared(); + obj->set_default_operator(t[0].cast()); + return obj; + })); +} + +void ZVecPyParams::bind_options(py::module_ &m) { // binding collection options + py::class_(m, "CollectionOption", R"pbdoc( +Options for opening or creating a collection. + +Attributes: + read_only (bool): Whether the collection is opened in read-only mode. + Default is False. + enable_mmap (bool): Whether to use memory-mapped I/O for data files. + Default is True. + +Examples: + >>> opt = CollectionOption(read_only=True, enable_mmap=False) + >>> print(opt.read_only) + True +)pbdoc") + .def(py::init(), py::arg("read_only") = false, + py::arg("enable_mmap") = true, + R"pbdoc( +Constructs a CollectionOption instance. + +Args: + read_only (bool, optional): Open collection in read-only mode. + Defaults to False. + enable_mmap (bool, optional): Enable memory-mapped I/O. + Defaults to True. +)pbdoc") + .def_property_readonly( + "enable_mmap", + [](const CollectionOptions &self) { return self.enable_mmap_; }) + .def_property_readonly( + "read_only", + [](const CollectionOptions &self) { return self.read_only_; }) + .def("__repr__", + [](const CollectionOptions &self) -> std::string { + return "{" + "\"enable_mmap\":" + + std::to_string(self.enable_mmap_) + + ", \"read_only\":" + std::to_string(self.read_only_) + "}"; + }) + .def(py::pickle( + [](const CollectionOptions &self) { + return py::make_tuple(self.read_only_, self.enable_mmap_, + self.max_buffer_size_); + }, + [](py::tuple t) { + if (t.size() != 3) + throw std::runtime_error( + "Invalid pickle data for CollectionOptions"); + CollectionOptions obj{}; + obj.read_only_ = t[0].cast(); + obj.enable_mmap_ = t[1].cast(); + obj.max_buffer_size_ = t[2].cast(); + return obj; + })); + + // SegmentOptions + py::class_(m, "SegmentOption", R"pbdoc( +Options for segment-level operations. + +Currently, this class mirrors CollectionOption and is used internally. +It supports read-only mode, memory mapping, and buffer configuration. + +Note: + This class is primarily for internal use. Most users should use + CollectionOption instead. + +Examples: + >>> opt = SegmentOption() + >>> print(opt.enable_mmap) + True +)pbdoc") + .def(py::init<>(), "Constructs a SegmentOption with default settings.") + .def_property_readonly( + "enable_mmap", + [](const SegmentOptions &self) { return self.enable_mmap_; }, + "bool: Whether memory-mapped I/O is enabled.") + .def_property_readonly( + "read_only", + [](const SegmentOptions &self) { return self.read_only_; }, + "bool: Whether the segment is read-only.") + .def_property_readonly( + "max_buffer_size", + [](const SegmentOptions &self) { return self.max_buffer_size_; }, + "int: Maximum buffer size in bytes (internal use).") + .def("__repr__", + [](const SegmentOptions &self) -> std::string { + return "{" + "\"enable_mmap\":" + + std::to_string(self.enable_mmap_) + + ", \"read_only\":" + std::to_string(self.read_only_) + + ", \"max_buffer_size\":" + + std::to_string(self.max_buffer_size_) + "}"; + }) + .def(py::pickle( + [](const SegmentOptions &self) { + return py::make_tuple(self.read_only_, self.enable_mmap_, + self.max_buffer_size_); + }, + [](py::tuple t) { + if (t.size() != 3) + throw std::runtime_error( + "Invalid pickle data for SegmentOptions"); + SegmentOptions obj{}; + obj.read_only_ = t[0].cast(); + obj.enable_mmap_ = t[1].cast(); + obj.max_buffer_size_ = t[2].cast(); + return obj; + })); + + // CreateIndexOptions + py::class_(m, "IndexOption", + R"pbdoc( +Options for creating an index. + +Attributes: + concurrency (int): Number of threads to use during index creation. + If 0, the system will choose an optimal value automatically. + Default is 0. + +Examples: + >>> opt = IndexOption(concurrency=4) + >>> print(opt.concurrency) + 4 +)pbdoc") + .def(py::init(), py::arg("concurrency") = 0, + R"pbdoc( +Constructs an IndexOption instance. + +Args: + concurrency (int, optional): Number of concurrent threads. + 0 means auto-detect. Defaults to 0. +)pbdoc") + .def_property_readonly( + "concurrency", + [](const CreateIndexOptions &self) { return self.concurrency_; }, + "int: Number of threads used for index creation (0 = auto).") + .def(py::pickle( + [](const CreateIndexOptions &self) { + return py::make_tuple(self.concurrency_); + }, + [](py::tuple t) { + if (t.size() != 1) + throw std::runtime_error( + "Invalid pickle data for CreateIndexOptions"); + CreateIndexOptions obj{}; + obj.concurrency_ = t[0].cast(); + return obj; + })); + + // OptimizeOptions + py::class_(m, "OptimizeOption", R"pbdoc( +Options for optimizing a collection (e.g., merging segments). + +Attributes: + concurrency (int): Number of threads to use during optimization. + If 0, the system will choose an optimal value automatically. + Default is 0. + +Examples: + >>> opt = OptimizeOption(concurrency=2) + >>> print(opt.concurrency) + 2 +)pbdoc") + .def(py::init(), py::arg("concurrency") = 0, + R"pbdoc( +Constructs an OptimizeOption instance. + +Args: + concurrency (int, optional): Number of concurrent threads. + 0 means auto-detect. Defaults to 0. +)pbdoc") + .def_property_readonly( + "concurrency", + [](const OptimizeOptions &self) { return self.concurrency_; }, + "int: Number of threads used for optimization (0 = auto).") + .def(py::pickle( + [](const OptimizeOptions &self) { + return py::make_tuple(self.concurrency_); + }, + [](py::tuple t) { + if (t.size() != 1) + throw std::runtime_error( + "Invalid pickle data for OptimizeOptions"); + OptimizeOptions obj{}; + obj.concurrency_ = t[0].cast(); + return obj; + })); + + // AddColumnOptions + py::class_(m, "AddColumnOption", + R"pbdoc( +Options for adding a new column to a collection. + +Attributes: + concurrency (int): Number of threads to use when backfilling data + for the new column. If 0, auto-detect is used. Default is 0. + +Examples: + >>> opt = AddColumnOption(concurrency=1) + >>> print(opt.concurrency) + 1 +)pbdoc") + .def(py::init(), py::arg("concurrency") = 0, + R"pbdoc( +Constructs an AddColumnOption instance. + +Args: + concurrency (int, optional): Number of threads for data backfill. + 0 means auto-detect. Defaults to 0. +)pbdoc") + .def_property_readonly( + "concurrency", + [](const AddColumnOptions &self) { return self.concurrency_; }, + "int: Number of threads used when adding a column (0 = auto).") + .def(py::pickle( + [](const AddColumnOptions &self) { + return py::make_tuple(self.concurrency_); + }, + [](py::tuple t) { + if (t.size() != 1) + throw std::runtime_error( + "Invalid pickle data for AddColumnOptions"); + AddColumnOptions obj{}; + obj.concurrency_ = t[0].cast(); + return obj; + })); + + // AlterColumnOptions + py::class_(m, "AlterColumnOption", R"pbdoc( +Options for altering an existing column (e.g., changing index settings). + +Attributes: + concurrency (int): Number of threads to use during the alteration process. + If 0, the system will choose an optimal value automatically. + Default is 0. + +Examples: + >>> opt = AlterColumnOption(concurrency=1) + >>> print(opt.concurrency) + 1 +)pbdoc") + .def(py::init(), py::arg("concurrency") = 0, + R"pbdoc( +Constructs an AlterColumnOption instance. + +Args: + concurrency (int, optional): Number of threads for column alteration. + 0 means auto-detect. Defaults to 0. +)pbdoc") + .def_property_readonly( + "concurrency", + [](const AlterColumnOptions &self) { return self.concurrency_; }, + "int: Number of threads used when altering a column (0 = auto).") + .def(py::pickle( + [](const AlterColumnOptions &self) { + return py::make_tuple(self.concurrency_); + }, + [](py::tuple t) { + if (t.size() != 1) + throw std::runtime_error( + "Invalid pickle data for AlterColumnOptions"); + AlterColumnOptions obj{}; + obj.concurrency_ = t[0].cast(); + return obj; + })); +} + +void ZVecPyParams::bind_vector_query(py::module_ &m) { + // bind Fts + py::class_(m, "_Fts") + .def(py::init<>()) + .def_readwrite("query_string", &FtsClause::query_string_) + .def_readwrite("match_string", &FtsClause::match_string_) + .def(py::pickle( + [](const FtsClause &self) { + return py::make_tuple(self.query_string_, self.match_string_); + }, + [](py::tuple t) { + if (t.size() != 2) + throw std::runtime_error("Invalid pickle data for Fts"); + FtsClause obj{}; + obj.query_string_ = t[0].cast(); + obj.match_string_ = t[1].cast(); + return obj; + })); + + // Bind SubQuery (used by MultiQuery) + py::class_(m, "_SubQuery") + .def(py::init<>()) + .def_readwrite("num_candidates", &SubQuery::num_candidates_) + .def_static( + "from_search_query", + [](const SearchQuery &sq) { + SubQuery sub; + sub.num_candidates_ = sq.topk_; + sub.target_ = sq.target_; + // SubQuery is copied by value into MultiQuery. Materialize + // non-owning vector views so the copied SubQuery does not depend on + // the original _SearchQuery keep-alive relationship. + if (auto *vvc = sub.target_.get_vector_view_clause()) { + VectorClause vc{std::string(vvc->query_vector_), + std::string(vvc->sparse_indices_), + std::string(vvc->sparse_values_)}; + sub.target_.clause_ = std::move(vc); + } + return sub; + }, + py::arg("search_query"), + "Create a SubQuery from a single-target search query."); + + // _SearchQuery is the Python class name; it wraps the + // single-target SearchQuery so external Python code keeps working unchanged. + py::class_(m, "_SearchQuery") + .def(py::init<>()) + // properties + .def_readwrite("topk", &SearchQuery::topk_) + .def_property( + "field_name", + [](const SearchQuery &s) { return s.target_.field_name_; }, + [](SearchQuery &s, std::string v) { + s.target_.field_name_ = std::move(v); + }) + .def_readwrite("filter", &SearchQuery::filter_) + .def_readwrite("include_vector", &SearchQuery::include_vector_) + .def_property( + "query_params", + [](const SearchQuery &s) { return s.target_.query_params_; }, + [](SearchQuery &s, QueryParams::Ptr p) { + s.target_.query_params_ = std::move(p); + }) + .def_readwrite("output_fields", &SearchQuery::output_fields_) + .def_property( + "fts", + [](const SearchQuery &self) -> py::object { + const auto *fc = self.target_.get_fts_clause(); + if (fc != nullptr) { + return py::cast(*fc); + } + return py::none(); + }, + [](SearchQuery &self, const py::object &obj) { + if (obj.is_none()) { + // Clearing FTS resets the target to an empty vector clause. + self.target_.clause_ = VectorClause{}; + } else { + self.target_.clause_ = obj.cast(); + } + }) + // vector + .def( + "set_vector", + [](SearchQuery &self, const FieldSchema &field_schema, + const py::object &obj) { + const DataType data_type = field_schema.data_type(); + + // Dense vector data is referenced by the query object. Callers + // must not modify the source data until the query returns. + if (FieldSchema::is_dense_vector_field(data_type)) { + if (!py::isinstance(obj)) { + throw py::type_error("Dense vector[" + field_schema.name() + + "] expects a ndarray, got " + + std::string(py::str(py::type::of(obj)))); + } + const auto arr = obj.cast(); + if (arr.ndim() != 1) { + throw py::type_error("Dense vector expects 1D array, got " + + std::to_string(arr.ndim()) + "D"); + } + const auto buf = arr.request(); + self.target_.clause_ = VectorViewClause{ + std::string_view( + static_cast(buf.ptr), + static_cast(buf.size) * buf.itemsize), + {}, + {}}; + return; + } + // sparse vector + if (FieldSchema::is_sparse_vector_field(data_type)) { + if (!py::isinstance(obj)) { + throw py::type_error("Sparse vector[" + field_schema.name() + + "] expects a Python dict, got " + + std::string(py::str(py::type::of(obj)))); + } + const auto sparse = obj.cast(); + + switch (data_type) { + case DataType::SPARSE_VECTOR_FP16: { + auto [indices, values] = + serialize_sparse_vector( + sparse, [](const py::handle &h, size_t idx) { + float f = checked_cast( + h, "Sparse value[" + std::to_string(idx) + "]", + "FLOAT"); + return ailego::Float16(f); + }); + self.target_.set_sparse_vector(std::move(indices), + std::move(values)); + break; + } + case DataType::SPARSE_VECTOR_FP32: { + auto [indices, values] = serialize_sparse_vector( + sparse, [](const py::handle &h, size_t idx) { + return checked_cast( + h, "Sparse value[" + std::to_string(idx) + "]", + "FLOAT"); + }); + self.target_.set_sparse_vector(std::move(indices), + std::move(values)); + break; + } + default: + throw py::type_error( + "Unsupported sparse vector type: " + + std::to_string(static_cast(data_type))); + } + return; + } + + throw py::type_error("Unsupported vector field type for field: " + + field_schema.name()); + }, + py::arg("field_schema"), py::arg("obj"), py::keep_alive<1, 3>(), + "Set query vector. Dense vector source data must not be modified " + "until the query finishes.") + .def( + "get_vector", + [](const SearchQuery &self, + const FieldSchema &field_schema) -> py::object { + DataType data_type = field_schema.data_type(); + // get_vector_view() works for both VectorClause and + // VectorViewClause. + auto vv = self.target_.get_vector_view(); + if (FieldSchema::is_dense_vector_field(data_type)) { + if (!vv || vv->query_vector_.empty()) { + throw std::runtime_error("No dense vector has been set"); + } + + size_t byte_size = vv->query_vector_.size(); + const void *data = vv->query_vector_.data(); + + switch (data_type) { + case DataType::VECTOR_FP32: { + if (byte_size % sizeof(float) != 0) { + throw std::runtime_error( + "Invalid buffer size for VECTOR_FP32"); + } + size_t dim = byte_size / sizeof(float); + return py::array_t({dim}, {sizeof(float)}, + static_cast(data)); + } + case DataType::VECTOR_FP64: { + if (byte_size % sizeof(double) != 0) { + throw std::runtime_error( + "Invalid buffer size for VECTOR_FP64"); + } + size_t dim = byte_size / sizeof(double); + return py::array_t({dim}, {sizeof(double)}, + static_cast(data)); + } + case DataType::VECTOR_INT8: { + if (byte_size % sizeof(int8_t) != 0) { + throw std::runtime_error( + "Invalid buffer size for VECTOR_INT8"); + } + size_t dim = byte_size / sizeof(int8_t); + return py::array_t({dim}, {sizeof(int8_t)}, + static_cast(data)); + } + case DataType::VECTOR_FP16: { + if (byte_size % 2 != 0) { + throw std::runtime_error( + "Invalid buffer size for VECTOR_FP16"); + } + size_t dim = byte_size / 2; + return py::array(py::dtype("float16"), {dim}, {2}, data); + } + + default: + throw py::type_error( + "Unsupported dense vector type for get_vector: " + + std::to_string(static_cast(data_type))); + } + } + if (FieldSchema::is_sparse_vector_field(data_type)) { + if (!vv || vv->sparse_indices_.empty()) { + return py::dict(); + } + + // Deserialize indices: stored as uint32_t[] + size_t indices_byte_size = vv->sparse_indices_.size(); + if (indices_byte_size % sizeof(uint32_t) != 0) { + throw std::runtime_error( + "Sparse indices buffer size not aligned to uint32_t"); + } + size_t n = indices_byte_size / sizeof(uint32_t); + const uint32_t *indices = reinterpret_cast( + vv->sparse_indices_.data()); + + // Deserialize values + switch (data_type) { + case DataType::SPARSE_VECTOR_FP32: { + if (vv->sparse_values_.size() != n * sizeof(float)) { + throw std::runtime_error( + "Sparse FP32 values buffer size mismatch"); + } + const float *values = reinterpret_cast( + vv->sparse_values_.data()); + py::dict result; + for (size_t i = 0; i < n; ++i) { + result[py::int_(indices[i])] = py::float_(values[i]); + } + return result; + } + case DataType::SPARSE_VECTOR_FP16: { + if (vv->sparse_values_.size() != n * sizeof(uint16_t)) { + throw std::runtime_error( + "Sparse FP16 values buffer size mismatch"); + } + const uint16_t *raw_bits = reinterpret_cast( + vv->sparse_values_.data()); + py::dict result; + for (size_t i = 0; i < n; ++i) { + float f = ailego::FloatHelper::ToFP32(raw_bits[i]); + result[py::int_(indices[i])] = py::float_(f); + } + return result; + } + default: + throw py::type_error("Unsupported sparse vector type..."); + } + } + + + throw py::type_error("Unsupported vector field type: " + + field_schema.name()); + }, + py::arg("field_schema")) + .def(py::pickle( + [](const SearchQuery &self) { + auto vv = self.target_.get_vector_view(); + const auto *fc = self.target_.get_fts_clause(); + return py::make_tuple( + self.topk_, self.target_.field_name_, + vv ? std::string(vv->query_vector_) : std::string(), + vv ? std::string(vv->sparse_indices_) : std::string(), + vv ? std::string(vv->sparse_values_) : std::string(), + self.filter_, self.include_vector_, self.output_fields_, + self.target_.query_params_ + ? py::cast(self.target_.query_params_) + : py::none(), + fc ? py::cast(*fc) : py::none()); + }, + [](py::tuple t) { + if (t.size() != 10) + throw std::runtime_error("Invalid pickle data for _SearchQuery"); + + SearchQuery obj{}; + obj.topk_ = t[0].cast(); + obj.target_.field_name_ = t[1].cast(); + // A vector clause and an FTS clause are mutually exclusive in the + // variant target; restore whichever the pickle carried. + if (!t[9].is_none()) { + obj.target_.clause_ = t[9].cast(); + } else { + obj.target_.clause_ = VectorClause{t[2].cast(), + t[3].cast(), + t[4].cast()}; + } + obj.filter_ = t[5].cast(); + obj.include_vector_ = t[6].cast(); + + if (!t[7].is_none()) { + obj.output_fields_ = t[7].cast>(); + } + if (!t[8].is_none()) { + obj.target_.query_params_ = t[8].cast(); + } + return obj; + })); +} +} // namespace zvec diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc new file mode 100644 index 0000000..b1311f1 --- /dev/null +++ b/src/binding/python/model/python_collection.cc @@ -0,0 +1,316 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "python_collection.h" +#include +#include + +namespace zvec { + +inline void throw_if_error(const Status &status) { + switch (status.code()) { + case StatusCode::OK: + return; + case StatusCode::NOT_FOUND: + throw py::key_error(status.message()); + case StatusCode::INVALID_ARGUMENT: + throw py::value_error(status.message()); + case StatusCode::INTERNAL_ERROR: + case StatusCode::ALREADY_EXISTS: + case StatusCode::NOT_SUPPORTED: + case StatusCode::PERMISSION_DENIED: + case StatusCode::FAILED_PRECONDITION: + case StatusCode::UNKNOWN: + default: + throw std::runtime_error(status.message()); + } +} + + +template +T unwrap_expected(const tl::expected &exp) { + if (exp.has_value()) { + return exp.value(); + } + throw_if_error(exp.error()); + return T{}; +} + +void ZVecPyCollection::Initialize(pybind11::module_ &m) { + py::class_ collection(m, "_Collection"); + bind_db_methods(collection); + bind_ddl_methods(collection); + bind_dml_methods(collection); + bind_dql_methods(collection); + collection.def(py::pickle( + [](const Collection &c) { + return py::make_tuple(c.Path(), c.Schema(), c.Options()); + }, + [](py::tuple t) { + if (t.size() != 3) { + throw std::runtime_error("Invalid tuple size for Collection pickle"); + } + std::string path = t[0].cast(); + auto schema = t[1].cast(); + CollectionOptions options = t[2].cast(); + auto result = Collection::Open(path, options); + // auto result = Collection::CreateAndOpen(path, schema, options); + return unwrap_expected(result); + })); +} + +void ZVecPyCollection::bind_db_methods( + py::class_ &col) { + col.def_static("CreateAndOpen", + [](const std::string &path, const CollectionSchema &schema, + const CollectionOptions &options) { + Result result; + { + py::gil_scoped_release release; + result = Collection::CreateAndOpen(path, schema, options); + } + return unwrap_expected(result); + }) + .def_static("Open", [](const std::string &path, + const CollectionOptions &options) { + Result result; + { + py::gil_scoped_release release; + result = Collection::Open(path, options); + } + return unwrap_expected(result); + }); +} + + +void ZVecPyCollection::bind_ddl_methods( + py::class_ &col) { + // bind collection properties + col.def("Path", + [](const Collection &self) { + auto ret = self.Path(); + return unwrap_expected(ret); + }) + .def("Options", + [](const Collection &self) { + auto ret = self.Options(); + return unwrap_expected(ret); + }) + .def("Schema", + [](const Collection &self) { + auto ret = self.Schema(); + return unwrap_expected(ret); + }) + .def("Stats", [](const Collection &self) { + auto ret = self.Stats(); + return unwrap_expected(ret); + }); + + // bind collection ddl methods + col.def("Destroy", + [](Collection &self) { + Status status; + { + py::gil_scoped_release release; + status = self.Destroy(); + } + throw_if_error(status); + }) + .def("Flush", [](Collection &self) { + Status status; + { + py::gil_scoped_release release; + status = self.Flush(); + } + throw_if_error(status); + }); + + // binding index ddl methods + col.def("CreateIndex", + [](Collection &self, const std::string &column_name, + const IndexParams::Ptr &index_options, + const CreateIndexOptions &options) { + Status status; + { + py::gil_scoped_release release; + status = self.CreateIndex(column_name, index_options, options); + } + throw_if_error(status); + }) + .def("DropIndex", + [](Collection &self, const std::string &column_name) { + Status status; + { + py::gil_scoped_release release; + status = self.DropIndex(column_name); + } + throw_if_error(status); + }) + .def("Optimize", [](Collection &self, const OptimizeOptions &options) { + Status status; + { + py::gil_scoped_release release; + status = self.Optimize(options); + } + throw_if_error(status); + }); + + // binding column ddl methods + col.def("AddColumn", + [](Collection &self, const FieldSchema::Ptr &column_schema, + const std::string &expression, const AddColumnOptions &options) { + Status status; + { + py::gil_scoped_release release; + status = self.AddColumn(column_schema, expression, options); + } + throw_if_error(status); + }) + .def("DropColumn", + [](Collection &self, std::string &column_name) { + Status status; + { + py::gil_scoped_release release; + status = self.DropColumn(column_name); + } + throw_if_error(status); + }) + .def("AlterColumn", [](Collection &self, std::string &column_name, + const std::string &rename, + const FieldSchema::Ptr &new_column_schema, + const AlterColumnOptions &options) { + Status status; + { + py::gil_scoped_release release; + status = + self.AlterColumn(column_name, rename, new_column_schema, options); + } + throw_if_error(status); + }); +} + +void ZVecPyCollection::bind_dml_methods( + py::class_ &col) { + // bind collection upsert/insert/update/delete methods + col.def("Insert", + [](Collection &self, std::vector &docs) { + Result result; + { + py::gil_scoped_release release; + result = self.Insert(docs); + } + return unwrap_expected(result); + }) + .def("Update", + [](Collection &self, std::vector &docs) { + Result result; + { + py::gil_scoped_release release; + result = self.Update(docs); + } + return unwrap_expected(result); + }) + .def("Upsert", + [](Collection &self, std::vector &docs) { + Result result; + { + py::gil_scoped_release release; + result = self.Upsert(docs); + } + return unwrap_expected(result); + }) + .def("Delete", + [](Collection &self, const std::vector &pks) { + Result result; + { + py::gil_scoped_release release; + result = self.Delete(pks); + } + return unwrap_expected(result); + }) + .def("DeleteByFilter", [](Collection &self, const std::string &filter) { + Status status; + { + py::gil_scoped_release release; + status = self.DeleteByFilter(filter); + } + throw_if_error(status); + }); +} + +void ZVecPyCollection::bind_dql_methods( + py::class_ &col) { + col.def("Query", + [](const Collection &self, const SearchQuery &query) { + Result result; + { + py::gil_scoped_release release; + result = self.Query(query); + } + // return DocPtrList + return unwrap_expected(result); + }) + // MultiQuery: multi query with reranker + .def( + "Query", + [](const Collection &self, const MultiQuery &query) { + Result result; + { + py::gil_scoped_release release; + result = self.Query(query); + } + // return DocPtrList + return unwrap_expected(result); + }, + py::arg("query"), "Execute a multi query with re-ranking.") + .def("GroupByQuery", + [](const Collection &self, const GroupByVectorQuery &query) { + Result result; + { + py::gil_scoped_release release; + result = self.GroupByQuery(query); + } + // return GroupResults + return unwrap_expected(result); + }) + .def( + "Fetch", + [](const Collection &self, const std::vector &pks, + const std::optional> &output_fields, + bool include_vector) { + Result result; + { + py::gil_scoped_release release; + result = self.Fetch(pks, output_fields, include_vector); + } + // return DocPtrMap + return unwrap_expected(result); + }, + py::arg("pks"), py::arg("output_fields") = py::none(), + py::arg("include_vector") = true) + .def( + "_debug_hnsw_storage_mode", + [](const Collection &self, const std::string &column_name) { + const auto result = self.DebugGetHnswStorageMode(column_name); + return unwrap_expected(result); + }, + py::arg("column_name"), + "Debug-only: returns the storage mode of the HNSW entity on the " + "given vector column. One of 'mmap', 'buffer_pool', 'contiguous'. " + "Raises KeyError if no HNSW index exists on the column, or " + "ValueError if the column's index is not an HNSW index. Intended " + "for introspection and testing only; not part of the stable API."); +} + +} // namespace zvec \ No newline at end of file diff --git a/src/binding/python/model/python_doc.cc b/src/binding/python/model/python_doc.cc new file mode 100644 index 0000000..6803197 --- /dev/null +++ b/src/binding/python/model/python_doc.cc @@ -0,0 +1,452 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "python_doc.h" +#include +#include + +namespace zvec { + +template +T checked_cast(const py::object &obj, const std::string &field, + const std::string &expected_type) { + try { + return obj.cast(); + } catch (const py::cast_error &e) { + std::string actual_type = std::string(py::str(py::type::of(obj))); + std::string msg = "Field '" + field + "': expected " + expected_type + + ", got " + actual_type; + throw py::type_error(msg); + } +} + +void ZVecPyDoc::Initialize(pybind11::module_ &m) { + bind_doc_operator(m); + bind_doc(m); +} + + +void ZVecPyDoc::bind_doc_operator(py::module_ &m) { + py::enum_(m, "_DocOp") + .value("INSERT", Operator::INSERT) + .value("UPDATE", Operator::UPDATE) + .value("DELETE", Operator::DELETE) + .value("UPSERT", Operator::UPSERT); +} + + +void ZVecPyDoc::bind_doc(py::module_ &m) { + // binding doc + py::class_ doc(m, "_Doc"); + + doc.def(py::init([]() { return std::make_shared(); })) + .def("set_pk", &Doc::set_pk) + .def("pk", &Doc::pk) + .def("set_score", &Doc::set_score) + .def("score", &Doc::score) + .def("has_field", &Doc::has) + .def("field_names", &Doc::field_names) + .def(py::pickle( + [](const Doc &d) { + std::vector data = d.serialize(); + return py::bytes(reinterpret_cast(data.data()), + data.size()); + }, + [](py::bytes b) { + py::buffer_info info(py::buffer(b).request()); + const uint8_t *buf = reinterpret_cast(info.ptr); + size_t size = static_cast(info.size); + Doc::Ptr d = Doc::deserialize(buf, size); + if (!d) throw std::runtime_error("Failed to unpickle Doc"); + return d; + })); + + + // binding doc set field + doc.def( + "set_any", + [](Doc &self, const std::string &field, const FieldSchema &field_schema, + const py::object &obj) -> bool { + if (obj.is_none()) { + if (field_schema.nullable()) { + self.set_null(field); + return true; + } + throw py::value_error("Field '" + field + + "': expected non-nullable type"); + } + switch (field_schema.data_type()) { + // base datatypes + case DataType::STRING: + return self.set(field, + checked_cast(obj, field, "STRING")); + case DataType::BOOL: + return self.set(field, checked_cast(obj, field, "BOOL")); + case DataType::INT32: + return self.set(field, checked_cast(obj, field, "INT32")); + case DataType::INT64: + return self.set(field, checked_cast(obj, field, "INT64")); + case DataType::UINT32: + return self.set(field, + checked_cast(obj, field, "UINT32")); + case DataType::UINT64: + return self.set(field, + checked_cast(obj, field, "UINT64")); + case DataType::FLOAT: + return self.set(field, checked_cast(obj, field, "FLOAT")); + case DataType::DOUBLE: + return self.set(field, checked_cast(obj, field, "DOUBLE")); + + // array datatypes + case DataType::ARRAY_STRING: + return self.set(field, checked_cast>( + obj, field, "ARRAY_STRING")); + case DataType::ARRAY_BOOL: + return self.set(field, checked_cast>( + obj, field, "ARRAY_BOOL")); + case DataType::ARRAY_INT32: + return self.set(field, checked_cast>( + obj, field, "ARRAY_INT32")); + case DataType::ARRAY_UINT32: + return self.set(field, checked_cast>( + obj, field, "ARRAY_UINT32")); + case DataType::ARRAY_INT64: + return self.set(field, checked_cast>( + obj, field, "ARRAY_INT64")); + case DataType::ARRAY_UINT64: + return self.set(field, checked_cast>( + obj, field, "ARRAY_UINT64")); + case DataType::ARRAY_FLOAT: + return self.set(field, checked_cast>( + obj, field, "ARRAY_FLOAT")); + case DataType::ARRAY_DOUBLE: + return self.set(field, checked_cast>( + obj, field, "ARRAY_DOUBLE")); + + // dense vector datatypes + case DataType::VECTOR_FP16: { + const auto value = checked_cast( + obj, field, "VECTOR_FP16 (list of numbers)"); + std::vector new_value; + new_value.reserve(value.size()); + for (const auto &item : value) { + try { + new_value.emplace_back(item.cast()); + } catch (const py::cast_error &e) { + throw py::type_error("Vector '" + field + + "': expected VECTOR_FP16, got " + + std::string(py::str(py::type::of(obj)))); + } + } + return self.set(field, new_value); + } + case DataType::VECTOR_FP32: + return self.set(field, checked_cast>( + obj, field, "VECTOR_FP32")); + case DataType::VECTOR_FP64: + return self.set(field, checked_cast>( + obj, field, "VECTOR_FP64")); + case DataType::VECTOR_INT8: + return self.set(field, checked_cast>( + obj, field, "VECTOR_INT8")); + + // sparse vector datatypes + case DataType::SPARSE_VECTOR_FP32: { + const auto sparse_dict = + checked_cast(obj, field, "SPARSE_VECTOR_FP32 (dict)"); + std::vector indices; + std::vector values; + for (const auto &item : sparse_dict) { + try { + indices.push_back(item.first.cast()); + values.push_back(item.second.cast()); + } catch (const py::cast_error &e) { + throw py::type_error( + "Vector '" + field + + "': sparse vector key/value must be (uint32, float), " + "got key=" + + std::string(py::str(py::type::of(item.first))) + + ", value=" + + std::string(py::str(py::type::of(item.second)))); + } + } + const std::pair, std::vector> + sparse_vector{std::move(indices), std::move(values)}; + return self.set(field, sparse_vector); + } + case DataType::SPARSE_VECTOR_FP16: { + const auto sparse_dict = + checked_cast(obj, field, "SPARSE_VECTOR_FP16 (dict)"); + std::vector indices; + std::vector values; + for (const auto &item : sparse_dict) { + try { + indices.push_back(item.first.cast()); + values.push_back(ailego::Float16(item.second.cast())); + } catch (const py::cast_error &e) { + throw py::type_error( + "Field '" + field + + "': sparse vector key/value must be (uint32, float), " + "got key=" + + std::string(py::str(py::type::of(item.first))) + + ", value=" + + std::string(py::str(py::type::of(item.second)))); + } + } + const std::pair, std::vector> + sparse_vector{std::move(indices), std::move(values)}; + return self.set(field, sparse_vector); + } + default: + throw py::type_error("Unsupported type for field: " + field); + } + }); + + // binding doc get field + doc.def( + "get_any", + [](Doc &self, const std::string &field, + const DataType &type) -> py::object { + switch (type) { + // base datatypes + case DataType::STRING: + return py::cast(self.get(field)); + case DataType::BOOL: + return py::cast(self.get(field)); + case DataType::INT32: + return py::cast(self.get(field)); + case DataType::UINT32: + return py::cast(self.get(field)); + case DataType::INT64: + return py::cast(self.get(field)); + case DataType::UINT64: + return py::cast(self.get(field)); + case DataType::FLOAT: + return py::cast(self.get(field)); + case DataType::DOUBLE: + return py::cast(self.get(field)); + + // array datatypes + case DataType::ARRAY_STRING: + return py::cast(self.get>(field)); + case DataType::ARRAY_INT32: + return py::cast(self.get>(field)); + case DataType::ARRAY_INT64: + return py::cast(self.get>(field)); + case DataType::ARRAY_UINT32: + return py::cast(self.get>(field)); + case DataType::ARRAY_UINT64: + return py::cast(self.get>(field)); + case DataType::ARRAY_FLOAT: + return py::cast(self.get>(field)); + case DataType::ARRAY_DOUBLE: + return py::cast(self.get>(field)); + case DataType::ARRAY_BOOL: + return py::cast(self.get>(field)); + + // vector datatypes + case DataType::VECTOR_INT8: + return py::cast(self.get>(field)); + case DataType::VECTOR_FP16: { + auto value = self.get>(field); + if (value.has_value()) { + std::vector new_value; + new_value.reserve(value.value().size()); + for (auto &item : value.value()) { + new_value.push_back(static_cast(item)); + } + return py::cast(new_value); + } + return py::none(); + } + case DataType::VECTOR_FP32: + return py::cast(self.get>(field)); + case DataType::VECTOR_FP64: + return py::cast(self.get>(field)); + case DataType::SPARSE_VECTOR_FP16: { + auto vector = self.get< + std::pair, std::vector>>( + field); + const auto &indices = vector->first; + const auto &values = vector->second; + py::dict d; + for (size_t i = 0; i < indices.size(); ++i) { + d[py::int_(indices[i])] = + py::float_(static_cast(values[i])); + } + return d; + } + case DataType::SPARSE_VECTOR_FP32: { + auto vector = + self.get, std::vector>>( + field); + const auto &indices = vector->first; + const auto &values = vector->second; + py::dict d; + for (size_t i = 0; i < indices.size(); ++i) { + d[py::int_(indices[i])] = py::float_(values[i]); + } + return d; + } + default: + throw py::type_error("Unsupported type for field: " + field); + } + }); + doc.def( + "get_all", + [](Doc &self, const CollectionSchema &schema) -> py::tuple { + py::tuple result(4); + // 1. set doc id and score + result[0] = py::str(self.pk()); + result[1] = py::float_(self.score()); + + if (self.is_empty()) { + result[2] = py::none(); + result[3] = py::none(); + return result; + } + // 2. set scalar fields + py::dict fields; + for (const auto &field_meta : schema.forward_fields()) { + const std::string &field = field_meta->name(); + if (!self.has_value(field)) { + continue; + } + + try { + auto val = [&]() -> py::object { + switch (field_meta->data_type()) { + // base datatypes + case DataType::STRING: + return py::str(self.get(field).value()); + case DataType::BOOL: + return py::cast(self.get(field)); + case DataType::INT32: + return py::cast(self.get(field)); + case DataType::UINT32: + return py::cast(self.get(field)); + case DataType::INT64: + return py::cast(self.get(field)); + case DataType::UINT64: + return py::cast(self.get(field)); + case DataType::FLOAT: + return py::cast(self.get(field)); + case DataType::DOUBLE: + return py::cast(self.get(field)); + + // array datatypes + case DataType::ARRAY_STRING: + return py::cast(self.get>(field)); + case DataType::ARRAY_INT32: + return py::cast(self.get>(field)); + case DataType::ARRAY_INT64: + return py::cast(self.get>(field)); + case DataType::ARRAY_UINT32: + return py::cast(self.get>(field)); + case DataType::ARRAY_UINT64: + return py::cast(self.get>(field)); + case DataType::ARRAY_FLOAT: + return py::cast(self.get>(field)); + case DataType::ARRAY_DOUBLE: + return py::cast(self.get>(field)); + case DataType::ARRAY_BOOL: + return py::cast(self.get>(field)); + default: + throw py::type_error("Unsupported type for field: " + field); + } + }(); + fields[py::str(field)] = val; + } catch (const std::exception &e) { + fields[py::str(field)] = py::none(); + } + } + if (!fields.empty()) { + result[2] = fields; + } else { + result[2] = py::none(); + } + // 3. set vector fields + py::dict vectors; + for (const auto &vec_meta : schema.vector_fields()) { + const std::string &vec = vec_meta->name(); + if (!self.has_value(vec)) continue; + + try { + auto array = [&]() -> py::object { + switch (vec_meta->data_type()) { + case DataType::VECTOR_INT8: + return py::cast(self.get>(vec)); + case DataType::VECTOR_FP16: { + auto value = self.get>(vec); + if (value.has_value()) { + std::vector new_value; + new_value.reserve(value.value().size()); + for (auto &item : value.value()) { + new_value.push_back(static_cast(item)); + } + return py::cast(new_value); + } + return py::none(); + } + case DataType::VECTOR_FP32: + return py::cast(self.get>(vec)); + case DataType::VECTOR_FP64: + return py::cast(self.get>(vec)); + case DataType::SPARSE_VECTOR_FP16: { + auto vector = + self.get, + std::vector>>(vec); + const auto &indices = vector->first; + const auto &values = vector->second; + py::dict d; + for (size_t i = 0; i < indices.size(); ++i) { + d[py::int_(indices[i])] = + py::float_(static_cast(values[i])); + } + return d; + } + case DataType::SPARSE_VECTOR_FP32: { + auto vector = self.get< + std::pair, std::vector>>( + vec); + const auto &indices = vector->first; + const auto &values = vector->second; + py::dict d; + for (size_t i = 0; i < indices.size(); ++i) { + d[py::int_(indices[i])] = py::float_(values[i]); + } + return d; + } + default: + throw py::type_error("Unsupported type for field: " + vec); + } + }(); + vectors[py::str(vec)] = array; + } catch (const std::exception &e) { + vectors[py::str(vec)] = py::none(); + } + } + if (!vectors.empty()) { + result[3] = vectors; + } else { + result[3] = py::none(); + } + return result; + }, + py::arg("schema"), + "Get all fields and vectors as a tuple: (id, score, fields, vectors). " + "Vectors are zero-copy numpy arrays (dense: ndarray, sparse: (indices, " + "values) tuple)."); +} +} // namespace zvec \ No newline at end of file diff --git a/src/binding/python/model/python_reranker.cc b/src/binding/python/model/python_reranker.cc new file mode 100644 index 0000000..ff5753f --- /dev/null +++ b/src/binding/python/model/python_reranker.cc @@ -0,0 +1,91 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "python_reranker.h" +#include +#include +#include +#include + +namespace zvec { + +void ZVecPyReranker::Initialize(py::module_ &m) { + // Bind RrfParams + py::class_(m, "_RrfParams") + .def(py::init(), py::arg("rank_constant") = 60) + .def_readwrite("rank_constant", &reranker::RrfParams::rank_constant); + + // Bind WeightedParams + py::class_(m, "_WeightedParams") + .def(py::init>(), py::arg("weights")) + .def_readwrite("weights", &reranker::WeightedParams::weights); + + // Bind CallbackParams + py::class_(m, "_CallbackParams") + .def(py::init(), py::arg("callback")); + + // Standalone rerank execution function + m.def( + "_reranker_rerank", + [](py::object params, const std::vector &results, + const std::vector &fields, int topn) -> DocPtrList { + reranker::RerankParams strategy; + if (py::isinstance(params)) { + strategy = params.cast(); + } else if (py::isinstance(params)) { + strategy = params.cast(); + } else if (py::isinstance(params)) { + strategy = params.cast(); + } else { + throw py::type_error( + "params must be _RrfParams, _WeightedParams, or _CallbackParams"); + } + auto result = reranker::rerank(strategy, results, fields, topn); + if (!result.has_value()) { + throw std::runtime_error(result.error().message()); + } + return std::move(result).value(); + }, + py::arg("params"), py::arg("results"), py::arg("fields"), + py::arg("topn")); + + // Bind MultiQuery struct + py::class_(m, "_MultiQuery") + .def(py::init<>()) + .def_readwrite("queries", &MultiQuery::queries) + .def_readwrite("topk", &MultiQuery::topk) + .def_readwrite("filter", &MultiQuery::filter) + .def_readwrite("include_vector", &MultiQuery::include_vector) + .def_readwrite("output_fields", &MultiQuery::output_fields) + .def( + "set_rerank_rrf", + [](MultiQuery &q, int rank_constant) { + q.rerank = reranker::RrfParams{rank_constant}; + }, + py::arg("rank_constant") = 60) + .def( + "set_rerank_weighted", + [](MultiQuery &q, std::vector weights) { + q.rerank = reranker::WeightedParams{std::move(weights)}; + }, + py::arg("weights")) + .def( + "set_rerank_callback", + [](MultiQuery &q, reranker::CallbackParams::Callback callback) { + q.rerank = reranker::CallbackParams{std::move(callback)}; + }, + py::arg("callback")); +} + +} // namespace zvec diff --git a/src/binding/python/model/schema/python_schema.cc b/src/binding/python/model/schema/python_schema.cc new file mode 100644 index 0000000..fcf3b60 --- /dev/null +++ b/src/binding/python/model/schema/python_schema.cc @@ -0,0 +1,162 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "python_schema.h" +#include +#include +#include + +namespace zvec { + +void ZVecPySchemas::Initialize(pybind11::module_ &parent) { + auto m = + parent.def_submodule("schema", "This module contains the schema of Zvec"); + + bind_field_schema(m); + bind_collection_schema(m); + bind_collection_stats(m); +} + +void ZVecPySchemas::bind_field_schema(py::module_ &m) { + py::class_(m, "_FieldSchema") + .def(py::init(), + py::arg("name"), py::arg("data_type"), py::arg("dimension") = 0, + py::arg("nullable") = false, py::arg("index_param") = nullptr) + .def_property_readonly("name", &FieldSchema::name) + .def_property_readonly("data_type", &FieldSchema::data_type) + .def_property_readonly("nullable", &FieldSchema::nullable) + .def_property_readonly("dimension", &FieldSchema::dimension) + .def_property_readonly("is_dense_vector", &FieldSchema::is_dense_vector) + .def_property_readonly("is_sparse_vector", &FieldSchema::is_sparse_vector) + .def_property_readonly("index_type", + [](const FieldSchema &self) { + return self.index_params() + ? self.index_type() + : IndexType::UNDEFINED; + }) + .def_property_readonly("index_param", + [](const FieldSchema &self) -> py::object { + if (self.index_params()) { + return py::cast(self.index_params()); + } + return py::none(); + }) + .def("__eq__", &FieldSchema::operator==) + .def("__ne__", &FieldSchema::operator!=) + .def(py::pickle( + [](const FieldSchema &self) { + return py::make_tuple(self.name(), self.data_type(), + self.dimension(), self.nullable(), + self.index_params() + ? py::cast(self.index_params()) + : py::none()); + }, + [](py::tuple t) { + if (t.size() != 5) { + throw std::runtime_error( + "Invalid tuple size for FieldSchema pickle"); + } + std::string name = t[0].cast(); + DataType dtype = t[1].cast(); + uint32_t dim = t[2].cast(); + bool nullable = t[3].cast(); + + IndexParams::Ptr idx_params = nullptr; + if (!t[4].is_none()) { + idx_params = t[4].cast(); + } + + return std::make_shared(name, dtype, dim, nullable, + idx_params); + })); +} + +void ZVecPySchemas::bind_collection_schema(py::module_ &m) { + py::class_(m, "_CollectionSchema") + .def(py::init(), + py::arg("name"), py::arg("fields"), + "Construct with name and list of fields") + .def_property_readonly("name", &CollectionSchema::name) + .def("has_field", &CollectionSchema::has_field, py::arg("field_name"), + "Check if a field exists.") + .def( + "get_field", + [](const CollectionSchema &self, const std::string &name) + -> const FieldSchema * { return self.get_field(name); }, + py::arg("field_name"), py::return_value_policy::reference_internal, + "Get field by name (const pointer), returns None if not found.") + .def( + "get_forward_field", + [](const CollectionSchema &self, const std::string &name) + -> const FieldSchema * { return self.get_forward_field(name); }, + py::arg("field_name"), py::return_value_policy::reference_internal, + "Get forward field (used for filtering).") + .def( + "get_vector_field", + [](const CollectionSchema &self, const std::string &name) + -> const FieldSchema * { return self.get_vector_field(name); }, + py::arg("field_name"), py::return_value_policy::reference_internal, + "Get vector field by name.") + .def("fields", &CollectionSchema::fields, + "Return list of all field schemas.", py::return_value_policy::copy) + .def("forward_fields", &CollectionSchema::forward_fields, + "Return list of forward-indexed fields.", + py::return_value_policy::copy) + .def("vector_fields", &CollectionSchema::vector_fields, + "Return list of vector fields.", py::return_value_policy::copy) + .def("__eq__", &CollectionSchema::operator==) + .def("__ne__", &CollectionSchema::operator!=) + .def(py::pickle( + [](const CollectionSchema &cs) { + return py::make_tuple(cs.name(), cs.fields(), + cs.max_doc_count_per_segment()); + }, + [](py::tuple t) { + if (t.size() != 3) + throw std::runtime_error("Invalid state for CollectionSchema!"); + + auto name = t[0].cast(); + auto fields = t[1].cast(); + auto max_docs = t[2].cast(); + + auto cs = std::make_shared(name, fields); + cs->set_max_doc_count_per_segment(max_docs); + return cs; + })); +} + +void ZVecPySchemas::bind_collection_stats(py::module_ &m) { + pybind11::class_(m, "CollectionStats") + .def(pybind11::init<>()) + .def_property_readonly( + "doc_count", [](const CollectionStats &c) { return c.doc_count; }) + .def_property_readonly( + "index_completeness", + [](const CollectionStats &c) { return c.index_completeness; }) + .def("__repr__", [](const CollectionStats &c) { + std::string map_str = "{"; + bool first = true; + for (const auto &[k, v] : c.index_completeness) { + if (!first) map_str += ", "; + map_str += "\"" + k + "\":" + std::to_string(v); + first = false; + } + map_str += "}"; + return "{\"doc_count\":" + std::to_string(c.doc_count) + + ", \"index_completeness\":" + map_str + "}"; + }); +} + +} // namespace zvec \ No newline at end of file diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc new file mode 100644 index 0000000..e6b9bd7 --- /dev/null +++ b/src/binding/python/typing/python_type.cc @@ -0,0 +1,229 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "python_type.h" + +namespace zvec { + +void ZVecPyTyping::Initialize(pybind11::module_ &parent) { + auto m = parent.def_submodule( + "typing", "This module contains the basic data types of Zvec"); + // binding base types + bind_datatypes(m); + bind_index_types(m); + bind_metric_types(m); + bind_quantize_types(m); + bind_status(m); +} + +void ZVecPyTyping::bind_datatypes(pybind11::module_ &m) { + py::enum_(m, "DataType", R"pbdoc( +Enumeration of supported data types in Zvec. + +Includes scalar types, dense/sparse vector types, and array types. + +Examples: + >>> from zvec.typing import DataType + >>> print(DataType.FLOAT) + DataType.FLOAT + >>> print(DataType.VECTOR_FP32) + DataType.VECTOR_FP32 +)pbdoc") + // field type + .value("STRING", DataType::STRING) + .value("BOOL", DataType::BOOL) + .value("INT32", DataType::INT32) + .value("INT64", DataType::INT64) + .value("FLOAT", DataType::FLOAT) + .value("DOUBLE", DataType::DOUBLE) + .value("UINT32", DataType::UINT32) + .value("UINT64", DataType::UINT64) + + + // dense vector type + .value("VECTOR_FP16", DataType::VECTOR_FP16) + .value("VECTOR_FP32", DataType::VECTOR_FP32) + .value("VECTOR_FP64", DataType::VECTOR_FP64) + .value("VECTOR_INT8", DataType::VECTOR_INT8) + + + // sparse vector type + .value("SPARSE_VECTOR_FP32", DataType::SPARSE_VECTOR_FP32) + .value("SPARSE_VECTOR_FP16", DataType::SPARSE_VECTOR_FP16) + + + // array type [not support bool/bytes] + .value("ARRAY_STRING", DataType::ARRAY_STRING) + .value("ARRAY_INT32", DataType::ARRAY_INT32) + .value("ARRAY_INT64", DataType::ARRAY_INT64) + .value("ARRAY_FLOAT", DataType::ARRAY_FLOAT) + .value("ARRAY_DOUBLE", DataType::ARRAY_DOUBLE) + .value("ARRAY_BOOL", DataType::ARRAY_BOOL) + .value("ARRAY_UINT32", DataType::ARRAY_UINT32) + .value("ARRAY_UINT64", DataType::ARRAY_UINT64) + + + // non support + // .value("BINARY", DataType::BINARY) + // .value("ARRAY_BINARY", DataType::ARRAY_BINARY) + // .value("VECTOR_INT4", DataType::VECTOR_INT4) + // .value("VECTOR_INT16", DataType::VECTOR_INT16) + // .value("VECTOR_BINARY32", DataType::VECTOR_BINARY32) + // .value("VECTOR_BINARY64", DataType::VECTOR_BINARY64) + // .value("UNDEFINED", DataType::UNDEFINED) + ; +} + +void ZVecPyTyping::bind_index_types(pybind11::module_ &m) { + py::enum_(m, "IndexType", R"pbdoc( +Enumeration of supported index types in Zvec. + +Examples: + >>> from zvec.typing import IndexType + >>> print(IndexType.HNSW) + IndexType.HNSW +)pbdoc") + .value("UNDEFINED", IndexType::UNDEFINED) + .value("HNSW", IndexType::HNSW) + .value("HNSW_RABITQ", IndexType::HNSW_RABITQ) + .value("IVF", IndexType::IVF) + .value("FLAT", IndexType::FLAT) + .value("VAMANA", IndexType::VAMANA) + .value("INVERT", IndexType::INVERT); +} + +void ZVecPyTyping::bind_metric_types(pybind11::module_ &m) { + py::enum_(m, "MetricType", R"pbdoc( +Enumeration of supported distance/similarity metrics. + +- COSINE: Cosine similarity. +- IP: Inner product (dot product). +- L2: Euclidean distance (L2 norm). + +Examples: + >>> from zvec.typing import MetricType + >>> print(MetricType.COSINE) + MetricType.COSINE +)pbdoc") + .value("COSINE", MetricType::COSINE) + .value("IP", MetricType::IP) + .value("L2", MetricType::L2); +} + +void ZVecPyTyping::bind_quantize_types(py::module_ &m) { + py::enum_(m, "QuantizeType", R"pbdoc( +Enumeration of supported quantization types for vector compression. + +Examples: + >>> from zvec.typing import QuantizeType + >>> print(QuantizeType.INT8) + QuantizeType.INT8 +)pbdoc") + .value("UNDEFINED", QuantizeType::UNDEFINED) + .value("FP16", QuantizeType::FP16) + .value("INT8", QuantizeType::INT8) + .value("INT4", QuantizeType::INT4) + .value("RABITQ", QuantizeType::RABITQ); +} + +void ZVecPyTyping::bind_status(py::module_ &m) { + // bind status code + py::enum_(m, "StatusCode", R"pbdoc( +Enumeration of possible status codes for Zvec operations. + +Used by the `Status` class to indicate success or failure reason. +)pbdoc") + .value("OK", StatusCode::OK) + .value("NOT_FOUND", StatusCode::NOT_FOUND) + .value("ALREADY_EXISTS", StatusCode::ALREADY_EXISTS) + .value("INVALID_ARGUMENT", StatusCode::INVALID_ARGUMENT) + .value("PERMISSION_DENIED", StatusCode::PERMISSION_DENIED) + .value("FAILED_PRECONDITION", StatusCode::FAILED_PRECONDITION) + .value("RESOURCE_EXHAUSTED", StatusCode::RESOURCE_EXHAUSTED) + .value("UNAVAILABLE", StatusCode::UNAVAILABLE) + .value("INTERNAL_ERROR", StatusCode::INTERNAL_ERROR) + .value("NOT_SUPPORTED", StatusCode::NOT_SUPPORTED) + .value("UNKNOWN", StatusCode::UNKNOWN); + + // bind status + py::class_(m, "Status", R"pbdoc( +Represents the outcome of a Zvec operation. + +A `Status` object is either OK (success) or carries an error code and message. + +Examples: + >>> from zvec.typing import Status, StatusCode + >>> s = Status() + >>> print(s.ok()) + True + >>> s = Status(StatusCode.INVALID_ARGUMENT, "Field not found") + >>> print(s.code() == StatusCode.INVALID_ARGUMENT) + True + >>> print(s.message()) + Field not found +)pbdoc") + .def(py::init<>()) + .def(py::init(), py::arg("code"), + py::arg("message") = "", R"pbdoc( +Construct a status with the given code and optional message. + +Args: + code (StatusCode): The status code. + message (str, optional): Error message. Defaults to empty string. +)pbdoc") + .def("ok", &Status::ok, "bool: Returns True if the status is OK.") + .def("code", &Status::code, "StatusCode: Returns the status code.") + .def("message", &Status::message, + "str: Returns the error message (may be empty).") + .def_static("OK", &Status::OK, "Create an OK status.") + .def_static( + "InvalidArgument", + [](const std::string &msg) { return Status::InvalidArgument(msg); }, + py::arg("message")) + .def_static( + "NotFound", + [](const std::string &msg) { return Status::NotFound(msg); }, + py::arg("message")) + .def_static( + "AlreadyExists", + [](const std::string &msg) { return Status::AlreadyExists(msg); }, + py::arg("message")) + .def_static( + "InternalError", + [](const std::string &msg) { return Status::InternalError(msg); }, + py::arg("message")) + .def_static( + "PermissionDenied", + [](const std::string &msg) { return Status::PermissionDenied(msg); }, + py::arg("message")) + .def("__eq__", [](const Status &self, + const Status &other) { return self == other; }) + .def("__ne__", [](const Status &self, + const Status &other) { return self != other; }) + .def("__repr__", [](const Status &self) { + std::string result = + "{" + "\"code\":" + + std::to_string(static_cast(self.code())); + + if (!self.message().empty()) { + result += ", \"message\":\"" + self.message() + "\""; + } + + result += "}"; + return result; + }); +} + +} // namespace zvec \ No newline at end of file diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt new file mode 100644 index 0000000..c52c3b8 --- /dev/null +++ b/src/core/CMakeLists.txt @@ -0,0 +1,87 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(RABITQ_SUPPORTED AND AUTO_DETECT_ARCH) + set(HNSW_RABITQ_FILES + hnsw_rabitq_query_algorithm.cc + hnsw_rabitq_streamer.cc + hnsw_rabitq_searcher.cc + hnsw_rabitq_entity.cc + rabitq_reformer.cc + rabitq_converter.cc + ) + set(HNSW_RABITQ_FILES_FULL ${HNSW_RABITQ_FILES}) + list(TRANSFORM HNSW_RABITQ_FILES_FULL PREPEND "algorithm/hnsw_rabitq/") + foreach(FILE ${HNSW_RABITQ_FILES_FULL}) + set_source_files_properties( + ${FILE} + PROPERTIES + COMPILE_FLAGS "${RABITQ_ARCH_FLAG}" + ) + endforeach() +endif() + +# utility/block_heap.cc uses AVX2 intrinsics guarded by __AVX2__. When the +# host toolchain supports it, compile this source with an AVX2-capable +# -march so AVX2 codegen is emitted. zvec_core glob-collects this source +# too, so per-file flags must be set here as well (in addition to the +# core_utility target in utility/CMakeLists.txt). Callers runtime-gate +# invocation of BlockHeap paths on CpuFeatures::AVX2. +if(NOT ANDROID AND AUTO_DETECT_ARCH) + if(HOST_ARCH MATCHES "^(x86|x64)$") + setup_compiler_march_for_x86( + _BLOCK_HEAP_MARCH_SSE _BLOCK_HEAP_MARCH_AVX2 + _BLOCK_HEAP_MARCH_AVX512 _BLOCK_HEAP_MARCH_AVX512FP16) + if(_BLOCK_HEAP_MARCH_AVX2) + set_source_files_properties( + utility/block_heap.cc + PROPERTIES + COMPILE_FLAGS "${_BLOCK_HEAP_MARCH_AVX2}" + ) + endif() + endif() +endif() + +cc_directory(framework) +cc_directory(algorithm) +cc_directory(metric) +cc_directory(quantizer) +cc_directory(utility) +cc_directory(interface) +cc_directory(mixed_reducer) +cc_directory(plugin) + +git_version(GIT_SRCS_VER ${CMAKE_CURRENT_SOURCE_DIR}) +file(GLOB_RECURSE ALL_CORE_SRCS *.cc *.c *.h) + +# Remove algorithm/hnsw_rabitq implementation files if not supported. +# interface/indexes/hnsw_rabitq_index.cc is kept because it provides the vtable +# for HNSWRabitqIndex and guards rabitqlib usage with #if RABITQ_SUPPORTED. +if(NOT RABITQ_SUPPORTED) + list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/hnsw_rabitq/.*") +endif() + +# Always exclude algorithm/diskann implementation files from zvec_core. +# The DiskAnn algorithm is provided by the separate core_knn_diskann library +# (real on Linux x86_64, stub on other platforms). Including them here causes +# duplicate symbols and missing -laio when test binaries link both zvec_core +# (via zvec) and core_knn_diskann. + +list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/diskann/.*") +if(NOT DISKANN_SUPPORTED) + list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/interface/indexes/diskann_index\\.cc") +endif() + +set(ZVEC_CORE_LIBS zvec_ailego zvec_turbo sparsehash magic_enum rabitqlib) +# The plugin loader uses dlopen/dlsym, so link libdl on Linux. +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + list(APPEND ZVEC_CORE_LIBS ${CMAKE_DL_LIBS}) +endif() + +cc_library( + NAME zvec_core STATIC STRICT PACKED + SRCS ${ALL_CORE_SRCS} + LIBS ${ZVEC_CORE_LIBS} + INCS . ${PROJECT_ROOT_DIR}/src/core + VERSION "${GIT_SRCS_VER}" +) diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt new file mode 100644 index 0000000..bed7729 --- /dev/null +++ b/src/core/algorithm/CMakeLists.txt @@ -0,0 +1,67 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +cc_directory(cluster) +cc_directory(flat) +cc_directory(flat_sparse) +cc_directory(ivf) +cc_directory(hnsw) +cc_directory(hnsw_sparse) +cc_directory(vamana) + +if(DISKANN_SUPPORTED) + message(STATUS "build diskann") + cc_directory(diskann) +else() + message(STATUS "not build diskann") + # Empty stub library for unsupported platforms + file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc + "// Stub implementation for unsupported platforms\n" + "// DiskAnn only supports Linux x86_64\n" + "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" + ) + + if(MSVC) + # MSVC: STATIC-only stub to avoid creating an empty DLL with no exports + # (MSVC linker fails to produce an import library when there are zero exports) + cc_library( + NAME core_knn_diskann + STATIC STRICT ALWAYS_LINK + SRCS ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc + LIBS core_framework + INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + VERSION "${PROXIMA_ZVEC_VERSION}" + ) + else() + cc_library( + NAME core_knn_diskann + STATIC SHARED STRICT ALWAYS_LINK + SRCS ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc + LIBS core_framework + INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + VERSION "${PROXIMA_ZVEC_VERSION}" + ) + endif() +endif() + +if(RABITQ_SUPPORTED) + message(STATUS "BUILD RABITQ") + cc_directory(hnsw_rabitq) +else() + message(STATUS "NOT BUILD RABITQ") + # Empty stub library for unsupported platforms + file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/rabitq_stub.cc + "// Stub implementation for unsupported platforms\n" + "// RaBitQ only supports Linux x86_64\n" + "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" + ) + + cc_library( + NAME core_knn_hnsw_rabitq + STATIC SHARED STRICT ALWAYS_LINK + SRCS ${CMAKE_CURRENT_BINARY_DIR}/rabitq_stub.cc + LIBS core_framework + INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + VERSION "${PROXIMA_ZVEC_VERSION}" + ) +endif() diff --git a/src/core/algorithm/cluster/CMakeLists.txt b/src/core/algorithm/cluster/CMakeLists.txt new file mode 100644 index 0000000..d954b0a --- /dev/null +++ b/src/core/algorithm/cluster/CMakeLists.txt @@ -0,0 +1,19 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +# --exclude-libs is GNU ld / LLVM lld only; Apple ld does not support it. +# On macOS (Mach-O), symbol interposition works differently and the +# Arrow/Parquet double-free issue does not apply. +if(NOT APPLE) + set(CORE_KNN_CLUSTER_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_cluster STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS zvec_ailego core_framework + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/cluster + LDFLAGS "${CORE_KNN_CLUSTER_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/cluster/cluster_params.h b/src/core/algorithm/cluster/cluster_params.h new file mode 100644 index 0000000..497792f --- /dev/null +++ b/src/core/algorithm/cluster/cluster_params.h @@ -0,0 +1,213 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +//! General +static const std::string GENERAL_CLUSTER_COUNT = "zvec.general.cluster.count"; +static const std::string GENERAL_THREAD_COUNT = + "zvec.general.cluster.thread_count"; + +//! Optimize K-means +static const std::string OPTKMEANS_CLUSTER_COUNT = + "zvec.optkmeans.cluster.count"; +static const std::string OPTKMEANS_CLUSTER_MAX_ITERATIONS = + "zvec.optkmeans.cluster.max_iterations"; +static const std::string OPTKMEANS_CLUSTER_EPSILON = + "zvec.optkmeans.cluster.epsilon"; +static const std::string OPTKMEANS_CLUSTER_SHARD_FACTOR = + "zvec.optkmeans.cluster.shard_factor"; +static const std::string OPTKMEANS_CLUSTER_PURGE_EMPTY = + "zvec.optkmeans.cluster.purge_empty"; +static const std::string OPTKMEANS_CLUSTER_MARKOV_CHAIN_LENGTH = + "zvec.optkmeans.cluster.markov_chain_length"; +static const std::string OPTKMEANS_CLUSTER_ASSUMPTION_FREE = + "zvec.optkmeans.cluster.assumption_free"; + +//! K-means +static const std::string KMEANS_CLUSTER_COUNT = "zvec.kmeans.cluster.count"; +static const std::string KMEANS_CLUSTER_SHARD_FACTOR = + "zvec.kmeans.cluster.shard_factor"; +static const std::string KMEANS_CLUSTER_EPSILON = "zvec.kmeans.cluster.epsilon"; +static const std::string KMEANS_CLUSTER_MAX_ITERATIONS = + "zvec.kmeans.cluster.max_iterations"; +static const std::string KMEANS_CLUSTER_PURGE_EMPTY = + "zvec.kmeans.cluster.purge_empty"; +static const std::string KMEANS_CLUSTER_BATCH = "zvec.kmeans.cluster.batch"; +static const std::string KMEANS_CLUSTER_SEEKER_CLASS = + "zvec.kmeans.cluster.seeker_class"; +static const std::string KMEANS_CLUSTER_SEEKER_PARAMS = + "zvec.kmeans.cluster.seeker_params"; + +//! Mini Batch K-means +static const std::string MINIBATCHKMEANS_CLUSTER_COUNT = + "zvec.minibatchkmeans.cluster.count"; +static const std::string MINIBATCHKMEANS_CLUSTER_SHARD_FACTOR = + "zvec.minibatchkmeans.cluster.shard_factor"; +static const std::string MINIBATCHKMEANS_CLUSTER_EPSILON = + "zvec.minibatchkmeans.cluster.epsilon"; +static const std::string MINIBATCHKMEANS_CLUSTER_MAX_ITERATIONS = + "zvec.minibatchkmeans.cluster.max_iterations"; +static const std::string MINIBATCHKMEANS_CLUSTER_PURGE_EMPTY = + "zvec.minibatchkmeans.cluster.purge_empty"; +static const std::string MINIBATCHKMEANS_CLUSTER_TRY_COUNT = + "zvec.minibatchkmeans.cluster.try_count"; +static const std::string MINIBATCHKMEANS_CLUSTER_BATCH_COUNT = + "zvec.minibatchkmeans.cluster.batch_count"; +static const std::string MINIBATCHKMEANS_CLUSTER_SEEKER_CLASS = + "zvec.minibatchkmeans.cluster.seeker_class"; +static const std::string MINIBATCHKMEANS_CLUSTER_SEEKER_PARAMS = + "zvec.minibatchkmeans.cluster.seeker_params"; + +//! K-means++ +static const std::string KMEANSPP_CLUSTER_COUNT = "zvec.kmeanspp.cluster.count"; +static const std::string KMEANSPP_CLUSTER_SHARD_FACTOR = + "zvec.kmeanspp.cluster.shard_factor"; +static const std::string KMEANSPP_CLUSTER_CLASS = "zvec.kmeanspp.cluster.class"; +static const std::string KMEANSPP_CLUSTER_PARAMS = + "zvec.kmeanspp.cluster.params"; + +//! K-MC2 +static const std::string KMC2_CLUSTER_COUNT = "zvec.kmc2.cluster.count"; +static const std::string KMC2_CLUSTER_SHARD_FACTOR = + "zvec.kmc2.cluster.shard_factor"; +static const std::string KMC2_CLUSTER_MARKOV_CHAIN_LENGTH = + "zvec.kmc2.cluster.markov_chain_length"; +static const std::string KMC2_CLUSTER_ASSUMPTION_FREE = + "zvec.kmc2.cluster.assumption_free"; +static const std::string KMC2_CLUSTER_CLASS = "zvec.kmc2.cluster.class"; +static const std::string KMC2_CLUSTER_PARAMS = "zvec.kmc2.cluster.params"; + +//! Bisecting K-means +static const std::string BIKMEANS_CLUSTER_COUNT = "zvec.bikmeans.cluster.count"; +static const std::string BIKMEANS_CLUSTER_INIT_COUNT = + "zvec.bikmeans.cluster.init_count"; +static const std::string BIKMEANS_CLUSTER_PURGE_EMPTY = + "zvec.bikmeans.cluster.purge_empty"; +static const std::string BIKMEANS_CLUSTER_FIRST_CLASS = + "zvec.bikmeans.cluster.first_class"; +static const std::string BIKMEANS_CLUSTER_SECOND_CLASS = + "zvec.bikmeans.cluster.second_class"; +static const std::string BIKMEANS_CLUSTER_FIRST_PARAMS = + "zvec.bikmeans.cluster.first_params"; +static const std::string BIKMEANS_CLUSTER_SECOND_PARAMS = + "zvec.bikmeans.cluster.second_params"; + +//! K-medoids +static const std::string KMEDOIDS_CLUSTER_COUNT = "zvec.kmedoids.cluster.count"; +static const std::string KMEDOIDS_CLUSTER_SHARD_FACTOR = + "zvec.kmedoids.cluster.shard_factor"; +static const std::string KMEDOIDS_CLUSTER_EPSILON = + "zvec.kmedoids.cluster.epsilon"; +static const std::string KMEDOIDS_CLUSTER_MAX_ITERATIONS = + "zvec.kmedoids.cluster.max_iterations"; +static const std::string KMEDOIDS_CLUSTER_PURGE_EMPTY = + "zvec.kmedoids.cluster.purge_empty"; +static const std::string KMEDOIDS_CLUSTER_BENCH_RATIO = + "zvec.kmedoids.cluster.bench_ratio"; +static const std::string KMEDOIDS_CLUSTER_ONLY_MEANS = + "zvec.kmedoids.cluster.only_means"; +static const std::string KMEDOIDS_CLUSTER_WITHOUT_MEANS = + "zvec.kmedoids.cluster.without_means"; +static const std::string KMEDOIDS_CLUSTER_SEEKER_CLASS = + "zvec.kmedoids.cluster.seeker_class"; +static const std::string KMEDOIDS_CLUSTER_SEEKER_PARAMS = + "zvec.kmedoids.cluster.seeker_params"; + +//! Stratified +static const std::string STRATIFIED_CLUSTER_COUNT = + "zvec.stratified.cluster.count"; +static const std::string STRATIFIED_CLUSTER_FIRST_CLASS = + "zvec.stratified.cluster.first_class"; +static const std::string STRATIFIED_CLUSTER_SECOND_CLASS = + "zvec.stratified.cluster.second_class"; +static const std::string STRATIFIED_CLUSTER_FIRST_COUNT = + "zvec.stratified.cluster.first_count"; +static const std::string STRATIFIED_CLUSTER_SECOND_COUNT = + "zvec.stratified.cluster.second_count"; +static const std::string STRATIFIED_CLUSTER_FIRST_PARAMS = + "zvec.stratified.cluster.first_params"; +static const std::string STRATIFIED_CLUSTER_SECOND_PARAMS = + "zvec.stratified.cluster.second_params"; +static const std::string STRATIFIED_CLUSTER_AUTO_TUNING = + "zvec.stratified.cluster.auto_tuning"; +static const std::string STRATIFIED_CLUSTER_SECOND_POOL_COUNT = + "zvec.stratified.cluster.second_pool_count"; + +//! Gap Statistics +static const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MIN = + "zvec.gapstats.cluster_estimater.k_min"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MAX = + "zvec.gapstats.cluster_estimater.k_max"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MIN_STEP = + "zvec.gapstats.cluster_estimater.k_min_step"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MAX_STEP = + "zvec.gapstats.cluster_estimater.k_max_step"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_TRY_COUNT = + "zvec.gapstats.cluster_estimater.try_count"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_SHARD_FACTOR = + "zvec.gapstats.cluster_estimater.shard_factor"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_ENABLE_MC2 = + "zvec.gapstats.cluster_estimater.enable_mc2"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_MARKOV_CHAIN_LENGTH = + "zvec.gapstats.cluster_estimater.markov_chain_length"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_CLUSTER_CLASS = + "zvec.gapstats.cluster_estimater.cluster_class"; + +static const std::string CLUSTER_TRAINER_SAMPLE_COUNT = + "zvec.cluster.trainer.sample_count"; +static const std::string CLUSTER_TRAINER_SAMPLE_RATIO = + "zvec.cluster.trainer.sample_ratio"; +static const std::string CLUSTER_TRAINER_THREAD_COUNT = + "zvec.cluster.trainer.thread_count"; +static const std::string CLUSTER_TRAINER_FILE_NAME = + "zvec.cluster.trainer.file_name"; +static const std::string CLUSTER_TRAINER_CLASS_NAME = + "zvec.cluster.trainer.class_name"; + +static const std::string STRATIFIED_TRAINER_SAMPLE_COUNT = + "zvec.stratified.trainer.sample_count"; +static const std::string STRATIFIED_TRAINER_SAMPLE_RATIO = + "zvec.stratified.trainer.sample_ratio"; +static const std::string STRATIFIED_TRAINER_THREAD_COUNT = + "zvec.stratified.trainer.thread_count"; +static const std::string STRATIFIED_TRAINER_FILE_NAME = + "zvec.stratified.trainer.file_name"; +static const std::string STRATIFIED_TRAINER_CLASS_NAME = + "zvec.stratified.trainer.class_name"; +static const std::string STRATIFIED_TRAINER_CLUSTER_COUNT = + "zvec.stratified.trainer.cluster_count"; +static const std::string STRATIFIED_TRAINER_AUTOAUNE = + "zvec.stratified.trainer.autotune"; +static const std::string STRATIFIED_TRAINER_PARAMS_IN_LEVEL_PREFIX = + "zvec.stratified.trainer.cluster_params_in_level_"; + +static const std::string MULTI_CHUNK_CLUSTER_COUNT = + "zvec.cluster.multi_chunk_cluster.count"; +static const std::string MULTI_CHUNK_CLUSTER_CHUNK_COUNT = + "zvec.cluster.multi_chunk_cluster.chunk_count"; +static const std::string MULTI_CHUNK_CLUSTER_THREAD_COUNT = + "zvec.cluster.multi_chunk_cluster.thread_count"; +static const std::string MULTI_CHUNK_CLUSTER_EPSILON = + "zvec.cluster.multi_chunk_cluster.epsilon"; +static const std::string MULTI_CHUNK_CLUSTER_MAX_ITERATIONS = + "zvec.cluster.multi_chunk_cluster.max_iterations"; +static const std::string MULTI_CHUNK_CLUSTER_MARKOV_CHAIN_LENGTH = + "zvec.cluster.multi_chunk_cluster.markov_chain_length"; +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/kmeans_cluster.cc b/src/core/algorithm/cluster/kmeans_cluster.cc new file mode 100644 index 0000000..c57a255 --- /dev/null +++ b/src/core/algorithm/cluster/kmeans_cluster.cc @@ -0,0 +1,917 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include +#include +#include +#include +#include +#include "cluster_params.h" +#include "linear_seeker.h" +#include "vector_mean.h" + +namespace zvec { +namespace core { + +/*! Kmeans Cluster + */ +class KmeansCluster : public IndexCluster { + public: + //! Constructor + KmeansCluster(void) {} + + //! Constructor + KmeansCluster(size_t iters, bool batch) + : max_iterations_(iters), batch_(batch) {} + + //! Constructor + KmeansCluster(bool batch) : batch_(batch) {} + + //! Destructor + ~KmeansCluster(void) override {} + + //! Initialize Cluster + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup Cluster + int cleanup(void) override; + + //! Reset Cluster + int reset(void) override; + + //! Update Cluster + int update(const ailego::Params ¶ms) override; + + //! Suggest dividing to K clusters + void suggest(uint32_t k) override; + + //! Mount features + int mount(IndexFeatures::Pointer feats) override; + + //! Cluster + int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + //! Classify + int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + //! Label + int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) override; + + protected: + //! Test if it is valid + bool is_valid(void) const; + + //! Cluster once + int clustering(IndexThreads *threads, IndexCluster::CentroidList ¢s, + double *cost); + + //! Update parameters + void update_params(const ailego::Params ¶ms); + + //! Init seeker + int init_seeker(void); + + //! Build seeker + int build_seeker(const IndexCluster::CentroidList ¢s); + + //! Check Centroids + bool check_centroids(const IndexCluster::CentroidList ¢s); + + //! Initialize Centroids + void init_centroids(size_t count, IndexCluster::CentroidList *out); + + //! Initialize Shard Containers + void init_containers(size_t shard_count); + + //! Initialize Shard Features Containers + void init_features_containers(size_t shard_count); + + //! Split Clusters + void split_clusters(IndexThreads *threads, + const IndexCluster::CentroidList ¢s); + + //! Update Centroids + void update_centroids(IndexThreads *threads, + IndexCluster::CentroidList ¢s); + + //! Update Clusters + void update_clusters(IndexThreads *threads, + const IndexCluster::CentroidList ¢s); + + //! Update Clusters' Features + void update_features(IndexThreads *threads, + IndexCluster::CentroidList ¢s); + + //! Update Labels + void update_labels(IndexThreads *threads, std::vector *labels); + + //! Split Clusters in Thread + void split_clusters_thread(size_t index_begin, size_t index_end, + const IndexThreads *threads); + + //! Update Centroid in Thread + void update_centroid_thread(size_t column, IndexCluster::CentroidList *out); + + //! Update Cluster in Thread + void update_cluster_thread(size_t index_begin, size_t index_end, + const IndexThreads *threads); + + //! Update Cluster's Features in Thread + void update_features_thread(size_t column, IndexCluster::CentroidList *out); + + //! Update Labels in Thread + void update_labels_thread(size_t index_begin, size_t index_end, + std::vector *labels); + + protected: + //! Members + IndexMeta meta_{}; + IndexFeatures::Pointer features_{}; + LinearSeeker::Pointer seeker_{}; + std::vector shard_cluster_scores_{}; + std::vector> shard_cluster_features_{}; + std::shared_ptr shard_cluster_means_{}; + std::shared_ptr batch_means_{}; + std::vector batch_scores_{}; + double epsilon_{std::numeric_limits::epsilon()}; + float shard_factor_{16.0f}; + uint32_t max_iterations_{20u}; + uint32_t cluster_count_{0u}; + uint32_t thread_count_{0u}; + bool batch_{false}; + bool purge_empty_{false}; +}; + +/*! Centroid Features + */ +class KmeansCentroidFeatures : public IndexFeatures { + public: + //! Constructor + KmeansCentroidFeatures(const IndexMeta &meta, + const IndexCluster::CentroidList ¢s) + : centroids_(cents), + feature_size_(meta.element_size()), + feature_dimension_(meta.dimension()), + data_type_(meta.data_type()) {} + + size_t count(void) const override { + return centroids_.size(); + } + + size_t dimension(void) const override { + return feature_dimension_; + } + + const void *element(size_t i) const override { + return centroids_[i].feature(); + } + + IndexMeta::DataType data_type(void) const override { + return data_type_; + } + + size_t element_size(void) const override { + return feature_size_; + } + + private: + const IndexCluster::CentroidList ¢roids_; + size_t feature_size_; + size_t feature_dimension_; + IndexMeta::DataType data_type_; +}; + +static inline std::shared_ptr NewVectorMean(const IndexMeta &meta) { + switch (meta.data_type()) { + case IndexMeta::DataType::DT_FP16: + return std::make_shared>( + meta.dimension()); + + case IndexMeta::DataType::DT_FP32: + return std::make_shared>(meta.dimension()); + + case IndexMeta::DataType::DT_FP64: + return std::make_shared>(meta.dimension()); + + case IndexMeta::DataType::DT_INT8: + return std::make_shared>(meta.dimension()); + + case IndexMeta::DataType::DT_INT4: + return std::make_shared>(meta.dimension()); + + case IndexMeta::DataType::DT_INT16: + return std::make_shared>(meta.dimension()); + + default: + break; + } + // As binary default + return std::make_shared(meta.dimension()); +} + +static inline std::shared_ptr NewVectorMeanArray( + const IndexMeta &meta) { + switch (meta.data_type()) { + case IndexMeta::DataType::DT_FP16: + return std::make_shared< + GeneralVectorMeanArray>>( + meta.dimension()); + + case IndexMeta::DataType::DT_FP32: + return std::make_shared< + GeneralVectorMeanArray>>(meta.dimension()); + + case IndexMeta::DataType::DT_FP64: + return std::make_shared< + GeneralVectorMeanArray>>( + meta.dimension()); + + case IndexMeta::DataType::DT_INT8: + return std::make_shared< + GeneralVectorMeanArray>>( + meta.dimension()); + + case IndexMeta::DataType::DT_INT4: + return std::make_shared< + GeneralVectorMeanArray>>(meta.dimension()); + + case IndexMeta::DataType::DT_INT16: + return std::make_shared< + GeneralVectorMeanArray>>( + meta.dimension()); + + default: + break; + } + // As binary default + return std::make_shared>( + meta.dimension()); +} + +static inline std::shared_ptr NewVectorMeanArray( + const IndexMeta &meta, const IndexCluster::CentroidList ¢s) { + switch (meta.data_type()) { + case IndexMeta::DataType::DT_FP16: { + auto ptr = std::make_shared< + GeneralVectorMeanArray>>( + meta.dimension()); + + for (const auto &it : cents) { + ptr->emplace(reinterpret_cast(it.feature()), + meta.dimension(), it.follows()); + } + return ptr; + } + + case IndexMeta::DataType::DT_FP32: { + auto ptr = + std::make_shared>>( + meta.dimension()); + + for (const auto &it : cents) { + ptr->emplace(reinterpret_cast(it.feature()), + meta.dimension(), it.follows()); + } + return ptr; + } + + case IndexMeta::DataType::DT_FP64: { + auto ptr = + std::make_shared>>( + meta.dimension()); + + for (const auto &it : cents) { + ptr->emplace(reinterpret_cast(it.feature()), + meta.dimension(), it.follows()); + } + return ptr; + } + + case IndexMeta::DataType::DT_INT8: { + auto ptr = + std::make_shared>>( + meta.dimension()); + + for (const auto &it : cents) { + ptr->emplace(reinterpret_cast(it.feature()), + meta.dimension(), it.follows()); + } + return ptr; + } + + case IndexMeta::DataType::DT_INT4: { + auto ptr = + std::make_shared>>( + meta.dimension()); + + for (const auto &it : cents) { + ptr->emplace(reinterpret_cast(it.feature()), + meta.dimension(), it.follows()); + } + return ptr; + } + + case IndexMeta::DataType::DT_INT16: { + auto ptr = std::make_shared< + GeneralVectorMeanArray>>( + meta.dimension()); + + for (const auto &it : cents) { + ptr->emplace(reinterpret_cast(it.feature()), + meta.dimension(), it.follows()); + } + return ptr; + } + + default: + break; + } + + // As binary default + auto ptr = std::make_shared>( + meta.dimension()); + + for (const auto &it : cents) { + ptr->emplace(it.feature(), meta.dimension(), it.follows()); + } + return ptr; +} + +static inline double CalculateSSE(const IndexCluster::CentroidList ¢s) { + double accum = 0.0; + for (const auto &it : cents) { + accum += it.score(); + } + return accum; +} + +static inline void PurgeCentroids(IndexCluster::CentroidList ¢s, + bool cutting) { + size_t index = 0; + size_t tamp = cents.size(); + + while (index < tamp) { + if (cents[index].follows() == 0) { + size_t last_index = tamp - 1; + + if (index != last_index) { + std::swap(cents[index], cents[last_index]); + } + tamp = last_index; + continue; + } + ++index; + } + if (cutting) { + cents.resize(tamp); + } +} + +int KmeansCluster::init(const IndexMeta &meta, const ailego::Params ¶ms) { + meta_ = meta; + this->update_params(params); + + return this->init_seeker(); +} + +int KmeansCluster::cleanup(void) { + features_.reset(); + shard_cluster_scores_.clear(); + shard_cluster_features_.clear(); + shard_cluster_means_.reset(); + batch_means_.reset(); + batch_scores_.clear(); + seeker_->cleanup(); + return 0; +} + +int KmeansCluster::reset(void) { + features_.reset(); + shard_cluster_scores_.clear(); + shard_cluster_features_.clear(); + shard_cluster_means_->clear(); + batch_means_->clear(); + batch_scores_.clear(); + seeker_->reset(); + return 0; +} + +int KmeansCluster::update(const ailego::Params ¶ms) { + this->update_params(params); + return 0; +} + +void KmeansCluster::suggest(uint32_t k) { + cluster_count_ = k; +} + +int KmeansCluster::mount(IndexFeatures::Pointer feats) { + if (!feats) { + return IndexError_InvalidArgument; + } + if (!feats->is_matched(meta_)) { + return IndexError_Mismatch; + } + + // Check dimension + auto data_type = meta_.data_type(); + switch (data_type) { + case IndexMeta::DataType::DT_INT4: + if (feats->dimension() % 2 != 0) { + LOG_ERROR( + "Unsupported feature dimension %zu (dimension of int4 " + "must be an integer multiple of 2).", + feats->dimension()); + return IndexError_Mismatch; + } + break; + case IndexMeta::DataType::DT_BINARY32: + if (feats->dimension() % 32 != 0) { + LOG_ERROR( + "Unsupported feature dimension %zu (dimension of binary32 " + "must be an integer multiple of 32).", + feats->dimension()); + return IndexError_Mismatch; + } + break; + case IndexMeta::DataType::DT_BINARY64: + if (feats->dimension() % 64 != 0) { + LOG_ERROR( + "Unsupported feature dimension %zu (dimension of binary64 " + "must be an integer multiple of 64).", + feats->dimension()); + return IndexError_Mismatch; + } + break; + default: + break; + } + + features_ = std::move(feats); + return 0; +} + +int KmeansCluster::cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + ailego::ElapsedTime stamp; + + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + if (cents.empty()) { + if (cluster_count_ == 0) { + LOG_ERROR("The count of cluster is unknown."); + return IndexError_NoReady; + } + this->init_centroids(cluster_count_, ¢s); + } + + if (batch_) { + batch_means_ = NewVectorMeanArray(meta_, cents); + batch_scores_.clear(); + for (const auto &it : cents) { + batch_scores_.push_back(it.score()); + } + } + + double cost = 0.0; + + // we need to do clustering and update the centroids' follows, even if + // cents.size() == 1. Otherwise, the centroid with empty follows will be + // removed if purge_empty enabled + for (uint32_t i = 0; (i < max_iterations_) && (cents.size() > 0); ++i) { + double new_cost, new_epsilon; + + int result = this->clustering(threads.get(), cents, &new_cost); + if (result != 0) { + LOG_ERROR("(%u) Failed to cluster.", i + 1); + return result; + } + + new_epsilon = new_cost - cost; + LOG_DEBUG("(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> %f = %f", + i, cents.size(), features_->count(), + (size_t)stamp.milli_seconds(), cost, new_cost, new_epsilon); + stamp.reset(); + + new_epsilon = std::abs(new_epsilon); + if (new_epsilon < epsilon_) { + break; + } + cost = new_cost; + } + + // Purge the empty centroids + PurgeCentroids(cents, purge_empty_); + return 0; +} + +int KmeansCluster::classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (cents.empty()) { + LOG_ERROR("The input centroid's list is empty."); + return IndexError_InvalidArgument; + } + + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + int result = this->build_seeker(cents); + if (result != 0) { + LOG_ERROR("Failed to build the seeker."); + return result; + } + + this->update_clusters(threads.get(), cents); + this->update_features(threads.get(), cents); + return 0; +} + +int KmeansCluster::label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) { + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (cents.empty()) { + LOG_ERROR("The input centroid's list is empty."); + return IndexError_InvalidArgument; + } + + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + int result = this->build_seeker(cents); + if (result != 0) { + LOG_ERROR("Failed to build the seeker."); + return result; + } + + this->update_labels(threads.get(), out); + return 0; +} + +bool KmeansCluster::is_valid(void) const { + if (!seeker_ || !features_ || !features_->count()) { + return false; + } + return true; +} + +int KmeansCluster::clustering(IndexThreads *threads, + IndexCluster::CentroidList ¢s, double *cost) { + int result = this->build_seeker(cents); + if (result != 0) { + LOG_ERROR("Failed to build the seeker."); + return result; + } + + this->split_clusters(threads, cents); + this->update_centroids(threads, cents); + *cost = CalculateSSE(cents); + return 0; +} + +void KmeansCluster::update_params(const ailego::Params ¶ms) { + params.get(GENERAL_THREAD_COUNT, &thread_count_); + params.get(GENERAL_CLUSTER_COUNT, &cluster_count_); + params.get(KMEANS_CLUSTER_COUNT, &cluster_count_); + params.get(KMEANS_CLUSTER_SHARD_FACTOR, &shard_factor_); + params.get(KMEANS_CLUSTER_EPSILON, &epsilon_); + params.get(KMEANS_CLUSTER_MAX_ITERATIONS, &max_iterations_); + params.get(KMEANS_CLUSTER_BATCH, &batch_); + params.get(KMEANS_CLUSTER_PURGE_EMPTY, &purge_empty_); +} + +int KmeansCluster::init_seeker(void) { + seeker_.reset(new (std::nothrow) LinearSeeker); + if (!seeker_) { + LOG_ERROR("Failed to create linear seeker."); + return IndexError_NoMemory; + } + + int result = seeker_->init(meta_); + if (result != 0) { + LOG_ERROR("Failed to initialize linear seeker."); + return result; + } + + return 0; +} + +int KmeansCluster::build_seeker(const IndexCluster::CentroidList ¢s) { + int result = + seeker_->mount(std::make_shared(meta_, cents)); + if (result != 0) { + LOG_ERROR("Failed to mount features for linear seeker."); + return result; + } + + return 0; +} + +bool KmeansCluster::check_centroids(const IndexCluster::CentroidList ¢s) { + for (const auto &it : cents) { + if (it.size() != meta_.element_size()) { + return false; + } + } + return true; +} + +void KmeansCluster::init_centroids(size_t count, + IndexCluster::CentroidList *out) { + size_t feature_size = features_->element_size(); + size_t features_count = features_->count(); + size_t sample_count = std::min(count, features_count); + + ailego::Reservoir sampler(sample_count); + for (size_t i = 0; i < features_count; ++i) { + sampler.fill(i); + } + + // Save centroids + out->reserve(sampler.pool().size()); + for (auto i : sampler.pool()) { + out->emplace_back(features_->element(i), feature_size); + } +} + +void KmeansCluster::init_containers(size_t shard_count) { + if (!shard_cluster_means_) { + shard_cluster_means_ = NewVectorMeanArray(meta_); + } + shard_cluster_means_->clear(); + shard_cluster_means_->resize(shard_count); + shard_cluster_scores_.clear(); + shard_cluster_scores_.resize(shard_count); +} + +void KmeansCluster::init_features_containers(size_t shard_count) { + shard_cluster_features_.resize(shard_count); + for (auto &features : shard_cluster_features_) { + features.clear(); + } +} + +void KmeansCluster::split_clusters(IndexThreads *threads, + const IndexCluster::CentroidList ¢s) { + // Initilize containers + this->init_containers(threads->count() * cents.size()); + auto task_group = threads->make_group(); + + // Initilize base information + size_t features_count = features_->count(); + size_t shard_count = std::max( + static_cast(std::ceil(threads->count() * shard_factor_)), 1u); + size_t fregment_count = (features_count + shard_count - 1) / shard_count; + + for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count); + ++i) { + size_t next_index = index + fregment_count; + if (next_index > features_count) { + next_index = features_count; + } + + // Process in work thread + task_group->submit( + ailego::Closure::New(this, &KmeansCluster::split_clusters_thread, index, + next_index, threads)); + + // Next index + index = next_index; + } + task_group->wait_finish(); +} + +void KmeansCluster::update_centroids(IndexThreads *threads, + IndexCluster::CentroidList ¢s) { + auto task_group = threads->make_group(); + for (size_t i = 0; i < cents.size(); ++i) { + task_group->submit(ailego::Closure::New( + this, &KmeansCluster::update_centroid_thread, i, ¢s)); + } + task_group->wait_finish(); +} + +void KmeansCluster::update_clusters(IndexThreads *threads, + const IndexCluster::CentroidList ¢s) { + // Initilize containers + this->init_features_containers(threads->count() * cents.size()); + auto task_group = threads->make_group(); + + size_t features_count = features_->count(); + size_t shard_count = std::max( + static_cast(std::ceil(threads->count() * shard_factor_)), 1u); + size_t fregment_count = (features_count + shard_count - 1) / shard_count; + + for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count); + ++i) { + size_t next_index = index + fregment_count; + if (next_index > features_count) { + next_index = features_count; + } + // Process in work thread + task_group->submit( + ailego::Closure::New(this, &KmeansCluster::update_cluster_thread, index, + next_index, threads)); + + // Next index + index = next_index; + } + task_group->wait_finish(); +} + +void KmeansCluster::update_features(IndexThreads *threads, + IndexCluster::CentroidList ¢s) { + auto task_group = threads->make_group(); + for (size_t i = 0; i < cents.size(); ++i) { + // Process in work thread + task_group->submit(ailego::Closure::New( + this, &KmeansCluster::update_features_thread, i, ¢s)); + } + task_group->wait_finish(); +} + +void KmeansCluster::update_labels(IndexThreads *threads, + std::vector *labels) { + size_t features_count = features_->count(); + size_t shard_count = std::max( + static_cast(std::ceil(threads->count() * shard_factor_)), 1u); + size_t fregment_count = (features_count + shard_count - 1) / shard_count; + auto task_group = threads->make_group(); + + // Prepare buffer + labels->resize(features_count); + + for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count); + ++i) { + size_t next_index = index + fregment_count; + if (next_index > features_count) { + next_index = features_count; + } + + // Process in work thread + task_group->submit(ailego::Closure::New( + this, &KmeansCluster::update_labels_thread, index, next_index, labels)); + + // Next index + index = next_index; + } + task_group->wait_finish(); +} + +void KmeansCluster::split_clusters_thread(size_t index_begin, size_t index_end, + const IndexThreads *threads) { + size_t feature_size = features_->element_size(); + size_t thread_offset = threads->indexof_this() * seeker_->original()->count(); + + for (size_t i = index_begin; i != index_end; ++i) { + const void *feat = features_->element(i); + LinearSeeker::Document result(0, std::numeric_limits::max()); + + // ignore error + seeker_->seek(feat, meta_.element_size(), &result); + + size_t sel_column = thread_offset + result.index; + shard_cluster_scores_[sel_column] += result.score; + shard_cluster_means_->at(sel_column).plus(feat, feature_size); + } +} + +void KmeansCluster::update_centroid_thread(size_t column, + IndexCluster::CentroidList *out) { + size_t cluster_count = out->size(); + double cluster_score = 0.0; + + // Create Accumulator + std::shared_ptr accum = NewVectorMean(meta_); + if (batch_) { + cluster_score += batch_scores_[column]; + accum->merge(batch_means_->at(column)); + } + + // Compute the score of centroid + for (size_t i = column; i < shard_cluster_scores_.size(); + i += cluster_count) { + cluster_score += shard_cluster_scores_[i]; + accum->merge(shard_cluster_means_->at(i)); + } + + // Update centroid + IndexCluster::Centroid *centroid = &(out->at(column)); + centroid->set_score(cluster_score); + centroid->set_follows(accum->count()); + accum->mean(centroid->mutable_buffer()); +} + +void KmeansCluster::update_cluster_thread(size_t index_begin, size_t index_end, + const IndexThreads *threads) { + size_t thread_offset = threads->indexof_this() * seeker_->original()->count(); + + for (size_t i = index_begin; i != index_end; ++i) { + const void *feat = features_->element(i); + LinearSeeker::Document result(0, std::numeric_limits::max()); + + // ignore error + seeker_->seek(feat, meta_.element_size(), &result); + + size_t sel_column = thread_offset + result.index; + shard_cluster_features_[sel_column].emplace_back(feat); + } +} + +void KmeansCluster::update_features_thread(size_t column, + IndexCluster::CentroidList *out) { + size_t cluster_count = out->size(); + size_t cluster_follows = 0u; + + // Compute the follows of cluster + for (size_t i = column; i < shard_cluster_features_.size(); + i += cluster_count) { + cluster_follows += shard_cluster_features_[i].size(); + } + + // Merge all features in cluster + std::vector &cluster_features = + *(out->at(column).mutable_similars()); + cluster_features.resize(cluster_follows); + + for (size_t i = column, j = 0; i < shard_cluster_features_.size(); + i += cluster_count) { + const std::vector &it = shard_cluster_features_[i]; + if (!it.empty()) { + std::memcpy(&cluster_features[j], it.data(), it.size() * sizeof(void *)); + j += it.size(); + } + } +} + +void KmeansCluster::update_labels_thread(size_t index_begin, size_t index_end, + std::vector *labels) { + for (size_t i = index_begin; i != index_end; ++i) { + const void *feat = features_->element(i); + LinearSeeker::Document result(0, std::numeric_limits::max()); + + // ignore error + seeker_->seek(feat, meta_.element_size(), &result); + (*labels)[i] = static_cast(result.index); + } +} + +INDEX_FACTORY_REGISTER_CLUSTER_ALIAS(KmeansCluster, KmeansCluster, false); +INDEX_FACTORY_REGISTER_CLUSTER_ALIAS(BatchKmeansCluster, KmeansCluster, true); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/linear_seeker.cc b/src/core/algorithm/cluster/linear_seeker.cc new file mode 100644 index 0000000..bcf55af --- /dev/null +++ b/src/core/algorithm/cluster/linear_seeker.cc @@ -0,0 +1,45 @@ + +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "linear_seeker.h" + +namespace zvec { +namespace core { + +int LinearSeeker::seek(const void *query, size_t len, Document *out) { + if (ailego_unlikely(!query || !out || meta_.element_size() != len)) { + return IndexError_InvalidArgument; + } + + float sel_score = std::numeric_limits::max(); + uint32_t sel_column = 0; + uint32_t total = static_cast(features_->count()); + + for (uint32_t i = 0; i < total; ++i) { + float score = 0.0f; + + distance_func_(features_->element(i), query, meta_.dimension(), &score); + if (score < sel_score) { + sel_score = score; + sel_column = i; + } + } + + out->index = sel_column; + out->score = sel_score; + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/linear_seeker.h b/src/core/algorithm/cluster/linear_seeker.h new file mode 100644 index 0000000..28f0cfa --- /dev/null +++ b/src/core/algorithm/cluster/linear_seeker.h @@ -0,0 +1,98 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "seeker.h" + +namespace zvec { +namespace core { + +/*! Linear Seeker + */ +class LinearSeeker : public Seeker { + public: + typedef std::shared_ptr Pointer; + + //! Constructor + LinearSeeker(void) : meta_(), metric_(), features_() {} + + //! Destructor + ~LinearSeeker(void) {} + + //! Initialize Seeker + int init(const IndexMeta &meta) override { + meta_ = meta; + + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Create Metric %s failed.", meta_.metric_name().c_str()); + + return IndexError_Unsupported; + } + int ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMetric init failed wit ret %d.", ret); + + return ret; + } + distance_func_ = metric_->distance_matrix(1, 1); + if (!distance_func_) { + LOG_ERROR("DistanceMatrix function is nullptr."); + + return IndexError_Unsupported; + } + return 0; + } + + //! Cleanup Seeker + int cleanup(void) override { + features_.reset(); + return 0; + } + + //! Reset Seeker + int reset(void) override { + features_.reset(); + return 0; + } + + //! Mount features + int mount(IndexFeatures::Pointer feats) override { + if (!feats) { + return IndexError_InvalidArgument; + } + if (!feats->is_matched(meta_)) { + return IndexError_Mismatch; + } + features_ = std::move(feats); + return 0; + } + + //! Seek (TOP 1 Document) + int seek(const void *query, size_t len, Document *out) override; + + //! Retrieve the original features + IndexFeatures::Pointer original(void) const override { + return features_; + } + + private: + IndexMeta meta_{}; + IndexMetric::Pointer metric_{}; + IndexFeatures::Pointer features_{}; + IndexMetric::MatrixDistance distance_func_{nullptr}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/multi_chunk_cluster.cc b/src/core/algorithm/cluster/multi_chunk_cluster.cc new file mode 100644 index 0000000..c194796 --- /dev/null +++ b/src/core/algorithm/cluster/multi_chunk_cluster.cc @@ -0,0 +1,440 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "multi_chunk_cluster.h" +#include + +namespace zvec { +namespace core { + +bool MultiChunkClusterAlgorithm::is_valid(void) const { + return features_ && features_->count(); +} + +bool MultiChunkClusterAlgorithm::check_centroids( + const IndexCluster::CentroidList & /*cents*/) { + return true; +} + +void MultiChunkClusterAlgorithm::init_centroids( + size_t count, IndexCluster::CentroidList *out) { + // Just resize, because the get random centroid step is done by cluster_once + out->resize(count); +} + +int MultiChunkClusterAlgorithm::reset(void) { + features_.reset(); + + return 0; +} + +int MultiChunkClusterAlgorithm::cleanup(void) { + features_.reset(); + + return 0; +} + +void MultiChunkClusterAlgorithm::suggest(uint32_t k) { + cluster_count_ = k; +} + +int MultiChunkClusterAlgorithm::update(const ailego::Params ¶ms) { + this->update_params(params); + return 0; +} + +//! MultiChunkCluster +int MultiChunkClusterAlgorithm::update_params(const ailego::Params ¶ms) { + params.get(GENERAL_THREAD_COUNT, &thread_count_); + params.get(GENERAL_CLUSTER_COUNT, &cluster_count_); + + params.get(MULTI_CHUNK_CLUSTER_THREAD_COUNT, &thread_count_); + params.get(MULTI_CHUNK_CLUSTER_COUNT, &cluster_count_); + params.get(MULTI_CHUNK_CLUSTER_CHUNK_COUNT, &chunk_count_); + params.get(MULTI_CHUNK_CLUSTER_EPSILON, &epsilon_); + params.get(MULTI_CHUNK_CLUSTER_MAX_ITERATIONS, &max_iterations_); + params.get(MULTI_CHUNK_CLUSTER_MARKOV_CHAIN_LENGTH, &markov_chain_length_); + + return 0; +} + +int MultiChunkClusterAlgorithm::init_distance_func() { + IndexMetric::Pointer metric_{}; + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Create metric %s failed.", meta_.metric_name().c_str()); + return IndexError_Unsupported; + } + + int ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMetric init failed wit ret %d.", ret); + return ret; + } + + distance_func_ = metric_->distance_matrix(1, 1); + if (!distance_func_) { + LOG_ERROR("DistanceMatrix function is nullptr."); + return IndexError_Unsupported; + } + return 0; +} + +int MultiChunkClusterAlgorithm::do_chunk() { + if (chunk_count_ == 0) { + LOG_ERROR("Invalid Chunk Count: %u", chunk_count_); + return IndexError_InvalidArgument; + } + + size_t large_chunk_count = meta_.dimension() % chunk_count_; + size_t base_chunk_dim_ = meta_.dimension() / chunk_count_; + + chunk_dims_.clear(); + + for (size_t i = 0; i < chunk_count_; ++i) { + if (i < large_chunk_count) { + chunk_dims_.push_back(base_chunk_dim_ + 1); + } else { + chunk_dims_.push_back(base_chunk_dim_); + } + } + + chunk_dim_offsets_.clear(); + chunk_dim_offsets_.push_back(0); + for (size_t i = 1; i < chunk_count_; ++i) { + chunk_dim_offsets_.push_back(chunk_dim_offsets_[i - 1] + + chunk_dims_[i - 1]); + } + chunk_dim_offsets_.push_back(meta_.dimension()); + + return 0; +} + +int MultiChunkClusterAlgorithm::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + meta_ = meta; + + int ret = update_params(params); + if (ret != 0) { + return ret; + } + + ret = do_chunk(); + if (ret != 0) { + return ret; + } + + ret = init_distance_func(); + if (ret != 0) { + return ret; + } + + return 0; +} + +int MultiChunkClusterAlgorithm::mount(IndexFeatures::Pointer features) { + if (!features) { + return IndexError_InvalidArgument; + } + + if (!features->is_matched(meta_)) { + return IndexError_Mismatch; + } + + auto data_type = meta_.data_type(); + if (data_type != IndexMeta::DataType::DT_FP32 && + data_type != IndexMeta::DataType::DT_FP16) { + LOG_ERROR("Unsupported meta type %u", data_type); + + return IndexError_Unsupported; + } + + features_ = std::move(features); + + return 0; +} + +//! cluster +int MultiChunkClusterAlgorithm::cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + if (chunk_count_ == 0) { + LOG_ERROR("Invalid Chunk Count: %u", chunk_count_); + + return IndexError_InvalidArgument; + } + + if (cluster_count_ == 0) { + LOG_ERROR("Invalid cluster Count: %u", cluster_count_); + + return IndexError_InvalidArgument; + } + + if (!threads) { + threads = std::make_shared(thread_count_, false); + } + + auto task_group = threads->make_group(); + if (!task_group) { + LOG_ERROR("Failed to create task group"); + return IndexError_Runtime; + } + + cents.clear(); + cents.resize(chunk_count_ * cluster_count_); + + std::atomic finished{0}; + + for (size_t i = 0; i < threads->count(); ++i) { + task_group->submit( + ailego::Closure::New(this, &MultiChunkClusterAlgorithm::do_cluster, i, + threads->count(), ¢s, &finished)); + } + + { + std::unique_lock lk(mutex_); + while (finished.load() < chunk_count_) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to cluster while waiting finish"); + return errcode_; + } + LOG_INFO("Finish Chunk Count %zu, Finished Percent %.3f%%", + finished.load(), finished.load() * 100.0f / chunk_count_); + } + } + + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to cluster while waiting finish"); + return errcode_; + } + + task_group->wait_finish(); + + return 0; +} + +//! Classify +int MultiChunkClusterAlgorithm::classify( + IndexThreads::Pointer /*threads*/, IndexCluster::CentroidList & /*cents*/) { + return IndexError_Unsupported; +} + +//! Label +int MultiChunkClusterAlgorithm::label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) { + if (chunk_count_ == 0) { + LOG_ERROR("Invalid Chunk Count: %u", chunk_count_); + + return IndexError_InvalidArgument; + } + + if (cents.empty()) { + LOG_ERROR("The input centroid's list is empty."); + return IndexError_InvalidArgument; + } + + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + if (cluster_count_ == 0) { + LOG_ERROR("Invalid cluster Count: %u", cluster_count_); + + return IndexError_InvalidArgument; + } + + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + + auto task_group = threads->make_group(); + if (!task_group) { + LOG_ERROR("Failed to create task group"); + return IndexError_Runtime; + } + + size_t features_count = features_->count(); + out->resize(features_count * chunk_count_); + + std::atomic finished{0}; + + for (size_t i = 0; i < threads->count(); ++i) { + task_group->submit( + ailego::Closure::New(this, &MultiChunkClusterAlgorithm::do_label, i, + threads->count(), cents, out, &finished)); + } + + { + std::unique_lock lk(mutex_); + while (finished.load() < features_count) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to cluster while waiting finish"); + return errcode_; + } + LOG_INFO("Finish label cnt %zu, finished percent %.3f%%", finished.load(), + finished.load() * 100.0f / features_count); + } + } + + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to cluster while waiting finish"); + return errcode_; + } + + task_group->wait_finish(); + + return 0; +} + +//! Cluster +int MultiChunkCluster::cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + return algorithm_->cluster(std::move(threads), cents); +} + +//! Classify +int MultiChunkCluster::classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + return algorithm_->classify(std::move(threads), cents); +} + +//! Label +int MultiChunkCluster::label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) { + return algorithm_->label(std::move(threads), cents, out); +} + +//! Update Cluster +int MultiChunkCluster::update(const ailego::Params ¶ms) { + return algorithm_->update(params); +} + +//! Reset Cluster +int MultiChunkCluster::reset(void) { + return algorithm_->reset(); +} + +//! Cleanup Cluster +int MultiChunkCluster::cleanup(void) { + return algorithm_->cleanup(); +} + +//! Suggest dividing to K clusters +void MultiChunkCluster::suggest(uint32_t k) { + algorithm_->suggest(k); +} + +int MultiChunkCluster::mount(IndexFeatures::Pointer feats) { + return algorithm_->mount(feats); +} + +int MultiChunkCluster::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + IndexMeta new_meta(meta.data_type(), meta.dimension()); + + if (meta.metric_name() == "Cosine") { + new_meta.set_dimension(meta.dimension() - 1); + new_meta.set_metric("InnerProduct", 0, ailego::Params()); + } + + auto data_type = new_meta.data_type(); + + if (new_meta.metric_name() == "InnerProduct") { + switch (data_type) { + case IndexMeta::DataType::DT_FP16: { + algorithm_.reset( + new (std::nothrow) + MultiChunkNumericalInnerProductAlgorithm); + break; + } + case IndexMeta::DataType::DT_FP32: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalInnerProductAlgorithm); + break; + } + case IndexMeta::DataType::DT_FP64: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalInnerProductAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT8: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalInnerProductAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT16: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalInnerProductAlgorithm); + break; + } + default: { + LOG_ERROR("Unsupported feature types %d.", data_type); + return IndexError_Mismatch; + } + } + } else { + switch (data_type) { + case IndexMeta::DataType::DT_FP16: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalAlgorithm); + break; + } + case IndexMeta::DataType::DT_FP32: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalAlgorithm); + break; + } + case IndexMeta::DataType::DT_FP64: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT8: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT16: { + algorithm_.reset(new (std::nothrow) + MultiChunkNumericalAlgorithm); + break; + } + default: { + LOG_ERROR("Unsupported feature types %d.", data_type); + return IndexError_Mismatch; + } + } + } + + algorithm_->init(new_meta, params); + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/multi_chunk_cluster.h b/src/core/algorithm/cluster/multi_chunk_cluster.h new file mode 100644 index 0000000..19096d7 --- /dev/null +++ b/src/core/algorithm/cluster/multi_chunk_cluster.h @@ -0,0 +1,482 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include "cluster_params.h" + +namespace zvec { +namespace core { + +//! MultiChunkClusterAlgorithm +class MultiChunkClusterAlgorithm { + public: + typedef std::shared_ptr Pointer; + + //! Constructor + MultiChunkClusterAlgorithm() = default; + + //! Destructor + virtual ~MultiChunkClusterAlgorithm() = default; + + //! Initialize Cluster + int init(const IndexMeta &meta, const ailego::Params ¶ms); + + //! Cleanup Cluster + int cleanup(); + + //! Reset Cluster + int reset(); + + //! Update Cluster + int update(const ailego::Params ¶ms); + + //! Suggest dividing to K clusters + void suggest(uint32_t k); + + //! Mount features + int mount(IndexFeatures::Pointer feats); + + //! Cluster + int cluster(IndexThreads::Pointer threads, IndexCluster::CentroidList ¢s); + + //! Classify + int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Label + int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out); + + + const std::vector &chunk_dims() const { + return chunk_dims_; + } + + const std::vector &chunk_dim_offsets() const { + return chunk_dim_offsets_; + } + + protected: + //! Check Centroids + bool check_centroids(const IndexCluster::CentroidList ¢s); + + //! Test if it is valid + bool is_valid() const; + + //! Do chunk + int do_chunk(); + + //! Update parameters + int update_params(const ailego::Params ¶ms); + + int init_distance_func(); + + //! cluster thread + virtual void do_cluster(size_t idx, size_t chunk_step, + IndexCluster::CentroidList *cents, + std::atomic *finished) = 0; + + //! label thread + virtual void do_label(size_t idx, size_t step, + const IndexCluster::CentroidList ¢s, + std::vector *out, + std::atomic *finished) = 0; + + //! Initialize Centroids + void init_centroids(size_t count, IndexCluster::CentroidList *out); + + private: + constexpr static uint32_t kDefaultLogIntervalSecs = 15U; + + protected: + uint32_t cluster_count_{0u}; + uint32_t thread_count_{0u}; + uint32_t chunk_count_{0u}; + uint32_t max_iterations_{20u}; + bool assumption_free_{false}; + uint32_t markov_chain_length_{32}; + double epsilon_{std::numeric_limits::epsilon()}; + + int errcode_{0}; + std::atomic_bool error_{false}; + uint32_t check_interval_secs_{kDefaultLogIntervalSecs}; + std::mutex mutex_{}; + std::condition_variable cond_{}; + + IndexMeta meta_{}; + IndexFeatures::Pointer features_{}; + + std::vector chunk_dims_; + std::vector chunk_dim_offsets_; + + IndexMetric::MatrixDistance distance_func_{nullptr}; +}; + +/*! Numerical cluster algorithm + */ +template +class MultiChunkNumericalAlgorithm : public MultiChunkClusterAlgorithm { + public: + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Check supporting type + static_assert(ailego::IsArithmetic::value, + "ValueType must be arithmetic"); + + //! Constructor + MultiChunkNumericalAlgorithm() = default; + + //! Destructor + ~MultiChunkNumericalAlgorithm() = default; + + protected: + //! cluster thread + void do_cluster(size_t idx, size_t chunk_step, + IndexCluster::CentroidList *cents, + std::atomic *finished); + + //! label thread + void do_label(size_t idx, size_t step, + const IndexCluster::CentroidList ¢s, + std::vector *out, std::atomic *finished); +}; + +//! cluster thread +template +void MultiChunkNumericalAlgorithm::do_cluster( + size_t idx, size_t chunk_step, IndexCluster::CentroidList *cents, + std::atomic *finished) { + for (size_t chunk = idx; chunk < chunk_count_; chunk += chunk_step) { + auto chunk_dim = chunk_dims_[chunk]; + + ailego::NumericalKmeans algorithm(cluster_count_, + chunk_dim); + + // mount features into algorithm + auto features_count = features_->count(); + + algorithm.feature_matrix_reserve(features_count); + + for (size_t i = 0; i < features_count; ++i) { + auto vec = reinterpret_cast(features_->element(i)); + algorithm.append(vec + chunk_dim_offsets_[chunk], chunk_dim); + } + + IndexThreads::Pointer local_threads = + std::make_shared(1, false); + if (!local_threads) { + error_ = IndexError_NoMemory; + return; + } + + ailego::Kmc2CentroidsGenerator, + IndexThreads> + cent_gen; + cent_gen.set_chain_length(markov_chain_length_); + cent_gen.set_assumption_free(assumption_free_); + cent_gen(&algorithm, *local_threads); + + double cost = 0.0; + + for (uint32_t i = 0; i < max_iterations_; ++i) { + double old_cost, new_epsilon; + old_cost = cost; + + bool result = algorithm.cluster_once(*local_threads, &cost); + if (result != true) { + LOG_ERROR("(%u) Failed to cluster.", i + 1); + errcode_ = -1; + + return; + } + + new_epsilon = std::abs(cost - old_cost); + if (new_epsilon < epsilon_) { + break; + } + } + + auto &chunk_cents = algorithm.centroids(); + + for (size_t i = 0; i < chunk_cents.count(); ++i) { + size_t global_cent_idx = chunk * cluster_count_ + i; + + IndexCluster::Centroid *centroid = &(cents->at(global_cent_idx)); + centroid->set_score(algorithm.context().clusters()[i].cost()); + centroid->set_follows(algorithm.context().clusters()[i].count()); + centroid->set_feature(algorithm.centroids()[i], + chunk_dim * meta_.unit_size()); + } + + LOG_INFO("(%zu) Chunk Done. Clusters Count: %zu, Features: %zu, Cost: %f", + chunk, algorithm.centroids().count(), features_->count(), cost); + + (*finished)++; + { + std::lock_guard lk(mutex_); + cond_.notify_one(); + } + } + + return; +} + +//! label thread +template +void MultiChunkNumericalAlgorithm::do_label( + size_t idx, size_t step, const IndexCluster::CentroidList ¢s, + std::vector *out, std::atomic *finished) { + for (size_t id = idx; id < features_->count(); id += step) { + const T *feat = reinterpret_cast(features_->element(id)); + + for (size_t chunk = 0; chunk < chunk_count_; ++chunk) { + size_t chunk_dim_offset = chunk_dim_offsets_[chunk]; + size_t chunk_dim = chunk_dims_[chunk]; + + uint32_t sel_index = 0; + float sel_score = std::numeric_limits::max(); + + for (uint32_t cluster = 0; cluster < cluster_count_; ++cluster) { + float score{0.0}; + + distance_func_(cents[chunk * cluster_count_ + cluster].feature(), + feat + chunk_dim_offset, chunk_dim, &score); + + if (score < sel_score) { + sel_score = score; + sel_index = cluster; + } + } + + (*out)[id * chunk_count_ + chunk] = static_cast(sel_index); + } + + (*finished)++; + { + std::lock_guard lk(mutex_); + cond_.notify_one(); + } + } +} + +/*! Inner Product Cluster Algorithm + */ +template +class MultiChunkNumericalInnerProductAlgorithm + : public MultiChunkClusterAlgorithm { + public: + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Check supporting type + static_assert(ailego::IsArithmetic::value, + "ValueType must be arithmetic"); + + //! Constructor + MultiChunkNumericalInnerProductAlgorithm() = default; + + //! Destructor + ~MultiChunkNumericalInnerProductAlgorithm() = default; + + protected: + //! cluster thread + void do_cluster(size_t idx, size_t chunk_step, + IndexCluster::CentroidList *cents, + std::atomic *finished); + + //! label thread + void do_label(size_t idx, size_t chunk_step, + const IndexCluster::CentroidList ¢s, + std::vector *out, std::atomic *finished); +}; + +//! cluster thread +template +void MultiChunkNumericalInnerProductAlgorithm::do_cluster( + size_t idx, size_t chunk_step, IndexCluster::CentroidList *cents, + std::atomic *finished) { + for (size_t chunk = idx; chunk < chunk_count_; chunk += chunk_step) { + auto chunk_dim = chunk_dims_[chunk]; + + ailego::NumericalInnerProductKmeans algorithm( + cluster_count_, chunk_dim); + + // mount features into algorithm + auto features_count = features_->count(); + + algorithm.feature_matrix_reserve(features_count); + + for (size_t i = 0; i < features_count; ++i) { + auto vec = reinterpret_cast(features_->element(i)); + algorithm.append(vec + chunk_dim_offsets_[chunk], chunk_dim); + } + + IndexThreads::Pointer local_threads = + std::make_shared(1, false); + if (!local_threads) { + error_ = IndexError_NoMemory; + return; + } + + ailego::Kmc2CentroidsGenerator< + ailego::NumericalInnerProductKmeans, IndexThreads> + cent_gen; + cent_gen.set_chain_length(markov_chain_length_); + cent_gen.set_assumption_free(assumption_free_); + cent_gen(&algorithm, *local_threads); + + double cost = 0.0; + + for (uint32_t i = 0; i < max_iterations_; ++i) { + double old_cost, new_epsilon; + old_cost = cost; + + bool result = algorithm.cluster_once(*local_threads, &cost); + if (result != true) { + LOG_ERROR("(%zu) Failed to cluster.", (size_t)(i + 1)); + errcode_ = -1; + + return; + } + + new_epsilon = std::abs(cost - old_cost); + if (new_epsilon < epsilon_) { + break; + } + } + + auto &chunk_cents = algorithm.centroids(); + + for (size_t i = 0; i < chunk_cents.count(); ++i) { + size_t global_cent_idx = chunk * cluster_count_ + i; + + IndexCluster::Centroid *centroid = &(cents->at(global_cent_idx)); + centroid->set_score(algorithm.context().clusters()[i].cost()); + centroid->set_follows(algorithm.context().clusters()[i].count()); + centroid->set_feature(algorithm.centroids()[i], + chunk_dim * meta_.unit_size()); + } + + LOG_INFO("(%zu) Chunk Done. Clusters Count: %zu, Features: %zu, Cost: %f", + chunk, algorithm.centroids().count(), features_->count(), cost); + + (*finished)++; + { + std::lock_guard lk(mutex_); + cond_.notify_one(); + } + } +} + +//! label thread +template +void MultiChunkNumericalInnerProductAlgorithm::do_label( + size_t idx, size_t step, const IndexCluster::CentroidList ¢s, + std::vector *out, std::atomic *finished) { + for (size_t id = idx; id < features_->count(); id += step) { + const T *feat = reinterpret_cast(features_->element(id)); + + for (size_t chunk = 0; chunk < chunk_count_; ++chunk) { + size_t chunk_dim_offset = chunk_dim_offsets_[chunk]; + size_t chunk_dim = chunk_dims_[chunk]; + + uint32_t sel_index = 0; + float sel_score = std::numeric_limits::max(); + + for (uint32_t cluster = 0; cluster < cluster_count_; ++cluster) { + float score{0.0}; + + distance_func_(cents[chunk * cluster_count_ + cluster].feature(), + feat + chunk_dim_offset, chunk_dim, &score); + + if (score < sel_score) { + sel_score = score; + sel_index = cluster; + } + } + + (*out)[id * chunk_count_ + chunk] = sel_index; + } + + (*finished)++; + { + std::lock_guard lk(mutex_); + cond_.notify_one(); + } + } +} + +//! MultiChunkCluster +class MultiChunkCluster { + public: + std::shared_ptr Pointer; + + //! Constructor + MultiChunkCluster() = default; + + //! Destructor + ~MultiChunkCluster() = default; + + //! Initialize Cluster + int init(const IndexMeta &meta, const ailego::Params ¶ms); + + //! Cleanup Cluster + int cleanup(); + + //! Reset Cluster + int reset(); + + //! Update Cluster + int update(const ailego::Params ¶ms); + + //! Suggest dividing to K clusters + void suggest(uint32_t k); + + //! Mount features + int mount(IndexFeatures::Pointer feats); + + //! Cluster + int cluster(IndexThreads::Pointer threads, IndexCluster::CentroidList ¢s); + + //! Classify + int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Label + int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out); + + const std::vector &chunk_dims() const { + return algorithm_->chunk_dims(); + } + + const std::vector &chunk_dim_offsets() const { + return algorithm_->chunk_dim_offsets(); + } + + protected: + MultiChunkClusterAlgorithm::Pointer algorithm_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/opt_kmeans_cluster.cc b/src/core/algorithm/cluster/opt_kmeans_cluster.cc new file mode 100644 index 0000000..449d0b9 --- /dev/null +++ b/src/core/algorithm/cluster/opt_kmeans_cluster.cc @@ -0,0 +1,1194 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include +#include +#include +#include +#include "cluster_params.h" + +namespace zvec { +namespace core { + +/*! Optimize K-Means cluster algorithm + */ +class OptKmeansAlgorithm : public IndexCluster { + public: + //! Constructor + OptKmeansAlgorithm(void) {} + + //! Destructor + ~OptKmeansAlgorithm(void) override {} + + //! Initialize Cluster + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Mount features + int mount(IndexFeatures::Pointer feats) override; + + //! Suggest dividing to K clusters + void suggest(uint32_t k) override; + + //! Classify + int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + //! Label + int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) override; + + //! Cluster + int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override = 0; + + //! Cleanup Cluster + int cleanup(void) override; + + //! Reset Cluster + int reset(void) override; + + //! Update Cluster + int update(const ailego::Params ¶ms) override; + + protected: + //! Update parameters + void update_params(const ailego::Params ¶ms); + + //! Init Kmeans Algorithm + int init_algorithm(); + + //! Init Distance function + int init_distance_func(); + + //! Check Centroids + bool check_centroids(const IndexCluster::CentroidList ¢s); + + //! Test if it is valid + bool is_valid(void) const; + + //! Update Clusters + void update_clusters(IndexThreads *threads, + const IndexCluster::CentroidList ¢s); + + //! Update Cluster in Thread + void update_cluster_thread(size_t index_begin, size_t index_end, + const IndexThreads *threads, + const IndexCluster::CentroidList ¢s); + + //! Initialize Shard Features Containers + void init_features_containers(size_t shard_count); + + //! Update Clusters' Features + void update_features(IndexThreads *threads, + IndexCluster::CentroidList ¢s); + + //! Update Cluster's Features in Thread + void update_features_thread(size_t column, IndexCluster::CentroidList *out); + + //! Update Labels + void update_labels(IndexThreads *threads, std::vector *labels, + const IndexCluster::CentroidList ¢s); + + //! Update Labels in Thread + void update_labels_thread(size_t index_begin, size_t index_end, + std::vector *labels, + const IndexCluster::CentroidList ¢s); + + //! Initialize Centroids + void init_centroids(size_t count, IndexCluster::CentroidList *out); + + protected: + uint32_t cluster_count_{0u}; + uint32_t thread_count_{0u}; + uint32_t max_iterations_{20u}; + double epsilon_{std::numeric_limits::epsilon()}; + float shard_factor_{16.0f}; + bool purge_empty_{false}; + bool assumption_free_{false}; + uint32_t markov_chain_length_{32}; + IndexMeta meta_{}; + IndexFeatures::Pointer features_{}; + std::vector> shard_cluster_features_{}; + IndexMetric::MatrixDistance distance_func_{nullptr}; +}; + +bool OptKmeansAlgorithm::is_valid(void) const { + if (!features_ || !features_->count()) { + return false; + } + return true; +} + +bool OptKmeansAlgorithm::check_centroids( + const IndexCluster::CentroidList ¢s) { + for (const auto &it : cents) { + if (it.size() != meta_.element_size()) { + return false; + } + } + return true; +} + +void OptKmeansAlgorithm::update_params(const ailego::Params ¶ms) { + params.get(GENERAL_THREAD_COUNT, &thread_count_); + params.get(GENERAL_CLUSTER_COUNT, &cluster_count_); + params.get(OPTKMEANS_CLUSTER_COUNT, &cluster_count_); + params.get(OPTKMEANS_CLUSTER_SHARD_FACTOR, &shard_factor_); + params.get(OPTKMEANS_CLUSTER_EPSILON, &epsilon_); + params.get(OPTKMEANS_CLUSTER_MAX_ITERATIONS, &max_iterations_); + params.get(OPTKMEANS_CLUSTER_PURGE_EMPTY, &purge_empty_); + params.get(OPTKMEANS_CLUSTER_MARKOV_CHAIN_LENGTH, &markov_chain_length_); + params.get(OPTKMEANS_CLUSTER_ASSUMPTION_FREE, &assumption_free_); +} + +int OptKmeansAlgorithm::init_distance_func() { + IndexMetric::Pointer metric_{}; + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Create Metric %s failed.", meta_.metric_name().c_str()); + return IndexError_Unsupported; + } + int ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMetric init failed wit ret %d.", ret); + return ret; + } + distance_func_ = metric_->distance_matrix(1, 1); + if (!distance_func_) { + LOG_ERROR("DistanceMatrix function is nullptr."); + return IndexError_Unsupported; + } + return 0; +} + +void OptKmeansAlgorithm::update_clusters( + IndexThreads *threads, const IndexCluster::CentroidList ¢s) { + // Initilize containers + this->init_features_containers(threads->count() * cents.size()); + auto task_group = threads->make_group(); + + size_t features_count = features_->count(); + size_t shard_count = std::max( + static_cast(std::ceil(threads->count() * shard_factor_)), 1u); + size_t fregment_count = (features_count + shard_count - 1) / shard_count; + + for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count); + ++i) { + size_t next_index = index + fregment_count; + if (next_index > features_count) { + next_index = features_count; + } + + // Process in work thread· + task_group->submit( + ailego::Closure::New(this, &OptKmeansAlgorithm::update_cluster_thread, + index, next_index, threads, cents)); + + // Next index + index = next_index; + } + task_group->wait_finish(); +} + +void OptKmeansAlgorithm::update_cluster_thread( + size_t index_begin, size_t index_end, const IndexThreads *threads, + const IndexCluster::CentroidList ¢s) { + size_t thread_offset = threads->indexof_this() * cents.size(); + + for (size_t i = index_begin; i != index_end; ++i) { + const void *feat = features_->element(i); + uint32_t sel_index = 0; + float sel_score = std::numeric_limits::max(); + + // todo: get min distance + uint32_t total = static_cast(cents.size()); + for (uint32_t j = 0; j < total; ++j) { + float score = 0.0f; + + distance_func_(cents[j].feature(), feat, meta_.dimension(), &score); + if (score < sel_score) { + sel_score = score; + sel_index = j; + } + } + + size_t sel_column = thread_offset + sel_index; + shard_cluster_features_[sel_column].emplace_back(feat); + } +} + +void OptKmeansAlgorithm::init_features_containers(size_t shard_count) { + shard_cluster_features_.resize(shard_count); + for (auto &features : shard_cluster_features_) { + features.clear(); + } +} + +void OptKmeansAlgorithm::update_features(IndexThreads *threads, + IndexCluster::CentroidList ¢s) { + auto task_group = threads->make_group(); + for (size_t i = 0; i < cents.size(); ++i) { + // Process in work thread + task_group->submit(ailego::Closure::New( + this, &OptKmeansAlgorithm::update_features_thread, i, ¢s)); + } + task_group->wait_finish(); +} + +void OptKmeansAlgorithm::update_labels( + IndexThreads *threads, std::vector *labels, + const IndexCluster::CentroidList ¢s) { + size_t features_count = features_->count(); + size_t shard_count = std::max( + static_cast(std::ceil(threads->count() * shard_factor_)), 1u); + size_t fregment_count = (features_count + shard_count - 1) / shard_count; + auto task_group = threads->make_group(); + + // Prepare buffer + labels->resize(features_count); + + for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count); + ++i) { + size_t next_index = index + fregment_count; + if (next_index > features_count) { + next_index = features_count; + } + + // Process in work thread + task_group->submit( + ailego::Closure::New(this, &OptKmeansAlgorithm::update_labels_thread, + index, next_index, labels, cents)); + + // Next index + index = next_index; + } + task_group->wait_finish(); +} + +void OptKmeansAlgorithm::update_labels_thread( + size_t index_begin, size_t index_end, std::vector *labels, + const IndexCluster::CentroidList ¢s) { + for (size_t i = index_begin; i != index_end; ++i) { + const void *feat = features_->element(i); + + uint32_t sel_index = 0; + float sel_score = std::numeric_limits::max(); + + // todo: get min distance + uint32_t total = static_cast(cents.size()); + for (uint32_t j = 0; j < total; ++j) { + float score = 0.0f; + + distance_func_(cents[j].feature(), feat, meta_.dimension(), &score); + if (score < sel_score) { + sel_score = score; + sel_index = j; + } + } + + (*labels)[i] = static_cast(sel_index); + } +} + +void OptKmeansAlgorithm::init_centroids(size_t count, + IndexCluster::CentroidList *out) { + // Just resize, because the get random centroid step is done by cluster_once + out->resize(count); +} + +void OptKmeansAlgorithm::update_features_thread( + size_t column, IndexCluster::CentroidList *out) { + size_t cluster_count = out->size(); + size_t cluster_follows = 0u; + + // Compute the follows of cluster + for (size_t i = column; i < shard_cluster_features_.size(); + i += cluster_count) { + cluster_follows += shard_cluster_features_[i].size(); + } + + // Merge all features in cluster + std::vector &cluster_features = + *(out->at(column).mutable_similars()); + cluster_features.resize(cluster_follows); + + for (size_t i = column, j = 0; i < shard_cluster_features_.size(); + i += cluster_count) { + const std::vector &it = shard_cluster_features_[i]; + if (!it.empty()) { + std::memcpy(&cluster_features[j], it.data(), it.size() * sizeof(void *)); + j += it.size(); + } + } +} + +static inline void PurgeCentroids(IndexCluster::CentroidList ¢s, + bool cutting) { + size_t index = 0; + size_t tamp = cents.size(); + + while (index < tamp) { + if (cents[index].follows() == 0) { + size_t last_index = tamp - 1; + + if (index != last_index) { + std::swap(cents[index], cents[last_index]); + } + tamp = last_index; + continue; + } + ++index; + } + if (cutting) { + cents.resize(tamp); + } +} + +int OptKmeansAlgorithm::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + meta_ = meta; + this->update_params(params); + + return init_distance_func(); +} + +int OptKmeansAlgorithm::mount(IndexFeatures::Pointer feats) { + if (!feats) { + return IndexError_InvalidArgument; + } + if (!feats->is_matched(meta_)) { + return IndexError_Mismatch; + } + + // Check dimension + auto type_ = meta_.data_type(); + switch (type_) { + case IndexMeta::DataType::DT_INT4: + if (feats->dimension() % 8 != 0) { + LOG_ERROR( + "Unsupported feature dimension %zu (dimension of int4 " + "must be an integer multiple of 8).", + feats->dimension()); + return IndexError_Mismatch; + } + break; + case IndexMeta::DataType::DT_INT8: + if (feats->dimension() % 4 != 0) { + LOG_ERROR( + "Unsupported feature dimension %zu (dimension of int8 " + "must be an integer multiple of 4).", + feats->dimension()); + return IndexError_Mismatch; + } + break; + default: + break; + } + + features_ = std::move(feats); + return 0; +} + +void OptKmeansAlgorithm::suggest(uint32_t k) { + cluster_count_ = k; +} + +int OptKmeansAlgorithm::classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (cents.empty()) { + LOG_ERROR("The input centroid's list is empty."); + return IndexError_InvalidArgument; + } + + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + this->update_clusters(threads.get(), cents); + this->update_features(threads.get(), cents); + return 0; +} + +int OptKmeansAlgorithm::label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) { + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (cents.empty()) { + LOG_ERROR("The input centroid's list is empty."); + return IndexError_InvalidArgument; + } + + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + this->update_labels(threads.get(), out, cents); + return 0; +} + +int OptKmeansAlgorithm::update(const ailego::Params ¶ms) { + this->update_params(params); + // algorithm_->reset(cluster_count_); + return 0; +} + +int OptKmeansAlgorithm::reset(void) { + features_.reset(); + shard_cluster_features_.clear(); + + return 0; +} + +int OptKmeansAlgorithm::cleanup(void) { + features_.reset(); + shard_cluster_features_.clear(); + + return 0; +} + + +/*! Numerical K-Means cluster algorithm + */ +template +class NumericalKmeansAlgorithm : public OptKmeansAlgorithm { + public: + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Check supporting type + static_assert(ailego::IsArithmetic::value, + "ValueType must be arithmetic"); + + //! Constructor + NumericalKmeansAlgorithm(void) {} + + //! Destructor + ~NumericalKmeansAlgorithm(void) override {} + + //! Cluster + int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + protected: + void update_centroids( + IndexCluster::CentroidList ¢s, + const ailego::NumericalKmeans &algorithm); +}; + +template +void NumericalKmeansAlgorithm::update_centroids( + IndexCluster::CentroidList ¢s, + const ailego::NumericalKmeans &algorithm) { + this->init_centroids(algorithm.centroids().count(), ¢s); + for (size_t i = 0; i < cents.size(); ++i) { + IndexCluster::Centroid *centroid = &(cents.at(i)); + centroid->set_score(algorithm.context().clusters()[i].cost()); + centroid->set_follows(algorithm.context().clusters()[i].count()); + centroid->set_feature(algorithm.centroids()[i], + meta_.dimension() * sizeof(T)); + } +} + +template +int NumericalKmeansAlgorithm::cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + ailego::ElapsedTime stamp; + + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + // get cluster algorithm + size_t centroid_count = + cents.empty() + ? std::min(cluster_count_, static_cast(features_->count())) + : cents.size(); + if (centroid_count == 0) { + LOG_ERROR("The count of cluster is unknown."); + return IndexError_NoReady; + } + ailego::NumericalKmeans algorithm(centroid_count, + meta_.dimension()); + + // mount features into algorithm + auto features_count = features_->count(); + auto dim = meta_.dimension(); + + algorithm.feature_matrix_reserve(features_count); + + for (size_t i = 0; i < features_count; ++i) { + auto vec = reinterpret_cast(features_->element(i)); + algorithm.append(vec, dim); + } + + if (!cents.empty()) { + auto centroids = algorithm.mutable_centroids(); + centroids->reserve(cents.size()); + for (const auto &it : cents) { + centroids->append(reinterpret_cast(it.feature()), + meta_.dimension()); + } + } else { + ailego::Kmc2CentroidsGenerator< + ailego::NumericalKmeans, IndexThreads> + g; + g.set_chain_length(markov_chain_length_); + g.set_assumption_free(assumption_free_); + algorithm.init_centroids(*threads, g); + } + + double cost = 0.0; + + for (uint32_t i = 0; i < max_iterations_; ++i) { + double old_cost, new_epsilon; + old_cost = cost; + + bool result = algorithm.cluster_once(*threads, &cost); + if (result != true) { + LOG_ERROR("(%u) Failed to cluster.", i + 1); + return -1; + } + + new_epsilon = std::abs(cost - old_cost); + LOG_DEBUG("(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> %f = %f", + i, algorithm.centroids().count(), features_->count(), + (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon); + stamp.reset(); + + if (new_epsilon < epsilon_) { + break; + } + } + + // update_centroids(cents); + update_centroids(cents, algorithm); + + // Purge the empty centroids + PurgeCentroids(cents, purge_empty_); + return 0; +} + +/*! Nibble K-Means cluster algorithm + */ +template +class NibbleKmeansAlgorithm : public OptKmeansAlgorithm { + public: + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Check supporting type + static_assert(ailego::IsArithmetic::value, + "ValueType must be arithmetic"); + + //! Constructor + NibbleKmeansAlgorithm(void) {} + + //! Destructor + ~NibbleKmeansAlgorithm(void) override {} + + //! Cluster + int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + protected: + //! update centroids + void update_centroids(IndexCluster::CentroidList ¢s, + const ailego::NibbleKmeans &algorithm); +}; + +template +void NibbleKmeansAlgorithm::update_centroids( + IndexCluster::CentroidList ¢s, + const ailego::NibbleKmeans &algorithm) { + this->init_centroids(algorithm.centroids().count(), ¢s); + for (size_t i = 0; i < cents.size(); ++i) { + IndexCluster::Centroid *centroid = &(cents.at(i)); + centroid->set_score(algorithm.context().clusters()[i].cost()); + centroid->set_follows(algorithm.context().clusters()[i].count()); + centroid->set_feature(algorithm.centroids()[i], (meta_.dimension() >> 1)); + } +} + +template +int NibbleKmeansAlgorithm::cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + ailego::ElapsedTime stamp; + + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + // get cluster algorithm + size_t centroid_count = + cents.empty() + ? std::min(cluster_count_, static_cast(features_->count())) + : cents.size(); + if (centroid_count == 0) { + LOG_ERROR("The count of cluster is unknown."); + return IndexError_NoReady; + } + ailego::NibbleKmeans algorithm(centroid_count, + meta_.dimension()); + + // mount features into algorithm + auto features_count = features_->count(); + auto dim = meta_.dimension(); + for (size_t i = 0; i < features_count; ++i) { + auto vec = reinterpret_cast::type *>( + features_->element(i)); + algorithm.append(vec, dim); + } + + if (!cents.empty()) { + auto centroids = algorithm.mutable_centroids(); + centroids->reserve(cents.size()); + for (const auto &it : cents) { + centroids->append( + reinterpret_cast::type *>( + it.feature()), + size_t(meta_.dimension())); + } + } else { + ailego::Kmc2CentroidsGenerator< + ailego::NibbleKmeans, IndexThreads> + g; + g.set_chain_length(markov_chain_length_); + g.set_assumption_free(assumption_free_); + algorithm.init_centroids(*threads, g); + } + + double cost = 0.0; + + for (uint32_t i = 0; i < max_iterations_; ++i) { + double old_cost, new_epsilon; + old_cost = cost; + + bool result = algorithm.cluster_once(*threads, &cost); + if (result != true) { + LOG_ERROR("(%u) Failed to cluster.", i + 1); + return -1; + } + + new_epsilon = std::abs(cost - old_cost); + LOG_DEBUG( + "(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> " + "%f = %f", + i, algorithm.centroids().count(), features_->count(), + (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon); + stamp.reset(); + + if (new_epsilon < epsilon_) { + break; + } + } + + // update centroids + update_centroids(cents, algorithm); + + // Purge the empty centroids + PurgeCentroids(cents, purge_empty_); + return 0; +} + +/*! Numerical K-Means cluster algorithm + */ +template +class NumericalInnerProductKmeansAlgorithm : public OptKmeansAlgorithm { + public: + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Check supporting type + static_assert(ailego::IsArithmetic::value, + "ValueType must be arithmetic"); + + //! Constructor + NumericalInnerProductKmeansAlgorithm(void) {} + + //! Destructor + ~NumericalInnerProductKmeansAlgorithm(void) override {} + + //! Cluster + int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + protected: + void update_centroids( + IndexCluster::CentroidList ¢s, + const ailego::NumericalInnerProductKmeans &algorithm); +}; + +template +void NumericalInnerProductKmeansAlgorithm::update_centroids( + IndexCluster::CentroidList ¢s, + const ailego::NumericalInnerProductKmeans &algorithm) { + this->init_centroids(algorithm.centroids().count(), ¢s); + for (size_t i = 0; i < cents.size(); ++i) { + IndexCluster::Centroid *centroid = &(cents.at(i)); + centroid->set_score(algorithm.context().clusters()[i].cost()); + centroid->set_follows(algorithm.context().clusters()[i].count()); + centroid->set_feature(algorithm.centroids()[i], + meta_.dimension() * sizeof(T)); + } +} + +template +int NumericalInnerProductKmeansAlgorithm::cluster( + IndexThreads::Pointer threads, IndexCluster::CentroidList ¢s) { + ailego::ElapsedTime stamp; + + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + // get cluster algorithm + size_t centroid_count = + cents.empty() + ? std::min(cluster_count_, static_cast(features_->count())) + : cents.size(); + if (centroid_count == 0) { + LOG_ERROR("The count of cluster is unknown."); + return IndexError_NoReady; + } + ailego::NumericalInnerProductKmeans algorithm( + centroid_count, meta_.dimension(), true); + + // mount features into algorithm + auto features_count = features_->count(); + auto dim = meta_.dimension(); + + algorithm.feature_matrix_reserve(features_count); + + for (size_t i = 0; i < features_count; ++i) { + auto vec = reinterpret_cast(features_->element(i)); + algorithm.append(vec, dim); + } + + if (!cents.empty()) { + auto centroids = algorithm.mutable_centroids(); + centroids->reserve(cents.size()); + for (const auto &it : cents) { + centroids->append(reinterpret_cast(it.feature()), + meta_.dimension()); + } + } else { + ailego::Kmc2CentroidsGenerator< + ailego::NumericalInnerProductKmeans, + IndexThreads> + g; + g.set_chain_length(markov_chain_length_); + g.set_assumption_free(assumption_free_); + algorithm.init_centroids(*threads, g); + } + + double cost = 0.0; + + for (uint32_t i = 0; i < max_iterations_; ++i) { + double old_cost, new_epsilon; + old_cost = cost; + + bool result = algorithm.cluster_once(*threads, &cost); + if (result != true) { + LOG_ERROR("(%u) Failed to cluster.", i + 1); + return -1; + } + + new_epsilon = std::abs(cost - old_cost); + LOG_DEBUG("(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> %f = %f", + i, algorithm.centroids().count(), features_->count(), + (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon); + stamp.reset(); + + if (new_epsilon < epsilon_) { + break; + } + } + + // update_centroids(cents); + update_centroids(cents, algorithm); + + // Purge the empty centroids + PurgeCentroids(cents, purge_empty_); + return 0; +} + +/*! Nibble Inner Product K-Means cluster algorithm + */ +template +class NibbleInnerProductKmeansAlgorithm : public OptKmeansAlgorithm { + public: + //! Type of value + using ValueType = typename std::remove_cv::type; + + // Check supporting type + static_assert(ailego::IsArithmetic::value, + "ValueType must be arithmetic"); + + //! Constructor + NibbleInnerProductKmeansAlgorithm(void) {} + + //! Destructor + ~NibbleInnerProductKmeansAlgorithm(void) override {} + + //! Cluster + int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + protected: + //! update centroids + void update_centroids( + IndexCluster::CentroidList ¢s, + const ailego::NibbleInnerProductKmeans &algorithm); +}; + +template +void NibbleInnerProductKmeansAlgorithm::update_centroids( + IndexCluster::CentroidList ¢s, + const ailego::NibbleInnerProductKmeans &algorithm) { + this->init_centroids(algorithm.centroids().count(), ¢s); + for (size_t i = 0; i < cents.size(); ++i) { + IndexCluster::Centroid *centroid = &(cents.at(i)); + centroid->set_score(algorithm.context().clusters()[i].cost()); + centroid->set_follows(algorithm.context().clusters()[i].count()); + centroid->set_feature(algorithm.centroids()[i], (meta_.dimension() >> 1)); + } +} + +template +int NibbleInnerProductKmeansAlgorithm::cluster( + IndexThreads::Pointer threads, IndexCluster::CentroidList ¢s) { + ailego::ElapsedTime stamp; + + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + // get cluster algorithm + size_t centroid_count = + cents.empty() + ? std::min(cluster_count_, static_cast(features_->count())) + : cents.size(); + if (centroid_count == 0) { + LOG_ERROR("The count of cluster is unknown."); + return IndexError_NoReady; + } + ailego::NibbleInnerProductKmeans algorithm( + centroid_count, meta_.dimension()); + + // mount features into algorithm + auto features_count = features_->count(); + auto dim = meta_.dimension(); + for (size_t i = 0; i < features_count; ++i) { + auto vec = reinterpret_cast::type *>( + features_->element(i)); + algorithm.append(vec, dim); + } + + if (!cents.empty()) { + auto centroids = algorithm.mutable_centroids(); + centroids->reserve(cents.size()); + for (const auto &it : cents) { + centroids->append( + reinterpret_cast::type *>( + it.feature()), + size_t(meta_.dimension())); + } + } else { + ailego::Kmc2CentroidsGenerator< + ailego::NibbleInnerProductKmeans, IndexThreads> + g; + g.set_chain_length(markov_chain_length_); + g.set_assumption_free(assumption_free_); + algorithm.init_centroids(*threads, g); + } + + double cost = 0.0; + + for (uint32_t i = 0; i < max_iterations_; ++i) { + double old_cost, new_epsilon; + old_cost = cost; + + bool result = algorithm.cluster_once(*threads, &cost); + if (result != true) { + LOG_ERROR("(%u) Failed to cluster.", i + 1); + return -1; + } + + new_epsilon = std::abs(cost - old_cost); + LOG_DEBUG( + "(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> " + "%f = %f", + i, algorithm.centroids().count(), features_->count(), + (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon); + stamp.reset(); + + if (new_epsilon < epsilon_) { + break; + } + } + + // update centroids + update_centroids(cents, algorithm); + + // Purge the empty centroids + PurgeCentroids(cents, purge_empty_); + return 0; +} + +/*! Kmeans Cluster + */ +class OptKmeansCluster : public IndexCluster { + public: + //! Constructor + OptKmeansCluster(void) {} + + //! Destructor + ~OptKmeansCluster(void) override {} + + //! Initialize Cluster + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup Cluster + int cleanup(void) override; + + //! Reset Cluster + int reset(void) override; + + //! Update Cluster + int update(const ailego::Params ¶ms) override; + + //! Suggest dividing to K clusters + void suggest(uint32_t k) override; + + //! Mount features + int mount(IndexFeatures::Pointer feats) override; + + //! Cluster + int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + //! Classify + int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + //! Label + int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) override; + + protected: + //! Members + IndexCluster::Pointer algorithm_{}; +}; + +//! Cluster +int OptKmeansCluster::cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + return algorithm_->cluster(std::move(threads), cents); +} + +//! Classify +int OptKmeansCluster::classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + return algorithm_->classify(std::move(threads), cents); +} + +//! Label +int OptKmeansCluster::label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) { + return algorithm_->label(std::move(threads), cents, out); +} + +//! Update Cluster +int OptKmeansCluster::update(const ailego::Params ¶ms) { + return algorithm_->update(params); +} + +//! Reset Cluster +int OptKmeansCluster::reset(void) { + return algorithm_->reset(); +} + +//! Cleanup Cluster +int OptKmeansCluster::cleanup(void) { + return algorithm_->cleanup(); +} + +//! Suggest dividing to K clusters +void OptKmeansCluster::suggest(uint32_t k) { + algorithm_->suggest(k); +} + +int OptKmeansCluster::mount(IndexFeatures::Pointer feats) { + return algorithm_->mount(feats); +} + +int OptKmeansCluster::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + auto type_ = meta.data_type(); + + if (meta.metric_name() == "InnerProduct") { + switch (type_) { + case IndexMeta::DataType::DT_FP16: { + algorithm_.reset( + new (std::nothrow) + NumericalInnerProductKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_FP32: { + algorithm_.reset(new (std::nothrow) + NumericalInnerProductKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_FP64: { + algorithm_.reset(new (std::nothrow) + NumericalInnerProductKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT8: { + algorithm_.reset(new (std::nothrow) + NumericalInnerProductKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT16: { + algorithm_.reset(new (std::nothrow) + NumericalInnerProductKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT4: { + algorithm_.reset(new (std::nothrow) + NibbleInnerProductKmeansAlgorithm); + break; + } + default: { + LOG_ERROR("Unsupported feature types %d.", type_); + return IndexError_Mismatch; + } + } + } else { + switch (type_) { + case IndexMeta::DataType::DT_FP16: { + algorithm_.reset(new (std::nothrow) + NumericalKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_FP32: { + algorithm_.reset(new (std::nothrow) NumericalKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_FP64: { + algorithm_.reset(new (std::nothrow) NumericalKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT8: { + algorithm_.reset(new (std::nothrow) NumericalKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT16: { + algorithm_.reset(new (std::nothrow) NumericalKmeansAlgorithm); + break; + } + case IndexMeta::DataType::DT_INT4: { + algorithm_.reset(new (std::nothrow) NibbleKmeansAlgorithm); + break; + } + default: { + LOG_ERROR("Unsupported feature types %d.", type_); + return IndexError_Mismatch; + } + } + } + + algorithm_->init(meta, params); + + return 0; +} + +INDEX_FACTORY_REGISTER_CLUSTER(OptKmeansCluster); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/seeker.h b/src/core/algorithm/cluster/seeker.h new file mode 100644 index 0000000..6f14f57 --- /dev/null +++ b/src/core/algorithm/cluster/seeker.h @@ -0,0 +1,72 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +class Seeker { + public: + struct Document { + uint32_t index; + float score; + + //! Constructor + Document(void) : index(0), score(0.0f) {} + + //! Constructor + Document(uint32_t i, float v) : index(i), score(v) {} + + //! Constructor + Document(const Document &rhs) : index(rhs.index), score(rhs.score) {} + + //! Assignment + Document &operator=(const Document &rhs) { + index = rhs.index; + score = rhs.score; + return *this; + } + + //! Less than + bool operator<(const Document &rhs) const { + return (this->score < rhs.score); + } + + //! Greater than + bool operator>(const Document &rhs) const { + return (this->score > rhs.score); + } + }; + + public: + //! Destructor + virtual ~Seeker(void) {} + + virtual int init(const IndexMeta &meta) = 0; + + virtual int cleanup(void) = 0; + + virtual int reset(void) = 0; + + virtual int mount(IndexFeatures::Pointer feats) = 0; + + virtual int seek(const void *query, size_t len, Document *out) = 0; + + virtual IndexFeatures::Pointer original(void) const = 0; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/stratified_cluster.cc b/src/core/algorithm/cluster/stratified_cluster.cc new file mode 100644 index 0000000..066868e --- /dev/null +++ b/src/core/algorithm/cluster/stratified_cluster.cc @@ -0,0 +1,502 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include +#include +#include +#include "cluster_params.h" + +namespace zvec { +namespace core { + +/*! Stratified Cluster + */ +class StratifiedCluster : public IndexCluster { + public: + //! Constructor + StratifiedCluster(void) {} + + //! Destructor + ~StratifiedCluster(void) override {} + + //! Initialize Cluster + int init(const IndexMeta &meta, const ailego::Params ¶ms) override { + meta_ = meta; + this->update_params(params); + return 0; + } + + //! Cleanup Cluster + int cleanup(void) override { + features_.reset(); + return 0; + } + + //! Reset Cluster + int reset(void) override { + features_.reset(); + return 0; + } + + //! Update Cluster + int update(const ailego::Params ¶ms) override { + this->update_params(params); + return 0; + } + + //! Suggest dividing to K clusters + void suggest(uint32_t k) override { + cluster_count_ = k; + } + + //! Mount features + int mount(IndexFeatures::Pointer feats) override { + if (!feats) { + return IndexError_InvalidArgument; + } + if (!feats->is_matched(meta_)) { + return IndexError_Mismatch; + } + features_ = std::move(feats); + return 0; + } + + //! Cluster + int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + //! Classify + int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) override; + + //! Label + int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) override; + + protected: + //! Test if it is valid + bool is_valid(void) const { + if (!features_ || !features_->count()) { + return false; + } + return true; + } + + //! Update parameters + void update_params(const ailego::Params ¶ms); + + //! Check Centroids + bool check_centroids(const IndexCluster::CentroidList ¢s); + + //! Initialize Sub Clusters + int init_sub_clusters(IndexCluster::Pointer *first, + IndexCluster::Pointer *second); + + //! Initialize First Cluster + int init_first_cluster(IndexCluster::Pointer *first); + + //! Initialize Second Cluster + int init_second_cluster(IndexCluster::Pointer *second, + IndexFeatures::Pointer features); + + private: + //! Members + IndexMeta meta_{}; + IndexFeatures::Pointer features_{}; + uint32_t cluster_count_{0u}; + uint32_t thread_count_{0u}; + uint32_t first_cluster_count_{0u}; + uint32_t second_cluster_count_{0u}; + bool auto_tuning_{false}; + std::string first_cluster_class_{"OptKmeansCluster"}; + std::string second_cluster_class_{"OptKmeansCluster"}; + ailego::Params first_cluster_params_{}; + ailego::Params second_cluster_params_{}; + + // TODO: Maybe optimize later + uint32_t second_threads_count_{10u}; // todo +}; + +int StratifiedCluster::cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + IndexCluster::Pointer first_cluster; + int result = init_first_cluster(&first_cluster); + if (result != 0) { + LOG_ERROR("Failed to initialize the first cluster."); + return result; + } + + if (first_cluster_count_) { + first_cluster->suggest(first_cluster_count_); + } + + // The first clustering + LOG_DEBUG("Clustering with first cluster: %s.", first_cluster_class_.c_str()); + result = first_cluster->cluster(threads, cents); + if (result != 0) { + LOG_ERROR("Failed to cluster in first cluster: %s.", + first_cluster_class_.c_str()); + return result; + } + + result = first_cluster->classify(threads, cents); + if (result != 0) { + LOG_ERROR("Failed to classify in first cluster: %s.", + first_cluster_class_.c_str()); + return result; + } + + // Cleanup for saving memory + first_cluster.reset(); + + // Calculate the total cluster count + uint32_t total_cluster_count = cents.size() * second_cluster_count_; + if (cluster_count_) { + total_cluster_count = cluster_count_; + } + + // Use thread_threads cluster instead + uint32_t tail_threads = threads->count() % second_threads_count_; + std::vector> threads_cluster; + + // TODO: reusing threads pool? + // Incase the threads count less than second threads count + if (threads->count() / second_threads_count_ == 0) { + for (size_t threads_idx = 0; threads_idx < tail_threads; threads_idx++) { + std::shared_ptr curr_threads = + std::make_shared(1, false); + threads_cluster.push_back(curr_threads); + } + } else { + for (size_t threads_idx = 0; threads_idx < second_threads_count_; + threads_idx++) { + uint32_t curr_threads_count = threads->count() / second_threads_count_; + if (threads_idx >= second_threads_count_ - tail_threads) { + curr_threads_count++; + } + std::shared_ptr curr_threads = + std::make_shared(curr_threads_count, false); + threads_cluster.push_back(curr_threads); + } + } + + auto task_group = threads->make_group(); + // The second clustering + for (size_t i = 0; i < cents.size(); ++i) { + if (cents[i].similars().empty()) { + continue; + } + + IndexThreads::Pointer &curr_threads = + threads_cluster[i % (threads_cluster.size())]; + + task_group->submit(ailego::Closure::New( + [this, &curr_threads, &total_cluster_count, ¢s](size_t index) { + auto &it = cents[index]; + IndexCluster::Pointer second_cluster; + std::shared_ptr features = + std::make_shared( + meta_, it.similars().data(), it.similars().size()); + + int ret = this->init_second_cluster(&second_cluster, features); + if (ret != 0) { + LOG_ERROR("Failed to initialize the second cluster."); + return; + } + + if (auto_tuning_) { + if (total_cluster_count) { + double factor = static_cast(it.similars().size()) / + static_cast(this->features_->count()); + second_cluster->suggest( + std::max(static_cast( + std::floor(total_cluster_count * factor)), + 1u)); + } + } else if (second_cluster_count_) { + second_cluster->suggest(second_cluster_count_); + } + + LOG_DEBUG("Clustering with second cluster: %s.", + second_cluster_class_.c_str()); + ret = second_cluster->cluster(curr_threads, *(it.mutable_subitems())); + if (ret != 0) { + LOG_ERROR("Failed to cluster in second cluster: %s.", + second_cluster_class_.c_str()); + } + }, + i)); + } + task_group->wait_finish(); + return 0; +} + +int StratifiedCluster::classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) { + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (cents.empty()) { + LOG_ERROR("The input centroid's list is empty."); + return IndexError_InvalidArgument; + } + + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + IndexCluster::Pointer first_cluster, second_cluster; + int result = init_sub_clusters(&first_cluster, &second_cluster); + if (result != 0) { + LOG_ERROR("Failed to initialize the subclusters."); + return result; + } + + // The first classifying + result = first_cluster->classify(threads, cents); + if (result != 0) { + LOG_ERROR("Failed to classify in first cluster: %s.", + first_cluster_class_.c_str()); + return result; + } + + // Cleanup for saving memory + first_cluster.reset(); + + std::shared_ptr shell = + std::make_shared(meta_); + + // The second classifying + for (IndexCluster::Centroid &it : cents) { + const auto &feats = it.similars(); + + if (feats.empty()) { + continue; + } + + shell->mount(feats.data(), feats.size()); + result = second_cluster->mount(shell); + if (result != 0) { + LOG_ERROR("Failed to mount features for second cluster: %s.", + second_cluster_class_.c_str()); + return result; + } + + result = second_cluster->classify(threads, *it.mutable_subitems()); + if (result != 0) { + LOG_ERROR("Failed to classify in second cluster: %s.", + second_cluster_class_.c_str()); + return result; + } + } + return 0; +} + +int StratifiedCluster::label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out) { + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + if (cents.empty()) { + LOG_ERROR("The input centroid's list is empty."); + return IndexError_InvalidArgument; + } + + if (!this->check_centroids(cents)) { + LOG_ERROR("The input centroid's list includes some invalid centroids."); + return IndexError_InvalidArgument; + } + + if (!this->is_valid()) { + LOG_ERROR("The cluster is not ready."); + return IndexError_NoReady; + } + + IndexCluster::Pointer first_cluster; + int result = init_first_cluster(&first_cluster); + if (result != 0) { + LOG_ERROR("Failed to initialize the subclusters."); + return result; + } + + result = first_cluster->label(threads, cents, out); + if (result != 0) { + LOG_ERROR("Failed to label in first cluster: %s.", + first_cluster_class_.c_str()); + return result; + } + return 0; +} + +void StratifiedCluster::update_params(const ailego::Params ¶ms) { + params.get(GENERAL_THREAD_COUNT, &thread_count_); + params.get(GENERAL_CLUSTER_COUNT, &cluster_count_); + params.get(STRATIFIED_CLUSTER_COUNT, &cluster_count_); + params.get(STRATIFIED_CLUSTER_FIRST_COUNT, &first_cluster_count_); + params.get(STRATIFIED_CLUSTER_SECOND_COUNT, &second_cluster_count_); + params.get(STRATIFIED_CLUSTER_FIRST_CLASS, &first_cluster_class_); + params.get(STRATIFIED_CLUSTER_SECOND_CLASS, &second_cluster_class_); + params.get(STRATIFIED_CLUSTER_FIRST_PARAMS, &first_cluster_params_); + params.get(STRATIFIED_CLUSTER_SECOND_PARAMS, &second_cluster_params_); + params.get(STRATIFIED_CLUSTER_AUTO_TUNING, &auto_tuning_); + params.get(STRATIFIED_CLUSTER_SECOND_POOL_COUNT, &second_threads_count_); +} + +bool StratifiedCluster::check_centroids( + const IndexCluster::CentroidList ¢s) { + for (const auto &it : cents) { + if (it.size() != meta_.element_size()) { + return false; + } + } + return true; +} + +int StratifiedCluster::init_sub_clusters(IndexCluster::Pointer *first, + IndexCluster::Pointer *second) { + IndexCluster::Pointer first_cluster = + IndexFactory::CreateCluster(first_cluster_class_); + + if (!first_cluster) { + LOG_ERROR("Failed to create first cluster: %s.", + first_cluster_class_.c_str()); + return IndexError_NoExist; + } + + IndexCluster::Pointer second_cluster = + IndexFactory::CreateCluster(second_cluster_class_); + + if (!second_cluster) { + LOG_ERROR("Failed to create second cluster: %s.", + first_cluster_class_.c_str()); + return IndexError_NoExist; + } + + int result = first_cluster->init(meta_, first_cluster_params_); + if (result != 0) { + LOG_ERROR("Failed to initialize first cluster: %s.", + first_cluster_class_.c_str()); + return result; + } + + result = second_cluster->init(meta_, second_cluster_params_); + if (result != 0) { + LOG_ERROR("Failed to initialize second cluster: %s.", + second_cluster_class_.c_str()); + return result; + } + + result = first_cluster->mount(features_); + if (result != 0) { + LOG_ERROR("Failed to mount features for first cluster: %s.", + first_cluster_class_.c_str()); + return result; + } + + *first = std::move(first_cluster); + *second = std::move(second_cluster); + return 0; +} + +int StratifiedCluster::init_first_cluster(IndexCluster::Pointer *first) { + IndexCluster::Pointer first_cluster = + IndexFactory::CreateCluster(first_cluster_class_); + + if (!first_cluster) { + LOG_ERROR("Failed to create first cluster: %s.", + first_cluster_class_.c_str()); + return IndexError_NoExist; + } + + int result = first_cluster->init(meta_, first_cluster_params_); + if (result != 0) { + LOG_ERROR("Failed to initialize first cluster: %s.", + first_cluster_class_.c_str()); + return result; + } + + result = first_cluster->mount(features_); + if (result != 0) { + LOG_ERROR("Failed to mount features for first cluster: %s.", + first_cluster_class_.c_str()); + return result; + } + + *first = std::move(first_cluster); + return 0; +} + +int StratifiedCluster::init_second_cluster(IndexCluster::Pointer *second, + IndexFeatures::Pointer features) { + IndexCluster::Pointer second_cluster = + IndexFactory::CreateCluster(second_cluster_class_); + + if (!second_cluster) { + LOG_ERROR("Failed to create second cluster: %s.", + second_cluster_class_.c_str()); + return IndexError_NoExist; + } + + int result = second_cluster->init(meta_, second_cluster_params_); + if (result != 0) { + LOG_ERROR("Failed to initialize second cluster: %s.", + second_cluster_class_.c_str()); + return result; + } + + result = second_cluster->mount(features); + if (result != 0) { + LOG_ERROR("Failed to mount features for second cluster: %s.", + second_cluster_class_.c_str()); + return result; + } + + *second = std::move(second_cluster); + return 0; +} + +INDEX_FACTORY_REGISTER_CLUSTER(StratifiedCluster); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/stratified_cluster_trainer.cc b/src/core/algorithm/cluster/stratified_cluster_trainer.cc new file mode 100644 index 0000000..2b47075 --- /dev/null +++ b/src/core/algorithm/cluster/stratified_cluster_trainer.cc @@ -0,0 +1,341 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "stratified_cluster_trainer.h" +#include +#include +#include +#include +#include +#include "cluster_params.h" + +namespace zvec { +namespace core { + +const std::string StratifiedClusterTrainer::SEP_TOKEN = "*"; +const std::string StratifiedClusterTrainer::DEFAULT_CLUSTER_CLASS = + "OptKmeansCluster"; + +int StratifiedClusterTrainer::init_params(const ailego::Params ¶ms) { + params.get(STRATIFIED_TRAINER_SAMPLE_COUNT, &sample_count_); + params.get(STRATIFIED_TRAINER_SAMPLE_RATIO, &sample_ratio_); + params.get(STRATIFIED_TRAINER_THREAD_COUNT, &thread_count_); + cluster_auto_tuning_ = params.get_as_bool(STRATIFIED_TRAINER_AUTOAUNE); + + std::string centroids_num = + params.get_as_string(STRATIFIED_TRAINER_CLUSTER_COUNT); + if (!centroids_num.empty()) { + ailego::StringHelper::Split(centroids_num, SEP_TOKEN, ¢roid_num_vec_); + for (size_t i = 0; i < centroid_num_vec_.size(); ++i) { + if (centroid_num_vec_[i] == 0) { + LOG_ERROR("Invalid centroid num %s", centroids_num.c_str()); + return IndexError_InvalidArgument; + } + } + } else { + LOG_ERROR("No centroids_num configed."); + return IndexError_InvalidArgument; + } + + size_t level_cnt = centroid_num_vec_.size(); + for (size_t i = 1; i <= level_cnt; ++i) { + std::string level_params_key = + STRATIFIED_TRAINER_PARAMS_IN_LEVEL_PREFIX + std::to_string(i); + ailego::Params level_params; + params.get(level_params_key, &level_params); + cluster_params_.push_back(level_params); + } + + std::string cluster_class(DEFAULT_CLUSTER_CLASS); + params.get(STRATIFIED_TRAINER_CLASS_NAME, &cluster_class); + ailego::StringHelper::Split(cluster_class, SEP_TOKEN, &cluster_class_); + if (cluster_class_.size() == 1) { + // repeat the cluster class to level_cnt + for (size_t i = 1; i < level_cnt; ++i) { + cluster_class_.push_back(cluster_class_[0]); + } + } else if (cluster_class_.size() != level_cnt) { + LOG_ERROR("Cluster class should be equal to level count"); + return IndexError_InvalidArgument; + } + return 0; +} + +int StratifiedClusterTrainer::init(const IndexMeta &index_meta, + const ailego::Params ¶ms) { + int err = init_params(params); + if (err != 0) { + LOG_ERROR("init params failed, errno:%d,%s", err, IndexError::What(err)); + return err; + } + + meta_ = index_meta; + ailego::Params cluster_params; + if (centroid_num_vec_.size() == 0) { + LOG_ERROR("invalid centroid num"); + return IndexError_InvalidArgument; + } else if (centroid_num_vec_.size() == 1) { + // one level clustering + class_name_ = cluster_class_[0]; + cluster_params = cluster_params_[0]; + suggest_centriod_cnt_ = centroid_num_vec_[0]; + } else if (centroid_num_vec_.size() == 2) { + // cluster level > 1 + class_name_ = "StratifiedCluster"; + int level_cnt = centroid_num_vec_.size(); + cluster_params.set(STRATIFIED_CLUSTER_FIRST_CLASS, + cluster_class_[level_cnt - 2]); + cluster_params.set(STRATIFIED_CLUSTER_SECOND_CLASS, + cluster_class_[level_cnt - 1]); + cluster_params.set(STRATIFIED_CLUSTER_FIRST_COUNT, + centroid_num_vec_[level_cnt - 2]); + cluster_params.set(STRATIFIED_CLUSTER_SECOND_COUNT, + centroid_num_vec_[level_cnt - 1]); + cluster_params.set(STRATIFIED_CLUSTER_FIRST_PARAMS, + cluster_params_[level_cnt - 2]); + cluster_params.set(STRATIFIED_CLUSTER_SECOND_PARAMS, + cluster_params_[level_cnt - 1]); + cluster_params.set(STRATIFIED_CLUSTER_AUTO_TUNING, cluster_auto_tuning_); + suggest_centriod_cnt_ = + centroid_num_vec_[level_cnt - 1] * centroid_num_vec_[level_cnt - 2]; + } else { + LOG_ERROR("Unsupported more than 2 level clustering."); + return IndexError_Unsupported; + } + + cluster_ = IndexFactory::CreateCluster(class_name_); + if (!cluster_) { + LOG_ERROR("Failed to create cluster[%s]", class_name_.c_str()); + return IndexError_InvalidArgument; + } + int result = cluster_->init(meta_, cluster_params); + if (result != 0) { + LOG_ERROR("Failed to initialize of cluster[%s], error: %d, %s", + class_name_.c_str(), result, IndexError::What(result)); + return result; + } + if (suggest_centriod_cnt_ > 0) { + cluster_->suggest(suggest_centriod_cnt_); + } + + return 0; +} + +int StratifiedClusterTrainer::cleanup(void) { + cluster_ = nullptr; + centroids_.clear(); + return 0; +} + +int StratifiedClusterTrainer::train(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) { + ailego::ElapsedTime timer; + if (!holder) { + return IndexError_InvalidArgument; + } + if (!holder->is_matched(meta_)) { + return IndexError_Mismatch; + } + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + + size_t train_sample_count = std::max( + sample_count_, static_cast(sample_ratio_ * holder->count())); + + IndexFeatures::Pointer features; + if (train_sample_count > 0) { + LOG_INFO( + "Train sampling, SampleCount=%u, SampleRatio=%f, HolderCount=%lu, " + "TrainCount=%lu", + sample_count_, sample_ratio_, holder->count(), train_sample_count); + + auto sampler = std::make_shared>( + meta_, train_sample_count); + size_t pre_reserve = train_sample_count < holder->count() + ? train_sample_count + : holder->count(); + sampler->reserve(pre_reserve); + for (auto iter = holder->create_iterator(); iter && iter->is_valid(); + iter->next()) { + sampler->emplace(iter->data()); + } + features = sampler; + stats_.set_trained_count(train_sample_count); + } else { + LOG_INFO( + "Do no sampling, SampleCount=%u, SampleRatio=%f, " + "HolderCount=%lu, TrainCount=%lu", + sample_count_, sample_ratio_, holder->count(), holder->count()); + + auto no_sampler = std::make_shared(meta_); + for (auto iter = holder->create_iterator(); iter && iter->is_valid(); + iter->next()) { + no_sampler->emplace(iter->data()); + } + + features = no_sampler; + stats_.set_trained_count(holder->count()); + } + stats_.set_discarded_count(0); + + // Holder is not needed, cleanup it. + holder.reset(); + + int result = cluster_->mount(features); + if (result != 0) { + LOG_ERROR("Failed to mount features of cluster[%s], error: %d, %s", + class_name_.c_str(), result, IndexError::What(result)); + return result; + } + + centroids_.clear(); + result = cluster_->cluster(std::move(threads), centroids_); + if (result != 0) { + LOG_ERROR("Failed to cluster features of cluster[%s], error: %d, %s", + class_name_.c_str(), result, IndexError::What(result)); + return result; + } + + // check build result + std::vector level_size; + std::function + cal_centroid_cnt = + [&cal_centroid_cnt, &level_size]( + const IndexCluster::CentroidList ¢s, size_t level) { + if (level > level_size.size()) { + level_size.resize(level); + } + level_size[level - 1] += cents.size(); + for (const auto &it : cents) { + if (!it.subitems().empty()) { + cal_centroid_cnt(it.subitems(), level + 1); + } + } + }; + cal_centroid_cnt(centroids_, 1); + + size_t centroids_num = level_size[level_size.size() - 1]; + if (centroids_num > suggest_centriod_cnt_) { + LOG_WARN( + "Built centroid(%zd level) count[%zd] bigger than expected " + "count[%d]", + level_size.size(), centroids_num, suggest_centriod_cnt_); + } else { + LOG_INFO("Built centroid(%zd level) count[%zd], expected count[%d]", + level_size.size(), centroids_num, suggest_centriod_cnt_); + } + + stats_.set_trained_costtime(timer.milli_seconds()); + + return 0; +} + +int StratifiedClusterTrainer::load(IndexStorage::Pointer cntr) { + if (!cntr) { + LOG_ERROR("IndexStorage is nullptr."); + return IndexError_InvalidArgument; + } + std::shared_ptr bundle = + std::make_shared(); + if (!bundle) { + LOG_ERROR("New MemoryInndexBundle failed."); + return IndexError_NoMemory; + } + + auto results = cntr->get_all(); + for (auto &it : results) { + IndexStorage::Segment::Pointer &seg = it.second; + if (!seg) { + LOG_ERROR("Get Segment %s failed.", it.first.c_str()); + return IndexError_InvalidArgument; + } + size_t data_size = seg->data_size(); + const void *data = nullptr; + size_t actual_size = seg->read(0, &data, data_size); + if (actual_size != data_size) { + LOG_ERROR("Read data failed expect %zu, actual %zu.", data_size, + actual_size); + return IndexError_ReadData; + } + bundle->set(it.first, data, data_size); + } + + int result = IndexHelper::DeserializeFromStorage(cntr.get(), &meta_); + if (result != 0) { + LOG_ERROR("Failed to deserialize meta from container"); + return result; + } + + result = IndexCluster::Deserialize(meta_, std::move(bundle), ¢roids_); + if (result != 0) { + LOG_ERROR("Failed to deserialize index: %d", result); + return result; + } + return 0; +} + +int StratifiedClusterTrainer::dump(const IndexDumper::Pointer &dumper) { + IndexBundle::Pointer bundle; + int result = IndexCluster::Serialize(meta_, centroids_, &bundle); + if (result != 0) { + LOG_ERROR("IndexCluster Serialize failed with ret %d.", result); + return result; + } + + result = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (result != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return result; + } + + for (const auto &it : bundle->all()) { + size_t data_size = it.second.size(); + result = dumper->append(it.first, data_size, 0, 0); + if (result != 0) { + LOG_ERROR("Dumper append meta %s %zu failed.", it.first.c_str(), + data_size); + return IndexError_PackIndex; + } + size_t actual_size = dumper->write(it.second.buffer(), data_size); + if (actual_size != data_size) { + LOG_ERROR("Dumper segment %s expect %zu, actual %zu.", it.first.c_str(), + data_size, actual_size); + return IndexError_PackIndex; + } + } + return 0; +} + +const IndexMeta &StratifiedClusterTrainer::meta(void) const { + return meta_; +} + +const IndexTrainer::Stats &StratifiedClusterTrainer::stats(void) const { + return stats_; +} + +IndexBundle::Pointer StratifiedClusterTrainer::indexes(void) const { + IndexBundle::Pointer bundle; + IndexCluster::Serialize(meta_, centroids_, &bundle); + return bundle; +} + +//! Register Cluster Trainer in Factory +INDEX_FACTORY_REGISTER_TRAINER(StratifiedClusterTrainer); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/stratified_cluster_trainer.h b/src/core/algorithm/cluster/stratified_cluster_trainer.h new file mode 100644 index 0000000..0a82c1e --- /dev/null +++ b/src/core/algorithm/cluster/stratified_cluster_trainer.h @@ -0,0 +1,84 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include + +namespace zvec { +namespace core { + +/*! Cluster Trainer + */ +class StratifiedClusterTrainer : public IndexTrainer { + public: + typedef std::shared_ptr Pointer; + + //! Constructor + StratifiedClusterTrainer(void) {} + + //! Destructor + ~StratifiedClusterTrainer(void) {} + + protected: + //! Initialize Trainer + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms); + + //! Cleanup Trainer + virtual int cleanup(void); + + //! Train the data + virtual int train(IndexThreads::Pointer threads, IndexHolder::Pointer holder); + + //! Load index from file path or dir + virtual int load(IndexStorage::Pointer cntr); + + //! Dump index into file path or dir + virtual int dump(const IndexDumper::Pointer &dumper); + + //! Retrieve Index Meta + virtual const IndexMeta &meta(void) const; + + //! Retrieve statistics + virtual const IndexTrainer::Stats &stats(void) const; + + //! Retrieve the output indexes + virtual IndexBundle::Pointer indexes(void) const; + + private: + int init_params(const ailego::Params ¶ms); + + private: + IndexMeta meta_{}; + uint32_t sample_count_{0u}; + float sample_ratio_{0.0}; + uint32_t thread_count_{0u}; + bool cluster_auto_tuning_{false}; + IndexCluster::Pointer cluster_{}; + IndexCluster::CentroidList centroids_{}; + + uint32_t suggest_centriod_cnt_{0u}; + std::string class_name_; + std::vector cluster_class_; + std::vector centroid_num_vec_; + std::vector cluster_params_; + IndexTrainer::Stats stats_{}; + + private: + static const std::string SEP_TOKEN; + static const std::string DEFAULT_CLUSTER_CLASS; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/cluster/vector_mean.h b/src/core/algorithm/cluster/vector_mean.h new file mode 100644 index 0000000..43c9153 --- /dev/null +++ b/src/core/algorithm/cluster/vector_mean.h @@ -0,0 +1,791 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +/*! Vector Mean + */ +struct VectorMean { + //! Destructor + virtual ~VectorMean(void) {} + + //! Reset accumulator + virtual void reset(void) = 0; + + //! Plus a vector + virtual bool plus(const void *vec, size_t len) = 0; + + //! Retrieve the mean of vectors + virtual bool mean(void *out, size_t len) const = 0; + + //! Retrieve the mean of vectors + virtual void mean(std::string *out) const = 0; + + //! Merge another vector mean + virtual bool merge(const VectorMean &rhs) = 0; + + //! Retrieve the count of vectors + virtual size_t count(void) const = 0; + + //! Retrieve the dimension of vectors + virtual size_t dimension(void) const = 0; +}; + +/*! Vector Mean Array + */ +struct VectorMeanArray { + //! Destructor + virtual ~VectorMeanArray(void) {} + + //! Operator [] + VectorMean &operator[](size_t i) { + return this->at(i); + } + + //! Operator [] + const VectorMean &operator[](size_t i) const { + return this->at(i); + } + + //! Resize accumulators + virtual void resize(size_t cnt) = 0; + + //! Clear accumulators + virtual void clear(void) = 0; + + //! Retrieve an accumulator + virtual VectorMean &at(size_t i) = 0; + + //! Retrieve an accumulator + virtual const VectorMean &at(size_t i) const = 0; + + //! Retrieve the count of accumulators + virtual size_t count(void) const = 0; + + //! Retrieve the dimension of accumulators + virtual size_t dimension(void) const = 0; +}; + +/*! General Vector Mean Array + */ +template ::type> +class GeneralVectorMeanArray : public VectorMeanArray { + public: + //! Constructor + GeneralVectorMeanArray(size_t dim) : dimension_(dim), array_() {} + + //! Constructor + GeneralVectorMeanArray(const GeneralVectorMeanArray &rhs) + : dimension_(rhs.dimension_), array_(rhs.array_) {} + + //! Constructor + GeneralVectorMeanArray(GeneralVectorMeanArray &&rhs) + : dimension_(rhs.dimension_), array_(std::move(rhs.array_)) {} + + //! Emplace an accumulator + template + bool emplace(TArgs &&...args) { + T accum(std::forward(args)...); + if (accum.dimension() != dimension_) { + return false; + } + array_.push_back(std::move(accum)); + return true; + } + + //! Resize accumulators + virtual void resize(size_t cnt) { + if (array_.size() < cnt) { + for (size_t i = array_.size(); i < cnt; ++i) { + array_.emplace_back(dimension_); + } + } else { + array_.resize(cnt); + } + } + + //! Clear accumulators + virtual void clear(void) { + array_.clear(); + } + + //! Retrieve an accumulator + virtual VectorMean &at(size_t i) { + return array_[i]; + } + + //! Retrieve an accumulator + virtual const VectorMean &at(size_t i) const { + return array_[i]; + } + + //! Retrieve the count of accumulators + virtual size_t count(void) const { + return array_.size(); + } + + //! Retrieve the dimension of accumulators + virtual size_t dimension(void) const { + return dimension_; + } + + private: + //! Disable them + GeneralVectorMeanArray(void) = delete; + + //! Members + size_t dimension_; + std::vector array_; +}; + +/*! Numerical Vector Mean + */ +template ::value>::type> +class NumericalVectorMean : public VectorMean { + public: + //! Constructor + NumericalVectorMean(void) : count_(0), accums_() {} + + //! Constructor + NumericalVectorMean(const NumericalVectorMean &rhs) + : count_(rhs.count_), accums_(rhs.accums_) {} + + //! Constructor + NumericalVectorMean(NumericalVectorMean &&rhs) + : count_(rhs.count_), accums_(std::move(rhs.accums_)) {} + + //! Constructor + NumericalVectorMean(size_t dim) : count_(0), accums_(dim) {} + + //! Constructor + NumericalVectorMean(const T *means, size_t dim, size_t cnt) + : count_(cnt), accums_(dim) { + for (size_t i = 0; i < dim; ++i) { + accums_[i] = static_cast(means[i]) * count_; + } + } + + //! Reset accumulator + void reset(size_t dim) { + count_ = 0u; + accums_.clear(); + accums_.resize(dim, 0.0); + } + + //! Reset accumulator + virtual void reset(void) { + this->reset(accums_.size()); + } + + //! Plus a vector + virtual bool plus(const void *vec, size_t len) { + size_t dim = accums_.size(); + if (dim * sizeof(T) != len) { + return false; + } + for (size_t i = 0; i < dim; ++i) { + accums_[i] += *(static_cast(vec) + i); + } + ++count_; + return true; + } + + //! Retrieve the mean of vectors + virtual bool mean(void *out, size_t len) const { + size_t dim = accums_.size(); + if (dim * sizeof(T) != len) { + return false; + } + for (size_t i = 0; i < dim; ++i) { + *(static_cast(out) + i) = FloatCast(accums_[i] / count_); + } + return true; + } + + //! Retrieve the mean of vectors + virtual void mean(std::string *out) const { + ailego::NumericalVector &vec = + *static_cast *>(out); + + size_t dim = accums_.size(); + vec.resize(dim); + for (size_t i = 0; i < dim; ++i) { + vec[i] = FloatCast(accums_[i] / count_); + } + } + + //! Merge another vector mean + virtual bool merge(const VectorMean &rhs) { + const NumericalVectorMean &src = + dynamic_cast &>(rhs); + + size_t dim = accums_.size(); + if (dim != src.accums_.size()) { + return false; + } + count_ += src.count_; + for (size_t i = 0; i < dim; ++i) { + accums_[i] += src.accums_[i]; + } + return true; + } + + //! Retrieve the count of vectors + virtual size_t count(void) const { + return count_; + } + + //! Retrieve dimension of accumulator + virtual size_t dimension(void) const { + return accums_.size(); + } + + protected: + //! Convert float type to another type + template + static auto FloatCast(const double &val) -> + typename std::enable_if::value, U>::type { + return static_cast(val); + } + + //! Convert float type to another type + template + static auto FloatCast(const double &val) -> + typename std::enable_if::value, U>::type { + return static_cast(std::round(val)); + } + + private: + //! Members + size_t count_; + std::vector accums_; +}; + +/*! Numerical Vector Harmonic Mean + */ +template ::value>::type> +class NumericalVectorHarmonicMean : public VectorMean { + public: + //! Constructor + NumericalVectorHarmonicMean(void) : count_(0), accums_() {} + + //! Constructor + NumericalVectorHarmonicMean(const NumericalVectorHarmonicMean &rhs) + : count_(rhs.count_), accums_(rhs.accums_) {} + + //! Constructor + NumericalVectorHarmonicMean(NumericalVectorHarmonicMean &&rhs) + : count_(rhs.count_), accums_(std::move(rhs.accums_)) {} + + //! Constructor + NumericalVectorHarmonicMean(size_t dim) : count_(0), accums_(dim) {} + + //! Constructor + NumericalVectorHarmonicMean(const T *means, size_t dim, size_t cnt) + : count_(cnt), accums_(dim) { + for (size_t i = 0; i < dim; ++i) { + accums_[i] = static_cast(count_) / static_cast(means[i]); + } + } + + //! Reset accumulator + void reset(size_t dim) { + count_ = 0u; + accums_.clear(); + accums_.resize(dim, 0.0); + } + + //! Reset accumulator + virtual void reset(void) { + this->reset(accums_.size()); + } + + //! Plus a vector (harmonic) + virtual bool plus(const void *vec, size_t len) { + size_t dim = accums_.size(); + if (dim * sizeof(T) != len) { + return false; + } + for (size_t i = 0; i < dim; ++i) { + accums_[i] += 1.0 / *(static_cast(vec) + i); + } + ++count_; + return true; + } + + //! Retrieve the mean of vectors (harmonic) + virtual bool mean(void *out, size_t len) const { + size_t dim = accums_.size(); + if (dim * sizeof(T) != len) { + return false; + } + for (size_t i = 0; i < dim; ++i) { + *(static_cast(out) + i) = FloatCast(count_ / accums_[i]); + } + return true; + } + + //! Retrieve the mean of vectors + virtual void mean(std::string *out) const { + ailego::NumericalVector &vec = + *static_cast *>(out); + + size_t dim = accums_.size(); + vec.resize(dim); + for (size_t i = 0; i < dim; ++i) { + vec[i] = FloatCast(count_ / accums_[i]); + } + } + + //! Merge another vector mean + virtual bool merge(const VectorMean &rhs) { + const NumericalVectorHarmonicMean &src = + dynamic_cast &>(rhs); + + size_t dim = accums_.size(); + if (dim != src.accums_.size()) { + return false; + } + count_ += src.count_; + for (size_t i = 0; i < dim; ++i) { + accums_[i] += src.accums_[i]; + } + return true; + } + + //! Retrieve the count of vectors + virtual size_t count(void) const { + return count_; + } + + //! Retrieve dimension of accumulator + virtual size_t dimension(void) const { + return accums_.size(); + } + + protected: + //! Convert float type to another type + template + static auto FloatCast(const double &val) -> + typename std::enable_if::value, U>::type { + return static_cast(val); + } + + //! Convert float type to another type + template + static auto FloatCast(const double &val) -> + typename std::enable_if::value, U>::type { + return static_cast(std::round(val)); + } + + private: + //! Members + size_t count_; + std::vector accums_; +}; + +/*! Numerical Vector Geometric Mean + */ +template ::value>::type> +class NumericalVectorGeometricMean : public VectorMean { + public: + //! Constructor + NumericalVectorGeometricMean(void) : count_(0), accums_() {} + + //! Constructor + NumericalVectorGeometricMean(const NumericalVectorGeometricMean &rhs) + : count_(rhs.count_), accums_(rhs.accums_) {} + + //! Constructor + NumericalVectorGeometricMean(NumericalVectorGeometricMean &&rhs) + : count_(rhs.count_), accums_(std::move(rhs.accums_)) {} + + //! Constructor + NumericalVectorGeometricMean(size_t dim) : count_(0), accums_(dim, 1.0) {} + + //! Constructor + NumericalVectorGeometricMean(const T *means, size_t dim, size_t cnt) + : count_(cnt), accums_(dim, 1.0) { + for (size_t i = 0; i < dim; ++i) { + accums_[i] = std::pow(static_cast(means[i]), count_); + } + } + + //! Reset accumulator + void reset(size_t dim) { + count_ = 0u; + accums_.clear(); + accums_.resize(dim, 1.0); + } + + //! Reset accumulator + virtual void reset(void) { + this->reset(accums_.size()); + } + + //! Plus a vector (geometric) + virtual bool plus(const void *vec, size_t len) { + size_t dim = accums_.size(); + if (dim * sizeof(T) != len) { + return false; + } + for (size_t i = 0; i < dim; ++i) { + accums_[i] *= *(static_cast(vec) + i); + } + ++count_; + return true; + } + + //! Retrieve the mean of vectors (geometric) + virtual bool mean(void *out, size_t len) const { + size_t dim = accums_.size(); + if (dim * sizeof(T) != len) { + return false; + } + for (size_t i = 0; i < dim; ++i) { + *(static_cast(out) + i) = + FloatCast(std::pow(accums_[i], 1.0 / count_)); + } + return true; + } + + //! Retrieve the mean of vectors + virtual void mean(std::string *out) const { + ailego::NumericalVector &vec = + *static_cast *>(out); + + size_t dim = accums_.size(); + vec.resize(dim); + for (size_t i = 0; i < dim; ++i) { + vec[i] = FloatCast(std::pow(accums_[i], 1.0 / count_)); + } + } + + //! Merge another vector mean + virtual bool merge(const VectorMean &rhs) { + const NumericalVectorGeometricMean &src = + dynamic_cast &>(rhs); + + size_t dim = accums_.size(); + if (dim != src.accums_.size()) { + return false; + } + count_ += src.count_; + for (size_t i = 0; i < dim; ++i) { + accums_[i] *= src.accums_[i]; + } + return true; + } + + //! Retrieve the count of vectors + virtual size_t count(void) const { + return count_; + } + + //! Retrieve dimension of accumulator + virtual size_t dimension(void) const { + return accums_.size(); + } + + protected: + //! Convert float type to another type + template + static auto FloatCast(const double &val) -> + typename std::enable_if::value, U>::type { + return static_cast(val); + } + + //! Convert float type to another type + template + static auto FloatCast(const double &val) -> + typename std::enable_if::value, U>::type { + return static_cast(std::round(val)); + } + + private: + //! Members + size_t count_; + std::vector accums_; +}; + +/*! Binary Vector Mean + */ +class BinaryVectorMean : public VectorMean { + public: + //! Constructor + BinaryVectorMean(void) : count_(0), accums_() {} + + //! Constructor + BinaryVectorMean(const BinaryVectorMean &rhs) + : count_(rhs.count_), accums_(rhs.accums_) {} + + //! Constructor + BinaryVectorMean(BinaryVectorMean &&rhs) + : count_(rhs.count_), accums_(std::move(rhs.accums_)) {} + + //! Constructor + BinaryVectorMean(size_t dim) : count_(0), accums_(((dim + 7) >> 3) << 3) {} + + //! Constructor + BinaryVectorMean(const void *means, size_t dim, size_t cnt) + : count_(cnt), accums_(((dim + 7) >> 3) << 3) { + const uint8_t *bits = reinterpret_cast(means); + for (size_t i = 0; i < dim; ++i) { + accums_[i] = (count_ >> 1); + + if (bits[i >> 3] & static_cast(1 << (i & 0x7))) { + accums_[i] += 1; + } + } + } + + //! Reset accumulator + void reset(size_t dim) { + count_ = 0u; + accums_.clear(); + accums_.resize(dim); + } + + //! Reset accumulator + virtual void reset(void) { + this->reset(accums_.size()); + } + + //! Plus a vector + virtual bool plus(const void *vec, size_t len) { + size_t dim = accums_.size(); + if (dim != (len << 3)) { + return false; + } + + const uint8_t *bits = reinterpret_cast(vec); + for (size_t i = 0; i < dim; ++i) { + if (bits[i >> 3] & static_cast(1 << (i & 0x7))) { + accums_[i] += 1; + } + } + ++count_; + return true; + } + + //! Retrieve the mean of vectors + virtual bool mean(void *out, size_t len) const { + size_t dim = accums_.size(); + if (dim != (len << 3)) { + return false; + } + memset(out, 0, len); + + uint8_t *bits = reinterpret_cast(out); + size_t half_count = count_ >> 1; + for (size_t i = 0; i < dim; ++i) { + if (accums_[i] > half_count) { + bits[i >> 3] |= static_cast(1 << (i & 0x7)); + } + } + return true; + } + + //! Retrieve the mean of vectors + virtual void mean(std::string *out) const { + size_t dim = accums_.size(); + out->clear(); + out->resize((dim + 7) / 8); + + uint8_t *bits = + reinterpret_cast(const_cast(out->data())); + size_t half_count = count_ >> 1; + for (size_t i = 0; i < dim; ++i) { + if (accums_[i] > half_count) { + bits[i >> 3] |= static_cast(1 << (i & 0x7)); + } + } + } + + //! Merge another vector mean + virtual bool merge(const VectorMean &rhs) { + const BinaryVectorMean &src = dynamic_cast(rhs); + + size_t dim = accums_.size(); + if (dim != src.accums_.size()) { + return false; + } + count_ += src.count_; + for (size_t i = 0; i < dim; ++i) { + accums_[i] += src.accums_[i]; + } + return true; + } + + //! Retrieve the count of vectors + virtual size_t count(void) const { + return count_; + } + + //! Retrieve dimension of accumulator + virtual size_t dimension(void) const { + return accums_.size(); + } + + private: + //! Members + size_t count_; + std::vector accums_; +}; + +/*! Numerical Vector Mean + */ +template ::value>::type> +class NibbleVectorMean : public VectorMean { + public: + //! Constructor + NibbleVectorMean(void) : count_(0), accums_() {} + + //! Constructor + NibbleVectorMean(const NibbleVectorMean &rhs) + : count_(rhs.count_), accums_(rhs.accums_) {} + + //! Constructor + NibbleVectorMean(NibbleVectorMean &&rhs) + : count_(rhs.count_), accums_(std::move(rhs.accums_)) {} + + //! Constructor + NibbleVectorMean(size_t dim) : count_(0), accums_(dim) {} + + //! Constructor + NibbleVectorMean(const void *means, size_t dim, size_t cnt) + : count_(cnt), accums_(dim) { + const uint8_t *arr = reinterpret_cast(means); + for (size_t i = 0; i != dim; i += 2) { + uint8_t val = arr[i >> 1]; + int lo = ((int8_t)(val << 4) >> 4); + int hi = ((int8_t)(val) >> 4); + accums_[i] = static_cast(lo) * count_; + accums_[i + 1] = static_cast(hi) * count_; + } + } + + //! Reset accumulator + void reset(size_t dim) { + count_ = 0u; + accums_.clear(); + accums_.resize(dim, 0.0); + } + + //! Reset accumulator + virtual void reset(void) { + this->reset(accums_.size()); + } + + //! Plus a vector + virtual bool plus(const void *vec, size_t len) { + size_t dim = accums_.size(); + if (dim != (len << 1)) { + return false; + } + + const uint8_t *arr = reinterpret_cast(vec); + for (size_t i = 0; i != dim; i += 2) { + uint8_t val = arr[i >> 1]; + accums_[i] += ((int8_t)(val << 4) >> 4); + accums_[i + 1] += ((int8_t)(val) >> 4); + } + ++count_; + return true; + } + + //! Retrieve the mean of vectors + virtual bool mean(void *out, size_t len) const { + size_t dim = accums_.size(); + if (dim != (len << 1)) { + return false; + } + memset(out, 0, len); + + uint8_t *arr = reinterpret_cast(out); + + for (size_t i = 0; i != dim; i += 2) { + int lo = static_cast(std::round(accums_[i] / count_)); + int hi = static_cast(std::round(accums_[i + 1] / count_)); + arr[i >> 1] = (uint8_t)((hi << 4) & 0xf0) | (uint8_t)(lo & 0xf); + } + + return true; + } + + //! Retrieve the mean of vectors + virtual void mean(std::string *out) const { + size_t dim = accums_.size(); + out->clear(); + out->resize(dim >> 1); + + uint8_t *arr = reinterpret_cast(const_cast(out->data())); + + for (size_t i = 0; i != dim; i += 2) { + int lo = static_cast(std::round(accums_[i] / count_)); + int hi = static_cast(std::round(accums_[i + 1] / count_)); + arr[i >> 1] = (uint8_t)((hi << 4) & 0xf0) | (uint8_t)(lo & 0xf); + } + } + + //! Merge another vector mean + virtual bool merge(const VectorMean &rhs) { + const NibbleVectorMean &src = dynamic_cast(rhs); + + size_t dim = accums_.size(); + if (dim != src.accums_.size()) { + return false; + } + count_ += src.count_; + for (size_t i = 0; i < dim; ++i) { + accums_[i] += src.accums_[i]; + } + return true; + } + + //! Retrieve the count of vectors + virtual size_t count(void) const { + return count_; + } + + //! Retrieve dimension of accumulator + virtual size_t dimension(void) const { + return accums_.size(); + } + + private: + //! Members + size_t count_; + std::vector accums_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/CMakeLists.txt b/src/core/algorithm/diskann/CMakeLists.txt new file mode 100644 index 0000000..1b22524 --- /dev/null +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -0,0 +1,57 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) + +file(GLOB_RECURSE ALL_SRCS *.cc *.c) + +# The DiskAnn plugin is loaded at runtime via zvec::LoadDiskAnnPlugin() with +# RTLD_GLOBAL, so its undefined references to internal zvec symbols +# (core_framework, core_knn_cluster, zvec_ailego, ...) are resolved at load +# time against the hosting binary (_zvec.so for the Python extension, the +# test executable for gtest, or libzvec_core for tools). +# +# As a consequence the plugin .so must NOT carry NEEDED entries for those +# libs: otherwise the dynamic loader would try to resolve them from disk +# before invoking dlopen's RTLD_GLOBAL symbol sharing, and the Python wheel +# (which only ships _zvec.so + the plugin) would fail to load. +# +# We therefore link the plugin only against its system-level dependency +# (libaio) and rely on the global include dirs plus -Wl,--unresolved-symbols +# =ignore-all to let the linker build a shared library with unresolved +# references to internal APIs. +set(CORE_KNN_DISKANN_LIBS "") + +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386") + list(APPEND CORE_KNN_DISKANN_LIBS aio) +endif() + +cc_library( + NAME core_knn_diskann + SHARED STRICT + SRCS *.cc + LIBS ${CORE_KNN_DISKANN_LIBS} + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + VERSION "${PROXIMA_ZVEC_VERSION}" +) + +# Internal zvec libs are referenced only for header availability; the actual +# symbol resolution happens at dlopen time. Expose their public include dirs +# without adding them as link dependencies. +foreach(_dep zvec_ailego core_framework core_knn_cluster) + if(TARGET ${_dep}) + target_include_directories(core_knn_diskann PRIVATE + $) + add_dependencies(core_knn_diskann ${_dep}) + endif() +endforeach() + +# Allow the plugin to have unresolved symbols that will be satisfied at +# dlopen(RTLD_NOW | RTLD_GLOBAL) time by the hosting binary. +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + target_link_options(core_knn_diskann PRIVATE + "LINKER:--unresolved-symbols=ignore-all") +endif() + +# Rename the artifact to libzvec_diskann_plugin.so so it is clearly identified +# as an optional plugin that users load at runtime via zvec::LoadDiskAnnPlugin(). +set_target_properties(core_knn_diskann PROPERTIES + OUTPUT_NAME zvec_diskann_plugin +) \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_algorithm.cc b/src/core/algorithm/diskann/diskann_algorithm.cc new file mode 100644 index 0000000..31ed820 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_algorithm.cc @@ -0,0 +1,324 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_algorithm.h" +#include +#include +#include +#include "diskann_util.h" + +namespace zvec { +namespace core { + +DiskAnnAlgorithm::DiskAnnAlgorithm(DiskAnnEntity &entity, uint32_t max_degree) + : entity_(entity), max_degree_(max_degree), lock_pool_(kLockCnt) {} + +std::vector DiskAnnAlgorithm::get_init_ids(DiskAnnContext *ctx) { + const auto &entity = ctx->get_entity(); + + std::vector init_ids; + + init_ids.emplace_back(entity.medoid()); + + return init_ids; +} + +int DiskAnnAlgorithm::add_node(diskann_id_t id, DiskAnnContext *ctx) { + const void *vec = entity_.get_vector(id); + + ctx->reset_query(vec); + + std::vector pruned_list; + + int ret = search_neighbor_and_prune(id, pruned_list, ctx); + if (ret != 0) { + return ret; + } + + { + auto lock = lock_for(id); + entity_.set_neighbors(id, pruned_list); + } + + return inter_insert(id, pruned_list, ctx); +} + +int DiskAnnAlgorithm::prune_node(diskann_id_t id, DiskAnnContext *ctx) { + DistCalculator &dc = ctx->dist_calculator(); + + auto neighbors = entity_.get_neighbors(id); + + if (neighbors.first > max_degree_) { + std::set dummy_visited; + std::vector dummy_pool(0); + std::vector new_out_neighbors; + + for (size_t i = 0; i < neighbors.first; ++i) { + diskann_id_t node_id = (neighbors.second)[i]; + + auto itr = dummy_visited.find(node_id); + if (itr == dummy_visited.end() && node_id != id) { + float dist = dc.dist(id, node_id); + + dummy_pool.emplace_back(Neighbor(node_id, dist)); + dummy_visited.insert(node_id); + } + } + + prune_neighbors(id, dummy_pool, new_out_neighbors, ctx); + + { + auto lock = lock_for(id); + entity_.set_neighbors(id, new_out_neighbors); + } + } + + return 0; +} + +int DiskAnnAlgorithm::inter_insert(diskann_id_t id, + std::vector &pruned_list, + DiskAnnContext *ctx) { + DistCalculator &dc = ctx->dist_calculator(); + + for (auto &des : pruned_list) { + std::vector new_neighbors; + bool need_prune = false; + + { + auto lock = lock_for(des); + + auto neighbors = entity_.get_neighbors(des); + + bool found = false; + for (size_t i = 0; i < neighbors.first; ++i) { + if ((neighbors.second)[i] == id) { + found = true; + break; + } + } + + if (!found) { + if (neighbors.first < + static_cast(DiskAnnEntity::kDefaultGraphSlackFactor * + max_degree_)) { + entity_.add_neighbor(des, id); + need_prune = false; + } else { + new_neighbors.resize(neighbors.first + 1); + memcpy(&new_neighbors[0], neighbors.second, + sizeof(diskann_id_t) * neighbors.first); + + new_neighbors[neighbors.first] = id; + + need_prune = true; + } + } + } + + if (need_prune) { + std::set new_visited; + std::vector new_pool(0); + + size_t reserve_size = static_cast(std::ceil( + 1.05 * DiskAnnEntity::kDefaultGraphSlackFactor * max_degree_)); + + new_pool.reserve(reserve_size); + + for (auto node_id : new_neighbors) { + if (new_visited.find(node_id) == new_visited.end() && node_id != des) { + float dist = dc.dist(des, node_id); + new_pool.emplace_back(Neighbor(node_id, dist)); + new_visited.insert(node_id); + } + } + + std::vector new_pruned_neighbors; + prune_neighbors(des, new_pool, new_pruned_neighbors, ctx); + + { + auto lock = lock_for(des); + entity_.set_neighbors(des, new_pruned_neighbors); + } + } + } + + return 0; +} + +int DiskAnnAlgorithm::iterate_to_fixed_point( + const std::vector &init_ids, DiskAnnContext *ctx) { + DistCalculator &dc = ctx->dist_calculator(); + std::vector &expanded_nodes = ctx->expanded_nodes(); + NeighborPriorityQueue &best_list_nodes = ctx->best_list_nodes(); + VisitFilter &visit = ctx->visit_filter(); + + best_list_nodes.reserve(ctx->list_size()); + + for (auto id : init_ids) { + const void *vec = entity_.get_vector(id); + + float distance = dc.dist(vec); + + Neighbor nn = Neighbor(id, distance); + best_list_nodes.insert(nn); + } + + while (best_list_nodes.has_unexpanded_node()) { + auto neighbor = best_list_nodes.closest_unexpanded(); + auto node_id = neighbor.id; + + expanded_nodes.emplace_back(neighbor); + + std::vector id_scratch; + { + auto lock = lock_for(node_id); + auto neighbors = entity_.get_neighbors(node_id); + + for (size_t i = 0; i < neighbors.first; ++i) { + diskann_id_t neighbor_id = (neighbors.second)[i]; + + if (!visit.visited(neighbor_id)) { + id_scratch.push_back(neighbor_id); + + visit.set_visited(neighbor_id); + } + } + } + + for (size_t i = 0; i < id_scratch.size(); ++i) { + diskann_id_t id = id_scratch[i]; + + const void *vec = entity_.get_vector(id); + float dist = dc.dist(vec); + + best_list_nodes.insert(Neighbor(id, dist)); + } + } + + return 0; +} + +int DiskAnnAlgorithm::occlude_list(diskann_id_t id, std::vector &pool, + std::vector &result, + DiskAnnContext *ctx) { + if (pool.size() == 0) return 0; + + DistCalculator &dc = ctx->dist_calculator(); + + ailego_assert(std::is_sorted(pool.begin(), pool.end())); + ailego_assert(result.size() == 0); + + if (pool.size() > max_candidate_size_) { + pool.resize(max_candidate_size_); + } + + std::vector &occlude_factor = ctx->occlude_factor(); + + occlude_factor.clear(); + occlude_factor.insert(occlude_factor.end(), pool.size(), 0.0f); + + float cur_alpha = 1; + while (cur_alpha <= alpha_ && result.size() < max_degree_) { + for (auto iter = pool.begin(); + result.size() < max_degree_ && iter != pool.end(); ++iter) { + if (occlude_factor[iter - pool.begin()] > cur_alpha) { + continue; + } + + occlude_factor[iter - pool.begin()] = std::numeric_limits::max(); + + if (iter->id != id) { + result.push_back(iter->id); + } + + for (auto iter2 = iter + 1; iter2 != pool.end(); iter2++) { + auto t = iter2 - pool.begin(); + if (occlude_factor[t] > alpha_) { + continue; + } + + float djk = dc.dist(iter2->id, iter->id); + + occlude_factor[t] = + (djk == 0) ? std::numeric_limits::max() + : std::max(occlude_factor[t], iter2->distance / djk); + } + } + cur_alpha *= 1.2f; + } + + return 0; +} + +int DiskAnnAlgorithm::prune_neighbors(diskann_id_t id, + std::vector &pool, + std::vector &pruned_list, + DiskAnnContext *ctx) { + if (pool.size() == 0) { + pruned_list.clear(); + return 0; + } + + std::sort(pool.begin(), pool.end()); + + pruned_list.clear(); + pruned_list.reserve(max_degree_); + + occlude_list(id, pool, pruned_list, ctx); + + ailego_assert(pruned_list.size() <= max_degree_); + + if (saturate_graph_ && alpha_ > 1) { + for (const auto &node : pool) { + if (pruned_list.size() >= max_degree_) { + break; + } + + if ((std::find(pruned_list.begin(), pruned_list.end(), node.id) == + pruned_list.end()) && + node.id != id) { + pruned_list.push_back(node.id); + } + } + } + + return 0; +} + +int DiskAnnAlgorithm::search_neighbor_and_prune( + diskann_id_t id, std::vector &pruned_list, + DiskAnnContext *ctx) { + const std::vector init_ids = get_init_ids(ctx); + + int ret = iterate_to_fixed_point(init_ids, ctx); + if (ret != 0) { + return ret; + } + + auto &pool = ctx->expanded_nodes(); + pool.erase(std::remove_if(pool.begin(), pool.end(), + [id](const auto &node) { return node.id == id; }), + pool.end()); + + ret = prune_neighbors(id, pool, pruned_list, ctx); + if (ret != 0) { + return ret; + } + + return 0; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_algorithm.h b/src/core/algorithm/diskann/diskann_algorithm.h new file mode 100644 index 0000000..da64ffc --- /dev/null +++ b/src/core/algorithm/diskann/diskann_algorithm.h @@ -0,0 +1,72 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "diskann_context.h" + +namespace zvec { +namespace core { + +class DiskAnnAlgorithm { + public: + typedef std::unique_ptr UPointer; + + public: + DiskAnnAlgorithm(DiskAnnEntity &entity, uint32_t max_degree); + + public: + int add_node(diskann_id_t id, DiskAnnContext *ctx); + int prune_node(diskann_id_t id, DiskAnnContext *ctx); + + private: + int search_neighbor_and_prune(diskann_id_t id, + std::vector &pruned_list, + DiskAnnContext *ctx); + int iterate_to_fixed_point(const std::vector &init_ids, + DiskAnnContext *ctx); + int prune_neighbors(diskann_id_t id, std::vector &pool, + std::vector &pruned_list, + DiskAnnContext *ctx); + int inter_insert(diskann_id_t id, std::vector &pruned_list, + DiskAnnContext *ctx); + int occlude_list(diskann_id_t id, std::vector &pool, + std::vector &result, DiskAnnContext *ctx); + + std::vector get_init_ids(DiskAnnContext *ctx); + + private: + static constexpr uint32_t kLockCnt{1U << 16}; + static constexpr uint32_t kLockMask{kLockCnt - 1U}; + + DiskAnnEntity &entity_; + + uint32_t max_degree_{DiskAnnEntity::kDefaultMaxDegree}; + uint32_t max_candidate_size_{DiskAnnEntity::kDefaultMaxOcclusionSize}; + + std::vector lock_pool_{}; + + //! Acquire the striped lock guarding the given id (RAII). + [[nodiscard]] std::unique_lock lock_for(diskann_id_t id) { + return std::unique_lock(lock_pool_[id & kLockMask]); + } + + float alpha_{DiskAnnEntity::kDefaultAlpha}; + bool saturate_graph_{true}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc new file mode 100644 index 0000000..e344844 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -0,0 +1,638 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_builder.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "algorithm/cluster/vector_mean.h" +#include "diskann_context.h" +#include "diskann_params.h" + +namespace zvec { +namespace core { + +int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { + LOG_INFO("Begin DiskAnnBuilder::init"); + + params.get(PARAM_DISKANN_BUILDER_MAX_DEGREE, &max_degree_); + params.get(PARAM_DISKANN_BUILDER_LIST_SIZE, &list_size_); + params.get(PARAM_DISKANN_BUILDER_THREAD_COUNT, &build_thread_count_); + + if (build_thread_count_ == 0) { + build_thread_count_ = std::thread::hardware_concurrency(); + } + + if (build_thread_count_ > std::thread::hardware_concurrency()) { + LOG_WARN("Build thread count [%s] greater than cpu cores %u", + PARAM_DISKANN_BUILDER_THREAD_COUNT.c_str(), + std::thread::hardware_concurrency()); + } + + uint32_t max_pq_chunk_num{0}; + if (params.get(PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM, &max_pq_chunk_num)) { + if (max_pq_chunk_num > meta.dimension()) { + LOG_ERROR( + "PQ Chunk Num larger than dimension, PQ Chunk Num: %d, Dimension: %d", + max_pq_chunk_num, meta.dimension()); + return IndexError_InvalidArgument; + } + + max_pq_chunk_num_ = max_pq_chunk_num; + } + + if (params.has(PARAM_DISKANN_BUILDER_MEMORY_LIMIT)) { + params.get(PARAM_DISKANN_BUILDER_MEMORY_LIMIT, &memory_limit_); + if (memory_limit_ <= 0) { + LOG_ERROR("Invalid memory limit: %lf", memory_limit_); + return IndexError_InvalidArgument; + } + + memory_limit_set_ = true; + } + + if (params.has(PARAM_DISKANN_BUILDER_MAX_TRAIN_SAMPLE_COUNT)) { + params.get(PARAM_DISKANN_BUILDER_MAX_TRAIN_SAMPLE_COUNT, + &max_train_sample_count_); + } + + if (params.has(PARAM_DISKANN_BUILDER_TRAIN_SAMPLE_RATIO)) { + params.get(PARAM_DISKANN_BUILDER_TRAIN_SAMPLE_RATIO, &train_sample_ratio_); + } + + raw_meta_ = meta; + + build_meta_ = meta; + if (meta.metric_name() == "InnerProduct") { + build_meta_.set_metric("SquaredEuclidean", 0, ailego::Params()); + } else if (meta.metric_name() == "Cosine") { + build_meta_.set_metric("SquaredEuclidean", 0, ailego::Params()); + + if (meta.data_type() == IndexMeta::DataType::DT_FP32) { + build_meta_.set_dimension(meta.dimension() - 1); + } else { + build_meta_.set_dimension(meta.dimension() - 2); + } + } + + metric_ = IndexFactory::CreateMetric(build_meta_.metric_name()); + if (!metric_) { + LOG_ERROR("CreateMetric failed, name: %s", + build_meta_.metric_name().c_str()); + return IndexError_NoExist; + } + + int ret = metric_->init(build_meta_, build_meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMeasure init failed, ret=%d", ret); + return ret; + } + + raw_meta_.set_builder("DiskAnnBuilder", DiskAnnEntity::kRevision, params); + + ret = entity_.init(meta, max_degree_, list_size_, memory_limit_, + build_thread_count_); + if (ret != 0) { + return ret; + } + + algo_ = + DiskAnnAlgorithm::UPointer(new DiskAnnAlgorithm(entity_, max_degree_)); + + trainer_ = + DiskAnnPqTrainer::UPointer(new DiskAnnPqTrainer(max_train_sample_count_)); + + state_ = BUILD_STATE_INITED; + + return 0; +} + +int DiskAnnBuilder::cleanup(void) { + return 0; +} + +int DiskAnnBuilder::calculate_entry_point() { + size_t dimension = build_meta_.dimension(); + + if (build_meta_.data_type() != IndexMeta::DataType::DT_FP32 && + build_meta_.data_type() != IndexMeta::DataType::DT_FP16) { + LOG_ERROR("Data type not supported"); + return IndexError_InvalidArgument; + } + + std::vector centroid_fp32; + std::vector centroid_fp16; + + switch (build_meta_.data_type()) { + case IndexMeta::DataType::DT_FP32: { + centroid_fp32.resize(dimension); + NumericalVectorMean accumulator(dimension); + for (size_t id = 0; id < entity_.doc_cnt(); id++) { + accumulator.plus(entity_.get_vector(id), dimension * sizeof(float)); + } + accumulator.mean(centroid_fp32.data(), dimension * sizeof(float)); + break; + } + case IndexMeta::DataType::DT_FP16: { + centroid_fp16.resize(dimension); + NumericalVectorMean accumulator(dimension); + for (size_t id = 0; id < entity_.doc_cnt(); id++) { + accumulator.plus(entity_.get_vector(id), + dimension * sizeof(ailego::Float16)); + } + accumulator.mean(centroid_fp16.data(), + dimension * sizeof(ailego::Float16)); + break; + } + default: + return IndexError_Unsupported; + } + + // compute all to one distance + diskann_id_t medoid_id = kInvalidId; + float min_dist = std::numeric_limits::max(); + + switch (build_meta_.data_type()) { + case IndexMeta::DataType::DT_FP32: + for (size_t id = 0; id < entity_.doc_cnt(); id++) { + const float *data_ptr = + reinterpret_cast(entity_.get_vector(id)); + + float dist = 0.0f; + ailego::SquaredEuclideanDistanceMatrix::Compute( + centroid_fp32.data(), data_ptr, dimension, &dist); + + if (dist < min_dist) { + min_dist = dist; + medoid_id = id; + } + } + break; + case IndexMeta::DataType::DT_FP16: + for (size_t id = 0; id < entity_.doc_cnt(); id++) { + const ailego::Float16 *data_ptr = + reinterpret_cast(entity_.get_vector(id)); + + float dist = 0.0f; + ailego::SquaredEuclideanDistanceMatrix::Compute( + centroid_fp16.data(), data_ptr, dimension, &dist); + + if (dist < min_dist) { + min_dist = dist; + medoid_id = id; + } + } + break; + default: + return IndexError_Unsupported; + } + + (*entity_.mutable_medoid()) = medoid_id; + + LOG_INFO("Medoid Calculation Done. ID: %zu", (size_t)medoid_id); + + return 0; +} + +int DiskAnnBuilder::calculate_pq_chunk_num() { + size_t doc_cnt = holder_->count(); + if (doc_cnt == 0) { + LOG_ERROR("Invalid Input. Empty Vecs."); + + return IndexError_InvalidLength; + } + + if (memory_limit_set_) { + size_t memory_limit_bytes = get_memory_in_bytes(memory_limit_); + size_t pq_chunk_num = std::floor(memory_limit_bytes / doc_cnt); + if (pq_chunk_num <= 0) { + LOG_ERROR("Insufficient memory limit for vec, memory: %zu, vec num: %zu", + memory_limit_bytes, doc_cnt); + return IndexError_InvalidArgument; + } + } + + pq_chunk_num_ = + pq_chunk_num_ < max_pq_chunk_num_ ? pq_chunk_num_ : max_pq_chunk_num_; + + // A chunk num of 0 (public API default) or the internal sentinel means + // "auto": quantize into half the dimensions. Resolve it before the + // upper-bound check so the default never reaches the divide-by-chunk path. + if (pq_chunk_num_ == 0 || pq_chunk_num_ == kDefaultPqChunkNum) { + pq_chunk_num_ = build_meta_.dimension() / 2; + LOG_INFO( + "No Chunk Num input. Quantizing %u dimension data into %u dimension.", + build_meta_.dimension(), pq_chunk_num_); + } + + if (pq_chunk_num_ > build_meta_.dimension()) { + LOG_ERROR("PQ Chunk Num is more than dimension, chunk num: %u, dim: %u", + pq_chunk_num_, build_meta_.dimension()); + return IndexError_InvalidArgument; + } + + LOG_INFO("Quantizing %u dimension data into %u bytes.", + build_meta_.dimension(), pq_chunk_num_); + + return 0; +} + +int DiskAnnBuilder::build_internal(IndexThreads::Pointer threads) { + auto task_group = threads->make_group(); + if (!task_group) { + LOG_ERROR("Failed to create task group"); + return IndexError_Runtime; + } + + std::atomic finished{0}; + for (size_t i = 0; i < threads->count(); ++i) { + task_group->submit(ailego::Closure ::New(this, &DiskAnnBuilder::do_build, i, + threads->count(), &finished)); + } + + { + std::unique_lock lk(mutex_); + while (finished.load() < entity_.doc_cnt()) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to build index while waiting finish"); + return errcode_; + } + LOG_INFO("Built cnt %zu, finished percent %.3f%%", + (size_t)finished.load(), + finished.load() * 100.0f / entity_.doc_cnt()); + } + } + + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to build index while waiting finish"); + return errcode_; + } + task_group->wait_finish(); + + return 0; +} + +int DiskAnnBuilder::prune_internal(IndexThreads::Pointer threads) { + auto task_group = threads->make_group(); + if (!task_group) { + LOG_ERROR("Failed to create task group"); + return IndexError_Runtime; + } + + std::atomic finished{0}; + for (size_t i = 0; i < threads->count(); ++i) { + task_group->submit(ailego::Closure ::New(this, &DiskAnnBuilder::do_prune, i, + threads->count(), &finished)); + } + + { + std::unique_lock lk(mutex_); + while (finished.load() < entity_.doc_cnt()) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to prune index while waiting finish"); + return errcode_; + } + LOG_INFO("Prune cnt %zu, finished percent %.3f%%", + (size_t)finished.load(), + finished.load() * 100.0f / entity_.doc_cnt()); + } + } + + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to prune index while waiting finish"); + return errcode_; + } + task_group->wait_finish(); + + return 0; +} + +int DiskAnnBuilder::train_quantized_data(IndexThreads::Pointer threads) { + LOG_INFO("Starting Train: Chunk Num: %u", pq_chunk_num_); + + ailego::ElapsedTime timer; + int ret = trainer_->train_quantized_data( + threads, holder_, build_meta_, entity_.pq_full_pivot_data(), + entity_.pq_centroid(), entity_.pq_chunk_offsets(), pq_chunk_num_); + if (ret != 0) { + LOG_ERROR("Train Quantized Data Error, ret=%d", ret); + return ret; + } + + size_t pq_time = timer.milli_seconds(); + LOG_INFO("Train Quantized Data Done, time: %zu ms", pq_time); + + (*entity_.mutable_pq_meta()).full_pivot_data_size = + entity_.pq_full_pivot_data().size(); + (*entity_.mutable_pq_meta()).centroid_data_size = + entity_.pq_centroid().size(); + (*entity_.mutable_pq_meta()).chunk_num = pq_chunk_num_; + + return 0; +} + +int DiskAnnBuilder::generate_quantized_data(IndexThreads::Pointer threads) { + LOG_INFO("Starting PQ Generate: Query Memory Limit: %lf, Chunk Num: %u", + memory_limit_, pq_chunk_num_); + + ailego::ElapsedTime timer; + int ret = trainer_->generate_quantized_data( + threads, holder_, build_meta_, entity_.pq_centroid(), + entity_.block_compressed_data(), pq_chunk_num_); + if (ret != 0) { + LOG_ERROR("Generate Quantized Data Error, ret=%d", ret); + return ret; + } + + size_t pq_time = timer.milli_seconds(); + LOG_INFO("Generate Quantized Data Done, time: %zu ms", pq_time); + + return 0; +} + +void DiskAnnBuilder::do_build(uint64_t idx, size_t step_size, + std::atomic *finished) { + AILEGO_DEFER([&]() { + std::lock_guard latch(mutex_); + cond_.notify_one(); + }); + + DiskAnnContext *ctx = new (std::nothrow) DiskAnnContext( + build_meta_, metric_, + std::shared_ptr(&entity_, [](DiskAnnEntity *) {})); + + if (ailego_unlikely(ctx == nullptr)) { + if (!error_.exchange(true)) { + LOG_ERROR("Failed to create context"); + errcode_ = IndexError_NoMemory; + } + return; + } + + DiskAnnContext::Pointer auto_ptr(ctx); + ctx->init(DiskAnnContext::kBuilderContext, max_degree_, pq_chunk_num_, + build_meta_.element_size()); + ctx->set_list_size(list_size_); + + for (uint64_t id = idx; id < entity_.doc_cnt(); id += step_size) { + ctx->reset_query(entity_.get_vector(id)); + int ret = algo_->add_node(id, ctx); + if (ailego_unlikely(ret != 0)) { + if (!error_.exchange(true)) { + LOG_ERROR("DiskAnn graph add node failed"); + errcode_ = ret; + } + return; + } + ctx->clear(); + (*finished)++; + } +} + +void DiskAnnBuilder::do_prune(uint64_t idx, size_t step_size, + std::atomic *finished) { + AILEGO_DEFER([&]() { + std::lock_guard latch(mutex_); + cond_.notify_one(); + }); + + DiskAnnContext *ctx = new (std::nothrow) DiskAnnContext( + build_meta_, metric_, + std::shared_ptr(&entity_, [](DiskAnnEntity *) {})); + + if (ailego_unlikely(ctx == nullptr)) { + if (!error_.exchange(true)) { + LOG_ERROR("Failed to create context"); + errcode_ = IndexError_NoMemory; + } + return; + } + + DiskAnnContext::Pointer auto_ptr(ctx); + ctx->init(DiskAnnContext::kBuilderContext, max_degree_, pq_chunk_num_, + build_meta_.element_size()); + ctx->set_list_size(list_size_); + + for (uint64_t id = idx; id < entity_.doc_cnt(); id += step_size) { + ctx->reset_query(entity_.get_vector(id)); + int ret = algo_->prune_node(id, ctx); + if (ailego_unlikely(ret != 0)) { + if (!error_.exchange(true)) { + LOG_ERROR("DiskAnn graph add node failed"); + errcode_ = ret; + } + return; + } + ctx->clear(); + (*finished)++; + } +} + +int DiskAnnBuilder::train(const IndexTrainer::Pointer & /*trainer*/) { + if (state_ != BUILD_STATE_INITED) { + LOG_ERROR("Init the builder before DiskAnnBuilder::train"); + return IndexError_NoReady; + } + + LOG_INFO("Begin DiskAnnBuilder::train by trainer"); + + stats_.set_trained_count(0UL); + stats_.set_trained_costtime(0UL); + state_ = BUILD_STATE_TRAINED; + + LOG_INFO("End DiskAnnBuilder::train by trainer"); + + return 0; +} + +int DiskAnnBuilder::train(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) { + if (state_ != BUILD_STATE_INITED) { + LOG_ERROR("Init the builder before DiskAnnBuilder::train"); + return IndexError_NoReady; + } + + LOG_INFO("Begin DiskAnnBuilder::train"); + + auto start_time = ailego::Monotime::MilliSeconds(); + + holder_ = std::move(holder); + + LOG_INFO("Start to calculate chunk num"); + int ret = calculate_pq_chunk_num(); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + if (!threads) { + threads = + std::make_shared(build_thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + + ret = train_quantized_data(threads); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + stats_.set_trained_count(holder_->count()); + + stats_.set_trained_costtime(ailego::Monotime::MilliSeconds() - start_time); + + state_ = BUILD_STATE_TRAINED; + + holder_.reset(); + + LOG_INFO("End DiskAnnBuilder::train"); + + return 0; +} + +int DiskAnnBuilder::do_norm(const void *data_ptr, std::string *norm_data) { + size_t dimension = build_meta_.dimension(); + const float *float_data_ptr = reinterpret_cast(data_ptr); + + norm_data->resize(dimension * sizeof(float)); + float *output_buf = reinterpret_cast(&((*norm_data)[0])); + std::memcpy(output_buf, float_data_ptr, dimension * sizeof(float)); + + float norm = 0.0f; + ailego::Normalizer::L2(output_buf, dimension, &norm); + + return 0; +} + +int DiskAnnBuilder::build(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) { + LOG_INFO("Start DiskAnnBuilder::build"); + + auto start_time = ailego::Monotime::MilliSeconds(); + + holder_ = holder; + + if (!threads) { + threads = + std::make_shared(build_thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + + auto iter = holder->create_iterator(); + if (!iter) { + LOG_ERROR("Create iterator for holder failed"); + return IndexError_Runtime; + } + + if (ailego_unlikely(holder->count() == 0)) { + LOG_ERROR("Holder is empty"); + return IndexError_Runtime; + } + + int ret = entity_.reserve_space(holder->count()); + + error_ = false; + while (iter->is_valid()) { + ret = entity_.add_vector(iter->key(), iter->data()); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + iter->next(); + } + + LOG_INFO("Finished saving vector"); + + LOG_INFO("Start to calculate entrypoint"); + ret = calculate_entry_point(); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + LOG_INFO("Start to build vamana graph"); + ret = build_internal(threads); + if (ret != 0) { + return ret; + } + + LOG_INFO("Start final cleanup.."); + ret = prune_internal(threads); + if (ret != 0) { + return ret; + } + + LOG_INFO("Start to generate quantized data"); + ret = generate_quantized_data(threads); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + state_ = BUILD_STATE_BUILT; + + stats_.set_built_count(entity_.doc_cnt()); + stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time); + + LOG_INFO("End DiskAnnBuilder::build"); + + return 0; +} + +int DiskAnnBuilder::dump(const IndexDumper::Pointer &dumper) { + if (state_ != BUILD_STATE_BUILT) { + LOG_INFO("Build the index before DiskAnnBuilder::dump"); + return IndexError_NoReady; + } + + LOG_INFO("Begin DiskAnnBuilder::dump"); + + raw_meta_.set_searcher("DiskAnnSearcher", 0, ailego::Params()); + auto start_time = ailego::Monotime::MilliSeconds(); + + int ret = IndexHelper::SerializeToDumper(raw_meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + + ret = entity_.dump(holder_, raw_meta_, dumper); + if (ret != 0) { + LOG_ERROR("Index dump failed, ret: %u", ret); + + return IndexError_Runtime; + } + + stats_.set_dumped_count(holder_->count()); + stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time); + + LOG_INFO("DiskAnnBuilder::dump"); + + return 0; +} + +INDEX_FACTORY_REGISTER_BUILDER(DiskAnnBuilder); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_builder.h b/src/core/algorithm/diskann/diskann_builder.h new file mode 100644 index 0000000..b8df7f5 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_builder.h @@ -0,0 +1,131 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "diskann_algorithm.h" +#include "diskann_builder_entity.h" +#include "diskann_pq_trainer.h" + +namespace zvec { +namespace core { + +class DiskAnnBuilder : public IndexBuilder { + public: + //! Constructor + DiskAnnBuilder() = default; + + //! Initialize the builder + virtual int init(const IndexMeta &meta, + const ailego::Params ¶ms) override; + + //! Cleanup the builder + virtual int cleanup(void) override; + + //! Train the data + virtual int train(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) override; + + //! Train the data + virtual int train(const IndexTrainer::Pointer &trainer) override; + + //! Build the index + virtual int build(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) override; + + //! Dump index into storage + virtual int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + virtual const Stats &stats(void) const override { + return stats_; + } + + int do_norm(const void *data_ptr, std::string *norm_data); + + private: + int train_quantized_data(IndexThreads::Pointer threads); + int generate_quantized_data(IndexThreads::Pointer threads); + int build_internal(IndexThreads::Pointer threads); + int prune_internal(IndexThreads::Pointer threads); + + void do_build(uint64_t idx, size_t step_size, + std::atomic *finished); + + void do_prune(uint64_t idx, size_t step_size, + std::atomic *finished); + + int calculate_entry_point(); + + int calculate_pq_chunk_num(); + + double get_memory_in_bytes(double search_ram_budget) { + return search_ram_budget * 1024 * 1024 * 1024; + } + + private: + enum BUILD_STATE { + BUILD_STATE_INIT = 0, + BUILD_STATE_INITED = 1, + BUILD_STATE_TRAINED = 2, + BUILD_STATE_BUILT = 3 + }; + + constexpr static uint32_t kDefaultLogIntervalSecs = 15U; + constexpr static uint32_t kDefaultListSize = 50U; + constexpr static uint32_t kDefaultMaxDegree = 100U; + constexpr static uint32_t kDefaultPqChunkNum = -1U; + + std::string data_file_; + + uint32_t max_degree_{kDefaultMaxDegree}; + uint32_t list_size_{kDefaultListSize}; + double memory_limit_{0.0}; + bool memory_limit_set_{false}; + uint32_t max_pq_chunk_num_{kDefaultPqChunkNum}; + uint32_t pq_chunk_num_{kDefaultPqChunkNum}; + uint32_t build_thread_count_{0}; + uint32_t max_train_sample_count_{PQTable::kMaxTrainSampleCount}; + double train_sample_ratio_{PQTable::kTrainSampleRatio}; + std::string universal_label_{""}; + std::string codebook_prefix_{""}; + std::string index_path_prefix_{"./diskann"}; + + BUILD_STATE state_; + Stats stats_; + + int errcode_{0}; + std::atomic_bool error_{false}; + + IndexMetric::Pointer metric_{}; + + std::mutex mutex_{}; + std::condition_variable cond_{}; + + IndexMeta raw_meta_; + IndexMeta build_meta_; + + DiskAnnBuilderEntity entity_{}; + + IndexHolder::Pointer holder_; + + DiskAnnAlgorithm::UPointer algo_; + DiskAnnPqTrainer::UPointer trainer_; + + uint32_t check_interval_secs_{kDefaultLogIntervalSecs}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_builder_entity.cc b/src/core/algorithm/diskann/diskann_builder_entity.cc new file mode 100644 index 0000000..4381293 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_builder_entity.cc @@ -0,0 +1,643 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_builder_entity.h" +#include +#include "diskann_algorithm.h" +#include "diskann_util.h" + +namespace zvec { +namespace core { + +int DiskAnnBuilderEntity::init(const IndexMeta &meta, uint32_t max_degree, + uint32_t list_size, double memory_limit, + uint32_t build_threads) { + meta_ = meta; + + max_degree_ = max_degree; + list_size_ = list_size; + + memory_limit_ = memory_limit; + + num_threads_ = build_threads; + + max_build_degree_ = max_degree_ * kDefaultGraphSlackFactor; + + neighbor_size_ = sizeof(uint32_t) + max_build_degree_ * sizeof(diskann_id_t); + + return 0; +} + +int DiskAnnBuilderEntity::reserve_space(uint32_t docs) { + vectors_buffer_.reserve(meta_.element_size() * docs); + keys_buffer_.reserve(sizeof(diskann_key_t) * docs); + neighbors_buffer_.reserve(neighbor_size_ * docs); + + return 0; +} + +int DiskAnnBuilderEntity::add_vector(diskann_key_t key, const void *vec) { + vectors_buffer_.append(reinterpret_cast(vec), + meta_.element_size()); + keys_buffer_.append(reinterpret_cast(&key), sizeof(key)); + + uint32_t neighbor_cnt = 0; + std::vector neighbor{max_build_degree_, 0}; + + neighbors_buffer_.append(reinterpret_cast(&neighbor_cnt), + sizeof(uint32_t)); + neighbors_buffer_.append(reinterpret_cast(neighbor.data()), + sizeof(diskann_id_t) * max_build_degree_); + + (*mutable_doc_cnt())++; + + return 0; +} + +const void *DiskAnnBuilderEntity::get_vector(diskann_id_t id) const { + size_t offset = (size_t)id * meta_.element_size(); + return vectors_buffer_.data() + offset; +} + +diskann_key_t DiskAnnBuilderEntity::get_key(diskann_id_t id) const { + size_t offset = (size_t)id * sizeof(diskann_key_t); + + return *( + reinterpret_cast(keys_buffer_.data() + offset)); +} + +//! Get vector local id by key +diskann_id_t DiskAnnBuilderEntity::get_id(diskann_key_t /*key*/) const { + LOG_ERROR("DiskAnnBuilderEntity::get_id not implemented."); + return kInvalidId; +} + +std::pair DiskAnnBuilderEntity::get_neighbors( + diskann_id_t id) const { + size_t offset = (size_t)id * neighbor_size_; + + const uint8_t *start_ptr = + reinterpret_cast(neighbors_buffer_.data()) + offset; + + uint32_t neighbor_cnt = *(reinterpret_cast(start_ptr)); + + const diskann_id_t *neighbors = + reinterpret_cast(start_ptr + sizeof(uint32_t)); + + return std::make_pair(neighbor_cnt, neighbors); +} + +int DiskAnnBuilderEntity::set_neighbors( + diskann_id_t id, const std::vector &neighbor_ids) { + size_t offset = (size_t)id * neighbor_size_; + + uint8_t *start_ptr = + reinterpret_cast(&neighbors_buffer_[0]) + offset; + + uint32_t neighbor_cnt = neighbor_ids.size(); + + memcpy(start_ptr + sizeof(uint32_t), neighbor_ids.data(), + sizeof(diskann_id_t) * neighbor_cnt); + memcpy(start_ptr, &neighbor_cnt, sizeof(uint32_t)); + + if (max_observed_degree_ < neighbor_cnt) { + max_observed_degree_ = neighbor_cnt; + } + + return 0; +} + +int DiskAnnBuilderEntity::add_neighbor(diskann_id_t id, + diskann_id_t neighbor_id) { + size_t offset = (size_t)id * neighbor_size_; + + uint8_t *start_ptr = + reinterpret_cast(&neighbors_buffer_[0]) + offset; + + uint32_t neighbor_cnt = *reinterpret_cast(start_ptr); + + memcpy(start_ptr + sizeof(uint32_t) + sizeof(diskann_id_t) * neighbor_cnt, + &neighbor_id, sizeof(diskann_id_t)); + + neighbor_cnt += 1; + + memcpy(start_ptr, &neighbor_cnt, sizeof(uint32_t)); + + if (max_observed_degree_ < neighbor_cnt) { + max_observed_degree_ = neighbor_cnt; + } + + return 0; +} + +int64_t DiskAnnBuilderEntity::dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, + const void *data, + size_t size) const { + size_t len = dumper->write(data, size); + if (len != size) { + LOG_ERROR("Dump segment %s data failed, expect: %lu, actual: %lu", + segment_id.c_str(), size, len); + return IndexError_WriteData; + } + + size_t padding_size = AlignSize(size) - size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + uint32_t crc = ailego::Crc32c::Hash(data, size); + int ret = dumper->append(segment_id, size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", segment_id.c_str(), ret); + return ret; + } + + return 0; +} + +int DiskAnnBuilderEntity::dump_pq_meta_segment( + const IndexDumper::Pointer &dumper) const { + uint32_t crc = 0U; + + // write meta + size_t size_pq_meta = dumper->write(&pq_meta_, sizeof(DiskAnnPqMeta)); + if (size_pq_meta != sizeof(DiskAnnPqMeta)) { + LOG_ERROR("Failed to dump PQ meta data, expect: %lu, actual: %lu", + sizeof(DiskAnnPqMeta), size_pq_meta); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(&pq_meta_, sizeof(DiskAnnPqMeta), crc); + + // write full pivot data + size_t size_full_pivot_data = + dumper->write(pq_full_pivot_data_.data(), pq_meta_.full_pivot_data_size); + if (size_full_pivot_data != pq_meta_.full_pivot_data_size) { + LOG_ERROR("Failed to dump full pivot data, expect: %zu, actual: %zu", + (size_t)pq_meta_.full_pivot_data_size, size_full_pivot_data); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(pq_full_pivot_data_.data(), + pq_meta_.full_pivot_data_size, crc); + + // write centroid num + size_t size_centroid = + dumper->write(pq_centroid_.data(), pq_meta_.centroid_data_size); + if (size_centroid != pq_meta_.centroid_data_size) { + LOG_ERROR("Failed to dump centroid num, expect: %zu, actual: %zu", + (size_t)pq_meta_.centroid_data_size, size_centroid); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(pq_centroid_.data(), pq_meta_.centroid_data_size, + crc); + + // write chunk offset + size_t size_chunk_offset = dumper->write( + pq_chunk_offsets_.data(), (pq_meta_.chunk_num + 1) * sizeof(uint32_t)); + if (size_chunk_offset != (pq_meta_.chunk_num + 1) * sizeof(uint32_t)) { + LOG_ERROR("Failed to dump centroid num, expect: %zu, actual: %zu", + (size_t)((pq_meta_.chunk_num + 1) * sizeof(uint32_t)), + size_chunk_offset); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(pq_chunk_offsets_.data(), + (pq_meta_.chunk_num + 1) * sizeof(uint32_t), crc); + + // write size + size_t size_total = + size_pq_meta + size_full_pivot_data + size_centroid + size_chunk_offset; + + // write pad + size_t padding_size = AlignSize(size_total) - size_total; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + int ret = + dumper->append(kDiskAnnPqMetaSegmentId, size_total, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump PQ segment failed, ret %d", ret); + return ret; + } + + return 0; +} + +int DiskAnnBuilderEntity::dump_pq_data_segment( + const IndexDumper::Pointer &dumper) const { + uint64_t doc_cnt = meta_header_.doc_cnt; + uint64_t chunk_num = pq_meta_.chunk_num; + + uint32_t crc = 0U; + + // write pq data + size_t size_total = + dumper->write(block_compressed_data_.data(), doc_cnt * chunk_num); + + if (size_total != doc_cnt * chunk_num) { + LOG_ERROR("Failed to dump block compressed data, expect: %zu, actual: %zu", + (size_t)(doc_cnt * chunk_num), size_total); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(block_compressed_data_.data(), doc_cnt * chunk_num, + crc); + + // write pad + size_t padding_size = AlignSize(size_total) - size_total; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + int ret = + dumper->append(kDiskAnnPqDataSegmentId, size_total, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump PQ data segment failed, ret %d", ret); + return ret; + } + + return 0; +} + +int DiskAnnBuilderEntity::dump_dummy_segment( + const IndexDumper::Pointer &dumper) const { + // to make offset aligned with 4K + size_t dumper_header_size = dumper->size(); + + size_t dummy_size = + DiskAnnUtil::round_up(dumper_header_size, DiskAnnUtil::kSectorSize) - + dumper_header_size; + + if (dummy_size != 0) { + std::string dummpy_data(dummy_size, '\0'); + if (dumper->write(dummpy_data.data(), dummy_size) != dummy_size) { + LOG_ERROR("write dummy failed, size %lu", dummy_size); + return IndexError_WriteData; + } + + int ret = dumper->append(kDiskAnnDummpySegmentId, dummy_size, 0, 0); + if (ret != 0) { + LOG_ERROR("Dump dummy data segment failed, ret %d", ret); + return ret; + } + } + + return 0; +} + +int DiskAnnBuilderEntity::dump_key_segment( + const IndexDumper::Pointer &dumper) const { + //! Dump keys + size_t key_segment_size = doc_cnt() * sizeof(diskann_key_t); + int64_t keys_size = dump_segment(dumper, kDiskAnnKeySegmentId, + keys_buffer_.data(), key_segment_size); + if (keys_size < 0) { + return keys_size; + } + + return 0; +} + +int DiskAnnBuilderEntity::dump_key_mapping_segment( + const IndexDumper::Pointer &dumper) const { + std::vector mapping(doc_cnt()); + + const diskann_key_t *keys = reinterpret_cast( + const_cast(keys_buffer_.data())); + + std::iota(mapping.begin(), mapping.end(), 0U); + std::sort(mapping.begin(), mapping.end(), + [&](diskann_id_t i, diskann_id_t j) { return keys[i] < keys[j]; }); + + size_t size = mapping.size() * sizeof(diskann_id_t); + int64_t ret = + dump_segment(dumper, kDiskAnnKeyMappingSegmentId, mapping.data(), size); + + if (ret != 0) { + LOG_ERROR("Dump vectors segment failed"); + + return ret; + } + + return 0; +} + +int DiskAnnBuilderEntity::dump_entrypoint_segment( + const IndexDumper::Pointer &dumper) const { + std::string entrypoint_buffer; + + size_t size = sizeof(uint32_t) + entrypoints_.size() * sizeof(diskann_id_t); + entrypoint_buffer.resize(size); + + uint8_t *buffer_ptr = reinterpret_cast(&entrypoint_buffer[0]); + + uint32_t point_cnt = entrypoints_.size(); + memcpy(buffer_ptr, &point_cnt, sizeof(uint32_t)); + memcpy(buffer_ptr + sizeof(uint32_t), entrypoints_.data(), + entrypoints_.size() * sizeof(diskann_id_t)); + + int64_t ret = dump_segment(dumper, kDiskAnnEntryPointSegmentId, + entrypoint_buffer.data(), size); + + if (ret != 0) { + LOG_ERROR("Dump entrypoint segment failed"); + + return ret; + } + + return 0; +} + +int DiskAnnBuilderEntity::dump(IndexHolder::Pointer holder, IndexMeta &meta, + const IndexDumper::Pointer &dumper) { + uint64_t doc_cnt = holder->count(); + uint64_t max_node_size = + (uint64_t)max_observed_degree_ * sizeof(diskann_id_t) + sizeof(uint32_t) + + meta_.element_size(); + uint64_t node_per_sector = + DiskAnnUtil::kSectorSize / + max_node_size; // 0 if max_node_size > DiskAnnUtil::kSectorSize + + std::string node_buf; + node_buf.resize(max_node_size); + + diskann_id_t *neighbor_buf = + (diskann_id_t *)(node_buf.data() + (meta_.element_size()) + + sizeof(uint32_t)); + + LOG_INFO( + "Dump Data, medoid: %zu, max node size: %zu, node per sector: %zu, " + "max observed degree: %zu", + (size_t)medoid(), (size_t)max_node_size, (size_t)node_per_sector, + (size_t)max_observed_degree_); + + // write a dummy segment to make data align + int ret = dump_dummy_segment(dumper); + if (ret != 0) { + LOG_ERROR("Dump dummy segment failed"); + + return ret; + } + + // dump data by sector + size_t write_size = 0; + uint32_t crc = 0U; + size_t len = 0; + + // no need to write first sector + auto iter = holder->create_iterator(); + if (!iter) { + LOG_ERROR("Create iterator for holder failed"); + return IndexError_Runtime; + } + + uint64_t index_size = 0; + uint32_t neighbor_num; + if (node_per_sector > 0) { + uint64_t sector_num = + DiskAnnUtil::round_up(doc_cnt, node_per_sector) / node_per_sector; + + diskann_id_t cur_node_id = 0; + + std::string sector_buf; + sector_buf.resize(DiskAnnUtil::kSectorSize); + + for (uint64_t sector = 0; sector < sector_num; sector++) { + if (sector != 0 && sector % 100000 == 0) { + LOG_INFO("Sector #%zu written", (size_t)sector); + } + + memset(&(sector_buf[0]), 0, DiskAnnUtil::kSectorSize); + + for (uint64_t sector_node_id = 0; + sector_node_id < node_per_sector && cur_node_id < doc_cnt; + sector_node_id++) { + memset(&(node_buf[0]), 0, max_node_size); + + auto neighbors = get_neighbors(cur_node_id); + neighbor_num = neighbors.first; + + ailego_assert(neighbor_num > 0); + ailego_assert(neighbor_num <= max_observed_degree_); + + memcpy(&(neighbor_buf[0]), neighbors.second, + neighbors.first * sizeof(diskann_id_t)); + + if (iter->is_valid()) { + const void *vec = iter->data(); + memcpy(&(node_buf[0]), vec, meta.element_size()); + + iter->next(); + } else { + return IndexError_Runtime; + } + + // write neighbor num + *(uint32_t *)(node_buf.data() + meta_.element_size()) = neighbor_num; + + // write neighbor buffer + memcpy(&(node_buf[0]) + meta_.element_size() + sizeof(uint32_t), + neighbor_buf, neighbor_num * sizeof(diskann_id_t)); + + // get offset into sector_buf + char *sector_node_buf = §or_buf[sector_node_id * max_node_size]; + + // copy node buf into sector_node_buf + memcpy(sector_node_buf, node_buf.data(), max_node_size); + + cur_node_id++; + } + + // flush sector to disk + len = dumper->write(sector_buf.data(), DiskAnnUtil::kSectorSize); + if (len != DiskAnnUtil::kSectorSize) { + LOG_ERROR("Write Vector Error, write=%zu, expect=%zu", len, + (size_t)DiskAnnUtil::kSectorSize); + + return IndexError_WriteData; + } + write_size += DiskAnnUtil::kSectorSize; + crc = ailego::Crc32c::Hash(sector_buf.data(), DiskAnnUtil::kSectorSize, + crc); + } + + LOG_INFO("Total Sector #%zu written", (size_t)sector_num); + + index_size = sector_num * DiskAnnUtil::kSectorSize; + } else { + // Write multi-sector nodes + std::string multisector_buf; + multisector_buf.resize( + DiskAnnUtil::round_up(max_node_size, DiskAnnUtil::kSectorSize)); + + uint64_t sector_num_per_node = + DiskAnnUtil::div_round_up(max_node_size, DiskAnnUtil::kSectorSize); + + for (uint64_t i = 0; i < doc_cnt; i++) { + if (i != 0 && (i * sector_num_per_node) % 100000 == 0) { + LOG_INFO("Sector # %zu written", (size_t)(i * sector_num_per_node)); + } + + memset(&(multisector_buf[0]), 0, + sector_num_per_node * DiskAnnUtil::kSectorSize); + memset(&(node_buf[0]), 0, max_node_size); + + auto neighbors = get_neighbors(i); + neighbor_num = neighbors.first; + + ailego_assert(neighbor_num > 0); + ailego_assert(neighbor_num <= max_observed_degree_); + + // read node's nhood + memcpy((char *)neighbor_buf, neighbors.second, + neighbor_num * sizeof(diskann_id_t)); + + if (iter->is_valid()) { + const void *vec = iter->data(); + memcpy(&(multisector_buf[0]), vec, meta.element_size()); + + iter->next(); + } else { + return IndexError_Runtime; + } + + // write neighbor + *(uint32_t *)(&(multisector_buf[0]) + meta_.element_size()) = + neighbor_num; + + // write nhood next + memcpy(&(multisector_buf[0]) + meta_.element_size() + sizeof(uint32_t), + neighbor_buf, neighbor_num * sizeof(diskann_id_t)); + + // flush sector to disk + len = dumper->write(multisector_buf.data(), + sector_num_per_node * DiskAnnUtil::kSectorSize); + if (len != sector_num_per_node * DiskAnnUtil::kSectorSize) { + LOG_ERROR("Write Vector Error, write=%zu, expect=%zu", len, + (size_t)(sector_num_per_node * DiskAnnUtil::kSectorSize)); + + return IndexError_WriteData; + } + + write_size += sector_num_per_node * DiskAnnUtil::kSectorSize; + + crc = ailego::Crc32c::Hash(multisector_buf.data(), + sector_num_per_node * DiskAnnUtil::kSectorSize, + crc); + } + + LOG_INFO("Total Sector #%zu written", + (size_t)(doc_cnt * sector_num_per_node)); + + index_size = doc_cnt * sector_num_per_node * DiskAnnUtil::kSectorSize; + } + + size_t padding_size = AlignSize(write_size) - write_size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + ret = dumper->append(kDiskAnnVectorSegmentId, write_size, 0UL, crc); + if (ret != 0) { + LOG_ERROR("Dump vectors segment failed, ret %d", ret); + return ret; + } + + // dump diskann meta + meta_header_.doc_cnt = doc_cnt; + meta_header_.ndims = meta_.dimension(); + meta_header_.medoid = medoid(); + meta_header_.max_node_size = max_node_size; + meta_header_.max_degree = max_observed_degree_; + meta_header_.node_per_sector = node_per_sector; + meta_header_.vamana_frozen_num = 0; + meta_header_.vamana_frozen_loc = medoid(); + meta_header_.append_reorder_data = 0; + meta_header_.index_size = index_size; + + ret = dump_segment(dumper, kDiskAnnMetaSegmentId, &meta_header_, + sizeof(DiskAnnMetaHeader)); + if (ret != 0) { + LOG_ERROR("Dump vectors segment failed"); + + return ret; + } + + // dump pq meta + ret = dump_pq_meta_segment(dumper); + if (ret != 0) { + LOG_ERROR("Dump pq meta segment failed"); + + return ret; + } + + // dump pq data + ret = dump_pq_data_segment(dumper); + if (ret != 0) { + LOG_ERROR("Dump pq data segment failed"); + + return ret; + } + + // dump key + ret = dump_key_segment(dumper); + if (ret != 0) { + LOG_ERROR("Dump key segment failed"); + + return ret; + } + + // dump key mapping + ret = dump_key_mapping_segment(dumper); + if (ret != 0) { + LOG_ERROR("Dump key mapping segment failed"); + + return ret; + } + + // dump entrypint + ret = dump_entrypoint_segment(dumper); + if (ret != 0) { + LOG_ERROR("Dump entrypoint segment failed"); + + return ret; + } + + LOG_INFO("DiskAnn Index File Dumped"); + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_builder_entity.h b/src/core/algorithm/diskann/diskann_builder_entity.h new file mode 100644 index 0000000..d2223ab --- /dev/null +++ b/src/core/algorithm/diskann/diskann_builder_entity.h @@ -0,0 +1,107 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "diskann_entity.h" + +namespace zvec { +namespace core { + +// wrapper class aligned with diskann +class DiskAnnBuilderEntity : public DiskAnnEntity { + public: + using Pointer = std::shared_ptr; + + DiskAnnBuilderEntity() = default; + virtual ~DiskAnnBuilderEntity() = default; + + public: + int add_vector(diskann_key_t key, const void *vec) override; + + std::pair get_neighbors( + diskann_id_t id) const override; + + int set_neighbors(diskann_id_t id, + const std::vector &neighbor_ids) override; + + int add_neighbor(diskann_id_t id, diskann_id_t neighbor_id) override; + + diskann_id_t get_id(diskann_key_t key) const override; + diskann_key_t get_key(diskann_id_t id) const override; + const void *get_vector(diskann_id_t id) const override; + + public: + int init(const IndexMeta &meta, uint32_t max_degree, uint32_t list_size, + double memory_limit, uint32_t build_threads); + + int dump(IndexHolder::Pointer holder, IndexMeta &meta, + const IndexDumper::Pointer &dumper); + + int64_t dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, const void *data, + size_t size) const; + int dump_dummy_segment(const IndexDumper::Pointer &dumper) const; + int dump_pq_meta_segment(const IndexDumper::Pointer &dumper) const; + int dump_pq_data_segment(const IndexDumper::Pointer &dumper) const; + int dump_key_mapping_segment(const IndexDumper::Pointer &dumper) const; + int dump_entrypoint_segment(const IndexDumper::Pointer &dumper) const; + int dump_key_segment(const IndexDumper::Pointer &dumper) const; + + int reserve_space(uint32_t docs); + + std::vector &pq_full_pivot_data() { + return pq_full_pivot_data_; + } + + std::vector &pq_centroid() { + return pq_centroid_; + } + + std::vector &pq_chunk_offsets() { + return pq_chunk_offsets_; + } + + std::vector &block_compressed_data() { + return block_compressed_data_; + } + + private: + uint32_t max_degree_{0}; + uint32_t list_size_{0}; + double memory_limit_{0}; + uint32_t num_threads_{0}; + uint32_t max_build_degree_{0}; + uint32_t max_observed_degree_{0}; + uint32_t neighbor_size_{0}; + + std::string mem_index_file_{""}; + std::string index_path_prefix_{""}; + + std::string vectors_buffer_{}; + std::string keys_buffer_{}; + std::string neighbors_buffer_{}; + std::vector entrypoints_{}; + + IndexMeta meta_; + + std::vector pq_full_pivot_data_; + std::vector pq_centroid_; + std::vector pq_chunk_offsets_; + std::vector block_compressed_data_; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_context.cc b/src/core/algorithm/diskann/diskann_context.cc new file mode 100644 index 0000000..f13affb --- /dev/null +++ b/src/core/algorithm/diskann/diskann_context.cc @@ -0,0 +1,142 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_context.h" +#include +#include "diskann_params.h" +#include "diskann_pq_table.h" +#include "diskann_util.h" + +namespace zvec { +namespace core { + +DiskAnnContext::DiskAnnContext(const IndexMeta &meta, + const IndexMetric::Pointer &measure, + const DiskAnnEntity::Pointer &entity) + : dc_(entity.get(), measure, meta.dimension()), entity_{entity} {} + +int DiskAnnContext::init(ContextType type, uint32_t graph_degree, + uint32_t pq_chunk_num, uint32_t element_size) { + type_ = type; + element_size_ = element_size; + pq_chunk_num_ = pq_chunk_num; + + DiskAnnUtil::alloc_aligned((void **)&query_, element_size_, 32); + DiskAnnUtil::alloc_aligned((void **)&query_rotated_, element_size_, 32); + + int ret; + switch (type) { + case kBuilderContext: + ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(), + entity_->doc_cnt(), negative_probility_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + break; + + case kSearcherContext: + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), + entity_->doc_cnt(), negative_probility_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + + DiskAnnUtil::alloc_aligned( + (void **)&pq_table_dist_buffer_, + PQTable::kPQCentroidNum * pq_chunk_num_ * sizeof(float), 256); + DiskAnnUtil::alloc_aligned((void **)&pq_coord_buffer_, + graph_degree * pq_chunk_num_ * sizeof(uint8_t), + 256); + DiskAnnUtil::alloc_aligned((void **)&coord_buffer_, element_size_, 256); + DiskAnnUtil::alloc_aligned( + (void **)§or_buffer_, + DiskAnnUtil::kMaxSectorReadNum * DiskAnnUtil::kSectorSize, + DiskAnnUtil::kSectorSize); + + ret = setup_io_ctx(io_ctx_); + if (ret != 0) { + LOG_ERROR("setup io ctx error, ret=%d", ret); + return ret; + } + break; + + default: + LOG_ERROR("Init context failed"); + return IndexError_Runtime; + } + + return 0; +} + +DiskAnnContext::~DiskAnnContext() { + free(query_); + free(query_rotated_); + free(pq_table_dist_buffer_); + free(pq_coord_buffer_); + free(coord_buffer_); + free(sector_buffer_); + + if (type_ == kSearcherContext) { + destroy_io_ctx(io_ctx_); + } +} + +int DiskAnnContext::update(const ailego::Params ¶ms) { + uint32_t list_size = list_size_; + params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size); + list_size_ = list_size; + return 0; +} + +int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &measure, + const DiskAnnEntity::Pointer &entity, + uint32_t magic_num) { + if (ailego_unlikely(type != type_)) { + LOG_ERROR( + "DiskAnnContext does not support shared by different type, " + "src=%u dst=%u", + type_, type); + return IndexError_Unsupported; + } + + magic_ = kInvalidMgic; + + switch (type) { + case kBuilderContext: + LOG_ERROR("BuildContext does not support update"); + return IndexError_NotImplemented; + + case kSearcherContext: + break; + + case kReducerContext: + break; + + default: + LOG_ERROR("update context failed"); + return IndexError_Runtime; + } + + entity_ = entity; + dc_.update(measure, meta.dimension()); + magic_ = magic_num; + + return 0; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h new file mode 100644 index 0000000..dd824ff --- /dev/null +++ b/src/core/algorithm/diskann/diskann_context.h @@ -0,0 +1,378 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "utility/topk_result_builder.h" +#include "diskann_dist_calculator.h" +#include "diskann_entity.h" +#include "diskann_file_reader.h" +#include "diskann_visit_filter.h" + +namespace zvec { +namespace core { + +struct SearchStats { + public: + float total_us = 0; + float io_us = 0; + float cpu_us = 0; + uint64_t disk_page_reads = 0; + uint64_t io_num = 0; + uint64_t dist_num = 0; + uint64_t cache_hits = 0; + uint64_t hop_num = 0; +}; + +class DiskAnnContext : public IndexContext, + public TopkResultBuilder { + friend class TopkResultBuilder; + + public: + //! Index Context Pointer + using Pointer = std::unique_ptr; + + enum ContextType { + kUnknownContext = 0, + kSearcherContext = 1, + kBuilderContext = 2, + kReducerContext = 3 + }; + + //! Construct + DiskAnnContext(const IndexMeta &meta, const IndexMetric::Pointer &measure, + const DiskAnnEntity::Pointer &entity); + + //! Destructor + virtual ~DiskAnnContext(); + + public: + //! Init + int init(ContextType type, uint32_t graph_degree, uint32_t pq_chunk_num, + uint32_t element_size); + + //! Update context, the context may be shared by different searcher/streamer + int update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &measure, + const DiskAnnEntity::Pointer &entity, uint32_t magic_num); + + //! Retrieve search result + virtual const IndexDocumentList &result(void) const override { + return results_[0]; + } + + //! Retrieve search result + virtual const IndexDocumentList &result(size_t idx) const override { + return results_[idx]; + } + + //! Retrieve result object for output + virtual IndexDocumentList *mutable_result(size_t idx) override { + ailego_assert_with(idx < results_.size(), "invalid idx"); + return &results_[idx]; + } + + //! Retrieve search group result with index + virtual const IndexGroupDocumentList &group_result(void) const override { + return group_results_[0]; + } + + //! Retrieve search group result with index + virtual const IndexGroupDocumentList &group_result( + size_t idx) const override { + return group_results_[idx]; + } + + virtual IndexGroupDocumentList *mutable_group_result(void) override { + return &group_results_[0]; + } + + virtual IndexGroupDocumentList *mutable_group_result(size_t idx) override { + return &group_results_[idx]; + } + + virtual uint32_t magic(void) const override { + return magic_; + } + + void set_magic(uint32_t magic) { + magic_ = magic; + } + + //! Set mode of debug + virtual void set_debug_mode(bool enable) override { + debug_mode_ = enable; + } + + //! Retrieve mode of debug + virtual bool debug_mode(void) const override { + return debug_mode_; + } + + //! Retrieve string of debug + virtual std::string debug_string(void) const override { + return std::string(""); + } + + //! Update the parameters of context + int update(const ailego::Params ¶ms) override; + + inline DistCalculator &dist_calculator() { + return dc_; + } + + public: + //! Set topk of search result + void set_topk(uint32_t val) override { + topk_ = val; + topk_heap_.limit(val); + } + + void set_list_size(uint32_t list_size) { + list_size_ = list_size; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + + //! Get topk + inline uint32_t topk() const override { + return topk_; + } + + inline uint32_t list_size() const { + return list_size_; + } + + inline void reset_query(const void *query) { + memcpy(query_, query, element_size_); + memcpy(query_rotated_, query, element_size_); + + dc_.reset_query(query); + } + + inline TopkHeap &topk_heap() { + return topk_heap_; + } + + inline void *query() { + return query_; + } + + inline void *query_rotated() { + return query_rotated_; + } + + inline float *pq_table_dist_buffer() { + return pq_table_dist_buffer_; + } + + inline void *pq_coord_buffer() { + return pq_coord_buffer_; + } + + inline void *coord_buffer() { + return coord_buffer_; + } + + inline void *sector_buffer() { + return sector_buffer_; + } + + inline IOContext &io_ctx() { + return io_ctx_; + } + + inline void resize_results(size_t size) { + if (group_by_search()) { + group_results_.resize(size); + } else { + results_.resize(size); + } + } + + inline bool error() const { + return has_error_; + } + + inline void set_error(bool err) { + has_error_ = err; + } + + inline void clear() { + for (auto &it : results_) { + it.clear(); + } + + best_list_nodes_.clear(); + expanded_nodes_.clear(); + visit_filter_.clear(); + has_error_ = false; + } + + SearchStats &query_stats() { + return query_stats_; + } + + const DiskAnnEntity &get_entity() const { + return *entity_; + } + + NeighborPriorityQueue &best_list_nodes() { + return best_list_nodes_; + } + + std::vector &expanded_nodes() { + return expanded_nodes_; + } + + std::vector &occlude_factor() { + return occlude_factor_; + } + + VisitFilter &visit_filter() { + return visit_filter_; + } + + //! Reset context + void reset(void) override { + set_filter(nullptr); + reset_threshold(); + set_fetch_vector(false); + set_group_params(0, 0); + reset_group_by(); + } + + inline std::map &group_topk_heaps() { + return group_topk_heaps_; + } + + //! Get group topk + inline uint32_t group_topk() const { + return group_topk_; + } + + //! Get group num + inline uint32_t group_num() const { + return group_num_; + } + + //! Get if group by search + inline bool group_by_search() { + return group_num_ > 0; + } + + //! Set group params + void set_group_params(uint32_t group_num, uint32_t group_topk) override { + group_num_ = group_num; + group_topk_ = group_topk; + + topk_ = group_topk_ * group_num_; + + topk_heap_.limit(topk_); + + group_topk_heaps_.clear(); + } + + // topk_to_result / topk_to_single_result / topk_to_group_result are provided + // by TopkResultBuilder via the hooks below (result_score, + // emplace_result_doc). + using TopkResultBuilder::topk_to_result; + + void set_to_result(uint32_t idx, const std::vector &result_ids, + const std::vector &result_dists) { + if (result_ids.size() != result_dists.size()) { + return; + } + + uint32_t size = result_ids.size(); + + for (uint32_t i = 0; i < size; ++i) { + results_[idx].emplace_back(result_ids[i], result_dists[i], 0); + } + } + + private: + //! Hooks required by TopkResultBuilder. + std::vector &mutable_results() { + return results_; + } + + std::vector &mutable_group_results() { + return group_results_; + } + + float result_score(const VectorInfo &info) const { + return info.dist_; + } + + void emplace_result_doc(IndexDocumentList &docs, diskann_id_t id, float score, + const VectorInfo &info) { + if (fetch_vector_) { + docs.emplace_back(entity_->get_key(id), score, id, info.vec_); + } else { + docs.emplace_back(entity_->get_key(id), score, id); + } + } + + private: + constexpr static uint32_t kInvalidMgic = -1U; + + uint32_t type_{kUnknownContext}; + + DistCalculator dc_; + DiskAnnEntity::Pointer entity_; + + uint32_t topk_{0}; + uint32_t magic_{0U}; + bool debug_mode_{false}; + uint32_t pq_chunk_num_{0}; + uint32_t element_size_{0}; + uint32_t element_rotated_size_{0}; + uint32_t list_size_{0}; + + TopkHeap topk_heap_{}; + + uint32_t group_topk_{0}; + uint32_t group_num_{0}; + std::map group_topk_heaps_{}; + + IOContext io_ctx_{0}; + SearchStats query_stats_; + + float *pq_table_dist_buffer_{nullptr}; + void *pq_coord_buffer_{nullptr}; + void *query_{nullptr}; + void *query_rotated_{nullptr}; + void *coord_buffer_{nullptr}; + void *sector_buffer_{nullptr}; + + std::vector results_{}; + std::vector group_results_{}; + + bool fetch_vector_{false}; + bool has_error_{false}; + + NeighborPriorityQueue best_list_nodes_; + std::vector expanded_nodes_; + std::vector occlude_factor_; + + VisitFilter visit_filter_{}; + uint32_t filter_mode_{VisitFilter::ByteMap}; + float negative_probility_{DiskAnnEntity::kDefaultBFNegativeProbility}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_dist_calculator.h b/src/core/algorithm/diskann/diskann_dist_calculator.h new file mode 100644 index 0000000..e64b6f3 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_dist_calculator.h @@ -0,0 +1,150 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "diskann_entity.h" + +namespace zvec { +namespace core { + +class DistCalculator { + public: + typedef std::shared_ptr Pointer; + + public: + //! Constructor + DistCalculator(const DiskAnnEntity *entity, + const IndexMetric::Pointer &measure, uint32_t dim) + : entity_(entity), + distance_(measure->distance()), + query_(nullptr), + dim_(dim), + compare_cnt_(0) {} + + void update(const IndexMetric::Pointer &measure, uint32_t dim) { + distance_ = measure->distance(); + dim_ = dim; + } + + inline void update_distance(const IndexMetric::MatrixDistance &distance) { + distance_ = distance; + } + + //! Reset query vector data + inline void reset_query(const void *query) { + error_ = false; + query_ = query; + } + + //! Returns distance + inline dist_t dist(const void *vec_lhs, const void *vec_rhs) { + if (ailego_unlikely(vec_lhs == nullptr || vec_rhs == nullptr)) { + LOG_ERROR("Nullptr of dense vector"); + + error_ = true; + return 0.0f; + } + + float score{0.0f}; + distance_(vec_lhs, vec_rhs, dim_, &score); + + return score; + } + + //! Returns distance between query and vec. + inline dist_t dist(const void *vec) { + compare_cnt_++; + + return dist(vec, query_); + } + + inline dist_t dist(diskann_id_t id) { + compare_cnt_++; + + const void *vec = entity_->get_vector(id); + if (ailego_unlikely(vec == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + + return dist(vec, query_); + } + + inline dist_t dist(diskann_id_t lhs, diskann_id_t rhs) { + compare_cnt_++; + + const void *vec_lhs = entity_->get_vector(lhs); + if (ailego_unlikely(vec_lhs == nullptr)) { + LOG_ERROR("Get nullptr vector, lhs id=%u", lhs); + error_ = true; + return 0.0f; + } + + const void *vec_rhs = entity_->get_vector(rhs); + if (ailego_unlikely(vec_rhs == nullptr)) { + LOG_ERROR("Get nullptr vector, rhs id=%u", rhs); + error_ = true; + return 0.0f; + } + + return dist(vec_lhs, vec_rhs); + } + + dist_t operator()(const void *vec) { + return dist(vec); + } + + inline void clear() { + compare_cnt_ = 0; + error_ = false; + } + + inline void clear_compare_cnt() { + compare_cnt_ = 0; + } + + inline bool error() const { + return error_; + } + + //! Get distances compute times + inline uint32_t compare_cnt() const { + return compare_cnt_; + } + + inline uint32_t dimension() const { + return dim_; + } + + private: + DistCalculator(const DistCalculator &) = delete; + DistCalculator &operator=(const DistCalculator &) = delete; + + private: + const DiskAnnEntity *entity_; + + IndexMetric::MatrixDistance distance_; + const void *query_; + uint32_t dim_; + + uint32_t compare_cnt_; // record distance compute times + bool error_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_entity.cc b/src/core/algorithm/diskann/diskann_entity.cc new file mode 100644 index 0000000..56d63fc --- /dev/null +++ b/src/core/algorithm/diskann/diskann_entity.cc @@ -0,0 +1,32 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_entity.h" + +namespace zvec { +namespace core { + +const std::string DiskAnnEntity::kDiskAnnVectorSegmentId = "diskann.vector"; +const std::string DiskAnnEntity::kDiskAnnMetaSegmentId = "diskann.meta"; +const std::string DiskAnnEntity::kDiskAnnPqMetaSegmentId = "diskann.pq_meta"; +const std::string DiskAnnEntity::kDiskAnnPqDataSegmentId = "diskann.pq_data"; +const std::string DiskAnnEntity::kDiskAnnDummpySegmentId = "diskann.dummy"; +const std::string DiskAnnEntity::kDiskAnnKeyMappingSegmentId = + "diskann.key_mapping"; +const std::string DiskAnnEntity::kDiskAnnEntryPointSegmentId = + "diskann.entrypoint"; +const std::string DiskAnnEntity::kDiskAnnKeySegmentId = "diskann.key"; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_entity.h b/src/core/algorithm/diskann/diskann_entity.h new file mode 100644 index 0000000..af30229 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_entity.h @@ -0,0 +1,244 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include + +namespace zvec { +namespace core { + +using dist_t = float; +using diskann_key_t = uint64_t; +using diskann_id_t = uint32_t; + +constexpr diskann_id_t kInvalidId = static_cast(-1); +constexpr diskann_key_t kInvalidKey = static_cast(-1); + +struct VectorInfo { + float dist_; + std::string vec_; + + VectorInfo() = default; + VectorInfo(float dist, const std::string &vec) : dist_{dist}, vec_{vec} {} + VectorInfo(float dist, std::string &&vec) + : dist_{dist}, vec_{std::move(vec)} {} +}; + +/*! Key Value Vecotr Heap Comparer + */ +struct KeyValueVectorHeapComparer { + //! Function call + bool operator()(const std::pair &lhs, + const std::pair &rhs) const { + return compare_(lhs.second.dist_, rhs.second.dist_); + } + + private: + std::less compare_; +}; + +/*! Key Value Vector Heap + */ +using TopkHeap = ailego::Heap, + KeyValueVectorHeapComparer>; + +struct DiskAnnMetaHeader { + public: + uint64_t doc_cnt; + uint64_t ndims; + uint64_t medoid; + uint64_t max_node_size; + uint64_t max_degree; + uint64_t node_per_sector; + uint64_t vamana_frozen_num; + uint64_t vamana_frozen_loc; + uint64_t append_reorder_data; + uint64_t index_size; + uint8_t reserved[4016]; /// pad DiskAnnMetaHeader to 4096 bytes + + DiskAnnMetaHeader() { + clear(); + } + + DiskAnnMetaHeader(const DiskAnnMetaHeader &header) { + memcpy(this, &header, sizeof(header)); + } + + DiskAnnMetaHeader &operator=(const DiskAnnMetaHeader &header) { + if (this != &header) { + memcpy(this, &header, sizeof(header)); + } + return *this; + } + + void reset() { + doc_cnt = 0U; + } + + void clear() { + memset(this, 0, sizeof(DiskAnnMetaHeader)); + } +}; + +struct DiskAnnPqMeta { + public: + uint64_t full_pivot_data_size{0}; + uint64_t centroid_data_size{0}; + uint64_t chunk_offsets_size{0}; + uint64_t chunk_num{0}; + uint8_t reserved[128]; + + DiskAnnPqMeta() { + clear(); + } + + DiskAnnPqMeta(const DiskAnnPqMeta &meta) { + memcpy(this, &meta, sizeof(meta)); + } + + DiskAnnPqMeta &operator=(const DiskAnnPqMeta &meta) { + memcpy(this, &meta, sizeof(meta)); + return *this; + } + + void clear() { + memset(this, 0, sizeof(DiskAnnPqMeta)); + } +}; + +static_assert(sizeof(DiskAnnMetaHeader) == 4096, + "DiskAnnMetaHeader size must stay 4096 bytes (on-disk format)"); + +static_assert(sizeof(DiskAnnPqMeta) % 32 == 0, + "DiskAnnPqMeta size must be a multiple of 32 bytes"); + +class DiskAnnEntity { + public: + DiskAnnEntity() = default; + virtual ~DiskAnnEntity() = default; + + //! Constructor + DiskAnnEntity(const DiskAnnMetaHeader &meta_header, + const DiskAnnPqMeta &pq_meta) { + meta_header_ = meta_header; + pq_meta_ = pq_meta; + } + + //! DiskAnnEntity Pointerd; + typedef std::shared_ptr Pointer; + + public: + static size_t AlignSize(size_t size) { + return (size + 0xFFF) & (~0xFFF); + } + + public: + virtual int add_vector(diskann_key_t /*key*/, const void * /*vec*/) { + return IndexError_NotImplemented; + } + + virtual const void *get_vector(diskann_id_t /*id*/) const { + return nullptr; + } + + virtual std::pair get_neighbors( + diskann_id_t /*id*/) const { + return std::make_pair(0, nullptr); + } + + virtual int set_neighbors( + diskann_id_t /*id*/, const std::vector & /*neighbor_ids*/) { + return IndexError_NotImplemented; + } + + virtual int add_neighbor(diskann_id_t /*id*/, diskann_id_t /*neighbor_id*/) { + return IndexError_NotImplemented; + } + + //! Get node id of primary key + virtual diskann_id_t get_id(diskann_key_t key) const = 0; + + //! Get primary key of the node id + virtual diskann_key_t get_key(diskann_id_t id) const = 0; + + public: + uint64_t max_node_size() const { + return meta_header_.max_node_size; + } + + uint64_t medoid() const { + return meta_header_.medoid; + } + + uint64_t *mutable_medoid() { + return &meta_header_.medoid; + } + + uint64_t node_per_sector() const { + return meta_header_.node_per_sector; + } + + uint64_t pq_chunk_num() { + return pq_meta_.chunk_num; + } + + uint64_t doc_cnt() const { + return meta_header_.doc_cnt; + } + + uint64_t *mutable_doc_cnt() { + return &meta_header_.doc_cnt; + } + + uint64_t max_degree() { + return meta_header_.max_degree; + } + + DiskAnnPqMeta *mutable_pq_meta() { + return &pq_meta_; + } + + public: + virtual const DiskAnnEntity::Pointer clone() const { + LOG_ERROR("Update neighbors not implemented"); + return DiskAnnEntity::Pointer(); + } + + public: + const static std::string kDiskAnnVectorSegmentId; + const static std::string kDiskAnnMetaSegmentId; + const static std::string kDiskAnnPqMetaSegmentId; + const static std::string kDiskAnnPqDataSegmentId; + const static std::string kDiskAnnDummpySegmentId; + const static std::string kDiskAnnMappingSegmentId; + const static std::string kDiskAnnKeyMappingSegmentId; + const static std::string kDiskAnnEntryPointSegmentId; + const static std::string kDiskAnnKeySegmentId; + + constexpr static float kDefaultBFNegativeProbility = 0.001f; + constexpr static float kDefaultGraphSlackFactor = 1.3f; + constexpr static float kDefaultAlpha = 1.2f; + constexpr static uint32_t kDefaultMaxOcclusionSize = 750; + constexpr static uint32_t kDefaultMaxDegree = 100; + constexpr static uint32_t kDefaultCompressBatchSize = 5000000; + constexpr static uint32_t kRevision = 0U; + + protected: + DiskAnnMetaHeader meta_header_; + DiskAnnPqMeta pq_meta_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc new file mode 100644 index 0000000..4268d51 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -0,0 +1,299 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_file_reader.h" +#include +#include +#include +#include +#include +#include + +#define MAX_EVENTS 1024 + +namespace zvec { +namespace core { + +#if (defined(__linux) || defined(__linux__)) +typedef struct io_event io_event_t; +typedef struct iocb iocb_t; +#endif + +int setup_io_ctx(IOContext &ctx) { +#if (defined(__linux) || defined(__linux__)) + int ret = io_setup(MAX_EVENTS, &ctx); + + return ret; +#else + return 0; +#endif +} + +int destroy_io_ctx(IOContext &ctx) { +#if (defined(__linux) || defined(__linux__)) + int ret = io_destroy(ctx); + + return ret; +#else + return 0; +#endif +} + +static int execute_io_pread(int fd, std::vector &read_reqs) { + for (auto &req : read_reqs) { + ssize_t bytes_read = ::pread(fd, req.buf, req.len, req.offset); + if (bytes_read < 0) { + LOG_ERROR("pread failed; errno=%d, %s, offset=%lu, len=%lu", errno, + ::strerror(errno), (unsigned long)req.offset, + (unsigned long)req.len); + return IndexError_Runtime; + } + if ((size_t)bytes_read != req.len) { + LOG_ERROR("pread short read; got=%zd, expected=%lu", bytes_read, + (unsigned long)req.len); + return IndexError_Runtime; + } + } + return 0; +} + +int execute_io(IOContext ctx, int fd, std::vector &read_reqs, + uint64_t n_retries = 0) { +#if (defined(__linux) || defined(__linux__)) + uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), MAX_EVENTS); + + for (uint64_t iter = 0; iter < iters; iter++) { + uint64_t n_ops = std::min((uint64_t)read_reqs.size() - (iter * MAX_EVENTS), + (uint64_t)MAX_EVENTS); + + std::vector cbs(n_ops, nullptr); + std::vector evts(n_ops); + std::vector cb(n_ops); + for (uint64_t j = 0; j < n_ops; j++) { + io_prep_pread(cb.data() + j, fd, read_reqs[j + iter * MAX_EVENTS].buf, + read_reqs[j + iter * MAX_EVENTS].len, + read_reqs[j + iter * MAX_EVENTS].offset); + } + + for (uint64_t i = 0; i < n_ops; i++) { + cbs[i] = cb.data() + i; + } + + size_t n_tries = 0; + // Phase 1: io_submit with retry. + while (true) { + int ret = io_submit(ctx, (int64_t)n_ops, cbs.data()); + if (ret == (int)n_ops) { + break; + } + if ((ret == -EAGAIN || ret == -EINTR) && n_tries < n_retries) { + n_tries++; + continue; + } + LOG_WARN( + "io_submit failed; returned: %d, expected=%lu. falling back to " + "pread", + ret, n_ops); + return execute_io_pread(fd, read_reqs); + } + + // Phase 2: io_getevents with retry (never re-submits). + n_tries = 0; + while (true) { + int ret = io_getevents(ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), + nullptr); + if (ret == (int)n_ops) { + break; + } + if (ret == -EINTR && n_tries < n_retries) { + n_tries++; + continue; + } + LOG_WARN( + "io_getevents failed; returned: %d, expected=%lu, errno=%d, %s, " + "falling back to pread", + ret, n_ops, errno, ::strerror(-ret)); + return execute_io_pread(fd, read_reqs); + } + + // Phase 3: verify each completed read (res must equal requested length). + bool all_ok = true; + for (uint64_t i = 0; i < n_ops; i++) { + int64_t expected_len = read_reqs[i + iter * MAX_EVENTS].len; + if ((int64_t)evts[i].res != expected_len) { + LOG_WARN("aio request %zu failed: res=%ld, expected=%ld, offset=%zu", + (size_t)i, (long)evts[i].res, (long)expected_len, + (size_t)read_reqs[i + iter * MAX_EVENTS].offset); + all_ok = false; + } + } + if (!all_ok) { + return execute_io_pread(fd, read_reqs); + } + } + + return 0; +#else + return execute_io_pread(fd, read_reqs); +#endif +} + +LinuxAlignedFileReader::LinuxAlignedFileReader(int file_desc) { + this->file_desc = file_desc; +} + +LinuxAlignedFileReader::LinuxAlignedFileReader() { + this->file_desc = -1; +} + +LinuxAlignedFileReader::~LinuxAlignedFileReader() { + deregister_all_threads(); + if (file_desc >= 0) { + ::close(file_desc); + file_desc = -1; + } +} + +IOContext &LinuxAlignedFileReader::get_ctx() { + std::unique_lock lk(ctx_mut); + auto it = ctx_map.find(std::this_thread::get_id()); + if (it == ctx_map.end()) { + LOG_ERROR("bad thread access; returning -1 as io_context_t"); + return this->bad_ctx; + } else { + return it->second; + } +} + +void LinuxAlignedFileReader::register_thread() { +#if (defined(__linux) || defined(__linux__)) + auto thread_id = std::this_thread::get_id(); + std::unique_lock lk(ctx_mut); + if (ctx_map.find(thread_id) != ctx_map.end()) { + LOG_ERROR("multiple calls to register_thread from the same thread"); + + return; + } + + IOContext ctx = nullptr; + + int ret = io_setup(MAX_EVENTS, &ctx); + if (ret != 0) { + lk.unlock(); + if (ret == -EAGAIN) { + LOG_ERROR( + "io_setup failed with EAGAIN: Consider increasing " + "/proc/sys/fs/aio-max-nr"); + } else { + LOG_ERROR("io_setup failed; returned: %d, %s", ret, ::strerror(-ret)); + ; + } + } else { + LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); + + ctx_map[thread_id] = ctx; + } + + lk.unlock(); +#endif +} + +void LinuxAlignedFileReader::deregister_thread() { +#if (defined(__linux) || defined(__linux__)) + auto thread_id = std::this_thread::get_id(); + IOContext ctx; + + { + std::lock_guard lk(ctx_mut); + auto it = ctx_map.find(thread_id); + if (it == ctx_map.end()) { + LOG_ERROR("deregister_thread: thread not registered"); + return; + } + ctx = it->second; + ctx_map.erase(it); + } + + // io_destroy is a syscall; keep it outside the lock to avoid blocking others + io_destroy(ctx); + LOG_INFO("returned ctx from thread"); +#endif +} + +void LinuxAlignedFileReader::deregister_all_threads() { +#if (defined(__linux) || defined(__linux__)) + std::unique_lock lk(ctx_mut); + for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { + IOContext ctx = x->second; + io_destroy(ctx); + } + ctx_map.clear(); +#endif +} + +void LinuxAlignedFileReader::open(const std::string &fname) { + int flags = O_RDONLY; + +#if defined(__linux__) || defined(__linux) + flags |= O_DIRECT | O_LARGEFILE; +#endif + + this->file_desc = ::open(fname.c_str(), flags); + +#if defined(__linux__) || defined(__linux) + // O_DIRECT may not be supported on all filesystems (e.g. tmpfs, overlay). + // Fall back to regular buffered I/O when it fails. + if (this->file_desc == -1) { + LOG_WARN( + "open with O_DIRECT failed for %s (errno=%d: %s), " + "falling back to buffered I/O", + fname.c_str(), errno, ::strerror(errno)); + this->file_desc = ::open(fname.c_str(), O_RDONLY | O_LARGEFILE); + } +#endif + + if (this->file_desc == -1) { + LOG_ERROR("Failed to open file: %s (errno=%d: %s)", fname.c_str(), errno, + ::strerror(errno)); + } + + LOG_INFO("Opened file : %s", fname.c_str()); +} + +void LinuxAlignedFileReader::close() { + if (file_desc >= 0) { + ::close(file_desc); + file_desc = -1; + } +} + +int LinuxAlignedFileReader::read(std::vector &read_reqs, + IOContext &ctx, bool async) { + if (async == true) { + LOG_WARN("Async currently not supported"); + } + + if (this->file_desc == -1) { + LOG_ERROR("Attempt to read from invalid file descriptor"); + return IndexError_Runtime; + } + + int ret = execute_io(ctx, this->file_desc, read_reqs); + + return ret; +} + + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h new file mode 100644 index 0000000..432247e --- /dev/null +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -0,0 +1,103 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#define MAX_IO_DEPTH 128 + +#include + +#if (defined(__linux) || defined(__linux__)) +#include +#endif + +#include +#include +#include +#include +#include "diskann_util.h" + +namespace zvec { +namespace core { + +#if (defined(__linux) || defined(__linux__)) +typedef io_context_t IOContext; +#else +typedef uint32_t IOContext; +#endif + +int setup_io_ctx(IOContext &ctx); +int destroy_io_ctx(IOContext &ctx); + +struct AlignedRead { + uint64_t offset; + uint64_t len; + void *buf; + + AlignedRead() : offset(0), len(0), buf(nullptr) {} + + AlignedRead(uint64_t offset, uint64_t len, void *buf) + : offset(offset), len(len), buf(buf) { + ailego_assert(static_cast(offset) % 512 == 0); + ailego_assert(static_cast(len) % 512 == 0); + ailego_assert(reinterpret_cast(buf) % 512 == 0); + } +}; + +class AlignedFileReader { + protected: + std::map ctx_map; + std::mutex ctx_mut; + + public: + virtual IOContext &get_ctx() = 0; + + virtual ~AlignedFileReader() {} + + virtual void register_thread() = 0; + virtual void deregister_thread() = 0; + virtual void deregister_all_threads() = 0; + + virtual void open(const std::string &fname) = 0; + virtual void close() = 0; + + virtual int read(std::vector &read_reqs, IOContext &ctx, + bool async = false) = 0; +}; + +class LinuxAlignedFileReader : public AlignedFileReader { + private: + int file_desc; + + IOContext bad_ctx = (IOContext)-1; + + public: + LinuxAlignedFileReader(); + LinuxAlignedFileReader(int file_desc); + ~LinuxAlignedFileReader(); + + public: + IOContext &get_ctx(); + + void register_thread(); + void deregister_thread(); + void deregister_all_threads(); + void open(const std::string &fname); + void close(); + + int read(std::vector &read_reqs, IOContext &ctx, + bool async = false); +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_holder.h b/src/core/algorithm/diskann/diskann_holder.h new file mode 100644 index 0000000..beb422b --- /dev/null +++ b/src/core/algorithm/diskann/diskann_holder.h @@ -0,0 +1,325 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "diskann_entity.h" + +namespace zvec { +namespace core { + +struct DiskAnnIndexHolderMeta { + uint32_t element_size_; + uint32_t key_size_; + uint32_t sector_size_; + uint32_t doc_cnt_; + uint8_t reserve_[]; +}; + +class DiskAnnIndexHolder : public IndexHolder { + public: + typedef std::shared_ptr Pointer; + + public: + enum Status { STATUS_UNINITED = 0, STATUS_WRITE = 1, STATUS_READ = 2 }; + + public: + static constexpr uint32_t kDataSectorSize = 128 * 1024; + static constexpr uint32_t kMetaSectorSize = 4096; + + public: + inline static uint32_t get_sector_id(uint32_t id, uint32_t sector_vec_num) { + return id / sector_vec_num; + } + + inline static uint32_t get_sector_offset(uint32_t id, uint32_t sector_vec_num, + uint32_t data_size) { + return (id % sector_vec_num) * data_size; + } + + public: + /*! Random Access Index Holder Iterator + */ + class Iterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(DiskAnnIndexHolder *owner) + : holder_(owner), sector_id_{0}, sector_offset_{0} { + path_ = holder_->path(); + data_size_ = holder_->data_size(); + data_sector_size_ = holder_->data_sector_size(); + meta_sector_size_ = holder_->meta_sector_size(); + + sector_buffer_.resize(data_sector_size_); + + sector_vec_num_ = data_sector_size_ / data_size_; + } + + //! Destructor + virtual ~Iterator(void) { + if (file_.is_open()) { + file_.close(); + } + } + + int init() { + file_.open(path_, std::ios::in); + if (!file_.is_open()) { + LOG_ERROR("file can not create, %s", path_.c_str()); + return IndexError_OpenFile; + } + + file_.seekg(meta_sector_size_); + + read_sector(); + + return 0; + } + + //! Retrieve pointer of data + const void *data(void) const override { + const uint8_t *data_ptr = + reinterpret_cast(sector_buffer_.data()); + return data_ptr + sector_offset_ + sizeof(diskann_key_t); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return id_ < holder_->count(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + const uint8_t *data_ptr = + reinterpret_cast(sector_buffer_.data()); + uint64_t key = + *reinterpret_cast(data_ptr + sector_offset_); + + return key; + } + + //! Next iterator + void next(void) override { + ++id_; + + uint32_t sector_id = get_sector_id(id_, sector_vec_num_); + if (sector_id > sector_id_) { + file_.seekg(sector_id * data_sector_size_ + meta_sector_size_); + read_sector(); + sector_id_ = sector_id; + } + + sector_offset_ = get_sector_offset(id_, sector_vec_num_, data_size_); + } + + int read_sector() { + file_.read(&((sector_buffer_)[0]), data_sector_size_); + if (!file_) { + LOG_ERROR("Failed to read sector from file: %s", path_.c_str()); + return IndexError_ReadData; + } + + return 0; + } + + private: + //! Members + DiskAnnIndexHolder *holder_{nullptr}; + std::string path_; + std::ifstream file_; + uint32_t sector_id_{0}; + std::string sector_buffer_; + uint32_t id_{0}; + uint32_t data_size_{0}; + uint32_t sector_offset_{0}; + uint32_t data_sector_size_{0}; + uint32_t meta_sector_size_{0}; + uint32_t sector_vec_num_{0}; + }; + + public: + DiskAnnIndexHolder(IndexMeta &meta, std::string &path) { + path_ = path; + + data_size_ = meta.element_size() + sizeof(diskann_key_t); + dimension_ = meta.dimension(); + type_ = meta.data_type(); + + element_size_ = meta.element_size(); + sector_vec_num_ = data_sector_size_ / data_size_; + padding_size_ = data_sector_size_ - sector_vec_num_ * data_size_; + sector_buffer_.resize(data_sector_size_); + sector_internal_id_ = 0; + } + + ~DiskAnnIndexHolder() override { + if (file_.is_open()) { + file_.close(); + } + } + + //! Init + int init() { + file_.open(path_, std::ios::out | std::ios::trunc); + + if (!file_.is_open()) { + LOG_ERROR("file can not create, %s", path_.c_str()); + return IndexError_OpenFile; + } + + DiskAnnIndexHolderMeta holder_meta; + holder_meta.element_size_ = element_size_; + holder_meta.key_size_ = sizeof(diskann_key_t); + holder_meta.sector_size_ = data_sector_size_; + + std::vector empty_sector; + empty_sector.resize(meta_sector_size_); + + std::memset(&(empty_sector[0]), 0, meta_sector_size_); + std::memcpy(&(empty_sector[0]), &holder_meta, + sizeof(DiskAnnIndexHolderMeta)); + + file_.write(reinterpret_cast(&(empty_sector[0])), + meta_sector_size_); + if (!file_) { + LOG_ERROR("Failed to write meta sector to file: %s", path_.c_str()); + return IndexError_WriteData; + } + + status_ = STATUS_WRITE; + + return 0; + } + + int close() { + if (sector_internal_id_ != 0) { + file_.write(reinterpret_cast(&(sector_buffer_[0])), + data_sector_size_); + } + + file_.close(); + + return 0; + } + + //! Retrieve count of elements in holder (-1 indicates unknown) + size_t count(void) const override { + return count_; + } + + //! Retrieve dimension + size_t dimension(void) const override { + return dimension_; + } + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return type_; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return element_size_; + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return true; + } + + //! Create a new iterator + IndexHolder::Iterator::Pointer create_iterator(void) override { + auto pointer = std::make_unique(this); + + if (pointer->init() != 0) { + return nullptr; + } + + return pointer; + } + + int emplace(uint64_t pkey, const void *vec) { + if (status_ != STATUS_WRITE) { + return IndexError_NoReady; + } + + uint8_t *data_ptr = reinterpret_cast(&(sector_buffer_[0])) + + sector_internal_id_ * data_size_; + std::memcpy(data_ptr, &pkey, sizeof(diskann_key_t)); + std::memcpy(data_ptr + sizeof(diskann_key_t), vec, element_size_); + + sector_internal_id_++; + if (sector_internal_id_ >= sector_vec_num_) { + std::memset(data_ptr + data_size_, 0, padding_size_); + + file_.write(reinterpret_cast(&(sector_buffer_[0])), + data_sector_size_); + + sector_internal_id_ = 0; + sector_id_++; + } + + count_++; + + return 0; + } + + uint32_t data_sector_size() { + return data_sector_size_; + } + + uint32_t meta_sector_size() { + return meta_sector_size_; + } + + uint32_t data_size() { + return data_size_; + } + + uint32_t *mutable_sector_id() { + return §or_id_; + } + + uint32_t sector_id() { + return sector_id_; + } + + std::string &path() { + return path_; + } + + private: + std::string path_; + std::ofstream file_; + uint32_t element_size_{0}; + uint32_t dimension_{0}; + IndexMeta::DataType type_{IndexMeta::DataType::DT_UNDEFINED}; + uint32_t sector_vec_num_{0}; + uint32_t data_size_{0}; + uint32_t padding_size_{0}; + uint32_t meta_sector_size_{DiskAnnIndexHolder::kMetaSectorSize}; + uint32_t data_sector_size_{DiskAnnIndexHolder::kDataSectorSize}; + std::string sector_buffer_; + uint32_t sector_internal_id_{0}; + uint32_t sector_id_{0}; + uint32_t count_{0}; + uint32_t status_{STATUS_UNINITED}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_index_provider.h b/src/core/algorithm/diskann/diskann_index_provider.h new file mode 100644 index 0000000..1fd8754 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_index_provider.h @@ -0,0 +1,120 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "diskann_entity.h" + +namespace zvec { +namespace core { + +//! IndexProvider implementation backed by a DiskAnn entity. +//! +//! Used by ``MixedStreamerReducer`` during segment merge: the reducer needs +//! to walk every vector held by a source DiskAnn streamer and feed it into +//! the merge target. Vectors are read on demand from the entity's on-disk +//! vector segment via ``DiskAnnEntity::get_vector(id)``. +class DiskAnnIndexProvider : public IndexProvider { + public: + DiskAnnIndexProvider(const IndexMeta &meta, + const DiskAnnEntity::Pointer &entity, + const std::string &owner) + : meta_(meta), entity_(entity), owner_class_(owner) {} + + DiskAnnIndexProvider(const DiskAnnIndexProvider &) = delete; + DiskAnnIndexProvider &operator=(const DiskAnnIndexProvider &) = delete; + + public: + IndexProvider::Iterator::Pointer create_iterator() override { + return IndexProvider::Iterator::Pointer(new (std::nothrow) + Iterator(entity_)); + } + + size_t count(void) const override { + return entity_->doc_cnt(); + } + + size_t dimension(void) const override { + return meta_.dimension(); + } + + IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + size_t element_size(void) const override { + return meta_.element_size(); + } + + const void *get_vector(uint64_t key) const override { + diskann_id_t id = entity_->get_id(static_cast(key)); + if (id == kInvalidId) { + return nullptr; + } + return entity_->get_vector(id); + } + + const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class Iterator : public IndexProvider::Iterator { + public: + explicit Iterator(const DiskAnnEntity::Pointer &entity) + : entity_(entity), cur_id_(0U) { + cur_id_ = next_valid_id(0U); + } + + const void *data(void) const override { + return entity_->get_vector(cur_id_); + } + + bool is_valid(void) const override { + return cur_id_ < static_cast(entity_->doc_cnt()); + } + + uint64_t key(void) const override { + return static_cast(entity_->get_key(cur_id_)); + } + + void next(void) override { + cur_id_ = next_valid_id(cur_id_ + 1); + } + + private: + //! Skip ids that map to ``kInvalidKey`` (deleted / never populated slots). + diskann_id_t next_valid_id(diskann_id_t start_id) const { + const auto total = static_cast(entity_->doc_cnt()); + for (diskann_id_t i = start_id; i < total; ++i) { + if (entity_->get_key(i) != kInvalidKey) { + return i; + } + } + return total; + } + + DiskAnnEntity::Pointer entity_; + diskann_id_t cur_id_; + }; + + IndexMeta meta_; + DiskAnnEntity::Pointer entity_; + std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc new file mode 100644 index 0000000..32e7ff6 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -0,0 +1,1205 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_indexer.h" +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +DiskAnnIndexer::DiskAnnIndexer(const IndexMeta &meta) { + meta_ = meta; +} + +DiskAnnIndexer::~DiskAnnIndexer() { + destroy_io_ctx(init_ctx_); + if (centroid_data_) { + free(centroid_data_); + } + DiskAnnUtil::free_aligned(coord_cache_buf_); +} + +int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { + entity_ = &entity; + + auto storage = entity.get_storage(); + auto vector_segment = entity.get_vector_segment(); + + pq_table_ = entity.get_pq_table(); + + index_segment_offset_ = vector_segment->data_offset(); + + reader_.reset(new LinuxAlignedFileReader()); + + auto file_path = storage->file_path(); + reader_->open(file_path); + + storage->cleanup(); + + int ret = setup_io_ctx(init_ctx_); + if (ret != 0) { + LOG_ERROR("setup io ctx error"); + return ret; + } + + max_node_size_ = entity.max_node_size(); + disk_bytes_per_point_ = meta_.element_size(); + + node_per_sector_ = entity.node_per_sector(); + aligned_dim_ = meta_.dimension(); + + pq_chunk_num_ = entity.pq_chunk_num(); + + medoid_ = entity.medoid(); + + entrypoints_.push_back(medoid_); + auto &entrypoints = entity.entrypoints(); + for (size_t i = 0; i < entrypoints.size(); ++i) { + entrypoints_.push_back(entrypoints[i]); + } + + doc_cnt_ = entity.doc_cnt(); + + max_degree_ = entity.max_degree(); + + sector_num_per_node_ = + DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); + if (beam_width_ > sector_num_per_node_ * DiskAnnUtil::kMaxSectorReadNum) { + LOG_ERROR("Beamwidth can not be higher than kMaxSectorReadNum"); + + return IndexError_InvalidArgument; + } + + DiskAnnUtil::alloc_aligned((void **)(¢roid_data_), + entrypoints_.size() * aligned_dim_ * sizeof(float), + 32); + + use_medroids_data_as_centroids(); + + return 0; +} + +int DiskAnnIndexer::use_medroids_data_as_centroids() { + LOG_INFO("Loading centroid data from medoid vector data"); + + std::vector nodes_to_read; + std::vector medoid_bufs; + std::vector> neighbor_bufs; + + std::vector centroid_buffer; + + size_t dim = meta_.dimension(); + centroid_buffer.resize(dim); + + nodes_to_read.push_back(medoid_); + medoid_bufs.push_back(&(centroid_buffer[0])); + neighbor_bufs.emplace_back(0, nullptr); + + auto read_status = read_nodes(nodes_to_read, medoid_bufs, neighbor_bufs); + + if (read_status[0] == true) { + for (uint32_t i = 0; i < dim; i++) centroid_data_[i] = centroid_buffer[i]; + } else { + LOG_ERROR("Failed to read medoid"); + return IndexError_Runtime; + } + + return 0; +} + +diskann_key_t DiskAnnIndexer::get_key(diskann_id_t id) const { + return entity_->get_key(id); +} + +diskann_id_t DiskAnnIndexer::get_id(diskann_key_t key) const { + return entity_->get_id(key); +} + +std::vector DiskAnnIndexer::read_nodes( + const std::vector &node_ids, + std::vector &coord_buffers, + std::vector> &neighbor_buffers) { + std::vector read_reqs; + std::vector retval(node_ids.size(), true); + + uint8_t *buf = nullptr; + auto sector_num = + node_per_sector_ > 0 + ? 1 + : DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); + DiskAnnUtil::alloc_aligned( + (void **)&buf, node_ids.size() * sector_num * DiskAnnUtil::kSectorSize, + DiskAnnUtil::kSectorSize); + + for (size_t i = 0; i < node_ids.size(); ++i) { + auto node_id = node_ids[i]; + + AlignedRead read; + read.len = sector_num * DiskAnnUtil::kSectorSize; + read.buf = buf + i * sector_num * DiskAnnUtil::kSectorSize; + read.offset = + index_segment_offset_ + + DiskAnnUtil::get_node_sector(node_per_sector_, max_node_size_, + DiskAnnUtil::kSectorSize, node_id) * + DiskAnnUtil::kSectorSize; + read_reqs.push_back(read); + } + + int read_ret = reader_->read(read_reqs, init_ctx_); + if (read_ret != 0) { + LOG_ERROR("read_nodes: reader_->read failed, ret=%d", read_ret); + for (size_t i = 0; i < retval.size(); i++) { + retval[i] = false; + } + DiskAnnUtil::free_aligned(buf); + return retval; + } + + for (uint32_t i = 0; i < read_reqs.size(); i++) { + uint8_t *node_buf = + DiskAnnUtil::offset_to_node(node_per_sector_, max_node_size_, + (uint8_t *)read_reqs[i].buf, node_ids[i]); + + if (coord_buffers[i] != nullptr) { + void *node_coords = node_buf; + memcpy(coord_buffers[i], node_coords, disk_bytes_per_point_); + } + + if (neighbor_buffers[i].second != nullptr) { + uint32_t *node_neighbor = + DiskAnnUtil::offset_to_node_neighbor(node_buf, meta_.element_size()); + uint32_t neighbor_num = *node_neighbor; + + neighbor_buffers[i].first = neighbor_num; + memcpy(neighbor_buffers[i].second, node_neighbor + 1, + neighbor_num * sizeof(diskann_id_t)); + } + } + + DiskAnnUtil::free_aligned(buf); + + return retval; +} + +int DiskAnnIndexer::load_cache_list( + const std::vector &node_list) { + LOG_INFO("Loading the cache list into memory"); + + size_t num_cached_nodes = node_list.size(); + + neighbor_cache_buffer_.resize(num_cached_nodes * (max_degree_ + 1), 0); + + size_t coord_cache_buf_len = num_cached_nodes * aligned_dim_; + DiskAnnUtil::alloc_aligned((void **)&coord_cache_buf_, + coord_cache_buf_len * meta_.unit_size(), + 8 * meta_.unit_size()); + + memset(coord_cache_buf_, 0, coord_cache_buf_len * meta_.unit_size()); + + constexpr size_t BLOCK_SIZE = 8; + size_t num_blocks = DiskAnnUtil::div_round_up(num_cached_nodes, BLOCK_SIZE); + for (size_t block = 0; block < num_blocks; block++) { + size_t start_idx = block * BLOCK_SIZE; + size_t end_idx = std::min(num_cached_nodes, (block + 1) * BLOCK_SIZE); + + std::vector nodes_to_read; + std::vector coord_buffers; + std::vector> neighbor_buffers; + for (size_t node_idx = start_idx; node_idx < end_idx; node_idx++) { + nodes_to_read.push_back(node_list[node_idx]); + coord_buffers.push_back(reinterpret_cast(coord_cache_buf_) + + node_idx * meta_.element_size()); + neighbor_buffers.emplace_back( + 0, neighbor_cache_buffer_.data() + node_idx * (max_degree_ + 1)); + } + + auto read_status = + read_nodes(nodes_to_read, coord_buffers, neighbor_buffers); + + for (size_t i = 0; i < read_status.size(); i++) { + if (read_status[i] == true) { + coord_cache_.insert(std::make_pair(nodes_to_read[i], coord_buffers[i])); + neighbor_cache_.insert( + std::make_pair(nodes_to_read[i], neighbor_buffers[i])); + } + } + } + + LOG_INFO("Load Cache List Done"); + + return 0; +} + +void DiskAnnIndexer::cache_bfs_levels(uint64_t num_nodes_to_cache, + std::vector &node_list) { + std::set node_set; + + size_t tenp_cnt = static_cast(std::round(doc_cnt_ * 0.1)); + if (num_nodes_to_cache > tenp_cnt) { + LOG_WARN( + "Reducing nodes to cache from: %zu, to: (10 percent of total nodes: " + "%zu)", + (size_t)num_nodes_to_cache, (size_t)tenp_cnt); + + num_nodes_to_cache = tenp_cnt == 0 ? 1 : tenp_cnt; + } + + LOG_INFO("Begin to cache %zu Nodes", (size_t)num_nodes_to_cache); + + std::unordered_set cur_level; + std::unordered_set prev_level; + + for (uint64_t iter = 0; + iter < entrypoints_.size() && cur_level.size() < num_nodes_to_cache; + iter++) { + cur_level.insert(entrypoints_[iter]); + } + + uint64_t level = 1; + uint64_t prev_node_set_size = 0; + while ((node_set.size() + cur_level.size() < num_nodes_to_cache) && + cur_level.size() != 0) { + prev_level.swap(cur_level); + + cur_level.clear(); + + std::vector nodes_to_expand; + nodes_to_expand.reserve(prev_level.size()); + + for (const diskann_id_t &id : prev_level) { + if (node_set.find(id) != node_set.end()) { + continue; + } + + node_set.insert(id); + nodes_to_expand.push_back(id); + } + + std::sort(nodes_to_expand.begin(), nodes_to_expand.end()); + + bool finish_flag = false; + + constexpr uint64_t BLOCK_SIZE = 1024; + uint64_t nblocks = + DiskAnnUtil::div_round_up(nodes_to_expand.size(), BLOCK_SIZE); + for (size_t block = 0; block < nblocks && !finish_flag; block++) { + size_t start = block * BLOCK_SIZE; + size_t end = std::min((uint64_t)((block + 1) * BLOCK_SIZE), + (uint64_t)(nodes_to_expand.size())); + const size_t block_size = end - start; + + std::vector nodes_to_read(nodes_to_expand.begin() + start, + nodes_to_expand.begin() + end); + std::vector coord_buffers(block_size, nullptr); + + std::vector>> + neighbor_buffers; + neighbor_buffers.reserve(block_size); + + for (size_t i = 0; i < block_size; i++) { + neighbor_buffers.emplace_back( + 0, std::vector(max_degree_ + 1)); + } + + std::vector> neighbor_buffers_ptr; + neighbor_buffers_ptr.reserve(block_size); + for (size_t i = 0; i < block_size; i++) { + neighbor_buffers_ptr.emplace_back(neighbor_buffers[i].first, + neighbor_buffers[i].second.data()); + } + + auto read_status = + read_nodes(nodes_to_read, coord_buffers, neighbor_buffers_ptr); + + for (uint32_t i = 0; i < read_status.size(); i++) { + if (read_status[i] == false) { + continue; + } else { + neighbor_buffers[i].first = neighbor_buffers_ptr[i].first; + uint32_t neighbor_num = neighbor_buffers[i].first; + diskann_id_t *neighbors = neighbor_buffers[i].second.data(); + + for (uint32_t j = 0; j < neighbor_num && !finish_flag; j++) { + if (node_set.find(neighbors[j]) == node_set.end()) { + cur_level.insert(neighbors[j]); + } + if (cur_level.size() + node_set.size() >= num_nodes_to_cache) { + finish_flag = true; + } + } + } + } + } + + size_t total_size = node_set.size(); + + LOG_INFO("Level: %zu, Cached Size: %zu, Total Cached Size: %zu", + (size_t)level, (size_t)(total_size - prev_node_set_size), + total_size); + + prev_node_set_size = total_size; + level++; + } + + ailego_assert(node_set.size() + cur_level.size() == num_nodes_to_cache || + cur_level.size() == 0); + + node_list.clear(); + node_list.reserve(node_set.size() + cur_level.size()); + + for (auto node : node_set) { + node_list.push_back(node); + } + + for (auto node : cur_level) { + node_list.push_back(node); + } + + size_t total_size = node_list.size(); + LOG_INFO("Level: %zu, Cached Size: %zu, Total Cached Size: %zu", + (size_t)level, (size_t)(total_size - prev_node_set_size), + (size_t)total_size); + + return; +} + +int DiskAnnIndexer::linear_search(DiskAnnContext *ctx) { + auto &stats = ctx->query_stats(); + auto &dc = ctx->dist_calculator(); + auto &topk_heap = ctx->topk_heap(); + + topk_heap.clear(); + + IOContext &io_ctx = ctx->io_ctx(); + void *aligned_query_raw = ctx->query(); + + void *data_buf = reinterpret_cast(ctx->coord_buffer()); + + uint8_t *sector_buffer = reinterpret_cast(ctx->sector_buffer()); + + const uint64_t sector_num_per_node = + node_per_sector_ > 0 + ? 1 + : DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); + + ailego::ElapsedTime io_timer; + ailego::ElapsedTime query_timer; + ailego::ElapsedTime cpu_timer; + + std::vector frontier; + frontier.reserve(2 * beam_width_); + + std::vector> frontier_neighbors; + frontier_neighbors.reserve(2 * beam_width_); + + std::vector frontier_read_reqs; + frontier_read_reqs.reserve(2 * beam_width_); + + std::vector> + cached_neighbors; + cached_neighbors.reserve(2 * beam_width_); + + uint64_t sector_buffer_idx = 0; + + diskann_id_t id = 0; + while (id < doc_cnt_) { + while (frontier.size() < beam_width_) { + if (!ctx->filter().is_valid() || !ctx->filter()(get_key(id))) { + auto iter = neighbor_cache_.find(id); + if (iter != neighbor_cache_.end()) { + cached_neighbors.push_back( + std::make_tuple(id, iter->second.first, iter->second.second)); + stats.cache_hits++; + } else { + frontier.push_back(id); + } + } + + id++; + if (id >= doc_cnt_) { + break; + } + } + + if (!frontier.empty()) { + for (uint64_t i = 0; i < frontier.size(); i++) { + diskann_id_t cur_id = frontier[i]; + + std::pair frontier_neighbor; + frontier_neighbor.first = cur_id; + frontier_neighbor.second = sector_buffer + sector_num_per_node * + sector_buffer_idx * + DiskAnnUtil::kSectorSize; + frontier_neighbors.push_back(frontier_neighbor); + + sector_buffer_idx++; + + frontier_read_reqs.emplace_back( + index_segment_offset_ + + DiskAnnUtil::get_node_sector(node_per_sector_, max_node_size_, + DiskAnnUtil::kSectorSize, cur_id) * + DiskAnnUtil::kSectorSize, + sector_num_per_node * DiskAnnUtil::kSectorSize, + frontier_neighbor.second); + + stats.disk_page_reads++; + stats.io_num++; + } + + io_timer.reset(); + + int read_ret = reader_->read(frontier_read_reqs, io_ctx); + stats.io_us += io_timer.micro_seconds(); + if (read_ret != 0) { + LOG_ERROR("linear_search: reader_->read failed, ret=%d", read_ret); + ctx->set_error(true); + return IndexError_Runtime; + } + } + + for (auto &cached_neighbor : cached_neighbors) { + auto global_cache_iter = coord_cache_.find(std::get<0>(cached_neighbor)); + void *node_fp_coords_copy = global_cache_iter->second; + + float cur_expanded_dist = dc.dist(aligned_query_raw, node_fp_coords_copy); + + topk_heap.emplace( + std::get<0>(cached_neighbor), + VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords_copy))); + } + + for (auto &frontier_neighbor : frontier_neighbors) { + uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( + node_per_sector_, max_node_size_, frontier_neighbor.second, + frontier_neighbor.first); + + void *node_fp_coords = node_disk_buf; + memcpy(data_buf, node_fp_coords, disk_bytes_per_point_); + + float cur_expanded_dist = dc.dist(aligned_query_raw, data_buf); + + topk_heap.emplace( + frontier_neighbor.first, + VectorInfo(cur_expanded_dist, make_vector_copy(data_buf))); + + stats.cpu_us += cpu_timer.micro_seconds(); + } + + frontier.clear(); + frontier_neighbors.clear(); + frontier_read_reqs.clear(); + cached_neighbors.clear(); + sector_buffer_idx = 0; + } + + stats.total_us += query_timer.micro_seconds(); + + return 0; +} + +int DiskAnnIndexer::keys_search(const std::vector &keys, + DiskAnnContext *ctx) { + auto &stats = ctx->query_stats(); + auto &dc = ctx->dist_calculator(); + auto &topk_heap = ctx->topk_heap(); + + topk_heap.clear(); + + IOContext &io_ctx = ctx->io_ctx(); + void *aligned_query_raw = ctx->query(); + + void *data_buf = reinterpret_cast(ctx->coord_buffer()); + + uint8_t *sector_buffer = reinterpret_cast(ctx->sector_buffer()); + + const uint64_t sector_num_per_node = + node_per_sector_ > 0 + ? 1 + : DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); + + ailego::ElapsedTime query_timer; + ailego::ElapsedTime io_timer; + ailego::ElapsedTime cpu_timer; + + std::vector frontier; + frontier.reserve(2 * beam_width_); + + std::vector> frontier_neighbors; + frontier_neighbors.reserve(2 * beam_width_); + + std::vector frontier_read_reqs; + frontier_read_reqs.reserve(2 * beam_width_); + + std::vector> + cached_neighbors; + cached_neighbors.reserve(2 * beam_width_); + + uint64_t sector_buffer_idx = 0; + + size_t idx = 0; + while (idx < keys.size()) { + while (frontier.size() < beam_width_) { + if (!ctx->filter().is_valid() || !ctx->filter()(keys[idx])) { + diskann_id_t id = get_id(keys[idx]); + + auto iter = neighbor_cache_.find(id); + if (iter != neighbor_cache_.end()) { + cached_neighbors.push_back( + std::make_tuple(id, iter->second.first, iter->second.second)); + stats.cache_hits++; + } else { + frontier.push_back(id); + } + } + + idx++; + if (idx >= keys.size()) { + break; + } + } + + if (!frontier.empty()) { + for (uint64_t i = 0; i < frontier.size(); i++) { + diskann_id_t cur_id = frontier[i]; + + std::pair frontier_neighbor; + frontier_neighbor.first = cur_id; + frontier_neighbor.second = sector_buffer + sector_num_per_node * + sector_buffer_idx * + DiskAnnUtil::kSectorSize; + frontier_neighbors.push_back(frontier_neighbor); + + sector_buffer_idx++; + + frontier_read_reqs.emplace_back( + index_segment_offset_ + + DiskAnnUtil::get_node_sector(node_per_sector_, max_node_size_, + DiskAnnUtil::kSectorSize, cur_id) * + DiskAnnUtil::kSectorSize, + sector_num_per_node * DiskAnnUtil::kSectorSize, + frontier_neighbor.second); + + stats.disk_page_reads++; + stats.io_num++; + } + + io_timer.reset(); + + int read_ret = reader_->read(frontier_read_reqs, io_ctx); + stats.io_us += io_timer.micro_seconds(); + if (read_ret != 0) { + LOG_ERROR("keys_search: reader_->read failed, ret=%d", read_ret); + ctx->set_error(true); + return IndexError_Runtime; + } + } + + for (auto &cached_neighbor : cached_neighbors) { + auto global_cache_iter = coord_cache_.find(std::get<0>(cached_neighbor)); + void *node_fp_coords_copy = global_cache_iter->second; + + float cur_expanded_dist = dc.dist(aligned_query_raw, node_fp_coords_copy); + + topk_heap.emplace( + std::get<0>(cached_neighbor), + VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords_copy))); + } + + for (auto &frontier_neighbor : frontier_neighbors) { + uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( + node_per_sector_, max_node_size_, frontier_neighbor.second, + frontier_neighbor.first); + + void *node_fp_coords = node_disk_buf; + memcpy(data_buf, node_fp_coords, disk_bytes_per_point_); + + float cur_expanded_dist = dc.dist(aligned_query_raw, data_buf); + + topk_heap.emplace( + frontier_neighbor.first, + VectorInfo(cur_expanded_dist, make_vector_copy(data_buf))); + + stats.cpu_us += cpu_timer.micro_seconds(); + } + + frontier.clear(); + frontier_neighbors.clear(); + frontier_read_reqs.clear(); + cached_neighbors.clear(); + sector_buffer_idx = 0; + } + + stats.total_us += query_timer.micro_seconds(); + + return 0; +} + +int DiskAnnIndexer::get_vector(diskann_id_t id, IndexContext::Pointer &context, + std::string &vector) { + DiskAnnContext *ctx = dynamic_cast(context.get()); + + auto &stats = ctx->query_stats(); + + IOContext &io_ctx = ctx->io_ctx(); + + uint8_t *sector_buffer = reinterpret_cast(ctx->sector_buffer()); + + const uint64_t sector_num_per_node = + node_per_sector_ > 0 + ? 1 + : DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); + + ailego::ElapsedTime query_timer; + ailego::ElapsedTime io_timer; + ailego::ElapsedTime cpu_timer; + + std::vector frontier; + frontier.reserve(2 * beam_width_); + + std::vector> frontier_neighbors; + frontier_neighbors.reserve(2 * beam_width_); + + std::vector frontier_read_reqs; + frontier_read_reqs.reserve(2 * beam_width_); + + std::vector> + cached_neighbors; + cached_neighbors.reserve(2 * beam_width_); + + auto iter = neighbor_cache_.find(id); + if (iter != neighbor_cache_.end()) { + void *node_fp_coords_copy = iter->second.second; + + vector.resize(meta_.element_size()); + ::memcpy(&(vector[0]), node_fp_coords_copy, meta_.element_size()); + + return 0; + } else { + std::pair frontier_neighbor; + frontier_neighbor.first = id; + frontier_neighbor.second = sector_buffer; + frontier_neighbors.push_back(frontier_neighbor); + + frontier_read_reqs.emplace_back( + index_segment_offset_ + + DiskAnnUtil::get_node_sector(node_per_sector_, max_node_size_, + DiskAnnUtil::kSectorSize, id) * + DiskAnnUtil::kSectorSize, + sector_num_per_node * DiskAnnUtil::kSectorSize, + frontier_neighbor.second); + + stats.disk_page_reads++; + stats.io_num++; + + io_timer.reset(); + + reader_->read(frontier_read_reqs, io_ctx); + stats.io_us += io_timer.micro_seconds(); + + uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( + node_per_sector_, max_node_size_, frontier_neighbor.second, id); + + void *node_fp_coords = node_disk_buf; + + vector.resize(meta_.element_size()); + ::memcpy(&(vector[0]), node_fp_coords, meta_.element_size()); + + stats.cpu_us += cpu_timer.micro_seconds(); + } + + return 0; +} + +int DiskAnnIndexer::knn_search(DiskAnnContext *ctx) { + int ret = cached_beam_search(ctx); + if (ret != 0) { + return ret; + } + + if (ctx->group_by_search()) { + ret = cached_beam_search_by_group(ctx); + if (ret != 0) { + return ret; + } + } + + return 0; +} + +int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { + auto &stats = ctx->query_stats(); + auto &dc = ctx->dist_calculator(); + auto &topk_heap = ctx->topk_heap(); + auto &visit_filter = ctx->visit_filter(); + + topk_heap.clear(); + + IOContext &io_ctx = ctx->io_ctx(); + + uint8_t *sector_buffer = reinterpret_cast(ctx->sector_buffer()); + + const uint64_t sector_num_per_node = + node_per_sector_ > 0 + ? 1 + : DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); + + pq_table_->preprocess_pq_dist_table(ctx->query_rotated(), + ctx->pq_table_dist_buffer()); + + ailego::ElapsedTime query_timer; + ailego::ElapsedTime io_timer; + ailego::ElapsedTime cpu_timer; + + NeighborPriorityQueue candidates; + + candidates.reserve(ctx->list_size()); + + diskann_id_t best_medoid = 0; + float best_dist = (std::numeric_limits::max)(); + for (uint64_t cur_m = 0; cur_m < entrypoints_.size(); cur_m++) { + float cur_expanded_dist = + dc.dist(ctx->query(), centroid_data_ + aligned_dim_ * cur_m); + + if (cur_expanded_dist < best_dist) { + best_medoid = entrypoints_[cur_m]; + best_dist = cur_expanded_dist; + } + } + + float dist; + pq_table_->compute_dists(1, &best_medoid, pq_chunk_num_, + ctx->pq_table_dist_buffer(), ctx->pq_coord_buffer(), + &dist); + candidates.insert(Neighbor(best_medoid, dist)); + visit_filter.set_visited(best_medoid); + + uint32_t num_ios = 0; + + std::vector frontier; + frontier.reserve(2 * beam_width_); + + std::vector> frontier_neighbors; + frontier_neighbors.reserve(2 * beam_width_); + + std::vector frontier_read_reqs; + frontier_read_reqs.reserve(2 * beam_width_); + + std::vector> + cached_neighbors; + cached_neighbors.reserve(2 * beam_width_); + + while (candidates.has_unexpanded_node() && num_ios < io_limit_) { + frontier.clear(); + frontier_neighbors.clear(); + frontier_read_reqs.clear(); + cached_neighbors.clear(); + + uint64_t sector_buffer_idx = 0; + + uint32_t num_seen = 0; + while (candidates.has_unexpanded_node() && frontier.size() < beam_width_ && + num_seen < beam_width_) { + auto neighbor = candidates.closest_unexpanded(); + num_seen++; + + auto iter = neighbor_cache_.find(neighbor.id); + if (iter != neighbor_cache_.end()) { + cached_neighbors.push_back(std::make_tuple( + neighbor.id, iter->second.first, iter->second.second)); + stats.cache_hits++; + } else { + frontier.push_back(neighbor.id); + } + } + + if (!frontier.empty()) { + stats.hop_num++; + + for (uint64_t i = 0; i < frontier.size(); i++) { + diskann_id_t cur_id = frontier[i]; + + std::pair frontier_neighbor; + frontier_neighbor.first = cur_id; + frontier_neighbor.second = sector_buffer + sector_num_per_node * + sector_buffer_idx * + DiskAnnUtil::kSectorSize; + frontier_neighbors.push_back(frontier_neighbor); + + sector_buffer_idx++; + + frontier_read_reqs.emplace_back( + index_segment_offset_ + + DiskAnnUtil::get_node_sector(node_per_sector_, max_node_size_, + DiskAnnUtil::kSectorSize, cur_id) * + DiskAnnUtil::kSectorSize, + sector_num_per_node * DiskAnnUtil::kSectorSize, + frontier_neighbor.second); + + stats.disk_page_reads++; + stats.io_num++; + num_ios++; + } + + io_timer.reset(); + + int read_ret = reader_->read(frontier_read_reqs, io_ctx); + stats.io_us += io_timer.micro_seconds(); + if (read_ret != 0) { + LOG_ERROR("cached_beam_search: reader_->read failed, ret=%d", read_ret); + ctx->set_error(true); + return IndexError_Runtime; + } + } + + for (auto &cached_neighbor : cached_neighbors) { + auto global_cache_iter = coord_cache_.find(std::get<0>(cached_neighbor)); + void *node_fp_coords_copy = global_cache_iter->second; + + float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords_copy); + + if (!ctx->filter().is_valid() || + !ctx->filter()(get_key(std::get<0>(cached_neighbor)))) { + topk_heap.emplace(std::get<0>(cached_neighbor), + VectorInfo(cur_expanded_dist, + make_vector_copy(node_fp_coords_copy))); + } + + uint32_t neighbor_num = std::get<1>(cached_neighbor); + diskann_id_t *node_neighbors = std::get<2>(cached_neighbor); + + cpu_timer.reset(); + + std::vector distances(neighbor_num); + pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, + ctx->pq_table_dist_buffer(), + ctx->pq_coord_buffer(), distances.data()); + + stats.dist_num += neighbor_num; + stats.cpu_us += cpu_timer.micro_seconds(); + + for (uint64_t m = 0; m < neighbor_num; ++m) { + diskann_id_t id = node_neighbors[m]; + visit_filter.set_visited(id); + + Neighbor nn(id, distances[m]); + candidates.insert(nn); + } + } + + for (auto &frontier_neighbor : frontier_neighbors) { + uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( + node_per_sector_, max_node_size_, frontier_neighbor.second, + frontier_neighbor.first); + uint32_t *node_buf = DiskAnnUtil::offset_to_node_neighbor( + node_disk_buf, meta_.element_size()); + uint32_t neighbor_num = *node_buf; + + void *node_fp_coords = node_disk_buf; + + float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords); + + if (!ctx->filter().is_valid() || + !ctx->filter()(get_key(frontier_neighbor.first))) { + topk_heap.emplace( + frontier_neighbor.first, + VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords))); + } + + diskann_id_t *node_neighbors = + reinterpret_cast(node_buf + 1); + + cpu_timer.reset(); + std::vector distances(neighbor_num); + pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, + ctx->pq_table_dist_buffer(), + ctx->pq_coord_buffer(), distances.data()); + + stats.dist_num += neighbor_num; + stats.cpu_us += cpu_timer.micro_seconds(); + + cpu_timer.reset(); + for (uint64_t m = 0; m < neighbor_num; ++m) { + diskann_id_t id = node_neighbors[m]; + visit_filter.set_visited(id); + stats.dist_num++; + + Neighbor nn(id, distances[m]); + candidates.insert(nn); + } + + stats.cpu_us += cpu_timer.micro_seconds(); + } + } + + stats.total_us += query_timer.micro_seconds(); + + return 0; +} + +int DiskAnnIndexer::cached_beam_search_in_mem(DiskAnnContext * /*ctx*/) { + return IndexError_NotImplemented; +} + +int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { + if (!ctx->group_by().is_valid()) { + return 0; + } + + std::function group_by = [&](diskann_id_t id) { + return ctx->group_by()(get_key(id)); + }; + + // devide into groups + auto &topk_heap = ctx->topk_heap(); + auto &visit_filter = ctx->visit_filter(); + + std::map &group_topk_heaps = ctx->group_topk_heaps(); + + for (uint32_t i = 0; i < topk_heap.size(); ++i) { + diskann_id_t id = topk_heap[i].first; + auto info = topk_heap[i].second; + + std::string group_id = group_by(id); + + auto &group_topk_heap = group_topk_heaps[group_id]; + if (group_topk_heap.empty()) { + group_topk_heap.limit(ctx->group_topk()); + } + + topk_heap.emplace(id, info); + } + + // stage 2, expand to reach group num as possible + if (group_topk_heaps.size() < ctx->group_num()) { + NeighborPriorityQueue candidates; + + candidates.reserve(ctx->list_size()); + + for (uint32_t i = 0; i < topk_heap.size(); ++i) { + diskann_id_t id = topk_heap[i].first; + float score = topk_heap[i].second.dist_; + + visit_filter.set_visited(id); + candidates.insert(Neighbor(id, score)); + } + + ailego::ElapsedTime io_timer; + ailego::ElapsedTime query_timer; + ailego::ElapsedTime cpu_timer; + + auto &stats = ctx->query_stats(); + auto &dc = ctx->dist_calculator(); + + IOContext &io_ctx = ctx->io_ctx(); + + void *data_buf = reinterpret_cast(ctx->coord_buffer()); + uint8_t *sector_buffer = reinterpret_cast(ctx->sector_buffer()); + + const uint64_t sector_num_per_node = + node_per_sector_ > 0 ? 1 + : DiskAnnUtil::div_round_up( + max_node_size_, DiskAnnUtil::kSectorSize); + + pq_table_->preprocess_pq_dist_table(ctx->query_rotated(), + ctx->pq_table_dist_buffer()); + + uint32_t num_ios = 0; + + std::vector frontier; + frontier.reserve(2 * beam_width_); + std::vector> frontier_neighbors; + frontier_neighbors.reserve(2 * beam_width_); + std::vector frontier_read_reqs; + frontier_read_reqs.reserve(2 * beam_width_); + std::vector> + cached_neighbors; + cached_neighbors.reserve(2 * beam_width_); + + uint64_t sector_buffer_idx; + + while (candidates.has_unexpanded_node() && num_ios < io_limit_) { + frontier.clear(); + frontier_neighbors.clear(); + frontier_read_reqs.clear(); + cached_neighbors.clear(); + sector_buffer_idx = 0; + + uint32_t num_seen = 0; + while (candidates.has_unexpanded_node() && + frontier.size() < beam_width_ && num_seen < beam_width_) { + auto neighbor = candidates.closest_unexpanded(); + num_seen++; + + auto iter = neighbor_cache_.find(neighbor.id); + if (iter != neighbor_cache_.end()) { + cached_neighbors.push_back(std::make_tuple( + neighbor.id, iter->second.first, iter->second.second)); + stats.cache_hits++; + } else { + frontier.push_back(neighbor.id); + } + } + + if (!frontier.empty()) { + stats.hop_num++; + + for (uint64_t i = 0; i < frontier.size(); i++) { + diskann_id_t cur_id = frontier[i]; + + std::pair frontier_neighbor; + frontier_neighbor.first = cur_id; + frontier_neighbor.second = + sector_buffer + sector_num_per_node * sector_buffer_idx * + DiskAnnUtil::kSectorSize; + frontier_neighbors.push_back(frontier_neighbor); + + sector_buffer_idx++; + + frontier_read_reqs.emplace_back( + index_segment_offset_ + DiskAnnUtil::get_node_sector( + node_per_sector_, max_node_size_, + DiskAnnUtil::kSectorSize, cur_id) * + DiskAnnUtil::kSectorSize, + sector_num_per_node * DiskAnnUtil::kSectorSize, + frontier_neighbor.second); + + stats.disk_page_reads++; + stats.io_num++; + num_ios++; + } + + io_timer.reset(); + + reader_->read(frontier_read_reqs, io_ctx); // synchronous IO linux + stats.io_us += io_timer.micro_seconds(); + } + + for (auto &cached_neighbor : cached_neighbors) { + auto global_cache_iter = + coord_cache_.find(std::get<0>(cached_neighbor)); + void *node_fp_coords_copy = global_cache_iter->second; + + float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords_copy); + + if (!ctx->filter().is_valid() || + !ctx->filter()(get_key(std::get<0>(cached_neighbor)))) { + std::string group_id = group_by(std::get<0>(cached_neighbor)); + + auto &group_topk_heap = group_topk_heaps[group_id]; + if (group_topk_heap.empty()) { + group_topk_heap.limit(ctx->group_topk()); + } + + group_topk_heap.emplace_back( + std::get<0>(cached_neighbor), + VectorInfo(cur_expanded_dist, + make_vector_copy(node_fp_coords_copy))); + + if (group_topk_heaps.size() >= ctx->group_num()) { + break; + } + } + + uint64_t neighbor_num = std::get<1>(cached_neighbor); + diskann_id_t *node_neighbors = std::get<2>(cached_neighbor); + + cpu_timer.reset(); + + std::vector distances(neighbor_num); + pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, + ctx->pq_table_dist_buffer(), + ctx->pq_coord_buffer(), distances.data()); + + stats.dist_num += neighbor_num; + stats.cpu_us += cpu_timer.micro_seconds(); + + for (uint64_t m = 0; m < neighbor_num; ++m) { + diskann_id_t id = node_neighbors[m]; + visit_filter.set_visited(id); + + Neighbor nn(id, distances[m]); + candidates.insert(nn); + } + } + + for (auto &frontier_neighbor : frontier_neighbors) { + uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( + node_per_sector_, max_node_size_, frontier_neighbor.second, + frontier_neighbor.first); + uint32_t *node_buf = DiskAnnUtil::offset_to_node_neighbor( + node_disk_buf, meta_.element_size()); + uint32_t neighbor_num = *node_buf; + + void *node_fp_coords = node_disk_buf; + memcpy(data_buf, node_fp_coords, disk_bytes_per_point_); + + float cur_expanded_dist = dc.dist(ctx->query(), data_buf); + + if (!ctx->filter().is_valid() || + !ctx->filter()(get_key(frontier_neighbor.first))) { + std::string group_id = group_by(frontier_neighbor.first); + + auto &group_topk_heap = group_topk_heaps[group_id]; + if (group_topk_heap.empty()) { + group_topk_heap.limit(ctx->group_topk()); + } + + group_topk_heap.emplace_back( + frontier_neighbor.first, + VectorInfo(cur_expanded_dist, make_vector_copy(data_buf))); + + if (group_topk_heaps.size() >= ctx->group_num()) { + break; + } + } + + cpu_timer.reset(); + + std::vector distances(neighbor_num); + diskann_id_t *node_neighbors = + reinterpret_cast(node_buf + 1); + pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, + ctx->pq_table_dist_buffer(), + ctx->pq_coord_buffer(), distances.data()); + + stats.dist_num += neighbor_num; + stats.cpu_us += cpu_timer.micro_seconds(); + + cpu_timer.reset(); + for (uint64_t m = 0; m < neighbor_num; ++m) { + diskann_id_t id = node_neighbors[m]; + visit_filter.set_visited(id); + stats.dist_num++; + + Neighbor nn(id, distances[m]); + candidates.insert(nn); + } + + stats.cpu_us += cpu_timer.micro_seconds(); + } + } + + stats.total_us += query_timer.micro_seconds(); + } + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h new file mode 100644 index 0000000..c372d28 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -0,0 +1,109 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "diskann_context.h" +#include "diskann_file_reader.h" +#include "diskann_pq_table.h" +#include "diskann_searcher_entity.h" +#include "diskann_util.h" + +namespace zvec { +namespace core { + +class DiskAnnIndexer { + public: + typedef std::shared_ptr Pointer; + + public: + DiskAnnIndexer(const IndexMeta &meta); + ~DiskAnnIndexer(); + + public: + int init(DiskAnnSearcherEntity &entity); + int load_cache_list(const std::vector &node_list); + + void cache_bfs_levels(uint64_t num_nodes_to_cache, + std::vector &node_list); + + int cached_beam_search(DiskAnnContext *ctx); + int cached_beam_search_by_group(DiskAnnContext *ctx); + + int cached_beam_search_in_mem(DiskAnnContext *ctx); + + int knn_search(DiskAnnContext *ctx); + int linear_search(DiskAnnContext *ctx); + int keys_search(const std::vector &keys, DiskAnnContext *ctx); + + int get_vector(diskann_id_t id, IndexContext::Pointer &context, + std::string &vector); + + diskann_key_t get_key(diskann_id_t id) const; + diskann_id_t get_id(diskann_key_t key) const; + + //! Copy element_size() bytes from src into a new vector value string + std::string make_vector_copy(const void *src) const { + return std::string(static_cast(src), meta_.element_size()); + } + + std::vector read_nodes( + const std::vector &node_ids, + std::vector &coord_buffers, + std::vector> &nbr_buffers); + + protected: + int use_medroids_data_as_centroids(); + + private: + DiskAnnSearcherEntity *entity_; + + IndexStorage::Pointer storage_{}; + IndexMeta meta_; + + uint32_t max_degree_{0}; + uint32_t node_per_sector_{0}; + uint32_t max_node_size_{0}; + uint64_t pq_chunk_num_{0}; + uint64_t disk_bytes_per_point_{0}; + uint64_t aligned_dim_{0}; + uint64_t index_segment_offset_{0}; + uint64_t sector_num_per_node_{0}; + + float *centroid_data_{nullptr}; + + diskann_id_t medoid_; + std::vector entrypoints_; + + std::shared_ptr reader_{nullptr}; + + PQTable::Pointer pq_table_; + + IOContext init_ctx_{0}; + + std::vector neighbor_cache_buffer_; + void *coord_cache_buf_{nullptr}; + + std::map coord_cache_; + std::map> neighbor_cache_; + + uint32_t beam_width_{2}; + uint32_t io_limit_{std::numeric_limits::max()}; + + uint64_t doc_cnt_{0}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_params.h b/src/core/algorithm/diskann/diskann_params.h new file mode 100644 index 0000000..fac0dc6 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_params.h @@ -0,0 +1,53 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +static const std::string PARAM_DISKANN_BUILDER_MAX_DEGREE( + "zvec.diskann.builder.max_degree"); +static const std::string PARAM_DISKANN_BUILDER_LIST_SIZE( + "zvec.diskann.builder.list_size"); +static const std::string PARAM_DISKANN_BUILDER_MEMORY_LIMIT( + "zvec.diskann.builder.memory_limit"); +static const std::string PARAM_DISKANN_BUILDER_MEMORY_BUDGET( + "zvec.diskann.builder.memory_budget"); +static const std::string PARAM_DISKANN_BUILDER_DISK_PQ_DIM( + "zvec.diskann.builder.disk_pq_dim"); +static const std::string PARAM_DISKANN_BUILDER_THREAD_COUNT( + "zvec.diskann.builder.thread_count"); +static const std::string PARAM_DISKANN_BUILDER_MAX_TRAIN_SAMPLE_COUNT( + "zvec.diskann.builder.max_train_sample_count"); +static const std::string PARAM_DISKANN_BUILDER_TRAIN_SAMPLE_RATIO( + "zvec.diskann.builder.train_sample_ratio"); +static const std::string PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM( + "zvec.diskann.builder.max_pq_chunk_num"); + +static const std::string PARAM_DISKANN_SEARCHER_LIST_SIZE( + "zvec.diskann.searcher.list_size"); +static const std::string PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM( + "zvec.diskann.searcher.cache_node_num"); + +static const std::string PARAM_DISKANN_REDUCER_INDEX_NAME( + "zvec.diskann.reducer.index_name"); +static const std::string PARAM_DISKANN_REDUCER_WORKING_PATH( + "zvec.diskann.reducer.working_path"); +static const std::string PARAM_DISKANN_REDUCER_NUM_OF_ADD_THREADS( + "zvec.diskann.reducer.num_of_add_threads"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_pq_table.cc b/src/core/algorithm/diskann/diskann_pq_table.cc new file mode 100644 index 0000000..0c13e40 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_pq_table.cc @@ -0,0 +1,117 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_pq_table.h" +#include "diskann_entity.h" + +namespace zvec { +namespace core { + +PQTable::PQTable(const IndexMeta &meta, uint32_t chunk_num) + : chunk_num_(chunk_num) { + meta_ = meta; + + if (meta.metric_name() == "Cosine") { + if (meta.data_type() == IndexMeta::DataType::DT_FP32) { + meta_.set_dimension(meta.dimension() - 1); + } else { + meta_.set_dimension(meta.dimension() - 2); + } + } +} + +PQTable::~PQTable() {} + +int PQTable::init(std::vector &full_pivot_data, + std::vector ¢roid, + std::vector &chunk_offsets, + std::vector &pq_data) { + full_pivot_data_ = std::move(full_pivot_data); + centroid_ = std::move(centroid); + chunk_offsets_ = std::move(chunk_offsets); + pq_data_ = std::move(pq_data); + + // alloc and compute transpose + transposed_tables_.resize(kPQCentroidNum * meta_.element_size()); + + uint32_t dim = meta_.dimension(); + uint32_t type = meta_.data_type(); + + switch (type) { + case IndexMeta::DataType::DT_FP32: { + float *transposed_tables_ptr = + reinterpret_cast(&transposed_tables_[0]); + float *full_pivot_data_ptr = + reinterpret_cast(&full_pivot_data_[0]); + for (size_t i = 0; i < kPQCentroidNum; i++) { + for (size_t j = 0; j < dim; j++) { + transposed_tables_ptr[j * kPQCentroidNum + i] = + full_pivot_data_ptr[i * dim + j]; + } + } + break; + } + case IndexMeta::DataType::DT_FP16: { + ailego::Float16 *transposed_tables_ptr = + reinterpret_cast(&transposed_tables_[0]); + ailego::Float16 *full_pivot_data_ptr = + reinterpret_cast(&full_pivot_data_[0]); + for (size_t i = 0; i < kPQCentroidNum; i++) { + for (size_t j = 0; j < dim; j++) { + transposed_tables_ptr[j * kPQCentroidNum + i] = + full_pivot_data_ptr[i * dim + j]; + } + } + break; + } + default: + LOG_ERROR("unsupported type, type: %u", type); + return IndexError_Unsupported; + } + + return 0; +} + +void PQTable::aggregate_coords(uint32_t id_num, const diskann_id_t *ids, + const uint8_t *all_coords, size_t dim, + uint8_t *out) { + for (size_t i = 0; i < id_num; i++) { + memcpy(out + i * dim, all_coords + ids[i] * dim, dim * sizeof(uint8_t)); + } +} + +void PQTable::pq_dist_lookup(const uint8_t *pq_ids, size_t id_num, + size_t pq_nchunks, const float *pq_dist_buffer, + float *dists_out) { + ailego_prefetch(dists_out); + ailego_prefetch(pq_ids); + ailego_prefetch(pq_ids + 64); + ailego_prefetch(pq_ids + 128); + + memset(dists_out, 0, id_num * sizeof(float)); + + for (size_t chunk = 0; chunk < pq_nchunks; chunk++) { + const float *chunk_dists = pq_dist_buffer + kPQCentroidNum * chunk; + if (chunk < pq_nchunks - 1) { + ailego_prefetch(chunk_dists + kPQCentroidNum); + } + for (size_t idx = 0; idx < id_num; idx++) { + uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; + dists_out[idx] += chunk_dists[pq_centerid]; + } + } +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_pq_table.h b/src/core/algorithm/diskann/diskann_pq_table.h new file mode 100644 index 0000000..a998c29 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_pq_table.h @@ -0,0 +1,154 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "diskann_entity.h" + +namespace zvec { +namespace core { + +class PQTable { + public: + typedef std::shared_ptr Pointer; + + public: + static constexpr uint32_t kPQBitNum = 8; + static constexpr uint32_t kPQCentroidNum = 1 << kPQBitNum; + static constexpr uint32_t kMaxTrainSampleCount = 200000; + static constexpr double kTrainSampleRatio = 1.0; + static constexpr uint32_t kMeanIterNum = 12; + + public: + PQTable(const IndexMeta &meta, uint32_t chunk_num); + virtual ~PQTable(); + + int init(std::vector &table, std::vector ¢roid, + std::vector &chunk_offsets, std::vector &pq_data); + + + template + void compute_distance_table_ip(const T *query_vec, float *dist_vec) { + memset(dist_vec, 0, kPQCentroidNum * chunk_num_ * sizeof(float)); + + const T *transposed_tables_ptr = + reinterpret_cast(transposed_tables_.data()); + // chunk wise distance computation + for (size_t chunk = 0; chunk < chunk_num_; chunk++) { + // sum (q-c)^2 for the dimensions associated with this chunk + float *chunk_dists = dist_vec + (kPQCentroidNum * chunk); + + for (size_t j = chunk_offsets_[chunk]; j < chunk_offsets_[chunk + 1]; + j++) { + const T *centers_dim_vec = &transposed_tables_ptr[kPQCentroidNum * j]; + + for (size_t idx = 0; idx < kPQCentroidNum; idx++) { + float centor_data = centers_dim_vec[idx]; + float query_data = query_vec[j]; + float dim_score = centor_data * query_data; + chunk_dists[idx] += -dim_score; + } + } + } + } + + template + void compute_distance_table(const T *query_vec, float *dist_vec) { + memset(dist_vec, 0, kPQCentroidNum * chunk_num_ * sizeof(float)); + + const T *transposed_tables_ptr = + reinterpret_cast(transposed_tables_.data()); + // chunk wise distance computation + for (size_t chunk = 0; chunk < chunk_num_; chunk++) { + // sum (q-c)^2 for the dimensions associated with this chunk + float *chunk_dists = dist_vec + (kPQCentroidNum * chunk); + + for (size_t j = chunk_offsets_[chunk]; j < chunk_offsets_[chunk + 1]; + j++) { + const T *centers_dim_vec = &transposed_tables_ptr[kPQCentroidNum * j]; + + for (size_t idx = 0; idx < kPQCentroidNum; idx++) { + float diff = centers_dim_vec[idx] - query_vec[j]; + chunk_dists[idx] += (diff * diff); + } + } + } + } + + template + void preprocess_query(T *query) { + const T *centroid_ptr = reinterpret_cast(centroid_.data()); + for (size_t i = 0; i < meta_.dimension(); i++) { + query[i] -= centroid_ptr[i]; + } + } + + void aggregate_coords(uint32_t id_num, const diskann_id_t *ids, + const uint8_t *all_coords, size_t dim, uint8_t *out); + + void pq_dist_lookup(const uint8_t *pq_ids, size_t id_num, size_t pq_nchunks, + const float *pq_dist_buffer, float *dists); + + void compute_dists(uint32_t id_num, const diskann_id_t *ids, + uint32_t chunk_num, float *pq_dist_buffer, + void *coord_buffer, float *dists) { + uint8_t *coord_buffer_ptr = reinterpret_cast(coord_buffer); + + aggregate_coords(id_num, ids, this->pq_data(), chunk_num, coord_buffer_ptr); + + pq_dist_lookup(coord_buffer_ptr, id_num, chunk_num, pq_dist_buffer, dists); + + return; + } + + int preprocess_pq_dist_table(void *query_rotated, float *dist_buffer) { + switch (meta_.data_type()) { + case IndexMeta::DataType::DT_FP32: + preprocess_query(reinterpret_cast(query_rotated)); + compute_distance_table(reinterpret_cast(query_rotated), + dist_buffer); + + break; + case IndexMeta::DataType::DT_FP16: + preprocess_query(reinterpret_cast(query_rotated)); + compute_distance_table( + reinterpret_cast(query_rotated), dist_buffer); + break; + default: + LOG_ERROR("Unsupported Type: %u", meta_.data_type()); + return IndexError_Unsupported; + } + + return 0; + } + + public: + const uint8_t *pq_data() const { + return pq_data_.data(); + } + + private: + std::vector full_pivot_data_; + std::vector transposed_tables_; + + std::vector centroid_; + std::vector chunk_offsets_; + std::vector pq_data_; + + IndexMeta meta_; + uint64_t chunk_num_{0}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_pq_trainer.cc b/src/core/algorithm/diskann/diskann_pq_trainer.cc new file mode 100644 index 0000000..73ca016 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_pq_trainer.cc @@ -0,0 +1,417 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_pq_trainer.h" +#include "diskann_entity.h" +#include "diskann_util.h" + +namespace zvec { +namespace core { + +DiskAnnPqTrainer::DiskAnnPqTrainer(uint32_t max_train_sample_count) + : max_train_sample_count_{max_train_sample_count} {} + +DiskAnnPqTrainer::~DiskAnnPqTrainer() {} + +int DiskAnnPqTrainer::gen_random_sample(IndexHolder::Pointer holder, + const IndexMeta &meta, + std::string &sample_data, + size_t &sample_size) { + double train_sample_ratio = + max_train_sample_count_ < 1 ? max_train_sample_count_ : 1; + + uint32_t max_train_sample_count = train_sample_ratio * holder->count(); + max_train_sample_count = max_train_sample_count > max_train_sample_count_ + ? max_train_sample_count_ + : max_train_sample_count; + + std::vector> sample_vecs; + + // Use a fixed seed for deterministic sampling across runs. + uint32_t x = 456321; + std::mt19937 gen(x); + std::uniform_real_distribution dist(0, 1); + + uint32_t vec_size = meta.element_size(); + + auto iter = holder->create_iterator(); + if (!iter) { + LOG_ERROR("Create iterator for holder failed"); + return IndexError_Runtime; + } + + size_t sample_count = 0; + while (iter->is_valid() && sample_count < max_train_sample_count) { + float random = dist(gen); + + if (random < train_sample_ratio) { + const void *vec = iter->data(); + + std::vector temp_vec; + temp_vec.resize(vec_size); + + std::memcpy(reinterpret_cast(&temp_vec[0]), vec, vec_size); + + sample_vecs.push_back(std::move(temp_vec)); + + sample_count++; + } + + iter->next(); + } + + sample_size = sample_vecs.size(); + sample_data.reserve(sample_size * vec_size); + + for (size_t i = 0; i < sample_size; i++) { + sample_data.append(reinterpret_cast(sample_vecs[i].data()), + vec_size); + } + + return 0; +} + +template +int DiskAnnPqTrainer::prepare_pq_train_data( + const IndexMeta &meta, size_t num_train, std::string &train_data, + bool use_zero_mean, std::vector ¢roid, + std::shared_ptr &train_features) { + uint32_t dim = meta.dimension(); + uint32_t vec_size = meta.element_size(); + + std::string train_data_processed; + train_data_processed.resize(num_train * vec_size); + + std::memcpy(&(train_data_processed[0]), train_data.data(), + num_train * vec_size); + + // use fp32 to accumulate to avoid overflow + std::vector centroid_temp(dim); + for (uint64_t d = 0; d < dim; d++) { + centroid_temp[d] = 0; + } + + T *train_data_processed_ptr = reinterpret_cast(&train_data_processed[0]); + + if (use_zero_mean) { + for (uint64_t d = 0; d < dim; d++) { + for (uint64_t p = 0; p < num_train; p++) { + centroid_temp[d] += train_data_processed_ptr[p * dim + d]; + } + centroid_temp[d] /= num_train; + } + + for (uint64_t d = 0; d < dim; d++) { + for (uint64_t p = 0; p < num_train; p++) { + train_data_processed_ptr[p * dim + d] -= centroid_temp[d]; + } + } + } + + for (size_t i = 0; i < num_train; ++i) { + train_features->emplace(train_data_processed_ptr + i * dim); + } + + // copy the centroid out + centroid.resize(vec_size); + T *centroid_ptr = reinterpret_cast(centroid.data()); + for (uint64_t d = 0; d < dim; d++) { + centroid_ptr[d] = centroid_temp[d]; + } + + return 0; +} + +template +int DiskAnnPqTrainer::convert_pivot_data( + const IndexMeta &meta, uint32_t num_centers, uint32_t pq_chunk_num, + const std::vector &chunk_dims, + const std::vector &chunk_offsets, + IndexCluster::CentroidList ¢roids, + std::vector &full_pivot_data) { + uint32_t dim = meta.dimension(); + uint32_t element_size = meta.element_size(); + + full_pivot_data.resize(num_centers * element_size); + + for (size_t chunk = 0; chunk < pq_chunk_num; ++chunk) { + for (size_t cluster = 0; cluster < num_centers; ++cluster) { + size_t idx = chunk * num_centers + cluster; + + T *pivot_data_ptr = reinterpret_cast(&(full_pivot_data[0])) + + cluster * dim + chunk_offsets[chunk]; + const T *feature_ptr = + reinterpret_cast(centroids[idx].feature()); + for (size_t d = 0; d <= chunk_dims[chunk]; ++d) { + pivot_data_ptr[d] = feature_ptr[d]; + } + } + } + + return 0; +} + +int DiskAnnPqTrainer::train_pq(IndexThreads::Pointer threads, + const IndexMeta &meta, std::string &train_data, + size_t num_train, uint32_t num_centers, + uint32_t pq_chunk_num, uint32_t max_iterations, + bool use_zero_mean, + std::vector &full_pivot_data, + std::vector ¢roid, + std::vector &chunk_offsets) { + uint32_t dim = meta.dimension(); + if (pq_chunk_num > dim) { + LOG_ERROR("Error: number of chunks more than dimension. chunk: %u, dim: %u", + pq_chunk_num, dim); + return IndexError_InvalidArgument; + } + + std::shared_ptr train_features( + new CompactIndexFeatures(meta)); + + uint32_t type = meta.data_type(); + + int ret; + switch (type) { + case IndexMeta::DataType::DT_FP32: + ret = prepare_pq_train_data( + meta, num_train, train_data, use_zero_mean, centroid, train_features); + if (ret != 0) { + LOG_ERROR("Failed to prepare pq train data"); + return ret; + } + break; + + case IndexMeta::DataType::DT_FP16: + ret = prepare_pq_train_data( + meta, num_train, train_data, use_zero_mean, centroid, train_features); + if (ret != 0) { + LOG_ERROR("Failed to prepare pq train data"); + return ret; + } + break; + } + + // Do Train + ailego::Params params; + params.set(MULTI_CHUNK_CLUSTER_COUNT, num_centers); + params.set(MULTI_CHUNK_CLUSTER_CHUNK_COUNT, pq_chunk_num); + params.set(MULTI_CHUNK_CLUSTER_MAX_ITERATIONS, max_iterations); + + ret = chunk_cluster_.init(meta, params); + if (ret != 0) { + LOG_ERROR("Failed to get chunk cluster"); + return IndexError_InvalidArgument; + } + + ret = chunk_cluster_.mount(train_features); + if (ret != 0) { + LOG_ERROR("Cannot mount train features"); + return ret; + } + + + std::vector labels; + + ret = chunk_cluster_.cluster(threads, cluster_centroids_); + if (ret != 0) { + LOG_ERROR("Failed to cluster"); + return ret; + } + + chunk_offsets = chunk_cluster_.chunk_dim_offsets(); + auto chunk_dims = chunk_cluster_.chunk_dims(); + + switch (type) { + case IndexMeta::DataType::DT_FP32: + ret = convert_pivot_data(meta, num_centers, pq_chunk_num, + chunk_dims, chunk_offsets, + cluster_centroids_, full_pivot_data); + if (ret != 0) { + LOG_ERROR("Failed to convert pivot data"); + return ret; + } + break; + + case IndexMeta::DataType::DT_FP16: + ret = convert_pivot_data( + meta, num_centers, pq_chunk_num, chunk_dims, chunk_offsets, + cluster_centroids_, full_pivot_data); + if (ret != 0) { + LOG_ERROR("Failed to convert pivot data"); + return ret; + } + break; + } + + return 0; +} + +int DiskAnnPqTrainer::train_quantized_data( + IndexThreads::Pointer threads, IndexHolder::Pointer holder, + const IndexMeta &meta, std::vector &pq_full_pivot_data, + std::vector &pq_centroid, std::vector &pq_chunk_offsets, + size_t pq_chunk_num) { + size_t train_size; + std::string train_data; + + int ret = gen_random_sample(holder, meta, train_data, train_size); + if (ret != 0) { + LOG_ERROR("Get Random Sample Error, ret: %d", ret); + return ret; + } + + LOG_INFO("Training data with %zu samples loaded.", train_size); + + // bool use_zero_mean = (meta.metric_name() != "InnerProduct" ? true : + // false); + bool use_zero_mean = false; + + ret = train_pq(threads, meta, train_data, train_size, PQTable::kPQCentroidNum, + pq_chunk_num, PQTable::kMeanIterNum, use_zero_mean, + pq_full_pivot_data, pq_centroid, pq_chunk_offsets); + if (ret != 0) { + LOG_ERROR("Train PQ Error, ret: %d", ret); + return ret; + } + + return 0; +} + +int DiskAnnPqTrainer::generate_pq(IndexThreads::Pointer threads, + const IndexMeta &meta, + IndexHolder::Pointer holder, + uint32_t pq_chunk_num, + std::vector ¢roid, + std::vector &block_compressed_data) { + uint32_t type = meta.data_type(); + uint32_t dim = meta.dimension(); + + if (pq_chunk_num > dim) { + LOG_ERROR("Error: number of chunks more than dimension. chunk: %u, dim: %u", + pq_chunk_num, dim); + return IndexError_InvalidArgument; + } + + // Do Label + std::vector labels; + size_t num_vecs = holder->count(); + size_t batch_size = + num_vecs <= compress_batch_size_ ? num_vecs : compress_batch_size_; + + std::vector block_compressed_base(batch_size * pq_chunk_num); + + std::memset(&block_compressed_base[0], 0, + batch_size * pq_chunk_num * sizeof(uint32_t)); + + std::vector block_data(batch_size * meta.element_size()); + std::vector block_data_converted(batch_size * meta.element_size()); + + size_t block_num = DiskAnnUtil::div_round_up(num_vecs, batch_size); + + block_compressed_data.resize(num_vecs * pq_chunk_num); + + auto iter = holder->create_iterator(); + if (!iter) { + LOG_ERROR("Create iterator for holder failed"); + return IndexError_Runtime; + } + + for (size_t block = 0; block < block_num; block++) { + size_t start_id = block * batch_size; + size_t end_id = std::min((block + 1) * batch_size, num_vecs); + + size_t cur_block_size = end_id - start_id; + + for (size_t i = 0; i < cur_block_size && iter->is_valid(); i++) { + const void *vec = iter->data(); + std::memcpy( + reinterpret_cast(&block_data[0]) + i * meta.element_size(), + vec, meta.element_size()); + iter->next(); + } + + std::memcpy(block_data_converted.data(), block_data.data(), + cur_block_size * meta.element_size()); + + LOG_INFO("Processing Docs, Range: [%zu, %zu)..", start_id, end_id); + + std::shared_ptr block_features( + new CompactIndexFeatures(meta)); + + switch (type) { + case IndexMeta::DataType::DT_FP32: + DiskAnnUtil::convert_vector_to_residual( + reinterpret_cast(block_data_converted.data()), + cur_block_size, dim, centroid.data()); + break; + case IndexMeta::DataType::DT_FP16: + DiskAnnUtil::convert_vector_to_residual( + reinterpret_cast(block_data_converted.data()), + cur_block_size, dim, centroid.data()); + break; + default: + return IndexError_InvalidArgument; + } + + for (size_t i = 0; i < cur_block_size; i++) { + block_features->emplace(block_data_converted.data() + + i * meta.element_size()); + } + + int ret = chunk_cluster_.mount(block_features); + if (ret != 0) { + LOG_ERROR("Cannot mount block features"); + return ret; + } + + ret = chunk_cluster_.label(threads, cluster_centroids_, &labels); + if (ret != 0) { + LOG_ERROR("Failed to label"); + return ret; + } + + std::vector compressed_data(cur_block_size * pq_chunk_num); + + DiskAnnUtil::convert_types_uint32_to_uint8( + labels.data(), compressed_data.data(), cur_block_size, pq_chunk_num); + + memcpy(&(block_compressed_data[0]) + start_id * pq_chunk_num, + compressed_data.data(), cur_block_size * pq_chunk_num); + + LOG_INFO("Generate PQ Data Done."); + } + + return 0; +} + +int DiskAnnPqTrainer::generate_quantized_data( + IndexThreads::Pointer threads, IndexHolder::Pointer holder, + const IndexMeta &meta, std::vector &pq_centroid, + std::vector &block_compressed_data, size_t pq_chunk_num) { + // bool use_zero_mean = (meta.metric_name() != "InnerProduct" ? true : + // false); + + int ret = generate_pq(threads, meta, holder, pq_chunk_num, pq_centroid, + block_compressed_data); + if (ret != 0) { + LOG_ERROR("Generate PQ Error, ret: %d", ret); + return ret; + } + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_pq_trainer.h b/src/core/algorithm/diskann/diskann_pq_trainer.h new file mode 100644 index 0000000..8a83c28 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_pq_trainer.h @@ -0,0 +1,87 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "diskann_entity.h" +#include "diskann_pq_table.h" +#include "../cluster/multi_chunk_cluster.h" + +namespace zvec { +namespace core { + +class DiskAnnPqTrainer { + public: + typedef std::unique_ptr UPointer; + + public: + DiskAnnPqTrainer(uint32_t max_train_sample_count); + virtual ~DiskAnnPqTrainer(); + + public: + template + static int prepare_pq_train_data( + const IndexMeta &meta, size_t num_train, std::string &train_data, + bool use_zero_mean, std::vector ¢roid, + std::shared_ptr &train_features); + + template + static int convert_pivot_data(const IndexMeta &meta, uint32_t num_centers, + uint32_t pq_chunk_num, + const std::vector &chunk_dims, + const std::vector &chunk_offsets, + IndexCluster::CentroidList ¢roids, + std::vector &full_pivot_data); + + int gen_random_sample(IndexHolder::Pointer holder, const IndexMeta &meta, + std::string &sample_data, size_t &sample_size); + + int generate_quantized_data(IndexThreads::Pointer threads, + IndexHolder::Pointer holder, + const IndexMeta &meta, + std::vector &pq_centroid, + std::vector &block_compressed_data, + size_t num_pq_chunks); + + int generate_pq(IndexThreads::Pointer threads, const IndexMeta &meta, + IndexHolder::Pointer holder, uint32_t num_pq_chunks, + std::vector ¢roid, + std::vector &block_compressed_data); + + int train_quantized_data(IndexThreads::Pointer threads, + IndexHolder::Pointer holder, const IndexMeta &meta, + std::vector &pq_full_pivot_data, + std::vector &pq_centroid, + std::vector &pq_chunk_offsets, + size_t num_pq_chunks); + + int train_pq(IndexThreads::Pointer threads, const IndexMeta &meta, + std::string &train_data, size_t num_train, uint32_t num_centers, + uint32_t num_pq_chunks, uint32_t max_iterations, + bool use_zero_mean, std::vector &full_pivot_data, + std::vector ¢roid, + std::vector &chunk_offsets); + + private: + static constexpr uint32_t compress_batch_size_{ + DiskAnnEntity::kDefaultCompressBatchSize}; + + // pq cluster + MultiChunkCluster chunk_cluster_; + IndexCluster::CentroidList cluster_centroids_; + uint32_t max_train_sample_count_{PQTable::kMaxTrainSampleCount}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_reducer.cc b/src/core/algorithm/diskann/diskann_reducer.cc new file mode 100644 index 0000000..b0eabcc --- /dev/null +++ b/src/core/algorithm/diskann/diskann_reducer.cc @@ -0,0 +1,190 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_reducer.h" +#include +#include +#include +#include +#include +#include "diskann_params.h" + +namespace zvec { +namespace core { + +int DiskAnnReducer::init(const ailego::Params ¶ms) { + params.get(PARAM_DISKANN_REDUCER_WORKING_PATH, &working_path_); + if (working_path_.empty()) { + LOG_ERROR("Missing parameter. %s", + PARAM_DISKANN_REDUCER_WORKING_PATH.c_str()); + return IndexError_InvalidArgument; + } + + std::string index_name = + params.get_as_string(PARAM_DISKANN_REDUCER_INDEX_NAME); + if (index_name.empty()) { + index_name = std::to_string(std::clock()); + } + + reducer_file_path_ = ailego::StringHelper::Concat( + working_path_, "/", kReducerFileName, index_name); + + holder_file_path_ = ailego::StringHelper::Concat(working_path_, "/", + kHolderFileName, index_name); + + state_ = STATE_INITED; + return 0; +} + +int DiskAnnReducer::cleanup(void) { + return 0; +} + +//! Reduce operator with filter +int DiskAnnReducer::reduce(const IndexFilter &filter) { + if (entities_.empty() || state_ != STATE_FEED) { + LOG_ERROR("No container to reduce, feed first"); + return IndexError_NoReady; + } + + if (use_mem_holder_) { + mem_holder_ = std::make_shared(meta_); + for (auto entity : entities_) { + size_t doc_cnt = entity->doc_cnt(); + for (size_t id = 0; id < doc_cnt; ++id) { + diskann_key_t pkey = entity->get_key(id); + + if (filter.is_valid() && filter(pkey)) { + continue; + } + + const void *vec = entity->get_vector(id); + mem_holder_->emplace(pkey, vec); + } + } + } else { + disk_holder_ = + std::make_shared(meta_, holder_file_path_); + + int ret = disk_holder_->init(); + if (ret != 0) { + LOG_ERROR("DiskAnn Index Holder init failed"); + return ret; + } + + for (auto entity : entities_) { + size_t doc_cnt = entity->doc_cnt(); + for (size_t id = 0; id < doc_cnt; ++id) { + diskann_key_t pkey = entity->get_key(id); + + if (filter.is_valid() && filter(pkey)) { + continue; + } + + const void *vec = entity->get_vector(id); + disk_holder_->emplace(pkey, vec); + } + } + + disk_holder_->close(); + } + + builder_ = IndexFactory::CreateBuilder(kDiskAnnBuilderName); + if (!builder_) { + LOG_ERROR("Create builder failed. name[%s]", kDiskAnnBuilderName.c_str()); + return IndexError_Runtime; + } + + if (thread_pool_ == nullptr) { + LOG_ERROR( + "Only support multi-thread mode. Thread pool is not set for reducer."); + return IndexError_NoReady; + } + + LOG_INFO("Start diskann reduce"); + + ailego::ElapsedTime timer; + + auto params = meta_.builder_params(); + + int ret = builder_->init(meta_, params); + if (ret != 0) { + LOG_ERROR("Init proxima streamer failed. ret[%d]", ret); + return ret; + } + + if (use_mem_holder_) { + ret = builder_->train(mem_holder_); + if (ret != 0) { + LOG_ERROR("Diskann builder failed to train. ret[%d]", ret); + return ret; + } + + ret = builder_->build(mem_holder_); + if (ret != 0) { + LOG_ERROR("Diskann builder failed to build. ret[%d]", ret); + return ret; + } + } else { + ret = builder_->train(disk_holder_); + if (ret != 0) { + LOG_ERROR("Diskann builder failed to train. ret[%d]", ret); + return ret; + } + + ret = builder_->build(disk_holder_); + if (ret != 0) { + LOG_ERROR("Diskann builder failed to build. ret[%d]", ret); + return ret; + } + } + + auto &stats = builder_->stats(); + + stats_.set_reduced_costtime(timer.seconds()); + stats_.set_filtered_count(stats.discarded_count()); + + state_ = STATE_REDUCE; + + LOG_INFO("End DiskAnn reduce. cost time: [%zu]s", (size_t)timer.seconds()); + return 0; +} + +//! Dump index by dumper +int DiskAnnReducer::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("Begin diskann reducer dump"); + + if (state_ != STATE_REDUCE) { + LOG_WARN("Reduce first before dump."); + return IndexError_NoReady; + } + + ailego::ElapsedTime timer; + int ret = builder_->dump(dumper); + if (ret != 0) { + LOG_ERROR("diskann reducer dump failed. ret[%d]", ret); + return ret; + } + stats_.set_dumped_costtime(timer.seconds()); + + LOG_INFO("End diskann reducer dump, dump costtime=[%zu]s", + (size_t)(stats_.dumped_costtime())); + + return 0; +} + +INDEX_FACTORY_REGISTER_REDUCER(DiskAnnReducer); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_reducer.h b/src/core/algorithm/diskann/diskann_reducer.h new file mode 100644 index 0000000..b912ea1 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_reducer.h @@ -0,0 +1,85 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include "diskann_holder.h" +#include "diskann_reducer_entity.h" + +namespace zvec { +namespace core { + +class DiskAnnReducer : public IndexReducer { + public: + //! Constructor + DiskAnnReducer(void) = default; + + protected: + //! Initialize Reducer + virtual int init(const ailego::Params ¶ms) override; + + //! Cleanup Reducer + virtual int cleanup(void) override; + + //! Feed indexes from containers + // virtual int feed(IndexStorage::Pointer container) override; + + //! Reduce operator (with filter) + virtual int reduce(const IndexFilter &filter) override; + + //! Dump index by dumper + virtual int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + virtual const Stats &stats(void) const override { + return stats_; + } + + private: + enum State { + STATE_UNINITED = 0, + STATE_INITED = 1, + STATE_FEED = 2, + STATE_REDUCE = 3 + }; + + std::string working_path_{""}; + + IndexMeta meta_{}; + std::vector entities_{}; + + // bool use_mem_holder_{true}; + bool use_mem_holder_{false}; + RandomAccessIndexHolder::Pointer mem_holder_; + DiskAnnIndexHolder::Pointer disk_holder_; + + IndexBuilder::Pointer builder_{nullptr}; + std::string reducer_file_path_{""}; + std::string holder_file_path_{""}; + + Stats stats_{}; + State state_{STATE_UNINITED}; + + const std::string kDiskAnnBuilderName{"DiskAnnBuilder"}; + const std::string kReducerFileName{"diskann.reducer.builder."}; + const std::string kHolderFileName{"diskann.reducer.holder."}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_reducer_entity.cc b/src/core/algorithm/diskann/diskann_reducer_entity.cc new file mode 100644 index 0000000..4ccfb6d --- /dev/null +++ b/src/core/algorithm/diskann/diskann_reducer_entity.cc @@ -0,0 +1,215 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_reducer_entity.h" +#include +#include + +namespace zvec { +namespace core { + +int DiskAnnReducerEntity::load(const IndexStorage::Pointer &container, + bool check_crc) { + container_ = container; + + int ret = load_segments(check_crc); + if (ret != 0) { + return ret; + } + + sector_num_per_node_ = node_per_sector() > 0 + ? 1 + : DiskAnnUtil::div_round_up( + max_node_size(), DiskAnnUtil::kSectorSize); + + loaded_ = true; + + return 0; +} + +int DiskAnnReducerEntity::load_segments(bool /*check_crc*/) { + int ret; + ret = load_header_segment(); + if (ret != 0) { + LOG_ERROR("Load Header Segment Failed, ret = %d", ret); + + return ret; + } + + ret = load_key_segment(); + if (ret != 0) { + LOG_ERROR("Load Key Segment Failed, ret = %d", ret); + + return ret; + } + + ret = load_vector_segment(); + if (ret != 0) { + LOG_ERROR("Load Vector Segment Failed, ret = %d", ret); + + return ret; + } + + return 0; +} + +int DiskAnnReducerEntity::load_header_segment() { + const void *data = nullptr; + meta_segment_ = container_->get(kDiskAnnMetaSegmentId); + if (!meta_segment_ || + meta_segment_->data_size() < sizeof(DiskAnnMetaHeader)) { + LOG_ERROR("Miss or invalid segment %s", kDiskAnnMetaSegmentId.c_str()); + return IndexError_InvalidFormat; + } + if (meta_segment_->read(0, reinterpret_cast(&data), + sizeof(DiskAnnMetaHeader)) != + sizeof(DiskAnnMetaHeader)) { + LOG_ERROR("Read segment %s failed", kDiskAnnMetaSegmentId.c_str()); + return IndexError_ReadData; + } + + ::memcpy(reinterpret_cast(&meta_header_), data, + sizeof(DiskAnnMetaHeader)); + + return 0; +} + +int DiskAnnReducerEntity::load_vector_segment() { + vector_segment_ = container_->get(kDiskAnnVectorSegmentId); + if (!vector_segment_) { + LOG_ERROR("Miss or invalid segment %s", + DiskAnnEntity::kDiskAnnVectorSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + return 0; +} + +int DiskAnnReducerEntity::load_key_segment() { + // load key + key_segment_ = container_->get(kDiskAnnKeySegmentId); + if (!key_segment_) { + LOG_ERROR("Miss or invalid segment %s", + DiskAnnEntity::kDiskAnnKeySegmentId.c_str()); + return IndexError_InvalidFormat; + } + + size_t key_data_len = doc_cnt() * sizeof(key_t); + + // load key mapping + key_mapping_segment_ = container_->get(kDiskAnnKeyMappingSegmentId); + const void *data = nullptr; + if (key_mapping_segment_->read(0, reinterpret_cast(&data), + key_data_len) != key_data_len) { + LOG_ERROR("Read segment %s failed", kDiskAnnKeyMappingSegmentId.c_str()); + return IndexError_ReadData; + } + + key_buffer_.resize(key_data_len); + memcpy(&(key_buffer_[0]), data, key_data_len); + + return 0; +} + +bool DiskAnnReducerEntity::do_crc_check( + std::vector &segments) const { + constexpr size_t blk_size = 4096; + const void *data; + + for (auto &segment : segments) { + size_t offset = 0; + size_t rd_size; + uint32_t crc = 0; + while (offset < segment->data_size()) { + size_t size = std::min(blk_size, segment->data_size() - offset); + if ((rd_size = segment->read(offset, &data, size)) <= 0) { + break; + } + offset += rd_size; + crc = ailego::Crc32c::Hash(data, rd_size, crc); + } + if (crc != segment->data_crc()) { + return false; + } + } + return true; +} + +//! Get vector local id by key +diskann_id_t DiskAnnReducerEntity::get_id(diskann_key_t key) const { + const diskann_id_t *key_mapping_data_ptr = + reinterpret_cast(key_mapping_buffer_.data()); + const diskann_key_t *key_data_ptr = + reinterpret_cast(key_buffer_.data()); + + //! Do binary search + diskann_id_t start = 0UL; + diskann_id_t end = doc_cnt(); + diskann_id_t idx = 0u; + while (start < end) { + idx = start + (end - start) / 2; + diskann_id_t local_id = key_mapping_data_ptr[idx]; + + const diskann_key_t local_key = key_data_ptr[local_id]; + + if (local_key < key) { + start = idx + 1; + } else if (local_key > key) { + end = idx; + } else { + return local_id; + } + } + + return kInvalidId; +} + +diskann_key_t DiskAnnReducerEntity::get_key(diskann_id_t id) const { + const void *key; + if (ailego_unlikely(key_segment_->read(id * sizeof(diskann_key_t), &key, + sizeof(diskann_key_t)) != + sizeof(diskann_key_t))) { + LOG_ERROR("Read key from segment failed"); + return kInvalidKey; + } + + return *(reinterpret_cast(key)); +} + +const void *DiskAnnReducerEntity::get_vector(diskann_id_t id) const { + size_t read_size = sector_num_per_node_ * DiskAnnUtil::kSectorSize; + size_t sector_id = DiskAnnUtil::get_node_sector( + node_per_sector(), max_node_size(), DiskAnnUtil::kSectorSize, id); + size_t offset = sector_id * DiskAnnUtil::kSectorSize; + + if (sector_id != sector_id_) { + const void *sector_data; + if (ailego_unlikely(vector_segment_->read(offset, §or_data, + read_size) != read_size)) { + LOG_ERROR("Read vector from segment failed"); + return nullptr; + } + + sector_id_ = sector_id; + sector_buffer_.assign(reinterpret_cast(sector_data), + read_size); + } + + return DiskAnnUtil::offset_to_node_const( + node_per_sector(), max_node_size(), + reinterpret_cast(sector_buffer_.data()), id); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_reducer_entity.h b/src/core/algorithm/diskann/diskann_reducer_entity.h new file mode 100644 index 0000000..3ab8930 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_reducer_entity.h @@ -0,0 +1,69 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include "diskann_entity.h" +#include "diskann_file_reader.h" +#include "diskann_pq_table.h" +#include "diskann_util.h" + +namespace zvec { +namespace core { + +class DiskAnnReducerEntity : public DiskAnnEntity { + public: + using Pointer = std::shared_ptr; + using SegmentPointer = IndexStorage::Segment::Pointer; + + public: + DiskAnnReducerEntity() = default; + virtual ~DiskAnnReducerEntity() = default; + + int load(const IndexStorage::Pointer &container, bool check_crc); + int load_segments(bool check_crc); + int load_header_segment(); + int load_vector_segment(); + int load_key_segment(); + int load_key_mapping_segment(); + + bool do_crc_check(std::vector &segments) const; + + diskann_id_t get_id(diskann_key_t key) const override; + diskann_key_t get_key(diskann_id_t id) const override; + const void *get_vector(diskann_id_t id) const override; + + private: + IndexStorage::Pointer container_{}; + IndexStorage::Segment::Pointer meta_segment_{}; + IndexStorage::Segment::Pointer vector_segment_{}; + IndexStorage::Segment::Pointer key_segment_{}; + IndexStorage::Segment::Pointer key_mapping_segment_{}; + + std::string key_buffer_; + std::string key_mapping_buffer_; + + size_t sector_num_per_node_{0}; + + mutable size_t sector_id_{-1U}; + mutable std::string sector_buffer_; + + bool loaded_{false}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc new file mode 100644 index 0000000..905fdf2 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -0,0 +1,312 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_searcher.h" +#include "diskann_context.h" +#include "diskann_indexer.h" +#include "diskann_params.h" + +namespace zvec { +namespace core { + +DiskAnnSearcher::DiskAnnSearcher() {} + +DiskAnnSearcher::~DiskAnnSearcher() {} + +int DiskAnnSearcher::init(const ailego::Params &search_params) { + search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); + search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + return 0; +} + +void DiskAnnSearcher::print_debug_info() {} + +int DiskAnnSearcher::cleanup() { + LOG_INFO("Begin DiskAnnSearcher:cleanup"); + + LOG_INFO("End DiskAnnSearcher:cleanup"); + + return 0; +} + +int DiskAnnSearcher::load(IndexStorage::Pointer storage, + IndexMetric::Pointer measure) { + LOG_INFO("DiskAnnSearcher::load Begin"); + + auto start_time = ailego::Monotime::MilliSeconds(); + + int ret = IndexHelper::DeserializeFromStorage(storage.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize meta from storage"); + return ret; + } + + ret = entity_.load(meta_, storage); + if (ret != 0) { + LOG_INFO("Searcher Entity Load Failed"); + return ret; + } + + diskann_indexer_ = std::make_shared(meta_); + + int res = diskann_indexer_->init(entity_); + if (res != 0) { + return res; + } + + if (cache_nodes_num_ != 0) { + std::vector node_list; + LOG_INFO("Caching %u nodes around medoid(s)", cache_nodes_num_); + + diskann_indexer_->cache_bfs_levels(cache_nodes_num_, node_list); + + diskann_indexer_->load_cache_list(node_list); + + node_list.clear(); + node_list.shrink_to_fit(); + } + + if (measure) { + measure_ = measure; + } else { + measure_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!measure_) { + LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = measure_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMetric init failed, ret=%d", ret); + return ret; + } + if (measure_->query_metric()) { + measure_ = measure_->query_metric(); + } + } + + stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time); + state_ = STATE_LOADED; + + magic_ = IndexContext::GenerateMagic(); + + LOG_INFO("DiskAnnSearcher::load Done"); + + return 0; +} + +int DiskAnnSearcher::unload() { + LOG_INFO("DiskAnnSearcher unload index"); + + state_ = STATE_INITED; + + return 0; +} + +int DiskAnnSearcher::update_context(DiskAnnContext *ctx) const { + const DiskAnnEntity::Pointer entity = entity_.clone(); + if (!entity) { + LOG_ERROR("Failed to clone search context entity"); + return IndexError_Runtime; + } + + return ctx->update_context(DiskAnnContext::kSearcherContext, meta_, measure_, + entity, magic_); +} + +int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + // do search + if (ailego_unlikely(!query || !context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + + DiskAnnContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to DiskAnnContext failed"); + return IndexError_Cast; + } + + // Context is pooled per index type. When switching between DiskAnn indexes + // with different element sizes (e.g., fp16 vs fp32), the cached context has + // undersized buffers. Recreate it to ensure correct buffer allocations. + if (ctx->magic() != magic_) { + uint32_t saved_topk = ctx->topk(); + context = create_context(); + if (!context) { + LOG_ERROR("Failed to recreate context for current streamer"); + return IndexError_Runtime; + } + ctx = dynamic_cast(context.get()); + ctx->set_topk(saved_topk); + } + + ctx->clear(); + ctx->resize_results(count); + + for (uint32_t i = 0; i < count; i++) { + ctx->reset_query(query); + + diskann_indexer_->knn_search(ctx); + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + ctx->topk_to_result(i); + + query = static_cast(query) + qmeta.element_size(); + } + + return 0; +} + +int DiskAnnSearcher::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (ailego_unlikely(!query || !context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + + DiskAnnContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to DiskAnnContext failed"); + return IndexError_Cast; + } + + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer, recreate it + //! to ensure buffers are correctly sized for this index's parameters. + uint32_t saved_topk = ctx->topk(); + context = create_context(); + if (!context) { + LOG_ERROR("Failed to recreate context for current streamer"); + return IndexError_Runtime; + } + ctx = dynamic_cast(context.get()); + ctx->set_topk(saved_topk); + } + + ctx->clear(); + ctx->resize_results(count); + + for (size_t i = 0; i < count; ++i) { + ctx->reset_query(query); + + diskann_indexer_->linear_search(ctx); + + ctx->topk_to_result(i); + + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +int DiskAnnSearcher::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (ailego_unlikely(!query || !context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + + DiskAnnContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to DiskAnnContext failed"); + return IndexError_Cast; + } + + if (ailego_unlikely(p_keys.size() != count)) { + LOG_ERROR("The size of p_keys is not equal to count"); + return IndexError_InvalidArgument; + } + + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer, recreate it + //! to ensure buffers are correctly sized for this index's parameters. + uint32_t saved_topk = ctx->topk(); + context = create_context(); + if (!context) { + LOG_ERROR("Failed to recreate context for current streamer"); + return IndexError_Runtime; + } + ctx = dynamic_cast(context.get()); + ctx->set_topk(saved_topk); + } + + ctx->clear(); + ctx->resize_results(count); + + for (size_t i = 0; i < count; ++i) { + ctx->reset_query(query); + + diskann_indexer_->keys_search(p_keys[i], ctx); + + ctx->topk_to_result(i); + + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +int DiskAnnSearcher::get_vector(uint64_t key, Context::Pointer &context, + std::string &vector) const { + return diskann_indexer_->get_vector(key, context, vector); +} + +IndexSearcher::Context::Pointer DiskAnnSearcher::create_context() const { + const DiskAnnEntity::Pointer search_ctx_entity = entity_.clone(); + if (!search_ctx_entity) { + LOG_ERROR("Failed to create search context entity"); + return Context::Pointer(); + } + + DiskAnnContext *ctx = + new (std::nothrow) DiskAnnContext(meta_, measure_, search_ctx_entity); + if (ctx == nullptr) { + LOG_ERROR("Failed to allocate DiskAnn Context"); + return Context::Pointer(); + } + if (ailego_unlikely(ctx->init( + DiskAnnContext::kSearcherContext, search_ctx_entity->max_degree(), + search_ctx_entity->pq_chunk_num(), meta_.element_size())) != 0) { + LOG_ERROR("Init DiskAnn Context failed"); + delete ctx; + + return Context::Pointer(); + } + + ctx->set_list_size(list_size_); + ctx->set_magic(magic_); + + return Context::Pointer(ctx); +} + +INDEX_FACTORY_REGISTER_SEARCHER(DiskAnnSearcher); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_searcher.h b/src/core/algorithm/diskann/diskann_searcher.h new file mode 100644 index 0000000..99584fa --- /dev/null +++ b/src/core/algorithm/diskann/diskann_searcher.h @@ -0,0 +1,166 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "diskann_context.h" +#include "diskann_indexer.h" + +class LinuxAlignedFileReader; + +namespace zvec { +namespace core { + +class DiskAnnSearcher : public IndexSearcher { + public: + using ContextPointer = IndexSearcher::Context::Pointer; + + public: + DiskAnnSearcher(void); + ~DiskAnnSearcher(void); + + DiskAnnSearcher(const DiskAnnSearcher &) = delete; + DiskAnnSearcher &operator=(const DiskAnnSearcher &) = delete; + + protected: + //! Initialize Searcher + int init(const ailego::Params ¶ms) override; + + //! Cleanup Searcher + int cleanup(void) override; + + //! Load Index from storage + int load(IndexStorage::Pointer storage, IndexMetric::Pointer metric) override; + + //! Unload index from storage + int unload(void) override; + + //! KNN Search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_impl(query, qmeta, 1, context); + } + + //! KNN Search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, ContextPointer &context) const override; + + //! Linear Search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_impl(query, qmeta, 1, context); + } + + //! Linear Search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, ContextPointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(query, &sparse_count, sparse_indices, + sparse_query, p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl( + const void * /*query*/, const uint32_t * /*sparse_count*/, + const uint32_t * /*sparse_indices*/, const void * /*sparse_query*/, + const std::vector> & /*p_keys*/, + const IndexQueryMeta & /*qmeta*/, uint32_t /*count*/, + ContextPointer & /*context*/) const override { + return IndexError_NotImplemented; + } + + //! Get vector by key + int get_vector(uint64_t key, Context::Pointer &context, + std::string &vector) const override; + + //! Create a searcher context + ContextPointer create_context() const override; + + //! Create a new iterator + IndexSearcher::Provider::Pointer create_provider(void) const override { + return nullptr; + } + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve params of index + const ailego::Params ¶ms(void) const override { + return params_; + } + + void print_debug_info() override; + + private: + template + int search_disk_index(const std::string &query_file, + const uint32_t num_nodes_to_cache, + const uint32_t recall_at, const uint32_t beamwidth); + + //! To share ctx across streamer/searcher, we need to update the context for + //! current streamer/searcher + int update_context(DiskAnnContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; + + IndexMetric::Pointer measure_{}; + IndexMeta meta_{}; + ailego::Params params_{}; + + uint32_t list_size_{200}; + uint32_t cache_nodes_num_{0}; + + bool warm_up_{false}; + uint32_t beam_size_{2}; + + DiskAnnIndexer::Pointer diskann_indexer_{nullptr}; + DiskAnnSearcherEntity entity_{}; + + uint32_t magic_{0U}; + + Stats stats_; + State state_{STATE_INIT}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.cc b/src/core/algorithm/diskann/diskann_searcher_entity.cc new file mode 100644 index 0000000..c9e49de --- /dev/null +++ b/src/core/algorithm/diskann/diskann_searcher_entity.cc @@ -0,0 +1,441 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_searcher_entity.h" + +namespace zvec { +namespace core { + +const DiskAnnEntity::Pointer DiskAnnSearcherEntity::clone() const { + auto meta_segment = meta_segment_->clone(); + if (ailego_unlikely(!meta_segment)) { + LOG_ERROR("clone segment %s failed", kDiskAnnMetaSegmentId.c_str()); + return DiskAnnEntity::Pointer(); + } + + auto pq_meta_segment = pq_meta_segment_->clone(); + if (ailego_unlikely(!pq_meta_segment)) { + LOG_ERROR("clone segment %s failed", kDiskAnnPqMetaSegmentId.c_str()); + return DiskAnnEntity::Pointer(); + } + + auto pq_data_segment = pq_data_segment_->clone(); + if (ailego_unlikely(!pq_data_segment)) { + LOG_ERROR("clone segment %s failed", kDiskAnnPqDataSegmentId.c_str()); + return DiskAnnEntity::Pointer(); + } + + auto vector_segment = vector_segment_->clone(); + if (ailego_unlikely(!vector_segment)) { + LOG_ERROR("clone segment %s failed", kDiskAnnVectorSegmentId.c_str()); + return DiskAnnEntity::Pointer(); + } + + auto key_segment = key_segment_->clone(); + if (ailego_unlikely(!key_segment)) { + LOG_ERROR("clone segment %s failed", kDiskAnnKeySegmentId.c_str()); + return DiskAnnEntity::Pointer(); + } + + auto key_mapping_segment = key_mapping_segment_->clone(); + if (ailego_unlikely(!key_mapping_segment)) { + LOG_ERROR("clone segment %s failed", kDiskAnnKeyMappingSegmentId.c_str()); + return DiskAnnEntity::Pointer(); + } + + auto entrypoint_segment = entrypoint_segment_->clone(); + if (ailego_unlikely(!entrypoint_segment)) { + LOG_ERROR("clone segment %s failed", kDiskAnnEntryPointSegmentId.c_str()); + return DiskAnnEntity::Pointer(); + } + + DiskAnnSearcherEntity *entity = new (std::nothrow) DiskAnnSearcherEntity( + meta_header_, pq_meta_, meta_segment, pq_meta_segment, pq_data_segment, + vector_segment, key_segment, key_mapping_segment, entrypoint_segment, + num_threads_, list_size_, cache_nodes_num_, warm_up_, beam_size_, meta_, + pq_table_, key_buffer_, key_mapping_buffer_, entrypoints_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("DiskAnnSearcherEntity new failed"); + } + + return DiskAnnEntity::Pointer(entity); +} + +int DiskAnnSearcherEntity::load(const IndexMeta &meta, + IndexStorage::Pointer storage) { + meta_ = meta; + + storage_ = storage; + + int ret; + ret = load_header_segment(); + if (ret != 0) { + LOG_ERROR("Load Header Segment Failed, ret = %d", ret); + + return ret; + } + + ret = load_pq_segment(); + if (ret != 0) { + LOG_ERROR("Load PQ Meta Segment Failed, ret = %d", ret); + + return ret; + } + + ret = load_key_segment(); + if (ret != 0) { + LOG_ERROR("Load Key Segment Failed, ret = %d", ret); + + return ret; + } + + ret = load_key_mapping_segment(); + if (ret != 0) { + LOG_ERROR("Load Key Segment Failed, ret = %d", ret); + + return ret; + } + + ret = load_entrypoint_segment(); + if (ret != 0) { + LOG_WARN("Load EntryPoint Segment Failed, ret = %d", ret); + + return ret; + } + + ret = load_vector_segment(); + if (ret != 0) { + LOG_ERROR("Load Vector Segment Failed, ret = %d", ret); + + return ret; + } + + return 0; +} + +int DiskAnnSearcherEntity::load_pq_segment() { + const void *data = nullptr; + + // load pq meta + pq_meta_segment_ = storage_->get(DiskAnnEntity::kDiskAnnPqMetaSegmentId); + if (!pq_meta_segment_) { + LOG_ERROR("Miss or invalid segment %s", + DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + size_t read_size; + size_t offset = 0; + + // 1. read pq meta + read_size = pq_meta_segment_->read(offset, &data, sizeof(DiskAnnPqMeta)); + if (read_size != sizeof(DiskAnnPqMeta)) { + LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", + DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), + sizeof(DiskAnnPqMeta), read_size); + + return IndexError_ReadData; + } + + memcpy(reinterpret_cast(&pq_meta_), data, sizeof(DiskAnnPqMeta)); + offset += read_size; + + // 2. read full pivot data + std::vector full_pivot_data; + full_pivot_data.resize(pq_meta_.full_pivot_data_size); + + read_size = + pq_meta_segment_->read(offset, &data, pq_meta_.full_pivot_data_size); + if (read_size != pq_meta_.full_pivot_data_size) { + LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", + DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), + (size_t)(pq_meta_.full_pivot_data_size), (size_t)read_size); + return IndexError_ReadData; + } + memcpy(&(full_pivot_data[0]), data, read_size); + offset += read_size; + + // 3. read centroid + std::vector centroid; + centroid.resize(pq_meta_.centroid_data_size); + + read_size = + pq_meta_segment_->read(offset, &data, pq_meta_.centroid_data_size); + if (read_size != pq_meta_.centroid_data_size) { + LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", + DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), + (size_t)(pq_meta_.centroid_data_size), (size_t)read_size); + return IndexError_ReadData; + } + memcpy(&(centroid[0]), data, read_size); + offset += read_size; + + // 4. chunk offset + std::vector chunk_offsets; + chunk_offsets.resize(pq_meta_.chunk_num + 1); + + read_size = pq_meta_segment_->read( + offset, &data, (pq_meta_.chunk_num + 1) * sizeof(uint32_t)); + if (read_size != (pq_meta_.chunk_num + 1) * sizeof(uint32_t)) { + LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", + DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), + (size_t)((pq_meta_.chunk_num + 1) * sizeof(uint32_t)), + (size_t)read_size); + return IndexError_ReadData; + } + memcpy(&(chunk_offsets[0]), data, read_size); + + // load pq data + std::vector pq_data; + pq_data_segment_ = storage_->get(DiskAnnEntity::kDiskAnnPqDataSegmentId); + if (!pq_data_segment_) { + LOG_ERROR("Miss or invalid segment %s", + DiskAnnEntity::kDiskAnnPqDataSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + pq_data.resize(meta_header_.doc_cnt * pq_meta_.chunk_num); + + void *pq_data_ptr = &pq_data[0]; + read_size = pq_data_segment_->fetch( + 0, pq_data_ptr, meta_header_.doc_cnt * pq_meta_.chunk_num); + + if (read_size != meta_header_.doc_cnt * pq_meta_.chunk_num) { + LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", + DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), + (size_t)(meta_header_.doc_cnt * pq_meta_.chunk_num), + (size_t)read_size); + + return IndexError_ReadData; + } + + pq_table_ = std::make_shared(meta_, pq_meta_.chunk_num); + + pq_table_->init(full_pivot_data, centroid, chunk_offsets, pq_data); + + return 0; +} + +int DiskAnnSearcherEntity::load_header_segment() { + const void *data = nullptr; + meta_segment_ = storage_->get(kDiskAnnMetaSegmentId); + if (!meta_segment_ || + meta_segment_->data_size() < sizeof(DiskAnnMetaHeader)) { + LOG_ERROR("Miss or invalid segment %s", kDiskAnnMetaSegmentId.c_str()); + return IndexError_InvalidFormat; + } + if (meta_segment_->read(0, reinterpret_cast(&data), + sizeof(DiskAnnMetaHeader)) != + sizeof(DiskAnnMetaHeader)) { + LOG_ERROR("Read segment %s failed", kDiskAnnMetaSegmentId.c_str()); + return IndexError_ReadData; + } + memcpy(reinterpret_cast(&meta_header_), data, + sizeof(DiskAnnMetaHeader)); + + return 0; +} + +int DiskAnnSearcherEntity::load_vector_segment() { + vector_segment_ = storage_->get(kDiskAnnVectorSegmentId); + if (!vector_segment_) { + LOG_ERROR("Miss or invalid segment %s", + DiskAnnEntity::kDiskAnnVectorSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + return 0; +} + +int DiskAnnSearcherEntity::load_key_segment() { + // load key + key_segment_ = storage_->get(kDiskAnnKeySegmentId); + if (!key_segment_) { + LOG_ERROR("Miss or invalid segment %s", + DiskAnnEntity::kDiskAnnKeySegmentId.c_str()); + return IndexError_InvalidFormat; + } + + size_t key_data_len = doc_cnt() * sizeof(diskann_key_t); + + const void *data = nullptr; + if (key_segment_->read(0, reinterpret_cast(&data), + key_data_len) != key_data_len) { + LOG_ERROR("Read segment %s failed", kDiskAnnKeySegmentId.c_str()); + return IndexError_ReadData; + } + + key_buffer_.resize(key_data_len); + memcpy(&(key_buffer_[0]), data, key_data_len); + + return 0; +} + +int DiskAnnSearcherEntity::load_entrypoint_segment() { + entrypoint_segment_ = storage_->get(kDiskAnnEntryPointSegmentId); + if (!entrypoint_segment_) { + LOG_ERROR("Miss or invalid segment %s", + DiskAnnEntity::kDiskAnnEntryPointSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + const void *data = nullptr; + + if (entrypoint_segment_->read(0, reinterpret_cast(&data), + sizeof(uint32_t)) != sizeof(uint32_t)) { + LOG_ERROR("Read segment %s failed", kDiskAnnEntryPointSegmentId.c_str()); + return IndexError_ReadData; + } + + uint32_t entrypoint_cnt = 0; + memcpy(&entrypoint_cnt, data, sizeof(uint32_t)); + + if (entrypoint_cnt != 0) { + size_t entrypoint_data_len = entrypoint_cnt * sizeof(diskann_id_t); + + if (entrypoint_segment_->read(sizeof(uint32_t), + reinterpret_cast(&data), + entrypoint_data_len) != entrypoint_data_len) { + LOG_ERROR("Read segment %s failed", kDiskAnnEntryPointSegmentId.c_str()); + return IndexError_ReadData; + } + + entrypoints_.resize(entrypoint_cnt); + memcpy(&(entrypoints_[0]), data, entrypoint_data_len); + } + + return 0; +} + + +int DiskAnnSearcherEntity::load_key_mapping_segment() { + key_mapping_segment_ = storage_->get(kDiskAnnKeyMappingSegmentId); + if (!key_mapping_segment_) { + LOG_ERROR("Miss or invalid segment %s", + DiskAnnEntity::kDiskAnnKeyMappingSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + size_t key_mapping_data_len = doc_cnt() * sizeof(diskann_id_t); + + const void *data = nullptr; + if (key_mapping_segment_->read(0, reinterpret_cast(&data), + key_mapping_data_len) != + key_mapping_data_len) { + LOG_ERROR("Read segment %s failed", kDiskAnnKeyMappingSegmentId.c_str()); + return IndexError_ReadData; + } + + key_mapping_buffer_.resize(key_mapping_data_len); + memcpy(&(key_mapping_buffer_[0]), data, key_mapping_data_len); + + return 0; +} + +//! Get vector local id by key +diskann_id_t DiskAnnSearcherEntity::get_id(diskann_key_t key) const { + const diskann_id_t *key_mapping_data_ptr = + reinterpret_cast(key_mapping_buffer_.data()); + + const diskann_key_t *key_data_ptr = + reinterpret_cast(key_buffer_.data()); + + //! Do binary search + diskann_id_t start = 0UL; + diskann_id_t end = doc_cnt(); + diskann_id_t idx = 0u; + while (start < end) { + idx = start + (end - start) / 2; + diskann_id_t local_id = key_mapping_data_ptr[idx]; + + const diskann_key_t local_key = key_data_ptr[local_id]; + + if (local_key < key) { + start = idx + 1; + } else if (local_key > key) { + end = idx; + } else { + return local_id; + } + } + + return kInvalidId; +} + +diskann_key_t DiskAnnSearcherEntity::get_key(diskann_id_t id) const { + const diskann_key_t *key_data_ptr = + reinterpret_cast(key_buffer_.data()); + + return key_data_ptr[id]; +} + +const void *DiskAnnSearcherEntity::get_vector(diskann_id_t id) const { + if (!vector_segment_) { + LOG_ERROR("Vector segment is null"); + return nullptr; + } + + uint64_t sector_offset = + DiskAnnUtil::get_node_sector(node_per_sector(), max_node_size(), + DiskAnnUtil::kSectorSize, id) * + DiskAnnUtil::kSectorSize; + uint64_t within_sector_offset = + (node_per_sector() == 0 ? 0 : (id % node_per_sector()) * max_node_size()); + uint64_t total_offset = sector_offset + within_sector_offset; + + size_t read_size = meta_.element_size(); + const void *vec; + if (ailego_unlikely(vector_segment_->read(total_offset, &vec, read_size) != + read_size)) { + LOG_ERROR("Read vector from segment failed, id: %u, offset: %lu", id, + total_offset); + return nullptr; + } + + return vec; +} + +std::pair DiskAnnSearcherEntity::get_neighbors( + diskann_id_t id) const { + if (!vector_segment_) { + return std::make_pair(0, nullptr); + } + + uint64_t read_sector_offset = + DiskAnnUtil::get_node_sector(node_per_sector(), max_node_size(), + DiskAnnUtil::kSectorSize, id) * + DiskAnnUtil::kSectorSize; + uint64_t node_vec_offset = + read_sector_offset + + (node_per_sector() == 0 ? 0 : (id % node_per_sector()) * max_node_size()); + + const void *data; + if (ailego_unlikely( + vector_segment_->read(node_vec_offset, &data, max_node_size()) != + max_node_size())) { + LOG_ERROR("Read neighbors from segment failed"); + return {0, nullptr}; + } + + const uint8_t *data_ptr = reinterpret_cast(data); + const diskann_id_t *node_neighbor = + reinterpret_cast(data_ptr + meta_.element_size()); + + auto neighbor_num = *node_neighbor; + + return std::make_pair(neighbor_num, node_neighbor + 1); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.h b/src/core/algorithm/diskann/diskann_searcher_entity.h new file mode 100644 index 0000000..68ab498 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_searcher_entity.h @@ -0,0 +1,125 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "diskann_entity.h" +#include "diskann_file_reader.h" +#include "diskann_pq_table.h" + +namespace zvec { +namespace core { + +class DiskAnnSearcherEntity : public DiskAnnEntity { + public: + using Pointer = std::shared_ptr; + using SegmentPointer = IndexStorage::Segment::Pointer; + + public: + DiskAnnSearcherEntity() = default; + virtual ~DiskAnnSearcherEntity() = default; + + public: + virtual const DiskAnnEntity::Pointer clone() const override; + + int load(const IndexMeta &meta, IndexStorage::Pointer storage); + int load_pq_segment(); + int load_header_segment(); + int load_vector_segment(); + int load_key_segment(); + int load_key_mapping_segment(); + int load_entrypoint_segment(); + + PQTable::Pointer get_pq_table() { + return pq_table_; + } + + IndexStorage::Pointer get_storage() { + return storage_; + } + + SegmentPointer get_vector_segment() { + return vector_segment_; + } + + std::vector &entrypoints() { + return entrypoints_; + } + + std::pair get_neighbors( + diskann_id_t id) const override; + + diskann_id_t get_id(diskann_key_t key) const override; + diskann_key_t get_key(diskann_id_t id) const override; + const void *get_vector(diskann_id_t id) const override; + + private: + DiskAnnSearcherEntity( + const DiskAnnMetaHeader &meta_header, const DiskAnnPqMeta &pq_meta, + const SegmentPointer &meta_segment, const SegmentPointer &pq_meta_segment, + const SegmentPointer &pq_data_segment, + const SegmentPointer &vector_segment, const SegmentPointer &key_segment, + const SegmentPointer &key_mapping_segment, + const SegmentPointer &entrypoint_segment, uint32_t num_threads, + uint32_t list_size, uint32_t cache_nodes_num, bool warm_up, + uint32_t beam_size, const IndexMeta meta, PQTable::Pointer pq_table, + const std::string &key_buffer, const std::string &key_mapping_buffer, + const std::vector &entrypoints) + : DiskAnnEntity(meta_header, pq_meta), + meta_segment_(meta_segment), + pq_meta_segment_(pq_meta_segment), + pq_data_segment_(pq_data_segment), + vector_segment_(vector_segment), + key_segment_(key_segment), + key_mapping_segment_(key_mapping_segment), + entrypoint_segment_{entrypoint_segment}, + num_threads_{num_threads}, + list_size_{list_size}, + cache_nodes_num_{cache_nodes_num}, + warm_up_{warm_up}, + beam_size_{beam_size}, + meta_{meta}, + pq_table_{pq_table}, + key_buffer_{key_buffer}, + key_mapping_buffer_{key_mapping_buffer}, + entrypoints_{entrypoints} {} + + IndexStorage::Pointer storage_{}; + + SegmentPointer meta_segment_{nullptr}; + SegmentPointer pq_meta_segment_{nullptr}; + SegmentPointer pq_data_segment_{nullptr}; + SegmentPointer vector_segment_{nullptr}; + SegmentPointer key_segment_{nullptr}; + SegmentPointer key_mapping_segment_{nullptr}; + SegmentPointer entrypoint_segment_{nullptr}; + + uint32_t num_threads_{1}; + uint32_t list_size_{200}; + uint32_t cache_nodes_num_{0}; + + bool warm_up_{false}; + uint32_t beam_size_{2}; + + IndexMeta meta_; + + PQTable::Pointer pq_table_; + std::string key_buffer_; + std::string key_mapping_buffer_; + std::vector entrypoints_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc new file mode 100644 index 0000000..7a978b6 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -0,0 +1,345 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_streamer.h" +#include "diskann_context.h" +#include "diskann_index_provider.h" +#include "diskann_indexer.h" +#include "diskann_params.h" + +namespace zvec { +namespace core { + +DiskAnnStreamer::DiskAnnStreamer() {} + +DiskAnnStreamer::~DiskAnnStreamer() {} + +int DiskAnnStreamer::init(const IndexMeta &meta, + const ailego::Params &search_params) { + meta_ = meta; + search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); + search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + return 0; +} + +void DiskAnnStreamer::print_debug_info() {} + +int DiskAnnStreamer::cleanup() { + LOG_INFO("Begin DiskAnnStreamer:cleanup"); + + LOG_INFO("End DiskAnnStreamer:cleanup"); + + return 0; +} + +int DiskAnnStreamer::open(IndexStorage::Pointer storage) { + LOG_INFO("DiskAnnStreamer::load Begin"); + + auto start_time = ailego::Monotime::MilliSeconds(); + + int ret = IndexHelper::DeserializeFromStorage(storage.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize meta from storage"); + return ret; + } + + ret = entity_.load(meta_, storage); + if (ret != 0) { + LOG_INFO("Searcher Entity Load Failed"); + return ret; + } + + diskann_indexer_ = std::make_shared(meta_); + + int res = diskann_indexer_->init(entity_); + if (res != 0) { + return res; + } + + if (cache_nodes_num_ != 0) { + std::vector node_list; + LOG_INFO("Caching %u nodes around medoid(s)", cache_nodes_num_); + + diskann_indexer_->cache_bfs_levels(cache_nodes_num_, node_list); + + diskann_indexer_->load_cache_list(node_list); + + node_list.clear(); + node_list.shrink_to_fit(); + } + + measure_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!measure_) { + LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = measure_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMetric init failed, ret=%d", ret); + return ret; + } + if (measure_->query_metric()) { + measure_ = measure_->query_metric(); + } + + stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time); + state_ = STATE_LOADED; + + magic_ = IndexContext::GenerateMagic(); + + LOG_INFO("DiskAnnStreamer::load Done"); + + return 0; +} + +int DiskAnnStreamer::unload() { + LOG_INFO("DiskAnnStreamer unload index"); + + state_ = STATE_INITED; + + return 0; +} + +int DiskAnnStreamer::update_context(DiskAnnContext *ctx) const { + const DiskAnnEntity::Pointer entity = entity_.clone(); + if (!entity) { + LOG_ERROR("Failed to clone search context entity"); + return IndexError_Runtime; + } + + return ctx->update_context(DiskAnnContext::kSearcherContext, meta_, measure_, + entity, magic_); +} + +int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + // do search + if (ailego_unlikely(!query || !context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + + DiskAnnContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to DiskAnnContext failed"); + return IndexError_Cast; + } + + // Context is pooled per index type. When switching between DiskAnn indexes + // with different element sizes (e.g., fp16 vs fp32), the cached context has + // undersized buffers. Recreate it to ensure correct buffer allocations. + if (ctx->magic() != magic_) { + uint32_t saved_topk = ctx->topk(); + context = create_context(); + if (!context) { + LOG_ERROR("Failed to recreate context for current streamer"); + return IndexError_Runtime; + } + ctx = dynamic_cast(context.get()); + ctx->set_topk(saved_topk); + } + + ctx->clear(); + ctx->resize_results(count); + + for (uint32_t i = 0; i < count; i++) { + ctx->reset_query(query); + + diskann_indexer_->knn_search(ctx); + + ctx->topk_to_result(i); + + query = static_cast(query) + qmeta.element_size(); + } + + return 0; +} + +int DiskAnnStreamer::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (ailego_unlikely(!query || !context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + + DiskAnnContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to DiskAnnContext failed"); + return IndexError_Cast; + } + + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer, recreate it + //! to ensure buffers are correctly sized for this index's parameters. + uint32_t saved_topk = ctx->topk(); + context = create_context(); + if (!context) { + LOG_ERROR("Failed to recreate context for current streamer"); + return IndexError_Runtime; + } + ctx = dynamic_cast(context.get()); + ctx->set_topk(saved_topk); + } + + ctx->clear(); + ctx->resize_results(count); + + for (size_t i = 0; i < count; ++i) { + ctx->reset_query(query); + + diskann_indexer_->linear_search(ctx); + + ctx->topk_to_result(i); + + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +int DiskAnnStreamer::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (ailego_unlikely(!query || !context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + + DiskAnnContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to DiskAnnContext failed"); + return IndexError_Cast; + } + + if (ailego_unlikely(p_keys.size() != count)) { + LOG_ERROR("The size of p_keys is not equal to count"); + return IndexError_InvalidArgument; + } + + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer, recreate it + //! to ensure buffers are correctly sized for this index's parameters. + uint32_t saved_topk = ctx->topk(); + context = create_context(); + if (!context) { + LOG_ERROR("Failed to recreate context for current streamer"); + return IndexError_Runtime; + } + ctx = dynamic_cast(context.get()); + ctx->set_topk(saved_topk); + } + + ctx->clear(); + ctx->resize_results(count); + + for (size_t i = 0; i < count; ++i) { + ctx->reset_query(query); + + diskann_indexer_->keys_search(p_keys[i], ctx); + + ctx->topk_to_result(i); + + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +int DiskAnnStreamer::get_vector(uint64_t key, Context::Pointer &context, + std::string &vector) const { + return diskann_indexer_->get_vector(key, context, vector); +} + +const void *DiskAnnStreamer::get_vector_by_id(uint32_t /*id*/) const { + // DiskAnn vectors are stored on disk in sector format; + // a const void* access requires sector I/O via create_context + // Return nullptr to indicate this path is not supported. + return nullptr; +} + +int DiskAnnStreamer::get_vector_by_id(const uint32_t id, + IndexStorage::MemoryBlock &block) const { + // Lazily create a reusable context for fetch operations + if (!fetch_ctx_) { + fetch_ctx_ = create_context(); + if (!fetch_ctx_) { + LOG_ERROR("Failed to create context for get_vector_by_id"); + return IndexError_Runtime; + } + } + int ret = diskann_indexer_->get_vector(id, fetch_ctx_, fetch_vector_buffer_); + if (ret != 0) { + LOG_ERROR("Failed to get vector by id: %u", id); + return IndexError_Runtime; + } + block.reset((void *)fetch_vector_buffer_.data()); + return 0; +} + +IndexSearcher::Provider::Pointer DiskAnnStreamer::create_provider(void) const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Load the index first before creating a provider"); + return nullptr; + } + const DiskAnnEntity::Pointer entity = entity_.clone(); + if (!entity) { + LOG_ERROR("Failed to clone DiskAnn entity for provider"); + return nullptr; + } + return IndexProvider::Pointer(new (std::nothrow) DiskAnnIndexProvider( + meta_, entity, "DiskAnnStreamer")); +} + +IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { + const DiskAnnEntity::Pointer search_ctx_entity = entity_.clone(); + if (!search_ctx_entity) { + LOG_ERROR("Failed to create search context entity"); + return Context::Pointer(); + } + + DiskAnnContext *ctx = + new (std::nothrow) DiskAnnContext(meta_, measure_, search_ctx_entity); + if (ctx == nullptr) { + LOG_ERROR("Failed to allocate DiskAnn Context"); + return Context::Pointer(); + } + if (ailego_unlikely(ctx->init( + DiskAnnContext::kSearcherContext, search_ctx_entity->max_degree(), + search_ctx_entity->pq_chunk_num(), meta_.element_size())) != 0) { + LOG_ERROR("Init DiskAnn Context failed"); + delete ctx; + + return Context::Pointer(); + } + + ctx->set_list_size(list_size_); + + return Context::Pointer(ctx); +} + +INDEX_FACTORY_REGISTER_STREAMER(DiskAnnStreamer); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_streamer.h b/src/core/algorithm/diskann/diskann_streamer.h new file mode 100644 index 0000000..8d731f3 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_streamer.h @@ -0,0 +1,180 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "diskann_context.h" +#include "diskann_indexer.h" + +class LinuxAlignedFileReader; + +namespace zvec { +namespace core { + +class DiskAnnStreamer : public IndexStreamer { + public: + using ContextPointer = IndexStreamer::Context::Pointer; + + public: + DiskAnnStreamer(void); + ~DiskAnnStreamer(void); + + DiskAnnStreamer(const DiskAnnStreamer &) = delete; + DiskAnnStreamer &operator=(const DiskAnnStreamer &) = delete; + + protected: + //! Initialize Searcher + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup Searcher + int cleanup(void) override; + + //! Load Index from storage + int open(IndexStorage::Pointer storage) override; + + //! Unload index from storage + int unload(void) override; + + //! KNN Search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_impl(query, qmeta, 1, context); + } + + //! KNN Search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, ContextPointer &context) const override; + + //! Linear Search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_impl(query, qmeta, 1, context); + } + + //! Linear Search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, ContextPointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(query, &sparse_count, sparse_indices, + sparse_query, p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl( + const void * /*query*/, const uint32_t * /*sparse_count*/, + const uint32_t * /*sparse_indices*/, const void * /*sparse_query*/, + const std::vector> & /*p_keys*/, + const IndexQueryMeta & /*qmeta*/, uint32_t /*count*/, + ContextPointer & /*context*/) const override { + return IndexError_NotImplemented; + } + + //! Get vector by key + int get_vector(uint64_t key, Context::Pointer &context, + std::string &vector) const override; + + //! Fetch vector by id + const void *get_vector_by_id(uint32_t id) const override; + + //! Fetch vector by id into memory block + int get_vector_by_id(const uint32_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Create a searcher context + ContextPointer create_context() const override; + + //! Create a vector iterator backed by the on-disk vector segment. + //! Used by the merge code path (``MixedStreamerReducer``) to walk every + //! vector held by this streamer. + IndexSearcher::Provider::Pointer create_provider(void) const override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + virtual int flush(uint64_t /*check_point*/) override { + return 0; + } + + virtual int close(void) override { + return this->unload(); + } + + void print_debug_info() override; + + private: + template + int search_disk_index(const std::string &query_file, + const uint32_t num_nodes_to_cache, + const uint32_t recall_at, const uint32_t beamwidth); + + //! To share ctx across streamer/searcher, we need to update the context for + //! current streamer/searcher + int update_context(DiskAnnContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; + + IndexMetric::Pointer measure_{}; + IndexMeta meta_{}; + ailego::Params params_{}; + + uint32_t list_size_{200}; + uint32_t cache_nodes_num_{0}; + + bool warm_up_{false}; + uint32_t beam_size_{2}; + + DiskAnnIndexer::Pointer diskann_indexer_{nullptr}; + DiskAnnSearcherEntity entity_{}; + + // Mutable members for get_vector_by_id (caches context and buffer) + mutable ContextPointer fetch_ctx_{}; + mutable std::string fetch_vector_buffer_; + + uint32_t magic_{0U}; + + Stats stats_; + State state_{STATE_INIT}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_util.h b/src/core/algorithm/diskann/diskann_util.h new file mode 100644 index 0000000..a02130b --- /dev/null +++ b/src/core/algorithm/diskann/diskann_util.h @@ -0,0 +1,221 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "diskann_entity.h" + +namespace zvec { +namespace core { + +class DiskAnnUtil { + public: + static constexpr uint64_t kSectorSize = 4096; + static constexpr uint64_t kMaxSectorReadNum = 128; + + public: + static inline size_t div_round_up(size_t x, size_t y) { + return (x / y + (x % y != 0)); + } + + static inline size_t round_up(size_t x, size_t y) { + return div_round_up(x, y) * y; + } + + static inline void alloc_aligned(void **ptr, size_t size, size_t align) { + *ptr = ::aligned_alloc(align, size); + } + + static inline void free_aligned(void *ptr) { + if (ptr == nullptr) { + return; + } + + free(ptr); + } + + template + static inline void convert_vector_to_residual(T *data, uint32_t blocksize_, + size_t dim, void *centroid) { + const T *centroid_ptr = reinterpret_cast(centroid); + for (size_t i = 0; i < blocksize_; i++) { + for (uint64_t d = 0; d < dim; d++) { + float data_float = data[i * dim + d]; + data_float -= centroid_ptr[d]; + data[i * dim + d] = data_float; + } + } + } + + static inline void convert_types_uint32_to_uint8(const uint32_t *src, + uint8_t *dest, size_t npts, + size_t dim) { + for (size_t i = 0; i < npts; i++) { + for (size_t j = 0; j < dim; j++) { + dest[i * dim + j] = src[i * dim + j]; + } + } + } + + static inline uint64_t get_node_sector(uint32_t node_per_sector, + uint32_t max_nodesize_, + uint32_t sectorsize_, + diskann_id_t node_id) { + return (node_per_sector > 0 + ? node_id / node_per_sector + : node_id * div_round_up(max_nodesize_, sectorsize_)); + } + + static inline uint32_t *offset_to_node_neighbor(uint8_t *node_buf, + uint32_t elementsize_) { + return (uint32_t *)(node_buf + elementsize_); + } + + static inline uint8_t *offset_to_node(uint32_t node_per_sector, + uint32_t max_nodesize_, + uint8_t *sector_buf, + diskann_id_t node_id) { + return sector_buf + (node_per_sector == 0 + ? 0 + : (node_id % node_per_sector) * max_nodesize_); + } + + static inline const uint8_t *offset_to_node_const(uint32_t node_per_sector, + uint32_t max_nodesize_, + const uint8_t *sector_buf, + diskann_id_t node_id) { + return sector_buf + (node_per_sector == 0 + ? 0 + : (node_id % node_per_sector) * max_nodesize_); + } +}; + +//! Neighbor +struct Neighbor { + public: + Neighbor() = default; + + Neighbor(diskann_id_t id, float distance) + : id{id}, distance{distance}, expanded(false) {} + + inline bool operator<(const Neighbor &other) const { + return distance < other.distance || + (distance == other.distance && id < other.id); + } + + inline bool operator==(const Neighbor &other) const { + return (id == other.id); + } + + public: + diskann_id_t id; + float distance; + bool expanded; +}; + +//! NeighborPriorityQueue +class NeighborPriorityQueue { + public: + NeighborPriorityQueue() : size_(0), capacity_(0), cur_(0) {} + + explicit NeighborPriorityQueue(size_t capacity) + : size_(0), capacity_(capacity), cur_(0), data_(capacity + 1) {} + + void insert(const Neighbor &nbr) { + if (size_ == capacity_ && data_[size_ - 1] < nbr) { + return; + } + + size_t low = 0, high = size_; + while (low < high) { + size_t mid = (low + high) >> 1; + if (nbr < data_[mid]) { + high = mid; + } else if (data_[mid].id == nbr.id) { + return; + } else { + low = mid + 1; + } + } + + if (low < capacity_) { + std::memmove(&data_[low + 1], &data_[low], + (size_ - low) * sizeof(Neighbor)); + } + + data_[low] = {nbr.id, nbr.distance}; + if (size_ < capacity_) { + size_++; + } + + if (low < cur_) { + cur_ = low; + } + } + + Neighbor closest_unexpanded() { + data_[cur_].expanded = true; + size_t pre = cur_; + while (cur_ < size_ && data_[cur_].expanded) { + cur_++; + } + return data_[pre]; + } + + bool has_unexpanded_node() const { + return cur_ < size_; + } + + size_t size() const { + return size_; + } + + size_t capacity() const { + return capacity_; + } + + void reserve(size_t capacity) { + if (capacity + 1 > data_.size()) { + data_.resize(capacity + 1); + } + capacity_ = capacity; + } + + Neighbor &operator[](size_t i) { + return data_[i]; + } + + Neighbor operator[](size_t i) const { + return data_[i]; + } + + void sort() { + std::sort(data_.begin(), data_.begin() + size_); + } + + void clear() { + size_ = 0; + cur_ = 0; + } + + private: + size_t size_; + size_t capacity_; + size_t cur_; + std::vector data_; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_vecs_reader.h b/src/core/algorithm/diskann/diskann_vecs_reader.h new file mode 100644 index 0000000..0d80933 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_vecs_reader.h @@ -0,0 +1,320 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include + +namespace zvec { +namespace core { + +enum VecsBitMapIndex { + BITMAP_INDEX_KEY = 0, + BITMAP_INDEX_DENSE = 1, + BITMAP_INDEX_SPARSE = 2, + BITMAP_INDEX_PARTITION = 3, + BITMAP_INDEX_TAGLIST = 4 +}; + +#pragma pack(4) +struct VecsHeaderDetect { + uint64_t num_vecs; + uint16_t reserved; // in v0, it's meta size + uint16_t version; // 0 for v0 format, 1 for new format + uint8_t meta_buf[0]; +}; +#pragma pack() + +#pragma pack(4) +struct VecsHeader { + uint64_t num_vecs; + uint32_t meta_size; + uint8_t meta_buf[0]; +}; +#pragma pack() + +#pragma pack(4) +struct VecsHeaderV1 { + uint64_t num_vecs; + uint16_t meta_size_v1; + uint16_t version; + uint32_t meta_size; + uint64_t bitmap; // set for data section + uint64_t key_offset; // offset for key + uint64_t key_size; // size for key + uint64_t dense_offset; // offset for dense + uint64_t dense_size; // size for dense + uint64_t sparse_offset; // offset for sparse + uint64_t sparse_size; // size for sparse + uint64_t partition_offset; // offset for partition + uint64_t partition_size; // size for partition + uint64_t taglist_offset; // offset for taglist + uint64_t taglist_size; // size for taglist + uint8_t meta_buf[0]; +}; +#pragma pack() + +class DiskAnnVecsReader { + public: + DiskAnnVecsReader() + : mmap_file_(), + index_meta_(), + num_vecs_(0), + vector_base_(nullptr), + key_base_(nullptr), + sparse_base_meta_{nullptr}, + sparse_base_data_{nullptr}, + partition_base_{nullptr}, + taglist_base_meta_{nullptr}, + taglist_base_data_{nullptr}, + taglist_size_{0} {} + + void set_measure(const std::string &name, const IndexParams ¶ms) { + index_meta_.set_measure(name, 0, params); + } + + bool load(const std::string &fname) { + return load(fname.c_str()); + } + + bool load(const char *fname) { + if (!fname) { + std::cerr << "Load fname is nullptr" << std::endl; + return false; + } + if (!mmap_file_.open(fname, true)) { + std::cerr << "Open file error: " << fname << std::endl; + return false; + } + if (mmap_file_.size() < sizeof(VecsHeaderDetect)) { + std::cerr << "File size is too small: " << mmap_file_.size() << std::endl; + return false; + } + + const VecsHeaderDetect *header = + reinterpret_cast(mmap_file_.region()); + + if (header->version == 0) { + return load_v0(); + } else if (header->version == 1) { + return load_v1(); + } + + std::cerr << "Can not recognize version: " << header->version << std::endl; + + return false; + } + + bool load_v0() { + const VecsHeader *header = + reinterpret_cast(mmap_file_.region()); + // check + num_vecs_ = header->num_vecs; + + // deserialize + bool bret = index_meta_.deserialize(&header->meta_buf, header->meta_size); + if (!bret) { + std::cerr << "deserialize index meta error." << std::endl; + return false; + } + + if (!index_meta_.hybrid_vector()) { + if ((mmap_file_.size() - sizeof(*header) - header->meta_size) % + num_vecs_ != + 0) { + std::cerr << "input file foramt check error." << std::endl; + return false; + } + } + + if (!index_meta_.hybrid_vector()) { + vector_base_ = + reinterpret_cast(header + 1) + header->meta_size; + key_base_ = reinterpret_cast( + vector_base_ + num_vecs_ * index_meta_.element_size()); + } else { + vector_base_ = + reinterpret_cast(header + 1) + header->meta_size; + key_base_ = reinterpret_cast( + vector_base_ + num_vecs_ * index_meta_.element_size()); + sparse_base_meta_ = reinterpret_cast(key_base_ + num_vecs_); + sparse_base_data_ = reinterpret_cast( + sparse_base_meta_ + num_vecs_ * sizeof(uint64_t)); + } + + return true; + } + + bool load_v1() { + const VecsHeaderV1 *header = + reinterpret_cast(mmap_file_.region()); + // check + num_vecs_ = header->num_vecs; + + // deserialize + bool bret = index_meta_.deserialize(&header->meta_buf, header->meta_size); + if (!bret) { + std::cerr << "deserialize index meta error." << std::endl; + return false; + } + + const char *data_base_ptr = + reinterpret_cast(header + 1) + header->meta_size; + + vector_base_ = reinterpret_cast(data_base_ptr); + key_base_ = reinterpret_cast( + vector_base_ + num_vecs_ * index_meta_.element_size()); + + if (header->sparse_offset != -1LLU) { + sparse_base_meta_ = data_base_ptr + header->sparse_offset; + sparse_base_data_ = sparse_base_meta_ + num_vecs_ * sizeof(uint64_t); + } + + if (header->partition_offset != -1LLU) { + partition_base_ = reinterpret_cast( + data_base_ptr + header->partition_offset); + } + + if (header->taglist_offset != -1LLU) { + taglist_base_meta_ = data_base_ptr + header->taglist_offset; + taglist_base_data_ = taglist_base_meta_ + num_vecs_; + taglist_size_ = header->taglist_size; + } + + return true; + } + + size_t num_vecs() const { + return num_vecs_; + } + + const void *vector_base() const { + return vector_base_; + } + + const uint64_t *key_base() const { + return key_base_; + } + + const IndexMeta &index_meta() const { + return index_meta_; + } + + uint64_t get_key(size_t index) const { + return key_base_[index]; + } + + const void *get_vector(size_t index) const { + return vector_base_ + index * index_meta_.element_size(); + } + + uint32_t get_sparse_count(size_t index) const { + if (index_meta_.hybrid_vector()) { + auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t); + uint64_t sparse_offset = *((uint64_t *)sparse_data_meta); + uint32_t sparse_count = + *((uint32_t *)(sparse_base_data_ + sparse_offset)); + + return sparse_count; + } + + return 0; + } + + const uint32_t *get_sparse_indices(size_t index) const { + if (index_meta_.hybrid_vector()) { + auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t); + uint64_t sparse_offset = *((uint64_t *)sparse_data_meta); + uint32_t *sparse_indices = + (uint32_t *)(sparse_base_data_ + sparse_offset + sizeof(uint32_t)); + + return sparse_indices; + } + + return nullptr; + } + + const void *get_sparse_data(size_t index) const { + if (index_meta_.hybrid_vector()) { + auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t); + uint64_t sparse_offset = *((uint64_t *)sparse_data_meta); + uint32_t sparse_count = + *((uint32_t *)(sparse_base_data_ + sparse_offset)); + void *sparse_data = + (uint32_t *)(sparse_base_data_ + sparse_offset + sizeof(uint32_t) + + sparse_count * sizeof(uint32_t)); + + return sparse_data; + } + + return nullptr; + } + + size_t get_total_sparse_count(void) const { + size_t total_sparse_count = 0; + for (size_t i = 0; i < num_vecs_; ++i) { + total_sparse_count += get_sparse_count(i); + } + + return total_sparse_count; + } + + bool has_taglist(void) const { + return taglist_base_meta_ != nullptr; + } + + uint64_t get_taglist_count(size_t index) const { + if (!taglist_base_data_ || !taglist_base_meta_) { + return 0; + } + + uint64_t taglist_count = *reinterpret_cast( + taglist_base_data_ + taglist_base_meta_[index]); + return taglist_count; + } + + const uint64_t *get_taglist(size_t index) const { + if (!taglist_base_data_ || !taglist_base_meta_) { + return nullptr; + } + + return reinterpret_cast(taglist_base_data_ + + taglist_base_meta_[index]) + + 1; + } + + const void *get_taglist_data(size_t &size) const { + size = taglist_size_; + + return taglist_base_meta_; + } + + private: + ailego::MMapFile mmap_file_; + IndexMeta index_meta_; + size_t num_vecs_; + const char *vector_base_; + const uint64_t *key_base_; + const char *sparse_base_meta_; + const char *sparse_base_data_; + const uint32_t *partition_base_; + const char *taglist_base_meta_; + const char *taglist_base_data_; + uint64_t taglist_size_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_visit_filter.h b/src/core/algorithm/diskann/diskann_visit_filter.h new file mode 100644 index 0000000..dd4aff7 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_visit_filter.h @@ -0,0 +1,16 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "utility/visit_filter.h" diff --git a/src/core/algorithm/flat/CMakeLists.txt b/src/core/algorithm/flat/CMakeLists.txt new file mode 100644 index 0000000..6081496 --- /dev/null +++ b/src/core/algorithm/flat/CMakeLists.txt @@ -0,0 +1,18 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +#message(STATUS "PROJECT_ROOT_DIR = ${PROJECT_ROOT_DIR}") + +if(NOT APPLE) + set(CORE_KNN_FLAT_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_flat + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm ${PROJECT_ROOT_DIR}/src/core/framework + LDFLAGS "${CORE_KNN_FLAT_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/flat/flat_builder.cc b/src/core/algorithm/flat/flat_builder.cc new file mode 100644 index 0000000..be6810c --- /dev/null +++ b/src/core/algorithm/flat/flat_builder.cc @@ -0,0 +1,283 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_builder.h" +#include "flat_utility.h" + +namespace zvec { +namespace core { + +template +int FlatBuilder::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + meta_ = meta; + + // Set the major order + bool column_major_order = false; + if (params.get(PARAM_FLAT_COLUMN_MAJOR_ORDER, &column_major_order)) { + meta_.set_major_order(column_major_order ? IndexMeta::MO_COLUMN + : IndexMeta::MO_ROW); + } + + // Verify column major order + if (meta_.major_order() != IndexMeta::MO_ROW) { + IndexMeta::DataType dt = meta_.data_type(); + + bool support_column_major = false; + if ((dt != IndexMeta::DataType::DT_FP32 && + dt != IndexMeta::DataType::DT_FP16 && + dt != IndexMeta::DataType::DT_INT8 && dt != IndexMeta::DT_INT4 && + dt != IndexMeta::DT_BINARY32 && dt != IndexMeta::DT_BINARY64) || + (meta_.unit_size() != IndexMeta::UnitSizeof(dt))) { + if (meta_.major_order() == IndexMeta::MO_COLUMN) { + LOG_ERROR("Unsupported type %d with unit size %u.", dt, + meta_.unit_size()); + return IndexError_Unsupported; + } else { + support_column_major = false; + } + } + if (meta_.element_size() % IndexMeta::AlignSizeof(dt) != 0) { + if (meta_.major_order() == IndexMeta::MO_COLUMN) { + LOG_ERROR("Unsupported type %d with dimension %u.", dt, + meta_.dimension()); + return IndexError_Unsupported; + } else { + support_column_major = false; + } + } + + if (meta_.major_order() == IndexMeta::MO_UNDEFINED && + support_column_major) { + meta_.set_major_order(IndexMeta::MO_COLUMN); + } + } + + if (!VerifyMetric(meta_)) { + LOG_ERROR("Invalid index measure %s.", meta_.metric_name().c_str()); + return IndexError_InvalidArgument; + } + + std::string tag = std::to_string(BATCH_SIZE); + ailego::Params searcher_params; + searcher_params.set(PARAM_FLAT_BATCH_SIZE, BATCH_SIZE); + meta_.set_searcher("FlatSearcher" + tag, 0, searcher_params); + meta_.set_builder("FlatBuilder" + tag, 0, params); + return 0; +} + +template +int FlatBuilder::build(IndexThreads::Pointer, + IndexHolder::Pointer holder) { + ailego::ElapsedTime stamp; + if (!holder->is_matched(meta_)) { + LOG_ERROR("The holder is unmatched with initialized meta."); + return IndexError_Mismatch; + } + + holder_ = std::move(holder); + stats_.set_built_count(holder_->count()); + stats_.set_built_costtime(stamp.milli_seconds()); + return 0; +} + +template +int FlatBuilder::dump(const IndexDumper::Pointer &dumper) { + ailego::ElapsedTime stamp; + if (!holder_) { + return IndexError_NoReady; + } + + std::vector keys; + if (meta_.major_order() == IndexMeta::MO_COLUMN) { + int error_code = this->write_column_index(dumper.get(), &keys); + if (error_code != 0) { + return error_code; + } + } else { + int error_code = this->write_row_index(dumper.get(), &keys); + if (error_code != 0) { + return error_code; + } + } + + int error_code = this->write_keys(keys, dumper.get()); + if (error_code != 0) { + return error_code; + } + + error_code = this->write_mapping(keys, dumper.get()); + if (error_code != 0) { + return error_code; + } + + error_code = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (error_code != 0) { + return error_code; + } + + stats_.set_dumped_count(keys.size()); + stats_.set_dumped_costtime(stamp.milli_seconds()); + return 0; +} + +template +int FlatBuilder::write_keys(const std::vector &keys, + IndexDumper *dumper) { + size_t keys_size = keys.size() * sizeof(uint64_t); + size_t keys_padding_size = ailego_align(keys_size, 32) - keys_size; + if (dumper->write(keys.data(), keys_size) != keys_size) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + + // Write the padding if need + if (keys_padding_size) { + std::string padding(keys_padding_size, '\0'); + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + return dumper->append(FLAT_SEGMENT_KEYS_SEG_ID, keys_size, keys_padding_size, + 0); +} + +template +int FlatBuilder::write_mapping(const std::vector &keys, + IndexDumper *dumper) { + std::vector mapping(keys.size()); + std::iota(mapping.begin(), mapping.end(), 0); + std::sort( + mapping.begin(), mapping.end(), + [&keys](uint32_t lhs, uint32_t rhs) { return (keys[lhs] < keys[rhs]); }); + + size_t mapping_size = mapping.size() * sizeof(uint32_t); + size_t mapping_padding_size = ailego_align(mapping_size, 32) - mapping_size; + if (dumper->write(mapping.data(), mapping_size) != mapping_size) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + + // Write the padding if need + if (mapping_padding_size) { + std::string padding(mapping_padding_size, '\0'); + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + return dumper->append(FLAT_SEGMENT_MAPPING_SEG_ID, mapping_size, + mapping_padding_size, 0); +} + +template +template +int FlatBuilder::write_column_index(IndexDumper *dumper, + std::vector *keys) { + auto iter = holder_->create_iterator(); + if (!iter) { + LOG_ERROR("Failed to create iterator of holder"); + return IndexError_Runtime; + } + + // Write features + size_t element_size = holder_->element_size(); + size_t block_size = element_size * BATCH_SIZE; + std::string block1, block2; + block1.reserve(block_size); + block2.reserve(block_size); + + for (; iter->is_valid(); iter->next()) { + block1.append(reinterpret_cast(iter->data()), element_size); + keys->emplace_back(iter->key()); + + if (block1.size() == block_size) { + ailego::MatrixHelper::Transpose( + block1.data(), element_size / sizeof(T), (void *)block2.data()); + + if (dumper->write(block2.data(), block_size) != block_size) { + LOG_ERROR("Failed to write data into dumper %s", + dumper->name().c_str()); + return IndexError_WriteData; + } + block1.clear(); + } + } + + if (!block1.empty()) { + if (dumper->write(block1.data(), block1.size()) != block1.size()) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + + // Write the padding if need + size_t features_size = keys->size() * element_size; + size_t features_padding_size = + ailego_align(features_size, 32) - features_size; + if (features_padding_size) { + std::string padding(features_padding_size, '\0'); + + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + return dumper->append(FLAT_SEGMENT_FEATURES_SEG_ID, features_size, + features_padding_size, 0); +} + +template +int FlatBuilder::write_row_index(IndexDumper *dumper, + std::vector *keys) { + auto iter = holder_->create_iterator(); + if (!iter) { + LOG_ERROR("Failed to create iterator of holder"); + return IndexError_Runtime; + } + + // Write features + size_t element_size = holder_->element_size(); + for (; iter->is_valid(); iter->next()) { + if (dumper->write(iter->data(), element_size) != element_size) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + keys->emplace_back(iter->key()); + } + + // Write the padding if need + size_t features_size = keys->size() * element_size; + size_t features_padding_size = + ailego_align(features_size, 32) - features_size; + if (features_padding_size) { + std::string padding(features_padding_size, '\0'); + + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + return dumper->append(FLAT_SEGMENT_FEATURES_SEG_ID, features_size, + features_padding_size, 0); +} + +INDEX_FACTORY_REGISTER_BUILDER_ALIAS(LinearBuilder, FlatBuilder<32>); +INDEX_FACTORY_REGISTER_BUILDER_ALIAS(FlatBuilder, FlatBuilder<32>); +INDEX_FACTORY_REGISTER_BUILDER_ALIAS(FlatBuilder16, FlatBuilder<16>); +INDEX_FACTORY_REGISTER_BUILDER_ALIAS(FlatBuilder32, FlatBuilder<32>); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_builder.h b/src/core/algorithm/flat/flat_builder.h new file mode 100644 index 0000000..397a6c1 --- /dev/null +++ b/src/core/algorithm/flat/flat_builder.h @@ -0,0 +1,102 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "flat_utility.h" + +namespace zvec { +namespace core { + +/*! Flat Builder + */ +template +class FlatBuilder : public IndexBuilder { + public: + //! Destructor + ~FlatBuilder(void) override {} + + //! Initialize the builder + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup the builder + int cleanup(void) override { + holder_ = nullptr; + return 0; + } + + //! Train the data + int train(IndexThreads::Pointer, IndexHolder::Pointer) override { + stats_.set_trained_count(0u); + stats_.set_trained_costtime(0u); + return 0; + } + + //! Train the data + int train(const IndexTrainer::Pointer &) override { + stats_.set_trained_count(0u); + stats_.set_trained_costtime(0u); + return 0; + } + + //! Build the index + int build(IndexThreads::Pointer, IndexHolder::Pointer holder) override; + + //! Dump index into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const IndexBuilder::Stats &stats(void) const override { + return stats_; + } + + protected: + //! Dump index keys + int write_keys(const std::vector &keys, IndexDumper *dumper); + + //! Dump index keys mapping + int write_mapping(const std::vector &keys, IndexDumper *dumper); + + //! Dump index using column-major-order format + template + int write_column_index(IndexDumper *dumper, std::vector *keys); + + //! Dump index using column-major-order format + int write_column_index(IndexDumper *dumper, std::vector *keys) { + switch (IndexMeta::AlignSizeof(meta_.data_type())) { + case 2: + return this->write_column_index(dumper, keys); + case 4: + return this->write_column_index(dumper, keys); + case 8: + return this->write_column_index(dumper, keys); + default: + ailego_check_with(0, "BAD CASE"); + } + return IndexError_Runtime; + } + + //! Dump index using row-major-order format + int write_row_index(IndexDumper *dumper, std::vector *keys); + + private: + IndexMeta meta_{}; + IndexBuilder::Stats stats_{}; + IndexHolder::Pointer holder_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_distance_matrix.h b/src/core/algorithm/flat/flat_distance_matrix.h new file mode 100644 index 0000000..c8c3dad --- /dev/null +++ b/src/core/algorithm/flat/flat_distance_matrix.h @@ -0,0 +1,184 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once +#include "flat_utility.h" + +namespace zvec { +namespace core { + +/*! Brute Force Distance Tuple + */ +template +class FlatDistanceTuple; + +/*! Brute Force Distance Tuple + */ +template <> +class FlatDistanceTuple<1> { + public: + //! Retrieve non-zero if all distances are valid. + bool is_valid(void) const { + return !!distance_; + } + + //! Retrieve non-zero if a distance is valid. + bool is_valid(size_t m) const { + return m == 1 && !!distance_; + } + + //! Initialize the distance tuple + void initialize(const IndexMetric &measure) { + distance_ = measure.distance_matrix(1, 1); + } + + //! Initialize the distance tuple + void initialize(const IndexMetric &measure, size_t m) { + distance_ = measure.distance_matrix(m, 1); + } + + //! Compute the distance between matrix and query + template + auto distance(const void *m, const void *q, size_t dim, float *out) const -> + typename std::enable_if::type { + distance_(m, q, dim, out); + } + + private: + IndexMetric::MatrixDistance distance_{}; +}; + +/*! Brute Force Distance Tuple + */ +template +class FlatDistanceTuple< + K, typename std::enable_if::value>::type> { + public: + //! Retrieve non-zero if all distances are valid. + bool is_valid(void) const { + return (distance_tuple_.is_valid() && !!distance_); + } + + //! Retrieve non-zero if a distance is valid. + bool is_valid(size_t m) const { + return (m == K ? (!!distance_) + : (m < K ? distance_tuple_.is_valid(m) : false)); + } + + //! Initialize the distance tuple + void initialize(const IndexMetric &measure) { + distance_tuple_.initialize(measure); + distance_ = measure.distance_matrix(K, 1); + } + + //! Initialize the distance tuple + void initialize(const IndexMetric &measure, size_t m) { + distance_tuple_.initialize(measure, m); + distance_ = measure.distance_matrix(m, K); + } + + //! Compute the distance between matrix and query + template + auto distance(const void *m, const void *q, size_t dim, float *out) const -> + typename std::enable_if::type { + distance_(m, q, dim, out); + } + + //! Compute the distance between matrix and query + template + auto distance(const void *m, const void *q, size_t dim, float *out) const -> + typename std::enable_if<(K > M) && IsEqualPowerofTwo::value>::type { + distance_tuple_.template distance(m, q, dim, out); + } + + private: + FlatDistanceTuple<(K >> 1)> distance_tuple_{}; + IndexMetric::MatrixDistance distance_{}; +}; + +/*! Brute Force Distance Matrix + */ +template +class FlatDistanceMatrix; + +/*! Brute Force Distance Matrix + */ +template <> +class FlatDistanceMatrix<1> { + public: + //! Retrieve non-zero if all distances are valid. + bool is_valid(void) const { + return (!!distance_); + } + + //! Initialize the distance tuple + void initialize(const IndexMetric &measure) { + distance_ = measure.distance_matrix(1, 1); + } + + //! Compute the distance between matrix and query + template + auto distance(const void *m, const void *q, size_t dim, float *out) const -> + typename std::enable_if::type { + distance_(m, q, dim, out); + } + + private: + IndexMetric::MatrixDistance distance_{}; +}; + +/*! Brute Force Distance Matrix + */ +template +class FlatDistanceMatrix< + K, typename std::enable_if::value>::type> { + public: + //! Retrieve non-zero if all distances are valid. + bool is_valid(void) const { + return (tuple_h_.is_valid() && tuple_v_.is_valid()); + } + + //! Retrieve non-zero if a distance is valid. + bool is_valid(size_t m, size_t n) const { + return (m == K ? tuple_h_.is_valid(n) + : (m < K && n == 1 ? tuple_v_.is_valid(m) : false)); + } + + //! Initialize the distance tuple + void initialize(const IndexMetric &measure) { + tuple_h_.initialize(measure, K); + tuple_v_.initialize(measure); + } + + //! Compute the distance between matrix and query + template + auto distance(const void *m, const void *q, size_t dim, float *out) const -> + typename std::enable_if<(K == M) && (K >= N)>::type { + tuple_h_.template distance(m, q, dim, out); + } + + //! Compute the distance between matrix and query + template + auto distance(const void *m, const void *q, size_t dim, float *out) const -> + typename std::enable_if<(K > M) && (N == 1u)>::type { + tuple_v_.template distance(m, q, dim, out); + } + + private: + FlatDistanceTuple tuple_h_{}; + FlatDistanceTuple<(K >> 1)> tuple_v_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_index_format.h b/src/core/algorithm/flat/flat_index_format.h new file mode 100644 index 0000000..e9c1c8e --- /dev/null +++ b/src/core/algorithm/flat/flat_index_format.h @@ -0,0 +1,187 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +using node_id_t = uint32_t; +using key_t = uint64_t; +using level_t = int32_t; +using dist_t = float; +using TopkHeap = ailego::KeyValueHeap; +using CandidateHeap = + ailego::KeyValueHeap>; +constexpr node_id_t kInvalidNodeId = static_cast(-1); +constexpr key_t kInvalidKey = static_cast(-1); + +/*! Index Format of Linear Index Header + */ +struct LinearIndexHeader { + LinearIndexHeader() + : header_size(0), + total_vector_count(0), + linear_body_size(0), + linear_list_count(0), + block_vector_count(0), + block_size(0), + block_count(0), + index_meta_size(0) { + memset(reserved_, 0, sizeof(reserved_)); + } + uint32_t header_size{0}; + uint32_t total_vector_count{0}; + uint64_t linear_body_size{0}; + uint32_t linear_list_count{0}; + uint32_t block_vector_count{0}; + uint32_t block_size{0}; + uint32_t block_count{0}; + uint32_t index_meta_size{0}; + char reserved_[28] = {0}; +#ifdef _MSC_VER + char index_meta[]; +#else + char index_meta[0]; +#endif +}; + +/*! Index Format of Linear Index Meta for each Linear list + */ +struct LinearListMeta { + LinearListMeta() : offset(0), block_count(0), vector_count(0), id_offset(0) { + memset(reserved_, 0, sizeof(reserved_)); + } + uint64_t offset{0}; + uint32_t block_count{0}; + uint32_t vector_count{0}; + uint32_t id_offset{0}; + char reserved_[16] = {0}; +}; + +/*! Index Format of Location in Linear Index for each vector + */ +struct LinearVecLocation { + LinearVecLocation(size_t off, bool col) + : offset(off), column_major(col), reserved(0u) {} + + uint64_t offset : 48; // feature offset in posting block segment + uint64_t column_major : 1; // coloum major if true + uint64_t reserved : 15; +}; + +/*! Index Format of Integer Quantizer params for each linear list + */ +struct LinearIntegerQuantizerParams { + float scale{1.0}; + float bias{0.0}; +}; + +/*! Location of Vectors Block in Storage Segment + */ +struct BlockLocation { + uint32_t segment_id{0}; + uint32_t block_index{0}; +}; + +/*! The Header of a Block in Storage Segment + */ +struct BlockHeader { + BlockHeader() : vector_count(0u), column_major(0u), reserved(0u) {} + BlockLocation next; + uint16_t vector_count{0}; + uint16_t column_major : 1; + uint16_t reserved : 15; +}; + +struct DeletionMap { + void set(uint32_t index) { + bitset.set(index); + } + + void reset(uint32_t index) { + bitset.reset(index); + } + + bool test(uint32_t index) const { + return bitset.test(index); + } + + bool is_dirty() const { + return bitset.test_any(); + } + + ailego::FixedBitset<32> bitset{}; +}; + +static_assert(sizeof(DeletionMap) == 4, "DeletionMap must be 4 bytes"); + +/*! Meta Information of Streamer Entity + */ +struct StreamerLinearMeta { + StreamerLinearMeta() + : create_time(0), + update_time(0), + revision_id(0), + segment_count(0), + segment_size(0) { + memset(reserved_, 0, sizeof(reserved_)); + } + uint64_t create_time{0}; + uint64_t update_time{0}; + uint64_t revision_id{0}; + uint32_t segment_count{0}; + uint32_t segment_size{0}; + uint8_t reserved_[32] = {0}; + LinearIndexHeader header; +}; + +/*! Location of Vector in Storage Segment + */ +struct VectorLocation { + //! Constructor + VectorLocation(void) + : segment_id(0u), column_major(0u), reserved(0u), offset(0u) {} + + //! Constructor + VectorLocation(uint32_t id, bool col, uint32_t off) + : segment_id(id), column_major(col), reserved(0u), offset(off) {} + + uint32_t segment_id{0}; + uint16_t column_major : 1; + uint16_t reserved : 15; + uint32_t offset{0}; + + public: + bool operator==(const VectorLocation &other) const { + return segment_id == other.segment_id && + column_major == other.column_major && offset == other.offset; + } +}; + +// static_assert(sizeof(VectorLocation) == sizeof(uint64_t), +// "VectorLocation must be size of 8 bytes"); + +struct KeyInfo { + KeyInfo(void) : centroid_idx(0u) {} + KeyInfo(uint32_t idx, const VectorLocation &loc) + : centroid_idx(idx), location(loc) {} + KeyInfo(VectorLocation loc) : location(loc) {} + uint32_t centroid_idx{0}; + VectorLocation location; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_searcher.cc b/src/core/algorithm/flat/flat_searcher.cc new file mode 100644 index 0000000..ef4e7d4 --- /dev/null +++ b/src/core/algorithm/flat/flat_searcher.cc @@ -0,0 +1,252 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "flat_searcher.h" +#include +#include +#include "flat_distance_matrix.h" +#include "flat_searcher_context.h" +#include "flat_searcher_provider.h" +#include "flat_utility.h" + +namespace zvec { +namespace core { + +template +IndexProvider::Pointer FlatSearcher::create_provider(void) const { + std::lock_guard lock(mapping_mutex_); + + if (mapping_.empty()) { + auto mapping_segment = container_->get(FLAT_SEGMENT_MAPPING_SEG_ID); + if (!mapping_segment) { + LOG_ERROR("Failed to fetch segment %s", + FLAT_SEGMENT_MAPPING_SEG_ID.c_str()); + return nullptr; + } + + if (mapping_segment->data_size() % sizeof(uint32_t) != 0) { + LOG_ERROR("Invalid data size %zu of mapping segment", + mapping_segment->data_size()); + return nullptr; + } + + size_t mapping_count = mapping_segment->data_size() / sizeof(uint32_t); + if (mapping_count * meta_.element_size() != + features_segment_->data_size()) { + LOG_ERROR("Invalid data size %zd of mapping segment", + features_segment_->data_size()); + return nullptr; + } + + const uint32_t *mapping = nullptr; + if (mapping_segment->read(0, reinterpret_cast(&mapping), + mapping_segment->data_size()) != + mapping_segment->data_size()) { + LOG_ERROR("Failed to read data (%zu bytes) from mapping segment", + mapping_segment->data_size()); + return nullptr; + } + mapping_.clear(); + mapping_.reserve(mapping_count); + std::copy(mapping, mapping + mapping_count, std::back_inserter(mapping_)); + } + return IndexProvider::Pointer(new (std::nothrow) + FlatSearcherProvider(this)); +} + +template +int FlatSearcher::load(IndexStorage::Pointer cntr, + IndexMetric::Pointer measure) { + ailego::ElapsedTime stamp; + if (!cntr) { + return IndexError_InvalidArgument; + } + + int error_code = IndexHelper::DeserializeFromStorage(cntr.get(), &meta_); + if (error_code != 0) { + LOG_ERROR( + "Failed to deserialize index meta from container %s, error=%d, %s", + cntr->name().c_str(), error_code, IndexError::What(error_code)); + return error_code; + } + + if (!measure) { + error_code = InitializeMetric(meta_, &measure_); + if (error_code != 0) { + LOG_ERROR("Failed to initialize index measure %s, error=%d, %s", + meta_.metric_name().c_str(), error_code, + IndexError::What(error_code)); + return error_code; + } + if (measure_->query_metric()) { + measure_ = measure_->query_metric(); + } + } else { + if (!measure->is_matched(meta_)) { + LOG_ERROR( + "The index measure is unmatched with index meta from container."); + return IndexError_Mismatch; + } + measure_ = std::move(measure); + } + + column_major_order_ = (meta_.major_order() == IndexMeta::MO_COLUMN); + distance_matrix_.initialize(*measure_); + + if (column_major_order_) { + if (!distance_matrix_.is_valid()) { + LOG_ERROR("Lack of distance functions to support column index."); + return IndexError_Unsupported; + } + } else { + if (!distance_matrix_.is_valid(1, 1)) { + LOG_ERROR("Lack of distance functions to support row index."); + return IndexError_Unsupported; + } + } + + auto keys_segment = cntr->get(FLAT_SEGMENT_KEYS_SEG_ID); + if (!keys_segment) { + LOG_ERROR("Failed to fetch segment %s", FLAT_SEGMENT_KEYS_SEG_ID.c_str()); + return IndexError_NoExist; + } + features_segment_ = cntr->get(FLAT_SEGMENT_FEATURES_SEG_ID); + if (!features_segment_) { + LOG_ERROR("Failed to fetch segment %s", FLAT_SEGMENT_KEYS_SEG_ID.c_str()); + return IndexError_NoExist; + } + + if (keys_segment->data_size() % sizeof(uint64_t) != 0) { + LOG_ERROR("Invalid data size %zu of keys segment", + keys_segment->data_size()); + return IndexError_InvalidLength; + } + + size_t keys_count = keys_segment->data_size() / sizeof(uint64_t); + if (keys_count * meta_.element_size() != features_segment_->data_size()) { + LOG_ERROR("Invalid data size %zd of features segment", + features_segment_->data_size()); + return IndexError_Mismatch; + } + + if (keys_segment->read(0, reinterpret_cast(&keys_), + keys_segment->data_size()) != + keys_segment->data_size()) { + LOG_ERROR("Failed to read data (%zu bytes) from keys segment", + keys_segment->data_size()); + return IndexError_ReadData; + } + + for (size_t i = 0; i < keys_count; i++) { + key_id_mapping_[keys_[i]] = i; + } + + container_ = cntr; + magic_ = IndexContext::GenerateMagic(); + stats_.set_loaded_count(keys_count); + stats_.set_loaded_costtime(stamp.milli_seconds()); + return 0; +} + +template +int FlatSearcher::search_impl(const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + ailego_assert(query && !!context); + ailego_assert(measure_->is_matched(meta_, qmeta)); + + FlatSearcherContext *bf_context = + dynamic_cast *>(context.get()); + if (!bf_context) { + LOG_ERROR("Invalid brute-force searcher context"); + return IndexError_InvalidArgument; + } + + if (bf_context->magic() != magic_) { + bf_context->reset(this); + } + if (bf_context->group_by_search()) { + return bf_context->group_by_search_impl(query, qmeta, 1); + } else { + return (column_major_order_ ? bf_context->search_column(query, qmeta) + : bf_context->search_row(query, qmeta)); + } +} + +template +int FlatSearcher::search_impl(const void *query, + const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + ailego_assert(query && count && !!context); + ailego_assert(measure_->is_matched(meta_, qmeta)); + + FlatSearcherContext *bf_context = + dynamic_cast *>(context.get()); + if (!bf_context) { + LOG_ERROR("Invalid brute-force searcher context"); + return IndexError_InvalidArgument; + } + + if (bf_context->magic() != magic_) { + bf_context->reset(this); + } + + if (bf_context->group_by_search()) { + return bf_context->group_by_search_impl(query, qmeta, count); + } else { + return (column_major_order_ ? bf_context->search_column(query, qmeta, count) + : bf_context->search_row(query, qmeta, count)); + } +} + +template +int FlatSearcher::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + ailego_assert(query && count && !!context); + ailego_assert(measure_->is_matched(meta_, qmeta)); + + if (ailego_unlikely(p_keys.size() != count)) { + LOG_ERROR("The size of p_keys is not equal to count"); + return IndexError_InvalidArgument; + } + + FlatSearcherContext *bf_context = + dynamic_cast *>(context.get()); + if (!bf_context) { + LOG_ERROR("Invalid brute-force searcher context"); + return IndexError_InvalidArgument; + } + + if (bf_context->magic() != magic_) { + bf_context->reset(this); + } + + return bf_context->search_bf_by_p_keys_impl(query, p_keys, qmeta, count); +} + +template +IndexSearcher::Context::Pointer FlatSearcher::create_context( + void) const { + return IndexSearcher::Context::Pointer( + new FlatSearcherContext(this)); +} + +INDEX_FACTORY_REGISTER_SEARCHER_ALIAS(LinearSearcher, FlatSearcher<32>); +INDEX_FACTORY_REGISTER_SEARCHER_ALIAS(FlatSearcher, FlatSearcher<32>); +INDEX_FACTORY_REGISTER_SEARCHER_ALIAS(FlatSearcher16, FlatSearcher<16>); +INDEX_FACTORY_REGISTER_SEARCHER_ALIAS(FlatSearcher32, FlatSearcher<32>); +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_searcher.h b/src/core/algorithm/flat/flat_searcher.h new file mode 100644 index 0000000..2b0d0d9 --- /dev/null +++ b/src/core/algorithm/flat/flat_searcher.h @@ -0,0 +1,182 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once +#include +#include +#include +#include "flat_distance_matrix.h" +#include "flat_index_format.h" + +namespace zvec { +namespace core { + +/*! Flat Searcher + */ +template +class FlatSearcher : public IndexSearcher { + public: + //! Destructor + ~FlatSearcher(void) override = default; + + //! Initialize Searcher + int init(const ailego::Params &index_params) override { + params_ = index_params; + read_block_size_ = FLAT_DEFAULT_READ_BLOCK_SIZE; + index_params.get(PARAM_FLAT_READ_BLOCK_SIZE, &read_block_size_); + return 0; + } + + //! Cleanup Searcher + int cleanup(void) override { + return this->unload(); + } + + //! Load index from container + int load(IndexStorage::Pointer cntr, IndexMetric::Pointer measure) override; + + //! Unload index + int unload(void) override { + container_ = nullptr; + measure_ = nullptr; + features_segment_ = nullptr; + keys_ = nullptr; + key_id_mapping_.clear(); + return 0; + } + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override { + return this->search_impl(query, qmeta, context); + } + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override { + return this->search_impl(query, qmeta, count, context); + } + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const override { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const override; + + //! Retrieve statistics + const IndexSearcher::Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve params of index + const ailego::Params ¶ms(void) const override { + return params_; + } + + //! Create a searcher context + IndexSearcher::Context::Pointer create_context(void) const override; + + //! Create a searcher provider + IndexProvider::Pointer create_provider(void) const override; + + //! Retrieve magic number + uint32_t magic(void) const { + return magic_; + } + + //! Retrieve block size of data read + uint32_t read_block_size(void) const { + return read_block_size_; + } + + //! Retrieve primary key via index id + uint64_t key(size_t i) const { + return keys_[i]; + } + + // Retrieve index id via primary key + node_id_t get_id(key_t key) const { + auto it = key_id_mapping_.find(key); + if (it != key_id_mapping_.end()) { + return it->second; + } else { + return kInvalidNodeId; + } + } + + //! Retrieve primary key via index id + uint32_t local_index(size_t i) const { + return mapping_[i]; + } + + //! Retrieve primary key via index id + inline bool column_major_order(void) const { + return column_major_order_; + } + + //! Retrieve the distance matrix + const FlatDistanceMatrix &distance_matrix(void) const { + return distance_matrix_; + } + + //! Clone a features segment + IndexStorage::Segment::Pointer clone_features_segment(void) const { + return features_segment_->clone(); + } + + const void *get_vector(key_t key) const override { + auto provider = this->create_provider(); + return provider->get_vector(key); + } + + private: + //! Members + const uint64_t *keys_{nullptr}; + std::unordered_map key_id_mapping_; + uint32_t magic_{IndexContext::GenerateMagic()}; + uint32_t read_block_size_{FLAT_DEFAULT_READ_BLOCK_SIZE}; + bool column_major_order_{false}; + IndexMeta meta_{}; + IndexStorage::Pointer container_{}; + IndexMetric::Pointer measure_{}; + ailego::Params params_{}; + IndexStorage::Segment::Pointer features_segment_{}; + mutable std::vector mapping_{}; + mutable std::mutex mapping_mutex_{}; + FlatDistanceMatrix distance_matrix_{}; + IndexSearcher::Stats stats_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_searcher_context.h b/src/core/algorithm/flat/flat_searcher_context.h new file mode 100644 index 0000000..a04e53a --- /dev/null +++ b/src/core/algorithm/flat/flat_searcher_context.h @@ -0,0 +1,1114 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include "flat_index_format.h" +#include "flat_searcher.h" +#include "flat_utility.h" + + +namespace zvec { +namespace core { + +/*! Brute Force Searcher Context + */ +template +class FlatSearcherContext : public IndexSearcher::Context { + public: + //! Constructor + FlatSearcherContext(const FlatSearcher *owner) { + this->reset(owner); + } + + //! Destructor + ~FlatSearcherContext(void) override {} + + //! Set topk of search result + void set_topk(uint32_t topk) override { + topk_ = topk; + } + + //! Retrieve search result + const IndexDocumentList &result(void) const override { + return result_heaps_.at(0); + } + + //! Retrieve search result with index + const IndexDocumentList &result(size_t index) const override { + return result_heaps_.at(index); + } + + //! Retrieve result object for output + IndexDocumentList *mutable_result(size_t idx) override { + return &result_heaps_.at(idx); + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(void) const override { + return group_results_[0]; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(size_t idx) const override { + return group_results_[idx]; + } + + IndexGroupDocumentList *mutable_group_result(void) override { + return &group_results_[0]; + } + + IndexGroupDocumentList *mutable_group_result(size_t idx) override { + return &group_results_[idx]; + } + + //! Update the parameters of context + int update(const ailego::Params & /*params*/) override { + return 0; + } + + //! Retrieve magic number + uint32_t magic(void) const override { + return magic_; + } + + //! Get group topk + inline uint32_t group_topk() const { + return group_topk_; + } + + //! Get group num + inline uint32_t group_num() const { + return group_num_; + } + + inline std::map &group_topk_heaps() { + return group_topk_heaps_; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + + bool fetch_vector() const override { + return fetch_vector_; + } + + inline void resize_group_results(size_t size) { + if (group_by_search()) { + group_results_.resize(size); + } + } + + void topk_to_group_result(uint32_t idx) { + ailego_assert_with(idx < group_results_.size(), "invalid idx"); + + group_results_[idx].clear(); + + std::vector> group_topk_list; + std::vector> best_score_in_groups; + for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end(); + itr++) { + const std::string &group_id = (*itr).first; + auto &heap = (*itr).second; + heap.sort(); + + if (heap.size() > 0) { + float best_score = heap[0].second; + best_score_in_groups.push_back(std::make_pair(group_id, best_score)); + } + } + + std::sort(best_score_in_groups.begin(), best_score_in_groups.end(), + [](const std::pair &a, + const std::pair &b) -> int { + return a.second < b.second; + }); + + // truncate to group num + for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size(); + ++i) { + const std::string &group_id = best_score_in_groups[i].first; + + group_topk_list.emplace_back( + std::make_pair(group_id, group_topk_heaps_[group_id])); + } + + group_results_[idx].resize(group_topk_list.size()); + + for (uint32_t i = 0; i < group_topk_list.size(); ++i) { + const std::string &group_id = group_topk_list[i].first; + group_results_[idx][i].set_group_id(group_id); + + uint32_t size = std::min( + group_topk_, static_cast(group_topk_list[i].second.size())); + + for (uint32_t j = 0; j < size; ++j) { + auto score = group_topk_list[i].second[j].second; + if (score > this->threshold()) { + break; + } + + node_id_t id = group_topk_list[i].second[j].first; + + auto provider = owner_->create_provider(); + + if (fetch_vector_) { + group_results_[idx][i].mutable_docs()->emplace_back( + id, score, id, provider->get_vector(id)); + } else { + group_results_[idx][i].mutable_docs()->emplace_back(id, score, id); + } + } + } + } + + //! Get if group by search + bool group_by_search() { + return group_num_ > 0; + } + + //! Set group params + void set_group_params(uint32_t group_num, uint32_t group_topk) override { + group_num_ = group_num; + group_topk_ = group_topk; + group_topk_heaps_.clear(); + } + + void reset() override {} + + //! Reset the context + void reset(const FlatSearcher *owner) { + magic_ = owner->magic(); + feature_size_ = owner->meta().element_size(); + + uint32_t block_size = feature_size_ * BATCH_SIZE; + actual_read_size_ = + (owner->read_block_size() + block_size - 1) / block_size * block_size; + features_segment_ = owner->clone_features_segment(); + owner_ = owner; + } + + //! Similarity search + int search_row(const void *query, const IndexQueryMeta &qmeta) { + return (this->filter().is_valid() + ? this->search_row_filter(query, qmeta) + : this->search_row_nofilter(query, qmeta)); + } + + //! Similarity search + int search_row(const void *query, const IndexQueryMeta &qmeta, size_t count) { + return (this->filter().is_valid() + ? this->batch_search_row_filter(query, qmeta, count) + : this->batch_search_row_nofilter(query, qmeta, count)); + } + + //! Similarity search + int search_column(const void *query, const IndexQueryMeta &qmeta) { + return (this->filter().is_valid() + ? this->search_column_filter(query, qmeta) + : this->search_column_nofilter(query, qmeta)); + } + + //! Similarity search + int search_column(const void *query, const IndexQueryMeta &qmeta, + size_t count) { + return (this->filter().is_valid() + ? this->batch_search_column_filter(query, qmeta, count) + : this->batch_search_column_nofilter(query, qmeta, count)); + } + + int group_by_search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count); + + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count); + + protected: + //! Enqueue items into the search heaps (without filter) + template + auto batch_enqueue_nofilter(const void *block, size_t block_index, + size_t query_index, const IndexQueryMeta &qmeta, + size_t query_count) -> + typename std::enable_if::value>::type { + size_t query_batch_count = query_count / K; + + for (size_t i = 0; i != query_batch_count; ++i) { + owner_->distance_matrix().template distance( + block, &batch_queries_[query_index * qmeta.element_size()], + qmeta.dimension(), scores_); + + for (size_t k = 0; k != K; ++k) { + IndexDocumentHeap *heap = &result_heaps_[query_index++]; + for (size_t j = 0; j != BATCH_SIZE; ++j) { + heap->emplace(0, scores_[k * BATCH_SIZE + j], block_index + j); + } + } // end of for + } // end of for + + size_t query_left_count = query_count % K; + if (query_left_count != 0) { + this->batch_enqueue_nofilter<(K >> 1)>(block, block_index, query_index, + qmeta, query_left_count); + } + } + + //! Enqueue items into the search heaps (without filter) + template + auto batch_enqueue_nofilter(const void *block, size_t block_index, + size_t query_index, const IndexQueryMeta &qmeta, + size_t query_count) -> + typename std::enable_if::type { + ailego_assert(query_count == 1); + (void)query_count; + + owner_->distance_matrix().template distance( + block, &batch_queries_[query_index * qmeta.element_size()], + qmeta.dimension(), scores_); + + IndexDocumentHeap *heap = &result_heaps_[query_index]; + for (size_t i = 0; i != BATCH_SIZE; ++i) { + heap->emplace(0, scores_[i], block_index + i); + } + } + + //! Enqueue items into the search heaps (with filter) + template + auto batch_enqueue_filter(const void *block, size_t block_index, + size_t block_mask, size_t query_index, + const IndexQueryMeta &qmeta, size_t query_count) -> + typename std::enable_if::value>::type { + size_t query_batch_count = query_count / K; + + for (size_t i = 0; i != query_batch_count; ++i) { + owner_->distance_matrix().template distance( + block, &batch_queries_[query_index * qmeta.element_size()], + qmeta.dimension(), scores_); + + for (size_t k = 0; k != K; ++k) { + IndexDocumentHeap *heap = &result_heaps_[query_index++]; + for (size_t j = 0; j != BATCH_SIZE; ++j) { + if ((block_mask & (1 << j)) != 0) { + heap->emplace(0, scores_[k * BATCH_SIZE + j], block_index + j); + } + } + } // end of for + } // end of for + + size_t query_left_count = query_count % K; + if (query_left_count != 0) { + this->batch_enqueue_filter<(K >> 1)>( + block, block_index, block_mask, query_index, qmeta, query_left_count); + } + } + + //! Enqueue items into the search heaps (with filter) + template + auto batch_enqueue_filter(const void *block, size_t block_index, + size_t block_mask, size_t query_index, + const IndexQueryMeta &qmeta, size_t query_count) -> + typename std::enable_if::type { + ailego_assert(query_count == 1); + (void)query_count; + + owner_->distance_matrix().template distance( + block, &batch_queries_[query_index * qmeta.element_size()], + qmeta.dimension(), scores_); + + IndexDocumentHeap *heap = &result_heaps_[query_index]; + for (size_t i = 0; i != BATCH_SIZE; ++i) { + if ((block_mask & (1 << i)) != 0) { + heap->emplace(0, scores_[i], block_index + i); + } + } + } + + //! Enqueue items into the search heaps (without filter) + template + auto single_enqueue_nofilter(const void *feature, size_t feature_index, + size_t query_index, const IndexQueryMeta &qmeta, + size_t query_count) -> + typename std::enable_if::value>::type { + size_t query_batch_count = query_count / K; + + for (size_t i = 0; i != query_batch_count; ++i) { + owner_->distance_matrix().template distance( + &batch_queries_[query_index * qmeta.element_size()], feature, + qmeta.dimension(), scores_); + + for (size_t k = 0; k != K; ++k) { + result_heaps_[query_index++].emplace(0, scores_[k], feature_index); + } + } + size_t query_left_count = query_count % K; + if (query_left_count != 0) { + this->single_enqueue_nofilter<(K >> 1)>( + feature, feature_index, query_index, qmeta, query_left_count); + } + } + + //! Enqueue items into the search heaps (without filter) + template + auto single_enqueue_nofilter(const void *feature, size_t feature_index, + size_t query_index, const IndexQueryMeta &qmeta, + size_t query_count) -> + typename std::enable_if::type { + ailego_assert(query_count == 1); + (void)query_count; + + owner_->distance_matrix().template distance<1>( + feature, &batch_queries_[query_index * qmeta.element_size()], + qmeta.dimension(), scores_); + result_heaps_[query_index].emplace(0, scores_[0], feature_index); + } + + protected: + //! Similarity search (1 column without filter) + int search_column_nofilter(const void *query, const IndexQueryMeta &qmeta); + + //! Similarity search (1 column with filter) + int search_column_filter(const void *query, const IndexQueryMeta &qmeta); + + //! Similarity search (1 row without filter) + int search_row_nofilter(const void *query, const IndexQueryMeta &qmeta); + + //! Similarity search (1 row with filter) + int search_row_filter(const void *query, const IndexQueryMeta &qmeta); + + //! Similarity search (n columns without filter) + int batch_search_column_nofilter(const void *query, + const IndexQueryMeta &qmeta, + size_t query_count); + + //! Similarity search (n columns with filter) + int batch_search_column_filter(const void *query, const IndexQueryMeta &qmeta, + size_t query_count); + + //! Similarity search (n rows without filter) + int batch_search_row_nofilter(const void *query, const IndexQueryMeta &qmeta, + size_t query_count); + + //! Similarity search (n rows with filter) + int batch_search_row_filter(const void *query, const IndexQueryMeta &qmeta, + size_t query_count); + + private: + const FlatSearcher *owner_{nullptr}; + uint32_t magic_{0}; + uint32_t topk_{0}; + uint32_t feature_size_{0}; + uint32_t actual_read_size_{0}; + IndexStorage::Segment::Pointer features_segment_{}; + std::vector result_heaps_{1}; + std::string batch_queries_{}; + float scores_[BATCH_SIZE * BATCH_SIZE]; + bool fetch_vector_{false}; + + // group + uint32_t group_num_{0}, group_topk_{0}; + std::map group_topk_heaps_{}; + std::vector group_results_{}; +}; + +template +int FlatSearcherContext::search_column_nofilter( + const void *query, const IndexQueryMeta &qmeta) { + IndexDocumentHeap *heap = &result_heaps_[0]; + heap->clear(); + heap->limit(topk_); + heap->set_threshold(this->threshold()); + + size_t left_size = features_segment_->data_size(); + size_t block_size = feature_size_ * BATCH_SIZE; + size_t read_offset = 0; + size_t feature_index = 0; + auto matrix = this->owner_->distance_matrix(); + + while (left_size >= actual_read_size_) { + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, actual_read_size_) != + actual_read_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + actual_read_size_); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < actual_read_size_; offset += block_size) { + matrix.template distance( + (const char *)data + offset, query, qmeta.dimension(), scores_); + + for (size_t i = 0; i != BATCH_SIZE; ++i) { + heap->emplace(0, scores_[i], feature_index++); + } + } + read_offset += actual_read_size_; + left_size -= actual_read_size_; + } + + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, left_size) != left_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + left_size); + return IndexError_ReadData; + } + + // Process left block features + size_t left_size_aligned = left_size / block_size * block_size; + for (size_t offset = 0; offset != left_size_aligned; offset += block_size) { + matrix.template distance((const char *)data + offset, query, + qmeta.dimension(), scores_); + + for (size_t i = 0; i != BATCH_SIZE; ++i) { + heap->emplace(0, scores_[i], feature_index++); + } + } + + // Process left single features + for (size_t offset = left_size_aligned; offset < left_size; + offset += feature_size_) { + float score; + matrix.template distance<1>((const char *)data + offset, query, + qmeta.dimension(), &score); + heap->emplace(0, score, feature_index++); + } + + for (auto &it : *heap) { + it.set_key(owner_->key(it.index())); + } + heap->sort(); + return 0; +} + +template +int FlatSearcherContext::search_column_filter( + const void *query, const IndexQueryMeta &qmeta) { + IndexDocumentHeap *heap = &result_heaps_[0]; + heap->clear(); + heap->limit(topk_); + heap->set_threshold(this->threshold()); + + size_t left_size = features_segment_->data_size(); + size_t block_size = feature_size_ * BATCH_SIZE; + size_t read_offset = 0; + size_t feature_index = 0; + auto matrix = owner_->distance_matrix(); + + while (left_size >= actual_read_size_) { + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, actual_read_size_) != + actual_read_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + actual_read_size_); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < actual_read_size_; offset += block_size) { + matrix.template distance( + (const char *)data + offset, query, qmeta.dimension(), scores_); + + for (size_t i = 0; i != BATCH_SIZE; ++i) { + uint64_t feature_key = owner_->key(feature_index); + + if (!this->filter()(feature_key)) { + if (group_by_search()) { + } + heap->emplace(feature_key, scores_[i], feature_index); + } + feature_index += 1; + } + } + read_offset += actual_read_size_; + left_size -= actual_read_size_; + } + + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, left_size) != left_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + left_size); + return IndexError_ReadData; + } + + // Process left block features + size_t left_size_aligned = left_size / block_size * block_size; + for (size_t offset = 0; offset != left_size_aligned; offset += block_size) { + matrix.template distance((const char *)data + offset, query, + qmeta.dimension(), scores_); + + for (size_t i = 0; i != BATCH_SIZE; ++i) { + uint64_t feature_key = owner_->key(feature_index); + + if (!this->filter()(feature_key)) { + heap->emplace(feature_key, scores_[i], feature_index); + } + feature_index += 1; + } + } + + // Process left single features + for (size_t offset = left_size_aligned; offset < left_size; + offset += feature_size_) { + uint64_t feature_key = owner_->key(feature_index); + if (!this->filter()(feature_key)) { + float score; + matrix.template distance<1>((const char *)data + offset, query, + qmeta.dimension(), &score); + heap->emplace(feature_key, score, feature_index); + } + feature_index += 1; + } + heap->sort(); + return 0; +} + +template +int FlatSearcherContext::search_row_nofilter( + const void *query, const IndexQueryMeta &qmeta) { + IndexDocumentHeap *heap = &result_heaps_[0]; + heap->clear(); + heap->limit(topk_); + heap->set_threshold(this->threshold()); + + size_t left_size = features_segment_->data_size(); + size_t read_offset = 0; + size_t feature_index = 0; + auto matrix = owner_->distance_matrix(); + + while (left_size >= actual_read_size_) { + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, actual_read_size_) != + actual_read_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + actual_read_size_); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < actual_read_size_; + offset += feature_size_) { + float score; + matrix.template distance<1>((const char *)data + offset, query, + qmeta.dimension(), &score); + heap->emplace(0, score, feature_index++); + } + read_offset += actual_read_size_; + left_size -= actual_read_size_; + } + + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, left_size) != left_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + left_size); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < left_size; offset += feature_size_) { + float score; + matrix.template distance<1>((const char *)data + offset, query, + qmeta.dimension(), &score); + heap->emplace(0, score, feature_index++); + } + for (auto &it : *heap) { + it.set_key(owner_->key(it.index())); + } + heap->sort(); + return 0; +} + +template +int FlatSearcherContext::search_row_filter( + const void *query, const IndexQueryMeta &qmeta) { + IndexDocumentHeap *heap = &result_heaps_[0]; + heap->clear(); + heap->limit(topk_); + heap->set_threshold(this->threshold()); + + size_t left_size = features_segment_->data_size(); + size_t read_offset = 0; + size_t feature_index = 0; + auto matrix = owner_->distance_matrix(); + + while (left_size >= actual_read_size_) { + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, actual_read_size_) != + actual_read_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + actual_read_size_); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < actual_read_size_; + offset += feature_size_) { + uint64_t feature_key = owner_->key(feature_index); + if (!this->filter()(feature_key)) { + float score; + matrix.template distance<1>((const char *)data + offset, query, + qmeta.dimension(), &score); + heap->emplace(feature_key, score, feature_index); + } + feature_index += 1; + } + read_offset += actual_read_size_; + left_size -= actual_read_size_; + } + + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, left_size) != left_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + left_size); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < left_size; offset += feature_size_) { + uint64_t feature_key = owner_->key(feature_index); + if (!this->filter()(feature_key)) { + float score; + matrix.template distance<1>((const char *)data + offset, query, + qmeta.dimension(), &score); + heap->emplace(feature_key, score, feature_index); + } + feature_index += 1; + } + heap->sort(); + return 0; +} + +template +int FlatSearcherContext::batch_search_column_nofilter( + const void *query, const IndexQueryMeta &qmeta, size_t query_count) { + // Initialize resources + result_heaps_.resize(query_count); + for (auto &heap : result_heaps_) { + heap.clear(); + heap.limit(topk_); + heap.set_threshold(this->threshold()); + } + + // Transpose queries + batch_queries_.clear(); + batch_queries_.reserve(query_count * qmeta.element_size()); + TransposeQueries(query, qmeta, query_count, &batch_queries_); + + size_t left_size = features_segment_->data_size(); + size_t block_size = feature_size_ * BATCH_SIZE; + size_t read_offset = 0; + size_t block_index = 0; + + // Process feature blocks + while (left_size >= actual_read_size_) { + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, actual_read_size_) != + actual_read_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + actual_read_size_); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < actual_read_size_; offset += block_size) { + this->batch_enqueue_nofilter( + (const char *)data + offset, block_index, 0, qmeta, query_count); + block_index += BATCH_SIZE; + } + read_offset += actual_read_size_; + left_size -= actual_read_size_; + } + + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, left_size) != left_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + left_size); + return IndexError_ReadData; + } + + // Process left block features + size_t left_size_aligned = left_size / block_size * block_size; + for (size_t offset = 0; offset != left_size_aligned; offset += block_size) { + this->batch_enqueue_nofilter( + (const char *)data + offset, block_index, 0, qmeta, query_count); + block_index += BATCH_SIZE; + } + + // Process left single features + for (size_t offset = left_size_aligned; offset < left_size; + offset += feature_size_) { + this->single_enqueue_nofilter( + (const char *)data + offset, block_index, 0, qmeta, query_count); + block_index += 1; + } + + // Normalize results + for (auto &heap : result_heaps_) { + for (auto &it : heap) { + it.set_key(owner_->key(it.index())); + } + heap.sort(); + } + return 0; +} + +template +int FlatSearcherContext::batch_search_column_filter( + const void *query, const IndexQueryMeta &qmeta, size_t query_count) { + // Initialize resources + result_heaps_.resize(query_count); + for (auto &heap : result_heaps_) { + heap.clear(); + heap.limit(topk_); + heap.set_threshold(this->threshold()); + } + + // Transpose queries + batch_queries_.clear(); + batch_queries_.reserve(query_count * qmeta.element_size()); + TransposeQueries(query, qmeta, query_count, &batch_queries_); + + size_t left_size = features_segment_->data_size(); + size_t block_size = feature_size_ * BATCH_SIZE; + size_t read_offset = 0; + size_t block_index = 0; + + // Process feature blocks + while (left_size >= actual_read_size_) { + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, actual_read_size_) != + actual_read_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + actual_read_size_); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < actual_read_size_; offset += block_size) { + size_t block_mask = 0; + for (size_t i = 0; i != BATCH_SIZE; ++i) { + if (!this->filter()(this->owner_->key(block_index + i))) { + block_mask |= (1 << i); + } + } + if (block_mask != 0) { + this->batch_enqueue_filter((const char *)data + offset, + block_index, block_mask, 0, + qmeta, query_count); + } + block_index += BATCH_SIZE; + } + read_offset += actual_read_size_; + left_size -= actual_read_size_; + } + + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, left_size) != left_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + left_size); + return IndexError_ReadData; + } + + // Process left block features + size_t left_size_aligned = left_size / block_size * block_size; + for (size_t offset = 0; offset != left_size_aligned; offset += block_size) { + size_t block_mask = 0; + for (size_t i = 0; i != BATCH_SIZE; ++i) { + if (!this->filter()(this->owner_->key(block_index + i))) { + block_mask |= (1 << i); + } + } + if (block_mask != 0) { + this->batch_enqueue_filter((const char *)data + offset, + block_index, block_mask, 0, qmeta, + query_count); + } + block_index += BATCH_SIZE; + } + + // Process left single features + for (size_t offset = left_size_aligned; offset < left_size; + offset += feature_size_) { + if (!this->filter()(owner_->key(block_index))) { + this->single_enqueue_nofilter( + (const char *)data + offset, block_index, 0, qmeta, query_count); + } + block_index += 1; + } + + // Normalize results + for (auto &heap : result_heaps_) { + for (auto &it : heap) { + it.set_key(owner_->key(it.index())); + } + heap.sort(); + } + return 0; +} + +template +int FlatSearcherContext::batch_search_row_nofilter( + const void *query, const IndexQueryMeta &qmeta, size_t query_count) { + // Initialize resources + result_heaps_.resize(query_count); + for (auto &heap : result_heaps_) { + heap.clear(); + heap.limit(topk_); + heap.set_threshold(this->threshold()); + } + + size_t left_size = features_segment_->data_size(); + size_t read_offset = 0; + size_t feature_index = 0; + auto matrix = owner_->distance_matrix(); + + // Process feature blocks + while (left_size >= actual_read_size_) { + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, actual_read_size_) != + actual_read_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + actual_read_size_); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < actual_read_size_; + offset += feature_size_) { + size_t query_offset = 0; + const void *feature = (const char *)data + offset; + + for (auto &heap : result_heaps_) { + float score; + matrix.template distance<1>(feature, (const char *)query + query_offset, + qmeta.dimension(), &score); + heap.emplace(0, score, feature_index); + query_offset += qmeta.element_size(); + } + feature_index += 1; + } + read_offset += actual_read_size_; + left_size -= actual_read_size_; + } + + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, left_size) != left_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + left_size); + return IndexError_ReadData; + } + + // Process left features + for (size_t offset = 0; offset < left_size; offset += feature_size_) { + size_t query_offset = 0; + const void *feature = (const char *)data + offset; + + for (auto &heap : result_heaps_) { + float score; + matrix.template distance<1>(feature, (const char *)query + query_offset, + qmeta.dimension(), &score); + heap.emplace(0, score, feature_index); + query_offset += qmeta.element_size(); + } + feature_index += 1; + } + + // Normalize results + for (auto &heap : result_heaps_) { + for (auto &it : heap) { + it.set_key(owner_->key(it.index())); + } + heap.sort(); + } + return 0; +} + +template +int FlatSearcherContext::batch_search_row_filter( + const void *query, const IndexQueryMeta &qmeta, size_t query_count) { + // Initialize resources + result_heaps_.resize(query_count); + for (auto &heap : result_heaps_) { + heap.clear(); + heap.limit(topk_); + heap.set_threshold(this->threshold()); + } + + size_t left_size = features_segment_->data_size(); + size_t read_offset = 0; + size_t feature_index = 0; + auto matrix = owner_->distance_matrix(); + + // Process feature blocks + while (left_size >= actual_read_size_) { + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, actual_read_size_) != + actual_read_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + actual_read_size_); + return IndexError_ReadData; + } + + for (size_t offset = 0; offset < actual_read_size_; + offset += feature_size_) { + uint64_t feature_key = owner_->key(feature_index); + + if (!this->filter()(feature_key)) { + size_t query_offset = 0; + const void *feature = (const char *)data + offset; + + for (auto &heap : result_heaps_) { + float score; + matrix.template distance<1>(feature, + (const char *)query + query_offset, + qmeta.dimension(), &score); + heap.emplace(feature_key, score, feature_index); + query_offset += qmeta.element_size(); + } + } + feature_index += 1; + } + read_offset += actual_read_size_; + left_size -= actual_read_size_; + } + + const void *data = nullptr; + if (features_segment_->read(read_offset, &data, left_size) != left_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + left_size); + return IndexError_ReadData; + } + + // Process left features + for (size_t offset = 0; offset < left_size; offset += feature_size_) { + uint64_t feature_key = owner_->key(feature_index); + + if (!this->filter()(feature_key)) { + size_t query_offset = 0; + const void *feature = (const char *)data + offset; + + for (auto &heap : result_heaps_) { + float score; + matrix.template distance<1>(feature, (const char *)query + query_offset, + qmeta.dimension(), &score); + heap.emplace(feature_key, score, feature_index); + query_offset += qmeta.element_size(); + } + } + feature_index += 1; + } + + // Normalize results + for (auto &heap : result_heaps_) { + heap.sort(); + } + return 0; +} + +template +int FlatSearcherContext::group_by_search_impl( + const void *query, const IndexQueryMeta &qmeta, uint32_t count) { + this->resize_group_results(count); + if (!this->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](uint64_t key) { + return this->group_by()(key); + }; + + auto provider = owner_->create_provider(); + + for (size_t q = 0; q < count; ++q) { + this->group_topk_heaps().clear(); + + for (node_id_t id = 0; id < provider->count(); ++id) { + if (!this->filter().is_valid() || !this->filter()(owner_->key(id))) { + dist_t dist = 0; + owner_->distance_matrix().template distance<1>( + query, provider->get_vector(owner_->key(id)), provider->dimension(), + &dist); + + std::string group_id = group_by(owner_->key(id)); + auto &topk_heap = this->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(this->group_topk()); + } + topk_heap.emplace(id, dist); + } + } + this->topk_to_group_result(q); + query = static_cast(query) + qmeta.element_size(); + } + return 0; +} + +template +int FlatSearcherContext::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count) { + auto provider = owner_->create_provider(); + if (this->group_by_search()) { + this->resize_group_results(count); + if (!this->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](uint64_t key) { + return this->group_by()(key); + }; + + for (size_t q = 0; q < count; ++q) { + this->group_topk_heaps().clear(); + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t pk = p_keys[q][idx]; + if (!this->filter().is_valid() || !this->filter()(pk)) { + dist_t dist = 0; + owner_->distance_matrix().template distance<1>( + query, provider->get_vector(pk), provider->dimension(), &dist); + + std::string group_id = group_by(pk); + auto &topk_heap = this->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(this->group_topk()); + } + topk_heap.emplace(owner_->get_id(pk), dist); + } + } + this->topk_to_group_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } else { + result_heaps_.resize(count); + for (auto &heap : result_heaps_) { + heap.clear(); + heap.limit(topk_); + heap.set_threshold(this->threshold()); + } + for (size_t q = 0; q < count; ++q) { + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t pk = p_keys[q][idx]; + if (!this->filter().is_valid() || !this->filter()(pk)) { + dist_t dist = 0; + owner_->distance_matrix().template distance<1>( + query, provider->get_vector(pk), provider->dimension(), &dist); + result_heaps_[q].emplace(pk, dist, owner_->get_id(pk)); + } + } + query = static_cast(query) + qmeta.element_size(); + } + for (auto &heap : result_heaps_) { + heap.sort(); + } + } + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_searcher_provider.h b/src/core/algorithm/flat/flat_searcher_provider.h new file mode 100644 index 0000000..758f89c --- /dev/null +++ b/src/core/algorithm/flat/flat_searcher_provider.h @@ -0,0 +1,220 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include "flat_distance_matrix.h" +#include "flat_searcher.h" +// #include "flat_streamer.h" +#include "flat_utility.h" + +namespace zvec { +namespace core { + +/*! Brute Force Searcher Provider + */ +template +class FlatSearcherProvider : public IndexProvider { + public: + //! Constructor + FlatSearcherProvider(const FlatSearcher *owner) { + feature_size_ = owner->meta().element_size(); + features_segment_ = owner->clone_features_segment(); + total_vector_count_ = + features_segment_->data_size() / owner->meta().element_size(); + owner_ = owner; + block_buffer_.resize(BATCH_SIZE * feature_size_); + } + + //! Create a new iterator + IndexProvider::Iterator::Pointer create_iterator(void) override { + return IndexProvider::Iterator::Pointer( + new (std::nothrow) FlatSearcherProvider::Iterator(owner_)); + } + + //! Retrieve count of vectors + size_t count(void) const override { + return total_vector_count_; + } + + //! Retrieve dimension of vector + size_t dimension(void) const override { + return owner_->meta().dimension(); + } + + //! Retrieve type of vector + IndexMeta::DataType data_type(void) const override { + return owner_->meta().data_type(); + } + + //! Retrieve vector size in bytes + size_t element_size(void) const override { + return owner_->meta().element_size(); + } + + //! Retrieve a vector using a primary key + const void *get_vector(uint64_t key) const override { + return this->get_vector_by_index(owner_->get_id(key)); + } + + //! Retrieve the owner class + const std::string &owner_class(void) const override { + return owner_->name(); + } + + protected: + /*! Brute Force Provider Iterator + */ + class Iterator : public IndexProvider::Iterator { + public: + //! Constructor + Iterator(const FlatSearcher *owner) { + block_buffer_.resize(BATCH_SIZE * owner->meta().element_size()); + feature_size_ = owner->meta().element_size(); + features_segment_ = owner->clone_features_segment(); + total_vector_count_ = + features_segment_->data_size() / owner->meta().element_size(); + owner_ = owner; + cursor_index_ = 0; + offset_ = 0; + this->next_block(); + } + + //! Retrieve pointer of data + //! NOTICE: the vec feature will be changed after iterating to next, so + //! the caller need to keep a copy of it before iterator to next vector + const void *data(void) const override { + return data_; + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return (!invalid_ && cursor_index_ < total_vector_count_); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return owner_->key(cursor_index_); + } + + //! Next iterator + void next(void) override { + ++cursor_index_; + + if ((cursor_index_ % BATCH_SIZE) != 0) { + data_ += feature_size_; + } else { + this->next_block(); + } + } + + protected: + //! Read a block of data + void next_block(void) { + const void *read_data = nullptr; + size_t read_size = 0; + + if (cursor_index_ >= total_vector_count_) { + invalid_ = true; + return; + } + + if (cursor_index_ + BATCH_SIZE < total_vector_count_) { + read_size = BATCH_SIZE * feature_size_; + } else { + read_size = (total_vector_count_ - cursor_index_) * feature_size_; + } + if (features_segment_->read(offset_, &read_data, read_size) != + read_size) { + LOG_ERROR("Failed to read data (%zu bytes) from features segment", + read_size); + invalid_ = true; + return; + } + offset_ += read_size; + + // The order of data may be a column format, convert it to the row format. + if (owner_->column_major_order() && + read_size == BATCH_SIZE * feature_size_) { + uint32_t align_size = + IndexMeta::AlignSizeof(owner_->meta().data_type()); + ReverseTranspose(align_size, read_data, + feature_size_ / align_size, + &block_buffer_[0]); + data_ = block_buffer_.data(); + } else { + data_ = reinterpret_cast(read_data); + } + } + + private: + const FlatSearcher *owner_{nullptr}; + IndexStorage::Segment::Pointer features_segment_{}; + uint32_t total_vector_count_{0}; + uint32_t feature_size_{0}; + std::vector block_buffer_{}; + const uint8_t *data_{nullptr}; + uint64_t offset_{0}; + uint32_t cursor_index_{0}; + bool invalid_{false}; + }; + + //! Retrieve a vector via local index + const void *get_vector_by_index(uint32_t index) const { + const void *read_data = nullptr; + if (index == kInvalidNodeId) { + LOG_ERROR("Failed to get vector by Invalid Id."); + return nullptr; + } + + if (owner_->column_major_order() && + index < (total_vector_count_ - (total_vector_count_ % BATCH_SIZE))) { + uint32_t block_size = feature_size_ * BATCH_SIZE; + uint64_t offset = (index - (index % BATCH_SIZE)) * feature_size_; + + if (features_segment_->read(offset, &read_data, block_size) != + block_size) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + block_size); + return nullptr; + } + + uint32_t align_size = IndexMeta::AlignSizeof(owner_->meta().data_type()); + ReverseTranspose( + align_size, read_data, feature_size_ / align_size, &block_buffer_[0]); + read_data = block_buffer_.data() + ((index % BATCH_SIZE) * feature_size_); + + } else { + if (features_segment_->read(index * feature_size_, &read_data, + feature_size_) != feature_size_) { + LOG_ERROR("Failed to read data (%u bytes) from features segment", + feature_size_); + return nullptr; + } + } + return read_data; + } + + private: + //! Members + const FlatSearcher *owner_{nullptr}; + IndexStorage::Segment::Pointer features_segment_{}; + uint32_t feature_size_{0}; + uint32_t total_vector_count_{0}; + mutable std::vector block_buffer_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc new file mode 100644 index 0000000..9f5b9ec --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -0,0 +1,491 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_streamer.h" +#include +#include "flat_streamer_context.h" +#include "flat_streamer_dumper.h" +#include "flat_streamer_provider.h" + +namespace zvec { +namespace core { + +#define WRITE_LOCK_GUARD(MUTEX, LOCK_NAME) \ + ailego::WriteLock write_lock(MUTEX); \ + std::unique_lock LOCK_NAME(write_lock); + +#define READ_LOCK_GUARD_DEFER(MUTEX, LOCK_NAME) \ + ailego::ReadLock read_lock(MUTEX); \ + std::unique_lock LOCK_NAME(read_lock, std::defer_lock); + +template +FlatStreamer::FlatStreamer() : entity_(stats_) {} + +template +FlatStreamer::~FlatStreamer() { + if (state_ == STATE_INITED || state_ == STATE_OPENED) { + this->cleanup(); + } +} + +template +int FlatStreamer::init(const IndexMeta &imeta, + const ailego::Params ¶ms) { + meta_ = imeta; + meta_.set_streamer("FlatStreamer", 0U, params); + + int error_code = InitializeMetric(meta_, &metric_); + if (error_code != 0) { + LOG_ERROR("Failed to initialize index metric %s, error=%d, %s", + meta_.metric_name().c_str(), error_code, + IndexError::What(error_code)); + return error_code; + } + if (metric_->query_metric()) { + metric_ = metric_->query_metric(); + } + + // 参数设置 + if (params.get(PARAM_FLAT_COLUMN_MAJOR_ORDER, &column_major_order_)) { + meta_.set_major_order(column_major_order_ ? IndexMeta::MO_COLUMN + : IndexMeta::MO_ROW); + } + // Verify column major order + if (meta_.major_order() != IndexMeta::MO_ROW) { + IndexMeta::DataType ft = meta_.data_type(); + + bool support_column_major = true; + if ((ft != IndexMeta::DT_FP32 && ft != IndexMeta::DT_FP16 && + ft != IndexMeta::DT_INT8 && ft != IndexMeta::DT_INT4 && + ft != IndexMeta::DT_BINARY32 && ft != IndexMeta::DT_BINARY64) || + (meta_.unit_size() != IndexMeta::UnitSizeof(ft))) { + if (meta_.major_order() == IndexMeta::MO_COLUMN) { + LOG_ERROR("Unsupported type %d with unit size %u.", ft, + meta_.unit_size()); + return IndexError_Unsupported; + } else { + support_column_major = false; + } + } + if (meta_.element_size() % IndexMeta::AlignSizeof(ft) != 0) { + if (meta_.major_order() == IndexMeta::MO_COLUMN) { + LOG_ERROR("Unsupported type %d with dimension %u.", ft, + meta_.dimension()); + return IndexError_Unsupported; + } else { + support_column_major = false; + } + } + + if (meta_.major_order() == IndexMeta::MO_UNDEFINED && + support_column_major) { + meta_.set_major_order(IndexMeta::MO_ROW); + } + } + + if (!VerifyMetric(meta_)) { + LOG_ERROR("Invalid index metric %s.", meta_.metric_name().c_str()); + return IndexError_InvalidArgument; + } + + read_block_size_ = FLAT_DEFAULT_READ_BLOCK_SIZE; + params.get(PARAM_FLAT_READ_BLOCK_SIZE, &read_block_size_); + params.get(PARAM_FLAT_USE_ID_MAP, &use_key_info_map_); + + // entity init + uint32_t block_vector_count = kDefaultBlockVecCount; + uint32_t segment_size = kDefaultSegmentSize; + bool filter_same_key = true; + entity_.set_block_vector_count(block_vector_count); + entity_.set_segment_size(segment_size); + entity_.enable_filter_same_key(filter_same_key); + entity_.set_linear_list_count(1); + entity_.set_use_key_info_map(use_key_info_map_); + *entity_.mutable_meta() = meta_; + + state_ = STATE_INITED; + + return 0; +} + +template +int FlatStreamer::cleanup() { + if (state_ == STATE_OPENED) { + this->close(); + } + + LOG_DEBUG("FlatStreamer cleanup"); + state_ = STATE_INIT; + return 0; +} + +template +int FlatStreamer::open(IndexStorage::Pointer stg) { + if (!stg) { + LOG_ERROR("Failed to open for invalid storage"); + return IndexError_InvalidArgument; + } + if (ailego_unlikely(state_ != STATE_INITED)) { + LOG_ERROR("Open storage failed, init streamer first!"); + return IndexError_NoReady; + } + + LOG_DEBUG("FlatStreamer open with %s", stg->name().c_str()); + + int ret = entity_.open(std::move(stg), meta_); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to open storage"); + return ret; + } + magic_ = IndexContext::GenerateMagic(); + + state_ = STATE_OPENED; + + return 0; +} + +template +int FlatStreamer::close(void) { + LOG_DEBUG("FlatStreamer close"); + + entity_.flush_linear_meta(); + + stats_.clear(); + + int ret = entity_.close(); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + state_ = STATE_INITED; + return 0; +} + +template +int FlatStreamer::flush(uint64_t checkpoint) { + LOG_INFO("FlatStreamer flush with checkpoint %zu", (size_t)checkpoint); + return entity_.flush(checkpoint); +} + +template +int FlatStreamer::dump(const IndexDumper::Pointer &dumper) { + std::string searcher_name = "FlatSearcher"; + if constexpr (BATCH_SIZE == 16) { + searcher_name = "FlatSearcher16"; + } + meta_.set_searcher(searcher_name, 0U, ailego::Params()); + WRITE_LOCK_GUARD(dump_mutex_, dump_lock); + std::shared_ptr> bf_dumper = + std::make_shared>(this); + int ret = bf_dumper->dump(dumper); + *(stats_.mutable_dumped_size()) += bf_dumper->dump_size(); + return ret; +} + +template +IndexStreamer::Context::UPointer FlatStreamer::create_context( + void) const { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to create Context, open storage first!"); + return Context::UPointer(); + } + return IndexStreamer::Context::Pointer( + new FlatStreamerContext(this)); +} + +template +IndexProvider::Pointer FlatStreamer::create_provider(void) const { + return IndexProvider::Pointer(new (std::nothrow) + FlatStreamerProvider(this)); +} + +template +int FlatStreamer::add_impl(uint64_t pkey, const void *query, + const IndexQueryMeta &qmeta, + Context::UPointer &context) { + if (!query || qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size()) { + LOG_ERROR( + "Failed to add for invalid arguments, query=%p, qmeta(type=%u " + "dim=%u size=%u) vs meta(type=%u dim=%u size=%u)", + query, qmeta.data_type(), qmeta.dimension(), qmeta.element_size(), + meta_.data_type(), meta_.dimension(), meta_.element_size()); + (*stats_.mutable_discarded_count())++; + return IndexError_InvalidArgument; + } + + auto *ctx = dynamic_cast *>(context.get()); + if (!ctx) { + LOG_ERROR("Failed to cast FlatStreamerContext"); + (*stats_.mutable_discarded_count())++; + return IndexError_Cast; + } + + READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock); + + if (!dump_lock.try_lock()) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + + // IndexQueryMeta iv_qmeta; + // int ret = entity_.convert(query, qmeta, &query, &iv_qmeta); + // if (ret != 0) { + // LOG_ERROR("Failed to convert record for %s", + // IndexError::What(ret)); + // (*stats_.mutable_discarded_count())++; + // return ret; + // } + + int ret = entity_.add(pkey, query, qmeta.element_size()); + if (ret != 0) { + LOG_ERROR("Failed to add record for %s", IndexError::What(ret)); + (*stats_.mutable_discarded_count())++; + return ret; + } + + return 0; +} + +template +int FlatStreamer::add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + if (!query || qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size()) { + LOG_ERROR( + "Failed to add for invalid arguments, query=%p, qmeta(type=%u " + "dim=%u size=%u) vs meta(type=%u dim=%u size=%u)", + query, qmeta.data_type(), qmeta.dimension(), qmeta.element_size(), + meta_.data_type(), meta_.dimension(), meta_.element_size()); + (*stats_.mutable_discarded_count())++; + return IndexError_InvalidArgument; + } + + auto *ctx = dynamic_cast *>(context.get()); + if (!ctx) { + LOG_ERROR("Failed to cast FlatStreamerContext"); + (*stats_.mutable_discarded_count())++; + return IndexError_Cast; + } + + READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock); + + if (!dump_lock.try_lock()) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + + int ret = entity_.add_vector_with_id(id, query, qmeta.element_size()); + if (ret != 0) { + LOG_ERROR("Failed to add record for %s", IndexError::What(ret)); + (*stats_.mutable_discarded_count())++; + return ret; + } + + return 0; +} + +template +int FlatStreamer::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + ailego_assert(query && count && !!context); + ailego_assert(metric_->is_matched(meta_, qmeta)); + + FlatStreamerContext *bf_context = + dynamic_cast *>(context.get()); + if (!bf_context) { + LOG_ERROR("Invalid brute-force streamer context"); + return IndexError_InvalidArgument; + } + + if (bf_context->magic() != magic_) { + bf_context->reset(this); + } + + if (bf_context->group_by_search()) { + return group_by_search_impl(query, qmeta, count, context); + } + + bf_context->reset_results(count); + auto &filter = bf_context->filter(); + + for (size_t q = 0; q < count; ++q) { + auto *heap = bf_context->result_heap(); + auto *context_stats = bf_context->mutable_stats(q); + uint32_t scan_count = 0; + int ret = entity_.search(query, filter, &scan_count, heap, context_stats); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to search for %s", IndexError::What(ret)); + return ret; + } + heap->sort(); + bf_context->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + return 0; +} + +template +int FlatStreamer::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + ailego_assert(query && count && !!context); + ailego_assert(metric_->is_matched(meta_, qmeta)); + + FlatStreamerContext *bf_context = + dynamic_cast *>(context.get()); + if (!bf_context) { + LOG_ERROR("Invalid brute-force streamer context"); + return IndexError_InvalidArgument; + } + + if (bf_context->magic() != magic_) { + bf_context->reset(this); + } + + if (bf_context->group_by_search()) { + return group_by_search_p_keys_impl(query, p_keys, qmeta, count, context); + } + + bf_context->reset_results(count); + auto &filter = bf_context->filter(); + + for (size_t q = 0; q < count; ++q) { + auto *heap = bf_context->result_heap(); + for (node_id_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t key = p_keys[q][idx]; + if (!filter.is_valid() || !filter(key)) { + dist_t dist = 0; + IndexStorage::MemoryBlock block; + if (entity_.get_vector_by_key(key, block) != 0) continue; + entity_.row_major_distance(query, block.data(), 1, &dist); + heap->emplace(key, dist); + } + } + heap->sort(); + bf_context->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + return 0; +} + +template +int FlatStreamer::group_by_search_impl( + const void *query, const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + FlatStreamerContext *bf_context = + dynamic_cast *>(context.get()); + if (!bf_context) { + LOG_ERROR("Invalid brute-force streamer context"); + return IndexError_InvalidArgument; + } + + bf_context->resize_group_results(count); + if (!bf_context->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](uint64_t key) { + return bf_context->group_by()(key); + }; + + auto iterator = entity_.creater_iterator(); + + for (size_t q = 0; q < count; ++q) { + bf_context->group_topk_heaps().clear(); + for (node_id_t id = 0; id < entity_.vector_count(); ++id) { + uint64_t key = entity_.key(id); + if (!bf_context->filter().is_valid() || !bf_context->filter()(key)) { + dist_t dist = 0; + IndexStorage::MemoryBlock block; + if (entity_.get_vector_by_key(key, block) != 0) continue; + entity_.row_major_distance(query, block.data(), 1, &dist); + + std::string group_id = group_by(key); + auto &topk_heap = bf_context->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(bf_context->group_topk()); + } + topk_heap.emplace(key, dist); + } + } + bf_context->topk_to_group_result(q); + query = static_cast(query) + qmeta.element_size(); + } + return 0; +} + +template +int FlatStreamer::group_by_search_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + FlatStreamerContext *bf_context = + dynamic_cast *>(context.get()); + if (!bf_context) { + LOG_ERROR("Invalid brute-force streamer context"); + return IndexError_InvalidArgument; + } + + bf_context->resize_group_results(count); + if (!bf_context->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](uint64_t key) { + return bf_context->group_by()(key); + }; + + auto iterator = entity_.creater_iterator(); + + for (size_t q = 0; q < count; ++q) { + bf_context->group_topk_heaps().clear(); + for (node_id_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t key = p_keys[q][idx]; + if (!bf_context->filter().is_valid() || !bf_context->filter()(key)) { + dist_t dist = 0; + IndexStorage::MemoryBlock block; + if (entity_.get_vector_by_key(key, block) != 0) continue; + entity_.row_major_distance(query, block.data(), 1, &dist); + + std::string group_id = group_by(key); + auto &topk_heap = bf_context->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(bf_context->group_topk()); + } + topk_heap.emplace(key, dist); + } + } + bf_context->topk_to_group_result(q); + query = static_cast(query) + qmeta.element_size(); + } + return 0; +} + +INDEX_FACTORY_REGISTER_STREAMER_ALIAS(LinearStreamer, FlatStreamer<32>); +INDEX_FACTORY_REGISTER_STREAMER_ALIAS(FlatStreamer, FlatStreamer<32>); +INDEX_FACTORY_REGISTER_STREAMER_ALIAS(FlatStreamer16, FlatStreamer<16>); +INDEX_FACTORY_REGISTER_STREAMER_ALIAS(FlatStreamer32, FlatStreamer<32>); +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_streamer.h b/src/core/algorithm/flat/flat_streamer.h new file mode 100644 index 0000000..9248e7c --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer.h @@ -0,0 +1,190 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include "flat_streamer_entity.h" +#include "flat_utility.h" + +namespace zvec { +namespace core { + +/*! Flat Streamer + */ +template +class FlatStreamer : public IndexStreamer { + public: + using ContextPointer = IndexStreamer::Context::UPointer; + + FlatStreamer(void); + ~FlatStreamer(void) override; + + FlatStreamer(const FlatStreamer &streamer) = delete; + FlatStreamer &operator=(const FlatStreamer &streamer) = delete; + + public: + //! Initialize Streamer + int init(const IndexMeta &, const ailego::Params &) override; + + //! Cleanup Streamer + int cleanup(void) override; + + //! Create a context + IndexStreamer::Context::UPointer create_context(void) const override; + + //! Create a new iterator + IndexProvider::Pointer create_provider(void) const override; + + //! Add a vector into index + int add_impl(uint64_t pkey, const void *query, const IndexQueryMeta &qmeta, + Context::UPointer &context) override; + + int add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::UPointer &context) const override { + return search_bf_impl(query, qmeta, 1, context); + } + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::UPointer &context) const override { + return search_bf_impl(query, qmeta, count, context); + } + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::UPointer &context) const override { + return search_bf_impl(query, qmeta, 1, context); + } + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::UPointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + Context::UPointer &context) const override { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::UPointer &context) const override; + + int group_by_search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::UPointer &context) const; + + int group_by_search_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const; + + //! Open index from file path + int open(IndexStorage::Pointer stg) override; + + //! Close file + int close(void) override; + + //! flush file + int flush(uint64_t checkpoint) override; + + //! Dump index into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + const FlatStreamerEntity &entity(void) const { + return entity_; + } + + const void *get_vector(uint64_t key) const override { + return this->get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return this->get_vector_by_key(key, block); + } + + const void *get_vector_by_key(uint64_t key) const { + return entity_.get_vector_by_key(key); + } + + int get_vector_by_key(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_.get_vector_by_key(key, block); + } + const void *get_vector_by_id(uint32_t id) const override { + return get_vector_by_key(id); + } + + int get_vector_by_id(const uint32_t id, + IndexStorage::MemoryBlock &block) const override { + return get_vector_by_key(id, block); + } + + uint32_t magic(void) const { + return magic_; + } + + //! Retrieve block size of data read + uint32_t read_block_size(void) const { + return read_block_size_; + } + + private: + //! Constants + static constexpr uint32_t kDefaultBlockVecCount = 32u; + static constexpr uint32_t kDefaultSegmentSize = 4 * 1024 * 1024u; + static constexpr float kDefaultDocsSoftLimitRatio = 0.9f; + + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 }; + + //! Members + uint32_t magic_{0}; + uint32_t docs_hard_limit_{std::numeric_limits::max()}; + uint32_t docs_soft_limit_{0}; + IndexMeta meta_{}; + std::vector> data_; + IndexStreamer::Stats stats_{}; + IndexMetric::Pointer metric_{}; + State state_{STATE_INIT}; + mutable std::mutex mapping_mutex_{}; + ailego::SharedMutex dump_mutex_{}; + bool column_major_order_{false}; + bool use_key_info_map_{true}; + uint32_t read_block_size_{0}; + FlatStreamerEntity entity_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_streamer_context.h b/src/core/algorithm/flat/flat_streamer_context.h new file mode 100644 index 0000000..6561787 --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer_context.h @@ -0,0 +1,259 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "flat_streamer.h" + +namespace zvec { +namespace core { + +/*! Brute Force Streamer Context + */ +template +class FlatStreamerContext : public IndexStreamer::Context { + public: + //! Constructor + FlatStreamerContext(const FlatStreamer *owner) { + this->reset(owner); + } + + //! Destructor + ~FlatStreamerContext(void) override = default; + + //! Set topk of search result + void set_topk(uint32_t topk) override { + topk_ = topk; + result_heap_.limit(topk); + } + + //! Retrieve search result + const IndexDocumentList &result(void) const override { + return results_[0]; + } + + //! Retrieve search result with index + const IndexDocumentList &result(size_t idx) const override { + return results_[idx]; + } + + //! Retrieve result object for output + IndexDocumentList *mutable_result(size_t idx) override { + ailego_assert_with(idx < results_.size(), "invalid idx"); + return &results_[idx]; + } + + inline IndexDocumentHeap *result_heap() { + return &result_heap_; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(void) const override { + return group_results_[0]; + } + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(size_t idx) const override { + return group_results_[idx]; + } + + IndexGroupDocumentList *mutable_group_result(void) override { + return &group_results_[0]; + } + + IndexGroupDocumentList *mutable_group_result(size_t idx) override { + return &group_results_[idx]; + } + + //! Update the parameters of context + int update(const ailego::Params & /*params*/) override { + return 0; + } + + //! Retrieve magic number + uint32_t magic(void) const override { + return magic_; + } + + //! Get group topk + inline uint32_t group_topk() const { + return group_topk_; + } + //! Get group num + inline uint32_t group_num() const { + return group_num_; + } + inline std::map &group_topk_heaps() { + return group_topk_heaps_; + } + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + bool fetch_vector() const override { + return fetch_vector_; + } + inline void resize_group_results(size_t size) { + if (group_by_search()) { + group_results_.resize(size); + } + } + + void topk_to_result(uint32_t idx) { + if (ailego_unlikely(result_heap_.size() == 0)) { + return; + } + + ailego_assert_with(idx < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(result_heap_.size())); + result_heap_.sort(); + results_[idx].clear(); + for (int i = 0; i < size; ++i) { + auto score = result_heap_[i].score(); + if (score > this->threshold()) { + break; + } + + key_t key = result_heap_[i].key(); + if (fetch_vector_) { + IndexStorage::MemoryBlock block; + owner_->entity().get_vector_by_key(key, block); + results_[idx].emplace_back(key, score, key, block); + } else { + results_[idx].emplace_back(key, score, key); + } + } + } + + void topk_to_group_result(uint32_t idx) { + ailego_assert_with(idx < group_results_.size(), "invalid idx"); + group_results_[idx].clear(); + std::vector> group_topk_list; + std::vector> best_score_in_groups; + for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end(); + itr++) { + const std::string &group_id = (*itr).first; + auto &heap = (*itr).second; + heap.sort(); + if (heap.size() > 0) { + float best_score = heap[0].second; + best_score_in_groups.push_back(std::make_pair(group_id, best_score)); + } + } + std::sort(best_score_in_groups.begin(), best_score_in_groups.end(), + [](const std::pair &a, + const std::pair &b) -> int { + return a.second < b.second; + }); + // truncate to group num + for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size(); + ++i) { + const std::string &group_id = best_score_in_groups[i].first; + group_topk_list.emplace_back( + std::make_pair(group_id, group_topk_heaps_[group_id])); + } + group_results_[idx].resize(group_topk_list.size()); + for (uint32_t i = 0; i < group_topk_list.size(); ++i) { + const std::string &group_id = group_topk_list[i].first; + group_results_[idx][i].set_group_id(group_id); + uint32_t size = std::min( + group_topk_, static_cast(group_topk_list[i].second.size())); + for (uint32_t j = 0; j < size; ++j) { + auto score = group_topk_list[i].second[j].second; + if (score > this->threshold()) { + break; + } + node_id_t id = group_topk_list[i].second[j].first; + auto provider = owner_->create_provider(); + if (fetch_vector_) { + IndexStorage::MemoryBlock block; + provider->get_vector(id, block); + group_results_[idx][i].mutable_docs()->emplace_back(id, score, id, + block); + } else { + group_results_[idx][i].mutable_docs()->emplace_back(id, score, id); + } + } + } + } + + //! Get if group by search + bool group_by_search() { + return group_num_ > 0; + } + //! Set group params + void set_group_params(uint32_t group_num, uint32_t group_topk) override { + group_num_ = group_num; + group_topk_ = group_topk; + group_topk_heaps_.clear(); + } + + void reset() override { + for (auto &it : results_) { + it.clear(); + } + for (auto &it : group_results_) { + it.clear(); + } + } + + //! Reset the context + void reset(const FlatStreamer *owner) { + this->reset(); + magic_ = owner->magic(); + feature_size_ = owner->meta().element_size(); + + uint32_t block_size = feature_size_ * BATCH_SIZE; + actual_read_size_ = + (owner->read_block_size() + block_size - 1) / block_size * block_size; + owner_ = owner; + } + + //! Reset all the query results + void reset_results(size_t qnum) { + results_.resize(qnum); + stats_vec_.resize(qnum); + for (size_t i = 0; i < qnum; ++i) { + results_[i].clear(); + stats_vec_[i].clear(); + } + result_heap_.clear(); + result_heap_.limit(topk_); + result_heap_.set_threshold(this->threshold()); + } + + Stats *mutable_stats(size_t idx = 0) { + ailego_assert_with(stats_vec_.size() > idx, "invalid index"); + return &stats_vec_[idx]; + } + + private: + const FlatStreamer *owner_{nullptr}; + std::vector stats_vec_{}; + uint32_t magic_{0}; + uint32_t topk_{0}; + uint32_t feature_size_{0}; + uint32_t actual_read_size_{0}; + IndexDocumentHeap result_heap_; + std::vector results_{}; + std::string batch_queries_{}; + float scores_[BATCH_SIZE * BATCH_SIZE]; + bool fetch_vector_{false}; + // group + uint32_t group_num_{0}; + uint32_t group_topk_{0}; + std::map group_topk_heaps_{}; + std::vector group_results_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_streamer_dumper.h b/src/core/algorithm/flat/flat_streamer_dumper.h new file mode 100644 index 0000000..30b2994 --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer_dumper.h @@ -0,0 +1,244 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include "flat_streamer.h" +#include "flat_utility.h" + +namespace zvec { +namespace core { + +template +class FlatStreamerDumper { + public: + typedef std::unique_ptr Pointer; + + FlatStreamerDumper(const FlatStreamer *owner) { + owner_ = owner; + dump_size_ = 0; + } + + int dump(const IndexDumper::Pointer &dumper) { + ailego::ElapsedTime stamp; + + std::vector keys; + if (owner_->meta().major_order() == IndexMeta::MO_COLUMN) { + int error_code = this->write_column_index(dumper.get(), &keys); + if (error_code != 0) { + return error_code; + } + } else { + int error_code = this->write_row_index(dumper.get(), &keys); + if (error_code != 0) { + return error_code; + } + } + + int error_code = this->write_keys(keys, dumper.get()); + if (error_code != 0) { + return error_code; + } + + error_code = this->write_mapping(keys, dumper.get()); + if (error_code != 0) { + return error_code; + } + + error_code = IndexHelper::SerializeToDumper(owner_->meta(), dumper.get()); + if (error_code != 0) { + return error_code; + } + LOG_DEBUG("dumped_count: %zu, costtime: %zu", keys.size(), + (size_t)stamp.milli_seconds()); + return 0; + } + + size_t dump_size() { + return dump_size_; + } + + private: + int write_column_index(IndexDumper *dumper, std::vector *keys) { + switch (IndexMeta::AlignSizeof(owner_->meta().data_type())) { + case 2: + return this->write_column_index(dumper, keys); + case 4: + return this->write_column_index(dumper, keys); + case 8: + return this->write_column_index(dumper, keys); + default: + ailego_check_with(0, "BAD CASE"); + } + return IndexError_Runtime; + } + + template + int write_column_index(IndexDumper *dumper, std::vector *keys) { + auto iter = owner_->entity().creater_iterator(); + if (!iter) { + LOG_ERROR("Failed to create iterator"); + return IndexError_Runtime; + } + + // Write features + size_t element_size = owner_->meta().element_size(); + size_t block_size = element_size * BATCH_SIZE; + std::string block1, block2; + block1.reserve(block_size); + block2.reserve(block_size); + + for (; iter->is_valid(); iter->next()) { + block1.append(reinterpret_cast(iter->data()), element_size); + keys->emplace_back(iter->key()); + + if (block1.size() == block_size) { + ailego::MatrixHelper::Transpose( + block1.data(), element_size / sizeof(T), (void *)block2.data()); + + if (dumper->write(block2.data(), block_size) != block_size) { + LOG_ERROR("Failed to write data into dumper %s", + dumper->name().c_str()); + return IndexError_WriteData; + } + block1.clear(); + dump_size_ += block_size; + } + } + + if (!block1.empty()) { + if (dumper->write(block1.data(), block1.size()) != block1.size()) { + LOG_ERROR("Failed to write data into dumper %s", + dumper->name().c_str()); + return IndexError_WriteData; + } + dump_size_ += block1.size(); + } + + // Write the padding if need + size_t features_size = keys->size() * element_size; + size_t features_padding_size = + ailego_align(features_size, 32) - features_size; + if (features_padding_size) { + std::string padding(features_padding_size, '\0'); + + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", + dumper->name().c_str()); + return IndexError_WriteData; + } + dump_size_ += padding.size(); + } + return dumper->append(FLAT_SEGMENT_FEATURES_SEG_ID, features_size, + features_padding_size, 0); + } + + int write_row_index(IndexDumper *dumper, std::vector *keys) { + auto iter = owner_->entity().creater_iterator(); + if (!iter) { + LOG_ERROR("Failed to create iterator"); + return IndexError_Runtime; + } + + // Write features + size_t element_size = owner_->meta().element_size(); + for (; iter->is_valid(); iter->next()) { + if (dumper->write(iter->data(), element_size) != element_size) { + LOG_ERROR("Failed to write data into dumper %s", + dumper->name().c_str()); + return IndexError_WriteData; + } + dump_size_ += element_size; + keys->emplace_back(iter->key()); + } + + // Write the padding if need + size_t features_size = keys->size() * element_size; + size_t features_padding_size = + ailego_align(features_size, 32) - features_size; + if (features_padding_size) { + std::string padding(features_padding_size, '\0'); + + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", + dumper->name().c_str()); + return IndexError_WriteData; + } + dump_size_ += padding.size(); + } + return dumper->append(FLAT_SEGMENT_FEATURES_SEG_ID, features_size, + features_padding_size, 0); + } + + int write_keys(const std::vector &keys, IndexDumper *dumper) { + size_t keys_size = keys.size() * sizeof(uint64_t); + size_t keys_padding_size = ailego_align(keys_size, 32) - keys_size; + if (dumper->write(keys.data(), keys_size) != keys_size) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + dump_size_ += keys_size; + + // Write the padding if need + if (keys_padding_size) { + std::string padding(keys_padding_size, '\0'); + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", + dumper->name().c_str()); + return IndexError_WriteData; + } + dump_size_ += padding.size(); + } + return dumper->append(FLAT_SEGMENT_KEYS_SEG_ID, keys_size, + keys_padding_size, 0); + } + + int write_mapping(const std::vector &keys, IndexDumper *dumper) { + std::vector mapping(keys.size()); + std::iota(mapping.begin(), mapping.end(), 0); + std::sort(mapping.begin(), mapping.end(), + [&keys](uint32_t lhs, uint32_t rhs) { + return (keys[lhs] < keys[rhs]); + }); + + size_t mapping_size = mapping.size() * sizeof(uint32_t); + size_t mapping_padding_size = ailego_align(mapping_size, 32) - mapping_size; + if (dumper->write(mapping.data(), mapping_size) != mapping_size) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + dump_size_ += mapping_size; + + // Write the padding if need + if (mapping_padding_size) { + std::string padding(mapping_padding_size, '\0'); + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", + dumper->name().c_str()); + return IndexError_WriteData; + } + dump_size_ += padding.size(); + } + return dumper->append(FLAT_SEGMENT_MAPPING_SEG_ID, mapping_size, + mapping_padding_size, 0); + } + + private: + const FlatStreamer *owner_{nullptr}; + size_t dump_size_{0}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_streamer_entity.cc b/src/core/algorithm/flat/flat_streamer_entity.cc new file mode 100644 index 0000000..87d9a19 --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -0,0 +1,1004 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_streamer_entity.h" +#include +#include +#include "flat_utility.h" + +namespace zvec { +namespace core { + +FlatStreamerEntity::FlatStreamerEntity(IndexStreamer::Stats &stats) + : stats_(stats) {} + +int FlatStreamerEntity::open(IndexStorage::Pointer storage, + const IndexMeta & /*mt*/) { + if (storage_) { + LOG_ERROR("An storage instance is already opened"); + return IndexError_Duplicate; + } + // segments_[0] store the meta information of the linear list + ailego_assert_with(segments_.size() == 0, "Invalid Size"); + + key_info_map_lock_ = std::make_shared(); + key_info_map_.clear(); + id_key_vector_.clear(); + withid_key_info_map_.clear(); + withid_key_map_.clear(); + + vec_unit_size_ = IndexMeta::AlignSizeof(index_meta_.data_type()); + vec_cols_ = index_meta_.element_size() / vec_unit_size_; + meta_.header.block_size = + ailego_align(sizeof(BlockHeader) + sizeof(DeletionMap) + + (index_meta_.element_size() + sizeof(uint64_t)) * + meta_.header.block_vector_count, + 32); + + if (storage->get(FLAT_LINEAR_LIST_HEAD_SEG_ID) || + storage->get(FLAT_LINEAR_META_SEG_ID)) { + int ret = this->load_storage(storage); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to load storage index"); + return ret; + } + } else { + int ret = this->init_storage(storage); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to init storage"); + return ret; + } + } + + storage_ = storage; + + //! Create the distance calculator + auto metric = IndexFactory::CreateMetric(index_meta_.metric_name()); + if (!metric) { + LOG_ERROR("Failed to create metric %s", index_meta_.metric_name().c_str()); + return IndexError_NoExist; + } + int ret = metric->init(index_meta_, index_meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("Failed to initialize metric %s", + index_meta_.metric_name().c_str()); + return ret; + } + row_distance_ = metric->distance(); + column_distance_ = + metric->distance_matrix(meta_.header.block_vector_count, 1); + + LOG_DEBUG("Open storage %s done, metric=%s", storage_->name().c_str(), + index_meta_.metric_name().c_str()); + + return 0; +} + +int FlatStreamerEntity::close(void) { + segments_.clear(); + storage_.reset(); + key_info_map_lock_.reset(); + key_info_map_.clear(); + withid_key_info_map_.clear(); + withid_key_map_.clear(); + id_key_vector_.clear(); + meta_.create_time = 0; + meta_.update_time = 0; + meta_.segment_count = 0; + meta_.header.total_vector_count = 0; + meta_.header.block_count = 0; + meta_.header.block_size = 0; + meta_.header.linear_body_size = 0; + + return 0; +} + +int FlatStreamerEntity::flush_linear_meta(void) { + if (!storage_) { + return 0; + } + + meta_.update_time = ailego::Realtime::Seconds(); + meta_.revision_id = stats_.revision_id(); + stats_.set_update_time(meta_.update_time); + auto segment = storage_->get(FLAT_LINEAR_META_SEG_ID); + if (ailego_unlikely(!segment)) { + LOG_ERROR("Failed to get segment %s", FLAT_LINEAR_META_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) { + LOG_ERROR("Failed to write segment %s", FLAT_LINEAR_META_SEG_ID.c_str()); + return IndexError_WriteData; + } + + return 0; +} + +int FlatStreamerEntity::flush(uint64_t checkpoint) { + int ret = this->flush_linear_meta(); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + if (checkpoint != 0) { + storage_->refresh(checkpoint); + } + ret = storage_->flush(); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to refresh storage for %s", IndexError::What(ret)); + return ret; + } + if (checkpoint != 0) { + stats_.set_check_point(checkpoint); + } + + return 0; +} + +int FlatStreamerEntity::add(uint64_t key, const void *vec, size_t size) { + std::lock_guard lock(mutex_); + if (filter_same_key_) { + key_info_map_lock_->lock_shared(); + if (key_info_map_.find(key) != key_info_map_.end()) { + key_info_map_lock_->unlock_shared(); + LOG_WARN("Try to add duplicate key, drop it"); + return IndexError_Duplicate; + } + key_info_map_lock_->unlock_shared(); + } + if (size != static_cast(index_meta_.element_size())) { + LOG_ERROR("Failed to add, mismatch size %zu vs elemsize %u", size, + index_meta_.element_size()); + return IndexError_Mismatch; + } + + IndexStorage::MemoryBlock head_block; + this->get_head_block(head_block); + BlockLocation block; + { + const BlockLocation *bl = + reinterpret_cast(head_block.data()); + if (ailego_unlikely(bl == nullptr)) { + LOG_ERROR("Failed to get block loc"); + return IndexError_ReadData; + } + block = *bl; + } + // Release the head block reference early so that the buffer pool ref_count + // and memory budget held by it do not block subsequent acquire/evict in this + // function (alloc_block / add_to_block may compete for the same memory). + head_block.reset(nullptr); + + if (!this->is_valid_block(block)) { + int ret = this->alloc_block(block, &block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + ret = this->update_head_block(block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + } + + int ret = this->add_to_block(block, key, vec, size); + if (ret == IndexError_IndexFull) { + ret = this->alloc_block(block, &block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + ret = this->update_head_block(block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + ret = this->add_to_block(block, key, vec, size); + if (ailego_unlikely(ret != 0)) { + return ret; + } + } + if (ailego_unlikely(ret != 0)) { + return ret; + } + + (*stats_.mutable_added_count())++; + stats_.set_revision_id(meta_.revision_id + 1); + + return 0; +} + +int FlatStreamerEntity::search(const void *query, const IndexFilter &filter, + uint32_t *scan_count, IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const { + IndexStorage::MemoryBlock head_block; + this->get_head_block(head_block); + const BlockLocation *bl = + reinterpret_cast(head_block.data()); + if (ailego_unlikely(bl == nullptr)) { + LOG_ERROR("Failed to get block loc"); + return IndexError_ReadData; + } + + BlockLocation block = *bl; + + while (this->is_valid_block(block)) { + IndexStorage::MemoryBlock block_header_block; + this->get_block_header(block, block_header_block); + const BlockHeader *hd = + reinterpret_cast(block_header_block.data()); + if (ailego_unlikely(hd == nullptr)) { + LOG_ERROR("Failed to get block header"); + return IndexError_ReadData; + } + + if (hd->vector_count > 0) { + *scan_count += hd->vector_count; + IndexStorage::MemoryBlock deletion_map_block; + this->get_block_deletion_map(block, deletion_map_block); + const DeletionMap *deletion_map = + reinterpret_cast(deletion_map_block.data()); + if (filter.is_valid() || deletion_map->is_dirty()) { + this->search_block(query, block, hd, 1.0, filter, deletion_map, heap, + context_stats); + } else { + *(context_stats->mutable_dist_calced_count()) += hd->vector_count; + this->search_block(query, block, hd, 1.0, heap); + } + } + block = hd->next; + } + return 0; +} + +//! Search in a block +void FlatStreamerEntity::search_block(const void *query, + const BlockLocation &bl, + const BlockHeader *hd, float norm_val, + IndexDocumentHeap *heap) const { + std::vector distances(block_vector_count()); + IndexStorage::MemoryBlock vecs_block; + this->get_block_vectors(bl, vecs_block); + const char *vecs = reinterpret_cast(vecs_block.data()); + IndexStorage::MemoryBlock keys_block; + this->get_block_keys(bl, keys_block); + const uint64_t *keys = reinterpret_cast(keys_block.data()); + row_major_distance(query, vecs, hd->vector_count, distances.data()); + for (size_t k = 0; k < hd->vector_count; ++k) { + if (keys[k] != kInvalidKey) { + heap->emplace(keys[k], distances[k] * norm_val); + } + } +} + +//! Search in a block with filter +void FlatStreamerEntity::search_block( + const void *query, const BlockLocation &bl, const BlockHeader *hd, + float norm_val, const IndexFilter &filter, const DeletionMap *deletion_map, + IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const { + std::vector distances(block_vector_count()); + + IndexStorage::MemoryBlock vecs_block; + this->get_block_vectors(bl, vecs_block); + const char *vecs = reinterpret_cast(vecs_block.data()); + IndexStorage::MemoryBlock keys_block; + this->get_block_keys(bl, keys_block); + const uint64_t *keys = reinterpret_cast(keys_block.data()); + + DeletionMap keeps; + for (size_t k = 0; k < hd->vector_count; ++k) { + const bool condition1 = !deletion_map->test(k); + const bool condition2 = filter.is_valid() ? !filter(keys[k]) : true; + const bool condition3 = keys[k] != kInvalidKey; + if (condition1 && condition2 && condition3) { + keeps.set(k); + } + } + if (!keeps.is_dirty()) { + (*context_stats->mutable_filtered_count()) += hd->vector_count; + return; + } + for (size_t k = 0; k < hd->vector_count; ++k) { + if (keeps.test(k)) { + auto cur_vec = vecs + index_meta_.element_size() * k; + row_major_distance(query, cur_vec, 1, distances.data() + k); + ++(*context_stats->mutable_dist_calced_count()); + } + } + for (size_t k = 0; k < hd->vector_count; ++k) { + if (keeps.test(k)) { + heap->emplace(keys[k], distances[k] * norm_val); + } else { + ++(*context_stats->mutable_filtered_count()); + } + } +} + +int FlatStreamerEntity::search_bf(const void *query, const IndexFilter &filter, + IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const { + uint32_t scan_count; + return this->search(query, filter, &scan_count, heap, context_stats); +} + +FlatStreamerEntity::Pointer FlatStreamerEntity::clone(void) const { + std::vector segments; + segments.reserve(segments_.size()); + for (size_t i = 0; i < segments_.size(); ++i) { + segments.emplace_back(segments_[i]->clone()); + if (!segments[i]) { + LOG_ERROR("Failed to clone segment, index=%zu", i); + return nullptr; + } + } + auto entity = new (std::nothrow) FlatStreamerEntity(stats_); + if (!entity) { + LOG_ERROR("Failed to New FlatStreamerEntity object"); + return nullptr; + } + entity->index_meta_ = this->index_meta_; + entity->storage_ = this->storage_; + // entity->reformer_ = this->reformer_; + entity->segments_ = segments; + entity->meta_ = this->meta_; + entity->key_info_map_lock_ = this->key_info_map_lock_; + entity->key_info_map_ = this->key_info_map_; + entity->id_key_vector_ = this->id_key_vector_; + entity->withid_key_info_map_ = this->withid_key_info_map_; + entity->withid_key_map_ = this->withid_key_map_; + entity->filter_same_key_ = this->filter_same_key_; + entity->vec_unit_size_ = this->vec_unit_size_; + entity->vec_cols_ = this->vec_cols_; + return FlatStreamerEntity::Pointer(entity); +} + +const void *FlatStreamerEntity::get_vector_by_key(uint64_t key) const { + VectorLocation loc{}; + key_info_map_lock_->lock_shared(); + if (use_key_info_map_) { + auto iterator = key_info_map_.find(key); + if (iterator == key_info_map_.end()) { + key_info_map_lock_->unlock_shared(); + return nullptr; + } + loc = iterator->second; + } else { + if (key < withid_key_info_map_.size()) { + loc = withid_key_info_map_[key]; + } else { + key_info_map_lock_->unlock_shared(); + return nullptr; + } + } + key_info_map_lock_->unlock_shared(); + + auto segment = this->get_segment(loc.segment_id); + const void *data = nullptr; + if (segment->read(loc.offset, &data, index_meta_.element_size()) != + index_meta_.element_size()) { + LOG_ERROR("Failed to read segment, size=%u", index_meta_.element_size()); + return nullptr; + } + return data; +} + +int FlatStreamerEntity::get_vector_by_key( + const uint64_t key, IndexStorage::MemoryBlock &block) const { + VectorLocation loc{}; + key_info_map_lock_->lock_shared(); + if (use_key_info_map_) { + auto iterator = key_info_map_.find(key); + if (iterator == key_info_map_.end()) { + key_info_map_lock_->unlock_shared(); + return -1; + } + loc = iterator->second; + } else { + if (key < withid_key_info_map_.size()) { + loc = withid_key_info_map_[key]; + } else { + key_info_map_lock_->unlock_shared(); + return -1; + } + } + key_info_map_lock_->unlock_shared(); + + auto segment = this->get_segment(loc.segment_id); + if (segment->read(loc.offset, block, index_meta_.element_size()) != + index_meta_.element_size()) { + LOG_ERROR("Failed to read segment, size=%u", index_meta_.element_size()); + return -1; + } + return 0; +} + +IndexProvider::Iterator::Pointer FlatStreamerEntity::creater_iterator( + void) const { + auto entity = this->clone(); + if (!entity) { + LOG_ERROR("Failed to clone entity"); + return nullptr; + } + + return Iterator::Pointer(new (std::nothrow) + FlatStreamerEntity::Iterator(std::move(entity))); +} + +void FlatStreamerEntity::Iterator::read_next_block(void) { + auto block_size = entity_->linear_block_size(); + while (segment_id_ < entity_->segments_.size()) { + auto &segment = entity_->segments_[segment_id_]; + size_t off = block_index_ * block_size; + if (off + block_size > segment->data_size()) { + ++segment_id_; + block_index_ = 0; + continue; + } + if (segment->read(off, block_, block_size) != block_size) { + LOG_ERROR("Failed to read block, off=%zu", off); + break; + } + data_ = block_.data(); + auto hd = reinterpret_cast( + static_cast(data_) + block_size - sizeof(BlockHeader)); + if (hd->vector_count == 0) { + ++block_index_; + continue; + } + + block_vector_count_ = hd->vector_count; + block_vector_index_ = 0; + size_t elemsize = entity_->index_meta_.element_size(); + keys_ = reinterpret_cast( + reinterpret_cast(data_) + + elemsize * entity_->block_vector_count()); + return; + } + + is_valid_ = false; +} + +int FlatStreamerEntity::init_storage(IndexStorage::Pointer storage) { + // Init Linear Meta Segment + meta_.create_time = ailego::Realtime::Seconds(); + stats_.set_create_time(meta_.create_time); + meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(meta_.update_time); + meta_.segment_count = 0; + meta_.revision_id = 0; + + std::string str; + index_meta_.serialize(&str); + const size_t page = ailego::MemoryHelper::PageSize(); + + meta_.header.header_size = sizeof(LinearIndexHeader) + str.size(); + meta_.header.total_vector_count = 0; + meta_.header.linear_body_size = 0; + meta_.header.block_count = 0; + meta_.header.index_meta_size = str.size(); + meta_.header.linear_list_count = 1; + + AdjustSegmentSize(&meta_); + + LOG_DEBUG( + "Create Streamer Index, VecSize=%u, BlockSize=%u SegmentSize=%u " + "LinearListCount=%u", + index_meta_.element_size(), meta_.header.block_size, meta_.segment_size, + meta_.header.linear_list_count); + + size_t size = ailego_align(sizeof(meta_) + str.size(), page); + int ret = storage->append(FLAT_LINEAR_META_SEG_ID, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to append segment %s", FLAT_LINEAR_META_SEG_ID.c_str()); + return ret; + } + auto segment = storage->get(FLAT_LINEAR_META_SEG_ID); + if (ailego_unlikely(!segment)) { + LOG_ERROR("Failed to get segment %s", FLAT_LINEAR_META_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) { + LOG_ERROR("Failed to write segment data"); + return IndexError_WriteData; + } + if (segment->write(sizeof(meta_), str.data(), str.size()) != str.size()) { + LOG_ERROR("Failed to write segment data, size=%zu", str.size()); + return IndexError_WriteData; + } + + ret = storage->append("IndexMeta", str.size()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to append segment IndexMeta, code: %d", ret); + return ret; + } + auto index_meta_segment = storage->get("IndexMeta"); + if (index_meta_segment->write(0, str.data(), str.size()) != str.size()) { + LOG_ERROR("Failed to write segment data, size=%zu", str.size()); + return IndexError_WriteData; + } + *stats_.mutable_index_size() += size; + + // Init Linear List Head Segment + size = ailego_align(sizeof(BlockLocation) * linear_list_count(), page); + ret = storage->append(FLAT_LINEAR_LIST_HEAD_SEG_ID, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to append segment %s for %s, size=%zu", + FLAT_LINEAR_LIST_HEAD_SEG_ID.c_str(), IndexError::What(ret), + size); + return ret; + } + segment = storage->get(FLAT_LINEAR_LIST_HEAD_SEG_ID); + if (ailego_unlikely(!segment)) { + LOG_ERROR("Failed to get segment %s", FLAT_LINEAR_LIST_HEAD_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (segment->resize(size) != size) { + LOG_ERROR("Failed to resize segment, size=%zu", size); + return IndexError_WriteData; + } + segments_.emplace_back(std::move(segment)); + + *stats_.mutable_index_size() += size; + + return 0; +} + +int FlatStreamerEntity::load_linear_meta(IndexStorage::Pointer storage) { + AdjustSegmentSize(&meta_); + + // Load Meta Segment + auto segment = storage->get(FLAT_LINEAR_META_SEG_ID); + if (!segment || segment->data_size() < sizeof(meta_)) { + LOG_ERROR("Missing segment %s, or invalid segment size", + FLAT_LINEAR_META_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + IndexStorage::MemoryBlock data_block; + if (segment->read(0, data_block, segment->data_size()) != + segment->data_size()) { + LOG_ERROR("Failed to read storage, size=%zu", segment->data_size()); + return IndexError_InvalidFormat; + } + auto *mt = reinterpret_cast(data_block.data()); + if (mt->header.block_vector_count != meta_.header.block_vector_count) { + LOG_ERROR("Unmatched BlockVecCount Setting, Index %u vs Setting %u", + mt->header.block_vector_count, meta_.header.block_vector_count); + return IndexError_Mismatch; + } + if (mt->header.block_size != meta_.header.block_size) { + LOG_ERROR("Unmatched BlockSize Setting, Index %u vs Setting %u", + mt->header.block_size, meta_.header.block_size); + return IndexError_Mismatch; + } + if (mt->header.index_meta_size + sizeof(meta_) > segment->data_size()) { + LOG_ERROR("Invalid format, IndexMetaSize %u, SegmentSize %zu", + mt->header.index_meta_size, segment->data_size()); + return IndexError_InvalidFormat; + } + if (mt->header.linear_list_count != meta_.header.linear_list_count) { + LOG_ERROR("Unmatch LinearListCount, Index size %u vs Setting %u", + mt->header.linear_list_count, meta_.header.linear_list_count); + return IndexError_InvalidFormat; + } + IndexMeta index_meta; + if (!index_meta.deserialize(mt->header.index_meta, + mt->header.index_meta_size)) { + LOG_ERROR("Failed to deserialize IndexMeta, size=%u", + mt->header.index_meta_size); + return IndexError_InvalidFormat; + } + if (index_meta.data_type() != index_meta_.data_type() || + index_meta.dimension() != index_meta_.dimension() || + index_meta.element_size() != index_meta_.element_size() || + index_meta.metric_name() != index_meta_.metric_name()) { + LOG_ERROR( + "Unmatch IndexMeta, Index(type=%u dim=%u elemsize=%u " + "metric=%s) Setting(type=%u dim=%u elemsize=%u metric=%s)", + index_meta.data_type(), index_meta.dimension(), + index_meta.element_size(), index_meta.metric_name().c_str(), + index_meta_.data_type(), index_meta_.dimension(), + index_meta_.element_size(), index_meta_.metric_name().c_str()); + return IndexError_Mismatch; + } + // Segment Size can be reconfigurable + auto segment_size = meta_.segment_size; + std::memcpy(&meta_, mt, sizeof(meta_)); + meta_.segment_size = segment_size; + return 0; +} + +int FlatStreamerEntity::load_segment_keys_to_map(BlockLocation block) { + while (this->is_valid_block(block)) { + auto segment = this->get_segment(block.segment_id); + + IndexStorage::MemoryBlock block_header_block; + this->get_block_header(block, block_header_block); + const BlockHeader *hd = + reinterpret_cast(block_header_block.data()); + if (ailego_unlikely(hd == nullptr)) { + LOG_ERROR("Failed to get block header"); + return IndexError_ReadData; + } + IndexStorage::MemoryBlock keys_block; + this->get_block_keys(block, keys_block); + const uint64_t *keys = + reinterpret_cast(keys_block.data()); + IndexStorage::MemoryBlock deletion_map_block; + this->get_block_deletion_map(block, deletion_map_block); + const DeletionMap *deletion_map = + reinterpret_cast(deletion_map_block.data()); + + for (uint32_t vector_index = 0; vector_index < hd->vector_count; + ++vector_index) { + if (deletion_map->test(vector_index)) { + continue; + } + size_t vector_off = + this->get_block_vector_offset(block.block_index, vector_index); + key_info_map_[keys[vector_index]] = + VectorLocation(block.segment_id, false, vector_off); + id_key_vector_.push_back(keys[vector_index]); + } + block = hd->next; + } + return 0; +} + +int FlatStreamerEntity::load_segment_keys_to_vector() { + for (uint32_t i = 0; i < meta_.header.total_vector_count; i++) { + size_t block_id = i / block_vector_count(); + uint32_t vector_index = i % block_vector_count(); + + ailego_assert(segments_.size() > 1); + size_t segment_block_count = + segments_[1]->data_size() / linear_block_size(); + size_t segment_id = block_id / segment_block_count + 1; + size_t real_block_id = block_id % segment_block_count; + size_t vector_off = + this->get_block_vector_offset(real_block_id, vector_index); + + withid_key_info_map_.push_back( + VectorLocation(segment_id, false, vector_off)); + size_t key_off = get_block_key_offset(real_block_id, vector_index); + withid_key_map_.push_back(key_off); + } + return 0; +} + +int FlatStreamerEntity::load_storage(IndexStorage::Pointer storage) { + int ret = this->load_linear_meta(storage); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + // Load Linear List + auto hd_segment = storage->get(FLAT_LINEAR_LIST_HEAD_SEG_ID); + if (ailego_unlikely(!hd_segment)) { + LOG_ERROR("Failed to get segment %s", FLAT_LINEAR_LIST_HEAD_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (hd_segment->data_size() < linear_list_count() * sizeof(BlockLocation)) { + LOG_ERROR("Invalid segment size, LinearListCount=%zu, size=%zu", + linear_list_count(), hd_segment->data_size()); + return IndexError_InvalidFormat; + } + segments_.emplace_back(hd_segment); + + size_t index_size = hd_segment->capacity(); + for (size_t i = 1; i <= meta_.segment_count; ++i) { + std::string segment_id = + ailego::StringHelper::Concat(FLAT_SEGMENT_FEATURES_SEG_ID, i); + auto seg = storage->get(segment_id); + if (!seg || seg->data_size() < meta_.header.block_size) { + LOG_ERROR("Failed to get segment %s, or invalid segment size", + segment_id.c_str()); + return IndexError_InvalidFormat; + } + index_size += seg->capacity(); + segments_.emplace_back(std::move(seg)); + } + + for (size_t i = 0; i < linear_list_count(); i++) { + IndexStorage::MemoryBlock head_block; + this->get_head_block(head_block); + const BlockLocation *bl = + reinterpret_cast(head_block.data()); + if (ailego_unlikely(bl == nullptr)) { + LOG_ERROR("Failed to get block loc"); + return IndexError_ReadData; + } + BlockLocation block = *bl; + if (use_key_info_map_) { + ret = this->load_segment_keys_to_map(block); + } else { + ret = this->load_segment_keys_to_vector(); + } + if (ailego_unlikely(ret != 0)) { + return ret; + } + } + + char create_time[32]; + char update_time[32]; + ailego::Realtime::Gmtime(meta_.create_time, "%Y-%m-%d %H:%M:%S", create_time, + sizeof(create_time)); + ailego::Realtime::Gmtime(meta_.update_time, "%Y-%m-%d %H:%M:%S", update_time, + sizeof(update_time)); + LOG_DEBUG( + "Load Index, IndexSize=%zu SegmentCount=%u SegmentSize=%u " + "RevisionId=%zu BlockCount=%u BlockSize=%u " + "BlockVectorCount=%u LinearListCount=%u TotalVecCount=%zu " + "CreateTime=%s UpdateTime=%s", + index_size, meta_.segment_count, meta_.segment_size, + static_cast(meta_.revision_id), meta_.header.block_count, + meta_.header.block_size, meta_.header.block_vector_count, + meta_.header.linear_list_count, + static_cast(meta_.header.total_vector_count), create_time, + update_time); + + stats_.set_index_size(index_size); + stats_.set_check_point(storage->check_point()); + stats_.set_create_time(meta_.create_time); + stats_.set_revision_id(meta_.revision_id); + stats_.set_update_time(meta_.update_time); + stats_.set_loaded_count(meta_.header.total_vector_count); + + return 0; +} + +int FlatStreamerEntity::alloc_segment(void) { + size_t index = segments_.size(); + if (index == kMaxSegmentId) { + LOG_ERROR("Failed to alloc new segment, exceed max count %zu", + kMaxSegmentId); + return IndexError_IndexFull; + } + + std::string segment_id = + ailego::StringHelper::Concat(FLAT_SEGMENT_FEATURES_SEG_ID, index); + size_t size = + ailego_align(meta_.segment_size, ailego::MemoryHelper::PageSize()); + auto segment = storage_->get(segment_id); + if (segment) { + if (segment->padding_size() < linear_block_size()) { + LOG_ERROR( + "Unexpect segment, index=%zu, data_size=%zu " + "padding_size=%zu block_size=%zu", + index, segment->data_size(), segment->padding_size(), + linear_block_size()); + return IndexError_Runtime; + } + LOG_WARN("Alloc an existing segment=%s capacity=%zu", segment_id.c_str(), + segment->capacity()); + } else { + int ret = storage_->append(segment_id, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to alloc segment from storage"); + return ret; + } + segment = storage_->get(segment_id); + if (ailego_unlikely(!segment)) { + LOG_ERROR("Failed to get segment %s", segment_id.c_str()); + return IndexError_Runtime; + } + } + meta_.segment_count += 1; + meta_.header.linear_body_size += size; + segments_.emplace_back(std::move(segment)); + *stats_.mutable_index_size() += size; + + // Update meta information + auto meta_segment = storage_->get(FLAT_LINEAR_META_SEG_ID); + if (ailego_unlikely(!meta_segment)) { + LOG_ERROR("Failed to get segment %s", FLAT_LINEAR_META_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (meta_segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) { + LOG_ERROR("Failed to write meta segment"); + return IndexError_WriteData; + } + + return 0; +} + +int FlatStreamerEntity::alloc_block(const BlockLocation &next, + BlockLocation *block) { + if (segments_.size() <= 1 || + segments_.back()->padding_size() < linear_block_size()) { + int ret = this->alloc_segment(); + if (ailego_unlikely(ret != 0)) { + return ret; + } + } + + auto &segment = segments_.back(); + size_t block_index = segment->data_size() / linear_block_size(); + if (block_index == kMaxBlockId) { + LOG_ERROR("Failed to alloc block, exceed max count %zu per segment", + kMaxBlockId); + return IndexError_IndexFull; + } + + BlockHeader header; + header.next = next; + header.vector_count = 0; + header.column_major = false; + + size_t hd_off = segment->data_size() + linear_block_size() - sizeof(header); + if (segment->write(hd_off, &header, sizeof(header)) != sizeof(header)) { + LOG_ERROR("Failed to write block header"); + return IndexError_WriteData; + } + + size_t del_off = hd_off - sizeof(DeletionMap); + DeletionMap reset_del_map{}; + if (segment->write(del_off, &reset_del_map, sizeof(reset_del_map)) != + sizeof(reset_del_map)) { + LOG_ERROR("Failed to write block deletion map"); + return IndexError_WriteData; + } + + ++meta_.header.block_count; + block->segment_id = segments_.size() - 1; + block->block_index = (segment->data_size() / linear_block_size()) - 1; + + return 0; +} + +int FlatStreamerEntity::add_to_block(const BlockLocation &block, uint64_t key, + const void *data, size_t size) { + IndexStorage::MemoryBlock block_header_block; + this->get_block_header(block, block_header_block); + const BlockHeader *header = + reinterpret_cast(block_header_block.data()); + if (ailego_unlikely(header == nullptr)) { + LOG_ERROR("Failed to get header"); + return IndexError_ReadData; + } + + if (header->vector_count == block_vector_count()) { + return IndexError_IndexFull; + } + + auto &segment = segments_[block.segment_id]; + + size_t vector_off = + get_block_vector_offset(block.block_index, header->vector_count); + if (segment->write(vector_off, data, size) != size) { + LOG_ERROR("Failed to write vector, off=%zu size=%zu", vector_off, size); + return IndexError_WriteData; + } + + size_t key_off = + get_block_key_offset(block.block_index, header->vector_count); + if (segment->write(key_off, &key, sizeof(key)) != sizeof(key)) { + LOG_ERROR("Failed to write key, off=%zu", key_off); + return IndexError_WriteData; + } + + BlockHeader hd = *header; + hd.vector_count += 1; + size_t hd_off = get_block_header_offset(block.block_index); + if (segment->write(hd_off, &hd, sizeof(hd)) != sizeof(hd)) { + LOG_ERROR("Failed to write block header, off=%zu", hd_off); + return IndexError_WriteData; + } + + VectorLocation loc(block.segment_id, false, vector_off); + key_info_map_lock_->lock(); + key_info_map_[key] = loc; + id_key_vector_.push_back(key); + withid_key_info_map_.push_back(loc); + withid_key_map_.push_back(key_off); + key_info_map_lock_->unlock(); + + ++meta_.header.total_vector_count; + return 0; +} + +int FlatStreamerEntity::add_vector_with_id(const uint32_t id, const void *query, + const uint32_t size) { + std::lock_guard lock(mutex_); + // if (filter_same_key_) { + // key_info_map_lock_->lock_shared(); + // if (key_info_map_.find(id) != key_info_map_.end()) { + // key_info_map_lock_->unlock_shared(); + // LOG_WARN("Try to add duplicate key, drop it"); + // return IndexError_Duplicate; + // } + // key_info_map_lock_->unlock_shared(); + // } + + if (size != static_cast(index_meta_.element_size())) { + LOG_ERROR("Failed to add, mismatch size %u vs elemsize %u", size, + index_meta_.element_size()); + return IndexError_Mismatch; + } + + + if (id >= vector_count()) { + IndexStorage::MemoryBlock head_block; + this->get_head_block(head_block); + BlockLocation block = + *reinterpret_cast(head_block.data()); + // Release buffer-pool pin before any alloc_block() call that may trigger + // append_segment() and rebuild the pool (same reason as in add()). + head_block.reset(nullptr); + if (!this->is_valid_block(block)) { + int ret = this->alloc_block(block, &block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + ret = this->update_head_block(block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + } + for (size_t start_id = vector_count(); start_id < id; ++start_id) { + std::vector vec(size); + int ret = this->add_to_block(block, kInvalidKey, vec.data(), size); + if (ret == IndexError_IndexFull) { + ret = this->alloc_block(block, &block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + ret = this->update_head_block(block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + ret = this->add_to_block(block, kInvalidKey, vec.data(), size); + if (ailego_unlikely(ret != 0)) { + return ret; + } + } + } + + int ret = this->add_to_block(block, id, query, size); + if (ret == IndexError_IndexFull) { + ret = this->alloc_block(block, &block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + ret = this->update_head_block(block); + if (ailego_unlikely(ret != 0)) { + return ret; + } + ret = this->add_to_block(block, id, query, size); + if (ailego_unlikely(ret != 0)) { + return ret; + } + } + } else { + VectorLocation vector_loc = withid_key_info_map_[id]; + auto segment = this->get_segment(vector_loc.segment_id); + size_t vector_off = vector_loc.offset; + if (segment->write(vector_off, query, size) != size) { + LOG_ERROR("Failed to write vector, off=%zu size=%u", vector_off, size); + return IndexError_WriteData; + } + size_t key_off = withid_key_map_[id]; + uint64_t key = id; + if (segment->write(key_off, &key, sizeof(key)) != sizeof(key)) { + LOG_ERROR("Failed to write key, off=%zu", key_off); + return IndexError_WriteData; + } + key_info_map_lock_->lock(); + key_info_map_[key] = vector_loc; + key_info_map_lock_->unlock(); + } + (*stats_.mutable_added_count())++; + stats_.set_revision_id(meta_.revision_id + 1); + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_streamer_entity.h b/src/core/algorithm/flat/flat_streamer_entity.h new file mode 100644 index 0000000..7cf9154 --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer_entity.h @@ -0,0 +1,407 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include "flat_index_format.h" +#include "flat_utility.h" + +namespace zvec { +namespace core { + +/*! Flat Streamer Entity + */ +class FlatStreamerEntity { + public: + typedef std::shared_ptr Pointer; + + //! Constructor + explicit FlatStreamerEntity(IndexStreamer::Stats &stats); + + //! Destructor + virtual ~FlatStreamerEntity(void) = default; + + //! Open the entity with storage + int open(IndexStorage::Pointer storage, const IndexMeta &mt); + + //! Close the entity + int close(void); + + //! Flush Linear Meta information to storage + int flush_linear_meta(void); + + //! Flush linear index to storage + int flush(uint64_t checkpoint); + + //! Add vector to linear index + int add(uint64_t key, const void *vec, size_t size); + + //! Search in linear list with filter + int search(const void *query, const IndexFilter &filter, uint32_t *scan_count, + IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const; + + //! Search in a block + void search_block(const void *query, const BlockLocation &bl, + const BlockHeader *hd, float norm_val, + IndexDocumentHeap *heap) const; + + //! Search in a block with filter + void search_block(const void *query, const BlockLocation &bl, + const BlockHeader *hd, float norm_val, + const IndexFilter &filter, const DeletionMap *deletion_map, + IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const; + + //! Flat Search with filter + int search_bf(const void *query, const IndexFilter &filter, + IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const; + + //! Clone the entity + virtual FlatStreamerEntity::Pointer clone(void) const; + + //! Retrieve the total vectors in the index + size_t vector_count(void) const { + return meta_.header.total_vector_count; + } + + //! Retrieve the linear list count + size_t linear_list_count(void) const { + return meta_.header.linear_list_count; + } + + //! Retrieve block size of the linear vector + size_t linear_block_size(void) const { + return meta_.header.block_size; + } + + //! Retrieve the vectors count in one block + size_t block_vector_count(void) const { + // assert(meta_.header.block_vector_count == 32); + return meta_.header.block_vector_count; + } + + //! Retrieve IndexMeta of the linear index + const IndexMeta &meta(void) const { + return index_meta_; + } + + //! Retrieve mutable IndexMeta of the linear index + IndexMeta *mutable_meta(void) { + return &index_meta_; + } + + //! Retrieve vector by local id + const void *get_vector_by_key(uint64_t key) const; + + int get_vector_by_key(const uint64_t key, + IndexStorage::MemoryBlock &block) const; + + //! Create a new iterator + IndexProvider::Iterator::Pointer creater_iterator(void) const; + + + //! Set params + void set_block_vector_count(uint32_t count) { + meta_.header.block_vector_count = count; + } + + void set_use_key_info_map(bool use_id_map) { + use_key_info_map_ = use_id_map; + LOG_DEBUG("use_key_info_map_: %d", (int)use_key_info_map_); + } + + //! Set params + void set_segment_size(uint32_t size) { + meta_.segment_size = size; + } + + //! Set params + void set_linear_list_count(uint32_t count) { + meta_.header.linear_list_count = count; + } + + //! Set params + void enable_filter_same_key(bool enabled) { + filter_same_key_ = enabled; + } + + inline uint64_t key(uint32_t id) const { + if (id < id_key_vector_.size()) { + return id_key_vector_[id]; + } else { + return kInvalidKey; + } + } + + inline void row_major_distance(const void *query, const void *feature, + size_t fnum, float *out) const { + const uint8_t *cur_feature = reinterpret_cast(feature); + for (size_t f = 0; f < fnum; ++f) { + row_distance_(query, cur_feature, index_meta_.dimension(), out + f); + cur_feature += index_meta_.element_size(); + } + } + + int add_vector_with_id(const uint32_t id, const void *query, + const uint32_t element_size); + + private: + //! Disable them + FlatStreamerEntity(const FlatStreamerEntity &) = delete; + FlatStreamerEntity &operator=(const FlatStreamerEntity &) = delete; + + /*! Iterator of all the linear list + */ + class Iterator : public IndexProvider::Iterator { + public: + //! Constructor + Iterator(const FlatStreamerEntity::Pointer &entity) : entity_(entity) { + this->read_next_block(); + } + //! Retrieve pointer of data + const void *data(void) const override { + return reinterpret_cast(data_) + + block_vector_index_ * entity_->index_meta_.element_size(); + } + //! Test if the iterator is valid + bool is_valid(void) const override { + return is_valid_; + } + //! Retrieve primary key + uint64_t key(void) const override { + return keys_[block_vector_index_]; + } + //! Next iterator + void next(void) override { + if (++block_vector_index_ == block_vector_count_) { + ++block_index_; + this->read_next_block(); + } + } + + private: + //! Read next non-empty block + void read_next_block(void); + + //! Members + std::string buffer_{}; + const FlatStreamerEntity::Pointer entity_; + IndexStorage::MemoryBlock block_; + const void *data_{nullptr}; + const uint64_t *keys_{nullptr}; + uint32_t segment_id_{1u}; // The first segment is header info + uint32_t block_index_{0u}; + uint32_t block_vector_index_{0u}; + uint32_t block_vector_count_{0u}; + bool is_valid_{true}; + }; + + //! Retrive storage segment by index + const IndexStorage::Segment::Pointer get_segment(size_t index) const { + for (size_t i = segments_.size(); i <= index; ++i) { + auto segment_id = + ailego::StringHelper::Concat(FLAT_SEGMENT_FEATURES_SEG_ID, i); + auto segment = storage_->get(segment_id); + if (!segment) { + LOG_ERROR("Failed to get segment %s", segment_id.c_str()); + return IndexStorage::Segment::Pointer(); + } + segments_.emplace_back(std::move(segment)); + } + return segments_[index]; + } + + //! Rejust the segment size as to aligned by page size + void AdjustSegmentSize(StreamerLinearMeta *mt) { + if (mt->segment_size < mt->header.block_size) { + mt->segment_size = mt->header.block_size; + } + mt->segment_size = ailego_align( + mt->segment_size / mt->header.block_size * mt->header.block_size, + ailego::MemoryHelper::PageSize()); + } + + //! Init with an empty storage + int init_storage(IndexStorage::Pointer storage); + + //! Load linear meta information from storage + int load_linear_meta(IndexStorage::Pointer storage); + + //! Load keys to keys map + int load_segment_keys_to_map(BlockLocation block); + + //! Load keys to keys map + int load_segment_keys_to_vector(void); + + //! Load index from storage + int load_storage(IndexStorage::Pointer storage); + + //! Check whether the block is empty + bool is_valid_block(const BlockLocation &block) const { + return block.segment_id != 0; + } + + //! Update header block of an linear list + int update_head_block(const BlockLocation &block) { + ailego_assert_with(segments_.size() != 0, "Invalid Segments"); + + auto &hd_segment = segments_[0]; + if (hd_segment->write(0, &block, sizeof(block)) != sizeof(block)) { + LOG_ERROR("Failed to write head block location"); + return IndexError_WriteData; + } + + return 0; + } + + //! Alloc a new segment + int alloc_segment(void); + + //! Alloc a new block + int alloc_block(const BlockLocation &next, BlockLocation *block); + + //! Add a record to a block + int add_to_block(const BlockLocation &block, uint64_t key, const void *data, + size_t size); + + private: + size_t get_block_offset(uint32_t block_index) const { + return block_index * linear_block_size(); + } + + size_t get_block_header_offset(uint32_t block_index) const { + return get_block_offset(block_index) + linear_block_size() - + sizeof(BlockHeader); + } + + size_t get_block_deletion_map_offset(uint32_t block_index) const { + return get_block_header_offset(block_index) - sizeof(DeletionMap); + } + + size_t get_block_key_offset(uint32_t block_index, + uint32_t vector_index) const { + return get_block_offset(block_index) + + block_vector_count() * index_meta_.element_size() + + sizeof(uint64_t) * vector_index; + } + + size_t get_block_vector_offset(uint32_t block_index, + uint32_t vector_index) const { + return this->get_block_offset(block_index) + + vector_index * index_meta_.element_size(); + } + + //! Get header block of an linear list + int get_head_block(IndexStorage::MemoryBlock &header_block) const { + ailego_assert_with(segments_.size() != 0, "Invalid Segments"); + auto &hd_segment = segments_[0]; + if (hd_segment->read(0, header_block, sizeof(BlockLocation)) != + sizeof(BlockLocation)) { + LOG_ERROR("Failed to read head block location"); + return -1; + } + return 0; + } + + //! Get BlockHeader of the block + int get_block_header(const BlockLocation &block, + IndexStorage::MemoryBlock &header_block) const { + // The header is located in the end of a block to align features + auto &segment = this->get_segment(block.segment_id); + ailego_assert_with(segment != nullptr, "Index Overflow"); + size_t off = this->get_block_header_offset(block.block_index); + if (segment->read(off, header_block, sizeof(BlockHeader)) != + sizeof(BlockHeader)) { + LOG_ERROR("Failed to read block header, off=%zu", off); + return -1; + } + return 0; + } + int get_block_deletion_map( + const BlockLocation &block, + IndexStorage::MemoryBlock &deletion_map_block) const { + auto &segment = this->get_segment(block.segment_id); + ailego_assert_with(segment != nullptr, "Index Overflow"); + size_t off = this->get_block_deletion_map_offset(block.block_index); + if (segment->read(off, deletion_map_block, sizeof(DeletionMap)) != + sizeof(DeletionMap)) { + LOG_ERROR("Failed to read deletion map, off=%zu", off); + return -1; + } + return 0; + } + + int get_block_keys(const BlockLocation &block, + IndexStorage::MemoryBlock &keys_block) const { + auto &segment = this->get_segment(block.segment_id); + ailego_assert_with(segment != nullptr, "Index Overflow"); + size_t off = this->get_block_key_offset(block.block_index, 0); + if (segment->read(off, keys_block, + block_vector_count() * sizeof(uint64_t)) != + block_vector_count() * sizeof(uint64_t)) { + LOG_ERROR("Failed to read block header, off=%zu", off); + return -1; + } + return 0; + } + + int get_block_vectors(const BlockLocation &block, + IndexStorage::MemoryBlock &vector_block) const { + auto &segment = this->get_segment(block.segment_id); + ailego_assert_with(segment != nullptr, "Index Overflow"); + size_t off = this->get_block_vector_offset(block.block_index, 0); + if (segment->read(off, vector_block, + block_vector_count() * index_meta_.element_size()) != + block_vector_count() * index_meta_.element_size()) { + LOG_ERROR("Failed to read block header, off=%zu", off); + return -1; + } + return 0; + } + + private: + //! Constants + static constexpr size_t kMaxSegmentId = std::numeric_limits::max(); + static constexpr size_t kMaxBlockId = std::numeric_limits::max(); + + //! Members + std::mutex mutex_{}; + IndexMeta index_meta_{}; + IndexStorage::Pointer storage_{}; + IndexMetric::MatrixDistance row_distance_{}, column_distance_{}; + mutable std::vector segments_{}; + IndexStreamer::Stats &stats_; + mutable std::shared_ptr key_info_map_lock_{}; + std::unordered_map key_info_map_{}; + std::vector withid_key_info_map_{}; + std::vector withid_key_map_{}; + std::vector id_key_vector_{}; + bool filter_same_key_{false}; + bool use_key_info_map_{true}; + uint32_t vec_unit_size_{0}; + uint32_t vec_cols_{0}; + mutable std::string vec_buf_{}; + StreamerLinearMeta meta_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_streamer_provider.h b/src/core/algorithm/flat/flat_streamer_provider.h new file mode 100644 index 0000000..70a22ef --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer_provider.h @@ -0,0 +1,100 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "flat_distance_matrix.h" +#include "flat_searcher.h" +#include "flat_streamer.h" +#include "flat_utility.h" + +namespace zvec { +namespace core { + +/*! Brute Force Streamer Provider + */ + +template +class FlatStreamerProvider : public IndexProvider { + public: + //! Constructor + FlatStreamerProvider(const FlatStreamer *owner) { + feature_size_ = owner->meta().element_size(); + total_vector_count_ = owner->entity().vector_count(); + owner_ = owner; + block_buffer_.resize(BATCH_SIZE * feature_size_); + } + + //! Create a new iterator + IndexProvider::Iterator::Pointer create_iterator(void) override { + return owner_->entity().creater_iterator(); + } + + //! Retrieve count of vectors + size_t count(void) const override { + return total_vector_count_; + } + + //! Retrieve dimension of vector + size_t dimension(void) const override { + return owner_->meta().dimension(); + } + + //! Retrieve type of vector + IndexMeta::DataType data_type(void) const override { + return owner_->meta().data_type(); + } + + //! Retrieve vector size in bytes + size_t element_size(void) const override { + return owner_->meta().element_size(); + } + + //! Retrieve a vector using a primary key + const void *get_vector(uint64_t key) const override { + return this->get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return this->get_vector_by_key(key, block); + } + + //! Retrieve the owner class + const std::string &owner_class(void) const override { + return owner_->name(); + } + + protected: + //! Retrieve a vector via primary key + const void *get_vector_by_key(uint64_t key) const { + return owner_->get_vector_by_key(key); + } + + int get_vector_by_key(const uint64_t key, + IndexStorage::MemoryBlock &block) const { + return owner_->get_vector_by_key(key, block); + } + + private: + //! Members + const FlatStreamer *owner_{nullptr}; + IndexStorage::Segment::Pointer features_segment_{}; + uint32_t feature_size_{0}; + uint32_t total_vector_count_{0}; + mutable std::vector block_buffer_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat/flat_utility.h b/src/core/algorithm/flat/flat_utility.h new file mode 100644 index 0000000..9912c58 --- /dev/null +++ b/src/core/algorithm/flat/flat_utility.h @@ -0,0 +1,190 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +//! The default size of reading a block +static constexpr uint32_t FLAT_DEFAULT_READ_BLOCK_SIZE = 4 * 1024 * 1024; +static const std::string FLAT_LINEAR_META_SEG_ID = "flat.linear_meta"; +static const std::string FLAT_LINEAR_LIST_HEAD_SEG_ID = "flat.linear_list_head"; + +static const std::string FLAT_SEGMENT_KEYS_SEG_ID("flat.keys"); +static const std::string FLAT_SEGMENT_FEATURES_SEG_ID("flat.features"); +static const std::string FLAT_SEGMENT_MAPPING_SEG_ID("flat.mapping"); + +// index params +static const std::string PARAM_FLAT_COLUMN_MAJOR_ORDER( + "proxima.flat.column_major_order"); +static const std::string PARAM_FLAT_BATCH_SIZE("proxima.flat.batch_size"); +static const std::string PARAM_FLAT_READ_BLOCK_SIZE( + "proxima.flat.read_block_size"); +static const std::string PARAM_FLAT_USE_ID_MAP("proxima.flat.use_id_map"); + +//! Determines if a number is equal to two to the power of n. +template +struct IsEqualPowerofTwo + : std::integral_constant {}; + +//! Transpose a block +template +static inline void ReverseTranspose(size_t align_size, const void *src, + size_t dim, void *dst) { + switch (align_size) { + case 2: + ailego::MatrixHelper::ReverseTranspose(src, dim, dst); + break; + case 4: + ailego::MatrixHelper::ReverseTranspose(src, dim, dst); + break; + case 8: + ailego::MatrixHelper::ReverseTranspose(src, dim, dst); + break; + } +} + +static inline void ReverseTranspose(size_t align_size, const void *src, + size_t m, size_t dim, void *dst) { + switch (align_size) { + case 2: + ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst); + break; + case 4: + ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst); + break; + case 8: + ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst); + break; + } +} + +template +static inline void TransposeOne(const void *src, size_t M, size_t N, + void *dst) { + for (size_t i = 0; i < N; ++i) { + reinterpret_cast(dst)[i] = reinterpret_cast(src)[i * M]; + } +} + +static inline void Transpose(size_t align_size, const void *src, size_t m, + size_t dim, void *dst) { + switch (align_size) { + case 2: + ailego::MatrixHelper::Transpose(src, m, dim, dst); + break; + case 4: + ailego::MatrixHelper::Transpose(src, m, dim, dst); + break; + case 8: + ailego::MatrixHelper::Transpose(src, m, dim, dst); + break; + } +} + +//! Transpose queries +template +void TransposeQueries(const void *query, const IndexQueryMeta &qmeta, + size_t query_count, std::string *out) { + if constexpr (K <= 1) { + ailego_assert(query_count == 1); + (void)query_count; + out->append(reinterpret_cast(query) + out->size(), + qmeta.element_size()); + } else { + ailego_assert_with(IsEqualPowerofTwo::value, + "K must be equal to two to the power of n."); + + size_t query_batch_count = query_count / K; + size_t query_offset = out->size(); + out->resize(query_offset + query_batch_count * K * qmeta.element_size()); + + switch (IndexMeta::AlignSizeof(qmeta.data_type())) { + case 2: + for (size_t i = 0; i != query_batch_count; ++i) { + ailego::MatrixHelper::Transpose( + (const char *)query + query_offset, + qmeta.element_size() / sizeof(uint16_t), &((*out)[query_offset])); + query_offset += qmeta.element_size() * K; + } + break; + + case 4: + for (size_t i = 0; i != query_batch_count; ++i) { + ailego::MatrixHelper::Transpose( + (const char *)query + query_offset, + qmeta.element_size() / sizeof(uint32_t), &((*out)[query_offset])); + + query_offset += qmeta.element_size() * K; + } + break; + + case 8: + for (size_t i = 0; i != query_batch_count; ++i) { + ailego::MatrixHelper::Transpose( + (const char *)query + query_offset, + qmeta.element_size() / sizeof(uint64_t), &((*out)[query_offset])); + query_offset += qmeta.element_size() * K; + } + break; + + default: + ailego_check_with(0, "BAD CASE"); + } + size_t query_left_count = query_count % K; + if (query_left_count != 0) { + TransposeQueries<(K >> 1)>(query, qmeta, query_left_count, out); + } + } +} + +//! Create and initialize measure +static inline int InitializeMetric(const IndexMeta &mt, + IndexMetric::Pointer *out) { + IndexMetric::Pointer measure = IndexFactory::CreateMetric(mt.metric_name()); + if (!measure) { + return IndexError_NoExist; + } + + int error_code = measure->init(mt, mt.metric_params()); + if (error_code != 0) { + return error_code; + } + *out = measure; + return 0; +} + +//! Verify measure +static inline bool VerifyMetric(const IndexMeta &meta) { + IndexMetric::Pointer measure = IndexFactory::CreateMetric(meta.metric_name()); + if (!measure) { + return false; + } + int error_code = measure->init(meta, meta.metric_params()); + if (error_code != 0) { + return false; + } + return true; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/CMakeLists.txt b/src/core/algorithm/flat_sparse/CMakeLists.txt new file mode 100644 index 0000000..4476613 --- /dev/null +++ b/src/core/algorithm/flat_sparse/CMakeLists.txt @@ -0,0 +1,20 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +# --exclude-libs is GNU ld / LLVM lld only; Apple ld does not support it. +# On macOS (Mach-O), symbol interposition works differently and the +# Arrow/Parquet double-free issue does not apply. +if(NOT APPLE) + set(CORE_KNN_FLAT_SPARSE_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_flat_sparse + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_FLAT_SPARSE_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/flat_sparse/flat_sparse_builder.cc b/src/core/algorithm/flat_sparse/flat_sparse_builder.cc new file mode 100644 index 0000000..34fcda5 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_builder.cc @@ -0,0 +1,338 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_sparse_builder.h" +#include +#include +#include +#include +#include +#include +#include +#include "flat_sparse_index_format.h" +#include "flat_sparse_utility.h" + +namespace zvec { +namespace core { + +FlatSparseBuilder::FlatSparseBuilder() {} + +int FlatSparseBuilder::init(const IndexMeta &meta, + const ailego::Params & /*params*/) { + LOG_INFO("Begin FlatSparseBuilder::init"); + + meta_ = meta; + + state_ = BUILD_STATE_INITED; + LOG_INFO("End FlatSparseBuilder::init"); + return 0; +} + +int FlatSparseBuilder::cleanup(void) { + LOG_INFO("Begin FlatSparseBuilder::cleanup"); + + stats_.clear_attributes(); + stats_.set_trained_count(0UL); + stats_.set_built_count(0UL); + stats_.set_dumped_count(0UL); + stats_.set_discarded_count(0UL); + stats_.set_trained_costtime(0UL); + stats_.set_built_costtime(0UL); + stats_.set_dumped_costtime(0UL); + state_ = BUILD_STATE_INIT; + + LOG_INFO("End FlatSparseBuilder::cleanup"); + + return 0; +} + +int FlatSparseBuilder::train(IndexThreads::Pointer, + IndexSparseHolder::Pointer /*holder*/) { + if (state_ != BUILD_STATE_INITED) { + LOG_ERROR("Init the builder before FlatSparseBuilder::train"); + return IndexError_NoReady; + } + + LOG_INFO("Begin FlatSparseBuilder::train"); + + stats_.set_trained_count(0UL); + stats_.set_trained_costtime(0UL); + state_ = BUILD_STATE_TRAINED; + + LOG_INFO("End FlatSparseBuilder::train"); + + return 0; +} + +int FlatSparseBuilder::train(const IndexTrainer::Pointer & /*trainer*/) { + if (state_ != BUILD_STATE_INITED) { + LOG_ERROR("Init the builder before FlatSparseBuilder::train"); + return IndexError_NoReady; + } + + LOG_INFO("Begin FlatSparseBuilder::train by trainer"); + + stats_.set_trained_count(0UL); + stats_.set_trained_costtime(0UL); + state_ = BUILD_STATE_TRAINED; + + LOG_INFO("End FlatSparseBuilder::train by trainer"); + + return 0; +} + +int FlatSparseBuilder::build(IndexThreads::Pointer, + IndexSparseHolder::Pointer holder) { + LOG_INFO("Begin FlatSparseBuilder::build"); + + ailego::ElapsedTime stamp; + if (!holder) { + LOG_ERROR("Input holder is nullptr while building index"); + return IndexError_InvalidArgument; + } + + if (!holder->is_matched(meta_)) { + LOG_ERROR("Input holder doesn't match index meta while building index"); + return IndexError_Mismatch; + } + + holder_ = std::move(holder); + + stats_.set_built_count(holder_->count()); + stats_.set_built_costtime(stamp.milli_seconds()); + state_ = BUILD_STATE_BUILT; + + LOG_INFO("End FlatSparseBuilder::build"); + return 0; +} + +int FlatSparseBuilder::dump(const IndexDumper::Pointer &dumper) { + if (state_ != BUILD_STATE_BUILT || !holder_) { + LOG_INFO("Build the index before FlatSparseBuilder::dump"); + return IndexError_NoReady; + } + + LOG_INFO("Begin FlatSparseBuilder::dump"); + + auto start_time = ailego::Monotime::MilliSeconds(); + + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + + uint32_t dump_count; + ret = do_dump(dumper, &dump_count); + if (ret != 0) { + LOG_ERROR("Failed to dump index"); + return ret; + } + + holder_ = nullptr; + stats_.set_dumped_count(dump_count); + stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time); + + LOG_INFO("End FlatSparseBuilder::dump"); + return 0; +} + +int FlatSparseBuilder::do_dump(const IndexDumper::Pointer &dumper, + uint32_t *dump_count) { + // bf meta + int ret = dump_meta(dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to dump meta"); + return ret; + } + + std::vector keys; + ret = dump_vector_and_offset(dumper.get(), &keys); + if (ret != 0) { + LOG_ERROR("Failed to dump offset data"); + return ret; + } + + ret = dump_keys(keys, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to dump keys"); + return ret; + } + + ret = dump_mapping(keys, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to dump mapping"); + return ret; + } + + *dump_count = keys.size(); + + return 0; +} + +int FlatSparseBuilder::dump_meta(IndexDumper *dumper) { + FlatSparseMeta meta; + meta.create_time = ailego::Realtime::Seconds(); + meta.update_time = ailego::Realtime::Seconds(); + meta.doc_cnt = holder_->count(); + + if (dumper->write(&meta, sizeof(meta)) != sizeof(meta)) { + LOG_ERROR("Failed to write meta"); + return IndexError_WriteData; + } + + size_t meta_padding_size = ailego_align(sizeof(meta), 32) - sizeof(meta); + if (meta_padding_size) { + std::string padding(meta_padding_size, '\0'); + if (dumper->write(padding.data(), meta_padding_size) != meta_padding_size) { + LOG_ERROR("Failed to write meta padding"); + return IndexError_WriteData; + } + } + return dumper->append(PARAM_FLAT_SPARSE_META_SEG_ID, sizeof(meta), + meta_padding_size, 0); +} + +int FlatSparseBuilder::dump_vector_and_offset(IndexDumper *dumper, + std::vector *keys) { + // iterate the holder + auto iter = holder_->create_iterator(); + if (!iter) { + LOG_ERROR("Failed to create iterator"); + return IndexError_Runtime; + } + + uint64_t written_length{0U}; + + std::vector> offset_lens; + while (iter->is_valid()) { + keys->push_back(iter->key()); + + uint32_t length; + if (write_vector_data(iter->sparse_count(), iter->sparse_indices(), + iter->sparse_data(), dumper, &length) != 0) { + return IndexError_WriteData; + } + + offset_lens.push_back({written_length, length}); + written_length += length; + iter->next(); + } + + if (dumper->append(PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID, written_length, 0, + 0) != 0) { + LOG_ERROR("Failed to append offset data"); + return IndexError_WriteData; + } + + LOG_DEBUG("Data total written: %zu", (size_t)written_length); + + for (auto &offset_len : offset_lens) { + if (dumper->write(&offset_len.first, sizeof(offset_len.first)) != + sizeof(offset_len.first)) { + LOG_ERROR("Failed to write offset"); + return IndexError_WriteData; + } + + if (dumper->write(&offset_len.second, sizeof(offset_len.second)) != + sizeof(offset_len.second)) { + LOG_ERROR("Failed to write length"); + return IndexError_WriteData; + } + } + + if (dumper->append(PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID, + offset_lens.size() * (sizeof(uint64_t) + sizeof(uint32_t)), + 0, 0) != 0) { + LOG_ERROR("Failed to append offset data"); + return IndexError_WriteData; + } + + LOG_DEBUG("Offset total written: %zu", + offset_lens.size() * (sizeof(uint64_t) + sizeof(uint32_t))); + + return 0; +} + +int FlatSparseBuilder::write_vector_data(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_vec, + IndexDumper *dumper, + uint32_t *length) { + std::string sparse_buffer; + + SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_vec, + meta_.unit_size(), sparse_buffer); + + if (dumper->write(sparse_buffer.data(), sparse_buffer.size()) != + sparse_buffer.size()) { + LOG_ERROR("Failed to write sparse data"); + return IndexError_WriteData; + } + + *length = sparse_buffer.size(); + + return 0; +} + +int FlatSparseBuilder::dump_keys(const std::vector &keys, + IndexDumper *dumper) { + size_t keys_size = keys.size() * sizeof(uint64_t); + if (dumper->write(keys.data(), keys_size) != keys_size) { + LOG_ERROR("Failed to write keys to dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + size_t keys_padding_size = ailego_align(keys_size, 32) - keys_size; + if (keys_padding_size) { + std::string padding(keys_padding_size, '\0'); + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write padding to dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + return dumper->append(PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID, keys_size, + keys_padding_size, 0); +} + +int FlatSparseBuilder::dump_mapping(const std::vector &keys, + IndexDumper *dumper) { + std::vector mapping(keys.size()); + std::iota(mapping.begin(), mapping.end(), 0); + std::sort( + mapping.begin(), mapping.end(), + [&keys](uint32_t lhs, uint32_t rhs) { return (keys[lhs] < keys[rhs]); }); + + size_t mapping_size = mapping.size() * sizeof(uint32_t); + size_t mapping_padding_size = ailego_align(mapping_size, 32) - mapping_size; + if (dumper->write(mapping.data(), mapping_size) != mapping_size) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + + // Write the padding if need + if (mapping_padding_size) { + std::string padding(mapping_padding_size, '\0'); + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + return dumper->append(PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID, mapping_size, + mapping_padding_size, 0); +} + +INDEX_FACTORY_REGISTER_BUILDER(FlatSparseBuilder); + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/flat_sparse/flat_sparse_builder.h b/src/core/algorithm/flat_sparse/flat_sparse_builder.h new file mode 100644 index 0000000..de439d8 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_builder.h @@ -0,0 +1,105 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +/*! Brute Force Sparse Builder + */ +class FlatSparseBuilder : public IndexBuilder { + public: + //! Constructor + FlatSparseBuilder(); + + //! Initialize the builder + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup the builder + int cleanup(void) override; + + //! Train the data + int train(IndexThreads::Pointer, IndexSparseHolder::Pointer holder) override; + + //! Train the data + int train(const IndexTrainer::Pointer &trainer) override; + + int train(IndexThreads::Pointer /*threads*/, + IndexHolder::Pointer /*holder*/) override { + return IndexError_NotImplemented; + } + + int build(IndexThreads::Pointer /*threads*/, + IndexHolder::Pointer /*holder*/) override { + return IndexError_NotImplemented; + } + + //! Build the index + int build(IndexThreads::Pointer threads, + IndexSparseHolder::Pointer holder) override; + + //! Dump index into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + private: + int do_dump(const IndexDumper::Pointer &dumper, uint32_t *dump_count); + + int dump_meta(IndexDumper *dumper); + + int dump_keys(const std::vector &keys, IndexDumper *dumper); + + int dump_mapping(const std::vector &keys, IndexDumper *dumper); + + int dump_vector_and_offset(IndexDumper *dumper, std::vector *keys); + + int write_vector_data(const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_vec, + IndexDumper *dumper, uint32_t *length); + + private: + enum BUILD_STATE { + BUILD_STATE_INIT = 0, + BUILD_STATE_INITED = 1, + BUILD_STATE_TRAINED = 2, + BUILD_STATE_BUILT = 3 + }; + + IndexSparseHolder::Pointer holder_{}; + + std::atomic_bool error_{false}; + IndexMeta meta_{}; + IndexMetric::Pointer measure_{}; + std::mutex mutex_{}; + std::condition_variable cond_{}; + Stats stats_{}; + + BUILD_STATE state_{BUILD_STATE_INIT}; +}; + + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_context.cc b/src/core/algorithm/flat_sparse/flat_sparse_context.cc new file mode 100644 index 0000000..7e3a74b --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_context.cc @@ -0,0 +1,30 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_sparse_context.h" + +namespace zvec { +namespace core { + +const FlatSparseEntity *FlatSparseContext::entity() const { + if (context_type_ == kStreamerContext) { + return &streamer_owner_->entity(); + } else if (context_type_ == kSearcherContext) { + return &searcher_owner_->entity(); + } + return nullptr; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/flat_sparse/flat_sparse_context.h b/src/core/algorithm/flat_sparse/flat_sparse_context.h new file mode 100644 index 0000000..fa62397 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_context.h @@ -0,0 +1,225 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include "flat_sparse_entity.h" +#include "flat_sparse_searcher.h" +#include "flat_sparse_streamer.h" + +namespace zvec { +namespace core { + +class FlatSparseStreamer; +class FlatSparseSearcher; + +/*! Brute Force Sparse Streamer Context + */ +class FlatSparseContext : public IndexContext { + public: + //! Constructor + enum ContextType { + kUnknownContext = 0, + kSearcherContext = 1, + kStreamerContext = 3 + }; + FlatSparseContext(const FlatSparseStreamer *streamer_ptr) + : streamer_owner_(streamer_ptr), context_type_(kStreamerContext) {} + + FlatSparseContext(const FlatSparseSearcher *searcher_ptr) + : searcher_owner_(searcher_ptr), context_type_(kSearcherContext) {} + + //! Destructor + ~FlatSparseContext(void) override = default; + + //! Set topk of search result + void set_topk(uint32_t topk) override { + topk_ = topk; + result_heap_.limit(topk_); + result_heap_.set_threshold(this->threshold()); + } + + //! Retrieve search result + const IndexDocumentList &result(void) const override { + return results_.at(0); + } + + //! Retrieve search result with index + const IndexDocumentList &result(size_t index) const override { + return results_.at(index); + } + + //! Retrieve result object for output + IndexDocumentList *mutable_result(size_t idx) override { + return &results_.at(idx); + } + + inline IndexDocumentHeap *result_heap() { + return &result_heap_; + } + + //! Update the parameters of context + int update(const ailego::Params & /*params*/) override { + return 0; + } + + //! Retrieve magic number + uint32_t magic(void) const override { + return magic_; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + + bool fetch_vector() const override { + return fetch_vector_; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(void) const override { + return group_results_[0]; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(size_t idx) const override { + return group_results_[idx]; + } + + IndexGroupDocumentList *mutable_group_result(void) override { + return &group_results_[0]; + } + + IndexGroupDocumentList *mutable_group_result(size_t idx) override { + return &group_results_[idx]; + } + + //! Set group params + void set_group_params(uint32_t group_num, uint32_t group_topk) override { + group_num_ = group_num; + group_topk_ = group_topk; + result_group_heap_.clear(); + } + + //! Get if group by search + inline bool group_by_search() { + return group_num_ > 0; + } + + inline uint32_t group_topk() const { + return group_topk_; + } + + inline uint32_t group_num() const { + return group_num_; + } + + void reset() override {} + + //! Reset the context + void reset(const FlatSparseStreamer *streamer_ptr) { + magic_ = streamer_ptr->magic(); + streamer_owner_ = streamer_ptr; + context_type_ = kStreamerContext; + } + + void reset(const FlatSparseSearcher *searcher_ptr) { + magic_ = searcher_ptr->magic(); + searcher_owner_ = searcher_ptr; + context_type_ = kSearcherContext; + } + + //! Reset all the query results + void reset_results(size_t qnum) { + if (group_by_search()) { + group_results_.resize(qnum); + } else { + result_heap_.clear(); + result_heap_.limit(topk_); + result_heap_.set_threshold(this->threshold()); + results_.resize(qnum); + stats_vec_.resize(qnum); + for (size_t i = 0; i < results_.size(); ++i) { + results_[i].clear(); + stats_vec_[i].clear(); + } + } + } + + Stats *mutable_stats(size_t idx = 0) { + ailego_assert_with(stats_vec_.size() > idx, "invalid index"); + return &stats_vec_[idx]; + } + + inline void topk_to_result(uint32_t idx) { + if (ailego_unlikely(result_heap_.size() == 0)) { + return; + } + + ailego_assert_with(idx < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(result_heap_.size())); + result_heap_.sort(); + results_[idx].clear(); + for (int i = 0; i < size; ++i) { + auto score = result_heap_[i].score(); + if (score > this->threshold()) { + break; + } + + key_t key = result_heap_[i].key(); + if (fetch_vector_) { + node_id_t id = entity()->get_id(key); + IndexStorage::MemoryBlock vec_block; + entity()->get_sparse_vector(id, vec_block); + const void *sparse_data = vec_block.data(); + IndexSparseDocument sparse_doc; + if (sparse_data != nullptr) { + SparseUtility::ReverseSparseFormat(sparse_data, sparse_doc, + entity()->sparse_unit_size()); + } + results_[idx].emplace_back(key, score, id, nullptr, sparse_doc); + } else { + results_[idx].emplace_back(key, score); + } + } + } + + private: + const FlatSparseEntity *entity() const; + + private: + const FlatSparseStreamer *streamer_owner_{nullptr}; + const FlatSparseSearcher *searcher_owner_{nullptr}; + ContextType context_type_{kUnknownContext}; + std::vector stats_vec_{}; + uint32_t magic_{0}; + uint32_t topk_{0}; + IndexDocumentHeap result_heap_; + // std::string batch_queries_{}; + bool fetch_vector_{false}; + + // group + uint32_t group_num_{0}; + uint32_t group_topk_{0}; + std::map result_group_heap_{}; + std::vector results_{}; + std::vector group_results_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_entity.h b/src/core/algorithm/flat_sparse/flat_sparse_entity.h new file mode 100644 index 0000000..daa3f86 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_entity.h @@ -0,0 +1,246 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include "flat_sparse_index_format.h" + +namespace zvec { +namespace core { + +using node_id_t = uint32_t; +constexpr node_id_t kInvalidNodeId = static_cast(-1); + +/*! Flat Sparse Entity + */ +class FlatSparseEntity { + public: + typedef std::shared_ptr Pointer; + + //! Constructor + explicit FlatSparseEntity() {} + + //! Destructor + ~FlatSparseEntity() = default; + + //! Disable them + FlatSparseEntity(const FlatSparseEntity &) = delete; + FlatSparseEntity &operator=(const FlatSparseEntity &) = delete; + + //! Search in linear list with filter + int search(const std::string &sparse_vector, const IndexFilter &filter, + IndexDocumentHeap *heap) const { + for (node_id_t i = 0; i < doc_cnt(); i++) { + uint64_t key = get_key(i); + if (ailego_unlikely(key == kInvalidKey)) { + // LOG_ERROR("The key of node_id[%u] not found in keys map", i); + // return IndexError_Runtime; + continue; + } + if (!filter.is_valid() || !filter(key)) { + float dist = get_search_distance(sparse_vector, i); + heap->emplace(key, dist); + } + } + + return 0; + } + + //! Search in linear list with filter and target pkeys + int search_p_keys(const std::string &sparse_vector, + const std::vector &p_keys, + const IndexFilter &filter, IndexDocumentHeap *heap) const { + for (auto p_key : p_keys) { + if (!filter.is_valid() || !filter(p_key)) { + auto node_id = get_id(p_key); + if (node_id != kInvalidNodeId) { + float dist = get_search_distance(sparse_vector, node_id); + heap->emplace(p_key, dist); + } + } + } + + return 0; + } + + //! Group search in linear list with filter + int search_group( + const std::string &sparse_vector, const IndexFilter &filter, + const std::function &group_by_func, uint32_t topk, + std::unordered_map *heap) const { + for (node_id_t i = 0; i < doc_cnt(); i++) { + uint64_t key = get_key(i); + if (ailego_unlikely(key == kInvalidKey)) { + LOG_ERROR("The key of node_id[%u] not found in keys map", i); + return IndexError_Runtime; + } + if (!filter.is_valid() || !filter(key)) { + float dist = get_search_distance(sparse_vector, i); + + std::string group_id = group_by_func(key); + + auto &group_heap = (*heap)[group_id]; + if (group_heap.empty()) { + group_heap.limit(topk); + } + group_heap.emplace(key, dist); + } + } + + return 0; + } + + //! Group search in linear list with filter and target pkeys + int search_group_p_keys( + const std::string &sparse_vector, const std::vector &p_keys, + const IndexFilter &filter, + const std::function &group_by_func, uint32_t topk, + std::unordered_map *heap) const { + for (auto p_key : p_keys) { + if (!filter.is_valid() || !filter(p_key)) { + auto node_id = get_id(p_key); + if (node_id != kInvalidNodeId) { + float dist = get_search_distance(sparse_vector, node_id); + + std::string group_id = group_by_func(p_key); + + auto &group_heap = (*heap)[group_id]; + if (group_heap.empty()) { + group_heap.limit(topk); + } + group_heap.emplace(p_key, dist); + } + } + } + + return 0; + } + + //! Get sparse vector by key + int get_sparse_vector(uint64_t key, std::string *sparse_vector) const { + const void *sparse_vector_ptr; + uint32_t sparse_vector_len; + int ret = get_sparse_vector_ptr_by_key(key, &sparse_vector_ptr, + &sparse_vector_len); + if (ret != 0) { + return ret; + } + *sparse_vector = std::string(static_cast(sparse_vector_ptr), + sparse_vector_len); + return 0; + } + + //! Get sparse vector by node id + const void *get_sparse_vector(node_id_t id) const { + const void *sparse_vector_ptr; + uint32_t sparse_vector_len; + int ret = + get_sparse_vector_ptr_by_id(id, &sparse_vector_ptr, &sparse_vector_len); + if (ret != 0) { + return nullptr; + } + return sparse_vector_ptr; + } + + int get_sparse_vector_by_key(const uint64_t key, + std::string *sparse_vector) const { + uint32_t sparse_vector_len; + IndexStorage::MemoryBlock sparse_vector_block; + int ret = get_sparse_vector_ptr_by_key(key, sparse_vector_block, + &sparse_vector_len); + if (ret != 0) { + return ret; + } + *sparse_vector = + std::string(static_cast(sparse_vector_block.data()), + sparse_vector_len); + return 0; + } + + int get_sparse_vector(node_id_t id, + IndexStorage::MemoryBlock &sparse_vector_block) const { + uint32_t sparse_vector_len; + return get_sparse_vector_ptr_by_id(id, sparse_vector_block, + &sparse_vector_len); + } + + int get_sparse_vector_ptr_by_key(uint64_t key, const void **sparse_vector_ptr, + uint32_t *sparse_vector_len_ptr) const { + auto node_id = get_id(key); + if (node_id == kInvalidNodeId) { + return IndexError_NoExist; + } + + return get_sparse_vector_ptr_by_id(node_id, sparse_vector_ptr, + sparse_vector_len_ptr); + } + + int get_sparse_vector_ptr_by_key( + const uint64_t key, IndexStorage::MemoryBlock &sparse_vector_block, + uint32_t *sparse_vector_len_ptr) const { + auto node_id = get_id(key); + if (node_id == kInvalidNodeId) { + return IndexError_NoExist; + } + + return get_sparse_vector_ptr_by_id(node_id, sparse_vector_block, + sparse_vector_len_ptr); + } + + std::vector get_keys() const { + std::vector keys; + node_id_t doc_total_cnt = doc_cnt(); + for (node_id_t node_id = 0; node_id < doc_total_cnt; ++node_id) { + uint64_t key = get_key(node_id); + if (key == kInvalidKey) { + return {kInvalidKey}; + } else { + keys.push_back(key); + } + } + + return keys; + } + + + public: + virtual uint32_t doc_cnt() const = 0; + + virtual uint32_t total_sparse_count() const = 0; + + virtual node_id_t get_id(uint64_t key) const = 0; + + virtual uint64_t get_key(node_id_t id) const = 0; + + virtual int get_sparse_vector_ptr_by_id( + node_id_t id, const void **sparse_vector, + uint32_t *sparse_vector_len) const = 0; + + virtual int get_sparse_vector_ptr_by_id( + const node_id_t /*id*/, + IndexStorage::MemoryBlock & /*sparse_vector_block*/, + uint32_t * /*sparse_vector_len*/) const { + return IndexError_NotImplemented; + } + + + virtual float get_search_distance(const std::string &vector, + node_id_t target_node_id) const = 0; + virtual size_t sparse_unit_size() const = 0; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_index_format.h b/src/core/algorithm/flat_sparse/flat_sparse_index_format.h new file mode 100644 index 0000000..23b5974 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_index_format.h @@ -0,0 +1,45 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace core { + +static constexpr uint64_t kInvalidKey = std::numeric_limits::max(); +static constexpr uint32_t kDefaultOffsetChunkSize = 1024 * 1024; // 1MB +static constexpr uint32_t kDefaultDataChunkSize = 8 * 1024 * 1024; // 8MB + +struct FlatSparseMeta { + uint64_t create_time{0}; + uint64_t update_time{0}; + uint32_t doc_cnt{0}; + uint32_t total_sparse_count{0}; + uint8_t reserved[8] = {0}; +}; + +static_assert(sizeof(FlatSparseMeta) % 32 == 0, + "FlatSparseMeta must be aligned with 32 bytes"); + +struct FlatSparseStreamerMeta { + uint32_t offset_chunk_count{0}; + uint32_t offset_chunk_size{kDefaultOffsetChunkSize}; + uint32_t data_chunk_count{0}; + uint32_t data_chunk_size{kDefaultDataChunkSize}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_provider.h b/src/core/algorithm/flat_sparse/flat_sparse_provider.h new file mode 100644 index 0000000..8065b5f --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_provider.h @@ -0,0 +1,193 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include "flat_sparse_streamer_entity.h" + +namespace zvec { +namespace core { + +/*! Brute Force Sparse Streamer Provider + */ +// FlatSparseStreamerEntity or FlatSparseSearcherEntity +template +class FlatSparseIndexProvider : public IndexSparseProvider { + public: + //! Constructor + FlatSparseIndexProvider(const std::shared_ptr entity, + const IndexMeta &meta, const std::string &owner) + : entity_(entity), meta_(meta), owner_class_(owner) {} + + //! Create a new iterator + IndexSparseProvider::Iterator::Pointer create_iterator(void) override { + return IndexSparseProvider::Iterator::Pointer(new (std::nothrow) + Iterator(entity_, meta_)); + } + + //! Retrieve count of vectors + size_t count(void) const override { + return entity_->doc_cnt(); + } + + //! Retrieve type of vector + IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + //! Retrieve a vector using a primary key + int get_sparse_vector(uint64_t key, uint32_t *sparse_count, + std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override { + std::string sparse_data; + + int ret = entity_->get_sparse_vector_by_key(key, &sparse_data); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to get sparse vector, key=%zu, ret=%s", (size_t)key, + IndexError::What(ret)); + return ret; + } + + SparseUtility::ReverseSparseFormat(sparse_data, sparse_count, + sparse_indices_buffer, + sparse_values_buffer, meta_.unit_size()); + return 0; + } + + //! Retrieve the owner class + const std::string &owner_class(void) const override { + return owner_class_; + } + + size_t total_sparse_count() const override { + return entity_->total_sparse_count(); + } + + private: + class Iterator : public IndexSparseProvider::Iterator { + public: + Iterator(const std::shared_ptr &entity, + const IndexMeta &meta) + : entity_(entity), meta_(meta), cur_id_(0U), valid_(false) { + IndexStorage::MemoryBlock sparse_data_block; + entity_->get_sparse_vector(cur_id_, sparse_data_block); + const void *sparse_data = sparse_data_block.data(); + if (sparse_data != nullptr) { + valid_ = true; + + sparse_indices_buffer_.clear(); + sparse_data_buffer_.clear(); + + SparseUtility::ReverseSparseFormat( + sparse_data, &sparse_count_, &sparse_indices_buffer_, + &sparse_data_buffer_, meta.unit_size()); + } + } + + //! Retrieve sparse count + uint32_t sparse_count() const override { + return sparse_count_; + } + + //! Retrieve sparse indices + const uint32_t *sparse_indices() const override { + return reinterpret_cast(sparse_indices_buffer_.data()); + } + + //! Retrieve sparse data + const void *sparse_data() const override { + return reinterpret_cast(sparse_data_buffer_.data()); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt() && valid_; + } + + //! Retrieve primary key + uint64_t key(void) const override { + // std::cout << "iter key=" << cur_id_ << std::endl; + return entity_->get_key(cur_id_); + } + + //! Next iterator + void next(void) override { + cur_id_ = get_next_valid_id(cur_id_ + 1); + + if (cur_id_ < entity_->doc_cnt()) { + IndexStorage::MemoryBlock sparse_data_block; + entity_->get_sparse_vector(cur_id_, sparse_data_block); + const void *sparse_data = sparse_data_block.data(); + if (sparse_data != nullptr) { + valid_ = true; + + sparse_indices_buffer_.clear(); + sparse_data_buffer_.clear(); + + SparseUtility::ReverseSparseFormat( + sparse_data, &sparse_count_, &sparse_indices_buffer_, + &sparse_data_buffer_, meta_.unit_size()); + } else { + valid_ = false; + } + } + } + + //! Reset the iterator + void reset(void) { + cur_id_ = get_next_valid_id(0); + IndexStorage::MemoryBlock sparse_data_block; + entity_->get_sparse_vector(cur_id_, sparse_data_block); + const void *sparse_data = sparse_data_block.data(); + if (sparse_data != nullptr) { + valid_ = true; + + SparseUtility::ReverseSparseFormat( + sparse_data, &sparse_count_, &sparse_indices_buffer_, + &sparse_data_buffer_, meta_.unit_size()); + } + } + + private: + node_id_t get_next_valid_id(node_id_t start_id) { + for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) { + if (entity_->get_key(i) != kInvalidNodeId) { + return i; + } + } + return kInvalidNodeId; + } + + private: + const std::shared_ptr entity_{nullptr}; + const IndexMeta &meta_; + node_id_t cur_id_; + uint32_t sparse_count_; + std::string sparse_indices_buffer_; + std::string sparse_data_buffer_; + bool valid_{false}; + }; + + private: + const std::shared_ptr entity_{nullptr}; + const IndexMeta &meta_; + const std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_search.h b/src/core/algorithm/flat_sparse/flat_sparse_search.h new file mode 100644 index 0000000..11c5429 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_search.h @@ -0,0 +1,174 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include "flat_sparse_context.h" + +namespace zvec { +namespace core { + +static inline IndexGroupDocumentList ConvertGroupMapToResult( + std::unordered_map group_map, + uint32_t group_num) { + IndexGroupDocumentList result; + + std::vector> best_score_in_groups; + for (auto itr = group_map.begin(); itr != group_map.end(); itr++) { + const std::string &group_id = (*itr).first; + auto &heap = (*itr).second; + + if (heap.size() > 0) { + float best_score = heap[0].score(); + best_score_in_groups.push_back(std::make_pair(group_id, best_score)); + } + } + + std::sort(best_score_in_groups.begin(), best_score_in_groups.end(), + [](const std::pair &a, + const std::pair &b) -> int { + return a.second < b.second; + }); + + // truncate to group num + for (uint32_t i = 0; i < group_num && i < best_score_in_groups.size(); ++i) { + const std::string &group_id = best_score_in_groups[i].first; + + result.emplace_back( + GroupIndexDocument(group_id, std::move(group_map[group_id]))); + } + + return result; +} + +static inline int FlatSearch(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, bool with_p_keys, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + const IndexMeta, IndexContext::Pointer &context, + FlatSparseEntity *entity) { + int ret; + + FlatSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to FlatSparseContext failed"); + return IndexError_Cast; + } + + // reset context results + ctx->reset_results(count); + + const uint32_t *sparse_indices_tmp = sparse_indices; + const void *sparse_query_tmp = sparse_query; + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](uint64_t key) { + return ctx->group_by()(key); + }; + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + ailego::MinusInnerProductSparseMatrix::transform_sparse_format( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + sparse_query_buffer); + + std::unordered_map group_heap{}; + + if (with_p_keys) { + ret = entity->search_group_p_keys(sparse_query_buffer, p_keys[q], + ctx->filter(), group_by, + ctx->group_topk(), &group_heap); + } else { + ret = entity->search_group(sparse_query_buffer, ctx->filter(), group_by, + ctx->group_topk(), &group_heap); + } + + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to search group, ret=%s", IndexError::What(ret)); + return ret; + } + + // sort group heap + for (auto &group : group_heap) { + group.second.sort(); + } + + auto group_result = + ConvertGroupMapToResult(std::move(group_heap), ctx->group_num()); + + // Populate sparse vector data when fetch_vector is enabled + if (ctx->fetch_vector()) { + for (auto &group_doc : group_result) { + for (auto &doc : *group_doc.mutable_docs()) { + node_id_t id = entity->get_id(doc.key()); + if (id != kInvalidNodeId) { + IndexSparseDocument sparse_doc; + IndexStorage::MemoryBlock vec_block; + entity->get_sparse_vector(id, vec_block); + const void *sparse_data = vec_block.data(); + if (sparse_data != nullptr) { + SparseUtility::ReverseSparseFormat(sparse_data, sparse_doc, + entity->sparse_unit_size()); + } + // Reconstruct doc with sparse vector data + doc = IndexDocument(doc.key(), doc.score(), id, nullptr, + sparse_doc); + } + } + } + } + + ctx->mutable_group_result(q)->swap(group_result); + } + } else { + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + ailego::MinusInnerProductSparseMatrix::transform_sparse_format( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + sparse_query_buffer); + + auto heap = ctx->result_heap(); + + if (with_p_keys) { + ret = entity->search_p_keys(sparse_query_buffer, p_keys[q], + ctx->filter(), heap); + } else { + ret = entity->search(sparse_query_buffer, ctx->filter(), heap); + } + + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to search, ret=%s", IndexError::What(ret)); + return ret; + } + + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + } + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_searcher.cc b/src/core/algorithm/flat_sparse/flat_sparse_searcher.cc new file mode 100644 index 0000000..3e9fce0 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_searcher.cc @@ -0,0 +1,182 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_sparse_searcher.h" +#include +#include +#include "flat_sparse_context.h" +#include "flat_sparse_provider.h" +#include "flat_sparse_search.h" + +namespace zvec { +namespace core { + +const uint32_t FlatSparseSearcher::VERSION = 0U; + +FlatSparseSearcher::FlatSparseSearcher(void) {} + +FlatSparseSearcher::~FlatSparseSearcher(void) {} + +int FlatSparseSearcher::init(const ailego::Params & /*params*/) { + state_ = STATE_INITED; + + return 0; +} + +int FlatSparseSearcher::cleanup(void) { + this->unload(); + return 0; +} + +int FlatSparseSearcher::load(IndexStorage::Pointer container, + IndexMetric::Pointer /*measure*/) { + if (state_ != STATE_INITED) { + LOG_ERROR("Init the searcher first before load index"); + return IndexError_Runtime; + } + + LOG_INFO("Begin FlatSparseSearcher::load"); + + int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize meta from container"); + return ret; + } + + if (meta_.searcher_revision() != VERSION) { + LOG_ERROR("Unsupported searcher revision %u", meta_.searcher_revision()); + return IndexError_Unsupported; + } + + ret = entity_.load(container, meta_); + if (ret != 0) { + LOG_ERROR("FlatSparseSearcher load index failed"); + return ret; + } + + state_ = STATE_LOADED; + magic_ = IndexContext::GenerateMagic(); + + LOG_INFO("End FlatSparseSearcher::load"); + + return 0; +} + +int FlatSparseSearcher::unload(void) { + LOG_INFO("Begin FlatSparseSearcher::unload"); + + meta_.clear(); + entity_.unload(); + state_ = STATE_INITED; + + LOG_INFO("End FlatSparseSearcher::unload"); + + return 0; +} + +int FlatSparseSearcher::search_bf_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + return do_search(sparse_count, sparse_indices, sparse_query, false, {}, qmeta, + count, context); +} + +int FlatSparseSearcher::search_bf_by_p_keys_impl( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const { + return do_search(sparse_count, sparse_indices, sparse_query, true, p_keys, + qmeta, count, context); +} + +int FlatSparseSearcher::get_sparse_vector( + uint64_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Failed to get sparse vector, load container first!"); + return IndexError_NoIndexLoaded; + } + + std::string sparse_data; + + int ret = entity_.get_sparse_vector(key, &sparse_data); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to get sparse vector, key=%zu, ret=%s", (size_t)key, + IndexError::What(ret)); + return ret; + } + + SparseUtility::ReverseSparseFormat(sparse_data, sparse_count, + sparse_indices_buffer, + sparse_values_buffer, meta_.unit_size()); + + return 0; +} + +FlatSparseSearcher::ContextPointer FlatSparseSearcher::create_context() const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Failed to create Context, load container first!"); + return Context::UPointer(); + } + FlatSparseSearcherEntity::Pointer entity = entity_.clone(); + return FlatSparseSearcher::ContextPointer(new FlatSparseContext(this)); +} + +//! Create a new iterator +IndexSearcher::SparseProvider::Pointer +FlatSparseSearcher::create_sparse_provider(void) const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Failed to create provider, load container first!"); + return SparseProvider::Pointer(); + } + + auto entity = entity_.clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone entity failed"); + return SparseProvider::Pointer(); + } + return SparseProvider::Pointer( + new FlatSparseIndexProvider( + entity, meta_, "FlatSparseSearcher")); +} + +int FlatSparseSearcher::do_search( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, bool with_p_keys, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Failed to do search, load container first!"); + return IndexError_NoIndexLoaded; + } + + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + return FlatSearch(sparse_count, sparse_indices, sparse_query, with_p_keys, + p_keys, qmeta, count, meta_, context, + (FlatSparseEntity *)&entity_); +} + +INDEX_FACTORY_REGISTER_SEARCHER(FlatSparseSearcher); + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/flat_sparse/flat_sparse_searcher.h b/src/core/algorithm/flat_sparse/flat_sparse_searcher.h new file mode 100644 index 0000000..0e7f439 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_searcher.h @@ -0,0 +1,185 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "flat_sparse_searcher_entity.h" + +namespace zvec { +namespace core { + +class FlatSparseSearcher : public IndexSearcher { + public: + static const uint32_t VERSION; + + public: + using ContextPointer = IndexSearcher::Context::Pointer; + + public: + FlatSparseSearcher(void); + ~FlatSparseSearcher(void) override; + + FlatSparseSearcher(const FlatSparseSearcher &) = delete; + FlatSparseSearcher &operator=(const FlatSparseSearcher &) = delete; + + public: + //! Initialize Searcher + int init(const ailego::Params ¶ms) override; + + //! Cleanup Searcher + int cleanup(void) override; + + //! Load Index from storage + int load(IndexStorage::Pointer container, + IndexMetric::Pointer /*measure*/) override; + + //! Unload index from storage + int unload(void) override; + + int search_impl(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + Context::Pointer & /*context*/) const override { + return IndexError_NotImplemented; + } + + int search_impl(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + uint32_t /*count*/, + Context::Pointer & /*context*/) const override { + return IndexError_NotImplemented; + } + + int search_bf_impl(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + Context::Pointer & /*context*/) const override { + return IndexError_NotImplemented; + } + + int search_bf_impl(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + uint32_t /*count*/, + Context::Pointer & /*context*/) const override { + return IndexError_NotImplemented; + } + + //! Similarity search with sparse inputs + int search_impl(const uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override { + return search_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1, + context); + } + + //! Similarity search with sparse inputs + int search_impl(const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override { + return search_bf_impl(sparse_count, sparse_indices, sparse_query, qmeta, + count, context); + } + + //! Similarity brute force search with sparse inputs + int search_bf_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const override { + return search_bf_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1, + context); + } + + //! Similarity brute force search with sparse inputs + int search_bf_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(&sparse_count, sparse_indices, sparse_query, + p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + //! Fetch sparser vector by key + int get_sparse_vector(uint64_t key, uint32_t *sparse_count, + std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override; + + //! Create a searcher context + ContextPointer create_context() const override; + + //! Create a new iterator + IndexSearcher::SparseProvider::Pointer create_sparse_provider( + void) const override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve params of index + const ailego::Params ¶ms(void) const override { + return params_; + } + + const FlatSparseSearcherEntity &entity(void) const { + return entity_; + } + + uint32_t magic(void) const { + return magic_; + } + + private: + inline int check_params(const IndexQueryMeta &qmeta) const { + if (ailego_unlikely(qmeta.data_type() != meta_.data_type())) { + LOG_ERROR("Unsupported query meta"); + return IndexError_Mismatch; + } + return 0; + } + + int do_search(const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, bool with_p_keys, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; + + FlatSparseSearcherEntity entity_{}; + IndexMeta meta_{}; + ailego::Params params_{}; + uint32_t magic_{0U}; + + Stats stats_; + State state_{STATE_INIT}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.cc b/src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.cc new file mode 100644 index 0000000..91357a5 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.cc @@ -0,0 +1,241 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_sparse_searcher_entity.h" +#include +#include +#include "flat_sparse_utility.h" + +namespace zvec { +namespace core { + +FlatSparseSearcherEntity::FlatSparseSearcherEntity() {} + +int FlatSparseSearcherEntity::load(const IndexStorage::Pointer &container, + const IndexMeta &index_meta) { + if (container_) { + LOG_ERROR("An storage instance is already opened"); + return IndexError_Duplicate; + } + + int ret = this->load_container(container); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to load storage index"); + return ret; + } + + if (init_measure(index_meta) != 0) { + LOG_ERROR("Failed to init measure"); + return IndexError_InvalidFormat; + } + + container_ = container; + return 0; +} + +int FlatSparseSearcherEntity::init_measure(const IndexMeta &meta) { + measure_ = IndexFactory::CreateMetric(meta.metric_name()); + if (!measure_) { + LOG_ERROR("Failed to create measure %s", meta.metric_name().c_str()); + return IndexError_NoExist; + } + int ret = measure_->init(meta, meta.metric_params()); + if (ret != 0) { + LOG_ERROR("Failled to init measure, ret=%d", ret); + return ret; + } + + if (!measure_->sparse_distance()) { + LOG_ERROR("Invalid measure distance"); + return IndexError_InvalidArgument; + } + + search_sparse_distance_ = measure_->sparse_distance(); + + if (measure_->query_metric() && measure_->query_metric()->distance()) { + search_sparse_distance_ = measure_->query_metric()->sparse_distance(); + } + sparse_unit_size_ = meta.unit_size(); + + return 0; +} + +int FlatSparseSearcherEntity::load_container( + const IndexStorage::Pointer &container) { + // meta + auto segment = container->get(PARAM_FLAT_SPARSE_META_SEG_ID); + if (!segment || segment->data_size() < sizeof(meta_)) { + LOG_ERROR("Missing segment %s, or invalid segment size", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + const void *data; + if (ailego_unlikely(segment->read(0, &data, sizeof(meta_)) != + sizeof(meta_))) { + LOG_ERROR("Failed to read meta segment %s", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return IndexError_ReadData; + } + meta_ = *(reinterpret_cast(data)); + + // keys segment + keys_chunk_ = container->get(PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID); + if (!keys_chunk_) { + LOG_ERROR("Missing segment %s", PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + + // mapping segment + mapping_chunk_ = container->get(PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID); + if (!mapping_chunk_) { + LOG_ERROR("Missing segment %s", + PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + + // offset segment + sparse_offset_chunk_ = container->get(PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID); + if (!sparse_offset_chunk_) { + LOG_ERROR("Missing segment %s", + PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + + // data segment + sparse_data_chunk_ = container->get(PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID); + if (!sparse_data_chunk_) { + LOG_ERROR("Missing segment %s", PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + + return 0; +} + +int FlatSparseSearcherEntity::unload() { + container_.reset(); + sparse_data_chunk_.reset(); + sparse_offset_chunk_.reset(); + keys_chunk_.reset(); + mapping_chunk_.reset(); + + return 0; +} + +FlatSparseSearcherEntity::Pointer FlatSparseSearcherEntity::clone() const { + auto entity = new (std::nothrow) + FlatSparseSearcherEntity(meta_, sparse_data_chunk_, sparse_offset_chunk_, + keys_chunk_, mapping_chunk_); + return FlatSparseSearcherEntity::Pointer(entity); +} + +int FlatSparseSearcherEntity::get_sparse_vector_ptr_by_id( + node_id_t id, const void **sparse_vector_ptr, + uint32_t *sparse_vector_len_ptr) const { + uint32_t offset_chunk_offset = id * offset_size_per_node(); + + const void *offset_info = nullptr; + if (ailego_unlikely(sparse_offset_chunk_->read( + offset_chunk_offset, &offset_info, + offset_size_per_node()) != offset_size_per_node())) { + LOG_ERROR("Read offset info failed, offset=%u", offset_chunk_offset); + return IndexError_ReadData; + }; + + // sparse offset + uint64_t sparse_offset = *(uint64_t *)offset_info; + uint32_t sparse_vector_len = + *(uint32_t *)((uint8_t *)offset_info + sizeof(uint64_t)); + + if (sparse_vector_len > 0) { + const void *sparse_data = + get_sparse_vector_data(sparse_offset, sparse_vector_len); + if (ailego_unlikely(sparse_data == nullptr)) { + LOG_ERROR("Get nullptr sparse, offset=%zu, len=%u", (size_t)sparse_offset, + sparse_vector_len); + + return IndexError_ReadData; + } + *sparse_vector_ptr = sparse_data; + *sparse_vector_len_ptr = sparse_vector_len; + } + + return 0; +} + +const void *FlatSparseSearcherEntity::get_sparse_vector_data( + uint64_t offset, uint32_t length) const { + const void *data; + auto size = sparse_data_chunk_->read(offset, &data, length); + if (size != length) { + LOG_ERROR( + "read sparse vector data failed: offset=%zu, " + "length=%u, size=%zu", + (size_t)offset, length, size); + return nullptr; + } + return data; +} + + +node_id_t FlatSparseSearcherEntity::get_id(uint64_t key) const { + if (ailego_unlikely(!mapping_chunk_)) { + LOG_ERROR("Index missing mapping segment"); + return kInvalidNodeId; + } + + //! Do binary search + node_id_t start = 0UL; + node_id_t end = doc_cnt(); + const void *data; + node_id_t idx = 0u; + while (start < end) { + idx = start + (end - start) / 2; + if (ailego_unlikely(mapping_chunk_->read(idx * sizeof(node_id_t), &data, + sizeof(node_id_t)) != + sizeof(node_id_t))) { + LOG_ERROR("Read key from segment failed"); + return kInvalidNodeId; + } + const uint64_t *mkey; + node_id_t local_id = *reinterpret_cast(data); + if (ailego_unlikely(keys_chunk_->read( + local_id * sizeof(uint64_t), (const void **)(&mkey), + sizeof(uint64_t)) != sizeof(uint64_t))) { + LOG_ERROR("Read key from segment failed"); + return kInvalidNodeId; + } + if (*mkey < key) { + start = idx + 1; + } else if (*mkey > key) { + end = idx; + } else { + return local_id; + } + } + return kInvalidNodeId; +} + +uint64_t FlatSparseSearcherEntity::get_key(node_id_t id) const { + const void *key; + if (ailego_unlikely( + keys_chunk_->read(id * sizeof(uint64_t), &key, sizeof(uint64_t)) != + sizeof(uint64_t))) { + LOG_ERROR("Read key from segment failed"); + return kInvalidKey; + } + return *(reinterpret_cast(key)); +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.h b/src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.h new file mode 100644 index 0000000..f87c294 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.h @@ -0,0 +1,126 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include "flat_sparse_entity.h" +#include "flat_sparse_index_format.h" + +namespace zvec { +namespace core { + + +/*! Flat Sparse Searcher Entity + */ +class FlatSparseSearcherEntity : public FlatSparseEntity { + public: + typedef std::shared_ptr Pointer; + + using Chunk = IndexStorage::Segment; + + //! Constructor + explicit FlatSparseSearcherEntity(); + + //! Destructor + virtual ~FlatSparseSearcherEntity() = default; + + //! Disable them + FlatSparseSearcherEntity(const FlatSparseSearcherEntity &) = delete; + FlatSparseSearcherEntity &operator=(const FlatSparseSearcherEntity &) = + delete; + + //! Load the entity with container + int load(const IndexStorage::Pointer &container, const IndexMeta &index_meta); + + //! Unload the entity + int unload(); + + public: + inline uint32_t doc_cnt() const override { + return meta_.doc_cnt; + } + + inline uint32_t total_sparse_count() const override { + return meta_.total_sparse_count; + } + + size_t sparse_unit_size() const override { + return sparse_unit_size_; + } + + float get_search_distance(const std::string &vector, + node_id_t target_node_id) const override { + float dist; + const void *target_vector; + uint32_t target_vector_len; + get_sparse_vector_ptr_by_id(target_node_id, &target_vector, + &target_vector_len); + search_sparse_distance_(vector.c_str(), target_vector, &dist); + return dist; + } + + FlatSparseSearcherEntity::Pointer clone() const; + + node_id_t get_id(uint64_t key) const override; + + uint64_t get_key(node_id_t id) const override; + + int get_sparse_vector_ptr_by_id(node_id_t id, const void **sparse_vector, + uint32_t *sparse_vector_len) const override; + + private: + int load_container(const IndexStorage::Pointer &container); + + int init_measure(const IndexMeta &meta); + + inline uint32_t offset_size_per_node() const { + return sizeof(uint64_t) + sizeof(uint32_t); + } + + const void *get_sparse_vector_data(uint64_t offset, uint32_t length) const; + + private: + FlatSparseSearcherEntity(const FlatSparseMeta &meta, + Chunk::Pointer sparse_data_chunk, + Chunk::Pointer sparse_offset_chunk, + Chunk::Pointer keys_chunk, + Chunk::Pointer mapping_chunk) + : meta_(meta), + sparse_data_chunk_(sparse_data_chunk), + sparse_offset_chunk_(sparse_offset_chunk), + keys_chunk_(keys_chunk), + mapping_chunk_(mapping_chunk) {} + + private: + IndexStorage::Pointer container_{}; + + // meta + FlatSparseMeta meta_; + + // measure + IndexMetric::Pointer measure_{}; + IndexMetric::MatrixSparseDistance search_sparse_distance_{}; + + // chunk + Chunk::Pointer sparse_data_chunk_; + Chunk::Pointer sparse_offset_chunk_; + Chunk::Pointer keys_chunk_; + Chunk::Pointer mapping_chunk_; + + size_t sparse_unit_size_{0U}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_streamer.cc b/src/core/algorithm/flat_sparse/flat_sparse_streamer.cc new file mode 100644 index 0000000..bc5a3da --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_streamer.cc @@ -0,0 +1,411 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_sparse_streamer.h" +#include +#include +#include +#include +#include +#include "flat_sparse_context.h" +#include "flat_sparse_provider.h" +#include "flat_sparse_search.h" + +namespace zvec { +namespace core { + +const uint32_t FlatSparseStreamer::VERSION = 0U; + +FlatSparseStreamer::FlatSparseStreamer() : entity_(stats_) {} + +FlatSparseStreamer::~FlatSparseStreamer() { + this->close(); +} + +int FlatSparseStreamer::init(const IndexMeta &imeta, + const ailego::Params ¶ms) { + LOG_DEBUG("FlatSparseStreamer init"); + + meta_ = imeta; + meta_.set_streamer("FlatSparseStreamer", VERSION, params); + + state_ = STATE_INITED; + + return 0; +} + +int FlatSparseStreamer::cleanup() { + LOG_DEBUG("FlatSparseStreamer cleanup"); + + this->close(); + + meta_.clear(); + + return 0; +} + +int FlatSparseStreamer::open(IndexStorage::Pointer stg) { + LOG_DEBUG("FlatSparseStreamer open"); + + if (ailego_unlikely(state_ != STATE_INITED)) { + LOG_ERROR("Open storage failed, init streamer first!"); + return IndexError_NoReady; + } + + int ret = entity_.open(std::move(stg), meta_); + if (ret != 0) { + LOG_ERROR("FlatSparseStreamer entity failed to open storage"); + return ret; + } + + IndexMeta index_meta; + ret = entity_.get_index_sparse_meta(&index_meta); + if (ret == IndexError_NoExist) { + // Set IndexMeta for the new index + ret = entity_.set_index_sparse_meta(meta_); + if (ret != 0) { + LOG_ERROR("Failed to set index meta for %s", IndexError::What(ret)); + return ret; + } + } else { + if (index_meta.streamer_revision() != meta_.streamer_revision()) { + LOG_ERROR("Streamer revision mismatch, expect=%u, actual=%u", + meta_.streamer_revision(), index_meta.streamer_revision()); + return IndexError_Mismatch; + } + if (index_meta.metric_name() != meta_.metric_name() || + index_meta.data_type() != meta_.data_type()) { + LOG_ERROR("IndexMeta mismatch from the previous in index"); + return IndexError_Mismatch; + } + // The IndexMeasure Params may be updated like MipsSquaredEuclidean + auto metric_params = index_meta.metric_params(); + metric_params.merge(meta_.metric_params()); + meta_.set_metric(index_meta.metric_name(), 0, metric_params); + } + + state_ = STATE_OPENED; + magic_ = IndexContext::GenerateMagic(); + + return 0; +} + +int FlatSparseStreamer::close() { + if (state_ != STATE_OPENED) { + return 0; + } + + LOG_DEBUG("FlatSparseStreamer close"); + + stats_.clear(); + int ret = entity_.close(); + if (ret != 0) { + LOG_ERROR("Failed to close entity %s", IndexError::What(ret)); + return ret; + } + state_ = STATE_INITED; + return 0; +} + +int FlatSparseStreamer::flush(uint64_t checkpoint) { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to flush, open streamer first!"); + return IndexError_NoReady; + } + + LOG_INFO("FlatSparseStreamer flush, checkpoint=%zu", (size_t)checkpoint); + + return entity_.flush(checkpoint); +} + +int FlatSparseStreamer::dump(const IndexDumper::Pointer &dumper) { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to dump, open streamer first!"); + return IndexError_NoReady; + } + + LOG_INFO("FlatSparseStreamer dump"); + + shared_mutex_.lock(); + AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); + + meta_.set_searcher("FlatSparseSearcher", VERSION, ailego::Params()); + + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + + return entity_.dump(dumper); +} + +FlatSparseStreamer::ContextPointer FlatSparseStreamer::create_context() const { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to create Context, open streamer first!"); + return Context::UPointer(); + } + FlatSparseStreamerEntity::Pointer entity = entity_.clone(); + return FlatSparseStreamer::ContextPointer(new FlatSparseContext(this)); +} + +IndexStreamer::SparseProvider::Pointer +FlatSparseStreamer::create_sparse_provider(void) const { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to create provider, open streamer first!"); + return SparseProvider::Pointer(); + } + + auto entity = entity_.clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone entity failed"); + return SparseProvider::Pointer(); + } + return SparseProvider::Pointer( + new FlatSparseIndexProvider( + entity, meta_, "FlatSparseStreamerProvider")); +} + +int FlatSparseStreamer::add_impl(uint64_t pkey, const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to add_impl, open streamer first!"); + (*stats_.mutable_discarded_count())++; + return IndexError_NoReady; + } + + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(sparse_count > PARAM_FLAT_SPARSE_MAX_DIM_SIZE)) { + LOG_ERROR( + "Failed to add sparse vector: number of non-zero elements (%u) exceeds " + "maximum allowed (%u), key=%zu", + sparse_count, PARAM_FLAT_SPARSE_MAX_DIM_SIZE, (size_t)pkey); + (*stats_.mutable_discarded_count())++; + return IndexError_InvalidValue; + } + + // context is trivial here + FlatSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to FlatSparseContext failed"); + (*stats_.mutable_discarded_count())++; + return IndexError_Cast; + } + + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + // convert to sparse format and add to entity + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_query, + meta_.unit_size(), sparse_query_buffer); + + ret = entity_.add(pkey, sparse_query_buffer, sparse_count); + if (ret != 0) { + LOG_ERROR("Failed to add sparse vector, key=%zu, ret=%s", (size_t)pkey, + IndexError::What(ret)); + (*stats_.mutable_discarded_count())++; + return ret; + } + + (*stats_.mutable_added_count())++; + return 0; +} + +int FlatSparseStreamer::add_with_id_impl(uint32_t pkey, + const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to add_with_id_impl, open streamer first!"); + (*stats_.mutable_discarded_count())++; + return IndexError_NoReady; + } + + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(sparse_count > PARAM_FLAT_SPARSE_MAX_DIM_SIZE)) { + LOG_ERROR( + "Failed to add sparse vector: number of non-zero elements (%u) exceeds " + "maximum allowed (%u), key=%zu", + sparse_count, PARAM_FLAT_SPARSE_MAX_DIM_SIZE, (size_t)pkey); + (*stats_.mutable_discarded_count())++; + return IndexError_InvalidValue; + } + + // context is trivial here + FlatSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to FlatSparseContext failed"); + (*stats_.mutable_discarded_count())++; + return IndexError_Cast; + } + + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + // convert to sparse format and add to entity + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_query, + meta_.unit_size(), sparse_query_buffer); + + ret = entity_.add_vector_with_id(pkey, sparse_query_buffer, sparse_count); + if (ret != 0) { + LOG_ERROR("Failed to add sparse vector, key=%zu, ret=%s", (size_t)pkey, + IndexError::What(ret)); + (*stats_.mutable_discarded_count())++; + return ret; + } + + (*stats_.mutable_added_count())++; + return 0; +} + +//! Similarity search with sparse inputs +int FlatSparseStreamer::search_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1, + context); +} + +//! Similarity search with sparse inputs +int FlatSparseStreamer::search_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + return search_bf_impl(sparse_count, sparse_indices, sparse_query, qmeta, + count, context); +} + +//! Similarity brute force search with sparse inputs +int FlatSparseStreamer::search_bf_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_bf_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1, + context); +} + +//! Linear search by primary keys +int FlatSparseStreamer::search_bf_by_p_keys_impl( + const uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, ContextPointer &context) const { + return search_bf_by_p_keys_impl(&sparse_count, sparse_indices, sparse_query, + p_keys, qmeta, 1, context); +} + +//! Similarity brute force search with sparse inputs +int FlatSparseStreamer::search_bf_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + return do_search(sparse_count, sparse_indices, sparse_query, false, {}, qmeta, + count, context); +} + +//! Linear search by primary keys with sparse inputs +int FlatSparseStreamer::search_bf_by_p_keys_impl( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const { + return do_search(sparse_count, sparse_indices, sparse_query, true, p_keys, + qmeta, count, context); +} + +//! Fetch sparse vector by key +int FlatSparseStreamer::get_sparse_vector( + uint64_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to get_sparse_vector, open streamer first!"); + return IndexError_NoReady; + } + + std::string sparse_data; + + int ret = entity_.get_sparse_vector_by_key(key, &sparse_data); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to get sparse vector, key=%zu, ret=%s", (size_t)key, + IndexError::What(ret)); + return ret; + } + + SparseUtility::ReverseSparseFormat(sparse_data, sparse_count, + sparse_indices_buffer, + sparse_values_buffer, meta_.unit_size()); + + return 0; +} + +int FlatSparseStreamer::do_search( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, bool with_p_keys, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const { + if (state_ != STATE_OPENED) { + LOG_ERROR("Failed to do_search, open streamer first!"); + return IndexError_NoReady; + } + + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + FlatSparseContext *ctx = dynamic_cast(context.get()); + if (ctx->magic() != magic_) { + ctx->reset(this); + } + + return FlatSearch(sparse_count, sparse_indices, sparse_query, with_p_keys, + p_keys, qmeta, count, meta_, context, + (FlatSparseEntity *)&entity_); +} + +INDEX_FACTORY_REGISTER_STREAMER(FlatSparseStreamer); + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/flat_sparse/flat_sparse_streamer.h b/src/core/algorithm/flat_sparse/flat_sparse_streamer.h new file mode 100644 index 0000000..661a7be --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_streamer.h @@ -0,0 +1,172 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include "flat_sparse_streamer_entity.h" + +namespace zvec { +namespace core { + +/*! Flat Sparse Streamer + */ +class FlatSparseStreamer : public IndexStreamer { + public: + static const uint32_t VERSION; + + public: + using ContextPointer = IndexStreamer::Context::Pointer; + + FlatSparseStreamer(void); + ~FlatSparseStreamer(void) override; + + FlatSparseStreamer(const FlatSparseStreamer &streamer) = delete; + FlatSparseStreamer &operator=(const FlatSparseStreamer &streamer) = delete; + + public: + //! Initialize Streamer + int init(const IndexMeta &, const ailego::Params &) override; + + //! Cleanup Streamer + int cleanup(void) override; + + //! Open index from file path + int open(IndexStorage::Pointer stg) override; + + //! Close file + int close(void) override; + + //! flush file + int flush(uint64_t checkpoint) override; + + //! Dump index into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Create a context + ContextPointer create_context(void) const override; + + //! Create a new iterator + IndexStreamer::SparseProvider::Pointer create_sparse_provider( + void) const override; + + int add_impl(uint64_t pkey, const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, Context::Pointer &context) override; + + int add_with_id_impl(uint32_t pkey, const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Similarity search with sparse inputs + int search_impl(const uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search with sparse inputs + int search_impl(const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Similarity brute force search with sparse inputs + int search_bf_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search with sparse inputs + int search_bf_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override; + + //! Linear search by primary keys with sparse inputs + int search_bf_by_p_keys_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + //! Fetch sparse vector by key + int get_sparse_vector(uint64_t key, uint32_t *sparse_count, + std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override; + + int get_sparse_vector_by_id( + uint32_t id, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override { + return get_sparse_vector(id, sparse_count, sparse_indices_buffer, + sparse_values_buffer); + } + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + const FlatSparseStreamerEntity &entity(void) const { + return entity_; + } + + uint32_t magic(void) const { + return magic_; + } + + private: + inline int check_params(const IndexQueryMeta &qmeta) const { + if (ailego_unlikely(qmeta.data_type() != meta_.data_type())) { + LOG_ERROR("Unsupported query meta, type=%d, expected=%d", + qmeta.data_type(), meta_.data_type()); + return IndexError_Mismatch; + } + return 0; + } + + int do_search(const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, bool with_p_keys, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 }; + + IndexMeta meta_{}; + Stats stats_{}; + FlatSparseStreamerEntity entity_; + + uint32_t magic_{0U}; + State state_{STATE_INIT}; + + //! avoid add vector while dumping index + ailego::SharedMutex shared_mutex_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.cc b/src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.cc new file mode 100644 index 0000000..1930d43 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.cc @@ -0,0 +1,1083 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flat_sparse_streamer_entity.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "flat_sparse_index_format.h" +#include "flat_sparse_utility.h" + +namespace zvec { +namespace core { + +FlatSparseStreamerEntity::FlatSparseStreamerEntity(IndexStreamer::Stats &stats) + : stats_(stats) {} + +int FlatSparseStreamerEntity::open(IndexStorage::Pointer storage, + const IndexMeta &meta) { + if (storage_) { + LOG_ERROR("An storage instance is already opened"); + return IndexError_Duplicate; + } + + keys_map_lock_ = std::make_shared(); + if (!keys_map_lock_) { + LOG_ERROR("FlatSparseStreamerEntity new object failed"); + return IndexError_NoMemory; + } + keys_map_ = std::make_shared>(); + + if (storage->get(PARAM_FLAT_SPARSE_META_SEG_ID) || + storage->get(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID)) { + int ret = this->load_storage(storage, meta); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to load storage index"); + return ret; + } + } else { + int ret = this->init_storage(storage, meta); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to load storage index"); + return ret; + } + } + + if (init_metric(meta) != 0) { + LOG_ERROR("Failed to init metric"); + return IndexError_InvalidFormat; + } + + // reserve data chunk + meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_MAX_DATA_CHUNK_CNT, + &max_data_chunk_cnt_); + sparse_data_chunks_.reserve(max_data_chunk_cnt_); + + // reserve offset chunk + meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_MAX_DOC_CNT, + &max_doc_cnt_); + sparse_offset_chunks_.reserve(max_doc_cnt_ / doc_cnt_per_offset_chunk() + 1); + sparse_unit_size_ = meta.unit_size(); + + LOG_DEBUG( + "FlatSparseStreamerEntity open success, doc_count[%u], " + "data_chunk_size[%u], offset_chunk_size[%u], data_chunk_count[%zu], " + "offset_chunk_count[%zu]", + meta_.doc_cnt, streamer_meta_.data_chunk_size, + streamer_meta_.offset_chunk_size, sparse_data_chunks_.size(), + sparse_offset_chunks_.size()); + + storage_ = storage; + return 0; +} + +int FlatSparseStreamerEntity::init_metric(const IndexMeta &meta) { + metric_ = IndexFactory::CreateMetric(meta.metric_name()); + if (!metric_) { + LOG_ERROR("Failed to create metric %s", meta.metric_name().c_str()); + return IndexError_NoExist; + } + int ret = metric_->init(meta, meta.metric_params()); + if (ret != 0) { + LOG_ERROR("Failled to init metric, ret=%d", ret); + return ret; + } + + if (!metric_->sparse_distance()) { + LOG_ERROR("Invalid metric distance"); + return IndexError_InvalidArgument; + } + + search_sparse_distance_ = metric_->sparse_distance(); + + if (metric_->query_metric() && metric_->query_metric()->distance()) { + search_sparse_distance_ = metric_->query_metric()->sparse_distance(); + } + + return 0; +} + +int FlatSparseStreamerEntity::load_storage(IndexStorage::Pointer storage, + const IndexMeta &meta) { + size_t index_size{0}; + + // load meta + auto segment = storage->get(PARAM_FLAT_SPARSE_META_SEG_ID); + + if (!segment || segment->data_size() < sizeof(meta_)) { + LOG_ERROR("Missing segment %s, or invalid segment size", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + IndexStorage::MemoryBlock data_block; + if (ailego_unlikely(segment->read(0, data_block, sizeof(meta_)) != + sizeof(meta_))) { + LOG_ERROR("Failed to read meta segment %s", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return IndexError_ReadData; + } + meta_ = *(reinterpret_cast(data_block.data())); + index_size += segment->capacity(); + + // load streamer meta + segment = storage->get(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID); + if (!segment || segment->data_size() < sizeof(streamer_meta_)) { + LOG_ERROR("Missing segment %s, or invalid segment size", + PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + if (ailego_unlikely(segment->read(0, data_block, sizeof(streamer_meta_)) != + sizeof(streamer_meta_))) { + LOG_ERROR("Failed to read streamer meta segment %s", + PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str()); + return IndexError_ReadData; + } + streamer_meta_ = + *(reinterpret_cast(data_block.data())); + index_size += segment->capacity(); + + uint32_t meta_data_chunk_size{streamer_meta_.data_chunk_size}; + uint32_t meta_offset_chunk_size{streamer_meta_.offset_chunk_size}; + meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_DATA_CHUNK_SIZE, + &meta_data_chunk_size); + meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_OFFSET_CHUNK_SIZE, + &meta_offset_chunk_size); + if (streamer_meta_.data_chunk_size != meta_data_chunk_size || + streamer_meta_.offset_chunk_size != meta_offset_chunk_size) { + LOG_ERROR( + "Invalid streamer meta chunk size data[%u] offset[%u], expect data[%u] " + "offset[%u]", + streamer_meta_.data_chunk_size, streamer_meta_.offset_chunk_size, + meta_data_chunk_size, meta_offset_chunk_size); + return IndexError_InvalidFormat; + } + + // check chunk cnt + if (streamer_meta_.data_chunk_count > max_data_chunk_cnt_ || + meta_.doc_cnt > max_doc_cnt_) { + LOG_ERROR( + "Invalid data chunk count[%u] doc count[%u], expect less than " + "chunk count[%u] doc count[%u]", + streamer_meta_.data_chunk_count, meta_.doc_cnt, max_data_chunk_cnt_, + max_doc_cnt_); + return IndexError_InvalidFormat; + } + + // load offset chunks + for (size_t i = 0; i < streamer_meta_.offset_chunk_count; ++i) { + std::string segment_id = + ailego::StringHelper::Concat(PARAM_FLAT_SPARSE_OFFSET_SEG_ID_PREFIX, i); + segment = storage->get(segment_id); + if (!segment) { + LOG_ERROR("Missing segment %s", segment_id.c_str()); + return IndexError_InvalidFormat; + } + sparse_offset_chunks_.emplace_back(segment); + index_size += segment->capacity(); + } + // load data chunks + for (size_t i = 0; i < streamer_meta_.data_chunk_count; ++i) { + std::string segment_id = + ailego::StringHelper::Concat(PARAM_FLAT_SPARSE_DATA_SEG_ID_PREFIX, i); + segment = storage->get(segment_id); + if (!segment) { + LOG_ERROR("Missing segment %s", segment_id.c_str()); + } + sparse_data_chunks_.emplace_back(segment); + index_size += segment->capacity(); + } + + // load keys + for (node_id_t i = 0; i < meta_.doc_cnt; ++i) { + (*keys_map_)[get_key(i)] = i; + } + + stats_.set_index_size(index_size); + stats_.set_check_point(storage->check_point()); + stats_.set_create_time(meta_.create_time); + stats_.set_update_time(meta_.update_time); + stats_.set_loaded_count(keys_map_->size()); + + return 0; +} + +int FlatSparseStreamerEntity::init_storage(IndexStorage::Pointer storage, + const IndexMeta &meta) { + meta_.create_time = ailego::Realtime::Seconds(); + stats_.set_create_time(meta_.create_time); + meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(meta_.update_time); + meta_.doc_cnt = 0; + + meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_DATA_CHUNK_SIZE, + &streamer_meta_.data_chunk_size); + meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_OFFSET_CHUNK_SIZE, + &streamer_meta_.offset_chunk_size); + + // append meta segment + size_t size = ailego_align(sizeof(meta_), ailego::MemoryHelper::PageSize()); + int ret = storage->append(PARAM_FLAT_SPARSE_META_SEG_ID, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to append meta segment %s", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return ret; + } + auto segment = storage->get(PARAM_FLAT_SPARSE_META_SEG_ID); + if (ailego_unlikely(!segment)) { + LOG_ERROR("Failed to get meta segment %s", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) { + LOG_ERROR("Failed to write meta segment %s", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return IndexError_WriteData; + } + + *stats_.mutable_index_size() += size; + + // append streamer meta segment + size = ailego_align(sizeof(streamer_meta_), ailego::MemoryHelper::PageSize()); + ret = storage->append(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to append streamer meta segment %s", + PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str()); + return ret; + } + segment = storage->get(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID); + if (ailego_unlikely(!segment)) { + LOG_ERROR("Failed to get streamer meta segment %s", + PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (segment->write(0, &streamer_meta_, sizeof(streamer_meta_)) != + sizeof(streamer_meta_)) { + LOG_ERROR("Failed to write streamer meta segment %s", + PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str()); + return IndexError_WriteData; + } + + *stats_.mutable_index_size() += size; + + return 0; +} + +int FlatSparseStreamerEntity::close() { + storage_.reset(); + sparse_data_chunks_.clear(); + sparse_offset_chunks_.clear(); + + keys_map_lock_.reset(); + keys_map_.reset(); + + return 0; +} + +int FlatSparseStreamerEntity::flush(uint64_t checkpoint) { + // flush meta + meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(meta_.update_time); + auto segment = storage_->get(PARAM_FLAT_SPARSE_META_SEG_ID); + if (ailego_unlikely(!segment)) { + LOG_ERROR("Failed to get meta segment %s", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) { + LOG_ERROR("Failed to write meta segment %s", + PARAM_FLAT_SPARSE_META_SEG_ID.c_str()); + return IndexError_WriteData; + } + + // flush streamer meta + streamer_meta_.data_chunk_count = sparse_data_chunks_.size(); + streamer_meta_.offset_chunk_count = sparse_offset_chunks_.size(); + segment = storage_->get(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID); + if (ailego_unlikely(!segment)) { + LOG_ERROR("Failed to get streamer meta segment %s", + PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str()); + return IndexError_Runtime; + } + if (segment->write(0, &streamer_meta_, sizeof(streamer_meta_)) != + sizeof(streamer_meta_)) { + LOG_ERROR("Failed to write streamer meta segment %s", + PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str()); + return IndexError_WriteData; + } + + if (checkpoint != 0) { + storage_->refresh(checkpoint); + } + int ret = storage_->flush(); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Failed to flush storage for %s", IndexError::What(ret)); + return ret; + } + if (checkpoint != 0) { + stats_.set_check_point(checkpoint); + } + + return 0; +} + +int FlatSparseStreamerEntity::dump(const IndexDumper::Pointer &dumper) { + ailego::ElapsedTime stamp; + + int ret; + // meta + ret = dump_meta(dumper.get()); + if (ret != 0) { + return ret; + } + + auto duration_dump_meta = stamp.milli_seconds(); + + // offset & data + ret = dump_offset_data(dumper.get()); + if (ret != 0) { + return ret; + } + + auto duration_dump_offset_data = stamp.milli_seconds() - duration_dump_meta; + + // keys + std::vector keys = get_keys(); + ret = dump_keys(keys, dumper.get()); + if (ret != 0) { + return ret; + } + + auto duration_dump_keys = + stamp.milli_seconds() - duration_dump_offset_data - duration_dump_meta; + + // mapping + ret = dump_mapping(keys, dumper.get()); + if (ret != 0) { + return ret; + } + + auto duration_dump_mapping = stamp.milli_seconds() - + duration_dump_offset_data - duration_dump_meta - + duration_dump_keys; + + LOG_INFO( + "Dump index meta: %zu ms, offset & data: %zu ms, keys: %zu ms, " + "mapping: %zu ms", + (size_t)duration_dump_meta, (size_t)duration_dump_offset_data, + (size_t)duration_dump_keys, (size_t)duration_dump_mapping); + + return 0; +} + +int FlatSparseStreamerEntity::dump_offset_data(IndexDumper *dumper) { + ailego::ElapsedTime stamp; + + uint64_t init_offset = dump_size_; + std::vector> offset_length; + + // write data + int ret; + node_id_t total_doc_cnt = doc_cnt(); + for (node_id_t node_id = 0; node_id < total_doc_cnt; node_id++) { + uint32_t target_vector_len; + IndexStorage::MemoryBlock target_vector_block; + ret = get_sparse_vector_ptr_by_id(node_id, target_vector_block, + &target_vector_len); + if (ret != 0) { + LOG_ERROR("Failed to get vector, node_id=%u, error: %s", node_id, + IndexError::What(ret)); + return ret; + } + const void *target_vector = target_vector_block.data(); + ret = dump_sparse_vector_data(target_vector, target_vector_len, dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump sparse vector data, node_id=%u, error: %s", + node_id, IndexError::What(ret)); + return ret; + } + + offset_length.push_back({dump_size_ - init_offset, target_vector_len}); + dump_size_ += target_vector_len; + } + + // append data segment + if (dumper->append(PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID, + dump_size_ - init_offset, 0, 0) != 0) { + LOG_ERROR("append data segment failed"); + return IndexError_WriteData; + } + + auto duration_dump_data = stamp.milli_seconds(); + + // write offset + for (auto &offset_length_pair : offset_length) { + if (dumper->write(&offset_length_pair.first, + sizeof(offset_length_pair.first)) != + sizeof(offset_length_pair.first)) { + return IndexError_WriteData; + } + if (dumper->write(&offset_length_pair.second, + sizeof(offset_length_pair.second)) != + sizeof(offset_length_pair.second)) { + return IndexError_WriteData; + } + dump_size_ += + sizeof(offset_length_pair.first) + sizeof(offset_length_pair.second); + } + + // append offset segment + if (dumper->append( + PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID, + offset_length.size() * (sizeof(uint64_t) + sizeof(uint32_t)), 0, + 0) != 0) { + LOG_ERROR("append offset segment failed"); + return IndexError_WriteData; + } + + auto duration_dump_offset = stamp.milli_seconds() - duration_dump_data; + + LOG_INFO("Dump offset: %zu ms, data: %zu ms", (size_t)duration_dump_offset, + (size_t)duration_dump_data); + + return 0; +} + +int FlatSparseStreamerEntity::dump_sparse_vector_data(const void *data, + uint32_t length, + IndexDumper *dumper) { + if (dumper->write(data, length) != length) { + return IndexError_WriteData; + } + return 0; +} + +int FlatSparseStreamerEntity::dump_meta(IndexDumper *dumper) { + if (dumper->write(&meta_, sizeof(meta_)) != sizeof(meta_)) { + LOG_ERROR("write meta failed"); + return IndexError_WriteData; + } + + size_t meta_padding_size = ailego_align(sizeof(meta_), 32) - sizeof(meta_); + if (meta_padding_size) { + std::string padding(meta_padding_size, '\0'); + if (dumper->write(padding.data(), meta_padding_size) != meta_padding_size) { + LOG_ERROR("write meta padding failed"); + return IndexError_WriteData; + } + } + return dumper->append(PARAM_FLAT_SPARSE_META_SEG_ID, sizeof(meta_), + meta_padding_size, 0); +} + +int FlatSparseStreamerEntity::dump_keys(const std::vector &keys, + IndexDumper *dumper) { + if (keys.size() == 1 && keys.back() == kInvalidKey) { + return IndexError_Runtime; + } + + size_t keys_size = keys.size() * sizeof(uint64_t); + if (dumper->write(keys.data(), keys_size) != keys_size) { + LOG_ERROR("Failed to write keys to dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + size_t keys_padding_size = ailego_align(keys_size, 32) - keys_size; + if (keys_padding_size) { + std::string padding(keys_padding_size, '\0'); + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write padding to dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + return dumper->append(PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID, keys_size, + keys_padding_size, 0); +} + +int FlatSparseStreamerEntity::dump_mapping(const std::vector &keys, + IndexDumper *dumper) { + std::vector mapping(keys.size()); + std::iota(mapping.begin(), mapping.end(), 0); + std::sort( + mapping.begin(), mapping.end(), + [&keys](uint32_t lhs, uint32_t rhs) { return (keys[lhs] < keys[rhs]); }); + + size_t mapping_size = mapping.size() * sizeof(uint32_t); + size_t mapping_padding_size = ailego_align(mapping_size, 32) - mapping_size; + if (dumper->write(mapping.data(), mapping_size) != mapping_size) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + + // Write the padding if need + if (mapping_padding_size) { + std::string padding(mapping_padding_size, '\0'); + if (dumper->write(padding.data(), padding.size()) != padding.size()) { + LOG_ERROR("Failed to write data into dumper %s", dumper->name().c_str()); + return IndexError_WriteData; + } + } + return dumper->append(PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID, mapping_size, + mapping_padding_size, 0); +} + +FlatSparseStreamerEntity::Pointer FlatSparseStreamerEntity::clone() const { + auto entity = new (std::nothrow) FlatSparseStreamerEntity( + stats_, meta_, streamer_meta_, keys_map_lock_, keys_map_, + sparse_data_chunks_, sparse_offset_chunks_); + return FlatSparseStreamerEntity::Pointer(entity); +} + +int FlatSparseStreamerEntity::add(uint64_t key, + const std::string &sparse_vector, + const uint32_t sparse_count) { + uint32_t sparse_vector_len = sparse_vector.size(); + + sparse_vector_len = AlignSize(sparse_vector_len); + + if (sparse_vector_len > streamer_meta_.data_chunk_size) { + LOG_ERROR( + "Sparse Vector Length exceed the chunk size, sparse vec len: %u, chunk " + "size: %u", + sparse_vector_len, streamer_meta_.data_chunk_size); + return IndexError_InvalidArgument; + } + + std::lock_guard lock(mutex_); + node_id_t local_id = doc_cnt(); + + if (ailego_unlikely(local_id >= max_doc_cnt_)) { + LOG_ERROR("Add vector failed for exceed max doc count: %u", max_doc_cnt_); + return IndexError_IndexFull; + } + + // duplicate check + if (ailego_unlikely(get_id(key) != kInvalidNodeId)) { + LOG_WARN("Try to add duplicate key, ignore it"); + return IndexError_Duplicate; + } + + // get sparse data chunk and offset for write sparse vector + Chunk::Pointer sparse_data_chunk; + uint32_t sparse_data_chunk_offset = -1U; + uint32_t sparse_data_chunk_index = sparse_data_chunks_.size() - 1U; + if (sparse_data_chunk_index == -1U || + sparse_data_chunks_[sparse_data_chunk_index]->data_size() + + sparse_vector_len > + streamer_meta_.data_chunk_size) { + if (ailego_unlikely(sparse_data_chunks_.capacity() == + sparse_data_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + if (sparse_data_chunk_index != -1U) { + LOG_ERROR( + "capacity: %zu, chunk used size: %zu, chunk size: %u, " + "sparse_vector_len: %u", + sparse_data_chunks_.capacity(), + sparse_data_chunks_[sparse_data_chunk_index]->data_size(), + streamer_meta_.data_chunk_size, sparse_vector_len); + } + return IndexError_IndexFull; + } + + sparse_data_chunk = alloc_new_data_chunk(sparse_data_chunks_.size()); + if (ailego_unlikely(!sparse_data_chunk)) { + LOG_ERROR("allocate data chunk failed"); + return IndexError_NoMemory; + } + sparse_data_chunks_.emplace_back(sparse_data_chunk); + sparse_data_chunk_index = sparse_data_chunks_.size() - 1U; + sparse_data_chunk_offset = 0UL; + } else { + sparse_data_chunk = sparse_data_chunks_[sparse_data_chunk_index]; + sparse_data_chunk_offset = sparse_data_chunk->data_size(); + } + + // write sparse vector + if (sparse_vector.size() > 0) { + if (ailego_unlikely(write_sparse_vector_data( + sparse_data_chunk_index, sparse_data_chunk_offset, + sparse_vector.data(), sparse_vector.size()) != 0)) { + LOG_ERROR("write sparse vector failed"); + return IndexError_NoMemory; + } + } + + uint64_t sparse_offset = sparse_data_chunk_index; + sparse_offset = (sparse_offset << 32U) + sparse_data_chunk_offset; + + // get sparse offset chunk and offset for write new info + Chunk::Pointer sparse_offset_chunk; + uint32_t sparse_offset_chunk_offset = -1U; + uint32_t sparse_offset_chunk_index = sparse_offset_chunks_.size() - 1U; + if (sparse_offset_chunk_index == -1U || + sparse_offset_chunks_[sparse_offset_chunk_index]->data_size() + + offset_size_per_node() > + streamer_meta_.offset_chunk_size) { + // no space left and need to allocate new offset chunk + if (ailego_unlikely(sparse_offset_chunks_.capacity() == + sparse_offset_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + + sparse_offset_chunk = alloc_new_offset_chunk(sparse_offset_chunks_.size()); + if (ailego_unlikely(!sparse_offset_chunk)) { + LOG_ERROR("allocate offset chunk failed"); + return IndexError_NoMemory; + } + sparse_offset_chunks_.emplace_back(sparse_offset_chunk); + sparse_offset_chunk_index = sparse_offset_chunks_.size() - 1U; + sparse_offset_chunk_offset = 0UL; + } else { + sparse_offset_chunk = sparse_offset_chunks_[sparse_offset_chunk_index]; + sparse_offset_chunk_offset = sparse_offset_chunk->data_size(); + } + + // write offset + size_t size = sparse_offset_chunk->write(sparse_offset_chunk_offset, + &sparse_offset, sizeof(uint64_t)); + if (ailego_unlikely(size != sizeof(uint64_t))) { + LOG_ERROR("Chunk write sparse vec offset failed, ret=%zu", size); + return IndexError_WriteData; + } + + // write length + size = + sparse_offset_chunk->write(sparse_offset_chunk_offset + sizeof(uint64_t), + &sparse_vector_len, sizeof(uint32_t)); + if (ailego_unlikely(size != sizeof(uint32_t))) { + LOG_ERROR("Chunk write sparse vec len failed, ret=%zu", size); + return IndexError_WriteData; + } + + // write key + size = sparse_offset_chunk->write( + sparse_offset_chunk_offset + 2 * sizeof(uint64_t), &key, + sizeof(uint64_t)); + if (ailego_unlikely(size != sizeof(uint64_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + + // LOG_INFO("Write sparse vector, key=%lu, offset chunk=%u, offset=%u, + // len=%u", + // key, sparse_offset_chunk_index, sparse_offset_chunk_offset, + // offset_size_per_node()); + + // LOG_INFO("Write sparse vector, key=%lu, data chunk=%u, offset=%u, len=%u", + // key, sparse_data_chunk_index, sparse_data_chunk_offset, + // sparse_vector_len); + + // resize chunk + if (sparse_vector_len > 0) { + sparse_data_chunk_offset += sparse_vector_len; + if (ailego_unlikely(sparse_data_chunk->resize(sparse_data_chunk_offset) != + sparse_data_chunk_offset)) { + LOG_ERROR("Sparse Chunk resize to %u failed", sparse_data_chunk_offset); + return IndexError_Runtime; + } + } + + // persist in keys_map + { + keys_map_lock_->lock(); + (*keys_map_)[key] = local_id; + keys_map_lock_->unlock(); + } + + inc_doc_count(); + inc_total_sparse_count(sparse_count); + + return 0; +} + +int FlatSparseStreamerEntity::add_vector_with_id( + uint32_t id, const std::string &sparse_vector, + const uint32_t sparse_count) { + uint32_t sparse_vector_len = sparse_vector.size(); + + sparse_vector_len = AlignSize(sparse_vector_len); + + if (sparse_vector_len > streamer_meta_.data_chunk_size) { + LOG_ERROR( + "Sparse Vector Length exceed the chunk size, sparse vec len: %u, chunk " + "size: %u", + sparse_vector_len, streamer_meta_.data_chunk_size); + return IndexError_InvalidArgument; + } + + std::lock_guard lock(mutex_); + + if (id >= doc_cnt()) { + for (auto i = doc_cnt(); i <= id; i++) { + node_id_t local_id = doc_cnt(); + if (ailego_unlikely(local_id >= max_doc_cnt_)) { + LOG_ERROR("Add vector failed for exceed max doc count: %u", + max_doc_cnt_); + return IndexError_IndexFull; + } + uint32_t sparse_data_chunk_index, sparse_data_chunk_offset, + sparse_offset_chunk_index, sparse_offset_chunk_offset; + if (i < id) { + write_sparse_vector_to_chunk("", 0, sparse_data_chunk_index, + sparse_data_chunk_offset); + } else { + write_sparse_vector_to_chunk(sparse_vector, sparse_vector_len, + sparse_data_chunk_index, + sparse_data_chunk_offset); + } + uint64_t sparse_offset = + ((uint64_t)sparse_data_chunk_index << 32U) + sparse_data_chunk_offset; + get_new_sparse_offset_chunk(sparse_offset_chunk_index, + sparse_offset_chunk_offset); + uint64_t written_key = kInvalidKey; + if (i == id) { + written_key = i; + } + write_sparse_offset_to_chunk(sparse_offset_chunk_index, + sparse_offset_chunk_offset, sparse_offset, + sparse_vector_len, written_key); + { + keys_map_lock_->lock(); + (*keys_map_)[i] = written_key; + keys_map_lock_->unlock(); + } + inc_doc_count(); + } + } else { + uint32_t sparse_data_chunk_index, sparse_data_chunk_offset; + write_sparse_vector_to_chunk(sparse_vector, sparse_vector_len, + sparse_data_chunk_index, + sparse_data_chunk_offset); + uint64_t sparse_offset = + ((uint64_t)sparse_data_chunk_index << 32U) + sparse_data_chunk_offset; + uint32_t sparse_offset_chunk_index = + id / get_offset_info_number_per_chunk(); + uint32_t sparse_offset_chunk_offset = + id % get_offset_info_number_per_chunk() * offset_size_per_node(); + write_sparse_offset_to_chunk(sparse_offset_chunk_index, + sparse_offset_chunk_offset, sparse_offset, + sparse_vector_len, id); + { + keys_map_lock_->lock(); + (*keys_map_)[id] = id; + keys_map_lock_->unlock(); + } + } + inc_total_sparse_count(sparse_count); + return 0; +} + +int FlatSparseStreamerEntity::write_sparse_vector_to_chunk( + const std::string &sparse_vector, const uint32_t sparse_vector_len, + uint32_t &sparse_data_chunk_index, uint32_t &sparse_data_chunk_offset) { + // get sparse data chunk and offset for write sparse vector + Chunk::Pointer sparse_data_chunk; + sparse_data_chunk_offset = -1U; + sparse_data_chunk_index = sparse_data_chunks_.size() - 1U; + if (sparse_data_chunk_index == -1U || + sparse_data_chunks_[sparse_data_chunk_index]->data_size() + + sparse_vector_len > + streamer_meta_.data_chunk_size) { + if (ailego_unlikely(sparse_data_chunks_.capacity() == + sparse_data_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + if (sparse_data_chunk_index != -1U) { + LOG_ERROR( + "capacity: %zu, chunk used size: %zu, chunk size: %u, " + "sparse_vector_len: %u", + sparse_data_chunks_.capacity(), + sparse_data_chunks_[sparse_data_chunk_index]->data_size(), + streamer_meta_.data_chunk_size, sparse_vector_len); + } + return IndexError_IndexFull; + } + + sparse_data_chunk = alloc_new_data_chunk(sparse_data_chunks_.size()); + if (ailego_unlikely(!sparse_data_chunk)) { + LOG_ERROR("allocate data chunk failed"); + return IndexError_NoMemory; + } + sparse_data_chunks_.emplace_back(sparse_data_chunk); + sparse_data_chunk_index = sparse_data_chunks_.size() - 1U; + sparse_data_chunk_offset = 0UL; + } else { + sparse_data_chunk = sparse_data_chunks_[sparse_data_chunk_index]; + sparse_data_chunk_offset = sparse_data_chunk->data_size(); + } + + // write sparse vector + if (sparse_vector.size() > 0) { + if (ailego_unlikely(write_sparse_vector_data( + sparse_data_chunk_index, sparse_data_chunk_offset, + sparse_vector.data(), sparse_vector.size()) != 0)) { + LOG_ERROR("write sparse vector failed"); + return IndexError_NoMemory; + } + } + + // resize chunk + if (sparse_vector_len > 0) { + uint32_t sparse_data_chunk_size = + sparse_data_chunk_offset + sparse_vector_len; + if (ailego_unlikely(sparse_data_chunk->resize(sparse_data_chunk_size) != + sparse_data_chunk_size)) { + LOG_ERROR("Sparse Chunk resize to %u failed", sparse_data_chunk_size); + return IndexError_Runtime; + } + } + return 0; +} + +int FlatSparseStreamerEntity::get_new_sparse_offset_chunk( + uint32_t &sparse_offset_chunk_index, uint32_t &sparse_offset_chunk_offset) { + // get sparse offset chunk and offset for write new info + Chunk::Pointer sparse_offset_chunk; + sparse_offset_chunk_offset = -1U; + sparse_offset_chunk_index = sparse_offset_chunks_.size() - 1U; + if (sparse_offset_chunk_index == -1U || + sparse_offset_chunks_[sparse_offset_chunk_index]->data_size() + + offset_size_per_node() > + streamer_meta_.offset_chunk_size) { + // no space left and need to allocate new offset chunk + if (ailego_unlikely(sparse_offset_chunks_.capacity() == + sparse_offset_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + + sparse_offset_chunk = alloc_new_offset_chunk(sparse_offset_chunks_.size()); + if (ailego_unlikely(!sparse_offset_chunk)) { + LOG_ERROR("allocate offset chunk failed"); + return IndexError_NoMemory; + } + sparse_offset_chunks_.emplace_back(sparse_offset_chunk); + sparse_offset_chunk_index = sparse_offset_chunks_.size() - 1U; + sparse_offset_chunk_offset = 0UL; + } else { + sparse_offset_chunk = sparse_offset_chunks_[sparse_offset_chunk_index]; + sparse_offset_chunk_offset = sparse_offset_chunk->data_size(); + } + return 0; +} + +int FlatSparseStreamerEntity::write_sparse_offset_to_chunk( + const uint32_t sparse_offset_chunk_index, + const uint32_t sparse_offset_chunk_offset, const uint64_t sparse_offset, + const uint32_t sparse_vector_len, const uint64_t node_id) { + // write offset + Chunk::Pointer sparse_offset_chunk = + sparse_offset_chunks_[sparse_offset_chunk_index]; + size_t size = sparse_offset_chunk->write(sparse_offset_chunk_offset, + &sparse_offset, sizeof(uint64_t)); + if (ailego_unlikely(size != sizeof(uint64_t))) { + LOG_ERROR("Chunk write sparse vec offset failed, ret=%zu", size); + return IndexError_WriteData; + } + + // write length + size = + sparse_offset_chunk->write(sparse_offset_chunk_offset + sizeof(uint64_t), + &sparse_vector_len, sizeof(uint32_t)); + if (ailego_unlikely(size != sizeof(uint32_t))) { + LOG_ERROR("Chunk write sparse vec len failed, ret=%zu", size); + return IndexError_WriteData; + } + + // write key + size = sparse_offset_chunk->write( + sparse_offset_chunk_offset + 2 * sizeof(uint64_t), &node_id, + sizeof(uint64_t)); + if (ailego_unlikely(size != sizeof(uint64_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + return 0; +} + +uint64_t FlatSparseStreamerEntity::get_key(node_id_t node_id) const { + uint32_t offset_chunk_index = node_id / get_offset_info_number_per_chunk(); + uint32_t offset_chunk_key_offset = + node_id % get_offset_info_number_per_chunk() * offset_size_per_node() + + 2 * sizeof(uint64_t); + + IndexStorage::MemoryBlock block; + if (ailego_unlikely(sparse_offset_chunks_[offset_chunk_index]->read( + offset_chunk_key_offset, block, sizeof(uint64_t)) != + sizeof(uint64_t))) { + LOG_ERROR("Read key failed, offset=%u, node_id=%u", offset_chunk_key_offset, + node_id); + return kInvalidKey; + }; + + return *reinterpret_cast(block.data()); +} + +int FlatSparseStreamerEntity::get_sparse_vector_ptr_by_id( + node_id_t node_id, const void **sparse_vector_ptr, + uint32_t *sparse_vector_len_ptr) const { + uint32_t offset_chunk_index = node_id / get_offset_info_number_per_chunk(); + uint32_t offset_chunk_offset = + node_id % get_offset_info_number_per_chunk() * offset_size_per_node(); + + // LOG_DEBUG("Read sparse vector, offset chunk=%u, offset=%u, len=%u", + // offset_chunk_index, offset_chunk_offset, offset_size_per_node()); + + auto offset_chunk = sparse_offset_chunks_[offset_chunk_index]; + + const void *offset_info = nullptr; + size_t read_len = offset_chunk->read(offset_chunk_offset, &offset_info, + offset_size_per_node()); + if (ailego_unlikely(read_len != offset_size_per_node())) { + LOG_ERROR("Read offset info failed, offset=%u, read_len=%zu, expect=%u", + offset_chunk_offset, read_len, offset_size_per_node()); + return IndexError_ReadData; + }; + + // sparse offset + uint64_t sparse_offset = *(uint64_t *)offset_info; + uint32_t sparse_vector_len = + *(uint32_t *)((uint8_t *)offset_info + sizeof(uint64_t)); + + uint32_t sparse_data_chunk_index = + static_cast((sparse_offset >> 32) & 0xFFFFFFFF); + uint32_t sparse_data_chunk_offset = + static_cast(sparse_offset & 0xFFFFFFFF); + + if (sparse_vector_len > 0) { + const void *sparse_data = get_sparse_vector_data( + sparse_data_chunk_index, sparse_data_chunk_offset, sparse_vector_len); + if (ailego_unlikely(sparse_data == nullptr)) { + LOG_ERROR("Get nullptr sparse, offset=%zu, len=%u", (size_t)sparse_offset, + sparse_vector_len); + + return IndexError_ReadData; + } + *sparse_vector_ptr = sparse_data; + *sparse_vector_len_ptr = sparse_vector_len; + } + + // LOG_DEBUG("Read sparse vector, data chunk=%u, offset=%u, len=%u", + // sparse_data_chunk_index, sparse_data_chunk_offset, + // sparse_vector_len); + + return 0; +} + +int FlatSparseStreamerEntity::get_sparse_vector_ptr_by_id( + node_id_t node_id, IndexStorage::MemoryBlock &sparse_vector_block, + uint32_t *sparse_vector_len_ptr) const { + uint32_t offset_chunk_index = node_id / get_offset_info_number_per_chunk(); + uint32_t offset_chunk_offset = + node_id % get_offset_info_number_per_chunk() * offset_size_per_node(); + + // LOG_DEBUG("Read sparse vector, offset chunk=%u, offset=%u, len=%u", + // offset_chunk_index, offset_chunk_offset, offset_size_per_node()); + + auto offset_chunk = sparse_offset_chunks_[offset_chunk_index]; + + const void *offset_info = nullptr; + IndexStorage::MemoryBlock offset_info_block; + size_t read_len = offset_chunk->read(offset_chunk_offset, offset_info_block, + offset_size_per_node()); + if (ailego_unlikely(read_len != offset_size_per_node())) { + LOG_ERROR("Read offset info failed, offset=%u, read_len=%zu, expect=%u", + offset_chunk_offset, read_len, offset_size_per_node()); + return IndexError_ReadData; + }; + offset_info = offset_info_block.data(); + + // sparse offset + uint64_t sparse_offset = *(uint64_t *)offset_info; + uint32_t sparse_vector_len = + *(uint32_t *)((uint8_t *)offset_info + sizeof(uint64_t)); + + uint32_t sparse_data_chunk_index = + static_cast((sparse_offset >> 32) & 0xFFFFFFFF); + uint32_t sparse_data_chunk_offset = + static_cast(sparse_offset & 0xFFFFFFFF); + + if (sparse_vector_len > 0) { + get_sparse_vector_data(sparse_data_chunk_index, sparse_data_chunk_offset, + sparse_vector_len, sparse_vector_block); + if (ailego_unlikely(sparse_vector_block.data() == nullptr)) { + LOG_ERROR("Get nullptr sparse, offset=%zu, len=%u", (size_t)sparse_offset, + sparse_vector_len); + + return IndexError_ReadData; + } + *sparse_vector_len_ptr = sparse_vector_len; + } + + return 0; +} + +int FlatSparseStreamerEntity::write_sparse_vector_data(uint32_t chunk_index, + uint64_t offset, + const void *data, + uint32_t length) { + auto size = sparse_data_chunks_[chunk_index]->write(offset, data, length); + if (size != length) { + LOG_ERROR( + "write sparse vector data failed: chunk_index=%u, offset=%zu, " + "length=%u, size=%zu, chunk_data_size=%zu", + chunk_index, (size_t)offset, length, size, + sparse_data_chunks_[chunk_index]->data_size()); + return IndexError_WriteData; + } + // LOG_DEBUG( + // "write_sparse_vector_data: chunk_index=%u, offset=%lu, length=%u, " + // "data=%p", + // chunk_index, offset, length, data); + return 0; +} + +const void *FlatSparseStreamerEntity::get_sparse_vector_data( + uint32_t chunk_index, uint64_t offset, uint32_t length) const { + const void *data; + auto size = sparse_data_chunks_[chunk_index]->read(offset, &data, length); + if (size != length) { + LOG_ERROR( + "read sparse vector data failed: chunk_index=%u, offset=%zu, " + "length=%u, size=%zu", + chunk_index, (size_t)offset, length, size); + return nullptr; + } + // LOG_DEBUG( + // "get_sparse_vector_data: chunk_index=%u, offset=%lu, length=%u, " + // "data=%p", + // chunk_index, offset, length, data); + return data; +} + +int FlatSparseStreamerEntity::get_sparse_vector_data( + uint32_t chunk_index, uint64_t offset, uint32_t length, + IndexStorage::MemoryBlock &block) const { + auto size = sparse_data_chunks_[chunk_index]->read(offset, block, length); + if (size != length) { + LOG_ERROR( + "read sparse vector data failed: chunk_index=%u, offset=%zu, " + "length=%u, size=%zu", + chunk_index, (size_t)offset, length, size); + return IndexError_ReadData; + } + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.h b/src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.h new file mode 100644 index 0000000..6ef2f9f --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.h @@ -0,0 +1,265 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "flat_sparse_entity.h" +#include "flat_sparse_index_format.h" +#include "flat_sparse_utility.h" + +namespace zvec { +namespace core { + + +/*! Flat Sparse Streamer Entity + */ +class FlatSparseStreamerEntity : public FlatSparseEntity { + public: + typedef std::shared_ptr Pointer; + + using Chunk = IndexStorage::Segment; + + //! Constructor + explicit FlatSparseStreamerEntity(IndexStreamer::Stats &stats); + + //! Destructor + virtual ~FlatSparseStreamerEntity() = default; + + //! Disable them + FlatSparseStreamerEntity(const FlatSparseStreamerEntity &) = delete; + FlatSparseStreamerEntity &operator=(const FlatSparseStreamerEntity &) = + delete; + + //! Open the entity with storage + int open(IndexStorage::Pointer storage, const IndexMeta &meta); + + //! Close the entity + int close(); + + //! Flush linear index to storage + int flush(uint64_t checkpoint); + + //! Dump index by dumper + int dump(const IndexDumper::Pointer &dumper); + + //! Add sparse vector to linear index + int add(uint64_t key, const std::string &sparse_vector, + const uint32_t sparse_count); + + //! Add sparse vector to linear index with id + int add_vector_with_id(uint32_t id, const std::string &sparse_vector, + uint32_t sparse_count); + + //! Clone entity + FlatSparseStreamerEntity::Pointer clone() const; + + int get_index_sparse_meta(IndexMeta *meta) const { + return IndexHelper::DeserializeFromStorage(storage_.get(), meta); + } + + int set_index_sparse_meta(const IndexMeta &meta) const { + return IndexHelper::SerializeToStorage(meta, storage_.get()); + } + + public: + inline uint32_t doc_cnt() const override { + return meta_.doc_cnt; + } + + inline uint32_t total_sparse_count() const override { + return meta_.total_sparse_count; + } + + size_t sparse_unit_size() const override { + return sparse_unit_size_; + } + + inline node_id_t get_id(uint64_t key) const override { + keys_map_lock_->lock_shared(); + auto it = keys_map_->find(key); + keys_map_lock_->unlock_shared(); + return it == keys_map_->end() ? kInvalidNodeId : it->second; + } + + uint64_t get_key(node_id_t node_id) const override; + + int get_sparse_vector_ptr_by_id(node_id_t id, const void **sparse_vector, + uint32_t *sparse_vector_len) const override; + int get_sparse_vector_ptr_by_id( + const node_id_t id, IndexStorage::MemoryBlock &sparse_vector_block, + uint32_t *sparse_vector_len) const override; + + float get_search_distance(const std::string &vector, + node_id_t target_node_id) const override { + float dist; + const void *target_vector; + uint32_t target_vector_len; + get_sparse_vector_ptr_by_id(target_node_id, &target_vector, + &target_vector_len); + search_sparse_distance_(vector.c_str(), target_vector, &dist); + return dist; + } + + private: + void inc_doc_count() { + meta_.doc_cnt++; + } + void inc_total_sparse_count(uint32_t count) { + meta_.total_sparse_count += count; + } + + int init_metric(const IndexMeta &meta); + + int init_storage(IndexStorage::Pointer storage, const IndexMeta &meta); + + int load_storage(IndexStorage::Pointer storage, const IndexMeta &meta); + + static inline size_t AlignSize(size_t size) { + return (size + 0x1F) & (~0x1F); + } + + inline uint32_t offset_size_per_node() const { + return 3 * sizeof(uint64_t); + } + + inline uint32_t doc_cnt_per_offset_chunk() const { + return streamer_meta_.offset_chunk_size / offset_size_per_node(); + } + + Chunk::Pointer alloc_new_offset_chunk(uint32_t chunk_id) { + std::string segment_id = ailego::StringHelper::Concat( + PARAM_FLAT_SPARSE_OFFSET_SEG_ID_PREFIX, chunk_id); + // LOG_INFO("Alloc new offset chunk %s", segment_id.c_str()); + return alloc_new_chunk(segment_id, streamer_meta_.offset_chunk_size); + } + + Chunk::Pointer alloc_new_data_chunk(uint32_t chunk_id) { + std::string segment_id = ailego::StringHelper::Concat( + PARAM_FLAT_SPARSE_DATA_SEG_ID_PREFIX, chunk_id); + // LOG_INFO("Alloc new data chunk %s", segment_id.c_str()); + return alloc_new_chunk(segment_id, streamer_meta_.data_chunk_size); + } + + Chunk::Pointer alloc_new_chunk(const std::string &segment_id, uint32_t size) { + int ret = storage_->append(segment_id, size); + if (ailego_unlikely(ret != 0)) { + return nullptr; + } + *stats_.mutable_index_size() += size; + return storage_->get(segment_id); + } + + inline uint32_t get_offset_info_number_per_chunk() const { + return streamer_meta_.offset_chunk_size / offset_size_per_node(); + } + + int write_sparse_vector_to_chunk(const std::string &sparse_vector, + const uint32_t sparse_vector_len, + uint32_t &sparse_data_chunk_index, + uint32_t &sparse_data_chunk_offset); + + int get_new_sparse_offset_chunk(uint32_t &sparse_offset_chunk_index, + uint32_t &sparse_offset_chunk_offset); + + int write_sparse_offset_to_chunk(const uint32_t sparse_offset_chunk_index, + const uint32_t sparse_offset_chunk_offset, + const uint64_t sparse_offset, + const uint32_t sparse_vector_len, + const uint64_t node_id); + + int write_sparse_vector_data(uint32_t chunk_index, uint64_t offset, + const void *data, uint32_t length); + + const void *get_sparse_vector_data(uint32_t chunk_index, uint64_t offset, + uint32_t length) const; + + int get_sparse_vector_data(uint32_t chunk_index, uint64_t offset, + uint32_t length, + IndexStorage::MemoryBlock &block) const; + + int dump_sparse_vector_data(const void *data, uint32_t length, + IndexDumper *dumper); + + int dump_meta(IndexDumper *dumper); + + int dump_index_meta(IndexDumper *dumper); + + int dump_keys(const std::vector &keys, IndexDumper *dumper); + + int dump_mapping(const std::vector &keys, IndexDumper *dumper); + + int dump_offset_data(IndexDumper *dumper); + + + private: + FlatSparseStreamerEntity( + IndexStreamer::Stats &stats, const FlatSparseMeta &meta, + const FlatSparseStreamerMeta &streamer_meta, + std::shared_ptr keys_map_lock, + std::shared_ptr> keys_map, + std::vector sparse_data_chunks, + std::vector sparse_offset_chunks) + : stats_(stats), + meta_(meta), + streamer_meta_(streamer_meta), + keys_map_lock_(keys_map_lock), + keys_map_(keys_map), + sparse_data_chunks_(std::move(sparse_data_chunks)), + sparse_offset_chunks_(std::move(sparse_offset_chunks)) {} + + private: + IndexStorage::Pointer storage_{}; + IndexStreamer::Stats &stats_; + + // meta + FlatSparseMeta meta_; + FlatSparseStreamerMeta streamer_meta_; + + // metric + IndexMetric::Pointer metric_{}; + IndexMetric::MatrixSparseDistance search_sparse_distance_{}; + + std::mutex mutex_{}; + + // keys map + mutable std::shared_ptr keys_map_lock_{}; + std::shared_ptr> keys_map_{}; + + // chunks + mutable std::vector sparse_data_chunks_{}; + mutable std::vector sparse_offset_chunks_{}; + + // config + uint32_t max_doc_cnt_{1 << 24U}; // 16 million + uint32_t max_data_chunk_cnt_{ + 1 << 10U}; // 1024, default single_data_chunk_size = 8M, + // default_total_max = 1024 * 8M = 8G + + uint64_t dump_size_{0U}; + size_t sparse_unit_size_{0U}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/flat_sparse/flat_sparse_utility.h b/src/core/algorithm/flat_sparse/flat_sparse_utility.h new file mode 100644 index 0000000..a66e2f6 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_utility.h @@ -0,0 +1,58 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec { +namespace core { + +static constexpr uint32_t PARAM_FLAT_SPARSE_MAX_DIM_SIZE = 16384; +static const std::string PARAM_FLAT_SPARSE_META_SEG_ID = + "bruteforce_sparse_meta"; + +// streamer +static const std::string PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID = + "bruteforce_sparse_streamer_meta"; +static const std::string PARAM_FLAT_SPARSE_OFFSET_SEG_ID_PREFIX = + "bruteforce_sparse_streamer_offset_"; +static const std::string PARAM_FLAT_SPARSE_DATA_SEG_ID_PREFIX = + "bruteforce_sparse_streamer_data_"; + +// searcher +static const std::string PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID = + "bruteforce_sparse_searcher_offset_segment"; +static const std::string PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID = + "bruteforce_sparse_searcher_data_segment"; +static const std::string PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID = + "bruteforce_sparse_searcher_keys_segment"; +static const std::string PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID = + "bruteforce_sparse_searcher_mapping_segment"; + +// streamer +static const std::string PARAM_FLAT_SPARSE_STREAMER_OFFSET_CHUNK_SIZE( + "proxima.bruteforce.sparse_streamer.offset_chunk_size"); + +static const std::string PARAM_FLAT_SPARSE_STREAMER_DATA_CHUNK_SIZE( + "proxima.bruteforce.sparse_streamer.data_chunk_size"); + +static const std::string PARAM_FLAT_SPARSE_STREAMER_MAX_DOC_CNT( + "proxima.bruteforce.sparse_streamer.max_doc_cnt"); + +static const std::string PARAM_FLAT_SPARSE_STREAMER_MAX_DATA_CHUNK_CNT( + "proxima.bruteforce.sparse_streamer.max_data_chunk_cnt"); + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw/CMakeLists.txt b/src/core/algorithm/hnsw/CMakeLists.txt new file mode 100644 index 0000000..d59dc8f --- /dev/null +++ b/src/core/algorithm/hnsw/CMakeLists.txt @@ -0,0 +1,17 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(NOT APPLE) + set(CORE_KNN_HNSW_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_hnsw + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework core_utility sparsehash + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_HNSW_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.cc b/src/core/algorithm/hnsw/hnsw_algorithm.cc new file mode 100644 index 0000000..1fd6f65 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_algorithm.cc @@ -0,0 +1,699 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_algorithm.h" +#include + +namespace zvec { +namespace core { + +template +int HnswAlgorithm::add_node(node_id_t id, level_t level, + HnswContext *ctx) { + spin_lock_.lock(); + + auto cur_max_level = entity_.cur_max_level(); + auto entry_point = entity_.entry_point(); + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + entity_.update_ep_and_level(id, level); + spin_lock_.unlock(); + return 0; + } + spin_lock_.unlock(); + + if (ailego_unlikely(level > cur_max_level)) { + mutex_.lock(); + // re-check max level + cur_max_level = entity_.cur_max_level(); + entry_point = entity_.entry_point(); + if (level <= cur_max_level) { + mutex_.unlock(); + } + } + + level_t cur_level = cur_max_level; + dist_t dist = ctx->dist_calculator().batch_dist(entry_point); + for (; cur_level > level; --cur_level) { + select_entry_point(cur_level, &entry_point, &dist, ctx); + } + + for (; cur_level >= 0; --cur_level) { + search_neighbors(cur_level, &entry_point, &dist, ctx->level_topk(cur_level), + ctx, /*use_pool=*/false); + } + + // add neighbors from down level to top level, to avoid upper level visible + // to knn_search but the under layer level not ready + for (cur_level = 0; cur_level <= level; ++cur_level) { + add_neighbors(id, cur_level, ctx->level_topk(cur_level), ctx); + ctx->level_topk(cur_level).clear(); + } + + if (ailego_unlikely(level > cur_max_level)) { + spin_lock_.lock(); + entity_.update_ep_and_level(id, level); + spin_lock_.unlock(); + mutex_.unlock(); + } + + return 0; +} + +template +int HnswAlgorithm::search(HnswContext *ctx) const { + spin_lock_.lock(); + auto maxLevel = entity_.cur_max_level(); + auto entry_point = entity_.entry_point(); + spin_lock_.unlock(); + + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + return 0; + } + + dist_t dist = ctx->dist_calculator().dist(entry_point); + for (level_t cur_level = maxLevel; cur_level >= 1; --cur_level) { + select_entry_point(cur_level, &entry_point, &dist, ctx); + } + + auto &topk_heap = ctx->topk_heap(); + topk_heap.clear(); + search_neighbors(0, &entry_point, &dist, topk_heap, ctx, /*use_pool=*/true); + + if (ctx->group_by_search()) { + expand_neighbors_by_group(topk_heap, ctx); + } + + return 0; +} + +template +void HnswAlgorithm::select_entry_point(level_t level, + node_id_t *entry_point, + dist_t *dist, + HnswContext *ctx) const { + const auto &entity = static_cast(ctx->get_entity()); + HnswDistCalculator &dc = ctx->dist_calculator(); + while (true) { + const auto neighbors = entity.get_neighbors_typed(level, *entry_point); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + uint32_t size = neighbors.size(); + if (size == 0) { + break; + } + + std::vector neighbor_vec_blocks; + int ret = entity.get_vector_typed(&neighbors[0], size, neighbor_vec_blocks); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } + if (ailego_unlikely(ret != 0)) { + break; + } + + bool find_closer = false; + + std::vector dists(size); + std::vector neighbor_vecs(size); + for (uint32_t i = 0; i < size; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + + dc.batch_dist(neighbor_vecs.data(), size, dists.data()); + + for (uint32_t i = 0; i < size; ++i) { + dist_t cur_dist = dists[i]; + + if (cur_dist < *dist) { + *entry_point = neighbors[i]; + *dist = cur_dist; + find_closer = true; + } + } + + if (!find_closer) { + break; + } + } + + return; +} + +template +void HnswAlgorithm::add_neighbors(node_id_t id, level_t level, + TopkHeap &topk_heap, + HnswContext *ctx) { + if (ailego_unlikely(topk_heap.size() == 0)) { + return; + } + + HnswDistCalculator &dc = ctx->dist_calculator(); + + update_neighbors(dc, id, level, topk_heap); + + // reverse update neighbors + for (size_t i = 0; i < topk_heap.size(); ++i) { + reverse_update_neighbors(dc, topk_heap[i].first, level, id, + topk_heap[i].second, ctx->update_heap()); + } + + return; +} + +// ============================================================================ +// search_neighbors helper templates +// +// Two specialized inner loops, dispatched from search_neighbors(): +// +// fast_search_neighbors: mmap/contiguous with direct vector pointers. +// Uses BlockHeap (AVX2) or LinearPool (scalar) +// for visited tracking and top-k maintenance. +// dual_heap_search_neighbors: CandidateHeap + TopkHeap + VisitFilter. +// Used for add_node (use_pool=false), filtered +// search, upper levels, and BufferPool fallback. +// ============================================================================ + +// mmap/contiguous variant: resolve vectors via get_vector_ptr and use +// LinearPool or BlockHeap for visited tracking + top-k maintenance. +// HeapType must expose reset/set_visited/check_visited/push_block/has_next/pop. +template +void fast_search_neighbors(const EntityType &entity, HeapType &pool, + VisitFilter &visit, HnswDistCalculator &dc, + uint32_t topk, uint32_t ef, node_id_t entry_point, + dist_t entry_dist, uint32_t prefetch_lines, + uint32_t prefetch_offset) { + const uint32_t max_deg = entity.max_degree(0); // level 0 only + const uint32_t cap = std::max(topk, ef); + pool.reset(static_cast(cap), static_cast(max_deg)); + visit.clear(); + + visit.set_visited(entry_point); + pool.push_block(&entry_dist, &entry_point, 1); + + uint32_t buf_capacity = max_deg; + std::vector neighbor_ids(buf_capacity); + std::vector dists(buf_capacity); + std::vector neighbor_vecs(buf_capacity); + + while (pool.has_next()) { + auto current_node = pool.pop(); + + const auto neighbors = entity.get_neighbors_typed(0, current_node); + ailego_prefetch(neighbors.data); + + if (neighbors.size() > buf_capacity) { + buf_capacity = neighbors.size(); + neighbor_ids.resize(buf_capacity); + dists.resize(buf_capacity); + neighbor_vecs.resize(buf_capacity); + } + + const uint32_t po = + std::min(static_cast(neighbors.size()), prefetch_offset); + uint32_t unvisited_count = 0; + uint32_t i = 0; + + // Phase 1: scan first `po` neighbors with prefetch. + for (; i < po; ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) continue; + visit.set_visited(node); + const void *vec_ptr = entity.get_vector_ptr(node); + const char *p = reinterpret_cast(vec_ptr); + for (uint32_t cl = 0; cl < prefetch_lines; ++cl) { + ailego_prefetch(p + cl * 64); + } + neighbor_ids[unvisited_count] = node; + neighbor_vecs[unvisited_count] = vec_ptr; + unvisited_count++; + } + + // Phase 2: scan remaining neighbors. + for (; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) continue; + visit.set_visited(node); + neighbor_ids[unvisited_count] = node; + neighbor_vecs[unvisited_count] = entity.get_vector_ptr(node); + unvisited_count++; + } + + if (unvisited_count == 0) continue; + dc.batch_dist(neighbor_vecs.data(), unvisited_count, dists.data()); + + pool.push_block(dists.data(), neighbor_ids.data(), + static_cast(unvisited_count)); + } +} + +// ============================================================================ +// dual_heap_search_neighbors: shared core for the fallback dual-heap path. +// +// Maintains a candidate min-heap + topk heap + VisitFilter. Supports +// arbitrary levels, filters, and MemoryBlock types (BufferPool/Mmap). +// Also updates entry_point/dist for next-level continuation. +// ============================================================================ +template +void dual_heap_search_neighbors(const EntityType &entity, level_t level, + node_id_t *entry_point, dist_t *dist, + TopkHeap &topk, HnswContext *ctx, + HnswDistCalculator &dc, FilterFn &&filter) { + const uint32_t prefetch_offset = ctx->po(); + const uint32_t prefetch_lines = + ctx->pl() > 0 ? ctx->pl() : (entity.vector_size() + 63) / 64; + + uint32_t buf_capacity = entity.max_degree(level); + std::vector neighbor_ids(buf_capacity); + std::vector neighbor_vec_blocks; + neighbor_vec_blocks.reserve(buf_capacity); + std::vector dists(buf_capacity); + std::vector neighbor_vecs(buf_capacity); + + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + + candidates.clear(); + visit.clear(); + visit.set_visited(*entry_point); + if (!filter(*entry_point)) { + topk.emplace(*entry_point, *dist); + } + + candidates.emplace(*entry_point, *dist); + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t main_node = top->first; + dist_t main_dist = top->second; + + if (topk.full() && main_dist > topk[0].second) { + break; + } + + candidates.pop(); + const auto neighbors = entity.get_neighbors_typed(level, main_node); + ailego_prefetch(neighbors.data); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + + if (neighbors.size() > buf_capacity) { + buf_capacity = neighbors.size(); + neighbor_ids.resize(buf_capacity); + neighbor_vec_blocks.resize(buf_capacity); + dists.resize(buf_capacity); + neighbor_vecs.resize(buf_capacity); + } + + uint32_t size = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) { + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_visit_dup_cnt())++; + } + continue; + } + visit.set_visited(node); + neighbor_ids[size++] = node; + } + if (size == 0) { + continue; + } + + neighbor_vec_blocks.clear(); + int ret = + entity.get_vector_typed(neighbor_ids.data(), size, neighbor_vec_blocks); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } + if (ailego_unlikely(ret != 0)) { + break; + } + + // do prefetch + for (uint32_t i = 0; i < std::min(prefetch_offset, size); ++i) { + const char *base = + static_cast(neighbor_vec_blocks[i].data()); + for (uint32_t cl = 0; cl < prefetch_lines; ++cl) { + ailego_prefetch(base + cl * 64); + } + } + + for (uint32_t i = 0; i < size; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + + dc.batch_dist(neighbor_vecs.data(), size, dists.data()); + + for (uint32_t i = 0; i < size; ++i) { + node_id_t node = neighbor_ids[i]; + dist_t cur_dist = dists[i]; + + if ((!topk.full()) || cur_dist < topk[0].second) { + candidates.emplace(node, cur_dist); + // update entry_point for next level scan + if (cur_dist < *dist) { + *entry_point = node; + *dist = cur_dist; + } + if (!filter(node)) { + topk.emplace(node, cur_dist); + } + } + } + } +} + +// ============================================================================ +// search_neighbors: Dispatch to fast or dual-heap path. +// +// - add_node / filtered / upper levels → dual_heap_search_neighbors +// - level-0 unfiltered search: +// MmapMemoryBlock → fast_search_neighbors (BlockHeap/LinearPool) +// BufferPool → dual_heap_search_neighbors (fallback) +// ============================================================================ +template +void HnswAlgorithm::search_neighbors(level_t level, + node_id_t *entry_point, + dist_t *dist, TopkHeap &topk, + HnswContext *ctx, + bool use_pool) const { + const auto &entity = static_cast(ctx->get_entity()); + HnswDistCalculator &dc = ctx->dist_calculator(); + + if (!use_pool || ctx->filter().is_valid() || level != 0) { + // Dual-heap path: add_node, filtered search, or upper-level scan. + auto run_with_filter = [&](auto &&filter) { + dual_heap_search_neighbors( + entity, level, entry_point, dist, topk, ctx, dc, + std::forward(filter)); + }; + + if (ctx->filter().is_valid()) { + auto filter = [&](node_id_t id) { + return ctx->filter()(entity.get_key_typed(id)); + }; + run_with_filter(filter); + } else { + auto filter = [](node_id_t) { return false; }; + run_with_filter(filter); + } + } else { + // Pool-based path for level-0 unfiltered search. + if constexpr (std::is_same_v) { + const uint32_t prefetch_lines = + ctx->pl() > 0 ? ctx->pl() : (entity.vector_size() + 63) / 64; + + // Fast path: direct pointer access via get_vector_ptr. + // BlockHeap (AVX2) or LinearPool (scalar) for top-k tracking. + const uint32_t topk_v = static_cast(ctx->topk()); + const uint32_t ef_v = ctx->ef(); + const bool avx2_ok = + zvec::ailego::internal::CpuFeatures::static_flags_.AVX2; + + auto &visit = ctx->visit_filter(); + + if (avx2_ok) { + auto &bpool = ctx->block_pool(); + fast_search_neighbors(entity, bpool, visit, dc, topk_v, ef_v, + *entry_point, *dist, prefetch_lines, ctx->po()); + copy_pool_to_topk(bpool, topk); + } else { + auto &lpool = ctx->pool(); + fast_search_neighbors(entity, lpool, visit, dc, topk_v, ef_v, + *entry_point, *dist, prefetch_lines, ctx->po()); + copy_pool_to_topk(lpool, topk); + } + } else { + // BufferPool entities: fallback to dual-heap path. + auto filter = [](node_id_t) { return false; }; + dual_heap_search_neighbors( + entity, level, entry_point, dist, topk, ctx, dc, filter); + } + } +} + +template +void HnswAlgorithm::expand_neighbors_by_group( + TopkHeap &topk, HnswContext *ctx) const { + if (!ctx->group_by().is_valid()) { + return; + } + + const auto &entity = static_cast(ctx->get_entity()); + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity.get_key_typed(id)); + }; + + // devide into groups + std::map &group_topk_heaps = ctx->group_topk_heaps(); + for (uint32_t i = 0; i < topk.size(); ++i) { + node_id_t id = topk[i].first; + auto score = topk[i].second; + + std::string group_id = group_by(id); + + auto &topk_heap = group_topk_heaps[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, score); + } + + // stage 2, expand to reach group num as possible + if (group_topk_heaps.size() < ctx->group_num()) { + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + HnswDistCalculator &dc = ctx->dist_calculator(); + + std::function filter = [](node_id_t) { return false; }; + if (ctx->filter().is_valid()) { + filter = [&](node_id_t id) { + return ctx->filter()(entity.get_key_typed(id)); + }; + } + + // refill to get enough groups + candidates.clear(); + visit.clear(); + for (uint32_t i = 0; i < topk.size(); ++i) { + node_id_t id = topk[i].first; + float score = topk[i].second; + + visit.set_visited(id); + candidates.emplace_back(id, score); + } + + // do expand + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t main_node = top->first; + + candidates.pop(); + const auto neighbors = entity.get_neighbors_typed(0, main_node); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + + std::vector neighbor_ids(neighbors.size()); + uint32_t size = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) { + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_visit_dup_cnt())++; + } + continue; + } + visit.set_visited(node); + neighbor_ids[size++] = node; + } + if (size == 0) { + continue; + } + + std::vector neighbor_vec_blocks; + int ret = entity.get_vector_typed(neighbor_ids.data(), size, + neighbor_vec_blocks); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } + if (ailego_unlikely(ret != 0)) { + break; + } + + std::vector dists(size); + std::vector neighbor_vecs(size); + for (uint32_t i = 0; i < size; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + dc.batch_dist(neighbor_vecs.data(), size, dists.data()); + + for (uint32_t i = 0; i < size; ++i) { + node_id_t node = neighbor_ids[i]; + dist_t cur_dist = dists[i]; + + if (!filter(node)) { + std::string group_id = group_by(node); + + auto &topk_heap = group_topk_heaps[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(node, cur_dist); + + if (group_topk_heaps.size() >= ctx->group_num()) { + break; + } + } + + candidates.emplace(node, cur_dist); + } + } + } +} + +template +void HnswAlgorithm::update_neighbors(HnswDistCalculator &dc, + node_id_t id, level_t level, + TopkHeap &topk_heap) { + topk_heap.sort(); + + uint32_t max_neighbor_cnt = entity_.neighbor_cnt(level); + if (topk_heap.size() <= static_cast(entity_.prune_cnt())) { + if (topk_heap.size() <= static_cast(max_neighbor_cnt)) { + entity_.update_neighbors(level, id, topk_heap); + return; + } + } + + uint32_t cur_size = 0; + for (size_t i = 0; i < topk_heap.size(); ++i) { + node_id_t cur_node = topk_heap[i].first; + dist_t cur_node_dist = topk_heap[i].second; + bool good = true; + for (uint32_t j = 0; j < cur_size; ++j) { + dist_t tmp_dist = dc.dist(cur_node, topk_heap[j].first); + if (tmp_dist <= cur_node_dist) { + good = false; + break; + } + } + + if (good) { + topk_heap[cur_size].first = cur_node; + topk_heap[cur_size].second = cur_node_dist; + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } + } + } + + // when after-prune neighbor count is too seldom, + // we use this strategy to make-up enough edges + // not only just make-up out-degrees + // we also make-up enough in-degrees + uint32_t min_neighbors = entity_.min_neighbor_cnt(); + for (size_t k = cur_size; cur_size < min_neighbors && k < topk_heap.size(); + ++k) { + bool exist = false; + for (size_t j = 0; j < cur_size; ++j) { + if (topk_heap[j].first == topk_heap[k].first) { + exist = true; + break; + } + } + if (!exist) { + topk_heap[cur_size].first = topk_heap[k].first; + topk_heap[cur_size].second = topk_heap[k].second; + cur_size++; + } + } + + topk_heap.resize(cur_size); + entity_.update_neighbors(level, id, topk_heap); + + return; +} + +template +void HnswAlgorithm::reverse_update_neighbors( + HnswDistCalculator &dc, node_id_t id, level_t level, node_id_t link_id, + dist_t dist, TopkHeap &update_heap) { + const size_t max_neighbor_cnt = entity_.neighbor_cnt(level); + + uint32_t lock_idx = id & kLockMask; + lock_pool_[lock_idx].lock(); + const Neighbors neighbors = entity_.get_neighbors(level, id); + size_t size = neighbors.size(); + ailego_assert_with(size <= max_neighbor_cnt, "invalid neighbor size"); + if (size < max_neighbor_cnt) { + entity_.add_neighbor(level, id, size, link_id); + lock_pool_[lock_idx].unlock(); + return; + } + + update_heap.emplace(link_id, dist); + + for (size_t i = 0; i < size; ++i) { + node_id_t node = neighbors[i]; + dist_t cur_dist = dc.dist(id, node); + update_heap.emplace(node, cur_dist); + } + + //! TODO: optimize prune + //! prune edges + update_heap.sort(); + size_t cur_size = 0; + for (size_t i = 0; i < update_heap.size(); ++i) { + node_id_t cur_node = update_heap[i].first; + dist_t cur_node_dist = update_heap[i].second; + bool good = true; + for (size_t j = 0; j < cur_size; ++j) { + dist_t tmp_dist = dc.dist(cur_node, update_heap[j].first); + if (tmp_dist <= cur_node_dist) { + good = false; + break; + } + } + + if (good) { + update_heap[cur_size].first = cur_node; + update_heap[cur_size].second = cur_node_dist; + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } + } + } + + update_heap.resize(cur_size); + entity_.update_neighbors(level, id, update_heap); + + lock_pool_[lock_idx].unlock(); + + update_heap.clear(); + + return; +} + +// Explicit template instantiation +template class HnswAlgorithm; +template class HnswAlgorithm; +template class HnswAlgorithm; +template class HnswAlgorithm; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.h b/src/core/algorithm/hnsw/hnsw_algorithm.h new file mode 100644 index 0000000..7851b1c --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_algorithm.h @@ -0,0 +1,156 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include "hnsw_context.h" +#include "hnsw_dist_calculator.h" +#include "hnsw_entity.h" +#include "hnsw_streamer_entity.h" + +namespace zvec { +namespace core { + +//! Non-template base class for HnswAlgorithm +class HnswAlgorithmBase { + public: + typedef std::unique_ptr UPointer; + + virtual ~HnswAlgorithmBase() = default; + + virtual int cleanup() = 0; + virtual int add_node(node_id_t id, level_t level, HnswContext *ctx) = 0; + virtual int search(HnswContext *ctx) const = 0; + virtual int init() = 0; + virtual uint32_t get_random_level() const = 0; +}; + +//! hnsw graph algorithm implement, templated on EntityType +template +class HnswAlgorithm : public HnswAlgorithmBase { + public: + using MemBlockType = typename EntityType::MemoryBlock; + + //! Constructor + explicit HnswAlgorithm(EntityType &entity) + : entity_(entity), + mt_(std::chrono::system_clock::now().time_since_epoch().count()), + lock_pool_(kLockCnt) {} + + //! Destructor + ~HnswAlgorithm() override = default; + + //! Cleanup HnswAlgorithm + int cleanup() override { + return 0; + } + + //! Add a node to hnsw graph + //! @id: the node unique id + //! @level: a node will be add to graph in each level [0, level] + //! return 0 on success, or errCode in failure + int add_node(node_id_t id, level_t level, HnswContext *ctx) override; + + //! do knn search in graph + //! return 0 on success, or errCode in failure. results saved in ctx + int search(HnswContext *ctx) const override; + + //! Initiate HnswAlgorithm + int init() override { + level_probas_.clear(); + double level_mult = + 1 / std::log(static_cast(entity_.scaling_factor())); + for (int level = 0;; level++) { + // refers faiss get_random_level alg + double proba = + std::exp(-level / level_mult) * (1 - std::exp(-1 / level_mult)); + if (proba < 1e-9) { + break; + } + level_probas_.push_back(proba); + } + + return 0; + } + + //! Generate a random level + //! return graph level + uint32_t get_random_level() const override { + // gen rand float (0, 1) + double f = mt_() / static_cast(mt_.max()); + for (size_t level = 0; level < level_probas_.size(); level++) { + if (f < level_probas_[level]) { + return level; + } + f -= level_probas_[level]; + } + return level_probas_.size() - 1; + } + + private: + //! Select in upper layer to get entry point for next layer search + void select_entry_point(level_t level, node_id_t *entry_point, dist_t *dist, + HnswContext *ctx) const; + + //! update node id neighbors from topkHeap, and reverse link is also updated + void add_neighbors(node_id_t id, level_t level, TopkHeap &topk_heap, + HnswContext *ctx); + + //! Given a node id and level, search the nearest neighbors in graph. + //! Dispatches to fast_search_neighbors (pool-based, direct pointer) for + //! mmap/contiguous level-0 unfiltered search, or dual_heap_search_neighbors + //! (CandidateHeap + TopkHeap) for add_node, filtered search, upper levels, + //! and BufferPool fallback. + //! Note: entry_point and dist will be updated to current level nearest node. + void search_neighbors(level_t level, node_id_t *entry_point, dist_t *dist, + TopkHeap &topk, HnswContext *ctx, bool use_pool) const; + + //! Update the node's neighbors + void update_neighbors(HnswDistCalculator &dc, node_id_t id, level_t level, + TopkHeap &topk_heap); + + //! Checking linkId could be id's new neighbor, and add as neighbor if true + //! @dc distance calculator + //! @updateHeap temporary heap in updating neighbors + void reverse_update_neighbors(HnswDistCalculator &dc, node_id_t id, + level_t level, node_id_t link_id, dist_t dist, + TopkHeap &update_heap); + + //! expand neighbors until group nums are reached + void expand_neighbors_by_group(TopkHeap &topk, HnswContext *ctx) const; + + private: + HnswAlgorithm(const HnswAlgorithm &) = delete; + HnswAlgorithm &operator=(const HnswAlgorithm &) = delete; + + private: + static constexpr uint32_t kLockCnt{1U << 8}; + static constexpr uint32_t kLockMask{kLockCnt - 1U}; + + EntityType &entity_; + mutable std::mt19937 mt_{}; + std::vector level_probas_{}; + + mutable ailego::SpinMutex spin_lock_{}; // global spin lock + std::mutex mutex_{}; // global mutex + // TODO: spin lock? + std::vector lock_pool_{}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw/hnsw_chunk.cc b/src/core/algorithm/hnsw/hnsw_chunk.cc new file mode 100644 index 0000000..2d1ea24 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_chunk.cc @@ -0,0 +1,219 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_chunk.h" +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +int ChunkBroker::init_storage(uint32_t chunk_size) { + chunk_meta_.clear(); + chunk_meta_.chunk_size = chunk_size; + chunk_meta_.create_time = ailego::Realtime::Seconds(); + stats_.set_create_time(chunk_meta_.create_time); + chunk_meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(chunk_meta_.update_time); + + //! alloc meta chunk + size_t size = sizeof(HnswChunkMeta); + size = (size + page_mask_) & (~page_mask_); + const std::string segment_id = + make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId); + int ret = stg_->append(segment_id, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Storage append segment failed for %s", IndexError::What(ret)); + return ret; + } + chunk_meta_segment_ = get_chunk(CHUNK_TYPE_META, kDefaultChunkSeqId); + if (ailego_unlikely(!chunk_meta_segment_)) { + LOG_ERROR("Get meta segment failed"); + return IndexError_Runtime; + } + + //! update meta info and write to storage + chunk_meta_.chunk_cnts[CHUNK_TYPE_META] += 1; + chunk_meta_.total_size += size; + (*stats_.mutable_index_size()) += size; + size = chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswChunkMeta))) { + LOG_ERROR("Storage write data failed, wsize=%zu", size); + return IndexError_WriteData; + } + + return 0; +} + +int ChunkBroker::load_storage(uint32_t &chunk_size) { + IndexStorage::MemoryBlock data_block; + size_t size = chunk_meta_segment_->read(0UL, data_block, + chunk_meta_segment_->data_size()); + if (size != sizeof(HnswChunkMeta)) { + LOG_ERROR("Invalid hnsw meta chunk, read size=%zu chunk size=%zu", size, + chunk_meta_segment_->data_size()); + return IndexError_InvalidFormat; + } + std::memcpy(static_cast(&chunk_meta_), data_block.data(), size); + if (chunk_meta_.chunk_size != chunk_size) { + LOG_WARN( + "chunk_size mismatch: config=%u, index=%zu. " + "Using index value to keep compatibility.", + chunk_size, (size_t)chunk_meta_.chunk_size); + } + chunk_size = chunk_meta_.chunk_size; + + *stats_.mutable_check_point() = stg_->check_point(); + stats_.set_revision_id(chunk_meta_.revision_id); + stats_.set_update_time(chunk_meta_.update_time); + stats_.set_create_time(chunk_meta_.create_time); + + char create_time[32]; + char update_time[32]; + ailego::Realtime::Gmtime(chunk_meta_.create_time, "%Y-%m-%d %H:%M:%S", + create_time, sizeof(create_time)); + ailego::Realtime::Gmtime(chunk_meta_.update_time, "%Y-%m-%d %H:%M:%S", + update_time, sizeof(update_time)); + LOG_DEBUG( + "Load index, indexSize=%zu chunkSize=%zu nodeChunks=%zu " + "upperNeighborChunks=%zu revisionId=%zu " + "createTime=%s updateTime=%s", + (size_t)chunk_meta_.total_size, (size_t)chunk_meta_.chunk_size, + (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_NODE], + (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_UPPER_NEIGHBOR], + (size_t)chunk_meta_.revision_id, create_time, update_time); + + return 0; +} + +int ChunkBroker::open(IndexStorage::Pointer stg, uint32_t &chunk_size, + bool check_crc) { + if (ailego_unlikely(stg_)) { + LOG_ERROR("An storage instance is already opened"); + return IndexError_Duplicate; + } + stg_ = std::move(stg); + if (stg_->isHugePage()) { + page_mask_ = ailego::MemoryHelper::HugePageSize() - 1; + } else { + page_mask_ = ailego::MemoryHelper::PageSize() - 1; + } + check_crc_ = check_crc; + dirty_ = false; + + const std::string segment_id = + make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId); + chunk_meta_segment_ = stg_->get(segment_id); + if (!chunk_meta_segment_) { + LOG_DEBUG("Create new index"); + return init_storage(chunk_size); + } + + return load_storage(chunk_size); +} + +int ChunkBroker::close(void) { + flush(0UL); + + stg_.reset(); + check_crc_ = false; + dirty_ = false; + + return 0; +} + +int ChunkBroker::flush(uint64_t checkpoint) { + ailego_assert_with(chunk_meta_segment_, "invalid meta segment"); + + chunk_meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(chunk_meta_.update_time); + + size_t size = + chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswChunkMeta))) { + LOG_ERROR("Storage write data failed, wsize=%zu", size); + } + + stg_->refresh(checkpoint); + int ret = stg_->flush(); + if (ret == 0) { + (*stats_.mutable_check_point()) = checkpoint; + } else { + LOG_ERROR("Storage flush failed for %s", IndexError::What(ret)); + } + return ret; +} + +std::pair ChunkBroker::alloc_chunk(int type, + uint64_t seq_id, + size_t size) { + ailego_assert_with(type < CHUNK_TYPE_MAX, "chunk type overflow"); + + Chunk::Pointer chunk; + if (ailego_unlikely(!stg_)) { + LOG_ERROR("Init storage first"); + return std::make_pair(IndexError_Uninitialized, chunk); + } + + //! check exist a empty chunk with the same name + chunk = get_chunk(type, seq_id); + if (chunk) { + if (ailego_unlikely(chunk->capacity() == size && + chunk->data_size() == 0UL)) { + LOG_ERROR("Exist invalid chunk size %zu, expect size %zu", + chunk->capacity(), size); + chunk.reset(); + return std::make_pair(IndexError_Runtime, chunk); + } + return std::make_pair(0, chunk); + } + //! align to page size + size = (size + page_mask_) & (~page_mask_); + if (ailego_unlikely(chunk_meta_.total_size + size >= max_chunks_size_)) { + LOG_ERROR("No space to new a chunk, curIndexSize=%zu allocSize=%zu", + (size_t)chunk_meta_.total_size, size); + return std::make_pair(IndexError_IndexFull, chunk); + } + + std::string segment_id = make_segment_id(type, seq_id); + int ret = stg_->append(segment_id, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Storage append segment failed for %s", IndexError::What(ret)); + return std::make_pair(ret, chunk); + } + chunk_meta_.chunk_cnts[type] += 1; + chunk_meta_.total_size += size; + (*stats_.mutable_index_size()) += size; + + size = chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswChunkMeta))) { + LOG_ERROR("Storage append segment failed, wsize=%zu", size); + } + + chunk = get_chunk(type, seq_id); + return std::make_pair(chunk ? 0 : IndexError_NoMemory, chunk); +} + +Chunk::Pointer ChunkBroker::get_chunk(int type, uint64_t seq_id) const { + std::string segment_id = make_segment_id(type, seq_id); + return stg_->get(segment_id); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_chunk.h b/src/core/algorithm/hnsw/hnsw_chunk.h new file mode 100644 index 0000000..034a2dc --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_chunk.h @@ -0,0 +1,153 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +using Chunk = IndexStorage::Segment; + +class ChunkBroker { + public: + typedef std::shared_ptr Pointer; + + enum CHUNK_TYPE { + CHUNK_TYPE_HEADER = 1, + CHUNK_TYPE_META = 2, + CHUNK_TYPE_NODE = 3, + CHUNK_TYPE_UPPER_NEIGHBOR = 4, + CHUNK_TYPE_NEIGHBOR_INDEX = 5, + CHUNK_TYPE_SPARSE_NODE = 6, + CHUNK_TYPE_NEIGHBOR_DIST = 7, // Vamana: per-node neighbor distances + CHUNK_TYPE_MAX = 8 + }; + static constexpr size_t kDefaultChunkSeqId = 0UL; + + ChunkBroker(IndexStreamer::Stats &stats) : stats_(stats) {} + + //! Open storage + int open(IndexStorage::Pointer stg, uint32_t &chunk_size, bool check_crc); + + int close(void); + + int flush(uint64_t checkpoint); + + //! alloc a new chunk with size, not thread-safe + std::pair alloc_chunk(int type, uint64_t seq_id, + size_t size); + + //! alloc a new chunk with chunk size + inline std::pair alloc_chunk(int type, uint64_t seq_id) { + return alloc_chunk(type, seq_id, chunk_meta_.chunk_size); + } + + Chunk::Pointer get_chunk(int type, uint64_t seq_id) const; + + inline size_t get_chunk_cnt(int type) const { + ailego_assert_with(type < CHUNK_TYPE_MAX, "chunk type overflow"); + return chunk_meta_.chunk_cnts[type]; + } + + inline bool dirty(void) const { + return dirty_; + } + + inline void mark_dirty(void) { + if (!dirty_) { + dirty_ = true; + chunk_meta_.revision_id += 1; + stats_.set_revision_id(chunk_meta_.revision_id); + } + } + + const IndexStorage::Pointer storage(void) const { + return stg_; + } + + //! Set the maximum total size (bytes) that alloc_chunk() is allowed to + //! consume. MUST be called after open() and before any alloc_chunk() + //! invocation; if omitted, max_chunks_size_ remains 0 and every + //! alloc_chunk() call will immediately return IndexError_IndexFull. + //! + //! Typical call sequence: + //! 1. open(stg, chunk_size, check_crc) + //! 2. init_chunk_params(max_index_size, huge_page) + //! 3. set_max_chunks_size(max_index_size_) // <- must be here + //! 4. alloc_chunk(...) + void set_max_chunks_size(size_t max_chunks_size) { + max_chunks_size_ = max_chunks_size; + } + + private: + ChunkBroker(const ChunkBroker &) = delete; + ChunkBroker &operator=(const ChunkBroker &) = delete; + + struct HnswChunkMeta { + HnswChunkMeta(void) { + memset(static_cast(this), 0, sizeof(HnswChunkMeta)); + } + void clear() { + memset(static_cast(this), 0, sizeof(HnswChunkMeta)); + } + + uint64_t chunk_cnts[CHUNK_TYPE_MAX]; + uint64_t chunk_size; // size of per chunk + uint64_t total_size; // total size of allocated chunk + uint64_t revision_id; // index revision + uint64_t create_time; + uint64_t update_time; + uint64_t reserved[3]; + }; + + static_assert(sizeof(HnswChunkMeta) % 32 == 0, + "HnswChunkMeta must be aligned with 32 bytes"); + + //! Init the storage after open an empty index + int init_storage(uint32_t chunk_size); + + //! Load index from storage + int load_storage(uint32_t &chunk_size); + + static inline const std::string make_segment_id(int type, uint64_t seq_id) { + return "HnswT" + ailego::StringHelper::ToString(type) + "S" + + ailego::StringHelper::ToString(seq_id); + } + + private: + IndexStreamer::Stats &stats_; + HnswChunkMeta chunk_meta_{}; + size_t page_mask_{0UL}; + size_t max_chunks_size_{0UL}; + IndexStorage::Pointer stg_{}; + IndexStorage::Segment::Pointer chunk_meta_segment_{}; + bool check_crc_{false}; + bool dirty_{false}; // set as true if index is modified , the flag + // will not be cleared even if flushed +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_context.cc b/src/core/algorithm/hnsw/hnsw_context.cc new file mode 100644 index 0000000..4710284 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_context.cc @@ -0,0 +1,309 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_context.h" +#include +#include "hnsw_params.h" + +namespace zvec { +namespace core { + +HnswContext::HnswContext(size_t dimension, const IndexMetric::Pointer &metric, + const HnswEntity::Pointer &entity) + : IndexContext(metric), + entity_(entity), + dc_(entity_.get(), metric, dimension) {} + +HnswContext::HnswContext(const IndexMetric::Pointer &metric, + const HnswEntity::Pointer &entity) + : IndexContext(metric), entity_(entity), dc_(entity_.get(), metric) {} + +HnswContext::~HnswContext() { + visit_filter_.destroy(); +} + +int HnswContext::init(ContextType type) { + int ret; + uint32_t doc_cnt; + + type_ = type; + + switch (type) { + case kBuilderContext: + ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(), + max_scan_num_, negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + update_heap_.limit(entity_->l0_neighbor_cnt() + 1); + break; + + case kSearcherContext: + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_, + negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + break; + + case kStreamerContext: + // maxScanNum is unknown if inited from streamer, so the docCnt may + // change. we need to compute maxScanNum by scan ratio, and preserve + // max_doc_cnt space from visit filter + doc_cnt = entity_->doc_cnt(); + max_scan_num_ = compute_max_scan_num(doc_cnt); + reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt); + ret = visit_filter_.init(filter_mode_, reserve_max_doc_cnt_, + max_scan_num_, negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + + update_heap_.limit(entity_->l0_neighbor_cnt() + 1); + candidates_.limit(max_scan_num_); + + check_need_adjuct_ctx(); + break; + + default: + LOG_ERROR("Init context failed"); + return IndexError_Runtime; + } + + return 0; +} + +int HnswContext::update(const ailego::Params ¶ms) { + auto update_visit_filter_param = [&]() { + bool need_update = false; + std::string p; + switch (type_) { + case kSearcherContext: + p = PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_ENABLE; + break; + case kStreamerContext: + p = PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_ENABLE; + break; + } + + if (params.has(p)) { + bool bf_enabled = false; + params.get(p, &bf_enabled); + if (bf_enabled ^ (filter_mode_ == VisitFilter::BloomFilter)) { + need_update = true; + filter_mode_ = + bf_enabled ? VisitFilter::BloomFilter : VisitFilter::ByteMap; + } + } + + float prob = negative_probability_; + p.clear(); + switch (type_) { + case kSearcherContext: + p = PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB; + break; + case kStreamerContext: + p = PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB; + break; + } + params.get(p, &prob); + if (filter_mode_ == VisitFilter::BloomFilter && + std::abs(prob - negative_probability_) > 1e-6) { + need_update = true; + } + if (need_update) { + visit_filter_.destroy(); + int max_doc_cnt = 0; + if (type_ == kSearcherContext) { + max_doc_cnt = entity_->doc_cnt(); + } else { + max_doc_cnt = reserve_max_doc_cnt_; + } + int ret = visit_filter_.init(filter_mode_, max_doc_cnt, max_scan_num_, + negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + } + return 0; + }; + + switch (type_) { + case kSearcherContext: + if (params.has(PARAM_HNSW_SEARCHER_EF)) { + params.get(PARAM_HNSW_SEARCHER_EF, &ef_); + topk_heap_.limit(std::max(topk_, ef_)); + } + + if (params.has(PARAM_HNSW_SEARCHER_PO)) { + params.get(PARAM_HNSW_SEARCHER_PO, &po_); + } + + if (params.has(PARAM_HNSW_SEARCHER_PL)) { + params.get(PARAM_HNSW_SEARCHER_PL, &pl_); + } + + if (params.has(PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO)) { + params.get(PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_); + max_scan_num_ = + static_cast(max_scan_ratio_ * entity_->doc_cnt()); + max_scan_num_ = std::max(10000U, max_scan_num_); + } + + if (params.has(PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD)) { + params.get(PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + } + + return update_visit_filter_param(); + + case kStreamerContext: + if (params.has(PARAM_HNSW_STREAMER_EF)) { + params.get(PARAM_HNSW_STREAMER_EF, &ef_); + topk_heap_.limit(std::max(topk_, ef_)); + } + params.get(PARAM_HNSW_STREAMER_EF, &ef_); + params.get(PARAM_HNSW_STREAMER_PO, &po_); + params.get(PARAM_HNSW_STREAMER_PL, &pl_); + params.get(PARAM_HNSW_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_); + params.get(PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_); + params.get(PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_); + if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) { + LOG_ERROR("[%s] must be in range (0.0f,1.0f]", + PARAM_HNSW_STREAMER_MAX_SCAN_RATIO.c_str()); + return IndexError_InvalidArgument; + } + if (max_scan_limit_ < min_scan_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT.c_str(), + PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT.c_str()); + return IndexError_InvalidArgument; + } + + if (params.has(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD)) { + params.get(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + } + + return update_visit_filter_param(); + + default: + LOG_ERROR("update context failed, type=%u", type_); + return IndexError_Runtime; + } +} + +int HnswContext::update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const HnswEntity::Pointer &entity, + uint32_t magic_num) { + uint32_t doc_cnt; + + if (ailego_unlikely(static_cast(type) != type_)) { + LOG_ERROR( + "HnswContext doesn't support shared by different type, " + "src=%u dst=%u", + type_, type); + return IndexError_Unsupported; + } + + magic_ = kInvalidMgic; + + // TODO: support change filter mode? + switch (type) { + case kBuilderContext: + LOG_ERROR("BuildContext doesn't support update"); + return IndexError_NotImplemented; + + case kSearcherContext: + if (!visit_filter_.reset(entity->doc_cnt(), max_scan_num_)) { + LOG_ERROR("Reset filter failed, mode %d", visit_filter_.get_mode()); + return IndexError_Runtime; + } + + candidates_.limit(max_scan_num_); + topk_heap_.limit(std::max(topk_, ef_)); + break; + + case kStreamerContext: + doc_cnt = entity->doc_cnt(); + max_scan_num_ = compute_max_scan_num(doc_cnt); + reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt); + if (!visit_filter_.reset(reserve_max_doc_cnt_, max_scan_num_)) { + LOG_ERROR("Reset filter failed, mode %d", visit_filter_.get_mode()); + return IndexError_Runtime; + } + + update_heap_.limit(entity->l0_neighbor_cnt() + 1); + candidates_.limit(max_scan_num_); + topk_heap_.limit(std::max(topk_, ef_)); + break; + + default: + LOG_ERROR("update context failed"); + return IndexError_Runtime; + } + + entity_ = entity; + dc_.update(entity_.get(), metric, meta.dimension()); + if (vector_source_) { + entity_->set_vector_source(vector_source_); + } + magic_ = magic_num; + level_topks_.clear(); + + return 0; +} + +void HnswContext::fill_random_to_topk_full(void) { + static std::mt19937 mt( + std::chrono::system_clock::now().time_since_epoch().count()); + std::uniform_int_distribution dt(0, entity_->doc_cnt() - 1); + std::function gen; + node_id_t seqid; + std::function myfilter = [](node_id_t) { return false; }; + if (this->filter().is_valid()) { + myfilter = [&](node_id_t id) { + return this->filter()(entity_->get_key(id)); + }; + } + + if (topk_heap_.limit() < entity_->doc_cnt() / 2) { + gen = [&](void) { return dt(mt); }; + } else { + // If topk limit is big value, gen sequential id from an random initial + seqid = dt(mt); + gen = [&](void) { + seqid = seqid == (entity_->doc_cnt() - 1) ? 0 : (seqid + 1); + return seqid; + }; + } + + for (size_t i = 0; !topk_heap_.full() && i < entity_->doc_cnt(); ++i) { + const auto id = gen(); + if (!visit_filter_.visited(id) && !myfilter(id)) { + visit_filter_.set_visited(id); + topk_heap_.emplace(id, dc_.dist(id)); + } + } + return; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw/hnsw_context.h b/src/core/algorithm/hnsw/hnsw_context.h new file mode 100644 index 0000000..b7ac79c --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_context.h @@ -0,0 +1,596 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "utility/block_heap.h" +#include "utility/linear_pool.h" +#include "utility/sparse_utility.h" +#include "utility/visit_filter.h" +#include "hnsw_dist_calculator.h" +#include "hnsw_entity.h" + +namespace zvec { +namespace core { + +class HnswContext : public IndexContext { + public: + //! Index Context Pointer + typedef std::unique_ptr Pointer; + + enum ContextType { + kUnknownContext = 0, + kSearcherContext = 1, + kBuilderContext = 2, + kStreamerContext = 3 + }; + + //! Construct + HnswContext(size_t dimension, const IndexMetric::Pointer &metric, + const HnswEntity::Pointer &entity); + + //! Construct + HnswContext(const IndexMetric::Pointer &metric, + const HnswEntity::Pointer &entity); + + //! Destructor + ~HnswContext() override; + + public: + //! Set topk of search result + void set_topk(uint32_t val) override { + topk_ = val; + topk_heap_.limit(std::max(val, ef_)); + } + + //! Retrieve search result + const IndexDocumentList &result(void) const override { + return results_[0]; + } + + //! Retrieve search result + const IndexDocumentList &result(size_t idx) const override { + return results_[idx]; + } + + //! Retrieve result object for output + IndexDocumentList *mutable_result(size_t idx) override { + ailego_assert_with(idx < results_.size(), "invalid idx"); + return &results_[idx]; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(void) const override { + return group_results_[0]; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(size_t idx) const override { + return group_results_[idx]; + } + + IndexGroupDocumentList *mutable_group_result(void) override { + return &group_results_[0]; + } + + IndexGroupDocumentList *mutable_group_result(size_t idx) override { + return &group_results_[idx]; + } + + uint32_t magic(void) const override { + return magic_; + } + + //! Set mode of debug + void set_debug_mode(bool enable) override { + debug_mode_ = enable; + } + + //! Retrieve mode of debug + bool debug_mode(void) const override { + return this->debugging(); + } + + //! Retrieve string of debug + std::string debug_string(void) const override { + char buf[4096]; + size_t size = snprintf( + buf, sizeof(buf), + "scan_cnt=%zu,get_vector_cnt=%u,get_neighbors_cnt=%u,dup_node=%u", + get_scan_num(), stats_get_vector_cnt_, stats_get_neighbors_cnt_, + stats_visit_dup_cnt_); + return std::string(buf, size); + } + + //! Update the parameters of context + int update(const ailego::Params ¶ms) override; + + public: + //! Init context + int init(ContextType type); + + //! Update context, the context may be shared by different searcher/streamer + int update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const HnswEntity::Pointer &entity, uint32_t magic_num); + + inline const HnswEntity &get_entity() const { + return *entity_; + } + + //! Bind an external vector source to this context. It is stored so that it + //! can be re-applied after the entity is re-cloned inside update_context, + //! and immediately forwarded to the current entity clone. + inline void set_vector_source(const VectorSource *src) { + vector_source_ = src; + if (entity_) { + entity_->set_vector_source(src); + } + } + + inline const VectorSource *vector_source() const { + return vector_source_; + } + + inline void resize_results(size_t size) { + if (group_by_search()) { + group_results_.resize(size); + } else { + results_.resize(size); + } + } + + inline void topk_to_result() { + return topk_to_result(0); + } + + //! Construct result from topk heap, result will be normalized + inline void topk_to_result(uint32_t idx) { + if (group_by_search()) { + topk_to_group_result(idx); + } else { + topk_to_single_result(idx); + } + } + + inline void recal_topk_dist() { + TopkHeap heap(topk_heap_); + topk_heap_.clear(); + + for (size_t i = 0; i < heap.size(); ++i) { + node_id_t id = heap[i].first; + dist_t dist = dc_.dist(id); + topk_heap_.emplace_back(id, dist); + } + } + + inline void topk_to_single_result(uint32_t idx) { + if (force_padding_topk_ && !topk_heap_.full() && + topk_heap_.size() < entity_->doc_cnt()) { + this->fill_random_to_topk_full(); + } + if (ailego_unlikely(topk_heap_.size() == 0)) { + return; + } + + ailego_assert_with(idx < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(topk_heap_.size())); + topk_heap_.sort(); + results_[idx].clear(); + + for (int i = 0; i < size; ++i) { + auto score = topk_heap_[i].second; + if (score > this->threshold()) { + break; + } + + node_id_t id = topk_heap_[i].first; + if (fetch_vector_) { + IndexStorage::MemoryBlock block; + entity_->get_vector(id, block); + results_[idx].emplace_back(entity_->get_key(id), score, id, block); + } else { + results_[idx].emplace_back(entity_->get_key(id), score, id); + } + } + + return; + } + + //! Construct result from topk heap, result will be normalized + inline void topk_to_group_result(uint32_t idx) { + ailego_assert_with(idx < group_results_.size(), "invalid idx"); + + group_results_[idx].clear(); + + std::vector> group_topk_list; + std::vector> best_score_in_groups; + for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end(); + itr++) { + const std::string &group_id = (*itr).first; + auto &heap = (*itr).second; + heap.sort(); + + if (heap.size() > 0) { + float best_score = heap[0].second; + best_score_in_groups.push_back(std::make_pair(group_id, best_score)); + } + } + + std::sort(best_score_in_groups.begin(), best_score_in_groups.end(), + [](const std::pair &a, + const std::pair &b) -> int { + return a.second < b.second; + }); + + // truncate to group num + for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size(); + ++i) { + const std::string &group_id = best_score_in_groups[i].first; + + group_topk_list.emplace_back( + std::make_pair(group_id, group_topk_heaps_[group_id])); + } + + group_results_[idx].resize(group_topk_list.size()); + + for (uint32_t i = 0; i < group_topk_list.size(); ++i) { + const std::string &group_id = group_topk_list[i].first; + group_results_[idx][i].set_group_id(group_id); + + uint32_t size = std::min( + group_topk_, static_cast(group_topk_list[i].second.size())); + + for (uint32_t j = 0; j < size; ++j) { + auto score = group_topk_list[i].second[j].second; + if (score > this->threshold()) { + break; + } + + node_id_t id = group_topk_list[i].second[j].first; + + if (fetch_vector_) { + IndexStorage::MemoryBlock block; + entity_->get_vector(id, block); + group_results_[idx][i].mutable_docs()->emplace_back( + entity_->get_key(id), score, id, block); + } else { + group_results_[idx][i].mutable_docs()->emplace_back( + entity_->get_key(id), score, id); + } + } + } + } + + inline void reset_query(const void *query) { + if (auto query_preprocess_func = index_metric_->get_query_preprocess_func(); + query_preprocess_func != nullptr) { + size_t dim = dc_.dimension(); + preprocess_buffer_.resize(dim); + memcpy(preprocess_buffer_.data(), query, dim); + query_preprocess_func(preprocess_buffer_.data(), dim); + query = preprocess_buffer_.data(); + } + + dc_.reset_query(query); + dc_.clear_compare_cnt(); + } + + inline HnswDistCalculator &dist_calculator() { + return dc_; + } + + inline TopkHeap &topk_heap() { + return topk_heap_; + } + + inline TopkHeap &update_heap() { + return update_heap_; + } + + inline LinearPool &pool() { + return pool_; + } + + // Only accessed under a runtime CpuFeatures::AVX2 guard at call sites. + inline BlockHeap &block_pool() { + return block_pool_; + } + + inline VisitFilter &visit_filter() { + return visit_filter_; + } + + inline CandidateHeap &candidates() { + return candidates_; + } + + inline void set_max_scan_num(uint32_t max_scan_num) { + max_scan_num_ = max_scan_num; + } + + inline void set_max_scan_limit(uint32_t max_scan_limit) { + max_scan_limit_ = max_scan_limit; + } + + inline void set_min_scan_limit(uint32_t min_scan_limit) { + min_scan_limit_ = min_scan_limit; + } + + inline void set_ef(uint32_t v) { + ef_ = v; + } + + inline uint32_t ef(void) const { + return ef_; + } + + inline void set_po(uint32_t v) { + po_ = v; + } + + inline uint32_t po(void) const { + return po_; + } + + inline void set_pl(uint32_t v) { + pl_ = v; + } + + inline uint32_t pl(void) const { + return pl_; + } + + inline void set_filter_mode(uint32_t v) { + filter_mode_ = v; + } + + inline void set_filter_negative_probability(float v) { + negative_probability_ = v; + } + + inline void set_max_scan_ratio(float v) { + max_scan_ratio_ = v; + } + + virtual void set_magic(uint32_t v) { + magic_ = v; + } + + virtual void set_force_padding_topk(bool v) { + force_padding_topk_ = v; + } + + void set_bruteforce_threshold(uint32_t v) override { + bruteforce_threshold_ = v; + } + + inline uint32_t get_bruteforce_threshold() const { + return bruteforce_threshold_; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + + bool fetch_vector() const override { + return fetch_vector_; + } + + //! Reset context + void reset(void) override { + this->clear(); + set_filter(nullptr); + reset_threshold(); + set_fetch_vector(false); + set_group_params(0, 0); + reset_group_by(); + set_vector_source(nullptr); + } + + inline std::map &group_topk_heaps() { + return group_topk_heaps_; + } + + inline TopkHeap &level_topk(int level) { + if (ailego_unlikely(level_topks_.size() <= static_cast(level))) { + int cur_level = level_topks_.size(); + level_topks_.resize(level + 1); + for (; cur_level <= level; ++cur_level) { + size_t heap_size = std::max(entity_->neighbor_cnt(cur_level), + entity_->ef_construction()); + level_topks_[cur_level].clear(); + level_topks_[cur_level].limit(heap_size); + } + } + + return level_topks_[level]; + } + + inline void check_need_adjuct_ctx(void) { + check_need_adjuct_ctx(entity_->doc_cnt()); + } + + inline size_t compute_reserve_cnt(uint32_t cur_doc) const { + if (cur_doc > kMaxReserveDocCnt) { + return kMaxReserveDocCnt; + } else if (cur_doc < kMinReserveDocCnt) { + return kMinReserveDocCnt; + } + return cur_doc; + } + + //! candidates heap and visitfilter need to resize as doc cnt growing up + inline void check_need_adjuct_ctx(uint32_t doc_cnt) { + if (ailego_unlikely(doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_)) { + while (doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_) { + reserve_max_doc_cnt_ = + reserve_max_doc_cnt_ + compute_reserve_cnt(reserve_max_doc_cnt_); + } + uint32_t max_scan_cnt = compute_max_scan_num(reserve_max_doc_cnt_); + max_scan_num_ = max_scan_cnt; + visit_filter_.reset(reserve_max_doc_cnt_, max_scan_cnt); + candidates_.clear(); + candidates_.limit(max_scan_num_); + } + } + + inline uint32_t compute_max_scan_num(uint32_t max_doc_cnt) const { + uint32_t max_scan = max_doc_cnt * max_scan_ratio_; + if (max_scan < min_scan_limit_) { + max_scan = min_scan_limit_; + } else if (max_scan > max_scan_limit_) { + max_scan = max_scan_limit_; + } + return max_scan; + } + + inline size_t get_scan_num() const { + return dc_.compare_cnt(); + } + + inline uint64_t reach_scan_limit() const { + return dc_.compare_cnt() >= max_scan_num_; + } + + inline bool error() const { + return dc_.error(); + } + + inline void clear() { + dc_.clear(); + if (ailego_unlikely(this->debugging())) { + stats_get_neighbors_cnt_ = 0u; + stats_get_vector_cnt_ = 0u; + stats_visit_dup_cnt_ = 0u; + } + // do not clear results_ for the next query will need it + for (auto &it : results_) { + it.clear(); + } + for (auto &it : group_results_) { + it.clear(); + } + } + + uint32_t *mutable_stats_get_neighbors() { + return &stats_get_neighbors_cnt_; + } + + uint32_t *mutable_stats_get_vector() { + return &stats_get_vector_cnt_; + } + + uint32_t *mutable_stats_visit_dup_cnt() { + return &stats_visit_dup_cnt_; + } + + inline bool debugging(void) const { + return debug_mode_; + } + + inline void update_dist_caculator_distance( + const IndexMetric::MatrixDistance &distance, + const IndexMetric::MatrixBatchDistance &batch_distance) { + dc_.update_distance(distance, batch_distance); + } + + //! Get topk + inline uint32_t topk() const override { + return topk_; + } + + //! Get group topk + inline uint32_t group_topk() const { + return group_topk_; + } + + //! Get group num + inline uint32_t group_num() const { + return group_num_; + } + + //! Get if group by search + inline bool group_by_search() { + return group_num_ > 0; + } + + //! Set group params + void set_group_params(uint32_t group_num, uint32_t group_topk) override { + group_num_ = group_num; + group_topk_ = group_topk; + + topk_ = group_topk_ * group_num_; + + topk_heap_.limit(std::max(topk_, ef_)); + + group_topk_heaps_.clear(); + } + + private: + // Filling random nodes if topk not full + void fill_random_to_topk_full(void); + + constexpr static uint32_t kTriggerReserveCnt = 4096UL; + constexpr static uint32_t kMinReserveDocCnt = 4096UL; + constexpr static uint32_t kMaxReserveDocCnt = 128 * 1024UL; + constexpr static uint32_t kInvalidMgic = -1U; + + private: + HnswEntity::Pointer entity_; + HnswDistCalculator dc_; + IndexMetric::Pointer metric_; + const VectorSource *vector_source_{nullptr}; + + bool debug_mode_{false}; + bool force_padding_topk_{false}; + uint32_t max_scan_num_{0}; + uint32_t max_scan_limit_{0}; + uint32_t min_scan_limit_{0}; + uint32_t reserve_max_doc_cnt_{kMinReserveDocCnt}; + uint32_t topk_{0}; + uint32_t group_topk_{0}; + uint32_t filter_mode_{VisitFilter::ByteMap}; + float negative_probability_{HnswEntity::kDefaultBFNegativeProbability}; + uint32_t ef_{HnswEntity::kDefaultEf}; + uint32_t po_{8}; + uint32_t pl_{0}; + float max_scan_ratio_{HnswEntity::kDefaultScanRatio}; + uint32_t magic_{0U}; + std::vector results_{}; + std::vector group_results_{}; + TopkHeap topk_heap_{}; + TopkHeap update_heap_{}; + std::vector level_topks_{}; + CandidateHeap candidates_{}; + VisitFilter visit_filter_{}; + uint32_t bruteforce_threshold_{}; + bool fetch_vector_{false}; + + uint32_t group_num_{0}; + std::map group_topk_heaps_{}; + + uint32_t type_{kUnknownContext}; + //! debug stats info + uint32_t stats_get_neighbors_cnt_{0u}; + uint32_t stats_get_vector_cnt_{0u}; + uint32_t stats_visit_dup_cnt_{0u}; + std::string preprocess_buffer_; + + LinearPool pool_; + BlockHeap block_pool_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_dist_calculator.h b/src/core/algorithm/hnsw/hnsw_dist_calculator.h new file mode 100644 index 0000000..2e4b22d --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_dist_calculator.h @@ -0,0 +1,247 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "hnsw_entity.h" + +namespace zvec { +namespace core { + +class HnswDistCalculator { + public: + typedef std::shared_ptr Pointer; + + public: + enum DistType { + DIST_NONE = 0, + DIST_DENSE = 1, + DIST_HYBRID = 2, + DIST_SPARSE = 3 + }; + + public: + //! Constructor + HnswDistCalculator(const HnswEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(nullptr), + dim_(dim), + compare_cnt_(0) {} + + //! Constructor + HnswDistCalculator(const HnswEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim, + const void *query) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(query), + dim_(dim), + compare_cnt_(0) {} + + //! Constructor + HnswDistCalculator(const HnswEntity *entity, + const IndexMetric::Pointer &metric) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(nullptr), + dim_(0), + compare_cnt_(0) {} + + void update(const HnswEntity *entity, const IndexMetric::Pointer &metric) { + entity_ = entity; + distance_ = metric->distance(); + batch_distance_ = metric->batch_distance(); + } + + void update(const HnswEntity *entity, const IndexMetric::Pointer &metric, + uint32_t dim) { + entity_ = entity; + distance_ = metric->distance(); + batch_distance_ = metric->batch_distance(); + dim_ = dim; + } + + inline void update_distance( + const IndexMetric::MatrixDistance &distance, + const IndexMetric::MatrixBatchDistance &batch_distance) { + distance_ = distance; + batch_distance_ = batch_distance; + } + + //! Reset query vector data + inline void reset_query(const void *query) { + error_ = false; + query_ = query; + } + + //! Returns distance + inline dist_t dist(const void *vec_lhs, const void *vec_rhs) { + if (ailego_unlikely(vec_lhs == nullptr || vec_rhs == nullptr)) { + LOG_ERROR("Nullptr of dense vector"); + error_ = true; + return 0.0f; + } + + float score{0.0f}; + + distance_(vec_lhs, vec_rhs, dim_, &score); + + return score; + } + + //! Returns distance between query and vec. + inline dist_t dist(const void *vec) { + compare_cnt_++; + + return dist(vec, query_); + } + + //! Return distance between query and node id. + inline dist_t dist(node_id_t id) { + compare_cnt_++; + IndexStorage::MemoryBlock vec_block; + int ret = entity_->get_vector(id, vec_block); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + const void *feat = vec_block.data(); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + + return dist(feat, query_); + } + + //! Return dist node lhs between node rhs + inline dist_t dist(node_id_t lhs, node_id_t rhs) { + compare_cnt_++; + + + IndexStorage::MemoryBlock vec_block_feat; + int ret = entity_->get_vector(lhs, vec_block_feat); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Get nullptr vector, id=%u", lhs); + error_ = true; + return 0.0f; + } + const void *feat = vec_block_feat.data(); + + IndexStorage::MemoryBlock vec_block_query; + ret = entity_->get_vector(rhs, vec_block_query); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Get nullptr vector, id=%u", rhs); + error_ = true; + return 0.0f; + } + const void *query = vec_block_query.data(); + if (ailego_unlikely(feat == nullptr || query == nullptr)) { + LOG_ERROR("Get nullptr vector"); + error_ = true; + return 0.0f; + } + + return dist(feat, query); + } + + dist_t operator()(const void *vec) { + return dist(vec); + } + + dist_t operator()(node_id_t i) { + return dist(i); + } + + dist_t operator()(node_id_t lhs, node_id_t rhs) { + return dist(lhs, rhs); + } + + void batch_dist(const void **vecs, size_t num, dist_t *distances) { + compare_cnt_++; + + batch_distance_(vecs, query_, num, dim_, distances); + } + + inline dist_t batch_dist(node_id_t id) { + compare_cnt_++; + + IndexStorage::MemoryBlock vec_block; + int ret = entity_->get_vector(id, vec_block); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + const void *feat = vec_block.data(); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + dist_t score = 0; + batch_distance_(&feat, query_, 1, dim_, &score); + + return score; + } + + inline void clear() { + compare_cnt_ = 0; + error_ = false; + } + + inline void clear_compare_cnt() { + compare_cnt_ = 0; + } + + inline bool error() const { + return error_; + } + + //! Get distances compute times + inline uint32_t compare_cnt() const { + return compare_cnt_; + } + + inline uint32_t dimension() const { + return dim_; + } + + private: + HnswDistCalculator(const HnswDistCalculator &) = delete; + HnswDistCalculator &operator=(const HnswDistCalculator &) = delete; + + private: + const HnswEntity *entity_; + + IndexMetric::MatrixDistance distance_; + IndexMetric::MatrixBatchDistance batch_distance_; + + const void *query_; + uint32_t dim_; + + uint32_t compare_cnt_; // record distance compute times + // uint32_t compare_cnt_batch_; // record batch distance compute time + bool error_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_entity.cc b/src/core/algorithm/hnsw/hnsw_entity.cc new file mode 100644 index 0000000..aa0c0a3 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_entity.cc @@ -0,0 +1,364 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_entity.h" +#include +#include "utility/sparse_utility.h" + +namespace zvec { +namespace core { + +const std::string HnswEntity::kGraphHeaderSegmentId = "graph.header"; +const std::string HnswEntity::kGraphFeaturesSegmentId = "graph.features"; +const std::string HnswEntity::kGraphKeysSegmentId = "graph.keys"; +const std::string HnswEntity::kGraphNeighborsSegmentId = "graph.neighbors"; +const std::string HnswEntity::kGraphOffsetsSegmentId = "graph.offsets"; +const std::string HnswEntity::kGraphMappingSegmentId = "graph.mapping"; +const std::string HnswEntity::kHnswHeaderSegmentId = "hnsw.header"; +const std::string HnswEntity::kHnswNeighborsSegmentId = "hnsw.neighbors"; +const std::string HnswEntity::kHnswOffsetsSegmentId = "hnsw.offsets"; + +int HnswEntity::CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size) { + *padding_size = AlignSize(data_size) - data_size; + if (*padding_size == 0) { + return 0; + } + + std::string padding(*padding_size, '\0'); + if (dumper->write(padding.data(), *padding_size) != *padding_size) { + LOG_ERROR("Append padding failed, size %lu", *padding_size); + return IndexError_WriteData; + } + return 0; +} + +int64_t HnswEntity::dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, + const void *data, size_t size) const { + size_t len = dumper->write(data, size); + if (len != size) { + LOG_ERROR("Dump segment %s data failed, expect: %lu, actual: %lu", + segment_id.c_str(), size, len); + return IndexError_WriteData; + } + + size_t padding_size = AlignSize(size) - size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + uint32_t crc = ailego::Crc32c::Hash(data, size); + int ret = dumper->append(segment_id, size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", segment_id.c_str(), ret); + return ret; + } + + return len + padding_size; +} + +int64_t HnswEntity::dump_header(const IndexDumper::Pointer &dumper, + const HNSWHeader &hd) const { + //! dump basic graph header. header is aligned and does not need padding + int64_t graph_hd_size = + dump_segment(dumper, kGraphHeaderSegmentId, &hd.graph, hd.graph.size); + if (graph_hd_size < 0) { + return graph_hd_size; + } + + //! dump basic graph header. header is aligned and does not need padding + int64_t hnsw_hd_size = + dump_segment(dumper, kHnswHeaderSegmentId, &hd.hnsw, hd.hnsw.size); + if (hnsw_hd_size < 0) { + return hnsw_hd_size; + } + + return graph_hd_size + hnsw_hd_size; +} + +void HnswEntity::reshuffle_vectors( + const std::function & /*get_level*/, + std::vector * /*n2o_mapping*/, + std::vector * /*o2n_mapping*/, key_t * /*keys*/) const { + // TODO + return; +} + +int64_t HnswEntity::dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const { + std::vector mapping(doc_cnt()); + + std::iota(mapping.begin(), mapping.end(), 0U); + std::sort(mapping.begin(), mapping.end(), + [&](node_id_t i, node_id_t j) { return keys[i] < keys[j]; }); + + size_t size = mapping.size() * sizeof(node_id_t); + + return dump_segment(dumper, kGraphMappingSegmentId, mapping.data(), size); +} + +int64_t HnswEntity::dump_segments( + const IndexDumper::Pointer &dumper, key_t *keys, + const std::function &get_level) const { + HNSWHeader dump_hd(header()); + + dump_hd.graph.node_size = AlignSize(vector_size()); + + std::vector n2o_mapping; // map new id to origin id + std::vector o2n_mapping; // map origin id to new id + reshuffle_vectors(get_level, &n2o_mapping, &o2n_mapping, keys); + if (!o2n_mapping.empty()) { + dump_hd.hnsw.entry_point = o2n_mapping[entry_point()]; + } + + //! Dump header + int64_t hd_size = dump_header(dumper, dump_hd); + if (hd_size < 0) { + return hd_size; + } + + //! Dump vectors + int64_t vecs_size = dump_vectors(dumper, n2o_mapping); + if (vecs_size < 0) { + return vecs_size; + } + + //! Dump neighbors + auto neighbors_size = + dump_neighbors(dumper, get_level, n2o_mapping, o2n_mapping); + if (neighbors_size < 0) { + return neighbors_size; + } + //! free memory + n2o_mapping = std::vector(); + o2n_mapping = std::vector(); + + //! Dump keys + size_t key_segment_size = doc_cnt() * sizeof(key_t); + int64_t keys_size = + dump_segment(dumper, kGraphKeysSegmentId, keys, key_segment_size); + if (keys_size < 0) { + return keys_size; + } + + //! Dump mapping + int64_t mapping_size = dump_mapping_segment(dumper, keys); + if (mapping_size < 0) { + return mapping_size; + } + + return hd_size + keys_size + vecs_size + neighbors_size + mapping_size; +} + +int64_t HnswEntity::dump_vectors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + size_t vector_dump_size = vector_size(); + + size_t padding_size = AlignSize(vector_dump_size) - vector_dump_size; + + std::vector padding(padding_size); + memset(padding.data(), 0, sizeof(char) * padding_size); + const void *data = nullptr; + uint32_t crc = 0U; + size_t vecs_size = 0UL; + + //! dump vectors + for (node_id_t id = 0; id < doc_cnt(); ++id) { + data = get_vector(reorder_mapping.empty() ? id : reorder_mapping[id]); + if (ailego_unlikely(!data)) { + return IndexError_ReadData; + } + size_t len = dumper->write(data, vector_size()); + if (len != vector_size()) { + LOG_ERROR("Dump vectors failed, write=%zu expect=%zu", len, + vector_size()); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(data, vector_size(), crc); + vecs_size += vector_size(); + + if (padding_size == 0) { + continue; + } + + len = dumper->write(padding.data(), padding_size); + if (len != padding_size) { + LOG_ERROR("Dump vectors failed, write=%zu expect=%zu", len, padding_size); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(padding.data(), padding_size, crc); + vecs_size += padding_size; + } + + int ret = dumper->append(kGraphFeaturesSegmentId, vecs_size, 0UL, crc); + if (ret != 0) { + LOG_ERROR("Dump vectors segment meta failed, ret %d", ret); + return ret; + } + + return vecs_size; +} + +int64_t HnswEntity::dump_graph_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + std::vector graph_meta; + graph_meta.reserve(doc_cnt()); + size_t offset = 0; + uint32_t crc = 0; + std::vector mapping(l0_neighbor_cnt()); + + uint32_t min_neighbor_count = 10000; + uint32_t max_neighbor_count = 0; + size_t sum_neighbor_count = 0; + + for (node_id_t id = 0; id < doc_cnt(); ++id) { + const Neighbors neighbors = + get_neighbors(0, reorder_mapping.empty() ? id : reorder_mapping[id]); + ailego_assert_with(!!neighbors.data, "invalid neighbors"); + ailego_assert_with(neighbors.size() <= l0_neighbor_cnt(), + "invalid neighbors"); + + uint32_t neighbor_count = neighbors.size(); + if (neighbor_count < min_neighbor_count) { + min_neighbor_count = neighbor_count; + } + if (neighbor_count > max_neighbor_count) { + max_neighbor_count = neighbor_count; + } + sum_neighbor_count += neighbor_count; + + graph_meta.emplace_back(offset, neighbor_count); + size_t size = neighbors.size() * sizeof(node_id_t); + const node_id_t *data = &neighbors[0]; + if (!neighbor_mapping.empty()) { + for (node_id_t i = 0; i < neighbors.size(); ++i) { + mapping[i] = neighbor_mapping[neighbors[i]]; + } + data = mapping.data(); + } + if (dumper->write(data, size) != size) { + LOG_ERROR("Dump graph neighbor id=%u failed, size %lu", id, size); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(data, size, crc); + offset += size; + } + + uint32_t average_neighbor_count = 0; + if (doc_cnt() > 0) { + average_neighbor_count = sum_neighbor_count / doc_cnt(); + } + LOG_INFO( + "Dump hnsw graph: min_neighbor_count[%u] max_neighbor_count[%u] " + "average_neighbor_count[%u]", + min_neighbor_count, max_neighbor_count, average_neighbor_count); + + size_t padding_size = 0; + int ret = CalcAndAddPadding(dumper, offset, &padding_size); + if (ret != 0) { + return ret; + } + ret = dumper->append(kGraphNeighborsSegmentId, offset, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s failed, ret %d", + kGraphNeighborsSegmentId.c_str(), ret); + return ret; + } + + //! dump level 0 neighbors meta + auto len = dump_segment(dumper, kGraphOffsetsSegmentId, graph_meta.data(), + graph_meta.size() * sizeof(GraphNeighborMeta)); + if (len < 0) { + return len; + } + + return len + offset + padding_size; +} + +int64_t HnswEntity::dump_upper_neighbors( + const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + std::vector hnsw_meta; + hnsw_meta.reserve(doc_cnt()); + size_t offset = 0; + uint32_t crc = 0; + std::vector buffer(upper_neighbor_cnt() + 1); + for (node_id_t id = 0; id < doc_cnt(); ++id) { + node_id_t new_id = reorder_mapping.empty() ? id : reorder_mapping[id]; + auto level = get_level(new_id); + if (level == 0) { + hnsw_meta.emplace_back(0U, 0U); + continue; + } + hnsw_meta.emplace_back(offset, level); + ailego_assert_with((size_t)level < kMaxGraphLayers, "invalid level"); + for (level_t cur_level = 1; cur_level <= level; ++cur_level) { + const Neighbors neighbors = get_neighbors(cur_level, new_id); + ailego_assert_with(!!neighbors.data, "invalid neighbors"); + ailego_assert_with(neighbors.size() <= neighbor_cnt(cur_level), + "invalid neighbors"); + size_t buffer_bytes = buffer.size() * sizeof(node_id_t); + memset(buffer.data(), 0, buffer_bytes); + buffer[0] = neighbors.size(); + if (neighbor_mapping.empty()) { + memcpy(&buffer[1], &neighbors[0], neighbors.size() * sizeof(node_id_t)); + } else { + for (node_id_t i = 0; i < neighbors.size(); ++i) { + buffer[i + 1] = neighbor_mapping[neighbors[i]]; + } + } + if (dumper->write(buffer.data(), buffer_bytes) != buffer_bytes) { + LOG_ERROR("Dump graph neighbor id=%u failed, size %lu", id, + buffer_bytes); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(buffer.data(), buffer_bytes, crc); + offset += buffer_bytes; + } + } + size_t padding_size = 0; + int ret = CalcAndAddPadding(dumper, offset, &padding_size); + if (ret != 0) { + return ret; + } + + ret = dumper->append(kHnswNeighborsSegmentId, offset, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s failed, ret %d", kHnswNeighborsSegmentId.c_str(), + ret); + return ret; + } + + //! dump level 0 neighbors meta + auto len = dump_segment(dumper, kHnswOffsetsSegmentId, hnsw_meta.data(), + hnsw_meta.size() * sizeof(HnswNeighborMeta)); + if (len < 0) { + return len; + } + + return len + offset + padding_size; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_entity.h b/src/core/algorithm/hnsw/hnsw_entity.h new file mode 100644 index 0000000..639fc4a --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_entity.h @@ -0,0 +1,696 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +using node_id_t = uint32_t; +using key_t = uint64_t; +using level_t = int32_t; +using dist_t = float; +using TopkHeap = ailego::KeyValueHeap; +using CandidateHeap = + ailego::KeyValueHeap>; +constexpr node_id_t kInvalidNodeId = static_cast(-1); +constexpr key_t kInvalidKey = static_cast(-1); +class DistCalculator; + +struct GraphHeader { + uint32_t size; + uint32_t version; + uint32_t graph_type; + uint32_t doc_count; + uint32_t vector_size; + uint32_t node_size; + uint32_t l0_neighbor_count; + uint32_t prune_type; + uint32_t prune_neighbor_count; + uint32_t ef_construction; + uint32_t options; + uint32_t min_neighbor_count; + uint8_t reserved_[4080]; +}; + +static_assert(sizeof(GraphHeader) % 32 == 0, + "GraphHeader must be aligned with 32 bytes"); + +//! Hnsw upper neighbor header +struct HnswHeader { + uint32_t size; // header size + uint32_t revision; // current total docs of the graph + uint32_t upper_neighbor_count; + uint32_t ef_construction; + uint32_t scaling_factor; + uint32_t max_level; + uint32_t entry_point; + uint32_t options; + uint8_t reserved_[30]; +}; + +static_assert(sizeof(HnswHeader) % 32 == 0, + "GraphHeader must be aligned with 32 bytes"); + +//! Hnsw common header and upper neighbor header +struct HNSWHeader { + HNSWHeader() { + clear(); + } + + HNSWHeader(const HNSWHeader &header) { + memcpy(static_cast(this), &header, sizeof(header)); + } + + HNSWHeader &operator=(const HNSWHeader &header) { + memcpy(static_cast(this), &header, sizeof(header)); + return *this; + } + + //! Reset state to zero, and the params is untouched + void inline reset() { + graph.doc_count = 0U; + hnsw.entry_point = kInvalidNodeId; + hnsw.max_level = 0; + } + + //! Clear all fields to init value + void inline clear() { + memset(static_cast(this), 0, sizeof(HNSWHeader)); + hnsw.entry_point = kInvalidNodeId; + graph.size = sizeof(GraphHeader); + hnsw.size = sizeof(HnswHeader); + } + + size_t l0_neighbor_cnt() const { + return graph.l0_neighbor_count; + } + + size_t upper_neighbor_cnt() const { + return hnsw.upper_neighbor_count; + } + + size_t vector_size() const { + return graph.vector_size; + } + + size_t ef_construction() const { + return graph.ef_construction; + } + + size_t scaling_factor() const { + return hnsw.scaling_factor; + } + + size_t neighbor_prune_cnt() const { + return graph.prune_neighbor_count; + } + + node_id_t entry_point() const { + return hnsw.entry_point; + } + + node_id_t doc_cnt() const { + return graph.doc_count; + } + + GraphHeader graph; + HnswHeader hnsw; +}; + +struct NeighborsHeader { + uint32_t neighbor_cnt; +#ifdef _MSC_VER + node_id_t neighbors[]; +#else + node_id_t neighbors[0]; +#endif +}; + +struct Neighbors { + Neighbors() : cnt{0}, data{nullptr} {} + + Neighbors(uint32_t cnt_in, const node_id_t *data_in) + : cnt{cnt_in}, data{data_in} {} + + Neighbors(const IndexStorage::MemoryBlock &mem_block) + : neighbor_block{mem_block} { + auto hd = reinterpret_cast(neighbor_block.data()); + cnt = hd->neighbor_cnt; + data = hd->neighbors; + } + + size_t size(void) const { + return cnt; + } + + const node_id_t &operator[](size_t idx) const { + return data[idx]; + } + + uint32_t cnt; + const node_id_t *data; + IndexStorage::MemoryBlock neighbor_block; +}; + +//! Lightweight MemoryBlock for mmap mode: zero-cost construction/destruction +struct MmapMemoryBlock { + MmapMemoryBlock() = default; + explicit MmapMemoryBlock(void *data) : data_(data) {} + + MmapMemoryBlock(const MmapMemoryBlock &) = default; + MmapMemoryBlock &operator=(const MmapMemoryBlock &) = default; + MmapMemoryBlock(MmapMemoryBlock &&) = default; + MmapMemoryBlock &operator=(MmapMemoryBlock &&) = default; + ~MmapMemoryBlock() = default; + + const void *data() const { + return data_; + } + + void reset(void *data) { + data_ = data; + } + + void *data_{nullptr}; +}; + +//! Lightweight MemoryBlock for buffer pool mode: release on destruction +struct BufferPoolMemoryBlock { + BufferPoolMemoryBlock() = default; + + BufferPoolMemoryBlock(ailego::VecBufferPoolHandle *handle, size_t block_id, + void *data) + : buffer_pool_handle_(handle), buffer_block_id_(block_id), data_(data) {} + + static BufferPoolMemoryBlock MakeOwned(void *owned_data) { + BufferPoolMemoryBlock b; + b.owns_buffer_ = true; + b.data_ = owned_data; + return b; + } + + BufferPoolMemoryBlock(const BufferPoolMemoryBlock &rhs) + : buffer_pool_handle_(rhs.buffer_pool_handle_), + buffer_block_id_(rhs.buffer_block_id_), + data_(rhs.data_) { + if (rhs.owns_buffer_) { + owns_buffer_ = false; + buffer_pool_handle_ = nullptr; + } else if (buffer_pool_handle_) { + buffer_pool_handle_->acquire_one(buffer_block_id_); + } + } + + BufferPoolMemoryBlock &operator=(const BufferPoolMemoryBlock &rhs) { + if (this != &rhs) { + release(); + buffer_pool_handle_ = rhs.buffer_pool_handle_; + buffer_block_id_ = rhs.buffer_block_id_; + data_ = rhs.data_; + if (rhs.owns_buffer_) { + owns_buffer_ = false; + buffer_pool_handle_ = nullptr; + } else if (buffer_pool_handle_) { + buffer_pool_handle_->acquire_one(buffer_block_id_); + } + } + return *this; + } + + BufferPoolMemoryBlock(BufferPoolMemoryBlock &&rhs) noexcept + : buffer_pool_handle_(rhs.buffer_pool_handle_), + buffer_block_id_(rhs.buffer_block_id_), + owns_buffer_(rhs.owns_buffer_), + data_(rhs.data_) { + rhs.buffer_pool_handle_ = nullptr; + rhs.owns_buffer_ = false; + rhs.data_ = nullptr; + } + + BufferPoolMemoryBlock &operator=(BufferPoolMemoryBlock &&rhs) noexcept { + if (this != &rhs) { + release(); + buffer_pool_handle_ = rhs.buffer_pool_handle_; + buffer_block_id_ = rhs.buffer_block_id_; + owns_buffer_ = rhs.owns_buffer_; + data_ = rhs.data_; + rhs.buffer_pool_handle_ = nullptr; + rhs.owns_buffer_ = false; + rhs.data_ = nullptr; + } + return *this; + } + + ~BufferPoolMemoryBlock() { + release(); + } + + const void *data() const { + return data_; + } + + void reset(ailego::VecBufferPoolHandle *handle, size_t block_id, void *data) { + release(); + buffer_pool_handle_ = handle; + buffer_block_id_ = block_id; + data_ = data; + } + + private: + void release() { + if (owns_buffer_) { + if (data_) { + ailego_free(data_); + } + owns_buffer_ = false; + } else if (buffer_pool_handle_) { + buffer_pool_handle_->release_one(buffer_block_id_); + buffer_pool_handle_ = nullptr; + } + data_ = nullptr; + } + + ailego::VecBufferPoolHandle *buffer_pool_handle_{nullptr}; + size_t buffer_block_id_{0}; + bool owns_buffer_{false}; + void *data_{nullptr}; +}; + +//! Typed Neighbors: holds a typed MemoryBlock to avoid runtime branching +template +struct NeighborsT { + NeighborsT() : cnt{0}, data{nullptr} {} + + NeighborsT(uint32_t cnt_in, const node_id_t *data_in) + : cnt{cnt_in}, data{data_in} {} + + explicit NeighborsT(const MemBlockType &mem_block) + : neighbor_block{mem_block} { + auto hd = reinterpret_cast(neighbor_block.data()); + cnt = hd->neighbor_cnt; + data = hd->neighbors; + } + + explicit NeighborsT(MemBlockType &&mem_block) + : neighbor_block{std::move(mem_block)} { + auto hd = reinterpret_cast(neighbor_block.data()); + cnt = hd->neighbor_cnt; + data = hd->neighbors; + } + + size_t size(void) const { + return cnt; + } + + const node_id_t &operator[](size_t idx) const { + return data[idx]; + } + + uint32_t cnt; + const node_id_t *data; + MemBlockType neighbor_block; +}; + +//! level 0 neighbors offset +struct GraphNeighborMeta { + GraphNeighborMeta(size_t o, size_t cnt) : offset(o), neighbor_cnt(cnt) {} + + uint64_t offset : 48; + uint64_t neighbor_cnt : 16; +}; + +//! hnsw upper neighbors meta +struct HnswNeighborMeta { + HnswNeighborMeta(size_t o, size_t l) : offset(o), level(l) {} + + uint64_t offset : 48; // offset = idx * upper neighors size + uint64_t level : 16; +}; + +class HnswEntity { + public: + //! Constructor + HnswEntity() {} + + //! Constructor + HnswEntity(const HNSWHeader &hd) { + header_ = hd; + } + + //! Destructor + virtual ~HnswEntity() {} + + //! HnswEntity Pointerd; + typedef std::shared_ptr Pointer; + + //! Get max neighbor size of graph level + inline size_t neighbor_cnt(level_t level) const { + return level == 0 ? header_.graph.l0_neighbor_count + : header_.hnsw.upper_neighbor_count; + } + + //! get max neighbor size of graph level 0 + inline size_t l0_neighbor_cnt() const { + return header_.graph.l0_neighbor_count; + } + + //! get min neighbor size of graph + inline size_t min_neighbor_cnt() const { + return header_.graph.min_neighbor_count; + } + + //! get upper neighbor size of graph level other than 0 + inline size_t upper_neighbor_cnt() const { + return header_.hnsw.upper_neighbor_count; + } + + //! Get current total doc of the hnsw graph + inline node_id_t *mutable_doc_cnt() { + return &header_.graph.doc_count; + } + + inline node_id_t doc_cnt() const { + return header_.graph.doc_count; + } + + //! Get hnsw graph scaling params + inline size_t scaling_factor() const { + return header_.hnsw.scaling_factor; + } + + //! Get prune_size + inline size_t prune_cnt() const { + return header_.graph.prune_neighbor_count; + } + + //! Current entity of top level graph + inline node_id_t entry_point() const { + return header_.hnsw.entry_point; + } + + //! Current max graph level + inline level_t cur_max_level() const { + return header_.hnsw.max_level; + } + + //! Retrieve index vector size + size_t vector_size() const { + return header_.graph.vector_size; + } + + //! Retrieve node size + size_t node_size() const { + return header_.graph.node_size; + } + + //! Retrieve ef constuction + size_t ef_construction() const { + return header_.graph.ef_construction; + } + + void set_vector_size(size_t size) { + header_.graph.vector_size = size; + } + + void set_prune_cnt(size_t v) { + header_.graph.prune_neighbor_count = v; + } + + void set_scaling_factor(size_t val) { + header_.hnsw.scaling_factor = val; + } + + void set_l0_neighbor_cnt(size_t cnt) { + header_.graph.l0_neighbor_count = cnt; + } + + void set_min_neighbor_cnt(size_t cnt) { + header_.graph.min_neighbor_count = cnt; + } + + void set_upper_neighbor_cnt(size_t cnt) { + header_.hnsw.upper_neighbor_count = cnt; + } + + void set_ef_construction(size_t ef) { + header_.graph.ef_construction = ef; + } + + protected: + inline const HNSWHeader &header() const { + return header_; + } + + inline HNSWHeader *mutable_header() { + return &header_; + } + + inline size_t header_size() const { + return sizeof(header_); + } + + void set_node_size(size_t size) { + header_.graph.node_size = size; + } + + //! Dump all segment by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_segments( + const IndexDumper::Pointer &dumper, key_t *keys, + const std::function &get_level) const; + + private: + //! dump mapping segment, for get_vector_by_key in provider + int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const; + + //! dump hnsw head by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_header(const IndexDumper::Pointer &dumper, + const HNSWHeader &hd) const; + + //! dump vectors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_vectors(const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + //! dump hnsw neighbors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_neighbors(const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + auto len1 = dump_graph_neighbors(dumper, reorder_mapping, neighbor_mapping); + if (len1 < 0) { + return len1; + } + auto len2 = dump_upper_neighbors(dumper, get_level, reorder_mapping, + neighbor_mapping); + if (len2 < 0) { + return len2; + } + + return len1 + len2; + } + + //! dump segment by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, const void *data, + size_t size) const; + + //! Dump level 0 neighbors + //! Return dump size if success, errno(<0) in failure + int64_t dump_graph_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + //! Dump upper level neighbors + //! Return dump size if success, errno(<0) in failure + int64_t dump_upper_neighbors( + const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + public: + //! Cleanup the entity + virtual int cleanup(void) { + header_.clear(); + return 0; + } + + //! Make a copy of searcher entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + virtual const HnswEntity::Pointer clone() const { + LOG_ERROR("Update neighbors not implemented"); + return HnswEntity::Pointer(); + } + + //! Get primary key of the node id + virtual key_t get_key(node_id_t id) const = 0; + + //! Get vector feature data by key + virtual const void *get_vector(node_id_t id) const = 0; + + //! Get vectors feature data by keys + virtual int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const = 0; + + virtual int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const = 0; + virtual int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const = 0; + + //! Retrieve a vector using a primary key + virtual const void *get_vector_by_key(uint64_t /*key*/) const { + LOG_ERROR("get vector not implemented"); + return nullptr; + } + + virtual int get_vector_by_key(const key_t /*key*/, + IndexStorage::MemoryBlock & /*block*/) const { + return IndexError_NotImplemented; + } + + //! Get the node id's neighbors on graph level + //! Note: the neighbors cannot be modified, using the following + //! method to get WritableNeighbors if want to + virtual const Neighbors get_neighbors(level_t level, node_id_t id) const = 0; + + //! Add vector and key to hnsw entity, and local id will be saved in id + virtual int add_vector(level_t /*level*/, key_t /*key*/, const void * /*vec*/, + node_id_t * /*id*/) { + return IndexError_NotImplemented; + } + + //! Add vector and id to hnsw entity + virtual int add_vector_with_id(level_t /*level*/, node_id_t /*id*/, + const void * /*vec*/) { + return IndexError_NotImplemented; + } + + virtual int update_neighbors( + level_t /*level*/, node_id_t /*id*/, + const std::vector> & /*neighbors*/) { + LOG_ERROR("Update neighbors dense not implemented"); + + return 0; + } + + //! Append neighbor_id to node id neighbors on level, size is the current + //! neighbors size. Notice: the caller must be ensure the neighbors not full + virtual void add_neighbor(level_t /*level*/, node_id_t /*id*/, + uint32_t /*size*/, node_id_t /*neighbor_id*/) { + LOG_ERROR("Add neighbor not implemented"); + } + + //! Update entry point and max level + virtual void update_ep_and_level(node_id_t ep, level_t level) { + header_.hnsw.entry_point = ep; + header_.hnsw.max_level = level; + } + + //! Bind an external vector source to this entity. The default + //! implementation is a no-op; only entities that read vectors from an + //! external source (e.g. HnswExternalStreamerEntity) override it. + virtual void set_vector_source(const VectorSource * /*src*/) {} + + virtual int load(const IndexStorage::Pointer & /*container*/, + bool /*check_crc*/) { + LOG_ERROR("Load not implemented"); + return IndexError_NotImplemented; + } + + virtual int dump(const IndexDumper::Pointer & /*dumper*/) { + LOG_ERROR("Dump not implemented"); + return IndexError_NotImplemented; + } + + static int CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size); + + protected: + static inline size_t AlignSize(size_t size) { + return (size + 0x1F) & (~0x1F); + } + + static inline size_t AlignPageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::PageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + static inline size_t AlignHugePageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::HugePageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + //! rearrange vectors to improve cache locality + void reshuffle_vectors(const std::function &get_level, + std::vector *n2o_mapping, + std::vector *o2n_mapping, + key_t *keys) const; + + public: + const static std::string kGraphHeaderSegmentId; + const static std::string kGraphFeaturesSegmentId; + const static std::string kGraphKeysSegmentId; + const static std::string kGraphNeighborsSegmentId; + const static std::string kGraphOffsetsSegmentId; + const static std::string kGraphMappingSegmentId; + const static std::string kHnswHeaderSegmentId; + const static std::string kHnswNeighborsSegmentId; + const static std::string kHnswOffsetsSegmentId; + + constexpr static uint32_t kRevision = 0U; + constexpr static size_t kMaxGraphLayers = 15; + constexpr static uint32_t kDefaultEfConstruction = 500; + constexpr static uint32_t kDefaultEf = 500; + constexpr static uint32_t kDefaultUpperMaxNeighborCnt = 50; // M of HNSW + constexpr static uint32_t kDefaultL0MaxNeighborCnt = 100; + constexpr static uint32_t kMaxNeighborCnt = 65535; + constexpr static float kDefaultScanRatio = 0.1f; + constexpr static uint32_t kDefaultMinScanLimit = 10000; + constexpr static uint32_t kDefaultMaxScanLimit = + std::numeric_limits::max(); + constexpr static float kDefaultBFNegativeProbability = 0.001f; + constexpr static uint32_t kDefaultScalingFactor = 50U; + constexpr static uint32_t kDefaultBruteForceThreshold = 1000U; + constexpr static uint32_t kDefaultDocsHardLimit = 1 << 30U; // 1 billion + constexpr static float kDefaultDocsSoftLimitRatio = 0.9f; + constexpr static size_t kMaxChunkSize = 0xFFFFFFFF; + constexpr static size_t kDefaultChunkSize = 2 * 1024UL * 1024UL; + constexpr static size_t kDefaultMaxChunkCnt = 50000UL; + constexpr static float kDefaultNeighborPruneMultiplier = + 1.0f; // prune_cnt = upper_max_neighbor_cnt * multiplier + constexpr static float kDefaultL0MaxNeighborCntMultiplier = + 2.0f; // l0_max_neighbor_cnt = upper_max_neighbor_cnt * multiplier + + protected: + HNSWHeader header_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_index_hash.h b/src/core/algorithm/hnsw/hnsw_index_hash.h new file mode 100644 index 0000000..cc59e84 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_index_hash.h @@ -0,0 +1,226 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "hnsw_chunk.h" + +namespace zvec { +namespace core { + +//! Persistent hashmap implement through open addressing algorithm +template ::value>::type> +class HnswIndexHashMap { + using key_type = Key; + using val_type = Val; + + struct Iterator { + key_type first; + val_type second; + }; + typedef Iterator *iterator; + typedef Iterator Item; + typedef const Iterator *const_iterator; + + class Slot { + public: + Slot(Chunk::Pointer &&chunk, const void *data) + : chunk_(std::move(chunk)), + items_(reinterpret_cast(data)) {} + //! Return a empty loc or the key item loc + + Slot(Chunk::Pointer &&chunk, std::vector &&local_data) + : chunk_(std::move(chunk)), local_data_(std::move(local_data)) { + items_ = reinterpret_cast(local_data_.data()); + } + const_iterator find(key_type key, uint32_t max_items, uint32_t mask) const { + auto it = &items_[key & mask]; + for (auto i = 0U; i < max_items; ++i) { + if (it->first == key || it->second == EmptyVal) { + // LOG_DEBUG("i=%u", i); + return it; + } + ++it; + if (it == &items_[max_items]) { + it = &items_[0]; + } + } + return nullptr; + } + + bool update(const_iterator it) { + uint32_t offset = reinterpret_cast(it) - + reinterpret_cast(&items_[0]); + if (ailego_unlikely(chunk_->write(offset, it, sizeof(Item)) != + sizeof(Item))) { + LOG_ERROR("Chunk write failed"); + return false; + } + return true; + } + + private: + Chunk::Pointer chunk_{}; + const Item *items_{nullptr}; // point to local_data_ + std::vector local_data_{}; + }; + + public: + //! Init the hash + //! broker the index allocator + //! chunk_size the size of per chunk allocated, actual size may greater + //! factor factor = 1/ratio, ratio is the probability of a squence + //! number inserted to this container + //! max the max number key can be inserted + //! expansion_ratio memory expansion ratio + int init(ChunkBroker::Pointer &broker, uint32_t chunk_size, uint32_t factor, + size_t max, float expansion_ratio) { + ailego_assert_with(expansion_ratio > 1.0f, "ratio must > 1.0f"); + broker_ = broker; + + size_t items = std::ceil(chunk_size * 1.0f / sizeof(Item)); + slot_items_ = 1UL << static_cast((std::ceil(std::log2(items)))); + size_t range = slot_items_ * factor / expansion_ratio; + mask_bits_ = std::floor(std::log2(range)); + range = 1UL << mask_bits_; + size_t max_slots = std::ceil(max * 1.0f / range); + slots_.reserve(max_slots); + slot_loc_mask_ = slot_items_ - 1U; + int ret = load(); + if (ret != 0) { + return ret; + } + + LOG_DEBUG( + "HnswIndexHash init, chunkSize=%u factor=%u max=%zu " + "ratio=%f slotItems=%u maxSlots=%zu maskBits=%u " + "range=%zu", + chunk_size, factor, max, expansion_ratio, slot_items_, max_slots, + mask_bits_, range); + + return 0; + } + + int cleanup(void) { + slots_.clear(); + slots_.shrink_to_fit(); + broker_.reset(); + mask_bits_ = 0U; + slot_items_ = 0U; + slot_loc_mask_ = 0U; + + return 0; + } + + const_iterator end(void) const { + return nullptr; + } + + const_iterator find(const key_type key) const { + auto idx = key >> mask_bits_; + if (idx >= slots_.size()) { + return end(); + } + auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_); + return it && it->second != EmptyVal ? it : nullptr; + } + + bool insert(key_type key, val_type val) { + auto idx = key >> mask_bits_; + if (idx >= slots_.size()) { + if (ailego_unlikely(idx >= slots_.capacity())) { + return false; + } + for (auto i = slots_.size(); i <= idx; ++i) { + if (ailego_unlikely(!alloc_slot(i))) { + return false; + } + } + } + auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_); + if (ailego_unlikely(it == nullptr)) { + return false; + } + + //! TODO: write memory is ok? + const_cast(it)->first = key; + const_cast(it)->second = val; + + return slots_[idx].update(it); + } + + private: + bool alloc_slot(size_t idx) { + ailego_assert_with(idx == slots_.size(), "invalid idx"); + + size_t size = slot_items_ * sizeof(Item); + auto p = + broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, idx, size); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return false; + } + Chunk::Pointer chunk = p.second; + if (ailego_unlikely(chunk->resize(size) != size)) { + LOG_ERROR("Chunk resize failed, size=%zu", size); + return false; + } + //! Use a local zero-initialized buffer; new chunks contain all zeros, + //! so no buffer-pool read is needed and no ref_count is pinned. + std::vector local_buf(size, 0); + slots_.emplace_back(std::move(chunk), std::move(local_buf)); + return true; + } + + int load(void) { + size_t slots_cnt = + broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX); + for (size_t i = 0UL; i < slots_cnt; ++i) { + auto chunk = + broker_->get_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, i); + if (!chunk) { + LOG_ERROR("Get chunk failed, seq=%zu", i); + return IndexError_InvalidFormat; + } + size_t size = sizeof(Item) * slot_items_; + if (chunk->data_size() < size) { + LOG_ERROR( + "Hash params may be mismatch, seq=%zu, data_size=%zu " + "expect=%zu", + i, chunk->data_size(), size); + return IndexError_InvalidFormat; + } + //! Copy chunk data into a local buffer via fetch() so that no + //! buffer-pool block is pinned for the lifetime of the Slot. + std::vector local_buf(size); + if (ailego_unlikely(chunk->fetch(0U, local_buf.data(), size) != size)) { + LOG_ERROR("Chunk fetch failed, size=%zu", size); + return IndexError_InvalidFormat; + } + slots_.emplace_back(std::move(chunk), std::move(local_buf)); + } + return 0; + } + + private: + ChunkBroker::Pointer broker_{}; // chunk broker + std::vector slots_{}; + uint32_t mask_bits_{0U}; + uint32_t slot_items_{}; // must be a power of 2 + uint32_t slot_loc_mask_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_index_provider.h b/src/core/algorithm/hnsw/hnsw_index_provider.h new file mode 100644 index 0000000..be96e1a --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_index_provider.h @@ -0,0 +1,133 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "hnsw_entity.h" + +namespace zvec { +namespace core { + +class HnswIndexProvider : public IndexProvider { + public: + HnswIndexProvider(const IndexMeta &meta, const HnswEntity::Pointer &entity, + const std::string &owner) + : meta_(meta), entity_(entity), owner_class_(owner) {} + + HnswIndexProvider(const HnswIndexProvider &) = delete; + HnswIndexProvider &operator=(const HnswIndexProvider &) = delete; + + public: // holder interface + //! Create a new iterator + IndexProvider::Iterator::Pointer create_iterator() override { + return HnswIndexProvider::Iterator::Pointer(new (std::nothrow) + Iterator(entity_)); + } + + //! Retrieve count of vectors + size_t count(void) const override { + return entity_->doc_cnt(); + } + + //! Retrieve dimension of vector + size_t dimension(void) const override { + return meta_.dimension(); + } + + //! Retrieve type of vector + IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + //! Retrieve vector size in bytes + size_t element_size(void) const override { + return meta_.element_size(); + } + + public: // provider's unique interface + //! Retrieve a vector using a primary key + const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector_by_key(key, block); + } + + //! Retrieve the owner class + const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class Iterator : public IndexProvider::Iterator { + public: + Iterator(const HnswEntity::Pointer &entity) + : entity_(entity), cur_id_(0U) {} + + //! Retrieve pointer of data + //! NOTICE: the vec feature will be changed after iterating to next, so + //! the caller need to keep a copy of it before iterator to next vector + const void *data(void) const override { + return entity_->get_vector(cur_id_); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return entity_->get_key(cur_id_); + } + + //! Next iterator + void next(void) override { + // cur_id_ += 1; + cur_id_ = get_next_valid_id(cur_id_ + 1); + } + + //! Reset the iterator + void reset(void) { + cur_id_ = get_next_valid_id(0); + } + + private: + node_id_t get_next_valid_id(node_id_t start_id) { + for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) { + if (entity_->get_key(i) != kInvalidNodeId) { + cur_id_ = i; + return i; + } + } + return kInvalidNodeId; + } + + private: + const HnswEntity::Pointer entity_; + node_id_t cur_id_; + }; + + private: + const IndexMeta &meta_; + const HnswEntity::Pointer entity_; + const std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_params.h b/src/core/algorithm/hnsw/hnsw_params.h new file mode 100644 index 0000000..2cd2ecf --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_params.h @@ -0,0 +1,122 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +static const std::string PARAM_HNSW_BUILDER_THREAD_COUNT( + "proxima.hnsw.builder.thread_count"); +static const std::string PARAM_HNSW_BUILDER_MEMORY_QUOTA( + "proxima.hnsw.builder.memory_quota"); +static const std::string PARAM_HNSW_BUILDER_EFCONSTRUCTION( + "proxima.hnsw.builder.efconstruction"); +static const std::string PARAM_HNSW_BUILDER_SCALING_FACTOR( + "proxima.hnsw.builder.scaling_factor"); +static const std::string PARAM_HNSW_BUILDER_CHECK_INTERVAL_SECS( + "proxima.hnsw.builder.check_interval_secs"); +static const std::string PARAM_HNSW_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER( + "proxima.hnsw.builder.neighbor_prune_multiplier"); +static const std::string PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT( + "proxima.hnsw.builder.min_neighbor_count"); +static const std::string PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT( + "proxima.hnsw.builder.max_neighbor_count"); +static const std::string PARAM_HNSW_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER( + "proxima.hnsw.builder.l0_max_neighbor_count_multiplier"); + +static const std::string PARAM_HNSW_SEARCHER_EF("proxima.hnsw.searcher.ef"); +static const std::string PARAM_HNSW_SEARCHER_PO("proxima.hnsw.searcher.po"); +static const std::string PARAM_HNSW_SEARCHER_PL("proxima.hnsw.searcher.pl"); +static const std::string PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD( + "proxima.hnsw.searcher.brute_force_threshold"); +static const std::string PARAM_HNSW_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE( + "proxima.hnsw.searcher.neighbors_in_memory_enable"); +static const std::string PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO( + "proxima.hnsw.searcher.max_scan_ratio"); +static const std::string PARAM_HNSW_SEARCHER_CHECK_CRC_ENABLE( + "proxima.hnsw.searcher.check_crc_enable"); +static const std::string PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_ENABLE( + "proxima.hnsw.searcher.visit_bloomfilter_enable"); +static const std::string PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB( + "proxima.hnsw.searcher.visit_bloomfilter_negative_prob"); +static const std::string PARAM_HNSW_SEARCHER_FORCE_PADDING_RESULT_ENABLE( + "proxima.hnsw.searcher.force_padding_result_enable"); + +static const std::string PARAM_HNSW_STREAMER_MAX_SCAN_RATIO( + "proxima.hnsw.streamer.max_scan_ratio"); +static const std::string PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT( + "proxima.hnsw.streamer.min_scan_limit"); +static const std::string PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT( + "proxima.hnsw.streamer.max_scan_limit"); +static const std::string PARAM_HNSW_STREAMER_EF("proxima.hnsw.streamer.ef"); +static const std::string PARAM_HNSW_STREAMER_PO("proxima.hnsw.streamer.po"); +static const std::string PARAM_HNSW_STREAMER_PL("proxima.hnsw.streamer.pl"); +static const std::string PARAM_HNSW_STREAMER_EFCONSTRUCTION( + "proxima.hnsw.streamer.efconstruction"); +static const std::string PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT( + "proxima.hnsw.streamer.max_neighbor_count"); +static const std::string PARAM_HNSW_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER( + "proxima.hnsw.streamer.l0_max_neighbor_count_multiplier"); +static const std::string PARAM_HNSW_STREAMER_SCALING_FACTOR( + "proxima.hnsw.streamer.scaling_factor"); +static const std::string PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD( + "proxima.hnsw.streamer.brute_force_threshold"); +static const std::string PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT( + "proxima.hnsw.streamer.docs_hard_limit"); +static const std::string PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT( + "proxima.hnsw.streamer.docs_soft_limit"); +static const std::string PARAM_HNSW_STREAMER_MAX_INDEX_SIZE( + "proxima.hnsw.streamer.max_index_size"); +static const std::string PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_ENABLE( + "proxima.hnsw.streamer.visit_bloomfilter_enable"); +static const std::string PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB( + "proxima.hnsw.streamer.visit_bloomfilter_negative_prob"); +static const std::string PARAM_HNSW_STREAMER_CHECK_CRC_ENABLE( + "proxima.hnsw.streamer.check_crc_enable"); +static const std::string PARAM_HNSW_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER( + "proxima.hnsw.streamer.neighbor_prune_multiplier"); +static const std::string PARAM_HNSW_STREAMER_CHUNK_SIZE( + "proxima.hnsw.streamer.chunk_size"); +static const std::string PARAM_HNSW_STREAMER_FILTER_SAME_KEY( + "proxima.hnsw.streamer.filter_same_key"); +static const std::string PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE( + "proxima.hnsw.streamer.get_vector_enable"); +static const std::string PARAM_HNSW_STREAMER_MIN_NEIGHBOR_COUNT( + "proxima.hnsw.streamer.min_neighbor_count"); +static const std::string PARAM_HNSW_STREAMER_FORCE_PADDING_RESULT_ENABLE( + "proxima.hnsw.streamer.force_padding_result_enable"); +static const std::string PARAM_HNSW_STREAMER_ESTIMATE_DOC_COUNT( + "proxima.hnsw.streamer.estimate_doc_count"); +static const std::string PARAM_HNSW_STREAMER_USE_ID_MAP( + "proxima.hnsw.streamer.use_id_map"); + +static const std::string PARAM_HNSW_REDUCER_WORKING_PATH( + "proxima.hnsw.reducer.working_path"); +static const std::string PARAM_HNSW_REDUCER_NUM_OF_ADD_THREADS( + "proxima.hnsw.reducer.num_of_add_threads"); +static const std::string PARAM_HNSW_REDUCER_INDEX_NAME( + "proxima.hnsw.reducer.index_name"); +static const std::string PARAM_HNSW_REDUCER_EFCONSTRUCTION( + "proxima.hnsw.reducer.efconstruction"); + +static const std::string PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY( + "proxima.hnsw.streamer.use_contiguous_memory"); + +static const std::string PARAM_HNSW_STREAMER_USE_EXTERNAL_VECTOR( + "proxima.hnsw.streamer.use_external_vector"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc new file mode 100644 index 0000000..8bd03f2 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -0,0 +1,927 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_streamer.h" +#include +#include +#include +#include +#include "utility/sparse_utility.h" +#include "hnsw_algorithm.h" +#include "hnsw_context.h" +#include "hnsw_dist_calculator.h" +#include "hnsw_index_provider.h" + +namespace zvec { +namespace core { + +HnswStreamer::HnswStreamer() = default; + +HnswStreamer::~HnswStreamer() { + if (state_ == STATE_INITED || state_ == STATE_OPENED) { + this->cleanup(); + } +} + +int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { + meta_ = imeta; + meta_.set_streamer("HnswStreamer", HnswEntity::kRevision, params); + + params.get(PARAM_HNSW_STREAMER_MAX_INDEX_SIZE, &max_index_size_); + + params.get(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, &upper_max_neighbor_cnt_); + float multiplier = HnswEntity::kDefaultL0MaxNeighborCntMultiplier; + params.get(PARAM_HNSW_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, &multiplier); + l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_; + + multiplier = HnswEntity::kDefaultNeighborPruneMultiplier; + params.get(PARAM_HNSW_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier); + size_t prune_cnt = multiplier * upper_max_neighbor_cnt_; + scaling_factor_ = upper_max_neighbor_cnt_; + params.get(PARAM_HNSW_STREAMER_SCALING_FACTOR, &scaling_factor_); + + params.get(PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_); + params.get(PARAM_HNSW_STREAMER_EF, &ef_); + params.get(PARAM_HNSW_STREAMER_PO, &po_); + params.get(PARAM_HNSW_STREAMER_PL, &pl_); + params.get(PARAM_HNSW_STREAMER_EFCONSTRUCTION, &ef_construction_); + params.get(PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_); + params.get(PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB, + &bf_negative_prob_); + params.get(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, &bruteforce_threshold_); + params.get(PARAM_HNSW_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_); + params.get(PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_); + params.get(PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_); + params.get(PARAM_HNSW_STREAMER_CHECK_CRC_ENABLE, &check_crc_enabled_); + params.get(PARAM_HNSW_STREAMER_CHUNK_SIZE, &chunk_size_); + params.get(PARAM_HNSW_STREAMER_FILTER_SAME_KEY, &filter_same_key_); + params.get(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, &get_vector_enabled_); + params.get(PARAM_HNSW_STREAMER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_); + params.get(PARAM_HNSW_STREAMER_FORCE_PADDING_RESULT_ENABLE, + &force_padding_topk_enabled_); + params.get(PARAM_HNSW_STREAMER_USE_ID_MAP, &use_id_map_); + params.get(PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY, + &use_contiguous_memory_); + params.get(PARAM_HNSW_STREAMER_USE_EXTERNAL_VECTOR, &use_external_vector_); + + params.get(PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit_); + if (docs_soft_limit_ > 0 && docs_soft_limit_ > docs_hard_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT.c_str(), + PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT.c_str()); + return IndexError_InvalidArgument; + } else if (docs_soft_limit_ == 0UL) { + docs_soft_limit_ = + docs_hard_limit_ * HnswEntity::kDefaultDocsSoftLimitRatio; + } + + if (ef_ == 0U) { + ef_ = HnswEntity::kDefaultEf; + } + if (ef_construction_ == 0U) { + ef_construction_ = HnswEntity::kDefaultEfConstruction; + } + if (upper_max_neighbor_cnt_ == 0U) { + upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt; + } + if (upper_max_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) { + LOG_ERROR("[%s] must be in range (0,%d)", + PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT.c_str(), + HnswEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (l0_max_neighbor_cnt_ == 0U) { + l0_max_neighbor_cnt_ = HnswEntity::kDefaultL0MaxNeighborCnt; + } + if (l0_max_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) { + LOG_ERROR("MaxL0NeighborCnt must be in range (0,%d)", + HnswEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) { + LOG_ERROR("[%s]-[%u] must be <= [%s]-[%u]", + PARAM_HNSW_STREAMER_MIN_NEIGHBOR_COUNT.c_str(), min_neighbor_cnt_, + PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT.c_str(), + upper_max_neighbor_cnt_); + return IndexError_InvalidArgument; + } + + if (bf_negative_prob_ <= 0.0f || bf_negative_prob_ >= 1.0f) { + LOG_ERROR("[%s] must be in range (0,1)", + PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str()); + return IndexError_InvalidArgument; + } + + if (scaling_factor_ == 0U) { + scaling_factor_ = HnswEntity::kDefaultScalingFactor; + } + if (scaling_factor_ < 5 || scaling_factor_ > 1000) { + LOG_ERROR("[%s] must be in range [5,1000]", + PARAM_HNSW_STREAMER_SCALING_FACTOR.c_str()); + return IndexError_InvalidArgument; + } + + if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) { + LOG_ERROR("[%s] must be in range (0.0f,1.0f]", + PARAM_HNSW_STREAMER_MAX_SCAN_RATIO.c_str()); + return IndexError_InvalidArgument; + } + + if (max_scan_limit_ < min_scan_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT.c_str(), + PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT.c_str()); + return IndexError_InvalidArgument; + } + + if (prune_cnt == 0UL) { + prune_cnt = upper_max_neighbor_cnt_; + } + prune_cnt_ = prune_cnt; + if (chunk_size_ == 0UL) { + chunk_size_ = HnswEntity::kDefaultChunkSize; + } + if (chunk_size_ > HnswEntity::kMaxChunkSize) { + LOG_ERROR("[%s] must be < %zu", PARAM_HNSW_STREAMER_CHUNK_SIZE.c_str(), + HnswEntity::kMaxChunkSize); + return IndexError_InvalidArgument; + } + + LOG_DEBUG( + "Init params: maxIndexSize=%zu docsHardLimit=%zu docsSoftLimit=%zu " + "efConstruction=%u ef=%u upperMaxNeighborCnt=%u l0MaxNeighborCnt=%u " + "scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu " + "bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbability=%.5f " + "checkCrcEnabled=%d pruneSize=%u vectorSize=%u chunkSize=%zu " + "filterSameKey=%u getVectorEnabled=%u minNeighborCount=%u " + "forcePadding=%u ", + max_index_size_, docs_hard_limit_, docs_soft_limit_, ef_construction_, + ef_, upper_max_neighbor_cnt_, l0_max_neighbor_cnt_, scaling_factor_, + max_scan_ratio_, min_scan_limit_, max_scan_limit_, bf_enabled_, + bruteforce_threshold_, bf_negative_prob_, check_crc_enabled_, prune_cnt_, + meta_.element_size(), chunk_size_, filter_same_key_, get_vector_enabled_, + min_neighbor_cnt_, force_padding_topk_enabled_); + + state_ = STATE_INITED; + + return 0; +} + +int HnswStreamer::cleanup(void) { + if (state_ == STATE_OPENED) { + this->close(); + } + + LOG_INFO("HnswStreamer cleanup"); + + meta_.clear(); + metric_.reset(); + stats_.clear(); + if (entity_) { + entity_->cleanup(); + } + + if (alg_) { + alg_->cleanup(); + } + + max_index_size_ = 0UL; + docs_hard_limit_ = HnswEntity::kDefaultDocsHardLimit; + docs_soft_limit_ = 0UL; + upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt; + l0_max_neighbor_cnt_ = HnswEntity::kDefaultL0MaxNeighborCnt; + ef_ = HnswEntity::kDefaultEf; + ef_construction_ = HnswEntity::kDefaultEfConstruction; + bf_enabled_ = false; + scaling_factor_ = HnswEntity::kDefaultScalingFactor; + bruteforce_threshold_ = HnswEntity::kDefaultBruteForceThreshold; + max_scan_limit_ = HnswEntity::kDefaultMaxScanLimit; + min_scan_limit_ = HnswEntity::kDefaultMinScanLimit; + chunk_size_ = HnswEntity::kDefaultChunkSize; + bf_negative_prob_ = HnswEntity::kDefaultBFNegativeProbability; + max_scan_ratio_ = HnswEntity::kDefaultScanRatio; + state_ = STATE_INIT; + check_crc_enabled_ = false; + filter_same_key_ = false; + get_vector_enabled_ = false; + + return 0; +} + +int HnswStreamer::setup_entity() { + entity_->set_use_key_info_map(use_id_map_); + entity_->set_ef_construction(ef_construction_); + entity_->set_upper_neighbor_cnt(upper_max_neighbor_cnt_); + entity_->set_l0_neighbor_cnt(l0_max_neighbor_cnt_); + entity_->set_scaling_factor(scaling_factor_); + entity_->set_prune_cnt(prune_cnt_); + // For external-vector entities the per-node vector prefix is removed; set + // vector_size to 0 so all inherited offset computations (key / neighbors / + // node_size) are correct and add_vector writes no vector bytes. The distance + // dimension is taken from meta.dimension(), not from vector_size(). + entity_->set_vector_size(use_external_vector_ ? 0 : meta_.element_size()); + entity_->set_chunk_size(chunk_size_); + entity_->set_filter_same_key(filter_same_key_); + entity_->set_get_vector(get_vector_enabled_); + entity_->set_min_neighbor_cnt(min_neighbor_cnt_); + + int ret = entity_->init(docs_hard_limit_); + if (ret != 0) { + LOG_ERROR("Hnsw entity init failed for %s", IndexError::What(ret)); + } + return ret; +} + +int HnswStreamer::open(IndexStorage::Pointer stg) { + LOG_INFO("HnswStreamer open"); + + if (ailego_unlikely(state_ != STATE_INITED)) { + LOG_ERROR("Open storage failed, init streamer first!"); + return IndexError_NoReady; + } + + // Create entity based on storage type + switch (stg->memory_block_type()) { + case IndexStorage::MemoryBlock::MBT_BUFFERPOOL: { + entity_ = std::make_unique(stats_); + break; + } + default: { + if (use_external_vector_) { + entity_ = std::make_unique(stats_); + } else if (use_contiguous_memory_) { + entity_ = std::make_unique(stats_); + } else { + entity_ = std::make_unique(stats_); + } + break; + } + } + int ret = setup_entity(); + if (ret != 0) { + return ret; + } + + ret = entity_->open(std::move(stg), max_index_size_, check_crc_enabled_); + if (ret != 0) { + return ret; + } + IndexMeta index_meta; + ret = entity_->get_index_meta(&index_meta); + if (ret == IndexError_NoExist) { + // Set IndexMeta for the new index + ret = entity_->set_index_meta(meta_); + if (ret != 0) { + LOG_ERROR("Failed to set index meta for %s", IndexError::What(ret)); + return ret; + } + } else if (ret != 0) { + LOG_ERROR("Failed to get index meta for %s", IndexError::What(ret)); + return ret; + } else { + if (index_meta.dimension() != meta_.dimension() || + index_meta.element_size() != meta_.element_size() || + index_meta.metric_name() != meta_.metric_name() || + index_meta.data_type() != meta_.data_type()) { + LOG_ERROR("IndexMeta mismatch from the previous in index"); + return IndexError_Mismatch; + } + // The IndexMetric Params may be updated like MipsSquaredEuclidean + auto metric_params = index_meta.metric_params(); + metric_params.merge(meta_.metric_params()); + meta_.set_metric(index_meta.metric_name(), 0, metric_params); + + // Restore converter/reformer from the persisted meta + if (!index_meta.reformer_name().empty()) { + meta_.set_reformer(index_meta.reformer_name(), + index_meta.reformer_revision(), + index_meta.reformer_params()); + } + if (!index_meta.converter_name().empty()) { + meta_.set_converter(index_meta.converter_name(), + index_meta.converter_revision(), + index_meta.converter_params()); + } + } + + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Failed to create metric %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("Failed to init metric, ret=%d", ret); + return ret; + } + + if (!metric_->distance()) { + LOG_ERROR("Invalid metric distance"); + return IndexError_InvalidArgument; + } + + if (!metric_->batch_distance()) { + LOG_ERROR("Invalid metric batch distance"); + return IndexError_InvalidArgument; + } + + add_distance_ = metric_->distance(); + add_batch_distance_ = metric_->batch_distance(); + + search_distance_ = add_distance_; + search_batch_distance_ = add_batch_distance_; + + if (metric_->query_metric() && metric_->query_metric()->distance() && + metric_->query_metric()->batch_distance()) { + search_distance_ = metric_->query_metric()->distance(); + search_batch_distance_ = metric_->query_metric()->batch_distance(); + } + + // Create algorithm based on entity storage mode + switch (entity_->storage_mode()) { + case HnswStorageMode::kBufferPool: + alg_ = HnswAlgorithmBase::UPointer( + new HnswAlgorithm( + static_cast(*entity_))); + break; + case HnswStorageMode::kContiguous: { + auto &contiguous_entity = + static_cast(*entity_); + int build_ret = contiguous_entity.build_contiguous_memory(); + if (build_ret != 0) { + LOG_ERROR("Failed to build contiguous memory, ret=%d", build_ret); + return build_ret; + } + alg_ = HnswAlgorithmBase::UPointer( + new HnswAlgorithm(contiguous_entity)); + break; + } + case HnswStorageMode::kExternal: + alg_ = HnswAlgorithmBase::UPointer( + new HnswAlgorithm( + static_cast(*entity_))); + break; + default: + alg_ = + HnswAlgorithmBase::UPointer(new HnswAlgorithm( + static_cast(*entity_))); + break; + } + ret = alg_->init(); + if (ret != 0) { + return ret; + } + + state_ = STATE_OPENED; + magic_ = IndexContext::GenerateMagic(); + + return 0; +} + +int HnswStreamer::close(void) { + LOG_INFO("HnswStreamer close"); + + stats_.clear(); + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_->set_index_meta(meta_); + int ret = entity_->close(); + if (ret != 0) { + return ret; + } + state_ = STATE_INITED; + + return 0; +} + +int HnswStreamer::flush(uint64_t checkpoint) { + LOG_INFO("HnswStreamer flush checkpoint=%zu", (size_t)checkpoint); + + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_->set_index_meta(meta_); + return entity_->flush(checkpoint); +} + +int HnswStreamer::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("HnswStreamer dump"); + + shared_mutex_.lock(); + AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); + + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + return entity_->dump(dumper); +} + +IndexStreamer::Context::Pointer HnswStreamer::create_context(void) const { + if (ailego_unlikely(state_ != STATE_OPENED)) { + LOG_ERROR("Create context failed, open storage first!"); + return Context::Pointer(); + } + + HnswEntity::Pointer entity = entity_->clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("CreateContext clone init failed"); + return Context::Pointer(); + } + HnswContext *ctx = + new (std::nothrow) HnswContext(meta_.dimension(), metric_, entity); + if (ailego_unlikely(ctx == nullptr)) { + LOG_ERROR("Failed to new HnswContext"); + return Context::Pointer(); + } + ctx->set_ef(ef_); + ctx->set_po(po_); + ctx->set_pl(pl_); + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_filter_mode(bf_enabled_ ? VisitFilter::BloomFilter + : VisitFilter::ByteMap); + ctx->set_filter_negative_probability(bf_negative_prob_); + ctx->set_magic(magic_); + ctx->set_force_padding_topk(force_padding_topk_enabled_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + + if (ailego_unlikely(ctx->init(HnswContext::kStreamerContext)) != 0) { + LOG_ERROR("Init HnswContext failed"); + delete ctx; + return Context::Pointer(); + } + uint32_t estimate_doc_count = 0; + if (meta_.streamer_params().get(PARAM_HNSW_STREAMER_ESTIMATE_DOC_COUNT, + &estimate_doc_count)) { + LOG_DEBUG("HnswStreamer doc_count[%zu] estimate[%zu]", + (size_t)entity_->doc_cnt(), (size_t)estimate_doc_count); + } + ctx->check_need_adjuct_ctx(std::max(entity_->doc_cnt(), estimate_doc_count)); + + return Context::Pointer(ctx); +} + +IndexProvider::Pointer HnswStreamer::create_provider(void) const { + LOG_DEBUG("HnswStreamer create provider"); + + auto entity = entity_->clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone HnswEntity failed"); + return nullptr; + } + return Provider::Pointer( + new HnswIndexProvider(meta_, entity, "HnswStreamer")); +} + +int HnswStreamer::update_context(HnswContext *ctx) const { + const HnswEntity::Pointer entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone search context entity"); + return IndexError_Runtime; + } + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + return ctx->update_context(HnswContext::kStreamerContext, meta_, metric_, + entity, magic_); +} + +//! Add a vector with id into index +int HnswStreamer::add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + HnswContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + if (ailego_unlikely(entity_->doc_cnt() >= docs_soft_limit_)) { + if (entity_->doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed [%s]", entity_->doc_cnt(), + PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT.c_str()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %u exceed [%s]", entity_->doc_cnt(), + PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT.c_str()); + } + } + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_); + ctx->reset_query(query); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + if (metric_->support_train()) { + const std::lock_guard lk(mutex_); + ret = metric_->train(query, meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer metric train failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + } + + level_t level = alg_->get_random_level(); + ret = entity_->add_vector_with_id(level, id, query); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, level, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw steamer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + +//! Add a vector into index +int HnswStreamer::add_impl(uint64_t pkey, const void *query, + const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + HnswContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + if (ailego_unlikely(entity_->doc_cnt() >= docs_soft_limit_)) { + if (entity_->doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed [%s]", entity_->doc_cnt(), + PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT.c_str()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %u exceed [%s]", entity_->doc_cnt(), + PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT.c_str()); + } + } + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_); + ctx->reset_query(query); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + if (metric_->support_train()) { + const std::lock_guard lk(mutex_); + ret = metric_->train(query, meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer metric train failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + } + + level_t level = alg_->get_random_level(); + node_id_t id; + ret = entity_->add_vector(level, pkey, query, &id); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, level, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw steamer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + + +int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) const { + return search_impl(query, qmeta, 1, context); +} + +//! Similarity search +int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + IndexStreamer::Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + HnswContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswContext failed"); + return IndexError_Cast; + } + + if (entity_->doc_cnt() <= ctx->get_bruteforce_threshold()) { + return search_bf_impl(query, qmeta, count, context); + } + + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + ret = alg_->search(ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw searcher fast search failed"); + return ret; + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +void HnswStreamer::print_debug_info() { + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) { + continue; + } + Neighbors neighbours = entity_->get_neighbors(0, id); + std::cout << "node: " << id << "; "; + if (neighbours.size() == 0) std::cout << std::endl; + for (uint32_t i = 0; i < neighbours.size(); ++i) { + std::cout << neighbours[i]; + + if (i == neighbours.size() - 1) { + std::cout << std::endl; + } else { + std::cout << ", "; + } + } + } + + // entity_->print_key_map(); +} + +int HnswStreamer::search_bf_impl( + const void *query, const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) const { + return search_bf_impl(query, qmeta, 1, context); +} + +int HnswStreamer::search_bf_impl( + const void *query, const IndexQueryMeta &qmeta, uint32_t count, + IndexStreamer::Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + HnswContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity_->get_key(id)); + }; + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + ctx->group_topk_heaps().clear(); + + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) { + continue; + } + + if (!ctx->filter().is_valid() || !ctx->filter()(entity_->get_key(id))) { + dist_t dist = ctx->dist_calculator().batch_dist(id); + + std::string group_id = group_by(id); + + auto &topk_heap = ctx->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, dist); + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } else { + auto &filter = ctx->filter(); + auto &topk = ctx->topk_heap(); + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + topk.clear(); + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) { + continue; + } + + if (!filter.is_valid() || !filter(entity_->get_key(id))) { + dist_t dist = ctx->dist_calculator().batch_dist(id); + topk.emplace(id, dist); + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +int HnswStreamer::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + if (ailego_unlikely(p_keys.size() != count)) { + LOG_ERROR("The size of p_keys is not equal to count"); + return IndexError_InvalidArgument; + } + + HnswContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity_->get_key(id)); + }; + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + ctx->group_topk_heaps().clear(); + + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t pk = p_keys[q][idx]; + if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { + node_id_t id = entity_->get_id(pk); + if (id != kInvalidNodeId) { + dist_t dist = ctx->dist_calculator().batch_dist(id); + std::string group_id = group_by(id); + + auto &topk_heap = ctx->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, dist); + } + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } else { + auto &filter = ctx->filter(); + auto &topk = ctx->topk_heap(); + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + topk.clear(); + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + key_t pk = p_keys[q][idx]; + if (!filter.is_valid() || !filter(pk)) { + node_id_t id = entity_->get_id(pk); + if (id != kInvalidNodeId) { + dist_t dist = ctx->dist_calculator().batch_dist(id); + topk.emplace(id, dist); + } + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + + +INDEX_FACTORY_REGISTER_STREAMER(HnswStreamer); + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw/hnsw_streamer.h b/src/core/algorithm/hnsw/hnsw_streamer.h new file mode 100644 index 0000000..f06321b --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_streamer.h @@ -0,0 +1,244 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "hnsw_algorithm.h" +#include "hnsw_streamer_entity.h" + +namespace zvec { +namespace core { + +class HnswStreamer : public IndexStreamer { + public: + using ContextPointer = IndexStreamer::Context::Pointer; + + HnswStreamer(void); + ~HnswStreamer(void) override; + + HnswStreamer(const HnswStreamer &streamer) = delete; + HnswStreamer &operator=(const HnswStreamer &streamer) = delete; + + public: + //! Retrieve the storage mode of the underlying entity. Returns + //! HnswStorageMode::kMmap when the entity has not been initialized yet. + //! Intended for introspection and debug/testing usage. + HnswStorageMode storage_mode() const { + if (!entity_) { + return HnswStorageMode::kMmap; + } + return entity_->storage_mode(); + } + + protected: + //! Initialize Streamer + int init(const IndexMeta &imeta, const ailego::Params ¶ms) override; + + //! Cleanup Streamer + int cleanup(void) override; + + //! Create a context + Context::Pointer create_context(void) const override; + + //! Create a new iterator + IndexProvider::Pointer create_provider(void) const override; + + //! Add a vector into index + int add_impl(uint64_t pkey, const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Add a vector with id into index + int add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + //! Fetch vector by key + const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector_by_key(key, block); + } + + //! Fetch vector by id + const void *get_vector_by_id(uint32_t id) const override { + return entity_->get_vector(id); + } + + int get_vector_by_id(const uint32_t id, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector(id, block); + } + + //! Open index from file path + int open(IndexStorage::Pointer stg) override; + + //! Close file + int close(void) override; + + //! flush file + int flush(uint64_t checkpoint) override; + + //! Dump index into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + void print_debug_info() override; + + private: + inline int check_params(const void *query, + const IndexQueryMeta &qmeta) const { + if (ailego_unlikely(!query)) { + LOG_ERROR("null query"); + return IndexError_InvalidArgument; + } + if (ailego_unlikely(qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size())) { + LOG_ERROR("Unsupported query meta"); + return IndexError_Mismatch; + } + return 0; + } + + inline int check_sparse_count_is_zero(const uint32_t *sparse_count, + uint32_t count) const { + for (uint32_t i = 0; i < count; ++i) { + if (sparse_count[i] != 0) + LOG_ERROR("Sparse cout is not empty. Index: %u, Sparse Count: %u", i, + sparse_count[i]); + return IndexError_InvalidArgument; + } + + return 0; + } + + private: + //! Configure and initialize the entity with saved parameters + int setup_entity(); + + //! To share ctx across streamer/searcher, we need to update the context for + //! current streamer/searcher + int update_context(HnswContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 }; + class Stats : public IndexStreamer::Stats { + public: + void clear(void) { + set_revision_id(0u); + set_loaded_count(0u); + set_added_count(0u); + set_discarded_count(0u); + set_index_size(0u); + set_dumped_size(0u); + set_check_point(0u); + set_create_time(0u); + set_update_time(0u); + clear_attributes(); + } + }; + + std::unique_ptr entity_; + HnswAlgorithmBase::UPointer alg_; + IndexMeta meta_{}; + IndexMetric::Pointer metric_{}; + + IndexMetric::MatrixDistance add_distance_{}; + IndexMetric::MatrixDistance search_distance_{}; + + IndexMetric::MatrixBatchDistance add_batch_distance_{}; + IndexMetric::MatrixBatchDistance search_batch_distance_{}; + + Stats stats_{}; + std::mutex mutex_{}; + + size_t max_index_size_{0UL}; + size_t chunk_size_{HnswEntity::kDefaultChunkSize}; + size_t docs_hard_limit_{HnswEntity::kDefaultDocsHardLimit}; + size_t docs_soft_limit_{0UL}; + uint32_t min_neighbor_cnt_{0u}; + uint32_t prune_cnt_{0u}; + uint32_t upper_max_neighbor_cnt_{HnswEntity::kDefaultUpperMaxNeighborCnt}; + uint32_t l0_max_neighbor_cnt_{HnswEntity::kDefaultL0MaxNeighborCnt}; + uint32_t ef_{HnswEntity::kDefaultEf}; + uint32_t po_{8}; + uint32_t pl_{0}; + uint32_t ef_construction_{HnswEntity::kDefaultEfConstruction}; + uint32_t scaling_factor_{HnswEntity::kDefaultScalingFactor}; + size_t bruteforce_threshold_{HnswEntity::kDefaultBruteForceThreshold}; + size_t max_scan_limit_{HnswEntity::kDefaultMaxScanLimit}; + size_t min_scan_limit_{HnswEntity::kDefaultMinScanLimit}; + float bf_negative_prob_{HnswEntity::kDefaultBFNegativeProbability}; + float max_scan_ratio_{HnswEntity::kDefaultScanRatio}; + + uint32_t magic_{0U}; + State state_{STATE_INIT}; + bool bf_enabled_{false}; + bool check_crc_enabled_{false}; + bool filter_same_key_{false}; + bool get_vector_enabled_{false}; + bool force_padding_topk_enabled_{false}; + bool use_id_map_{true}; + bool use_contiguous_memory_{false}; + bool use_external_vector_{false}; + + //! avoid add vector while dumping index + ailego::SharedMutex shared_mutex_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc new file mode 100644 index 0000000..1ad1ebd --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc @@ -0,0 +1,1061 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "hnsw_streamer_entity.h" +#if defined(__linux__) || defined(__APPLE__) +#include +#endif +#include + +// #define DEBUG_PRINT + +namespace zvec { +namespace core { + +HnswStreamerEntity::HnswStreamerEntity(IndexStreamer::Stats &stats) + : stats_(stats) {} + +HnswStreamerEntity::~HnswStreamerEntity() {} + +int HnswStreamerEntity::init(size_t max_doc_cnt) { + if (std::pow(scaling_factor(), kMaxGraphLayers) < max_doc_cnt) { + LOG_ERROR("scalingFactor=%zu is too small", scaling_factor()); + return IndexError_InvalidArgument; + } + + std::lock_guard lock(mutex_); + broker_ = std::make_shared(stats_); + upper_neighbor_index_ = std::make_shared(); + upper_neighbor_rw_mutex_ = std::make_shared(); + keys_map_lock_ = std::make_shared(); + keys_map_ = std::make_shared>(); + if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) { + LOG_ERROR("HnswStreamerEntity new object failed"); + return IndexError_NoMemory; + } + keys_map_->set_empty_key(kInvalidKey); + + neighbor_size_ = neighbors_size(); + upper_neighbor_size_ = upper_neighbors_size(); + + //! vector + key + level 0 neighbors + size_t size = vector_size() + sizeof(key_t) + neighbor_size_; + + size = AlignSize(size); + set_node_size(size); + return 0; +} + +int HnswStreamerEntity::cleanup() { + std::lock_guard lock(mutex_); + mutable_header()->clear(); + chunk_size_ = kDefaultChunkSize; + node_index_mask_bits_ = 0U; + node_index_mask_ = 0U; + node_cnt_per_chunk_ = 0U; + neighbor_size_ = 0U; + upper_neighbor_size_ = 0U; + if (upper_neighbor_index_) { + upper_neighbor_index_->cleanup(); + } + if (keys_map_) { + keys_map_->clear(); + } + node_chunks_.clear(); + node_chunk_bases_.reset(); + upper_neighbor_chunks_.clear(); + upper_neighbor_chunk_bases_.reset(); + filter_same_key_ = false; + get_vector_enabled_ = false; + broker_.reset(); + + return 0; +} + +int HnswStreamerEntity::update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) { + std::vector buffer(neighbor_size_); + NeighborsHeader *hd = reinterpret_cast(buffer.data()); + hd->neighbor_cnt = neighbors.size(); + size_t i = 0; + for (; i < neighbors.size(); ++i) { + hd->neighbors[i] = neighbors[i].first; + } + + auto loc = get_neighbor_chunk_loc(level, id); + size_t size = reinterpret_cast(&hd->neighbors[i]) - buffer.data(); + size_t ret = loc.first->write(loc.second, hd, size); + if (ailego_unlikely(ret != size)) { + LOG_ERROR("Write neighbor header failed, ret=%zu", ret); + + return IndexError_Runtime; + } + + return 0; +} + +const Neighbors HnswStreamerEntity::get_neighbors(level_t level, + node_id_t id) const { + size_t offset = 0UL; + size_t neighbor_size = neighbor_size_; + IndexStorage::MemoryBlock neighbor_block; + + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + + // Fast path: use pre-cached stable base pointer (mmap backend). + // Bounds-check guards against new chunks added after clone() was taken. + if (node_chunk_bases_ && chunk_idx < node_chunk_bases_->size() && + (*node_chunk_bases_)[chunk_idx]) { + neighbor_block.reset((void *)((*node_chunk_bases_)[chunk_idx] + offset)); + } else { + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + Chunk *chunk = node_chunks_[chunk_idx].get(); + ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); + size_t size = chunk->read(offset, neighbor_block, neighbor_size); + if (ailego_unlikely(size != neighbor_size)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", size); + return Neighbors(); + } + return Neighbors(neighbor_block); + } + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + offset = p.second; + neighbor_size = upper_neighbor_size_; + + // Fast path: use pre-cached stable base pointer (mmap backend). + // Bounds-check guards against new chunks added after clone() was taken. + if (upper_neighbor_chunk_bases_ && + p.first < upper_neighbor_chunk_bases_->size() && + (*upper_neighbor_chunk_bases_)[p.first]) { + neighbor_block.reset( + (void *)((*upper_neighbor_chunk_bases_)[p.first] + offset)); + } else { + Chunk *chunk = upper_neighbor_chunks_[p.first].get(); + ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); + size_t size = chunk->read(offset, neighbor_block, neighbor_size); + if (ailego_unlikely(size != neighbor_size)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", size); + return Neighbors(); + } + return Neighbors(neighbor_block); + } + } + + return Neighbors(neighbor_block); +} + +//! Get vector data by key +const void *HnswStreamerEntity::get_vector(node_id_t id) const { + auto loc = get_vector_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + + // Fast path: mmap backend — direct pointer arithmetic. + // Bounds-check guards against new chunks added after clone() was taken. + if (node_chunk_bases_ && loc.first < node_chunk_bases_->size() && + (*node_chunk_bases_)[loc.first]) { + return (*node_chunk_bases_)[loc.first] + loc.second; + } + + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + const void *vec = nullptr; + size_t read_size = vector_size(); + size_t ret = node_chunks_[loc.first]->read(loc.second, &vec, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + } + return vec; +} + +int HnswStreamerEntity::get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const { + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + + // Fast path: mmap backend. + // Bounds-check guards against new chunks added after clone() was taken. + if (node_chunk_bases_ && loc.first < node_chunk_bases_->size() && + (*node_chunk_bases_)[loc.first]) { + vecs[i] = (*node_chunk_bases_)[loc.first] + loc.second; + continue; + } + + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t read_size = vector_size(); + size_t ret = node_chunks_[loc.first]->read(loc.second, &vecs[i], read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + } + return 0; +} + +int HnswStreamerEntity::get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const { + auto loc = get_vector_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + + // Fast path: mmap backend. + // Bounds-check guards against new chunks added after clone() was taken. + if (node_chunk_bases_ && loc.first < node_chunk_bases_->size() && + (*node_chunk_bases_)[loc.first]) { + block.reset((void *)((*node_chunk_bases_)[loc.first] + loc.second)); + return 0; + } + + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t read_size = vector_size(); + size_t ret = node_chunks_[loc.first]->read(loc.second, block, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + return 0; +} + +int HnswStreamerEntity::get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + + // Fast path: mmap backend. + // Bounds-check guards against new chunks added after clone() was taken. + if (node_chunk_bases_ && loc.first < node_chunk_bases_->size() && + (*node_chunk_bases_)[loc.first]) { + vec_blocks[i].reset( + (void *)((*node_chunk_bases_)[loc.first] + loc.second)); + continue; + } + + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t read_size = vector_size(); + size_t ret = + node_chunks_[loc.first]->read(loc.second, vec_blocks[i], read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + } + return 0; +} + +key_t HnswStreamerEntity::get_key(node_id_t id) const { + if (use_key_info_map_) { + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + + // Fast path: mmap backend. + // Bounds-check guards against new chunks added after clone() was taken. + if (node_chunk_bases_ && loc.first < node_chunk_bases_->size() && + (*node_chunk_bases_)[loc.first]) { + return *reinterpret_cast((*node_chunk_bases_)[loc.first] + + loc.second); + } + + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + IndexStorage::MemoryBlock key_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(key_block.data()); + } else { + return id; + } +} + +void HnswStreamerEntity::add_neighbor(level_t level, node_id_t id, + uint32_t size, node_id_t neighbor_id) { + auto loc = get_neighbor_chunk_loc(level, id); + size_t offset = + loc.second + sizeof(NeighborsHeader) + size * sizeof(node_id_t); + ailego_assert_with(size < neighbor_cnt(level), "invalid neighbor size"); + ailego_assert_with(offset < loc.first->data_size(), "invalid chunk offset"); + size_t ret = loc.first->write(offset, &neighbor_id, sizeof(node_id_t)); + if (ailego_unlikely(ret != sizeof(node_id_t))) { + LOG_ERROR("Write neighbor id failed, ret=%zu", ret); + return; + } + + uint32_t neighbors = size + 1; + ret = loc.first->write(loc.second, &neighbors, sizeof(uint32_t)); + if (ailego_unlikely(ret != sizeof(uint32_t))) { + LOG_ERROR("Write neighbor cnt failed, ret=%zu", ret); + } + + return; +} + +int HnswStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) { + if (header_chunk->data_size() < header_size()) { + LOG_ERROR("Invalid header chunk size"); + return IndexError_InvalidFormat; + } + IndexStorage::MemoryBlock header_block; + size_t size = header_chunk->read(0UL, header_block, header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Read header chunk failed"); + return IndexError_ReadData; + } + *mutable_header() = + *reinterpret_cast(header_block.data()); + + int ret = check_hnsw_index(&header()); + if (ret != 0) { + broker_->close(); + return ret; + } + + node_chunks_.resize(broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NODE)); + node_chunk_bases_ = std::make_shared>( + node_chunks_.size(), nullptr); + for (auto seq = 0UL; seq < node_chunks_.size(); ++seq) { + node_chunks_[seq] = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_NODE, seq); + if (!node_chunks_[seq]) { + LOG_ERROR("Missing hnsw streamer data chunk %zu th of %zu", seq, + node_chunks_.size()); + return IndexError_InvalidFormat; + } + (*node_chunk_bases_)[seq] = node_chunks_[seq]->base_data(); + } + + upper_neighbor_chunks_.resize( + broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR)); + upper_neighbor_chunk_bases_ = std::make_shared>( + upper_neighbor_chunks_.size(), nullptr); + for (auto seq = 0UL; seq < upper_neighbor_chunks_.size(); ++seq) { + upper_neighbor_chunks_[seq] = + broker_->get_chunk(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, seq); + if (!upper_neighbor_chunks_[seq]) { + LOG_ERROR("Missing hnsw streamer index chunk %zu th of %zu", seq, + upper_neighbor_chunks_.size()); + return IndexError_InvalidFormat; + } + (*upper_neighbor_chunk_bases_)[seq] = + upper_neighbor_chunks_[seq]->base_data(); + } + + return 0; +} + +int HnswStreamerEntity::open(IndexStorage::Pointer stg, uint64_t max_index_size, + bool check_crc) { + std::lock_guard lock(mutex_); + bool huge_page = stg->isHugePage(); + LOG_DEBUG("huge_page: %d", (int)huge_page); + int ret = broker_->open(std::move(stg), chunk_size_, check_crc); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Open index failed for %s", IndexError::What(ret)); + return ret; + } + ret = init_chunk_params(max_index_size, huge_page); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("init_chunk_params failed for %s", IndexError::What(ret)); + return ret; + } + broker_->set_max_chunks_size(max_index_size_); + + ret = upper_neighbor_index_->init(broker_, upper_neighbor_chunk_size_, + scaling_factor(), estimate_doc_capacity(), + kUpperHashMemoryInflateRatio); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Init neighbor hash map failed"); + return ret; + } + + //! init header + auto header_chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId); + if (!header_chunk) { // open empty index, create one + auto p = + broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId, header_size()); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc header chunk failed"); + return p.first; + } + size_t size = p.second->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + return 0; + } + + //! Open an exist hnsw index + ret = init_chunks(header_chunk); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + //! total docs including features wrote in index but neighbors may not ready + node_id_t total_vecs = 0; + if (node_chunks_.size() > 0) { + size_t last_idx = node_chunks_.size() - 1; + auto last_chunk = node_chunks_[last_idx]; + if (last_chunk->data_size() % node_size()) { + LOG_WARN("The index may broken"); + return IndexError_InvalidFormat; + } + total_vecs = last_idx * node_cnt_per_chunk_ + + node_chunks_[last_idx]->data_size() / node_size(); + } + + LOG_INFO( + "Open index, l0NeighborCnt=%zu upperNeighborCnt=%zu " + "efConstruction=%zu curDocCnt=%u totalVecs=%u maxLevel=%u", + l0_neighbor_cnt(), upper_neighbor_cnt(), ef_construction(), doc_cnt(), + total_vecs, cur_max_level()); + //! try to correct the docCnt if index not fully flushed + if (doc_cnt() != total_vecs) { + LOG_WARN("Index closed abnormally, using totalVecs as curDocCnt"); + *mutable_doc_cnt() = total_vecs; + } + if (filter_same_key_ || get_vector_enabled_) { + if (use_key_info_map_) { + for (node_id_t id = 0U; id < doc_cnt(); ++id) { + if (get_key(id) == kInvalidKey) { + continue; + } + (*keys_map_)[get_key(id)] = id; + } + } + } + + stats_.set_loaded_count(doc_cnt()); + + return 0; +} + +int HnswStreamerEntity::close() { + LOG_DEBUG("close index"); + + std::lock_guard lock(mutex_); + flush_header(); + mutable_header()->reset(); + upper_neighbor_index_->cleanup(); + keys_map_->clear(); + header_.clear(); + node_chunks_.clear(); + node_chunk_bases_.reset(); + upper_neighbor_chunks_.clear(); + upper_neighbor_chunk_bases_.reset(); + + return broker_->close(); +} + +int HnswStreamerEntity::flush(uint64_t checkpoint) { + LOG_INFO("Flush index, curDocs=%u", doc_cnt()); + + std::lock_guard lock(mutex_); + flush_header(); + int ret = broker_->flush(checkpoint); + if (ret != 0) { + return ret; + } + + return 0; +} + +int HnswStreamerEntity::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("Dump index, curDocs=%u", doc_cnt()); + + //! sort by keys, to support get_vector by key in searcher + std::vector keys(doc_cnt()); + for (node_id_t i = 0; i < doc_cnt(); ++i) { + keys[i] = get_key(i); + } + + //! dump neighbors + auto get_level = [&](node_id_t id) { + auto it = upper_neighbor_index_->find(id); + if (it == upper_neighbor_index_->end()) { + return 0U; + }; + auto meta = reinterpret_cast(&it->second); + return meta->bits.level; + }; + auto ret = dump_segments(dumper, keys.data(), get_level); + if (ailego_unlikely(ret < 0)) { + return ret; + } + *stats_.mutable_dumped_size() += ret; + + return 0; +} + +int HnswStreamerEntity::check_hnsw_index(const HNSWHeader *hd) const { + if (l0_neighbor_cnt() != hd->l0_neighbor_cnt() || + upper_neighbor_cnt() != hd->upper_neighbor_cnt()) { + LOG_ERROR("Param neighbor cnt: %zu:%zu mismatch index previous %zu:%zu", + l0_neighbor_cnt(), upper_neighbor_cnt(), hd->l0_neighbor_cnt(), + hd->upper_neighbor_cnt()); + return IndexError_Mismatch; + } + if (vector_size() != hd->vector_size()) { + LOG_ERROR("vector size %zu mismatch index previous %zu", vector_size(), + hd->vector_size()); + return IndexError_Mismatch; + } + if (ef_construction() != hd->ef_construction()) { + LOG_WARN("Param efConstruction %zu mismatch index previous %zu", + ef_construction(), hd->ef_construction()); + } + if (scaling_factor() != hd->scaling_factor()) { + LOG_WARN("Param scalingFactor %zu mismatch index previous %zu", + scaling_factor(), hd->scaling_factor()); + return IndexError_Mismatch; + } + if (prune_cnt() != hd->neighbor_prune_cnt()) { + LOG_WARN("Param pruneCnt %zu mismatch index previous %zu", prune_cnt(), + hd->neighbor_prune_cnt()); + return IndexError_Mismatch; + } + if ((hd->entry_point() != kInvalidNodeId && + hd->entry_point() >= hd->doc_cnt()) || + (hd->entry_point() == kInvalidNodeId && hd->doc_cnt() > 0U)) { + LOG_WARN("Invalid entryPoint %u, docCnt %u", hd->entry_point(), + hd->doc_cnt()); + return IndexError_InvalidFormat; + } + if (hd->entry_point() == kInvalidNodeId && + broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NODE) > 0) { + LOG_WARN("The index is broken, maybe it haven't flush"); + return IndexError_InvalidFormat; + } + + return 0; +} + +int HnswStreamerEntity::add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) { + Chunk::Pointer node_chunk; + // On MSVC, unsigned long is 32-bit, so -1UL is 0xFFFFFFFF not + // 0xFFFFFFFFFFFFFFFF. + size_t chunk_offset = static_cast(-1); + + std::lock_guard lock(mutex_); + // duplicate check + if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) { + LOG_WARN("Try to add duplicate key, ignore it"); + return IndexError_Duplicate; + } + + node_id_t local_id = static_cast(doc_cnt()); + uint32_t chunk_index = node_chunks_.size() - 1U; + if (chunk_index == -1U || + (node_chunks_[chunk_index]->data_size() >= + node_cnt_per_chunk_ * node_size())) { // no space left and need to alloc + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + chunk_index++; + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NODE, chunk_index, + chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + chunk_offset = 0UL; + node_chunks_.emplace_back(node_chunk); + } else { + node_chunk = node_chunks_[chunk_index]; + chunk_offset = node_chunk->data_size(); + } + + size_t size = node_chunk->write(chunk_offset, vec, vector_size()); + if (ailego_unlikely(size != vector_size())) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + //! level 0 neighbors is inited to zero by default + + int ret = add_upper_neighbor(level, local_id); + if (ret != 0) { + return ret; + } + + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + if (filter_same_key_ || get_vector_enabled_) { + if (use_key_info_map_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = local_id; + keys_map_lock_->unlock(); + } + } + + *mutable_doc_cnt() += 1; + broker_->mark_dirty(); + *id = local_id; + + return 0; +} + +int HnswStreamerEntity::add_vector_with_id(level_t level, node_id_t id, + const void *vec) { + Chunk::Pointer node_chunk; + size_t chunk_offset = static_cast(-1); + key_t key = id; + + std::lock_guard lock(mutex_); + + // duplicate check + if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) { + LOG_WARN("Try to add duplicate key, ignore it"); + return IndexError_Duplicate; + } + + // set node_chunk & chunk_offset if succeed + auto func_get_node_chunk_and_offset = [&](node_id_t node_id) -> int { + uint32_t chunk_index = node_id >> node_index_mask_bits_; + ailego_assert_with(chunk_index <= node_chunks_.size(), "invalid chunk idx"); + // belongs to next chunk + if (chunk_index == node_chunks_.size()) { + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NODE, chunk_index, + chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + node_chunks_.emplace_back(node_chunk); + } + + node_chunk = node_chunks_[chunk_index]; + chunk_offset = (node_id & node_index_mask_) * node_size(); + return 0; + }; + + for (size_t start_id = doc_cnt(); start_id < id; ++start_id) { + if (auto ret = func_get_node_chunk_and_offset(start_id); ret != 0) { + LOG_ERROR("func_get_node_chunk_and_offset failed"); + return ret; + } + size_t size = node_chunk->write(chunk_offset + vector_size(), &kInvalidKey, + sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (auto ret = func_get_node_chunk_and_offset(id); ret != 0) { + LOG_ERROR("func_get_node_chunk_and_offset failed"); + return ret; + } + + size_t size = node_chunk->write(chunk_offset, vec, vector_size()); + if (ailego_unlikely(size != vector_size())) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + + size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + //! level 0 neighbors is inited to zero by default + + int ret = add_upper_neighbor(level, id); + if (ret != 0) { + return ret; + } + + if (*mutable_doc_cnt() <= id) { + *mutable_doc_cnt() = id + 1; + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (filter_same_key_ || get_vector_enabled_) { + if (use_key_info_map_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = id; + keys_map_lock_->unlock(); + } + } + + broker_->mark_dirty(); + + return 0; +} + +void HnswStreamerEntity::update_ep_and_level(node_id_t ep, level_t level) { + HnswEntity::update_ep_and_level(ep, level); + flush_header(); + + return; +} + +const HnswEntity::Pointer HnswStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + HnswStreamerEntity *entity = new (std::nothrow) HnswStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_, + keys_map_, use_key_info_map_, std::move(node_chunks), + std::move(upper_neighbor_chunks), broker_, node_chunk_bases_, + upper_neighbor_chunk_bases_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswStreamerEntity new failed"); + } + return HnswEntity::Pointer(entity); +} + +const HnswEntity::Pointer HnswMmapStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswMmapStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswMmapStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) HnswMmapStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_, + keys_map_, use_key_info_map_, std::move(node_chunks), + std::move(upper_neighbor_chunks), broker_, nullptr, nullptr); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswMmapStreamerEntity new failed"); + } + return HnswEntity::Pointer(entity); +} + +const HnswEntity::Pointer HnswContiguousStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswContiguousStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswContiguousStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) HnswContiguousStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_, + keys_map_, use_key_info_map_, std::move(node_chunks), + std::move(upper_neighbor_chunks), broker_, nullptr, nullptr); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswContiguousStreamerEntity new failed"); + return HnswEntity::Pointer(); + } + + // Share contiguous memory with the clone (zero-copy) + entity->vector_memory_ = vector_memory_; + entity->vector_base_ = vector_base_; + entity->graph_memory_ = graph_memory_; + entity->graph_base_ = graph_base_; + entity->graph_stride_ = graph_stride_; + entity->upper_neighbor_memory_ = upper_neighbor_memory_; + entity->upper_neighbor_base_ = upper_neighbor_base_; + entity->upper_chunk_offsets_ = upper_chunk_offsets_; + + return HnswEntity::Pointer(entity); +} + +const HnswEntity::Pointer HnswExternalStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswExternalStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswExternalStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + // Note: vec_src_ is intentionally NOT shared with the clone; it stays null + // and is re-bound per add/search call via HnswContext::set_vector_source. + auto *entity = new (std::nothrow) HnswExternalStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_, + keys_map_, use_key_info_map_, std::move(node_chunks), + std::move(upper_neighbor_chunks), broker_, nullptr, nullptr); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswExternalStreamerEntity new failed"); + } + return HnswEntity::Pointer(entity); +} + +// ============================================================================ +// HnswContiguousStreamerEntity implementation +// ============================================================================ + +char *HnswContiguousStreamerEntity::allocate_contiguous(size_t size) { + if (size == 0) { + return nullptr; + } +#if defined(__linux__) + // Use mmap with MAP_ANONYMOUS for contiguous memory + void *ptr = ::mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (ptr == MAP_FAILED) { + LOG_ERROR("mmap failed for contiguous memory, size=%zu", size); + return nullptr; + } + // Request transparent huge pages + ::madvise(ptr, size, MADV_HUGEPAGE); + return static_cast(ptr); +#elif defined(__APPLE__) + // macOS: use mmap with MAP_ANONYMOUS + void *ptr = ::mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANON, -1, 0); + if (ptr == MAP_FAILED) { + LOG_ERROR("mmap failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#elif defined(_WIN32) + void *ptr = ::_aligned_malloc(size, ailego::MemoryHelper::PageSize()); + if (!ptr) { + LOG_ERROR("_aligned_malloc failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#else + // Fallback: aligned allocation + void *ptr = std::aligned_alloc(ailego::MemoryHelper::PageSize(), size); + if (!ptr) { + LOG_ERROR("aligned_alloc failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#endif +} + +int HnswContiguousStreamerEntity::build_contiguous_memory() { + vector_memory_.reset(); + vector_base_ = nullptr; + graph_memory_.reset(); + graph_base_ = nullptr; + upper_neighbor_memory_.reset(); + upper_neighbor_base_ = nullptr; + upper_chunk_offsets_.clear(); + + const uint32_t total_docs = doc_cnt(); + if (total_docs == 0) { + return 0; + } + + const size_t per_node = node_size(); + const size_t vec_size = vector_size(); + // graph_stride = key + L0 neighbors (everything except vector) + graph_stride_ = sizeof(key_t) + neighbor_size_; + + // --- Allocate flat vector array (stride = vector_size) --- + const size_t total_vec_data = static_cast(total_docs) * vec_size; + size_t vector_memory_size = AlignHugePageSize(total_vec_data); + char *raw_vec = allocate_contiguous(vector_memory_size); + if (!raw_vec) { + return IndexError_Runtime; + } + vector_memory_.reset(raw_vec, ContiguousDeleter{vector_memory_size}); + vector_base_ = raw_vec; + + // --- Allocate graph array (stride = sizeof(key_t) + neighbor_size) --- + const size_t total_graph_data = + static_cast(total_docs) * graph_stride_; + size_t graph_memory_size = AlignHugePageSize(total_graph_data); + char *raw_graph = allocate_contiguous(graph_memory_size); + if (!raw_graph) { + vector_memory_.reset(); + vector_base_ = nullptr; + return IndexError_Runtime; + } + graph_memory_.reset(raw_graph, ContiguousDeleter{graph_memory_size}); + graph_base_ = raw_graph; + + // Split node data from chunks into vector and graph arrays. + // Original node layout: [vector (vec_size) | key (8B) | L0 neighbors] + const auto &chunks = node_chunks_; + const uint32_t nodes_per_chunk = 1U << node_index_mask_bits_; + for (size_t chunk_idx = 0; chunk_idx < chunks.size(); ++chunk_idx) { + const void *chunk_data = nullptr; + size_t data_size = chunks[chunk_idx]->data_size(); + chunks[chunk_idx]->read(0, &chunk_data, data_size); + + uint32_t base_id = chunk_idx * nodes_per_chunk; + uint32_t count_in_chunk = std::min(nodes_per_chunk, total_docs - base_id); + + const char *src = static_cast(chunk_data); + for (uint32_t i = 0; i < count_in_chunk; ++i) { + const char *node_src = src + static_cast(i) * per_node; + size_t global_id = static_cast(base_id + i); + + // Copy vector to flat vector array + std::memcpy(vector_base_ + global_id * vec_size, node_src, vec_size); + + // Copy key + L0 neighbors to graph array + std::memcpy(graph_base_ + global_id * graph_stride_, node_src + vec_size, + graph_stride_); + } + } + + // --- Build contiguous upper neighbor memory --- + const auto &upper_chunks = upper_neighbor_chunks_; + if (upper_chunks.empty()) { + LOG_INFO( + "Built HNSW contiguous memory (split layout): " + "vector_mem=%zu graph_mem=%zu total_docs=%u node_chunks=%zu", + vector_memory_size, graph_memory_size, total_docs, chunks.size()); + return 0; + } + + // Sync all upper neighbor chunks + sync_upper_neighbor_chunks(upper_chunks.size() - 1); + + // Calculate cumulative offsets and total size + upper_chunk_offsets_.resize(upper_chunks.size()); + size_t total_upper_size = 0; + for (size_t i = 0; i < upper_chunks.size(); ++i) { + upper_chunk_offsets_[i] = total_upper_size; + total_upper_size += upper_chunks[i]->data_size(); + } + + size_t upper_memory_size = AlignHugePageSize(total_upper_size); + char *raw_upper = allocate_contiguous(upper_memory_size); + if (!raw_upper) { + vector_memory_.reset(); + vector_base_ = nullptr; + graph_memory_.reset(); + graph_base_ = nullptr; + return IndexError_Runtime; + } + upper_neighbor_memory_.reset(raw_upper, ContiguousDeleter{upper_memory_size}); + upper_neighbor_base_ = raw_upper; + + // Copy upper neighbor data from chunks into contiguous memory + for (size_t i = 0; i < upper_chunks.size(); ++i) { + const void *chunk_data = nullptr; + size_t data_size = upper_chunks[i]->data_size(); + upper_chunks[i]->read(0, &chunk_data, data_size); + std::memcpy(upper_neighbor_base_ + upper_chunk_offsets_[i], chunk_data, + data_size); + } + + LOG_INFO( + "Built HNSW contiguous memory (split layout): " + "vector_mem=%zu graph_mem=%zu upper_neighbor_mem=%zu " + "total_docs=%u node_chunks=%zu upper_chunks=%zu", + vector_memory_size, graph_memory_size, upper_memory_size, total_docs, + chunks.size(), upper_chunks.size()); + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.h b/src/core/algorithm/hnsw/hnsw_streamer_entity.h new file mode 100644 index 0000000..9c53aaf --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.h @@ -0,0 +1,1195 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#if defined(__linux__) || defined(__APPLE__) +#include +#endif +#include +#include +#include +#include +#include +#include "hnsw_chunk.h" +#include "hnsw_entity.h" +#include "hnsw_index_hash.h" +#include "hnsw_params.h" + +namespace zvec { +namespace core { + + +//! Storage mode for HnswStreamerEntity +enum class HnswStorageMode { + kMmap = 0, + kBufferPool = 1, + kContiguous = 2, + kExternal = 3 +}; + +//! HnswStreamerEntity manage vector data, pkey, and node's neighbors +class HnswStreamerEntity : public HnswEntity { + public: + //! Cleanup + //! return 0 on success, or errCode in failure + int cleanup() override; + + //! Make a copy of streamer entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + const HnswEntity::Pointer clone() const override; + + //! Get primary key of the node id + key_t get_key(node_id_t id) const override; + + //! Get vector feature data by key + const void *get_vector(node_id_t id) const override; + + //! Get vectors feature data by local ids + int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + + int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const override; + + //! Get the node id's neighbors on graph level + //! Note: the neighbors cannot be modified, using the following + //! method to get WritableNeighbors if want to + const Neighbors get_neighbors(level_t level, node_id_t id) const override; + + //! Add vector and key to hnsw entity, and local id will be saved in id + int add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) override; + + //! Add vector and id to hnsw entity + int add_vector_with_id(level_t level, node_id_t id, const void *vec) override; + + int update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) override; + + //! Append neighbor_id to node id neighbors on level + //! Notice: the caller must be ensure the neighbors not full + void add_neighbor(level_t level, node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + + //! Dump index by dumper + int dump(const IndexDumper::Pointer &dumper) override; + + void update_ep_and_level(node_id_t ep, level_t level) override; + + //! Get the storage mode of this entity + virtual HnswStorageMode storage_mode() const { + return HnswStorageMode::kMmap; + } + + void set_use_key_info_map(bool use_id_map) { + use_key_info_map_ = use_id_map; + LOG_DEBUG("use_key_info_map_: %d", (int)use_key_info_map_); + } + + public: + //! Constructor + HnswStreamerEntity(IndexStreamer::Stats &stats); + + //! Destructor + ~HnswStreamerEntity(); + + //! Get vector feature data by key + const void *get_vector_by_key(key_t key) const override { + auto id = get_id(key); + return id == kInvalidNodeId ? nullptr : get_vector(id); + } + + int get_vector_by_key(const key_t key, + IndexStorage::MemoryBlock &block) const override { + auto id = get_id(key); + if (id != kInvalidNodeId) { + return get_vector(id, block); + } else { + return IndexError_InvalidArgument; + } + } + + //! Init entity + int init(size_t max_doc_cnt); + + //! Flush graph entity to disk + //! return 0 on success, or errCode in failure + int flush(uint64_t checkpoint); + + //! Open entity from storage + //! return 0 on success, or errCode in failure + int open(IndexStorage::Pointer stg, uint64_t max_index_size, bool check_crc); + + //! Close entity + //! return 0 on success, or errCode in failure + int close(); + + //! Set meta information from entity + int set_index_meta(const IndexMeta &meta) const { + return IndexHelper::SerializeToStorage(meta, broker_->storage().get()); + } + + //! Get meta information from entity + int get_index_meta(IndexMeta *meta) const { + return IndexHelper::DeserializeFromStorage(broker_->storage().get(), meta); + } + + //! Set params: chunk size + inline void set_chunk_size(size_t val) { + chunk_size_ = val; + } + + //! Set params + inline void set_filter_same_key(bool val) { + filter_same_key_ = val; + } + + //! Set params + inline void set_get_vector(bool val) { + get_vector_enabled_ = val; + } + + //! Get vector local id by key + inline node_id_t get_id(key_t key) const { + if (use_key_info_map_) { + keys_map_lock_->lock_shared(); + auto it = keys_map_->find(key); + keys_map_lock_->unlock_shared(); + return it == keys_map_->end() ? kInvalidNodeId : it->second; + } else { + return key; + } + } + + void print_key_map() const { + std::cout << "key map begins" << std::endl; + + auto iter = keys_map_->begin(); + while (iter != keys_map_->end()) { + std::cout << "key: " << iter->first << ", id: " << iter->second + << std::endl; + ; + iter++; + } + + std::cout << "key map ends" << std::endl; + } + + //! Typed get_neighbors: returns NeighborsT without runtime + //! branching on MemoryBlock type. MmapMemoryBlock specialization uses + //! pointer-based read; BufferPoolMemoryBlock uses MemoryBlock-based read. + template + inline NeighborsT get_neighbors_typed(level_t level, + node_id_t id) const; + + //! Typed batch get_vector: fills vector without runtime branching + template + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const; + + //! Typed get_key: reads key using typed MemBlock + template + inline key_t get_key_typed(node_id_t id) const; + + //! Get l0 neighbors size + inline size_t neighbors_size() const { + return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); + } + + //! Get neighbors size for level > 0 + inline size_t upper_neighbors_size() const { + return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); + } + + inline size_t max_degree(level_t level) const { + return level == 0 ? neighbor_size_ : upper_neighbor_size_; + } + + + protected: + union UpperNeighborIndexMeta { + struct { + uint32_t level : 4; + uint32_t index : 28; // index is composite type: chunk idx, and the + // N th neighbors in chunk, they two composite + // the 28 bits location + } bits; + uint32_t data; + }; + + protected: + template + using HashMap = google::dense_hash_map>; + template + using HashMapPointer = std::shared_ptr>; + + template + using HashSet = google::dense_hash_set>; + template + using HashSetPointer = std::shared_ptr>; + + //! upper neighbor index hashmap + using NIHashMap = HnswIndexHashMap; + using NIHashMapPointer = std::shared_ptr; + + //! Clone construct, used by clone method in subclasses + HnswStreamerEntity( + IndexStreamer::Stats &stats, const HNSWHeader &hd, size_t chunk_size, + uint32_t node_index_mask_bits, uint32_t upper_neighbor_mask_bits, + bool filter_same_key, bool get_vector_enabled, + const NIHashMapPointer &upper_neighbor_index, + const std::shared_ptr &upper_neighbor_rw_mutex, + std::shared_ptr &keys_map_lock, + const HashMapPointer &keys_map, bool use_key_info_map, + std::vector &&node_chunks, + std::vector &&upper_neighbor_chunks, + const ChunkBroker::Pointer &broker, + std::shared_ptr> node_bases, + std::shared_ptr> upper_bases) + : stats_(stats), + upper_neighbor_rw_mutex_(upper_neighbor_rw_mutex), + chunk_size_(chunk_size), + node_index_mask_bits_(node_index_mask_bits), + node_cnt_per_chunk_(1UL << node_index_mask_bits_), + node_index_mask_(node_cnt_per_chunk_ - 1), + upper_neighbor_mask_bits_(upper_neighbor_mask_bits), + upper_neighbor_mask_((1U << upper_neighbor_mask_bits_) - 1), + filter_same_key_(filter_same_key), + get_vector_enabled_(get_vector_enabled), + use_key_info_map_(use_key_info_map), + upper_neighbor_index_(upper_neighbor_index), + keys_map_lock_(keys_map_lock), + keys_map_(keys_map), + node_chunks_(std::move(node_chunks)), + upper_neighbor_chunks_(std::move(upper_neighbor_chunks)), + broker_(broker) { + *mutable_header() = hd; + + neighbor_size_ = neighbors_size(); + upper_neighbor_size_ = upper_neighbors_size(); + + // Reuse the shared base-pointer arrays created by init_chunks(). + // All clones share the same arrays so hot HNSW hub-node chunks are + // collectively promoted to L1/L2 by every search thread instead of + // each clone warming its own private copy in L3. + node_chunk_bases_ = std::move(node_bases); + upper_neighbor_chunk_bases_ = std::move(upper_bases); + } + + //! Called only in searching procedure per context, so no need to lock + void sync_chunks(ChunkBroker::CHUNK_TYPE type, size_t idx, + std::vector *chunks) const { + if (ailego_likely(idx < chunks->size())) { + return; + } + for (size_t i = chunks->size(); i <= idx; ++i) { + auto chunk = broker_->get_chunk(type, i); + // the storage can ensure get chunk will success after the first get + ailego_assert_with(!!chunk, "get chunk failed"); + chunks->emplace_back(std::move(chunk)); + } + } + + //! return pair: chunk index + chunk offset + inline std::pair get_vector_chunk_loc( + node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size(); + + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + //! return pair: chunk index + chunk offset + inline std::pair get_key_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + inline std::pair get_upper_neighbor_chunk_loc( + level_t level, node_id_t id) const { + // Shared lock: concurrent readers are fine, but must synchronize with + // add_upper_neighbor's exclusive lock to avoid data-race on + // slots_.size() inside HnswIndexHashMap. + std::shared_lock lk(*upper_neighbor_rw_mutex_); + auto it = upper_neighbor_index_->find(id); + ailego_assert_abort(it != upper_neighbor_index_->end(), + "Get upper neighbor header failed"); + auto meta = reinterpret_cast(&it->second); + uint32_t chunk_idx = (meta->bits.index) >> upper_neighbor_mask_bits_; + uint32_t offset = + (((meta->bits.index) & upper_neighbor_mask_) + level - 1) * + upper_neighbor_size_; + sync_chunks(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, chunk_idx, + &upper_neighbor_chunks_); + ailego_assert_abort(chunk_idx < upper_neighbor_chunks_.size(), + "invalid chunk idx"); + ailego_assert_abort(offset < upper_neighbor_chunks_[chunk_idx]->data_size(), + "invalid chunk offset"); + return std::make_pair(chunk_idx, offset); + } + + //! return pair: chunk + chunk offset + inline std::pair get_neighbor_chunk_loc(level_t level, + node_id_t id) const { + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_abort(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_abort(offset < node_chunks_[chunk_idx]->data_size(), + "invalid chunk offset"); + return std::make_pair(node_chunks_[chunk_idx].get(), offset); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + return std::make_pair(upper_neighbor_chunks_[p.first].get(), p.second); + } + } + + //! Chunk hnsw index valid + int check_hnsw_index(const HNSWHeader *hd) const; + + size_t get_total_upper_neighbors_size(level_t level) const { + return level * upper_neighbor_size_; + } + + //! Add upper neighbor header and reserve space for upper neighbor + int add_upper_neighbor(level_t level, node_id_t id) { + if (level == 0) { + return 0; + } + // Exclusive lock: protects upper_neighbor_chunks_.emplace_back() and + // upper_neighbor_index_->insert() from racing with concurrent find() + // calls in get_upper_neighbor_chunk_loc(). + std::unique_lock lk(*upper_neighbor_rw_mutex_); + Chunk::Pointer chunk; + uint64_t chunk_offset = UINT64_MAX; + size_t neighbors_size = get_total_upper_neighbors_size(level); + uint64_t chunk_index = upper_neighbor_chunks_.size() - 1ULL; + if (chunk_index == UINT64_MAX || + (upper_neighbor_chunks_[chunk_index]->padding_size() < + neighbors_size)) { // no space left and need to alloc + chunk_index++; + if (ailego_unlikely(upper_neighbor_chunks_.capacity() == + upper_neighbor_chunks_.size())) { + LOG_ERROR("add upper neighbor failed for no memory quota"); + return IndexError_IndexFull; + } + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, + chunk_index, upper_neighbor_chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + chunk = p.second; + chunk_offset = 0UL; + upper_neighbor_chunks_.emplace_back(chunk); + } else { + chunk = upper_neighbor_chunks_[chunk_index]; + chunk_offset = chunk->data_size(); + } + ailego_assert_with((size_t)level < kMaxGraphLayers, "invalid level"); + ailego_assert_with(chunk_offset % upper_neighbor_size_ == 0, + "invalid offset"); + ailego_assert_with((chunk_offset / upper_neighbor_size_) < + (1U << upper_neighbor_mask_bits_), + "invalid offset"); + ailego_assert_with(chunk_index < (1U << (28 - upper_neighbor_mask_bits_)), + "invalid chunk index"); + UpperNeighborIndexMeta meta; + meta.bits.level = level; + meta.bits.index = (chunk_index << upper_neighbor_mask_bits_) | + (chunk_offset / upper_neighbor_size_); + size_t zero_start = chunk_offset; + chunk_offset += upper_neighbor_size_ * level; + + // IMPORTANT: order matters here. + // 1) resize so the chunk's data_size covers the new region. + // 2) zero-fill the new region: storage backends like BufferStorage do + // NOT zero on resize -- only metadata is updated, and the underlying + // page may contain stale content from a previously-evicted page. + // Without this step, NeighborsHeader::neighbor_cnt is garbage and + // select_entry_point()/search_neighbors() iterate over garbage + // node_ids, eventually triggering find()'s assertion in + // get_upper_neighbor_chunk_loc(). + // 3) ONLY THEN publish the entry to upper_neighbor_index_, so that any + // concurrent reader that finds this id already sees a properly + // zeroed upper-neighbor slot. + if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", (size_t)chunk_offset); + return IndexError_Runtime; + } + + // Use std::vector instead of a VLA: VLAs are a GNU extension and may + // produce different codegen / be rejected under clang/MSVC. + std::vector zeros(neighbors_size, 0); + if (ailego_unlikely(chunk->write(zero_start, zeros.data(), + neighbors_size) != neighbors_size)) { + LOG_ERROR("Chunk write zeros failed"); + return IndexError_Runtime; + } + + if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) { + LOG_ERROR("HashMap insert value failed"); + return IndexError_Runtime; + } + + return 0; + } + + size_t estimate_doc_capacity() const { + return node_chunks_.capacity() * node_cnt_per_chunk_; + } + + int init_chunk_params(size_t max_index_size, bool huge_page) { + node_cnt_per_chunk_ = std::max(1, chunk_size_ / node_size()); + //! align node cnt per chunk to pow of 2 + node_index_mask_bits_ = std::ceil(std::log2(node_cnt_per_chunk_)); + node_cnt_per_chunk_ = 1UL << node_index_mask_bits_; + if (huge_page) { + chunk_size_ = AlignHugePageSize(node_cnt_per_chunk_ * node_size()); + } else { + chunk_size_ = AlignPageSize(node_cnt_per_chunk_ * node_size()); + } + node_index_mask_ = node_cnt_per_chunk_ - 1; + + if (max_index_size == 0UL) { + max_index_size_ = chunk_size_ * kDefaultMaxChunkCnt; + } else { + max_index_size_ = max_index_size; + } + + //! To get a balanced upper neighbor chunk size. + //! If the upper chunk size is equal to node chunk size, it may waste + //! upper neighbor chunk space; if the upper neighbor chunk size is too + //! small, the will need large upper neighbor chunks index space. So to + //! get a balanced ratio be sqrt of the node/neighbor size ratio + float ratio = + std::sqrt(node_size() * scaling_factor() * 1.0f / upper_neighbor_size_); + if (huge_page) { + upper_neighbor_chunk_size_ = AlignHugePageSize( + std::max(get_total_upper_neighbors_size(kMaxGraphLayers), + static_cast(chunk_size_ / ratio))); + } else { + upper_neighbor_chunk_size_ = AlignPageSize( + std::max(get_total_upper_neighbors_size(kMaxGraphLayers), + static_cast(chunk_size_ / ratio))); + } + upper_neighbor_mask_bits_ = + std::ceil(std::log2(upper_neighbor_chunk_size_ / upper_neighbor_size_)); + upper_neighbor_mask_ = (1 << upper_neighbor_mask_bits_) - 1; + + size_t max_node_chunk_cnt = std::ceil(max_index_size_ / chunk_size_); + size_t max_upper_chunk_cnt = std::ceil( + (max_node_chunk_cnt * node_cnt_per_chunk_ * 1.0f / scaling_factor()) / + (upper_neighbor_chunk_size_ / upper_neighbor_size_)); + max_upper_chunk_cnt = + max_upper_chunk_cnt + std::ceil(max_upper_chunk_cnt / scaling_factor()); + + //! reserve space to avoid memmove in chunks vector emplace chunk, so + //! as to lock-free in reading chunk + node_chunks_.reserve(max_node_chunk_cnt); + upper_neighbor_chunks_.reserve(max_upper_chunk_cnt); + + LOG_DEBUG( + "Settings: nodeSize=%zu chunkSize=%u upperNeighborSize=%u " + "upperNeighborChunkSize=%u " + "nodeCntPerChunk=%u maxChunkCnt=%zu maxNeighborChunkCnt=%zu " + "maxIndexSize=%zu ratio=%.3f", + node_size(), chunk_size_, upper_neighbor_size_, + upper_neighbor_chunk_size_, node_cnt_per_chunk_, max_node_chunk_cnt, + max_upper_chunk_cnt, max_index_size_, ratio); + + return 0; + } + + //! Init node chunk and neighbor chunks + int init_chunks(const Chunk::Pointer &header_chunk); + + int flush_header(void) { + if (!broker_->dirty()) { + // do not need to flush + return 0; + } + auto header_chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId); + if (ailego_unlikely(!header_chunk)) { + LOG_ERROR("get header chunk failed"); + return IndexError_Runtime; + } + size_t size = header_chunk->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + + return 0; + } + + protected: + //! Expose sync_chunks for subclass use + inline void sync_node_chunks(size_t idx) const { + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, idx, &node_chunks_); + } + inline void sync_upper_neighbor_chunks(size_t idx) const { + sync_chunks(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, idx, + &upper_neighbor_chunks_); + } + + private: + HnswStreamerEntity(const HnswStreamerEntity &) = delete; + HnswStreamerEntity &operator=(const HnswStreamerEntity &) = delete; + static constexpr uint64_t kUpperHashMemoryInflateRatio = 2.0f; + + protected: + IndexStreamer::Stats &stats_; + std::mutex mutex_{}; + //! Guards upper_neighbor_index_ and upper_neighbor_chunks_ against + //! concurrent reads (find) and writes (insert/emplace_back). + //! Shared via shared_ptr so all clones synchronize on the SAME mutex. + mutable std::shared_ptr upper_neighbor_rw_mutex_{}; + size_t max_index_size_{0UL}; + uint32_t chunk_size_{kDefaultChunkSize}; + uint32_t upper_neighbor_chunk_size_{kDefaultChunkSize}; + uint32_t node_index_mask_bits_{0U}; + uint32_t node_cnt_per_chunk_{0U}; + uint32_t node_index_mask_{0U}; + uint32_t neighbor_size_{0U}; + uint32_t upper_neighbor_size_{0U}; + //! UpperNeighborIndex.index composite chunkIdx and offset in chunk by the + //! following mask + uint32_t upper_neighbor_mask_bits_{0U}; + uint32_t upper_neighbor_mask_{0U}; + bool filter_same_key_{false}; + bool get_vector_enabled_{false}; + bool use_key_info_map_{true}; + + NIHashMapPointer upper_neighbor_index_{}; + + mutable std::shared_ptr keys_map_lock_{}; + HashMapPointer keys_map_{}; + + //! the chunks will be changed in searcher, so need mutable + //! data chunk include: vector, key, level 0 neighbors + mutable std::vector node_chunks_{}; + + //! Flat cache of base_data() pointers for node_chunks_ and + //! upper_neighbor_chunks_. Non-empty only when the storage backend + //! returns a stable mmap pointer (base_data() != nullptr). Avoids + //! following the full shared_ptr -> Segment -> IndexMapping::Segment + //! pointer chain on every get_vector() / get_neighbors() call, which + //! is critical for small chunk sizes (e.g. 16 K) where node_chunks_ + //! can hold 100K+ entries and the metadata no longer fits in L2 cache. + //! + //! Shared across all clones (read-only after open) so that hot entries + //! (hub-node chunks near the HNSW entry point) are promoted to L1/L2 + //! by all search threads collectively, instead of each clone warming + //! its own private 250 KB copy in L3. + mutable std::shared_ptr> node_chunk_bases_{}; + + //! upper neighbor chunk inlude: UpperNeighborHeader + (1~level) neighbors + mutable std::vector upper_neighbor_chunks_{}; + mutable std::shared_ptr> + upper_neighbor_chunk_bases_{}; + + ChunkBroker::Pointer broker_{}; // chunk broker +}; + +// --- Template specializations for typed MemoryBlock access --- + +//! MmapMemoryBlock specialization: uses pointer-based Chunk::read +template <> +inline NeighborsT +HnswStreamerEntity::get_neighbors_typed(level_t level, + node_id_t id) const { + Chunk *chunk = nullptr; + size_t offset = 0UL; + size_t nbr_size = neighbor_size_; + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + chunk = node_chunks_[chunk_idx].get(); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + chunk = upper_neighbor_chunks_[p.first].get(); + offset = p.second; + nbr_size = upper_neighbor_size_; + } + ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); + const void *ptr = nullptr; + size_t ret = chunk->read(offset, &ptr, nbr_size); + if (ailego_unlikely(ret != nbr_size)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return NeighborsT(); + } + MmapMemoryBlock block(const_cast(ptr)); + return NeighborsT(std::move(block)); +} + +//! BufferPoolMemoryBlock specialization: uses MemoryBlock-based Chunk::read +template <> +inline NeighborsT +HnswStreamerEntity::get_neighbors_typed( + level_t level, node_id_t id) const { + Chunk *chunk = nullptr; + size_t offset = 0UL; + size_t nbr_size = neighbor_size_; + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + chunk = node_chunks_[chunk_idx].get(); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + chunk = upper_neighbor_chunks_[p.first].get(); + offset = p.second; + nbr_size = upper_neighbor_size_; + } + ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); + IndexStorage::MemoryBlock mem_block; + size_t ret = chunk->read(offset, mem_block, nbr_size); + if (ailego_unlikely(ret != nbr_size)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return NeighborsT(); + } + BufferPoolMemoryBlock block; + if (mem_block.type_ == IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH) { + block = BufferPoolMemoryBlock::MakeOwned(mem_block.data_); + mem_block.data_ = nullptr; + mem_block.type_ = IndexStorage::MemoryBlock::MBT_UNKNOWN; + } else { + block = BufferPoolMemoryBlock(mem_block.buffer_pool_handle_, + mem_block.buffer_block_id_, mem_block.data_); + mem_block.buffer_pool_handle_ = nullptr; + } + return NeighborsT(std::move(block)); +} + +//! MmapMemoryBlock specialization for batch get_vector +template <> +inline int HnswStreamerEntity::get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t read_size = vector_size(); + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + vec_blocks[i].reset(const_cast(ptr)); + } + return 0; +} + +//! BufferPoolMemoryBlock specialization for batch get_vector +template <> +inline int HnswStreamerEntity::get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t read_size = vector_size(); + IndexStorage::MemoryBlock mem_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, mem_block, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + vec_blocks[i] = [&]() { + if (mem_block.type_ == IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH) { + BufferPoolMemoryBlock b = + BufferPoolMemoryBlock::MakeOwned(mem_block.data_); + mem_block.data_ = nullptr; + mem_block.type_ = IndexStorage::MemoryBlock::MBT_UNKNOWN; + return b; + } + BufferPoolMemoryBlock b(mem_block.buffer_pool_handle_, + mem_block.buffer_block_id_, mem_block.data_); + mem_block.buffer_pool_handle_ = nullptr; + return b; + }(); + } + return 0; +} + +//! MmapMemoryBlock specialization for get_key +template <> +inline key_t HnswStreamerEntity::get_key_typed( + node_id_t id) const { + if (!use_key_info_map_) { + return id; + } + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(ptr); +} + +//! BufferPoolMemoryBlock specialization for get_key +template <> +inline key_t HnswStreamerEntity::get_key_typed( + node_id_t id) const { + if (!use_key_info_map_) { + return id; + } + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + IndexStorage::MemoryBlock key_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(key_block.data()); +} + +//! Typed entity subclass for mmap mode. +//! Caches chunk base addresses to eliminate virtual function calls on the +//! search hot path. For mmap mode, chunk data is memory-mapped at init time, +//! so we can directly compute pointers via base_addr + offset. +class HnswMmapStreamerEntity : public HnswStreamerEntity { + public: + using MemoryBlock = MmapMemoryBlock; + using TypedNeighbors = NeighborsT; + + using HnswStreamerEntity::HnswStreamerEntity; + + HnswStorageMode storage_mode() const override { + return HnswStorageMode::kMmap; + } + + //! Override clone to return correct subclass type, so that + //! static_cast in the algorithm is safe. + const HnswEntity::Pointer clone() const override; + + ailego_force_inline TypedNeighbors get_neighbors_typed(level_t level, + node_id_t id) const { + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + const char *base = get_node_chunk_base(chunk_idx); + MmapMemoryBlock block(const_cast(base + offset)); + return TypedNeighbors(std::move(block)); + } + // Upper level: use index to locate chunk and offset + auto it = upper_neighbor_index_->find(id); + ailego_assert_abort(it != upper_neighbor_index_->end(), + "Get upper neighbor header failed"); + auto meta = reinterpret_cast(&it->second); + uint32_t chunk_idx = (meta->bits.index) >> upper_neighbor_mask_bits_; + uint32_t offset = + (((meta->bits.index) & upper_neighbor_mask_) + level - 1) * + upper_neighbor_size_; + const char *base = get_upper_neighbor_chunk_base(chunk_idx); + MmapMemoryBlock block(const_cast(base + offset)); + return TypedNeighbors(std::move(block)); + } + + ailego_force_inline int get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + uint32_t chunk_idx = ids[i] >> node_index_mask_bits_; + uint32_t offset = (ids[i] & node_index_mask_) * node_size(); + const char *base = get_node_chunk_base(chunk_idx); + vec_blocks[i].reset(const_cast(base + offset)); + } + return 0; + } + + ailego_force_inline key_t get_key_typed(node_id_t id) const { + if (!use_key_info_map_) { + return id; + } + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + const char *base = get_node_chunk_base(chunk_idx); + return *reinterpret_cast(base + offset); + } + + //! Direct vector pointer access (no MemoryBlock wrapper). + //! For use in the merged search loop to avoid intermediate allocations. + ailego_force_inline const void *get_vector_ptr(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size(); + return get_node_chunk_base(chunk_idx) + offset; + } + + protected: + //! Get cached base address for a node chunk, syncing if needed + ailego_force_inline const char *get_node_chunk_base( + uint32_t chunk_idx) const { + if (ailego_unlikely(chunk_idx >= node_chunk_bases_.size())) { + sync_node_chunk_bases(chunk_idx); + } + return node_chunk_bases_[chunk_idx]; + } + + //! Get cached base address for an upper neighbor chunk, syncing if needed + ailego_force_inline const char *get_upper_neighbor_chunk_base( + uint32_t chunk_idx) const { + if (ailego_unlikely(chunk_idx >= upper_neighbor_chunk_bases_.size())) { + sync_upper_neighbor_chunk_bases(chunk_idx); + } + return upper_neighbor_chunk_bases_[chunk_idx]; + } + + //! Sync node chunk base addresses up to the given index + void sync_node_chunk_bases(uint32_t chunk_idx) const { + sync_node_chunks(chunk_idx); + const auto &chunks = node_chunks_; + for (size_t i = node_chunk_bases_.size(); i <= chunk_idx; ++i) { + const void *ptr = nullptr; + chunks[i]->read(0, &ptr, 1); + node_chunk_bases_.push_back(static_cast(ptr)); + } + } + + //! Sync upper neighbor chunk base addresses up to the given index + void sync_upper_neighbor_chunk_bases(uint32_t chunk_idx) const { + sync_upper_neighbor_chunks(chunk_idx); + const auto &chunks = upper_neighbor_chunks_; + for (size_t i = upper_neighbor_chunk_bases_.size(); i <= chunk_idx; ++i) { + const void *ptr = nullptr; + chunks[i]->read(0, &ptr, 1); + upper_neighbor_chunk_bases_.push_back(static_cast(ptr)); + } + } + + mutable std::vector node_chunk_bases_{}; + mutable std::vector upper_neighbor_chunk_bases_{}; +}; + +//! Typed entity subclass for buffer pool mode. +class HnswBufferPoolStreamerEntity : public HnswStreamerEntity { + public: + using MemoryBlock = BufferPoolMemoryBlock; + using TypedNeighbors = NeighborsT; + + using HnswStreamerEntity::HnswStreamerEntity; + + HnswStorageMode storage_mode() const override { + return HnswStorageMode::kBufferPool; + } + + inline TypedNeighbors get_neighbors_typed(level_t level, node_id_t id) const { + return HnswStreamerEntity::get_neighbors_typed(level, + id); + } + + inline int get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + return HnswStreamerEntity::get_vector_typed( + ids, count, vec_blocks); + } + + inline key_t get_key_typed(node_id_t id) const { + return HnswStreamerEntity::get_key_typed(id); + } +}; + +//! Typed entity subclass for contiguous memory mode. +//! Splits node data into two dense arrays during build: +//! 1. vector_base_: flat vector array (stride = vector_size) +//! 2. graph_base_: key + L0 neighbors (stride = graph_stride_) +//! Total memory = vector_size + graph_stride_ per node (same as original +//! node_size), but each access pattern gets optimal cache locality. +class HnswContiguousStreamerEntity : public HnswMmapStreamerEntity { + public: + using HnswMmapStreamerEntity::HnswMmapStreamerEntity; + + HnswStorageMode storage_mode() const override { + return HnswStorageMode::kContiguous; + } + + //! Override clone to return correct subclass type. + //! Cloned entity shares contiguous memory via shared_ptr. + const HnswEntity::Pointer clone() const override; + + ~HnswContiguousStreamerEntity() = default; + + //! Build contiguous memory from chunks after open. + //! Must be called after the entity is fully opened and all chunks are loaded. + int build_contiguous_memory(); + + //! Degrade to mmap mode by releasing contiguous memory and falling back + //! to chunk-based access. + void degrade_to_mmap() { + vector_memory_.reset(); + vector_base_ = nullptr; + graph_memory_.reset(); + graph_base_ = nullptr; + upper_neighbor_memory_.reset(); + upper_neighbor_base_ = nullptr; + upper_chunk_offsets_.clear(); + LOG_INFO("HNSW contiguous entity degraded to mmap mode for insertion"); + } + + bool is_contiguous() const { + return vector_base_ != nullptr; + } + + int add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) override { + if (ailego_unlikely(is_contiguous())) degrade_to_mmap(); + return HnswMmapStreamerEntity::add_vector(level, key, vec, id); + } + + int add_vector_with_id(level_t level, node_id_t id, + const void *vec) override { + if (ailego_unlikely(is_contiguous())) degrade_to_mmap(); + return HnswMmapStreamerEntity::add_vector_with_id(level, id, vec); + } + + ailego_force_inline TypedNeighbors get_neighbors_typed(level_t level, + node_id_t id) const { + if (ailego_likely(graph_base_ != nullptr)) { + if (level == 0UL) { + // graph layout: [key (sizeof(key_t)) | NeighborsHeader + neighbors] + const char *ptr = graph_base_ + + static_cast(id) * graph_stride_ + + sizeof(key_t); + MmapMemoryBlock block(const_cast(ptr)); + return TypedNeighbors(std::move(block)); + } + // Upper level: use index to locate global offset + auto it = upper_neighbor_index_->find(id); + ailego_assert_abort(it != upper_neighbor_index_->end(), + "Get upper neighbor header failed"); + auto meta = reinterpret_cast(&it->second); + uint32_t chunk_idx = (meta->bits.index) >> upper_neighbor_mask_bits_; + uint32_t local_idx = (meta->bits.index) & upper_neighbor_mask_; + size_t global_offset = + upper_chunk_offsets_[chunk_idx] + + static_cast(local_idx + level - 1) * upper_neighbor_size_; + const char *ptr = upper_neighbor_base_ + global_offset; + MmapMemoryBlock block(const_cast(ptr)); + return TypedNeighbors(std::move(block)); + } + return HnswMmapStreamerEntity::get_neighbors_typed(level, id); + } + + ailego_force_inline int get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + if (ailego_likely(vector_base_ != nullptr)) { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + const char *ptr = + vector_base_ + static_cast(ids[i]) * vector_size(); + vec_blocks[i].reset(const_cast(ptr)); + } + return 0; + } + return HnswMmapStreamerEntity::get_vector_typed(ids, count, vec_blocks); + } + + ailego_force_inline key_t get_key_typed(node_id_t id) const { + if (ailego_likely(graph_base_ != nullptr)) { + if (!use_key_info_map_) { + return id; + } + const char *ptr = graph_base_ + static_cast(id) * graph_stride_; + return *reinterpret_cast(ptr); + } + return HnswMmapStreamerEntity::get_key_typed(id); + } + + //! Direct vector pointer from flat vector array (stride = vector_size). + //! For use in the merged search loop to avoid intermediate allocations. + ailego_force_inline const void *get_vector_ptr(node_id_t id) const { + if (ailego_likely(vector_base_ != nullptr)) { + return vector_base_ + static_cast(id) * vector_size(); + } + // Fallback to mmap chunk-based access + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size(); + return get_node_chunk_base(chunk_idx) + offset; + } + + protected: + //! Custom deleter for contiguous memory (munmap / _aligned_free / free) + //! Used by shared_ptr to properly release mmap'd memory. + struct ContiguousDeleter { + size_t size; + void operator()(char *ptr) const { + if (!ptr) return; +#if defined(__linux__) || defined(__APPLE__) + ::munmap(ptr, size); +#elif defined(_WIN32) + ::_aligned_free(ptr); +#else + std::free(ptr); +#endif + } + }; + + //! Flat vector array: vectors stored densely (stride = vector_size). + std::shared_ptr vector_memory_{}; + char *vector_base_{nullptr}; + + //! Graph array: [key | L0 neighbors] stored densely (stride = graph_stride_). + std::shared_ptr graph_memory_{}; + char *graph_base_{nullptr}; + size_t graph_stride_{0}; // sizeof(key_t) + neighbor_size_ + + //! Shared ownership of upper neighbor contiguous memory + std::shared_ptr upper_neighbor_memory_{}; + char *upper_neighbor_base_{nullptr}; + + //! Cumulative offsets for each upper neighbor chunk in contiguous memory + std::vector upper_chunk_offsets_{}; + + private: + //! Allocate contiguous memory with hugepage/THP support + static char *allocate_contiguous(size_t size); +}; + +//! Typed entity subclass that reads vectors from an external vector source. +//! The graph structure (key + neighbors) is stored in chunks just like +//! HnswMmapStreamerEntity, but the per-node vector prefix is removed by +//! setting vector_size() to 0 (see HnswStreamer setup). With vector_size()==0 +//! all inherited offset computations (key at node start, L0 neighbors right +//! after the key, node_size == AlignSize(sizeof(key) + neighbor_size)) become +//! automatically correct, and base add_vector writes a zero-byte vector (i.e. +//! it skips vector storage). Vectors are instead read through the bound +//! VectorSource, which is supplied per add/search call. +class HnswExternalStreamerEntity : public HnswMmapStreamerEntity { + public: + using MemoryBlock = MmapMemoryBlock; + using TypedNeighbors = NeighborsT; + + using HnswMmapStreamerEntity::HnswMmapStreamerEntity; + + HnswStorageMode storage_mode() const override { + return HnswStorageMode::kExternal; + } + + //! Override clone to return the correct subclass type, so that + //! static_cast in the algorithm is safe. + //! The external vector source is NOT shared with the clone; it is re-applied + //! per add/search call (via HnswContext::set_vector_source). + const HnswEntity::Pointer clone() const override; + + //! Bind the external vector source for the current add/search call. + void set_vector_source(const VectorSource *src) override { + vec_src_ = src; + } + + //! Typed batch get_vector: zero-copy view into the external vector source. + //! Hides HnswMmapStreamerEntity::get_vector_typed (non-virtual, used by the + //! template algorithm via static_cast). + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + if (ailego_unlikely(vec_src_ == nullptr)) { + return IndexError_Runtime; + } + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + vec_blocks[i].reset(const_cast(vec_src_->get_vector(ids[i]))); + } + return 0; + } + + //! Virtual get_vector overrides (distance-calculator / provider paths). + const void *get_vector(node_id_t id) const override { + return vec_src_ ? vec_src_->get_vector(id) : nullptr; + } + + int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const override { + if (ailego_unlikely(vec_src_ == nullptr)) { + return IndexError_Runtime; + } + vec_src_->get_vectors(ids, count, vecs); + return 0; + } + + int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const override { + if (ailego_unlikely(vec_src_ == nullptr)) { + return IndexError_Runtime; + } + block.reset(const_cast(vec_src_->get_vector(id))); + return 0; + } + + int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const override { + if (ailego_unlikely(vec_src_ == nullptr)) { + return IndexError_Runtime; + } + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + vec_blocks[i].reset(const_cast(vec_src_->get_vector(ids[i]))); + } + return 0; + } + + private: + //! Transient, per-call vector source. Never shared across clones; bound by + //! HnswContext::set_vector_source before each add/search. + const VectorSource *vec_src_{nullptr}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/CMakeLists.txt b/src/core/algorithm/hnsw_rabitq/CMakeLists.txt new file mode 100644 index 0000000..09ce72f --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/CMakeLists.txt @@ -0,0 +1,27 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(AUTO_DETECT_ARCH) + foreach(FILE ${HNSW_RABITQ_FILES}) + set_source_files_properties( + ${FILE} + PROPERTIES + COMPILE_FLAGS "${RABITQ_ARCH_FLAG}" + ) + endforeach() +endif() + +if(NOT APPLE) + set(CORE_KNN_HNSW_RABITQ_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_hnsw_rabitq + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework rabitqlib sparsehash + INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_HNSW_RABITQ_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.cc new file mode 100644 index 0000000..c8379a8 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.cc @@ -0,0 +1,382 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_rabitq_algorithm.h" +#include +#include "hnsw_rabitq_entity.h" + +namespace zvec { +namespace core { + +HnswRabitqAlgorithm::HnswRabitqAlgorithm(HnswRabitqEntity &entity) + : entity_(entity), + mt_(std::chrono::system_clock::now().time_since_epoch().count()), + lock_pool_(kLockCnt) {} + +int HnswRabitqAlgorithm::cleanup() { + return 0; +} + +int HnswRabitqAlgorithm::add_node(node_id_t id, level_t level, + HnswRabitqContext *ctx) { + spin_lock_.lock(); + + // std::cout << "id: " << id << ", level: " << level << std::endl; + + auto cur_max_level = entity_.cur_max_level(); + auto entry_point = entity_.entry_point(); + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + entity_.update_ep_and_level(id, level); + spin_lock_.unlock(); + return 0; + } + spin_lock_.unlock(); + + if (ailego_unlikely(level > cur_max_level)) { + mutex_.lock(); + // re-check max level + cur_max_level = entity_.cur_max_level(); + entry_point = entity_.entry_point(); + if (level <= cur_max_level) { + mutex_.unlock(); + } + } + + level_t cur_level = cur_max_level; + ResultRecord dist = ctx->dist_calculator()(entry_point); + for (; cur_level > level; --cur_level) { + select_entry_point(cur_level, &entry_point, &dist, ctx); + } + + for (; cur_level >= 0; --cur_level) { + search_neighbors(cur_level, &entry_point, &dist, ctx->level_topk(cur_level), + ctx); + } + + // add neighbors from down level to top level, to avoid upper level visible + // to knn_search but the under layer level not ready + for (cur_level = 0; cur_level <= level; ++cur_level) { + add_neighbors(id, cur_level, ctx->level_topk(cur_level), ctx); + ctx->level_topk(cur_level).clear(); + } + + if (ailego_unlikely(level > cur_max_level)) { + spin_lock_.lock(); + entity_.update_ep_and_level(id, level); + spin_lock_.unlock(); + mutex_.unlock(); + } + + return 0; +} + +//! select_entry_point on hnsw level, ef = 1 +void HnswRabitqAlgorithm::select_entry_point(level_t level, + node_id_t *entry_point, + ResultRecord *dist, + HnswRabitqContext *ctx) const { + auto &entity = ctx->get_entity(); + HnswRabitqAddDistCalculator &dc = ctx->dist_calculator(); + while (true) { + const Neighbors neighbors = entity.get_neighbors(level, *entry_point); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + uint32_t size = neighbors.size(); + if (size == 0) { + break; + } + + std::vector neighbor_vec_blocks; + int ret = dc.get_vector(&neighbors[0], size, neighbor_vec_blocks); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } + if (ailego_unlikely(ret != 0)) { + break; + } + + bool find_closer = false; + + std::vector dists(size); + std::vector neighbor_vecs(size); + for (uint32_t i = 0; i < size; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + + dc.batch_dist(neighbor_vecs.data(), size, dists.data()); + + for (uint32_t i = 0; i < size; ++i) { + ResultRecord cur_dist = dists[i]; + + if (cur_dist < *dist) { + *entry_point = neighbors[i]; + *dist = cur_dist; + find_closer = true; + } + } + + if (!find_closer) { + break; + } + } + + return; +} + +void HnswRabitqAlgorithm::add_neighbors(node_id_t id, level_t level, + TopkHeap &topk_heap, + HnswRabitqContext *ctx) { + if (ailego_unlikely(topk_heap.size() == 0)) { + return; + } + + HnswRabitqAddDistCalculator &dc = ctx->dist_calculator(); + + update_neighbors(dc, id, level, topk_heap); + + // reverse update neighbors + for (size_t i = 0; i < topk_heap.size(); ++i) { + reverse_update_neighbors(dc, topk_heap[i].first, level, id, + topk_heap[i].second, ctx->update_heap()); + } + + return; +} + +void HnswRabitqAlgorithm::search_neighbors(level_t level, + node_id_t *entry_point, + ResultRecord *dist, TopkHeap &topk, + HnswRabitqContext *ctx) const { + const auto &entity = ctx->get_entity(); + HnswRabitqAddDistCalculator &dc = ctx->dist_calculator(); + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + std::function filter = [](node_id_t) { return false; }; + if (ctx->filter().is_valid()) { + filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); }; + } + + candidates.clear(); + visit.clear(); + visit.set_visited(*entry_point); + if (!filter(*entry_point)) { + topk.emplace(*entry_point, *dist); + } + + candidates.emplace(*entry_point, *dist); + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t main_node = top->first; + ResultRecord main_dist = top->second; + + if (topk.full() && main_dist > topk[0].second) { + break; + } + + candidates.pop(); + const Neighbors neighbors = entity.get_neighbors(level, main_node); + ailego_prefetch(neighbors.data); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + + std::vector neighbor_ids(neighbors.size()); + uint32_t size = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) { + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_visit_dup_cnt())++; + } + continue; + } + visit.set_visited(node); + neighbor_ids[size++] = node; + } + if (size == 0) { + continue; + } + + std::vector neighbor_vec_blocks; + int ret = dc.get_vector(neighbor_ids.data(), size, neighbor_vec_blocks); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } + if (ailego_unlikely(ret != 0)) { + break; + } + + // do prefetch + static constexpr node_id_t BATCH_SIZE = 12; + static constexpr node_id_t PREFETCH_STEP = 2; + for (uint32_t i = 0; i < std::min(BATCH_SIZE * PREFETCH_STEP, size); ++i) { + ailego_prefetch(neighbor_vec_blocks[i].data()); + } + // done + + std::vector dists(size); + std::vector neighbor_vecs(size); + + for (uint32_t i = 0; i < size; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + + dc.batch_dist(neighbor_vecs.data(), size, dists.data()); + + for (uint32_t i = 0; i < size; ++i) { + node_id_t node = neighbor_ids[i]; + ResultRecord cur_dist = dists[i]; + + if ((!topk.full()) || cur_dist < topk[0].second) { + candidates.emplace(node, cur_dist); + // update entry_point for next level scan + if (cur_dist < *dist) { + *entry_point = node; + *dist = cur_dist; + } + if (!filter(node)) { + topk.emplace(node, cur_dist); + } + } // end if + } // end for + } // while + + return; +} + +void HnswRabitqAlgorithm::update_neighbors(HnswRabitqAddDistCalculator &dc, + node_id_t id, level_t level, + TopkHeap &topk_heap) { + topk_heap.sort(); + + uint32_t max_neighbor_cnt = entity_.neighbor_cnt(level); + if (topk_heap.size() <= static_cast(entity_.prune_cnt())) { + if (topk_heap.size() <= static_cast(max_neighbor_cnt)) { + entity_.update_neighbors(level, id, topk_heap); + return; + } + } + + uint32_t cur_size = 0; + for (size_t i = 0; i < topk_heap.size(); ++i) { + node_id_t cur_node = topk_heap[i].first; + ResultRecord cur_node_dist = topk_heap[i].second; + bool good = true; + for (uint32_t j = 0; j < cur_size; ++j) { + ResultRecord tmp_dist = dc.dist(cur_node, topk_heap[j].first); + if (tmp_dist <= cur_node_dist) { + good = false; + break; + } + } + + if (good) { + topk_heap[cur_size].first = cur_node; + topk_heap[cur_size].second = cur_node_dist; + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } + } + } + + // when after-prune neighbor count is too seldom, + // we use this strategy to make-up enough edges + // not only just make-up out-degrees + // we also make-up enough in-degrees + uint32_t min_neighbors = entity_.min_neighbor_cnt(); + for (size_t k = cur_size; cur_size < min_neighbors && k < topk_heap.size(); + ++k) { + bool exist = false; + for (size_t j = 0; j < cur_size; ++j) { + if (topk_heap[j].first == topk_heap[k].first) { + exist = true; + break; + } + } + if (!exist) { + topk_heap[cur_size].first = topk_heap[k].first; + topk_heap[cur_size].second = topk_heap[k].second; + cur_size++; + } + } + + topk_heap.resize(cur_size); + entity_.update_neighbors(level, id, topk_heap); + + return; +} + +void HnswRabitqAlgorithm::reverse_update_neighbors( + HnswRabitqAddDistCalculator &dc, node_id_t id, level_t level, + node_id_t link_id, ResultRecord dist, TopkHeap &update_heap) { + const size_t max_neighbor_cnt = entity_.neighbor_cnt(level); + + uint32_t lock_idx = id & kLockMask; + lock_pool_[lock_idx].lock(); + const Neighbors neighbors = entity_.get_neighbors(level, id); + size_t size = neighbors.size(); + ailego_assert_with(size <= max_neighbor_cnt, "invalid neighbor size"); + if (size < max_neighbor_cnt) { + entity_.add_neighbor(level, id, size, link_id); + lock_pool_[lock_idx].unlock(); + return; + } + + update_heap.emplace(link_id, dist); + + for (size_t i = 0; i < size; ++i) { + node_id_t node = neighbors[i]; + ResultRecord cur_dist = dc.dist(id, node); + update_heap.emplace(node, cur_dist); + } + + //! TODO: optimize prune + //! prune edges + update_heap.sort(); + size_t cur_size = 0; + for (size_t i = 0; i < update_heap.size(); ++i) { + node_id_t cur_node = update_heap[i].first; + ResultRecord cur_node_dist = update_heap[i].second; + bool good = true; + for (size_t j = 0; j < cur_size; ++j) { + ResultRecord tmp_dist = dc.dist(cur_node, update_heap[j].first); + if (tmp_dist <= cur_node_dist) { + good = false; + break; + } + } + + if (good) { + update_heap[cur_size].first = cur_node; + update_heap[cur_size].second = cur_node_dist; + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } + } + } + + update_heap.resize(cur_size); + entity_.update_neighbors(level, id, update_heap); + + lock_pool_[lock_idx].unlock(); + + update_heap.clear(); + + return; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.h new file mode 100644 index 0000000..861ca46 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.h @@ -0,0 +1,124 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "hnsw_rabitq_context.h" +#include "hnsw_rabitq_dist_calculator.h" +#include "hnsw_rabitq_entity.h" + +namespace zvec { +namespace core { + +//! hnsw graph algorithm implement +class HnswRabitqAlgorithm { + public: + typedef std::unique_ptr UPointer; + + public: + //! Constructor + explicit HnswRabitqAlgorithm(HnswRabitqEntity &entity); + + //! Destructor + ~HnswRabitqAlgorithm() = default; + + //! Cleanup HnswRabitqAlgorithm + int cleanup(); + + //! Add a node to hnsw graph + //! @id: the node unique id + //! @level: a node will be add to graph in each level [0, level] + //! return 0 on success, or errCode in failure + int add_node(node_id_t id, level_t level, HnswRabitqContext *ctx); + + //! Initiate HnswRabitqAlgorithm + int init() { + level_probas_.clear(); + double level_mult = + 1 / std::log(static_cast(entity_.scaling_factor())); + for (int level = 0;; level++) { + // refers faiss get_random_level alg + double proba = + std::exp(-level / level_mult) * (1 - std::exp(-1 / level_mult)); + if (proba < 1e-9) { + break; + } + level_probas_.push_back(proba); + } + + return 0; + } + + //! Generate a random level + //! return graph level + uint32_t get_random_level() const { + // gen rand float (0, 1) + double f = mt_() / static_cast(mt_.max()); + for (size_t level = 0; level < level_probas_.size(); level++) { + if (f < level_probas_[level]) { + return level; + } + f -= level_probas_[level]; + } + return level_probas_.size() - 1; + } + + private: + //! Select in upper layer to get entry point for next layer search + void select_entry_point(level_t level, node_id_t *entry_point, + ResultRecord *dist, HnswRabitqContext *ctx) const; + + //! update node id neighbors from topkHeap, and reverse link is also updated + void add_neighbors(node_id_t id, level_t level, TopkHeap &topk_heap, + HnswRabitqContext *ctx); + + //! Given a node id and level, search the nearest neighbors in graph + //! Note: the nearest neighbors result keeps in topk, and entry_point and + //! dist will be updated to current level nearest node id and distance + void search_neighbors(level_t level, node_id_t *entry_point, + ResultRecord *dist, TopkHeap &topk, + HnswRabitqContext *ctx) const; + + //! Update the node's neighbors + void update_neighbors(HnswRabitqAddDistCalculator &dc, node_id_t id, + level_t level, TopkHeap &topk_heap); + + //! Checking linkId could be id's new neighbor, and add as neighbor if true + //! @dc distance calculator + //! @updateHeap temporary heap in updating neighbors + void reverse_update_neighbors(HnswRabitqAddDistCalculator &dc, node_id_t id, + level_t level, node_id_t link_id, + ResultRecord dist, TopkHeap &update_heap); + + private: + HnswRabitqAlgorithm(const HnswRabitqAlgorithm &) = delete; + HnswRabitqAlgorithm &operator=(const HnswRabitqAlgorithm &) = delete; + + private: + static constexpr uint32_t kLockCnt{1U << 8}; + static constexpr uint32_t kLockMask{kLockCnt - 1U}; + + HnswRabitqEntity &entity_; + mutable std::mt19937 mt_{}; + std::vector level_probas_{}; + + mutable ailego::SpinMutex spin_lock_{}; // global spin lock + std::mutex mutex_{}; // global mutex + // TODO: spin lock? + std::vector lock_pool_{}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.cc new file mode 100644 index 0000000..8833a61 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.cc @@ -0,0 +1,221 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_rabitq_chunk.h" +#include +#include +#include +#include +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_helper.h" +#include "zvec/core/framework/index_logger.h" +#include "zvec/core/framework/index_streamer.h" + +namespace zvec { +namespace core { + +int HnswRabitqChunkBroker::init_storage(size_t chunk_size) { + chunk_meta_.clear(); + chunk_meta_.chunk_size = chunk_size; + chunk_meta_.create_time = ailego::Realtime::Seconds(); + stats_.set_create_time(chunk_meta_.create_time); + chunk_meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(chunk_meta_.update_time); + + //! alloc meta chunk + size_t size = sizeof(HnswChunkMeta); + size = (size + page_mask_) & (~page_mask_); + const std::string segment_id = + make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId); + int ret = stg_->append(segment_id, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Storage append segment failed for %s", IndexError::What(ret)); + return ret; + } + chunk_meta_segment_ = get_chunk(CHUNK_TYPE_META, kDefaultChunkSeqId); + if (ailego_unlikely(!chunk_meta_segment_)) { + LOG_ERROR("Get meta segment failed"); + return IndexError_Runtime; + } + + //! update meta info and write to storage + chunk_meta_.chunk_cnts[CHUNK_TYPE_META] += 1; + chunk_meta_.total_size += size; + (*stats_.mutable_index_size()) += size; + size = chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswChunkMeta))) { + LOG_ERROR("Storage write data failed, wsize=%zu", size); + return IndexError_WriteData; + } + + return 0; +} + +int HnswRabitqChunkBroker::load_storage(size_t chunk_size) { + IndexStorage::MemoryBlock data_block; + size_t size = chunk_meta_segment_->read(0UL, data_block, + chunk_meta_segment_->data_size()); + if (size != sizeof(HnswChunkMeta)) { + LOG_ERROR("Invalid hnsw meta chunk, read size=%zu chunk size=%zu", size, + chunk_meta_segment_->data_size()); + return IndexError_InvalidFormat; + } + std::memcpy(static_cast(&chunk_meta_), data_block.data(), size); + if (chunk_meta_.chunk_size != chunk_size) { + LOG_ERROR( + "Params hnsw chunk size=%zu mismatch from previous %zu " + "in index", + chunk_size, (size_t)chunk_meta_.chunk_size); + return IndexError_Mismatch; + } + + *stats_.mutable_check_point() = stg_->check_point(); + stats_.set_revision_id(chunk_meta_.revision_id); + stats_.set_update_time(chunk_meta_.update_time); + stats_.set_create_time(chunk_meta_.create_time); + + char create_time[32]; + char update_time[32]; + ailego::Realtime::Gmtime(chunk_meta_.create_time, "%Y-%m-%d %H:%M:%S", + create_time, sizeof(create_time)); + ailego::Realtime::Gmtime(chunk_meta_.update_time, "%Y-%m-%d %H:%M:%S", + update_time, sizeof(update_time)); + LOG_DEBUG( + "Load index, indexSize=%zu chunkSize=%zu nodeChunks=%zu " + "upperNeighborChunks=%zu revisionId=%zu " + "createTime=%s updateTime=%s", + (size_t)chunk_meta_.total_size, (size_t)chunk_meta_.chunk_size, + (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_NODE], + (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_UPPER_NEIGHBOR], + (size_t)chunk_meta_.revision_id, create_time, update_time); + + return 0; +} + +int HnswRabitqChunkBroker::open(IndexStorage::Pointer stg, + size_t max_index_size, size_t chunk_size, + bool check_crc) { + if (ailego_unlikely(stg_)) { + LOG_ERROR("An storage instance is already opened"); + return IndexError_Duplicate; + } + stg_ = std::move(stg); + if (stg_->isHugePage()) { + page_mask_ = ailego::MemoryHelper::HugePageSize() - 1; + } else { + page_mask_ = ailego::MemoryHelper::PageSize() - 1; + } + check_crc_ = check_crc; + max_chunks_size_ = max_index_size; + dirty_ = false; + + const std::string segment_id = + make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId); + chunk_meta_segment_ = stg_->get(segment_id); + if (!chunk_meta_segment_) { + LOG_DEBUG("Create new index"); + return init_storage(chunk_size); + } + + return load_storage(chunk_size); +} + +int HnswRabitqChunkBroker::close(void) { + flush(0UL); + + stg_.reset(); + check_crc_ = false; + dirty_ = false; + + return 0; +} + +int HnswRabitqChunkBroker::flush(uint64_t checkpoint) { + ailego_assert_with(chunk_meta_segment_, "invalid meta segment"); + + chunk_meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(chunk_meta_.update_time); + + size_t size = + chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswChunkMeta))) { + LOG_ERROR("Storage write data failed, wsize=%zu", size); + } + + stg_->refresh(checkpoint); + int ret = stg_->flush(); + if (ret == 0) { + (*stats_.mutable_check_point()) = checkpoint; + } else { + LOG_ERROR("Storage flush failed for %s", IndexError::What(ret)); + } + return ret; +} + +std::pair HnswRabitqChunkBroker::alloc_chunk( + int type, uint64_t seq_id, size_t size) { + ailego_assert_with(type < CHUNK_TYPE_MAX, "chunk type overflow"); + + Chunk::Pointer chunk; + if (ailego_unlikely(!stg_)) { + LOG_ERROR("Init storage first"); + return std::make_pair(IndexError_Uninitialized, chunk); + } + + //! check exist a empty chunk with the same name + chunk = get_chunk(type, seq_id); + if (chunk) { + if (ailego_unlikely(chunk->capacity() == size && + chunk->data_size() == 0UL)) { + LOG_ERROR("Exist invalid chunk size %zu, expect size %zu", + chunk->capacity(), size); + chunk.reset(); + return std::make_pair(IndexError_Runtime, chunk); + } + return std::make_pair(0, chunk); + } + //! align to page size + size = (size + page_mask_) & (~page_mask_); + if (ailego_unlikely(chunk_meta_.total_size + size >= max_chunks_size_)) { + LOG_ERROR("No space to new a chunk, curIndexSize=%zu allocSize=%zu", + (size_t)chunk_meta_.total_size, size); + return std::make_pair(IndexError_IndexFull, chunk); + } + + std::string segment_id = make_segment_id(type, seq_id); + int ret = stg_->append(segment_id, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Storage append segment failed for %s", IndexError::What(ret)); + return std::make_pair(ret, chunk); + } + chunk_meta_.chunk_cnts[type] += 1; + chunk_meta_.total_size += size; + (*stats_.mutable_index_size()) += size; + + size = chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswChunkMeta))) { + LOG_ERROR("Storage append segment failed, wsize=%zu", size); + } + + chunk = get_chunk(type, seq_id); + return std::make_pair(chunk ? 0 : IndexError_NoMemory, chunk); +} + +Chunk::Pointer HnswRabitqChunkBroker::get_chunk(int type, + uint64_t seq_id) const { + std::string segment_id = make_segment_id(type, seq_id); + return stg_->get(segment_id); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.h new file mode 100644 index 0000000..f0bc339 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.h @@ -0,0 +1,140 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_logger.h" +#include "zvec/core/framework/index_storage.h" +#include "zvec/core/framework/index_streamer.h" + +namespace zvec { +namespace core { + +using Chunk = IndexStorage::Segment; + +class HnswRabitqChunkBroker { + public: + typedef std::shared_ptr Pointer; + + enum CHUNK_TYPE { + CHUNK_TYPE_HEADER = 1, + CHUNK_TYPE_META = 2, + CHUNK_TYPE_NODE = 3, + CHUNK_TYPE_UPPER_NEIGHBOR = 4, + CHUNK_TYPE_NEIGHBOR_INDEX = 5, + CHUNK_TYPE_SPARSE_NODE = 6, + CHUNK_TYPE_MAX = 8 + }; + static constexpr size_t kDefaultChunkSeqId = 0UL; + + HnswRabitqChunkBroker(IndexStreamer::Stats &stats) : stats_(stats) {} + + //! Open storage + int open(IndexStorage::Pointer stg, size_t max_index_size, size_t chunk_size, + bool check_crc); + + int close(void); + + int flush(uint64_t checkpoint); + + //! alloc a new chunk with size, not thread-safe + std::pair alloc_chunk(int type, uint64_t seq_id, + size_t size); + + //! alloc a new chunk with chunk size + inline std::pair alloc_chunk(int type, uint64_t seq_id) { + return alloc_chunk(type, seq_id, chunk_meta_.chunk_size); + } + + Chunk::Pointer get_chunk(int type, uint64_t seq_id) const; + + inline size_t get_chunk_cnt(int type) const { + ailego_assert_with(type < CHUNK_TYPE_MAX, "chunk type overflow"); + return chunk_meta_.chunk_cnts[type]; + } + + inline bool dirty(void) const { + return dirty_; + } + + inline void mark_dirty(void) { + if (!dirty_) { + dirty_ = true; + chunk_meta_.revision_id += 1; + stats_.set_revision_id(chunk_meta_.revision_id); + } + } + + const IndexStorage::Pointer storage(void) const { + return stg_; + } + + private: + HnswRabitqChunkBroker(const HnswRabitqChunkBroker &) = delete; + HnswRabitqChunkBroker &operator=(const HnswRabitqChunkBroker &) = delete; + + struct HnswChunkMeta { + HnswChunkMeta(void) { + memset(static_cast(this), 0, sizeof(HnswChunkMeta)); + } + void clear() { + memset(static_cast(this), 0, sizeof(HnswChunkMeta)); + } + + uint64_t chunk_cnts[CHUNK_TYPE_MAX]; + uint64_t chunk_size; // size of per chunk + uint64_t total_size; // total size of allocated chunk + uint64_t revision_id; // index revision + uint64_t create_time; + uint64_t update_time; + uint64_t reserved[3]; + }; + + static_assert(sizeof(HnswChunkMeta) % 32 == 0, + "HnswChunkMeta must be aligned with 32 bytes"); + + //! Init the storage after open an empty index + int init_storage(size_t chunk_size); + + //! Load index from storage + int load_storage(size_t chunk_size); + + static inline const std::string make_segment_id(int type, uint64_t seq_id) { + return "HnswT" + ailego::StringHelper::ToString(type) + "S" + + ailego::StringHelper::ToString(seq_id); + } + + private: + IndexStreamer::Stats &stats_; + HnswChunkMeta chunk_meta_{}; + size_t page_mask_{0UL}; + size_t max_chunks_size_{0UL}; + IndexStorage::Pointer stg_{}; + IndexStorage::Segment::Pointer chunk_meta_segment_{}; + bool check_crc_{false}; + bool dirty_{false}; // set as true if index is modified , the flag + // will not be cleared even if flushed +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.cc new file mode 100644 index 0000000..e723177 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.cc @@ -0,0 +1,297 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_rabitq_context.h" +#include +#include "hnsw_rabitq_params.h" + +namespace zvec { +namespace core { + +HnswRabitqContext::HnswRabitqContext(size_t dimension, + const IndexMetric::Pointer &metric, + const HnswRabitqEntity::Pointer &entity) + : IndexContext(metric), + entity_(entity), + add_dc_(entity_.get(), metric, dimension) {} + +HnswRabitqContext::HnswRabitqContext(const IndexMetric::Pointer &metric, + const HnswRabitqEntity::Pointer &entity) + : IndexContext(metric), entity_(entity), add_dc_(entity_.get(), metric) {} + +HnswRabitqContext::~HnswRabitqContext() { + visit_filter_.destroy(); +} + +int HnswRabitqContext::init(ContextType type) { + int ret; + uint32_t doc_cnt; + + type_ = type; + + switch (type) { + case kBuilderContext: + ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(), + max_scan_num_, negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + update_heap_.limit(entity_->l0_neighbor_cnt() + 1); + break; + + case kSearcherContext: + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_, + negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + break; + + case kStreamerContext: + // maxScanNum is unknown if inited from streamer, so the docCnt may + // change. we need to compute maxScanNum by scan ratio, and preserve + // max_doc_cnt space from visit filter + doc_cnt = entity_->doc_cnt(); + max_scan_num_ = compute_max_scan_num(doc_cnt); + reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt); + ret = visit_filter_.init(filter_mode_, reserve_max_doc_cnt_, + max_scan_num_, negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + + update_heap_.limit(entity_->l0_neighbor_cnt() + 1); + candidates_.limit(max_scan_num_); + + check_need_adjuct_ctx(); + break; + + default: + LOG_ERROR("Init context failed"); + return IndexError_Runtime; + } + + return 0; +} + +int HnswRabitqContext::update(const ailego::Params ¶ms) { + auto update_visit_filter_param = [&]() { + bool need_update = false; + std::string p; + switch (type_) { + case kSearcherContext: + p = PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_ENABLE; + break; + case kStreamerContext: + p = PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_ENABLE; + break; + } + + if (params.has(p)) { + bool bf_enabled = false; + params.get(p, &bf_enabled); + if (bf_enabled ^ (filter_mode_ == VisitFilter::BloomFilter)) { + need_update = true; + filter_mode_ = + bf_enabled ? VisitFilter::BloomFilter : VisitFilter::ByteMap; + } + } + + float prob = negative_probability_; + p.clear(); + switch (type_) { + case kSearcherContext: + p = PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB; + break; + case kStreamerContext: + p = PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB; + break; + } + params.get(p, &prob); + if (filter_mode_ == VisitFilter::BloomFilter && + std::abs(prob - negative_probability_) > 1e-6) { + need_update = true; + } + if (need_update) { + visit_filter_.destroy(); + int max_doc_cnt = 0; + if (type_ == kSearcherContext) { + max_doc_cnt = entity_->doc_cnt(); + } else { + max_doc_cnt = reserve_max_doc_cnt_; + } + int ret = visit_filter_.init(filter_mode_, max_doc_cnt, max_scan_num_, + negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + } + return 0; + }; + + switch (type_) { + case kSearcherContext: + if (params.has(PARAM_HNSW_RABITQ_SEARCHER_EF)) { + params.get(PARAM_HNSW_RABITQ_SEARCHER_EF, &ef_); + topk_heap_.limit(std::max(topk_, ef_)); + } + + if (params.has(PARAM_HNSW_RABITQ_SEARCHER_MAX_SCAN_RATIO)) { + params.get(PARAM_HNSW_RABITQ_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_); + max_scan_num_ = + static_cast(max_scan_ratio_ * entity_->doc_cnt()); + max_scan_num_ = std::max(10000U, max_scan_num_); + } + + if (params.has(PARAM_HNSW_RABITQ_SEARCHER_BRUTE_FORCE_THRESHOLD)) { + params.get(PARAM_HNSW_RABITQ_SEARCHER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + } + + return update_visit_filter_param(); + + case kStreamerContext: + if (params.has(PARAM_HNSW_RABITQ_STREAMER_EF)) { + params.get(PARAM_HNSW_RABITQ_STREAMER_EF, &ef_); + topk_heap_.limit(std::max(topk_, ef_)); + } + params.get(PARAM_HNSW_RABITQ_STREAMER_EF, &ef_); + params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_); + params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_); + params.get(PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_); + if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) { + LOG_ERROR("[%s] must be in range (0.0f,1.0f]", + PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO.c_str()); + return IndexError_InvalidArgument; + } + if (max_scan_limit_ < min_scan_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT.c_str(), + PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT.c_str()); + return IndexError_InvalidArgument; + } + + if (params.has(PARAM_HNSW_RABITQ_STREAMER_BRUTE_FORCE_THRESHOLD)) { + params.get(PARAM_HNSW_RABITQ_STREAMER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + } + + return update_visit_filter_param(); + + default: + LOG_ERROR("update context failed, type=%zu", static_cast(type_)); + return IndexError_Runtime; + } +} + +int HnswRabitqContext::update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const HnswRabitqEntity::Pointer &entity, + uint32_t magic_num) { + uint32_t doc_cnt; + + if (ailego_unlikely(static_cast(type) != type_)) { + LOG_ERROR( + "HnswRabitqContext doesn't support shared by different type, " + "src=%u dst=%u", + type_, type); + return IndexError_Unsupported; + } + + magic_ = kInvalidMgic; + + // TODO: support change filter mode? + switch (type) { + case kBuilderContext: + LOG_ERROR("BuildContext doesn't support update"); + return IndexError_NotImplemented; + + case kSearcherContext: + if (!visit_filter_.reset(entity->doc_cnt(), max_scan_num_)) { + LOG_ERROR("Reset filter failed, mode %d", visit_filter_.get_mode()); + return IndexError_Runtime; + } + + candidates_.limit(max_scan_num_); + topk_heap_.limit(std::max(topk_, ef_)); + break; + + case kStreamerContext: + doc_cnt = entity->doc_cnt(); + max_scan_num_ = compute_max_scan_num(doc_cnt); + reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt); + if (!visit_filter_.reset(reserve_max_doc_cnt_, max_scan_num_)) { + LOG_ERROR("Reset filter failed, mode %d", visit_filter_.get_mode()); + return IndexError_Runtime; + } + + update_heap_.limit(entity->l0_neighbor_cnt() + 1); + candidates_.limit(max_scan_num_); + topk_heap_.limit(std::max(topk_, ef_)); + break; + + default: + LOG_ERROR("update context failed"); + return IndexError_Runtime; + } + + entity_ = entity; + dc().update(entity_.get(), metric, meta.dimension()); + magic_ = magic_num; + level_topks_.clear(); + + return 0; +} + +void HnswRabitqContext::fill_random_to_topk_full(void) { + static std::mt19937 mt( + std::chrono::system_clock::now().time_since_epoch().count()); + std::uniform_int_distribution dt(0, entity_->doc_cnt() - 1); + std::function gen; + node_id_t seqid; + std::function myfilter = [](node_id_t) { return false; }; + if (this->filter().is_valid()) { + myfilter = [&](node_id_t id) { + return this->filter()(entity_->get_key(id)); + }; + } + + if (topk_heap_.limit() < entity_->doc_cnt() / 2) { + gen = [&](void) { return dt(mt); }; + } else { + // If topk limit is big value, gen sequential id from an random initial + seqid = dt(mt); + gen = [&](void) { + seqid = seqid == (entity_->doc_cnt() - 1) ? 0 : (seqid + 1); + return seqid; + }; + } + + for (size_t i = 0; !topk_heap_.full() && i < entity_->doc_cnt(); ++i) { + const auto id = gen(); + if (!visit_filter_.visited(id) && !myfilter(id)) { + visit_filter_.set_visited(id); + topk_heap_.emplace(id, dc().dist(id)); + } + } + return; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h new file mode 100644 index 0000000..ea2c86b --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h @@ -0,0 +1,548 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "utility/visit_filter.h" +#include "zvec/core/framework/index_context.h" +#include "zvec/core/framework/index_provider.h" +#include "hnsw_rabitq_dist_calculator.h" +#include "hnsw_rabitq_entity.h" + +namespace zvec { +namespace core { + +class HnswRabitqContext : public IndexContext { + public: + //! Index Context Pointer + typedef std::unique_ptr Pointer; + + enum ContextType { + kUnknownContext = 0, + kSearcherContext = 1, + kBuilderContext = 2, + kStreamerContext = 3 + }; + + //! Construct + HnswRabitqContext(size_t dimension, const IndexMetric::Pointer &metric, + const HnswRabitqEntity::Pointer &entity); + + //! Construct + HnswRabitqContext(const IndexMetric::Pointer &metric, + const HnswRabitqEntity::Pointer &entity); + + //! Destructor + ~HnswRabitqContext() override; + + public: + //! Set topk of search result + void set_topk(uint32_t val) override { + topk_ = val; + topk_heap_.limit(std::max(val, ef_)); + } + + //! Retrieve search result + const IndexDocumentList &result(void) const override { + return results_[0]; + } + + //! Retrieve search result + const IndexDocumentList &result(size_t idx) const override { + return results_[idx]; + } + + //! Retrieve result object for output + IndexDocumentList *mutable_result(size_t idx) override { + ailego_assert_with(idx < results_.size(), "invalid idx"); + return &results_[idx]; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(void) const override { + return group_results_[0]; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(size_t idx) const override { + return group_results_[idx]; + } + + IndexGroupDocumentList *mutable_group_result(void) override { + return &group_results_[0]; + } + + IndexGroupDocumentList *mutable_group_result(size_t idx) override { + return &group_results_[idx]; + } + + uint32_t magic(void) const override { + return magic_; + } + + //! Set mode of debug + void set_debug_mode(bool enable) override { + debug_mode_ = enable; + } + + //! Retrieve mode of debug + bool debug_mode(void) const override { + return this->debugging(); + } + + //! Retrieve string of debug + std::string debug_string(void) const override { + char buf[4096]; + size_t size = snprintf( + buf, sizeof(buf), + "scan_cnt=%zu,get_vector_cnt=%u,get_neighbors_cnt=%u,dup_node=%u", + get_scan_num(), stats_get_vector_cnt_, stats_get_neighbors_cnt_, + stats_visit_dup_cnt_); + return std::string(buf, size); + } + + //! Update the parameters of context + int update(const ailego::Params ¶ms) override; + + public: + //! Init context + int init(ContextType type); + + //! Update context, the context may be shared by different searcher/streamer + int update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const HnswRabitqEntity::Pointer &entity, + uint32_t magic_num); + + inline const HnswRabitqEntity &get_entity() const { + return *entity_; + } + + inline void resize_results(size_t size) { + if (group_by_search()) { + group_results_.resize(size); + } else { + results_.resize(size); + } + } + + inline void topk_to_result() { + return topk_to_result(0); + } + + //! Construct result from topk heap, result will be normalized + inline void topk_to_result(uint32_t idx) { + if (group_by_search()) { + topk_to_group_result(idx); + } else { + topk_to_single_result(idx); + } + } + + inline void topk_to_single_result(uint32_t idx) { + if (force_padding_topk_ && !topk_heap_.full() && + topk_heap_.size() < entity_->doc_cnt()) { + this->fill_random_to_topk_full(); + } + if (ailego_unlikely(topk_heap_.size() == 0)) { + return; + } + + ailego_assert_with(idx < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(topk_heap_.size())); + topk_heap_.sort(); + results_[idx].clear(); + + for (int i = 0; i < size; ++i) { + auto score = topk_heap_[i].second; + if (score.est_dist > this->threshold()) { + break; + } + + node_id_t id = topk_heap_[i].first; + if (fetch_vector_) { + results_[idx].emplace_back(entity_->get_key(id), score.est_dist, id, + entity_->get_vector(id)); + } else { + results_[idx].emplace_back(entity_->get_key(id), score.est_dist, id); + } + } + + return; + } + + //! Construct result from topk heap, result will be normalized + inline void topk_to_group_result(uint32_t idx) { + ailego_assert_with(idx < group_results_.size(), "invalid idx"); + + group_results_[idx].clear(); + + std::vector> group_topk_list; + std::vector> best_score_in_groups; + for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end(); + itr++) { + const std::string &group_id = (*itr).first; + auto &heap = (*itr).second; + heap.sort(); + + if (heap.size() > 0) { + ResultRecord best_score = heap[0].second; + best_score_in_groups.push_back(std::make_pair(group_id, best_score)); + } + } + + std::sort(best_score_in_groups.begin(), best_score_in_groups.end(), + [](const std::pair &a, + const std::pair &b) -> int { + return a.second < b.second; + }); + + // truncate to group num + for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size(); + ++i) { + const std::string &group_id = best_score_in_groups[i].first; + + group_topk_list.emplace_back( + std::make_pair(group_id, group_topk_heaps_[group_id])); + } + + group_results_[idx].resize(group_topk_list.size()); + + for (uint32_t i = 0; i < group_topk_list.size(); ++i) { + const std::string &group_id = group_topk_list[i].first; + group_results_[idx][i].set_group_id(group_id); + + uint32_t size = std::min( + group_topk_, static_cast(group_topk_list[i].second.size())); + + for (uint32_t j = 0; j < size; ++j) { + auto score = group_topk_list[i].second[j].second; + if (score > this->threshold()) { + break; + } + + node_id_t id = group_topk_list[i].second[j].first; + + if (fetch_vector_) { + group_results_[idx][i].mutable_docs()->emplace_back( + entity_->get_key(id), score.est_dist, id, + entity_->get_vector(id)); + } else { + group_results_[idx][i].mutable_docs()->emplace_back( + entity_->get_key(id), score.est_dist, id); + } + } + } + } + + inline void reset_query(const void *query) { + if (auto query_preprocess_func = index_metric_->get_query_preprocess_func(); + query_preprocess_func != nullptr) { + size_t dim = dc().dimension(); + preprocess_buffer_.resize(dim); + memcpy(preprocess_buffer_.data(), query, dim); + query_preprocess_func(preprocess_buffer_.data(), dim); + query = preprocess_buffer_.data(); + } + + dc().reset_query(query); + dc().clear_compare_cnt(); + query_ = query; + } + + inline HnswRabitqAddDistCalculator &dist_calculator() { + return dc(); + } + + inline TopkHeap &topk_heap() { + return topk_heap_; + } + + inline TopkHeap &update_heap() { + return update_heap_; + } + + inline VisitFilter &visit_filter() { + return visit_filter_; + } + + inline CandidateHeap &candidates() { + return candidates_; + } + + inline void set_max_scan_num(uint32_t max_scan_num) { + max_scan_num_ = max_scan_num; + } + + inline void set_max_scan_limit(uint32_t max_scan_limit) { + max_scan_limit_ = max_scan_limit; + } + + inline void set_min_scan_limit(uint32_t min_scan_limit) { + min_scan_limit_ = min_scan_limit; + } + + inline void set_ef(uint32_t v) { + ef_ = v; + } + + inline void set_filter_mode(uint32_t v) { + filter_mode_ = v; + } + + inline void set_filter_negative_probability(float v) { + negative_probability_ = v; + } + + inline void set_max_scan_ratio(float v) { + max_scan_ratio_ = v; + } + + virtual void set_magic(uint32_t v) { + magic_ = v; + } + + virtual void set_force_padding_topk(bool v) { + force_padding_topk_ = v; + } + + void set_bruteforce_threshold(uint32_t v) override { + bruteforce_threshold_ = v; + } + + inline uint32_t get_bruteforce_threshold() const { + return bruteforce_threshold_; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + + bool fetch_vector() const override { + return fetch_vector_; + } + + //! Reset context + void reset(void) override { + set_filter(nullptr); + reset_threshold(); + set_fetch_vector(false); + set_group_params(0, 0); + reset_group_by(); + } + + inline std::map &group_topk_heaps() { + return group_topk_heaps_; + } + + inline TopkHeap &level_topk(int level) { + if (ailego_unlikely(level_topks_.size() <= static_cast(level))) { + int cur_level = level_topks_.size(); + level_topks_.resize(level + 1); + for (; cur_level <= level; ++cur_level) { + size_t heap_size = std::max(entity_->neighbor_cnt(cur_level), + entity_->ef_construction()); + level_topks_[cur_level].clear(); + level_topks_[cur_level].limit(heap_size); + } + } + + return level_topks_[level]; + } + + inline void check_need_adjuct_ctx(void) { + check_need_adjuct_ctx(entity_->doc_cnt()); + } + + inline size_t compute_reserve_cnt(uint32_t cur_doc) const { + if (cur_doc > kMaxReserveDocCnt) { + return kMaxReserveDocCnt; + } else if (cur_doc < kMinReserveDocCnt) { + return kMinReserveDocCnt; + } + return cur_doc; + } + + //! candidates heap and visitfilter need to resize as doc cnt growing up + inline void check_need_adjuct_ctx(uint32_t doc_cnt) { + if (ailego_unlikely(doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_)) { + while (doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_) { + reserve_max_doc_cnt_ = + reserve_max_doc_cnt_ + compute_reserve_cnt(reserve_max_doc_cnt_); + } + uint32_t max_scan_cnt = compute_max_scan_num(reserve_max_doc_cnt_); + max_scan_num_ = max_scan_cnt; + visit_filter_.reset(reserve_max_doc_cnt_, max_scan_cnt); + candidates_.clear(); + candidates_.limit(max_scan_num_); + } + } + + inline uint32_t compute_max_scan_num(uint32_t max_doc_cnt) const { + uint32_t max_scan = max_doc_cnt * max_scan_ratio_; + if (max_scan < min_scan_limit_) { + max_scan = min_scan_limit_; + } else if (max_scan > max_scan_limit_) { + max_scan = max_scan_limit_; + } + return max_scan; + } + + inline size_t get_scan_num() const { + return dc().compare_cnt(); + } + + inline uint64_t reach_scan_limit() const { + return dc().compare_cnt() >= max_scan_num_; + } + + inline bool error() const { + return dc().error(); + } + + inline void clear() { + add_dc_.clear(); + if (ailego_unlikely(this->debugging())) { + stats_get_neighbors_cnt_ = 0u; + stats_get_vector_cnt_ = 0u; + stats_visit_dup_cnt_ = 0u; + } + // do not clear results_ for the next query will need it + for (auto &it : results_) { + it.clear(); + } + } + + uint32_t *mutable_stats_get_neighbors() { + return &stats_get_neighbors_cnt_; + } + + uint32_t *mutable_stats_get_vector() { + return &stats_get_vector_cnt_; + } + + uint32_t *mutable_stats_visit_dup_cnt() { + return &stats_visit_dup_cnt_; + } + + inline bool debugging(void) const { + return debug_mode_; + } + + inline void update_dist_caculator_distance( + const IndexMetric::MatrixDistance &distance, + const IndexMetric::MatrixBatchDistance &batch_distance) { + dc().update_distance(distance, batch_distance); + } + + //! Get topk + inline uint32_t topk() const override { + return topk_; + } + + //! Get group topk + inline uint32_t group_topk() const { + return group_topk_; + } + + //! Get group num + inline uint32_t group_num() const { + return group_num_; + } + + //! Get if group by search + inline bool group_by_search() { + return group_num_ > 0; + } + + //! Set group params + void set_group_params(uint32_t group_num, uint32_t group_topk) override { + group_num_ = group_num; + group_topk_ = group_topk; + + topk_ = group_topk_ * group_num_; + + topk_heap_.limit(std::max(topk_, ef_)); + + group_topk_heaps_.clear(); + } + + void set_provider(IndexProvider::Pointer provider) { + add_dc_.set_provider(std::move(provider)); + } + + const void *query() const { + return query_; + } + + private: + inline HnswRabitqAddDistCalculator &dc() { + return add_dc_; + } + + inline const HnswRabitqAddDistCalculator &dc() const { + return add_dc_; + } + + private: + // Filling random nodes if topk not full + void fill_random_to_topk_full(void); + + constexpr static uint32_t kTriggerReserveCnt = 4096UL; + constexpr static uint32_t kMinReserveDocCnt = 4096UL; + constexpr static uint32_t kMaxReserveDocCnt = 128 * 1024UL; + constexpr static uint32_t kInvalidMgic = -1U; + + private: + HnswRabitqEntity::Pointer entity_; + HnswRabitqAddDistCalculator add_dc_; + IndexMetric::Pointer metric_; + + bool debug_mode_{false}; + bool force_padding_topk_{false}; + uint32_t max_scan_num_{0}; + uint32_t max_scan_limit_{0}; + uint32_t min_scan_limit_{0}; + uint32_t reserve_max_doc_cnt_{kMinReserveDocCnt}; + uint32_t topk_{0}; + uint32_t group_topk_{0}; + uint32_t filter_mode_{VisitFilter::ByteMap}; + float negative_probability_{HnswRabitqEntity::kDefaultBFNegativeProbability}; + uint32_t ef_{HnswRabitqEntity::kDefaultEf}; + float max_scan_ratio_{HnswRabitqEntity::kDefaultScanRatio}; + uint32_t magic_{0U}; + std::vector results_{}; + std::vector group_results_{}; + TopkHeap topk_heap_{}; + TopkHeap update_heap_{}; + std::vector level_topks_{}; + CandidateHeap candidates_{}; + VisitFilter visit_filter_{}; + uint32_t bruteforce_threshold_{}; + bool fetch_vector_{false}; + + uint32_t group_num_{0}; + std::map group_topk_heaps_{}; + + uint32_t type_{kUnknownContext}; + //! debug stats info + uint32_t stats_get_neighbors_cnt_{0u}; + uint32_t stats_get_vector_cnt_{0u}; + uint32_t stats_visit_dup_cnt_{0u}; + std::string preprocess_buffer_; + const void *query_{nullptr}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.cc new file mode 100644 index 0000000..43ed854 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.cc @@ -0,0 +1,39 @@ +// Copyright 2025-present the centaurdb project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License + +#include "core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.h" +#include "zvec/core/framework/index_error.h" + +namespace zvec::core { + +int HnswRabitqAddDistCalculator::get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + for (uint32_t i = 0; i < count; ++i) { + const node_id_t id = ids[i]; + key_t key = entity_->get_key(id); + if (key == kInvalidKey) { + return IndexError_NoExist; + } + IndexStorage::MemoryBlock block; + int ret = provider_->get_vector(key, block); + if (ret != 0) { + return ret; + } + vec_blocks.push_back(std::move(block)); + } + return 0; +} + +} // namespace zvec::core diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.h new file mode 100644 index 0000000..faef466 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.h @@ -0,0 +1,239 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "zvec/core/framework/index_meta.h" +#include "zvec/core/framework/index_metric.h" +#include "zvec/core/framework/index_provider.h" +#include "hnsw_rabitq_entity.h" + +namespace zvec { +namespace core { + +//! HnswRabitqAddDistCalculator is only used for index construction +class HnswRabitqAddDistCalculator { + public: + typedef std::shared_ptr Pointer; + + public: + enum DistType { + DIST_NONE = 0, + DIST_DENSE = 1, + DIST_HYBRID = 2, + DIST_SPARSE = 3 + }; + + public: + //! Constructor + HnswRabitqAddDistCalculator(const HnswRabitqEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(nullptr), + dim_(dim), + compare_cnt_(0) {} + + //! Constructor + HnswRabitqAddDistCalculator(const HnswRabitqEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim, + const void *query) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(query), + dim_(dim), + compare_cnt_(0) {} + + //! Constructor + HnswRabitqAddDistCalculator(const HnswRabitqEntity *entity, + const IndexMetric::Pointer &metric) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(nullptr), + dim_(0), + compare_cnt_(0) {} + + void update(const HnswRabitqEntity *entity, + const IndexMetric::Pointer &metric) { + entity_ = entity; + distance_ = metric->distance(); + batch_distance_ = metric->batch_distance(); + } + + void update(const HnswRabitqEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim) { + entity_ = entity; + distance_ = metric->distance(); + batch_distance_ = metric->batch_distance(); + dim_ = dim; + } + + inline void update_distance( + const IndexMetric::MatrixDistance &distance, + const IndexMetric::MatrixBatchDistance &batch_distance) { + distance_ = distance; + batch_distance_ = batch_distance; + } + + //! Reset query vector data + inline void reset_query(const void *query) { + error_ = false; + query_ = query; + } + + //! Returns distance + inline dist_t dist(const void *vec_lhs, const void *vec_rhs) { + if (ailego_unlikely(vec_lhs == nullptr || vec_rhs == nullptr)) { + LOG_ERROR("Nullptr of dense vector"); + error_ = true; + return 0.0f; + } + + float score{0.0f}; + + distance_(vec_lhs, vec_rhs, dim_, &score); + + return score; + } + + //! Returns distance between query and vec. + inline dist_t dist(const void *vec) { + compare_cnt_++; + + return dist(vec, query_); + } + + //! Return distance between query and node id. + inline dist_t dist(node_id_t id) { + compare_cnt_++; + + const void *feat = get_vector(id); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + + return dist(feat, query_); + } + + //! Return dist node lhs between node rhs + inline dist_t dist(node_id_t lhs, node_id_t rhs) { + compare_cnt_++; + + const void *feat = get_vector(lhs); + const void *query = get_vector(rhs); + if (ailego_unlikely(feat == nullptr || query == nullptr)) { + LOG_ERROR("Get nullptr vector"); + error_ = true; + return 0.0f; + } + + return dist(feat, query); + } + + dist_t operator()(const void *vec) { + return dist(vec); + } + + dist_t operator()(id_t i) { + return dist(i); + } + + dist_t operator()(id_t lhs, id_t rhs) { + return dist(lhs, rhs); + } + + void batch_dist(const void **vecs, size_t num, dist_t *distances) { + compare_cnt_++; + + batch_distance_(vecs, query_, num, dim_, distances); + } + + inline dist_t batch_dist(node_id_t id) { + compare_cnt_++; + + const void *feat = get_vector(id); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + dist_t score = 0; + batch_distance_(&feat, query_, 1, dim_, &score); + + return score; + } + + inline void clear() { + compare_cnt_ = 0; + error_ = false; + } + + inline void clear_compare_cnt() { + compare_cnt_ = 0; + } + + inline bool error() const { + return error_; + } + + //! Get distances compute times + inline uint32_t compare_cnt() const { + return compare_cnt_; + } + + inline uint32_t dimension() const { + return dim_; + } + + void set_provider(IndexProvider::Pointer provider) { + provider_ = std::move(provider); + } + + int get_vector(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const; + + const void *get_vector(node_id_t id) const { + key_t key = entity_->get_key(id); + if (key == kInvalidKey) { + return nullptr; + } + return provider_->get_vector(key); + } + + private: + HnswRabitqAddDistCalculator(const HnswRabitqAddDistCalculator &) = delete; + HnswRabitqAddDistCalculator &operator=(const HnswRabitqAddDistCalculator &) = + delete; + + private: + const HnswRabitqEntity *entity_; + IndexMetric::MatrixDistance distance_; + IndexMetric::MatrixBatchDistance batch_distance_; + + const void *query_; + uint32_t dim_; + + uint32_t compare_cnt_; // record distance compute times + bool error_{false}; + + // get raw vector + IndexProvider::Pointer provider_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.cc new file mode 100644 index 0000000..d3c6a49 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.cc @@ -0,0 +1,382 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_rabitq_entity.h" +#include +#include "utility/sparse_utility.h" +#include "zvec/core/framework/index_stats.h" + +namespace zvec { +namespace core { + +const std::string HnswRabitqEntity::kGraphHeaderSegmentId = "graph.header"; +const std::string HnswRabitqEntity::kGraphFeaturesSegmentId = "graph.features"; +const std::string HnswRabitqEntity::kGraphKeysSegmentId = "graph.keys"; +const std::string HnswRabitqEntity::kGraphNeighborsSegmentId = + "graph.neighbors"; +const std::string HnswRabitqEntity::kGraphOffsetsSegmentId = "graph.offsets"; +const std::string HnswRabitqEntity::kGraphMappingSegmentId = "graph.mapping"; +const std::string HnswRabitqEntity::kHnswHeaderSegmentId = "hnsw.header"; +const std::string HnswRabitqEntity::kHnswNeighborsSegmentId = "hnsw.neighbors"; +const std::string HnswRabitqEntity::kHnswOffsetsSegmentId = "hnsw.offsets"; + +void HnswRabitqEntity::update_rabitq_params_and_vector_size( + uint32_t dimension) { + uint32_t padded_dim = ((dimension + 63) / 64) * 64; + header_.graph.padded_dim = padded_dim; + // BinDataMap layout: bin_code (padded_dim/8) + f_add + f_rescale + f_error + header_.graph.size_bin_data = + rabitqlib::BinDataMap::data_bytes(padded_dim); + // ExDataMap layout: ex_code (padded_dim*ex_bits/8) + f_add_ex + f_rescale_ex + header_.graph.size_ex_data = rabitqlib::ExDataMap::data_bytes( + padded_dim, header_.graph.ex_bits); + // quantized vector format: cluster_id + bin_data + ex_data + header_.graph.vector_size = + sizeof(uint32_t) + size_bin_data() + size_ex_data(); +} + +int HnswRabitqEntity::CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, + size_t *padding_size) { + *padding_size = AlignSize(data_size) - data_size; + if (*padding_size == 0) { + return 0; + } + + std::string padding(*padding_size, '\0'); + if (dumper->write(padding.data(), *padding_size) != *padding_size) { + LOG_ERROR("Append padding failed, size %zu", *padding_size); + return IndexError_WriteData; + } + return 0; +} + +int64_t HnswRabitqEntity::dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, + const void *data, size_t size) const { + size_t len = dumper->write(data, size); + if (len != size) { + LOG_ERROR("Dump segment %s data failed, expect: %zu, actual: %zu", + segment_id.c_str(), size, len); + return IndexError_WriteData; + } + + size_t padding_size = AlignSize(size) - size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %zu", padding_size); + return IndexError_WriteData; + } + } + + uint32_t crc = ailego::Crc32c::Hash(data, size); + int ret = dumper->append(segment_id, size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", segment_id.c_str(), ret); + return ret; + } + + return len + padding_size; +} + +int64_t HnswRabitqEntity::dump_header(const IndexDumper::Pointer &dumper, + const HNSWHeader &hd) const { + //! dump basic graph header. header is aligned and does not need padding + int64_t graph_hd_size = + dump_segment(dumper, kGraphHeaderSegmentId, &hd.graph, hd.graph.size); + if (graph_hd_size < 0) { + return graph_hd_size; + } + + //! dump basic graph header. header is aligned and does not need padding + int64_t hnsw_hd_size = + dump_segment(dumper, kHnswHeaderSegmentId, &hd.hnsw, hd.hnsw.size); + if (hnsw_hd_size < 0) { + return hnsw_hd_size; + } + + return graph_hd_size + hnsw_hd_size; +} + +void HnswRabitqEntity::reshuffle_vectors( + const std::function & /*get_level*/, + std::vector * /*n2o_mapping*/, + std::vector * /*o2n_mapping*/, key_t * /*keys*/) const { + // TODO + return; +} + +int64_t HnswRabitqEntity::dump_mapping_segment( + const IndexDumper::Pointer &dumper, const key_t *keys) const { + std::vector mapping(doc_cnt()); + + std::iota(mapping.begin(), mapping.end(), 0U); + std::sort(mapping.begin(), mapping.end(), + [&](node_id_t i, node_id_t j) { return keys[i] < keys[j]; }); + + size_t size = mapping.size() * sizeof(node_id_t); + + return dump_segment(dumper, kGraphMappingSegmentId, mapping.data(), size); +} + +int64_t HnswRabitqEntity::dump_segments( + const IndexDumper::Pointer &dumper, key_t *keys, + const std::function &get_level) const { + HNSWHeader dump_hd(header()); + + dump_hd.graph.node_size = AlignSize(vector_size()); + + std::vector n2o_mapping; // map new id to origin id + std::vector o2n_mapping; // map origin id to new id + reshuffle_vectors(get_level, &n2o_mapping, &o2n_mapping, keys); + if (!o2n_mapping.empty()) { + dump_hd.hnsw.entry_point = o2n_mapping[entry_point()]; + } + + //! Dump header + int64_t hd_size = dump_header(dumper, dump_hd); + if (hd_size < 0) { + return hd_size; + } + + //! Dump vectors + int64_t vecs_size = dump_vectors(dumper, n2o_mapping); + if (vecs_size < 0) { + return vecs_size; + } + + //! Dump neighbors + auto neighbors_size = + dump_neighbors(dumper, get_level, n2o_mapping, o2n_mapping); + if (neighbors_size < 0) { + return neighbors_size; + } + //! free memory + n2o_mapping = std::vector(); + o2n_mapping = std::vector(); + + //! Dump keys + size_t key_segment_size = doc_cnt() * sizeof(key_t); + int64_t keys_size = + dump_segment(dumper, kGraphKeysSegmentId, keys, key_segment_size); + if (keys_size < 0) { + return keys_size; + } + + //! Dump mapping + int64_t mapping_size = dump_mapping_segment(dumper, keys); + if (mapping_size < 0) { + return mapping_size; + } + + return hd_size + keys_size + vecs_size + neighbors_size + mapping_size; +} + +int64_t HnswRabitqEntity::dump_vectors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + size_t vector_dump_size = vector_size(); + + size_t padding_size = AlignSize(vector_dump_size) - vector_dump_size; + + std::vector padding(padding_size, 0); + const void *data = nullptr; + uint32_t crc = 0U; + size_t vecs_size = 0UL; + + //! dump vectors + for (node_id_t id = 0; id < doc_cnt(); ++id) { + data = get_vector(reorder_mapping.empty() ? id : reorder_mapping[id]); + if (ailego_unlikely(!data)) { + return IndexError_ReadData; + } + size_t len = dumper->write(data, vector_size()); + if (len != vector_size()) { + LOG_ERROR("Dump vectors failed, write=%zu expect=%zu", len, + vector_size()); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(data, vector_size(), crc); + vecs_size += vector_size(); + + if (padding_size == 0) { + continue; + } + + len = dumper->write(padding.data(), padding_size); + if (len != padding_size) { + LOG_ERROR("Dump vectors failed, write=%zu expect=%zu", len, padding_size); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(padding.data(), padding_size, crc); + vecs_size += padding_size; + } + + int ret = dumper->append(kGraphFeaturesSegmentId, vecs_size, 0UL, crc); + if (ret != 0) { + LOG_ERROR("Dump vectors segment meta failed, ret %d", ret); + return ret; + } + + return vecs_size; +} + +int64_t HnswRabitqEntity::dump_graph_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + std::vector graph_meta; + graph_meta.reserve(doc_cnt()); + size_t offset = 0; + uint32_t crc = 0; + std::vector mapping(l0_neighbor_cnt()); + + uint32_t min_neighbor_count = 10000; + uint32_t max_neighbor_count = 0; + size_t sum_neighbor_count = 0; + + for (node_id_t id = 0; id < doc_cnt(); ++id) { + const Neighbors neighbors = + get_neighbors(0, reorder_mapping.empty() ? id : reorder_mapping[id]); + ailego_assert_with(!!neighbors.data, "invalid neighbors"); + ailego_assert_with(neighbors.size() <= l0_neighbor_cnt(), + "invalid neighbors"); + + uint32_t neighbor_count = neighbors.size(); + if (neighbor_count < min_neighbor_count) { + min_neighbor_count = neighbor_count; + } + if (neighbor_count > max_neighbor_count) { + max_neighbor_count = neighbor_count; + } + sum_neighbor_count += neighbor_count; + + graph_meta.emplace_back(offset, neighbor_count); + size_t size = neighbors.size() * sizeof(node_id_t); + const node_id_t *data = &neighbors[0]; + if (!neighbor_mapping.empty()) { + for (node_id_t i = 0; i < neighbors.size(); ++i) { + mapping[i] = neighbor_mapping[neighbors[i]]; + } + data = mapping.data(); + } + if (dumper->write(data, size) != size) { + LOG_ERROR("Dump graph neighbor id=%zu failed, size %zu", + static_cast(id), size); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(data, size, crc); + offset += size; + } + + uint32_t average_neighbor_count = 0; + if (doc_cnt() > 0) { + average_neighbor_count = sum_neighbor_count / doc_cnt(); + } + LOG_INFO( + "Dump hnsw graph: min_neighbor_count[%u] max_neighbor_count[%u] " + "average_neighbor_count[%u]", + min_neighbor_count, max_neighbor_count, average_neighbor_count); + + size_t padding_size = 0; + int ret = CalcAndAddPadding(dumper, offset, &padding_size); + if (ret != 0) { + return ret; + } + ret = dumper->append(kGraphNeighborsSegmentId, offset, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s failed, ret %d", + kGraphNeighborsSegmentId.c_str(), ret); + return ret; + } + + //! dump level 0 neighbors meta + auto len = dump_segment(dumper, kGraphOffsetsSegmentId, graph_meta.data(), + graph_meta.size() * sizeof(GraphNeighborMeta)); + if (len < 0) { + return len; + } + + return len + offset + padding_size; +} + +int64_t HnswRabitqEntity::dump_upper_neighbors( + const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + std::vector hnsw_meta; + hnsw_meta.reserve(doc_cnt()); + size_t offset = 0; + uint32_t crc = 0; + std::vector buffer(upper_neighbor_cnt() + 1); + for (node_id_t id = 0; id < doc_cnt(); ++id) { + node_id_t new_id = reorder_mapping.empty() ? id : reorder_mapping[id]; + auto level = get_level(new_id); + if (level == 0) { + hnsw_meta.emplace_back(0U, 0U); + continue; + } + hnsw_meta.emplace_back(offset, level); + ailego_assert_with((size_t)level < kMaxGraphLayers, "invalid level"); + for (level_t cur_level = 1; cur_level <= level; ++cur_level) { + const Neighbors neighbors = get_neighbors(cur_level, new_id); + ailego_assert_with(!!neighbors.data, "invalid neighbors"); + ailego_assert_with(neighbors.size() <= neighbor_cnt(cur_level), + "invalid neighbors"); + size_t buffer_bytes = buffer.size() * sizeof(node_id_t); + memset(buffer.data(), 0, buffer_bytes); + buffer[0] = neighbors.size(); + if (neighbor_mapping.empty()) { + memcpy(&buffer[1], &neighbors[0], neighbors.size() * sizeof(node_id_t)); + } else { + for (node_id_t i = 0; i < neighbors.size(); ++i) { + buffer[i + 1] = neighbor_mapping[neighbors[i]]; + } + } + if (dumper->write(buffer.data(), buffer_bytes) != buffer_bytes) { + LOG_ERROR("Dump graph neighbor id=%zu failed, size %zu", + static_cast(id), buffer_bytes); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(buffer.data(), buffer_bytes, crc); + offset += buffer_bytes; + } + } + size_t padding_size = 0; + int ret = CalcAndAddPadding(dumper, offset, &padding_size); + if (ret != 0) { + return ret; + } + + ret = dumper->append(kHnswNeighborsSegmentId, offset, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s failed, ret %d", kHnswNeighborsSegmentId.c_str(), + ret); + return ret; + } + + //! dump level 0 neighbors meta + auto len = dump_segment(dumper, kHnswOffsetsSegmentId, hnsw_meta.data(), + hnsw_meta.size() * sizeof(HnswNeighborMeta)); + if (len < 0) { + return len; + } + + return len + offset + padding_size; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.h new file mode 100644 index 0000000..6d78eea --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.h @@ -0,0 +1,614 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include "zvec/core/framework/index_dumper.h" +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_storage.h" + +namespace zvec { +namespace core { + +using node_id_t = uint32_t; +using key_t = uint64_t; +using level_t = int32_t; +using dist_t = float; +struct EstimateRecord { + float ip_x0_qr; + float est_dist; + float low_dist; + + bool operator<(const EstimateRecord &other) const { + return this->est_dist < other.est_dist; + } +}; +struct ResultRecord { + float est_dist; + float low_dist; + ResultRecord() : est_dist(0.0f), low_dist(0.0f) {} + ResultRecord(float dist) : est_dist(dist), low_dist(dist) {} + explicit ResultRecord(const EstimateRecord &other) + : est_dist(other.est_dist), low_dist(other.low_dist) {} + ResultRecord(float est_dist, float low_dist) + : est_dist(est_dist), low_dist(low_dist) {} + bool operator<(const ResultRecord &other) const { + return this->est_dist < other.est_dist; + } + bool operator<=(const ResultRecord &other) const { + return this->est_dist <= other.est_dist; + } + bool operator>(const ResultRecord &other) const { + return this->est_dist > other.est_dist; + } +}; +using TopkHeap = ailego::KeyValueHeap; +using CandidateHeap = + ailego::KeyValueHeap>; +constexpr node_id_t kInvalidNodeId = static_cast(-1); +constexpr key_t kInvalidKey = static_cast(-1); +class DistCalculator; + +struct GraphHeader { + uint32_t size; + uint32_t version; + uint32_t graph_type; + uint32_t doc_count; + uint32_t vector_size; + uint32_t node_size; + uint32_t l0_neighbor_count; + uint32_t prune_type; + uint32_t prune_neighbor_count; + uint32_t ef_construction; + uint32_t options; + uint32_t min_neighbor_count; + uint32_t padded_dim; + uint32_t size_bin_data; + uint32_t size_ex_data; + uint8_t ex_bits; + uint8_t reserved_[4067]; +}; + +static_assert(sizeof(GraphHeader) % 32 == 0, + "GraphHeader must be aligned with 32 bytes"); + +//! Hnsw upper neighbor header +struct HnswHeader { + uint32_t size; // header size + uint32_t revision; // current total docs of the graph + uint32_t upper_neighbor_count; + uint32_t ef_construction; + uint32_t scaling_factor; + uint32_t max_level; + uint32_t entry_point; + uint32_t options; + uint8_t reserved_[30]; +}; + +static_assert(sizeof(HnswHeader) % 32 == 0, + "GraphHeader must be aligned with 32 bytes"); + +//! Hnsw common header and upper neighbor header +struct HNSWHeader { + HNSWHeader() { + clear(); + } + + HNSWHeader(const HNSWHeader &header) { + memcpy(static_cast(this), &header, sizeof(header)); + } + + HNSWHeader &operator=(const HNSWHeader &header) { + memcpy(static_cast(this), &header, sizeof(header)); + return *this; + } + + //! Reset state to zero, and the params is untouched + void inline reset() { + graph.doc_count = 0U; + hnsw.entry_point = kInvalidNodeId; + hnsw.max_level = 0; + } + + //! Clear all fields to init value + void inline clear() { + memset(static_cast(this), 0, sizeof(HNSWHeader)); + hnsw.entry_point = kInvalidNodeId; + graph.size = sizeof(GraphHeader); + hnsw.size = sizeof(HnswHeader); + } + + size_t l0_neighbor_cnt() const { + return graph.l0_neighbor_count; + } + + size_t upper_neighbor_cnt() const { + return hnsw.upper_neighbor_count; + } + + size_t vector_size() const { + return graph.vector_size; + } + + uint8_t ex_bits() const { + return graph.ex_bits; + } + + uint32_t padded_dim() const { + return graph.padded_dim; + } + + size_t ef_construction() const { + return graph.ef_construction; + } + + size_t scaling_factor() const { + return hnsw.scaling_factor; + } + + size_t neighbor_prune_cnt() const { + return graph.prune_neighbor_count; + } + + node_id_t entry_point() const { + return hnsw.entry_point; + } + + node_id_t doc_cnt() const { + return graph.doc_count; + } + + GraphHeader graph; + HnswHeader hnsw; +}; + +struct NeighborsHeader { + uint32_t neighbor_cnt; + node_id_t neighbors[0]; +}; + +struct Neighbors { + Neighbors() : cnt{0}, data{nullptr} {} + + Neighbors(uint32_t cnt_in, const node_id_t *data_in) + : cnt{cnt_in}, data{data_in} {} + + Neighbors(IndexStorage::MemoryBlock &&mem_block) + : neighbor_block{std::move(mem_block)} { + auto hd = reinterpret_cast(neighbor_block.data()); + cnt = hd->neighbor_cnt; + data = hd->neighbors; + } + + size_t size(void) const { + return cnt; + } + + const node_id_t &operator[](size_t idx) const { + return data[idx]; + } + + uint32_t cnt; + const node_id_t *data; + IndexStorage::MemoryBlock neighbor_block; +}; + +//! level 0 neighbors offset +struct GraphNeighborMeta { + GraphNeighborMeta(size_t o, size_t cnt) : offset(o), neighbor_cnt(cnt) {} + + uint64_t offset : 48; + uint64_t neighbor_cnt : 16; +}; + +//! hnsw upper neighbors meta +struct HnswNeighborMeta { + HnswNeighborMeta(size_t o, size_t l) : offset(o), level(l) {} + + uint64_t offset : 48; // offset = idx * upper neighors size + uint64_t level : 16; +}; + +class HnswRabitqEntity { + public: + //! Constructor + HnswRabitqEntity() {} + + //! Constructor + HnswRabitqEntity(const HNSWHeader &hd) { + header_ = hd; + } + + //! Destructor + virtual ~HnswRabitqEntity() {} + + //! HnswRabitqEntity Pointerd; + typedef std::shared_ptr Pointer; + + //! Get max neighbor size of graph level + inline size_t neighbor_cnt(level_t level) const { + return level == 0 ? header_.graph.l0_neighbor_count + : header_.hnsw.upper_neighbor_count; + } + + //! get max neighbor size of graph level 0 + inline size_t l0_neighbor_cnt() const { + return header_.graph.l0_neighbor_count; + } + + //! get min neighbor size of graph + inline size_t min_neighbor_cnt() const { + return header_.graph.min_neighbor_count; + } + + //! get upper neighbor size of graph level other than 0 + inline size_t upper_neighbor_cnt() const { + return header_.hnsw.upper_neighbor_count; + } + + //! Get current total doc of the hnsw graph + inline node_id_t *mutable_doc_cnt() { + return &header_.graph.doc_count; + } + + inline node_id_t doc_cnt() const { + return header_.graph.doc_count; + } + + //! Get hnsw graph scaling params + inline size_t scaling_factor() const { + return header_.hnsw.scaling_factor; + } + + //! Get prune_size + inline size_t prune_cnt() const { + return header_.graph.prune_neighbor_count; + } + + //! Current entity of top level graph + inline node_id_t entry_point() const { + return header_.hnsw.entry_point; + } + + //! Current max graph level + inline level_t cur_max_level() const { + return header_.hnsw.max_level; + } + + //! Retrieve index vector size + size_t vector_size() const { + return header_.graph.vector_size; + } + + //! Retrieve node size + size_t node_size() const { + return header_.graph.node_size; + } + + //! Retrieve ef constuction + size_t ef_construction() const { + return header_.graph.ef_construction; + } + + uint8_t ex_bits() const { + return header_.graph.ex_bits; + } + + uint32_t padded_dim() const { + return header_.graph.padded_dim; + } + + uint32_t size_bin_data() const { + return header_.graph.size_bin_data; + } + + uint32_t size_ex_data() const { + return header_.graph.size_ex_data; + } + + void update_rabitq_params_and_vector_size(uint32_t dimension); + + void set_ex_bits(uint8_t ex_bits) { + header_.graph.ex_bits = ex_bits; + } + + void set_prune_cnt(size_t v) { + header_.graph.prune_neighbor_count = v; + } + + void set_scaling_factor(size_t val) { + header_.hnsw.scaling_factor = val; + } + + void set_l0_neighbor_cnt(size_t cnt) { + header_.graph.l0_neighbor_count = cnt; + } + + void set_min_neighbor_cnt(size_t cnt) { + header_.graph.min_neighbor_count = cnt; + } + + void set_upper_neighbor_cnt(size_t cnt) { + header_.hnsw.upper_neighbor_count = cnt; + } + + void set_ef_construction(size_t ef) { + header_.graph.ef_construction = ef; + } + + protected: + inline const HNSWHeader &header() const { + return header_; + } + + inline HNSWHeader *mutable_header() { + return &header_; + } + + inline size_t header_size() const { + return sizeof(header_); + } + + void set_node_size(size_t size) { + header_.graph.node_size = size; + } + + //! Dump all segment by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_segments( + const IndexDumper::Pointer &dumper, key_t *keys, + const std::function &get_level) const; + + private: + //! dump mapping segment, for get_vector_by_key in provider + int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const; + + //! dump hnsw head by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_header(const IndexDumper::Pointer &dumper, + const HNSWHeader &hd) const; + + //! dump vectors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_vectors(const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + //! dump hnsw neighbors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_neighbors(const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + auto len1 = dump_graph_neighbors(dumper, reorder_mapping, neighbor_mapping); + if (len1 < 0) { + return len1; + } + auto len2 = dump_upper_neighbors(dumper, get_level, reorder_mapping, + neighbor_mapping); + if (len2 < 0) { + return len2; + } + + return len1 + len2; + } + + //! dump segment by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, const void *data, + size_t size) const; + + //! Dump level 0 neighbors + //! Return dump size if success, errno(<0) in failure + int64_t dump_graph_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + //! Dump upper level neighbors + //! Return dump size if success, errno(<0) in failure + int64_t dump_upper_neighbors( + const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + public: + //! Cleanup the entity + virtual int cleanup(void) { + header_.clear(); + return 0; + } + + //! Make a copy of searcher entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + virtual const HnswRabitqEntity::Pointer clone() const { + LOG_ERROR("Update neighbors not implemented"); + return HnswRabitqEntity::Pointer(); + } + + //! Get primary key of the node id + virtual key_t get_key(node_id_t id) const = 0; + + //! Get vector feature data by key + virtual const void *get_vector(node_id_t id) const = 0; + + //! Get vectors feature data by keys + virtual int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const = 0; + + virtual int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const = 0; + virtual int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const = 0; + + //! Retrieve a vector using a primary key + virtual const void *get_vector_by_key(uint64_t /*key*/) const { + LOG_ERROR("get vector not implemented"); + return nullptr; + } + + virtual int get_vector_by_key(const key_t /*key*/, + IndexStorage::MemoryBlock & /*block*/) const { + return IndexError_NotImplemented; + } + + //! Get the node id's neighbors on graph level + //! Note: the neighbors cannot be modified, using the following + //! method to get WritableNeighbors if want to + virtual const Neighbors get_neighbors(level_t level, node_id_t id) const = 0; + + //! Add vector and key to hnsw entity, and local id will be saved in id + virtual int add_vector(level_t /*level*/, key_t /*key*/, const void * /*vec*/, + node_id_t * /*id*/) { + return IndexError_NotImplemented; + } + + //! Add vector and id to hnsw entity + virtual int add_vector_with_id(level_t /*level*/, node_id_t /*id*/, + const void * /*vec*/) { + return IndexError_NotImplemented; + } + + virtual int update_neighbors( + level_t /*level*/, node_id_t /*id*/, + const std::vector> & /*neighbors*/) { + LOG_ERROR("Update neighbors dense not implemented"); + + return 0; + } + + //! Append neighbor_id to node id neighbors on level, size is the current + //! neighbors size. Notice: the caller must be ensure the neighbors not full + virtual void add_neighbor(level_t /*level*/, node_id_t /*id*/, + uint32_t /*size*/, node_id_t /*neighbor_id*/) { + LOG_ERROR("Add neighbor not implemented"); + } + + //! Update entry point and max level + virtual void update_ep_and_level(node_id_t ep, level_t level) { + header_.hnsw.entry_point = ep; + header_.hnsw.max_level = level; + } + + virtual int load(const IndexStorage::Pointer & /*container*/, + bool /*check_crc*/) { + LOG_ERROR("Load not implemented"); + return IndexError_NotImplemented; + } + + virtual int dump(const IndexDumper::Pointer & /*dumper*/) { + LOG_ERROR("Dump not implemented"); + return IndexError_NotImplemented; + } + + static int CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size); + + uint32_t get_cluster_id(const void *vec) const { + return *reinterpret_cast( + reinterpret_cast(vec) + cluster_id_offset()); + } + + const char *get_bin_data(const void *vec) const { + return reinterpret_cast(vec) + bin_data_offset(); + } + + const char *get_ex_data(const void *vec) const { + return reinterpret_cast(vec) + ex_data_offset(); + } + + uint32_t cluster_id_offset() const { + return 0; + } + + uint32_t bin_data_offset() const { + return cluster_id_offset() + sizeof(uint32_t); + } + + uint32_t ex_data_offset() const { + return bin_data_offset() + size_bin_data(); + } + + protected: + static inline size_t AlignSize(size_t size) { + return (size + 0x1F) & (~0x1F); + } + + static inline size_t AlignPageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::PageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + static inline size_t AlignHugePageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::HugePageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + //! rearrange vectors to improve cache locality + void reshuffle_vectors(const std::function &get_level, + std::vector *n2o_mapping, + std::vector *o2n_mapping, + key_t *keys) const; + + public: + const static std::string kGraphHeaderSegmentId; + const static std::string kGraphFeaturesSegmentId; + const static std::string kGraphKeysSegmentId; + const static std::string kGraphNeighborsSegmentId; + const static std::string kGraphOffsetsSegmentId; + const static std::string kGraphMappingSegmentId; + const static std::string kHnswHeaderSegmentId; + const static std::string kHnswNeighborsSegmentId; + const static std::string kHnswOffsetsSegmentId; + + constexpr static uint32_t kRevision = 0U; + constexpr static size_t kMaxGraphLayers = 15; + constexpr static uint32_t kDefaultEfConstruction = 500; + constexpr static uint32_t kDefaultEf = 500; + constexpr static uint32_t kDefaultUpperMaxNeighborCnt = 50; // M of HNSW + constexpr static uint32_t kDefaultL0MaxNeighborCnt = 100; + constexpr static uint32_t kMaxNeighborCnt = 65535; + constexpr static float kDefaultScanRatio = 0.1f; + constexpr static uint32_t kDefaultMinScanLimit = 10000; + constexpr static uint32_t kDefaultMaxScanLimit = + std::numeric_limits::max(); + constexpr static float kDefaultBFNegativeProbability = 0.001f; + constexpr static uint32_t kDefaultScalingFactor = 50U; + constexpr static uint32_t kDefaultBruteForceThreshold = 1000U; + constexpr static uint32_t kDefaultDocsHardLimit = 1 << 30U; // 1 billion + constexpr static float kDefaultDocsSoftLimitRatio = 0.9f; + constexpr static size_t kMaxChunkSize = 0xFFFFFFFF; + constexpr static size_t kDefaultChunkSize = 2UL * 1024UL * 1024UL; + constexpr static size_t kDefaultMaxChunkCnt = 50000UL; + constexpr static float kDefaultNeighborPruneMultiplier = + 1.0f; // prune_cnt = upper_max_neighbor_cnt * multiplier + constexpr static float kDefaultL0MaxNeighborCntMultiplier = + 2.0f; // l0_max_neighbor_cnt = upper_max_neighbor_cnt * multiplier + + protected: + HNSWHeader header_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_hash.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_hash.h new file mode 100644 index 0000000..bf3dc1e --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_hash.h @@ -0,0 +1,236 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "hnsw_rabitq_chunk.h" + +namespace zvec { +namespace core { + +//! Persistent hashmap implement through open addressing algorithm +template ::value>::type> +class HnswIndexHashMap { + using key_type = Key; + using val_type = Val; + + struct Iterator { + key_type first; + val_type second; + }; + typedef Iterator *iterator; + typedef Iterator Item; + typedef const Iterator *const_iterator; + + class Slot { + public: + Slot(Chunk::Pointer &&chunk, const void *data) + : chunk_(std::move(chunk)), + items_(reinterpret_cast(data)) {} + //! Return a empty loc or the key item loc + + Slot(Chunk::Pointer &&chunk, std::vector &&local_data) + : chunk_(std::move(chunk)), local_data_(std::move(local_data)) { + items_ = reinterpret_cast(local_data_.data()); + } + const_iterator find(key_type key, uint32_t max_items, uint32_t mask) const { + auto it = &items_[key & mask]; + for (auto i = 0U; i < max_items; ++i) { + if (it->first == key || it->second == EmptyVal) { + // LOG_DEBUG("i=%u", i); + return it; + } + ++it; + if (it == &items_[max_items]) { + it = &items_[0]; + } + } + return nullptr; + } + + bool update(const_iterator it) { + uint32_t offset = reinterpret_cast(it) - + reinterpret_cast(&items_[0]); + if (ailego_unlikely(chunk_->write(offset, it, sizeof(Item)) != + sizeof(Item))) { + LOG_ERROR("Chunk write failed"); + return false; + } + return true; + } + + private: + Chunk::Pointer chunk_{}; + const Item *items_{nullptr}; // point to local_data_ + std::vector local_data_{}; + }; + + public: + //! Init the hash + //! broker the index allocator + //! chunk_size the size of per chunk allocated, actual size may greater + //! factor factor = 1/ratio, ratio is the probability of a squence + //! number inserted to this container + //! max the max number key can be inserted + //! expansion_ratio memory expansion ratio + int init(HnswRabitqChunkBroker::Pointer &broker, uint32_t chunk_size, + uint32_t factor, size_t max, float expansion_ratio) { + ailego_assert_with(expansion_ratio > 1.0f, "ratio must > 1.0f"); + broker_ = broker; + + size_t items = std::ceil(chunk_size * 1.0f / sizeof(Item)); + slot_items_ = 1UL << static_cast((std::ceil(std::log2(items)))); + size_t range = slot_items_ * factor / expansion_ratio; + mask_bits_ = std::floor(std::log2(range)); + range = 1UL << mask_bits_; + size_t max_slots = std::ceil(max * 1.0f / range); + slots_.reserve(max_slots); + slot_loc_mask_ = slot_items_ - 1U; + int ret = load(); + if (ret != 0) { + return ret; + } + + LOG_DEBUG( + "HnswRabitqIndexHash init, chunkSize=%u factor=%u max=%zu " + "ratio=%f slotItems=%u maxSlots=%zu maskBits=%u " + "range=%zu", + chunk_size, factor, max, expansion_ratio, slot_items_, max_slots, + mask_bits_, range); + + return 0; + } + + int cleanup(void) { + broker_.reset(); + slots_.clear(); + slots_.shrink_to_fit(); + mask_bits_ = 0U; + slot_items_ = 0U; + slot_loc_mask_ = 0U; + + return 0; + } + + const_iterator end(void) const { + return nullptr; + } + + const_iterator find(const key_type key) const { + auto idx = key >> mask_bits_; + if (idx >= slots_.size()) { + return end(); + } + auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_); + return it && it->second != EmptyVal ? it : nullptr; + } + + bool insert(key_type key, val_type val) { + auto idx = key >> mask_bits_; + if (idx >= slots_.size()) { + if (ailego_unlikely(idx >= slots_.capacity())) { + LOG_ERROR("no space to insert"); + return false; + } + for (auto i = slots_.size(); i <= idx; ++i) { + if (ailego_unlikely(!alloc_slot(i))) { + return false; + } + } + } + auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_); + if (ailego_unlikely(it == nullptr)) { + LOG_ERROR("no space to insert"); + return false; + } + + //! TODO: write memory is ok? + const_cast(it)->first = key; + const_cast(it)->second = val; + + return slots_[idx].update(it); + } + + private: + bool alloc_slot(size_t idx) { + ailego_assert_with(idx == slots_.size(), "invalid idx"); + + size_t size = slot_items_ * sizeof(Item); + auto p = broker_->alloc_chunk( + HnswRabitqChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, idx, size); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return false; + } + Chunk::Pointer chunk = p.second; + if (ailego_unlikely(chunk->resize(size) != size)) { + LOG_ERROR("Chunk resize failed, size=%zu", size); + return false; + } + //! Use a local zero-initialized buffer; new chunks contain all zeros, + //! so no buffer-pool read is needed and no ref_count is pinned. + //! NOTE: Previously this used `chunk->read(0U, data_block, size)` which + //! returns a view into the underlying BufferPool page. That made the + //! Slot's `items_` pointer alias buffer-pool memory shared across + //! threads, which under clang -O3 release exposed a data race on + //! Slot::find()'s probing read of `it->second` (concurrent + //! const_cast writes from insert() were not reliably visible). Using a + //! private zero-initialized vector matches the HNSW (non-RABITQ) + //! implementation and avoids this race. + std::vector local_buf(size, 0); + slots_.emplace_back(std::move(chunk), std::move(local_buf)); + return true; + } + + int load(void) { + size_t slots_cnt = broker_->get_chunk_cnt( + HnswRabitqChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX); + for (size_t i = 0UL; i < slots_cnt; ++i) { + auto chunk = broker_->get_chunk( + HnswRabitqChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, i); + if (!chunk) { + LOG_ERROR("Get chunk failed, seq=%zu", i); + return IndexError_InvalidFormat; + } + size_t size = sizeof(Item) * slot_items_; + if (chunk->data_size() < size) { + LOG_ERROR( + "Hash params may be mismatch, seq=%zu, data_size=%zu " + "expect=%zu", + i, chunk->data_size(), size); + return IndexError_InvalidFormat; + } + //! Copy chunk data into a local buffer via fetch() so that no + //! buffer-pool block is pinned for the lifetime of the Slot. + std::vector local_buf(size); + if (ailego_unlikely(chunk->fetch(0U, local_buf.data(), size) != size)) { + LOG_ERROR("Chunk fetch failed, size=%zu", size); + return IndexError_InvalidFormat; + } + slots_.emplace_back(std::move(chunk), std::move(local_buf)); + } + return 0; + } + + private: + HnswRabitqChunkBroker::Pointer broker_{}; // chunk broker + std::vector slots_{}; + uint32_t mask_bits_{0U}; + uint32_t slot_items_{}; // must be a power of 2 + uint32_t slot_loc_mask_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_provider.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_provider.h new file mode 100644 index 0000000..7611bf1 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_provider.h @@ -0,0 +1,134 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "zvec/core/framework/index_provider.h" +#include "zvec/core/framework/index_searcher.h" +#include "zvec/core/framework/index_streamer.h" +#include "hnsw_rabitq_entity.h" + +namespace zvec { +namespace core { + +class HnswRabitqIndexProvider : public IndexProvider { + public: + HnswRabitqIndexProvider(const IndexMeta &meta, + const HnswRabitqEntity::Pointer &entity, + const std::string &owner) + : meta_(meta), entity_(entity), owner_class_(owner) {} + + HnswRabitqIndexProvider(const HnswRabitqIndexProvider &) = delete; + HnswRabitqIndexProvider &operator=(const HnswRabitqIndexProvider &) = delete; + + public: // holder interface + //! Create a new iterator + IndexProvider::Iterator::Pointer create_iterator() override { + return HnswRabitqIndexProvider::Iterator::Pointer(new (std::nothrow) + Iterator(entity_)); + } + + //! Retrieve count of vectors + size_t count(void) const override { + return entity_->doc_cnt(); + } + + //! Retrieve dimension of vector + size_t dimension(void) const override { + return meta_.dimension(); + } + + //! Retrieve type of vector + IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + //! Retrieve vector size in bytes + size_t element_size(void) const override { + return meta_.element_size(); + } + + public: // provider's unique interface + //! Retrieve a vector using a primary key + const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector_by_key(key, block); + } + + //! Retrieve the owner class + const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class Iterator : public IndexProvider::Iterator { + public: + Iterator(const HnswRabitqEntity::Pointer &entity) + : entity_(entity), cur_id_(0U) {} + + //! Retrieve pointer of data + //! NOTICE: the vec feature will be changed after iterating to next, so + //! the caller need to keep a copy of it before iterator to next vector + const void *data(void) const override { + return entity_->get_vector(cur_id_); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return entity_->get_key(cur_id_); + } + + //! Next iterator + void next(void) override { + // cur_id_ += 1; + cur_id_ = get_next_valid_id(cur_id_ + 1); + } + + //! Reset the iterator + void reset(void) { + cur_id_ = get_next_valid_id(0); + } + + private: + node_id_t get_next_valid_id(node_id_t start_id) { + for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) { + if (entity_->get_key(i) != kInvalidNodeId) { + cur_id_ = i; + return i; + } + } + return kInvalidNodeId; + } + + private: + const HnswRabitqEntity::Pointer entity_; + node_id_t cur_id_; + }; + + private: + const IndexMeta &meta_; + const HnswRabitqEntity::Pointer entity_; + const std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_params.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_params.h new file mode 100644 index 0000000..8b1c597 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_params.h @@ -0,0 +1,121 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +inline const std::string PARAM_HNSW_RABITQ_GENERAL_DIMENSION( + "proxima.hnsw_rabitq.general.dimension"); + +inline const std::string PARAM_HNSW_RABITQ_BUILDER_THREAD_COUNT( + "proxima.hnsw_rabitq.builder.thread_count"); +inline const std::string PARAM_HNSW_RABITQ_BUILDER_MEMORY_QUOTA( + "proxima.hnsw_rabitq.builder.memory_quota"); +inline const std::string PARAM_HNSW_RABITQ_BUILDER_EFCONSTRUCTION( + "proxima.hnsw_rabitq.builder.efconstruction"); +inline const std::string PARAM_HNSW_RABITQ_BUILDER_SCALING_FACTOR( + "proxima.hnsw_rabitq.builder.scaling_factor"); +inline const std::string PARAM_HNSW_RABITQ_BUILDER_CHECK_INTERVAL_SECS( + "proxima.hnsw_rabitq.builder.check_interval_secs"); +inline const std::string PARAM_HNSW_RABITQ_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER( + "proxima.hnsw_rabitq.builder.neighbor_prune_multiplier"); +inline const std::string PARAM_HNSW_RABITQ_BUILDER_MIN_NEIGHBOR_COUNT( + "proxima.hnsw_rabitq.builder.min_neighbor_count"); +inline const std::string PARAM_HNSW_RABITQ_BUILDER_MAX_NEIGHBOR_COUNT( + "proxima.hnsw_rabitq.builder.max_neighbor_count"); +inline const std::string + PARAM_HNSW_RABITQ_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER( + "proxima.hnsw_rabitq.builder.l0_max_neighbor_count_multiplier"); + +inline const std::string PARAM_HNSW_RABITQ_SEARCHER_EF( + "proxima.hnsw_rabitq.searcher.ef"); +inline const std::string PARAM_HNSW_RABITQ_SEARCHER_BRUTE_FORCE_THRESHOLD( + "proxima.hnsw_rabitq.searcher.brute_force_threshold"); +inline const std::string PARAM_HNSW_RABITQ_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE( + "proxima.hnsw_rabitq.searcher.neighbors_in_memory_enable"); +inline const std::string PARAM_HNSW_RABITQ_SEARCHER_MAX_SCAN_RATIO( + "proxima.hnsw_rabitq.searcher.max_scan_ratio"); +inline const std::string PARAM_HNSW_RABITQ_SEARCHER_CHECK_CRC_ENABLE( + "proxima.hnsw_rabitq.searcher.check_crc_enable"); +inline const std::string PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_ENABLE( + "proxima.hnsw_rabitq.searcher.visit_bloomfilter_enable"); +inline const std::string + PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB( + "proxima.hnsw_rabitq.searcher.visit_bloomfilter_negative_prob"); +inline const std::string PARAM_HNSW_RABITQ_SEARCHER_FORCE_PADDING_RESULT_ENABLE( + "proxima.hnsw_rabitq.searcher.force_padding_result_enable"); + +inline const std::string PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO( + "proxima.hnsw_rabitq.streamer.max_scan_ratio"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT( + "proxima.hnsw_rabitq.streamer.min_scan_limit"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT( + "proxima.hnsw_rabitq.streamer.max_scan_limit"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_EF( + "proxima.hnsw_rabitq.streamer.ef"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_EFCONSTRUCTION( + "proxima.hnsw_rabitq.streamer.efconstruction"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT( + "proxima.hnsw_rabitq.streamer.max_neighbor_count"); +inline const std::string + PARAM_HNSW_RABITQ_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER( + "proxima.hnsw_rabitq.streamer.l0_max_neighbor_count_multiplier"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_SCALING_FACTOR( + "proxima.hnsw_rabitq.streamer.scaling_factor"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_BRUTE_FORCE_THRESHOLD( + "proxima.hnsw_rabitq.streamer.brute_force_threshold"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT( + "proxima.hnsw_rabitq.streamer.docs_hard_limit"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT( + "proxima.hnsw_rabitq.streamer.docs_soft_limit"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_MAX_INDEX_SIZE( + "proxima.hnsw_rabitq.streamer.max_index_size"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_ENABLE( + "proxima.hnsw_rabitq.streamer.visit_bloomfilter_enable"); +inline const std::string + PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB( + "proxima.hnsw_rabitq.streamer.visit_bloomfilter_negative_prob"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_CHECK_CRC_ENABLE( + "proxima.hnsw_rabitq.streamer.check_crc_enable"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER( + "proxima.hnsw_rabitq.streamer.neighbor_prune_multiplier"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_CHUNK_SIZE( + "proxima.hnsw_rabitq.streamer.chunk_size"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_FILTER_SAME_KEY( + "proxima.hnsw_rabitq.streamer.filter_same_key"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_GET_VECTOR_ENABLE( + "proxima.hnsw_rabitq.streamer.get_vector_enable"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_MIN_NEIGHBOR_COUNT( + "proxima.hnsw_rabitq.streamer.min_neighbor_count"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_FORCE_PADDING_RESULT_ENABLE( + "proxima.hnsw_rabitq.streamer.force_padding_result_enable"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_ESTIMATE_DOC_COUNT( + "proxima.hnsw_rabitq.streamer.estimate_doc_count"); +inline const std::string PARAM_HNSW_RABITQ_STREAMER_USE_ID_MAP( + "proxima.hnsw_rabitq.streamer.use_id_map"); + +inline const std::string PARAM_HNSW_RABITQ_REDUCER_WORKING_PATH( + "proxima.hnsw_rabitq.reducer.working_path"); +inline const std::string PARAM_HNSW_RABITQ_REDUCER_NUM_OF_ADD_THREADS( + "proxima.hnsw_rabitq.reducer.num_of_add_threads"); +inline const std::string PARAM_HNSW_RABITQ_REDUCER_INDEX_NAME( + "proxima.hnsw_rabitq.reducer.index_name"); +inline const std::string PARAM_HNSW_RABITQ_REDUCER_EFCONSTRUCTION( + "proxima.hnsw_rabitq.reducer.efconstruction"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.cc new file mode 100644 index 0000000..d15cf11 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.cc @@ -0,0 +1,354 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_rabitq_query_algorithm.h" +#include +#include +#include +#include +#include "zvec/ailego/internal/platform.h" +#include "hnsw_rabitq_entity.h" +#include "hnsw_rabitq_query_entity.h" + +namespace zvec { +namespace core { + +HnswRabitqQueryAlgorithm::HnswRabitqQueryAlgorithm(HnswRabitqEntity &entity, + size_t num_clusters, + RabitqMetricType metric_type) + : entity_(entity), + mt_(std::chrono::system_clock::now().time_since_epoch().count()), + lock_pool_(kLockCnt), + num_clusters_(num_clusters), + metric_type_(metric_type) { + ex_bits_ = entity_.ex_bits(); + padded_dim_ = entity_.padded_dim(); + ip_func_ = rabitqlib::select_excode_ipfunc(ex_bits_); + LOG_INFO( + "Create query algorithm. num_clusters=%zu ex_bits=%zu padded_dim=%zu", + num_clusters_, ex_bits_, padded_dim_); +} + +int HnswRabitqQueryAlgorithm::cleanup() { + return 0; +} + +int HnswRabitqQueryAlgorithm::search(HnswRabitqQueryEntity *entity, + HnswRabitqContext *ctx) const { + spin_lock_.lock(); + auto maxLevel = entity_.cur_max_level(); + auto entry_point = entity_.entry_point(); + spin_lock_.unlock(); + + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + return 0; + } + + EstimateRecord curest; + get_bin_est(entity_.get_vector(entry_point), curest, *entity); + + for (level_t cur_level = maxLevel; cur_level >= 1; --cur_level) { + select_entry_point(cur_level, &entry_point, &curest, ctx, entity); + } + + auto &topk_heap = ctx->topk_heap(); + topk_heap.clear(); + search_neighbors(0, &entry_point, &curest, topk_heap, ctx, entity); + + if (ctx->group_by_search()) { + expand_neighbors_by_group(topk_heap, ctx, entity); + } + + return 0; +} + + +//! select_entry_point on hnsw level, ef = 1 +void HnswRabitqQueryAlgorithm::select_entry_point( + level_t level, node_id_t *entry_point, EstimateRecord *curest, + HnswRabitqContext *ctx, HnswRabitqQueryEntity *query_entity) const { + auto &entity = ctx->get_entity(); + while (true) { + const Neighbors neighbors = entity.get_neighbors(level, *entry_point); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + ailego_prefetch(neighbors.data); + uint32_t size = neighbors.size(); + if (size == 0) { + break; + } + + bool find_closer = false; + for (uint32_t i = 0; i < size; ++i) { + EstimateRecord candest; + get_bin_est(entity_.get_vector(neighbors[i]), candest, *query_entity); + + if (candest.est_dist < curest->est_dist) { + *curest = candest; + *entry_point = neighbors[i]; + find_closer = true; + } + } + + if (!find_closer) { + break; + } + } + + return; +} + +void HnswRabitqQueryAlgorithm::search_neighbors( + level_t level, node_id_t *entry_point, EstimateRecord *dist, TopkHeap &topk, + HnswRabitqContext *ctx, HnswRabitqQueryEntity *query_entity) const { + const auto &entity = ctx->get_entity(); + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + std::function filter = [](node_id_t) { return false; }; + if (ctx->filter().is_valid()) { + filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); }; + } + + candidates.clear(); + visit.clear(); + visit.set_visited(*entry_point); + if (!filter(*entry_point)) { + topk.emplace(*entry_point, ResultRecord(*dist)); + } + + candidates.emplace(*entry_point, ResultRecord(*dist)); + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t main_node = top->first; + auto main_dist = top->second; + + if (topk.full() && main_dist.est_dist > topk[0].second.est_dist) { + break; + } + + candidates.pop(); + const Neighbors neighbors = entity.get_neighbors(level, main_node); + ailego_prefetch(neighbors.data); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + + std::vector neighbor_ids(neighbors.size()); + uint32_t size = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) { + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_visit_dup_cnt())++; + } + continue; + } + visit.set_visited(node); + neighbor_ids[size++] = node; + } + if (size == 0) { + continue; + } + + for (uint32_t i = 0; i < size; ++i) { + node_id_t node = neighbor_ids[i]; + EstimateRecord candest; + auto *cand_vector = entity_.get_vector(node); + ailego_prefetch(cand_vector); + get_bin_est(cand_vector, candest, *query_entity); + + if (ex_bits_ > 0) { + // Check preliminary score against current worst full estimate. + bool flag_update_KNNs = + (!topk.full()) || candest.low_dist < topk[0].second.est_dist; + + if (flag_update_KNNs) { + // Compute the full estimate if promising. + get_full_est(cand_vector, candest, *query_entity); + } else { + continue; + } + } else { + // ex_bits_ == 0: est_dist is already the best estimate + if (topk.full() && candest.est_dist >= topk[0].second.est_dist) { + continue; + } + } + candidates.emplace(node, ResultRecord(candest)); + // update entry_point for next level scan + if (candest < *dist) { + *entry_point = node; + *dist = candest; + } + if (!filter(node)) { + topk.emplace(node, ResultRecord(candest)); + } + } // end for + } // while + + return; +} + +void HnswRabitqQueryAlgorithm::expand_neighbors_by_group( + TopkHeap &topk, HnswRabitqContext *ctx, + HnswRabitqQueryEntity *query_entity) const { + if (!ctx->group_by().is_valid()) { + return; + } + + const auto &entity = ctx->get_entity(); + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity.get_key(id)); + }; + + // devide into groups + std::map &group_topk_heaps = ctx->group_topk_heaps(); + for (uint32_t i = 0; i < topk.size(); ++i) { + node_id_t id = topk[i].first; + auto score = topk[i].second; + + std::string group_id = group_by(id); + + auto &topk_heap = group_topk_heaps[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, score); + } + + // stage 2, expand to reach group num as possible + if (group_topk_heaps.size() < ctx->group_num()) { + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + + std::function filter = [](node_id_t) { return false; }; + if (ctx->filter().is_valid()) { + filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); }; + } + + // refill to get enough groups + candidates.clear(); + visit.clear(); + for (uint32_t i = 0; i < topk.size(); ++i) { + node_id_t id = topk[i].first; + auto score = topk[i].second; + + visit.set_visited(id); + candidates.emplace_back(id, score); + } + + // do expand + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t main_node = top->first; + + candidates.pop(); + const Neighbors neighbors = entity.get_neighbors(0, main_node); + ailego_prefetch(neighbors.data); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + + std::vector neighbor_ids(neighbors.size()); + uint32_t size = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) { + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_visit_dup_cnt())++; + } + continue; + } + visit.set_visited(node); + neighbor_ids[size++] = node; + } + if (size == 0) { + continue; + } + + for (uint32_t i = 0; i < size; ++i) { + node_id_t node = neighbor_ids[i]; + EstimateRecord candest; + auto *cand_vector = entity_.get_vector(node); + ailego_prefetch(cand_vector); + get_full_est(cand_vector, candest, *query_entity); + + if (!filter(node)) { + std::string group_id = group_by(node); + + auto &topk_heap = group_topk_heaps[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(node, ResultRecord(candest)); + + if (group_topk_heaps.size() >= ctx->group_num()) { + break; + } + } + candidates.emplace(node, ResultRecord(candest)); + } // end for + } // end while + } // end if +} + +void HnswRabitqQueryAlgorithm::get_bin_est( + const void *vector, EstimateRecord &res, + HnswRabitqQueryEntity &entity) const { + const auto &q_to_centroids = entity.q_to_centroids; + auto &query_wrapper = *entity.query_wrapper; + uint32_t cluster_id = entity_.get_cluster_id(vector); + const char *bin_data = entity_.get_bin_data(vector); + if (metric_type_ == RabitqMetricType::kIP) { + float norm = q_to_centroids[cluster_id]; + float error = q_to_centroids[cluster_id + num_clusters_]; + rabitqlib::split_single_estdist(bin_data, query_wrapper, padded_dim_, + res.ip_x0_qr, res.est_dist, res.low_dist, + -norm, error); + } else { + // L2 distance + float norm = q_to_centroids[cluster_id]; + rabitqlib::split_single_estdist(bin_data, query_wrapper, padded_dim_, + res.ip_x0_qr, res.est_dist, res.low_dist, + norm * norm, norm); + } +} + +void HnswRabitqQueryAlgorithm::get_full_est( + const void *vector, EstimateRecord &res, + HnswRabitqQueryEntity &entity) const { + const auto &q_to_centroids = entity.q_to_centroids; + auto &query_wrapper = *entity.query_wrapper; + uint32_t cluster_id = entity_.get_cluster_id(vector); + const char *bin_data = entity_.get_bin_data(vector); + const char *ex_data = entity_.get_ex_data(vector); + + if (metric_type_ == RabitqMetricType::kIP) { + float norm = q_to_centroids[cluster_id]; + float error = q_to_centroids[cluster_id + num_clusters_]; + rabitqlib::split_single_fulldist(bin_data, ex_data, ip_func_, query_wrapper, + padded_dim_, ex_bits_, res.est_dist, + res.low_dist, res.ip_x0_qr, -norm, error); + } else { + // L2 distance + float norm = q_to_centroids[cluster_id]; + rabitqlib::split_single_fulldist( + bin_data, ex_data, ip_func_, query_wrapper, padded_dim_, ex_bits_, + res.est_dist, res.low_dist, res.ip_x0_qr, norm * norm, norm); + } +} + + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.h new file mode 100644 index 0000000..bcf6e42 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.h @@ -0,0 +1,136 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "hnsw_rabitq_context.h" +#include "hnsw_rabitq_dist_calculator.h" +#include "hnsw_rabitq_entity.h" +#include "rabitq_params.h" + +namespace zvec { +namespace core { + +struct HnswRabitqQueryEntity; + +//! hnsw graph algorithm implement +class HnswRabitqQueryAlgorithm { + public: + typedef std::unique_ptr UPointer; + + public: + //! Constructor + explicit HnswRabitqQueryAlgorithm(HnswRabitqEntity &entity, + size_t num_clusters, + RabitqMetricType metric_type); + + //! Destructor + ~HnswRabitqQueryAlgorithm() = default; + + //! Cleanup HnswRabitqQueryAlgorithm + int cleanup(); + + //! do knn search in graph + //! return 0 on success, or errCode in failure. results saved in ctx + int search(HnswRabitqQueryEntity *entity, HnswRabitqContext *ctx) const; + + //! Initiate HnswRabitqQueryAlgorithm + int init() { + level_probas_.clear(); + double level_mult = + 1 / std::log(static_cast(entity_.scaling_factor())); + for (int level = 0;; level++) { + // refers faiss get_random_level alg + double proba = + std::exp(-level / level_mult) * (1 - std::exp(-1 / level_mult)); + if (proba < 1e-9) { + break; + } + level_probas_.push_back(proba); + } + + return 0; + } + + //! Generate a random level + //! return graph level + uint32_t get_random_level() const { + // gen rand float (0, 1) + double f = mt_() / static_cast(mt_.max()); + for (size_t level = 0; level < level_probas_.size(); level++) { + if (f < level_probas_[level]) { + return level; + } + f -= level_probas_[level]; + } + return level_probas_.size() - 1; + } + void get_full_est(node_id_t id, EstimateRecord &res, + HnswRabitqQueryEntity &entity) const { + return get_full_est(entity_.get_vector(id), res, entity); + } + + private: + //! Select in upper layer to get entry point for next layer search + void select_entry_point(level_t level, node_id_t *entry_point, + EstimateRecord *dist, HnswRabitqContext *ctx, + HnswRabitqQueryEntity *entity) const; + + + //! Given a node id and level, search the nearest neighbors in graph + //! Note: the nearest neighbors result keeps in topk, and entry_point and + //! dist will be updated to current level nearest node id and distance + void search_neighbors(level_t level, node_id_t *entry_point, + EstimateRecord *dist, TopkHeap &topk, + HnswRabitqContext *ctx, + HnswRabitqQueryEntity *entity) const; + + + //! expand neighbors until group nums are reached + void expand_neighbors_by_group(TopkHeap &topk, HnswRabitqContext *ctx, + HnswRabitqQueryEntity *query_entity) const; + + void get_full_est(const void *vector, EstimateRecord &res, + HnswRabitqQueryEntity &entity) const; + void get_bin_est(const void *vector, EstimateRecord &res, + HnswRabitqQueryEntity &entity) const; + + private: + HnswRabitqQueryAlgorithm(const HnswRabitqQueryAlgorithm &) = delete; + HnswRabitqQueryAlgorithm &operator=(const HnswRabitqQueryAlgorithm &) = + delete; + + + private: + static constexpr uint32_t kLockCnt{1U << 8}; + static constexpr uint32_t kLockMask{kLockCnt - 1U}; + + HnswRabitqEntity &entity_; + mutable std::mt19937 mt_{}; + std::vector level_probas_{}; + + mutable ailego::SpinMutex spin_lock_{}; // global spin lock + std::mutex mutex_{}; // global mutex + // TODO: spin lock? + std::vector lock_pool_{}; + size_t num_clusters_{0}; + RabitqMetricType metric_type_{RabitqMetricType::kL2}; + size_t padded_dim_{0}; + size_t ex_bits_{0}; + float (*ip_func_)(const float *, const uint8_t *, size_t); +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_entity.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_entity.h new file mode 100644 index 0000000..760cdd2 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_entity.h @@ -0,0 +1,28 @@ +// Copyright 2025-present the centaurdb project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License + +#pragma once + +#include +#include +#include +namespace zvec::core { + +struct HnswRabitqQueryEntity { + std::vector rotated_query; + std::vector q_to_centroids; + std::unique_ptr> query_wrapper; +}; + +} // namespace zvec::core \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc new file mode 100644 index 0000000..4f48668 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc @@ -0,0 +1,24 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License +#include "hnsw_rabitq_streamer.h" +#include "rabitq_converter.h" +#include "rabitq_reformer.h" + +namespace zvec::core { + +INDEX_FACTORY_REGISTER_STREAMER(HnswRabitqStreamer); +INDEX_FACTORY_REGISTER_REFORMER_ALIAS(RabitqReformer, RabitqReformer); +INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(RabitqConverter, RabitqConverter); + +} // namespace zvec::core \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc new file mode 100644 index 0000000..6ebf57f --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc @@ -0,0 +1,995 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_rabitq_streamer.h" +#include +#include +#include +#include +#include +#include "algorithm/hnsw_rabitq/rabitq_reformer.h" +#include "zvec/ailego/container/params.h" +#include "zvec/ailego/logger/logger.h" +#include "hnsw_rabitq_algorithm.h" +#include "hnsw_rabitq_context.h" +#include "hnsw_rabitq_dist_calculator.h" +#include "hnsw_rabitq_index_provider.h" +#include "hnsw_rabitq_query_entity.h" +#include "rabitq_params.h" +#include "rabitq_utils.h" + +namespace zvec { +namespace core { +HnswRabitqStreamer::HnswRabitqStreamer() : entity_(stats_) {} + +HnswRabitqStreamer::HnswRabitqStreamer(IndexProvider::Pointer provider, + RabitqReformer::Pointer reformer) + : entity_(stats_), + reformer_(std::move(reformer)), + provider_(std::move(provider)) {} + +HnswRabitqStreamer::~HnswRabitqStreamer() { + if (state_ == STATE_INITED || state_ == STATE_OPENED) { + this->cleanup(); + } +} + +int HnswRabitqStreamer::init(const IndexMeta &imeta, + const ailego::Params ¶ms) { + meta_ = imeta; + meta_.set_streamer("HnswRabitqStreamer", HnswRabitqEntity::kRevision, params); + + params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_INDEX_SIZE, &max_index_size_); + + params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT, + &upper_max_neighbor_cnt_); + float multiplier = HnswRabitqEntity::kDefaultL0MaxNeighborCntMultiplier; + params.get(PARAM_HNSW_RABITQ_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, + &multiplier); + l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_; + + multiplier = HnswRabitqEntity::kDefaultNeighborPruneMultiplier; + params.get(PARAM_HNSW_RABITQ_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier); + size_t prune_cnt = multiplier * upper_max_neighbor_cnt_; + scaling_factor_ = upper_max_neighbor_cnt_; + params.get(PARAM_HNSW_RABITQ_STREAMER_SCALING_FACTOR, &scaling_factor_); + + params.get(PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_); + params.get(PARAM_HNSW_RABITQ_STREAMER_EF, &ef_); + params.get(PARAM_HNSW_RABITQ_STREAMER_EFCONSTRUCTION, &ef_construction_); + params.get(PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_); + params.get(PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB, + &bf_negative_prob_); + params.get(PARAM_HNSW_RABITQ_STREAMER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_); + params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_); + params.get(PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_); + params.get(PARAM_HNSW_RABITQ_STREAMER_CHECK_CRC_ENABLE, &check_crc_enabled_); + params.get(PARAM_HNSW_RABITQ_STREAMER_CHUNK_SIZE, &chunk_size_); + params.get(PARAM_HNSW_RABITQ_STREAMER_FILTER_SAME_KEY, &filter_same_key_); + params.get(PARAM_HNSW_RABITQ_STREAMER_GET_VECTOR_ENABLE, + &get_vector_enabled_); + params.get(PARAM_HNSW_RABITQ_STREAMER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_); + params.get(PARAM_HNSW_RABITQ_STREAMER_FORCE_PADDING_RESULT_ENABLE, + &force_padding_topk_enabled_); + params.get(PARAM_HNSW_RABITQ_STREAMER_USE_ID_MAP, &use_id_map_); + entity_.set_use_key_info_map(use_id_map_); + + params.get(PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit_); + if (docs_soft_limit_ > 0 && docs_soft_limit_ > docs_hard_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT.c_str(), + PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT.c_str()); + return IndexError_InvalidArgument; + } else if (docs_soft_limit_ == 0UL) { + docs_soft_limit_ = + docs_hard_limit_ * HnswRabitqEntity::kDefaultDocsSoftLimitRatio; + } + + if (ef_ == 0U) { + ef_ = HnswRabitqEntity::kDefaultEf; + } + if (ef_construction_ == 0U) { + ef_construction_ = HnswRabitqEntity::kDefaultEfConstruction; + } + if (upper_max_neighbor_cnt_ == 0U) { + upper_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt; + } + if (upper_max_neighbor_cnt_ > HnswRabitqEntity::kMaxNeighborCnt) { + LOG_ERROR("[%s] must be in range (0,%d)", + PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT.c_str(), + HnswRabitqEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (l0_max_neighbor_cnt_ == 0U) { + l0_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultL0MaxNeighborCnt; + } + if (l0_max_neighbor_cnt_ > HnswRabitqEntity::kMaxNeighborCnt) { + LOG_ERROR("MaxL0NeighborCnt must be in range (0,%d)", + HnswRabitqEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) { + LOG_ERROR("[%s]-[%zu] must be <= [%s]-[%zu]", + PARAM_HNSW_RABITQ_STREAMER_MIN_NEIGHBOR_COUNT.c_str(), + static_cast(min_neighbor_cnt_), + PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT.c_str(), + static_cast(upper_max_neighbor_cnt_)); + return IndexError_InvalidArgument; + } + + if (bf_negative_prob_ <= 0.0f || bf_negative_prob_ >= 1.0f) { + LOG_ERROR( + "[%s] must be in range (0,1)", + PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str()); + return IndexError_InvalidArgument; + } + + if (scaling_factor_ == 0U) { + scaling_factor_ = HnswRabitqEntity::kDefaultScalingFactor; + } + if (scaling_factor_ < 5 || scaling_factor_ > 1000) { + LOG_ERROR("[%s] must be in range [5,1000]", + PARAM_HNSW_RABITQ_STREAMER_SCALING_FACTOR.c_str()); + return IndexError_InvalidArgument; + } + + if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) { + LOG_ERROR("[%s] must be in range (0.0f,1.0f]", + PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO.c_str()); + return IndexError_InvalidArgument; + } + + if (max_scan_limit_ < min_scan_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT.c_str(), + PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT.c_str()); + return IndexError_InvalidArgument; + } + + if (prune_cnt == 0UL) { + prune_cnt = upper_max_neighbor_cnt_; + } + if (chunk_size_ == 0UL) { + chunk_size_ = HnswRabitqEntity::kDefaultChunkSize; + } + if (chunk_size_ > HnswRabitqEntity::kMaxChunkSize) { + LOG_ERROR("[%s] must be < %zu", + PARAM_HNSW_RABITQ_STREAMER_CHUNK_SIZE.c_str(), + HnswRabitqEntity::kMaxChunkSize); + return IndexError_InvalidArgument; + } + uint32_t total_bits = 0; + params.get(PARAM_RABITQ_TOTAL_BITS, &total_bits); + if (total_bits == 0) { + total_bits = kDefaultRabitqTotalBits; + } + if (total_bits < 1 || total_bits > 9) { + LOG_ERROR("Invalid total_bits: %zu, must be in [1, 9]", (size_t)total_bits); + return IndexError_InvalidArgument; + } + uint8_t ex_bits = total_bits - 1; + entity_.set_ex_bits(ex_bits); + + uint32_t dimension = 0; + params.get(PARAM_HNSW_RABITQ_GENERAL_DIMENSION, &dimension); + if (dimension == 0) { + LOG_ERROR("%s not set", PARAM_HNSW_RABITQ_GENERAL_DIMENSION.c_str()); + return IndexError_InvalidArgument; + } + if (dimension < kMinRabitqDimSize || dimension > kMaxRabitqDimSize) { + LOG_ERROR("Invalid dimension: %u, must be in [%d, %d]", dimension, + kMinRabitqDimSize, kMaxRabitqDimSize); + return IndexError_InvalidArgument; + } + entity_.update_rabitq_params_and_vector_size(dimension); + + entity_.set_ef_construction(ef_construction_); + entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_); + entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_); + entity_.set_scaling_factor(scaling_factor_); + entity_.set_prune_cnt(prune_cnt); + + entity_.set_chunk_size(chunk_size_); + entity_.set_filter_same_key(filter_same_key_); + entity_.set_get_vector(get_vector_enabled_); + entity_.set_min_neighbor_cnt(min_neighbor_cnt_); + + int ret = entity_.init(docs_hard_limit_); + if (ret != 0) { + LOG_ERROR("Hnsw entity init failed for %s", IndexError::What(ret)); + return ret; + } + + LOG_DEBUG( + "Init params: maxIndexSize=%zu docsHardLimit=%zu docsSoftLimit=%zu " + "efConstruction=%u ef=%u upperMaxNeighborCnt=%u l0MaxNeighborCnt=%u " + "scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu " + "bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbability=%.5f " + "checkCrcEnabled=%d pruneSize=%zu vectorSize=%u chunkSize=%zu " + "filterSameKey=%u getVectorEnabled=%u minNeighborCount=%u " + "forcePadding=%u ", + max_index_size_, docs_hard_limit_, docs_soft_limit_, ef_construction_, + ef_, upper_max_neighbor_cnt_, l0_max_neighbor_cnt_, scaling_factor_, + max_scan_ratio_, min_scan_limit_, max_scan_limit_, bf_enabled_, + bruteforce_threshold_, bf_negative_prob_, check_crc_enabled_, prune_cnt, + meta_.element_size(), chunk_size_, filter_same_key_, get_vector_enabled_, + min_neighbor_cnt_, force_padding_topk_enabled_); + + alg_ = HnswRabitqAlgorithm::UPointer(new HnswRabitqAlgorithm(entity_)); + + ret = alg_->init(); + if (ret != 0) { + return ret; + } + + state_ = STATE_INITED; + + return 0; +} + +int HnswRabitqStreamer::cleanup(void) { + if (state_ == STATE_OPENED) { + this->close(); + } + + LOG_INFO("HnswRabitqStreamer cleanup"); + + meta_.clear(); + metric_.reset(); + stats_.clear(); + entity_.cleanup(); + + if (alg_) { + alg_->cleanup(); + } + + max_index_size_ = 0UL; + docs_hard_limit_ = HnswRabitqEntity::kDefaultDocsHardLimit; + docs_soft_limit_ = 0UL; + upper_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt; + l0_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultL0MaxNeighborCnt; + ef_ = HnswRabitqEntity::kDefaultEf; + ef_construction_ = HnswRabitqEntity::kDefaultEfConstruction; + bf_enabled_ = false; + scaling_factor_ = HnswRabitqEntity::kDefaultScalingFactor; + bruteforce_threshold_ = HnswRabitqEntity::kDefaultBruteForceThreshold; + max_scan_limit_ = HnswRabitqEntity::kDefaultMaxScanLimit; + min_scan_limit_ = HnswRabitqEntity::kDefaultMinScanLimit; + chunk_size_ = HnswRabitqEntity::kDefaultChunkSize; + bf_negative_prob_ = HnswRabitqEntity::kDefaultBFNegativeProbability; + max_scan_ratio_ = HnswRabitqEntity::kDefaultScanRatio; + state_ = STATE_INIT; + check_crc_enabled_ = false; + filter_same_key_ = false; + get_vector_enabled_ = false; + + return 0; +} + +int HnswRabitqStreamer::open(IndexStorage::Pointer stg) { + LOG_INFO("HnswRabitqStreamer open"); + + if (ailego_unlikely(state_ != STATE_INITED)) { + LOG_ERROR("Open storage failed, init streamer first!"); + return IndexError_NoReady; + } + + // try to load reformer + if (reformer_ == nullptr) { + reformer_ = std::make_shared(); + ailego::Params reformer_params; + reformer_params.set(PARAM_RABITQ_METRIC_NAME, meta_.metric_name()); + int ret = reformer_->init(reformer_params); + if (ret != 0) { + LOG_ERROR("Failed to initialize RabitqReformer: %d", ret); + return ret; + } + + ret = reformer_->load(stg); + if (ret != 0) { + LOG_ERROR("Failed to load reformer, ret=%d", ret); + return ret; + } + } else { + if (!stg->has(RABITQ_CONVERTER_SEG_ID)) { + int ret = reformer_->dump(stg); + if (ret != 0) { + LOG_ERROR("Failed to dump reformer, ret=%d", ret); + return ret; + } + LOG_INFO("Dump reformer success."); + } + } + + int ret = entity_.open(std::move(stg), max_index_size_, check_crc_enabled_); + if (ret != 0) { + return ret; + } + + // Verify ex_bits consistency to avoid quantized data layout mismatch + if (reformer_->ex_bits() != entity_.ex_bits()) { + LOG_ERROR( + "ex_bits mismatch between reformer(%zu) and entity(%zu). " + "Reformer and entity must use the same total_bits configuration", + reformer_->ex_bits(), (size_t)entity_.ex_bits()); + return IndexError_Mismatch; + } + + IndexMeta index_meta; + ret = entity_.get_index_meta(&index_meta); + if (ret == IndexError_NoExist) { + // Set IndexMeta for the new index + ret = entity_.set_index_meta(meta_); + if (ret != 0) { + LOG_ERROR("Failed to set index meta for %s", IndexError::What(ret)); + return ret; + } + } else if (ret != 0) { + LOG_ERROR("Failed to get index meta for %s", IndexError::What(ret)); + return ret; + } else { + if (index_meta.dimension() != meta_.dimension() || + index_meta.element_size() != meta_.element_size() || + index_meta.metric_name() != meta_.metric_name() || + index_meta.data_type() != meta_.data_type()) { + LOG_ERROR("IndexMeta mismatch from the previous in index"); + return IndexError_Mismatch; + } + // The IndexMetric Params may be updated like MipsSquaredEuclidean + auto metric_params = index_meta.metric_params(); + metric_params.merge(meta_.metric_params()); + meta_.set_metric(index_meta.metric_name(), 0, metric_params); + } + + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Failed to create metric %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("Failed to init metric, ret=%d", ret); + return ret; + } + + if (!metric_->distance()) { + LOG_ERROR("Invalid metric distance"); + return IndexError_InvalidArgument; + } + + if (!metric_->batch_distance()) { + LOG_ERROR("Invalid metric batch distance"); + return IndexError_InvalidArgument; + } + + add_distance_ = metric_->distance(); + add_batch_distance_ = metric_->batch_distance(); + + search_distance_ = add_distance_; + search_batch_distance_ = add_batch_distance_; + + if (metric_->query_metric() && metric_->query_metric()->distance() && + metric_->query_metric()->batch_distance()) { + search_distance_ = metric_->query_metric()->distance(); + search_batch_distance_ = metric_->query_metric()->batch_distance(); + } + + state_ = STATE_OPENED; + magic_ = IndexContext::GenerateMagic(); + + query_alg_ = HnswRabitqQueryAlgorithm::UPointer(new HnswRabitqQueryAlgorithm( + entity_, reformer_->num_clusters(), reformer_->rabitq_metric_type())); + + return 0; +} + +int HnswRabitqStreamer::close(void) { + LOG_INFO("HnswRabitqStreamer close"); + + stats_.clear(); + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_.set_index_meta(meta_); + int ret = entity_.close(); + if (ret != 0) { + return ret; + } + state_ = STATE_INITED; + + return 0; +} + +int HnswRabitqStreamer::flush(uint64_t checkpoint) { + LOG_INFO("HnswRabitqStreamer flush checkpoint=%zu", (size_t)checkpoint); + + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_.set_index_meta(meta_); + return entity_.flush(checkpoint); +} + +int HnswRabitqStreamer::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("HnswRabitqStreamer dump"); + + shared_mutex_.lock(); + AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); + + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + ret = reformer_->dump(dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump reformer into dumper."); + return ret; + } + return entity_.dump(dumper); +} + +IndexStreamer::Context::Pointer HnswRabitqStreamer::create_context(void) const { + if (ailego_unlikely(state_ != STATE_OPENED)) { + LOG_ERROR("Create context failed, open storage first!"); + return Context::Pointer(); + } + + HnswRabitqEntity::Pointer entity = entity_.clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("CreateContext clone init failed"); + return Context::Pointer(); + } + HnswRabitqContext *ctx = + new (std::nothrow) HnswRabitqContext(meta_.dimension(), metric_, entity); + if (ailego_unlikely(ctx == nullptr)) { + LOG_ERROR("Failed to new HnswRabitqContext"); + return Context::Pointer(); + } + ctx->set_ef(ef_); + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_filter_mode(bf_enabled_ ? VisitFilter::BloomFilter + : VisitFilter::ByteMap); + ctx->set_filter_negative_probability(bf_negative_prob_); + ctx->set_magic(magic_); + ctx->set_force_padding_topk(force_padding_topk_enabled_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + + if (ailego_unlikely(ctx->init(HnswRabitqContext::kStreamerContext)) != 0) { + LOG_ERROR("Init HnswRabitqContext failed"); + delete ctx; + return Context::Pointer(); + } + uint32_t estimate_doc_count = 0; + if (meta_.streamer_params().get(PARAM_HNSW_RABITQ_STREAMER_ESTIMATE_DOC_COUNT, + &estimate_doc_count)) { + LOG_DEBUG("HnswRabitqStreamer doc_count[%zu] estimate[%zu]", + (size_t)entity_.doc_cnt(), (size_t)estimate_doc_count); + } + ctx->check_need_adjuct_ctx(std::max(entity_.doc_cnt(), estimate_doc_count)); + + return Context::Pointer(ctx); +} + +IndexProvider::Pointer HnswRabitqStreamer::create_provider(void) const { + LOG_DEBUG("HnswRabitqStreamer create provider"); + + auto entity = entity_.clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone HnswRabitqEntity failed"); + return nullptr; + } + return Provider::Pointer( + new HnswRabitqIndexProvider(meta_, entity, "HnswRabitqStreamer")); +} + +int HnswRabitqStreamer::update_context(HnswRabitqContext *ctx) const { + const HnswRabitqEntity::Pointer entity = entity_.clone(); + if (!entity) { + LOG_ERROR("Failed to clone search context entity"); + return IndexError_Runtime; + } + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + return ctx->update_context(HnswRabitqContext::kStreamerContext, meta_, + metric_, entity, magic_); +} + +//! Add a vector with id into index +int HnswRabitqStreamer::add_with_id_impl( + uint32_t id, const void *query, const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) { + if (!provider_) { + LOG_ERROR("Provider is nullptr, cannot add vector"); + return IndexError_InvalidArgument; + } + + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + HnswRabitqContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswRabitqContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) { + if (entity_.doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %zu exceed [%s]", + static_cast(entity_.doc_cnt()), + PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT.c_str()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %zu exceed [%s]", + static_cast(entity_.doc_cnt()), + PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT.c_str()); + } + } + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_); + ctx->reset_query(query); + ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + ctx->set_provider(provider_); + + if (metric_->support_train()) { + const std::lock_guard lk(mutex_); + ret = metric_->train(query, meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer metric train failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + } + + std::string converted_vector; + IndexQueryMeta converted_meta; + ret = reformer_->convert(query, qmeta, &converted_vector, &converted_meta); + if (ret != 0) { + LOG_ERROR("Rabitq hnsw convert failed, ret=%d", ret); + return ret; + } + + level_t level = alg_->get_random_level(); + ret = entity_.add_vector_with_id(level, id, converted_vector.data()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, level, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw steamer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + +//! Add a vector into index +int HnswRabitqStreamer::add_impl(uint64_t pkey, const void *query, + const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) { + if (!provider_) { + LOG_ERROR("Provider is nullptr, cannot add vector"); + return IndexError_InvalidArgument; + } + + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + HnswRabitqContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswRabitqContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) { + if (entity_.doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %zu exceed [%s]", + static_cast(entity_.doc_cnt()), + PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT.c_str()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %zu exceed [%s]", + static_cast(entity_.doc_cnt()), + PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT.c_str()); + } + } + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_); + ctx->reset_query(query); + ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + ctx->set_provider(provider_); + + if (metric_->support_train()) { + const std::lock_guard lk(mutex_); + ret = metric_->train(query, meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer metric train failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + } + + std::string converted_vector; + IndexQueryMeta converted_meta; + ret = reformer_->convert(query, qmeta, &converted_vector, &converted_meta); + if (ret != 0) { + LOG_ERROR("Rabitq hnsw convert failed, ret=%d", ret); + return ret; + } + + level_t level = alg_->get_random_level(); + node_id_t id; + ret = entity_.add_vector(level, pkey, converted_vector.data(), &id); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, level, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw steamer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + + +int HnswRabitqStreamer::search_impl( + const void *query, const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) const { + return search_impl(query, qmeta, 1, context); +} + +//! Similarity search +int HnswRabitqStreamer::search_impl( + const void *query, const IndexQueryMeta &qmeta, uint32_t count, + IndexStreamer::Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + HnswRabitqContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswRabitqContext failed"); + return IndexError_Cast; + } + + if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) { + return search_bf_impl(query, qmeta, count, context); + } + + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + for (size_t q = 0; q < count; ++q) { + HnswRabitqQueryEntity entity; + ret = reformer_->transform_to_entity(query, &entity); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw searcher transform failed"); + return ret; + } + ctx->reset_query(query); + ret = query_alg_->search(&entity, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw searcher fast search failed"); + return ret; + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +void HnswRabitqStreamer::print_debug_info() { + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + if (entity_.get_key(id) == kInvalidKey) { + continue; + } + Neighbors neighbours = entity_.get_neighbors(0, id); + std::cout << "node: " << id << "; "; + if (neighbours.size() == 0) std::cout << std::endl; + for (uint32_t i = 0; i < neighbours.size(); ++i) { + std::cout << neighbours[i]; + + if (i == neighbours.size() - 1) { + std::cout << std::endl; + } else { + std::cout << ", "; + } + } + } + + // entity_.print_key_map(); +} + +int HnswRabitqStreamer::search_bf_impl( + const void *query, const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) const { + return search_bf_impl(query, qmeta, 1, context); +} + +int HnswRabitqStreamer::search_bf_impl( + const void *query, const IndexQueryMeta &qmeta, uint32_t count, + IndexStreamer::Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + HnswRabitqContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswRabitqContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity_.get_key(id)); + }; + + for (size_t q = 0; q < count; ++q) { + HnswRabitqQueryEntity entity; + ret = reformer_->transform_to_entity(query, &entity); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw rabitq streamer transform failed"); + return ret; + } + ctx->reset_query(query); + ctx->group_topk_heaps().clear(); + + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + if (entity_.get_key(id) == kInvalidKey) { + continue; + } + + if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { + EstimateRecord dist; + query_alg_->get_full_est(id, dist, entity); + + std::string group_id = group_by(id); + + auto &topk_heap = ctx->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, dist); + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } else { + for (size_t q = 0; q < count; ++q) { + HnswRabitqQueryEntity entity; + ret = reformer_->transform_to_entity(query, &entity); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw rabitq streamer transform failed"); + return ret; + } + ctx->reset_query(query); + ctx->topk_heap().clear(); + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + if (entity_.get_key(id) == kInvalidKey) { + continue; + } + if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { + EstimateRecord dist; + query_alg_->get_full_est(id, dist, entity); + ctx->topk_heap().emplace(id, dist); + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +int HnswRabitqStreamer::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + if (ailego_unlikely(p_keys.size() != count)) { + LOG_ERROR("The size of p_keys is not equal to count"); + return IndexError_InvalidArgument; + } + + HnswRabitqContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswRabitqContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_InvalidArgument; + } + + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity_.get_key(id)); + }; + + for (size_t q = 0; q < count; ++q) { + HnswRabitqQueryEntity entity; + ret = reformer_->transform_to_entity(query, &entity); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw rabitq streamer transform failed"); + return ret; + } + ctx->reset_query(query); + ctx->group_topk_heaps().clear(); + + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t pk = p_keys[q][idx]; + if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { + node_id_t id = entity_.get_id(pk); + if (id != kInvalidNodeId) { + EstimateRecord dist; + query_alg_->get_full_est(id, dist, entity); + std::string group_id = group_by(id); + + auto &topk_heap = ctx->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, dist); + } + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } else { + for (size_t q = 0; q < count; ++q) { + HnswRabitqQueryEntity entity; + ret = reformer_->transform_to_entity(query, &entity); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw rabitq streamer transform failed"); + return ret; + } + ctx->reset_query(query); + ctx->topk_heap().clear(); + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + key_t pk = p_keys[q][idx]; + if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { + node_id_t id = entity_.get_id(pk); + if (id != kInvalidNodeId) { + EstimateRecord dist; + query_alg_->get_full_est(id, dist, entity); + ctx->topk_heap().emplace(id, dist); + } + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.h new file mode 100644 index 0000000..4d03b59 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.h @@ -0,0 +1,245 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "algorithm/hnsw_rabitq/rabitq_reformer.h" +#include "zvec/core/framework/index_framework.h" +#include "zvec/core/framework/index_provider.h" +#include "zvec/core/framework/index_reformer.h" +#include "hnsw_rabitq_algorithm.h" +#include "hnsw_rabitq_query_algorithm.h" +#include "hnsw_rabitq_streamer_entity.h" + +namespace zvec { +namespace core { + +class HnswRabitqStreamer : public IndexStreamer { + public: + using ContextPointer = IndexStreamer::Context::Pointer; + + HnswRabitqStreamer(); + explicit HnswRabitqStreamer(IndexProvider::Pointer provider, + RabitqReformer::Pointer reformer = nullptr); + ~HnswRabitqStreamer(void) override; + + HnswRabitqStreamer(const HnswRabitqStreamer &streamer) = delete; + HnswRabitqStreamer &operator=(const HnswRabitqStreamer &streamer) = delete; + + void set_provider(IndexProvider::Pointer provider) { + provider_ = std::move(provider); + } + + void set_reformer(IndexReformer::Pointer reformer) { + reformer_ = std::dynamic_pointer_cast(reformer); + } + + protected: + //! Initialize Streamer + int init(const IndexMeta &imeta, const ailego::Params ¶ms) override; + + //! Cleanup Streamer + int cleanup(void) override; + + //! Create a context + Context::Pointer create_context(void) const override; + + //! Create a new iterator + IndexProvider::Pointer create_provider(void) const override; + + //! Add a vector into index + int add_impl(uint64_t pkey, const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Add a vector with id into index + int add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + //! Fetch vector by key + const void *get_vector(uint64_t key) const override { + return entity_.get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_.get_vector_by_key(key, block); + } + + //! Fetch vector by id + const void *get_vector_by_id(uint32_t id) const override { + return entity_.get_vector(id); + } + + int get_vector_by_id(const uint32_t id, + IndexStorage::MemoryBlock &block) const override { + return entity_.get_vector(id, block); + } + + //! Open index from file path + int open(IndexStorage::Pointer stg) override; + + //! Close file + int close(void) override; + + //! flush file + int flush(uint64_t checkpoint) override; + + //! Dump index into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + void print_debug_info() override; + + private: + inline int check_params(const void *query, + const IndexQueryMeta &qmeta) const { + if (ailego_unlikely(!query)) { + LOG_ERROR("null query"); + return IndexError_InvalidArgument; + } + if (ailego_unlikely(qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size())) { + LOG_ERROR("Unsupported query meta"); + return IndexError_Mismatch; + } + return 0; + } + + inline int check_sparse_count_is_zero(const uint32_t *sparse_count, + uint32_t count) const { + for (uint32_t i = 0; i < count; ++i) { + if (sparse_count[i] != 0) + LOG_ERROR("Sparse cout is not empty. Index: %u, Sparse Count: %u", i, + sparse_count[i]); + return IndexError_InvalidArgument; + } + + return 0; + } + + private: + //! To share ctx across streamer/searcher, we need to update the context for + //! current streamer/searcher + int update_context(HnswRabitqContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 }; + class Stats : public IndexStreamer::Stats { + public: + void clear(void) { + set_revision_id(0u); + set_loaded_count(0u); + set_added_count(0u); + set_discarded_count(0u); + set_index_size(0u); + set_dumped_size(0u); + set_check_point(0u); + set_create_time(0u); + set_update_time(0u); + clear_attributes(); + } + }; + + Stats stats_{}; + HnswRabitqStreamerEntity entity_; + HnswRabitqAlgorithm::UPointer alg_; + IndexMeta meta_{}; + IndexMetric::Pointer metric_{}; + + IndexMetric::MatrixDistance add_distance_{}; + IndexMetric::MatrixDistance search_distance_{}; + + IndexMetric::MatrixBatchDistance add_batch_distance_{}; + IndexMetric::MatrixBatchDistance search_batch_distance_{}; + + RabitqReformer::Pointer reformer_{}; // RaBitQ reformer + HnswRabitqQueryAlgorithm::UPointer query_alg_; // query algorithm + // provider_ provides raw vector, which is used to build graph + IndexProvider::Pointer provider_{}; + std::mutex mutex_{}; + + size_t max_index_size_{0UL}; + size_t chunk_size_{HnswRabitqEntity::kDefaultChunkSize}; + size_t docs_hard_limit_{HnswRabitqEntity::kDefaultDocsHardLimit}; + size_t docs_soft_limit_{0UL}; + uint32_t min_neighbor_cnt_{0u}; + uint32_t upper_max_neighbor_cnt_{ + HnswRabitqEntity::kDefaultUpperMaxNeighborCnt}; + uint32_t l0_max_neighbor_cnt_{HnswRabitqEntity::kDefaultL0MaxNeighborCnt}; + uint32_t ef_{HnswRabitqEntity::kDefaultEf}; + uint32_t ef_construction_{HnswRabitqEntity::kDefaultEfConstruction}; + uint32_t scaling_factor_{HnswRabitqEntity::kDefaultScalingFactor}; + size_t bruteforce_threshold_{HnswRabitqEntity::kDefaultBruteForceThreshold}; + size_t max_scan_limit_{HnswRabitqEntity::kDefaultMaxScanLimit}; + size_t min_scan_limit_{HnswRabitqEntity::kDefaultMinScanLimit}; + float bf_negative_prob_{HnswRabitqEntity::kDefaultBFNegativeProbability}; + float max_scan_ratio_{HnswRabitqEntity::kDefaultScanRatio}; + + uint32_t magic_{0U}; + State state_{STATE_INIT}; + bool bf_enabled_{false}; + bool check_crc_enabled_{false}; + bool filter_same_key_{false}; + bool get_vector_enabled_{false}; + bool force_padding_topk_enabled_{false}; + bool use_id_map_{true}; + + //! avoid add vector while dumping index + ailego::SharedMutex shared_mutex_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.cc new file mode 100644 index 0000000..27643aa --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.cc @@ -0,0 +1,711 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "hnsw_rabitq_streamer_entity.h" +#include + +// #define DEBUG_PRINT + +namespace zvec { +namespace core { + +HnswRabitqStreamerEntity::HnswRabitqStreamerEntity(IndexStreamer::Stats &stats) + : stats_(stats) {} + +HnswRabitqStreamerEntity::~HnswRabitqStreamerEntity() {} + +int HnswRabitqStreamerEntity::init(size_t max_doc_cnt) { + if (std::pow(scaling_factor(), kMaxGraphLayers) < max_doc_cnt) { + LOG_ERROR("scalingFactor=%zu is too small", scaling_factor()); + return IndexError_InvalidArgument; + } + + std::lock_guard lock(mutex_); + broker_ = std::make_shared(stats_); + upper_neighbor_index_ = std::make_shared(); + upper_neighbor_rw_mutex_ = std::make_shared(); + keys_map_lock_ = std::make_shared(); + keys_map_ = std::make_shared>(); + if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) { + LOG_ERROR("HnswRabitqStreamerEntity new object failed"); + return IndexError_NoMemory; + } + keys_map_->set_empty_key(kInvalidKey); + + neighbor_size_ = neighbors_size(); + upper_neighbor_size_ = upper_neighbors_size(); + + //! vector + key + level 0 neighbors + size_t size = vector_size() + sizeof(key_t) + neighbor_size_; + + size = AlignSize(size); + set_node_size(size); + return 0; +} + +int HnswRabitqStreamerEntity::cleanup() { + std::lock_guard lock(mutex_); + mutable_header()->clear(); + chunk_size_ = kDefaultChunkSize; + node_index_mask_bits_ = 0U; + node_index_mask_ = 0U; + node_cnt_per_chunk_ = 0U; + neighbor_size_ = 0U; + upper_neighbor_size_ = 0U; + if (upper_neighbor_index_) { + upper_neighbor_index_->cleanup(); + } + if (keys_map_) { + keys_map_->clear(); + } + node_chunks_.clear(); + upper_neighbor_chunks_.clear(); + filter_same_key_ = false; + get_vector_enabled_ = false; + broker_.reset(); + + return 0; +} + +int HnswRabitqStreamerEntity::update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) { + std::vector buffer(neighbor_size_); + NeighborsHeader *hd = reinterpret_cast(buffer.data()); + hd->neighbor_cnt = neighbors.size(); + size_t i = 0; + for (; i < neighbors.size(); ++i) { + hd->neighbors[i] = neighbors[i].first; + } + + auto loc = get_neighbor_chunk_loc(level, id); + size_t size = reinterpret_cast(&hd->neighbors[i]) - &buffer[0]; + size_t ret = loc.first->write(loc.second, hd, size); + if (ailego_unlikely(ret != size)) { + LOG_ERROR("Write neighbor header failed, ret=%zu", ret); + + return IndexError_Runtime; + } + + return 0; +} + +const Neighbors HnswRabitqStreamerEntity::get_neighbors(level_t level, + node_id_t id) const { + Chunk *chunk = nullptr; + size_t offset = 0UL; + size_t neighbor_size = neighbor_size_; + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + + sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, chunk_idx, + &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + chunk = node_chunks_[chunk_idx].get(); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + chunk = upper_neighbor_chunks_[p.first].get(); + offset = p.second; + neighbor_size = upper_neighbor_size_; + } + + ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); + IndexStorage::MemoryBlock neighbor_block; + size_t size = chunk->read(offset, neighbor_block, neighbor_size); + if (ailego_unlikely(size != neighbor_size)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", size); + return Neighbors(); + } + return Neighbors(std::move(neighbor_block)); +} + +//! Get vector data by key +const void *HnswRabitqStreamerEntity::get_vector(node_id_t id) const { + auto loc = get_vector_chunk_loc(id); + const void *vec = nullptr; + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t read_size = vector_size(); + + size_t ret = node_chunks_[loc.first]->read(loc.second, &vec, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%zu, read size=%zu, ret=%zu", + static_cast(loc.second), read_size, ret); + } + + return vec; +} + +int HnswRabitqStreamerEntity::get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const { + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t read_size = vector_size(); + + size_t ret = node_chunks_[loc.first]->read(loc.second, &vecs[i], read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%zu, read size=%zu, ret=%zu", + static_cast(loc.second), read_size, ret); + return IndexError_ReadData; + } + } + return 0; +} + +int HnswRabitqStreamerEntity::get_vector( + const node_id_t id, IndexStorage::MemoryBlock &block) const { + auto loc = get_vector_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t read_size = vector_size(); + + size_t ret = node_chunks_[loc.first]->read(loc.second, block, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%zu, read size=%zu, ret=%zu", + static_cast(loc.second), read_size, ret); + return IndexError_ReadData; + } + return 0; +} + +int HnswRabitqStreamerEntity::get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t read_size = vector_size(); + + size_t ret = + node_chunks_[loc.first]->read(loc.second, vec_blocks[i], read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%zu, read size=%zu, ret=%zu", + static_cast(loc.second), read_size, ret); + return IndexError_ReadData; + } + } + return 0; +} + +key_t HnswRabitqStreamerEntity::get_key(node_id_t id) const { + if (use_key_info_map_) { + auto loc = get_key_chunk_loc(id); + IndexStorage::MemoryBlock key_block; + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t ret = + node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return kInvalidKey; + } + + return *reinterpret_cast(key_block.data()); + } else { + return id; + } +} + +void HnswRabitqStreamerEntity::add_neighbor(level_t level, node_id_t id, + uint32_t size, + node_id_t neighbor_id) { + auto loc = get_neighbor_chunk_loc(level, id); + size_t offset = + loc.second + sizeof(NeighborsHeader) + size * sizeof(node_id_t); + ailego_assert_with(size < neighbor_cnt(level), "invalid neighbor size"); + ailego_assert_with(offset < loc.first->data_size(), "invalid chunk offset"); + size_t ret = loc.first->write(offset, &neighbor_id, sizeof(node_id_t)); + if (ailego_unlikely(ret != sizeof(node_id_t))) { + LOG_ERROR("Write neighbor id failed, ret=%zu", ret); + return; + } + + uint32_t neighbors = size + 1; + ret = loc.first->write(loc.second, &neighbors, sizeof(uint32_t)); + if (ailego_unlikely(ret != sizeof(uint32_t))) { + LOG_ERROR("Write neighbor cnt failed, ret=%zu", ret); + } + + return; +} + +int HnswRabitqStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) { + if (header_chunk->data_size() < header_size()) { + LOG_ERROR("Invalid header chunk size"); + return IndexError_InvalidFormat; + } + IndexStorage::MemoryBlock header_block; + size_t size = header_chunk->read(0UL, header_block, header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Read header chunk failed"); + return IndexError_ReadData; + } + *mutable_header() = + *reinterpret_cast(header_block.data()); + + int ret = check_hnsw_index(&header()); + if (ret != 0) { + broker_->close(); + return ret; + } + + node_chunks_.resize( + broker_->get_chunk_cnt(HnswRabitqChunkBroker::CHUNK_TYPE_NODE)); + for (auto seq = 0UL; seq < node_chunks_.size(); ++seq) { + node_chunks_[seq] = + broker_->get_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, seq); + if (!node_chunks_[seq]) { + LOG_ERROR("Missing hnsw streamer data chunk %zu th of %zu", seq, + node_chunks_.size()); + return IndexError_InvalidFormat; + } + } + + upper_neighbor_chunks_.resize( + broker_->get_chunk_cnt(HnswRabitqChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR)); + for (auto seq = 0UL; seq < upper_neighbor_chunks_.size(); ++seq) { + upper_neighbor_chunks_[seq] = broker_->get_chunk( + HnswRabitqChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, seq); + if (!upper_neighbor_chunks_[seq]) { + LOG_ERROR("Missing hnsw streamer index chunk %zu th of %zu", seq, + upper_neighbor_chunks_.size()); + return IndexError_InvalidFormat; + } + } + + return 0; +} + +int HnswRabitqStreamerEntity::open(IndexStorage::Pointer stg, + uint64_t max_index_size, bool check_crc) { + std::lock_guard lock(mutex_); + bool huge_page = stg->isHugePage(); + LOG_DEBUG("huge_page: %d", (int)huge_page); + int ret = init_chunk_params(max_index_size, huge_page); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("init_chunk_params failed for %s", IndexError::What(ret)); + return ret; + } + ret = broker_->open(std::move(stg), max_index_size_, chunk_size_, check_crc); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Open index failed for %s", IndexError::What(ret)); + return ret; + } + ret = upper_neighbor_index_->init(broker_, upper_neighbor_chunk_size_, + scaling_factor(), estimate_doc_capacity(), + kUpperHashMemoryInflateRatio); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Init neighbor hash map failed"); + return ret; + } + + //! init header + auto header_chunk = + broker_->get_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_HEADER, + HnswRabitqChunkBroker::kDefaultChunkSeqId); + if (!header_chunk) { // open empty index, create one + auto p = broker_->alloc_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_HEADER, + HnswRabitqChunkBroker::kDefaultChunkSeqId, + header_size()); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc header chunk failed"); + return p.first; + } + size_t size = p.second->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + return 0; + } + + //! Open an exist hnsw index + ret = init_chunks(header_chunk); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + //! total docs including features wrote in index but neighbors may not ready + node_id_t total_vecs = 0; + if (node_chunks_.size() > 0) { + size_t last_idx = node_chunks_.size() - 1; + auto last_chunk = node_chunks_[last_idx]; + if (last_chunk->data_size() % node_size()) { + LOG_WARN("The index may broken"); + return IndexError_InvalidFormat; + } + total_vecs = last_idx * node_cnt_per_chunk_ + + node_chunks_[last_idx]->data_size() / node_size(); + } + + LOG_INFO( + "Open index, l0NeighborCnt=%zu upperNeighborCnt=%zu " + "efConstruction=%zu curDocCnt=%u totalVecs=%u maxLevel=%u", + l0_neighbor_cnt(), upper_neighbor_cnt(), ef_construction(), doc_cnt(), + total_vecs, cur_max_level()); + //! try to correct the docCnt if index not fully flushed + if (doc_cnt() != total_vecs) { + LOG_WARN("Index closed abnormally, using totalVecs as curDocCnt"); + *mutable_doc_cnt() = total_vecs; + } + if (filter_same_key_ || get_vector_enabled_) { + if (use_key_info_map_) { + for (node_id_t id = 0U; id < doc_cnt(); ++id) { + if (get_key(id) == kInvalidKey) { + continue; + } + (*keys_map_)[get_key(id)] = id; + } + } + } + + stats_.set_loaded_count(doc_cnt()); + + return 0; +} + +int HnswRabitqStreamerEntity::close() { + LOG_DEBUG("close index"); + + std::lock_guard lock(mutex_); + flush_header(); + mutable_header()->reset(); + upper_neighbor_index_->cleanup(); + keys_map_->clear(); + header_.clear(); + node_chunks_.clear(); + upper_neighbor_chunks_.clear(); + + return broker_->close(); +} + +int HnswRabitqStreamerEntity::flush(uint64_t checkpoint) { + LOG_INFO("Flush index, curDocs=%zu", static_cast(doc_cnt())); + + std::lock_guard lock(mutex_); + flush_header(); + int ret = broker_->flush(checkpoint); + if (ret != 0) { + return ret; + } + + return 0; +} + +int HnswRabitqStreamerEntity::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("Dump index, curDocs=%zu", static_cast(doc_cnt())); + + //! sort by keys, to support get_vector by key in searcher + std::vector keys(doc_cnt()); + for (node_id_t i = 0; i < doc_cnt(); ++i) { + keys[i] = get_key(i); + } + + //! dump neighbors + auto get_level = [&](node_id_t id) { + auto it = upper_neighbor_index_->find(id); + if (it == upper_neighbor_index_->end()) { + return 0U; + }; + auto meta = reinterpret_cast(&it->second); + return meta->bits.level; + }; + auto ret = dump_segments(dumper, keys.data(), get_level); + if (ailego_unlikely(ret < 0)) { + return ret; + } + *stats_.mutable_dumped_size() += ret; + + return 0; +} + +int HnswRabitqStreamerEntity::check_hnsw_index(const HNSWHeader *hd) const { + if (l0_neighbor_cnt() != hd->l0_neighbor_cnt() || + upper_neighbor_cnt() != hd->upper_neighbor_cnt()) { + LOG_ERROR("Param neighbor cnt: %zu:%zu mismatch index previous %zu:%zu", + l0_neighbor_cnt(), upper_neighbor_cnt(), hd->l0_neighbor_cnt(), + hd->upper_neighbor_cnt()); + return IndexError_Mismatch; + } + if (vector_size() != hd->vector_size()) { + LOG_ERROR("vector size %zu mismatch index previous %zu", vector_size(), + hd->vector_size()); + return IndexError_Mismatch; + } + if (ef_construction() != hd->ef_construction()) { + LOG_WARN("Param efConstruction %zu mismatch index previous %zu", + ef_construction(), hd->ef_construction()); + } + if (scaling_factor() != hd->scaling_factor()) { + LOG_WARN("Param scalingFactor %zu mismatch index previous %zu", + scaling_factor(), hd->scaling_factor()); + return IndexError_Mismatch; + } + if (prune_cnt() != hd->neighbor_prune_cnt()) { + LOG_WARN("Param pruneCnt %zu mismatch index previous %zu", prune_cnt(), + hd->neighbor_prune_cnt()); + return IndexError_Mismatch; + } + if ((hd->entry_point() != kInvalidNodeId && + hd->entry_point() >= hd->doc_cnt()) || + (hd->entry_point() == kInvalidNodeId && hd->doc_cnt() > 0U)) { + LOG_WARN("Invalid entryPoint %zu, docCnt %zu", + static_cast(hd->entry_point()), + static_cast(hd->doc_cnt())); + return IndexError_InvalidFormat; + } + if (hd->entry_point() == kInvalidNodeId && + broker_->get_chunk_cnt(HnswRabitqChunkBroker::CHUNK_TYPE_NODE) > 0) { + LOG_WARN("The index is broken, maybe it haven't flush"); + return IndexError_InvalidFormat; + } + + return 0; +} + +int HnswRabitqStreamerEntity::add_vector(level_t level, key_t key, + const void *vec, node_id_t *id) { + Chunk::Pointer node_chunk; + size_t chunk_offset = -1UL; + + std::lock_guard lock(mutex_); + // duplicate check + if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) { + LOG_WARN("Try to add duplicate key, ignore it"); + return IndexError_Duplicate; + } + + node_id_t local_id = static_cast(doc_cnt()); + uint32_t chunk_index = node_chunks_.size() - 1U; + if (chunk_index == -1U || + (node_chunks_[chunk_index]->data_size() >= + node_cnt_per_chunk_ * node_size())) { // no space left and need to alloc + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + chunk_index++; + auto p = broker_->alloc_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, + chunk_index, chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + chunk_offset = 0UL; + node_chunks_.emplace_back(node_chunk); + } else { + node_chunk = node_chunks_[chunk_index]; + chunk_offset = node_chunk->data_size(); + } + + size_t size = node_chunk->write(chunk_offset, vec, vector_size()); + if (ailego_unlikely(size != vector_size())) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + //! level 0 neighbors is inited to zero by default + + int ret = add_upper_neighbor(level, local_id); + if (ret != 0) { + return ret; + } + + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + if (filter_same_key_ || get_vector_enabled_) { + if (use_key_info_map_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = local_id; + keys_map_lock_->unlock(); + } + } + + *mutable_doc_cnt() += 1; + broker_->mark_dirty(); + *id = local_id; + + return 0; +} + +int HnswRabitqStreamerEntity::add_vector_with_id(level_t level, node_id_t id, + const void *vec) { + Chunk::Pointer node_chunk; + size_t chunk_offset = -1UL; + key_t key = id; + + std::lock_guard lock(mutex_); + + // duplicate check + if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) { + LOG_WARN("Try to add duplicate key, ignore it"); + return IndexError_Duplicate; + } + + // set node_chunk & chunk_offset if succeed + auto func_get_node_chunk_and_offset = [&](node_id_t node_id) -> int { + uint32_t chunk_index = node_id >> node_index_mask_bits_; + ailego_assert_with(chunk_index <= node_chunks_.size(), "invalid chunk idx"); + // belongs to next chunk + if (chunk_index == node_chunks_.size()) { + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + auto p = broker_->alloc_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, + chunk_index, chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + node_chunks_.emplace_back(node_chunk); + } + + node_chunk = node_chunks_[chunk_index]; + chunk_offset = (node_id & node_index_mask_) * node_size(); + return 0; + }; + + for (size_t start_id = doc_cnt(); start_id < id; ++start_id) { + if (auto ret = func_get_node_chunk_and_offset(start_id); ret != 0) { + LOG_ERROR("func_get_node_chunk_and_offset failed"); + return ret; + } + size_t size = node_chunk->write(chunk_offset + vector_size(), &kInvalidKey, + sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (auto ret = func_get_node_chunk_and_offset(id); ret != 0) { + LOG_ERROR("func_get_node_chunk_and_offset failed"); + return ret; + } + + size_t size = node_chunk->write(chunk_offset, vec, vector_size()); + if (ailego_unlikely(size != vector_size())) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + + size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + //! level 0 neighbors is inited to zero by default + + int ret = add_upper_neighbor(level, id); + if (ret != 0) { + return ret; + } + + if (*mutable_doc_cnt() <= id) { + *mutable_doc_cnt() = id + 1; + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (filter_same_key_ || get_vector_enabled_) { + if (use_key_info_map_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = id; + keys_map_lock_->unlock(); + } + } + + broker_->mark_dirty(); + + return 0; +} + +void HnswRabitqStreamerEntity::update_ep_and_level(node_id_t ep, + level_t level) { + HnswRabitqEntity::update_ep_and_level(ep, level); + flush_header(); + + return; +} + +const HnswRabitqEntity::Pointer HnswRabitqStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswRabitqStreamerEntity get chunk failed in clone"); + return HnswRabitqEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswRabitqStreamerEntity get chunk failed in clone"); + return HnswRabitqEntity::Pointer(); + } + } + + HnswRabitqStreamerEntity *entity = + new (std::nothrow) HnswRabitqStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_, + keys_map_, use_key_info_map_, std::move(node_chunks), + std::move(upper_neighbor_chunks), broker_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswRabitqStreamerEntity new failed"); + } + return HnswRabitqEntity::Pointer(entity); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.h new file mode 100644 index 0000000..511fbad --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.h @@ -0,0 +1,559 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include "zvec/core/framework/index_framework.h" +#include "hnsw_rabitq_chunk.h" +#include "hnsw_rabitq_entity.h" +#include "hnsw_rabitq_index_hash.h" +#include "hnsw_rabitq_params.h" + +namespace zvec { +namespace core { + +//! HnswRabitqStreamerEntity manage vector data, pkey, and node's neighbors +class HnswRabitqStreamerEntity : public HnswRabitqEntity { + public: + //! Cleanup + //! return 0 on success, or errCode in failure + int cleanup() override; + + //! Make a copy of streamer entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + const HnswRabitqEntity::Pointer clone() const override; + + //! Get primary key of the node id + key_t get_key(node_id_t id) const override; + + //! Get vector feature data by key + const void *get_vector(node_id_t id) const override; + + //! Get vectors feature data by local ids + int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + + int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const override; + + //! Get the node id's neighbors on graph level + //! Note: the neighbors cannot be modified, using the following + //! method to get WritableNeighbors if want to + const Neighbors get_neighbors(level_t level, node_id_t id) const override; + + //! Add vector and key to hnsw entity, and local id will be saved in id + int add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) override; + + //! Add vector and id to hnsw entity + int add_vector_with_id(level_t level, node_id_t id, const void *vec) override; + + int update_neighbors(level_t level, node_id_t id, + const std::vector> + &neighbors) override; + + //! Append neighbor_id to node id neighbors on level + //! Notice: the caller must be ensure the neighbors not full + void add_neighbor(level_t level, node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + + //! Dump index by dumper + int dump(const IndexDumper::Pointer &dumper) override; + + void update_ep_and_level(node_id_t ep, level_t level) override; + + void set_use_key_info_map(bool use_id_map) { + use_key_info_map_ = use_id_map; + LOG_DEBUG("use_key_info_map_: %d", (int)use_key_info_map_); + } + + public: + //! Constructor + HnswRabitqStreamerEntity(IndexStreamer::Stats &stats); + + //! Destructor + ~HnswRabitqStreamerEntity(); + + //! Get vector feature data by key + const void *get_vector_by_key(key_t key) const override { + auto id = get_id(key); + return id == kInvalidNodeId ? nullptr : get_vector(id); + } + + int get_vector_by_key(const key_t key, + IndexStorage::MemoryBlock &block) const override { + auto id = get_id(key); + if (id != kInvalidNodeId) { + return get_vector(id, block); + } else { + return IndexError_InvalidArgument; + } + } + + //! Init entity + int init(size_t max_doc_cnt); + + //! Flush graph entity to disk + //! return 0 on success, or errCode in failure + int flush(uint64_t checkpoint); + + //! Open entity from storage + //! return 0 on success, or errCode in failure + int open(IndexStorage::Pointer stg, uint64_t max_index_size, bool check_crc); + + //! Close entity + //! return 0 on success, or errCode in failure + int close(); + + //! Set meta information from entity + int set_index_meta(const IndexMeta &meta) const { + return IndexHelper::SerializeToStorage(meta, broker_->storage().get()); + } + + //! Get meta information from entity + int get_index_meta(IndexMeta *meta) const { + return IndexHelper::DeserializeFromStorage(broker_->storage().get(), meta); + } + + //! Set params: chunk size + inline void set_chunk_size(size_t val) { + chunk_size_ = val; + } + + //! Set params + inline void set_filter_same_key(bool val) { + filter_same_key_ = val; + } + + //! Set params + inline void set_get_vector(bool val) { + get_vector_enabled_ = val; + } + + //! Get vector local id by key + inline node_id_t get_id(key_t key) const { + if (use_key_info_map_) { + keys_map_lock_->lock_shared(); + auto it = keys_map_->find(key); + keys_map_lock_->unlock_shared(); + return it == keys_map_->end() ? kInvalidNodeId : it->second; + } else { + return key; + } + } + + void print_key_map() const { + std::cout << "key map begins" << std::endl; + + auto iter = keys_map_->begin(); + while (iter != keys_map_->end()) { + std::cout << "key: " << iter->first << ", id: " << iter->second + << std::endl; + ; + iter++; + } + + std::cout << "key map ends" << std::endl; + } + + //! Get l0 neighbors size + inline size_t neighbors_size() const { + return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); + } + + //! Get neighbors size for level > 0 + inline size_t upper_neighbors_size() const { + return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); + } + + + private: + union UpperNeighborIndexMeta { + struct { + uint32_t level : 4; + uint32_t index : 28; // index is composite type: chunk idx, and the + // N th neighbors in chunk, they two composite + // the 28 bits location + } bits; + uint32_t data; + }; + + template + using HashMap = google::dense_hash_map>; + template + using HashMapPointer = std::shared_ptr>; + + template + using HashSet = google::dense_hash_set>; + template + using HashSetPointer = std::shared_ptr>; + + //! upper neighbor index hashmap + using NIHashMap = HnswIndexHashMap; + using NIHashMapPointer = std::shared_ptr; + + //! Private construct, only be called by clone method + HnswRabitqStreamerEntity( + IndexStreamer::Stats &stats, const HNSWHeader &hd, size_t chunk_size, + uint32_t node_index_mask_bits, uint32_t upper_neighbor_mask_bits, + bool filter_same_key, bool get_vector_enabled, + const NIHashMapPointer &upper_neighbor_index, + const std::shared_ptr &upper_neighbor_rw_mutex, + std::shared_ptr &keys_map_lock, + const HashMapPointer &keys_map, bool use_key_info_map, + std::vector &&node_chunks, + std::vector &&upper_neighbor_chunks, + const HnswRabitqChunkBroker::Pointer &broker) + : stats_(stats), + chunk_size_(chunk_size), + node_index_mask_bits_(node_index_mask_bits), + node_cnt_per_chunk_(1UL << node_index_mask_bits_), + node_index_mask_(node_cnt_per_chunk_ - 1), + upper_neighbor_mask_bits_(upper_neighbor_mask_bits), + upper_neighbor_mask_((1U << upper_neighbor_mask_bits_) - 1), + filter_same_key_(filter_same_key), + get_vector_enabled_(get_vector_enabled), + use_key_info_map_(use_key_info_map), + upper_neighbor_rw_mutex_(upper_neighbor_rw_mutex), + upper_neighbor_index_(upper_neighbor_index), + keys_map_lock_(keys_map_lock), + keys_map_(keys_map), + node_chunks_(std::move(node_chunks)), + upper_neighbor_chunks_(std::move(upper_neighbor_chunks)), + broker_(broker) { + *mutable_header() = hd; + + neighbor_size_ = neighbors_size(); + upper_neighbor_size_ = upper_neighbors_size(); + } + + //! Called only in searching procedure per context, so no need to lock + void sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE type, size_t idx, + std::vector *chunks) const { + if (ailego_likely(idx < chunks->size())) { + return; + } + for (size_t i = chunks->size(); i <= idx; ++i) { + auto chunk = broker_->get_chunk(type, i); + // the storage can ensure get chunk will success after the first get + ailego_assert_with(!!chunk, "get chunk failed"); + chunks->emplace_back(std::move(chunk)); + } + } + + //! return pair: chunk index + chunk offset + inline std::pair get_vector_chunk_loc( + node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size(); + + sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, chunk_idx, + &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + //! return pair: chunk index + chunk offset + inline std::pair get_key_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + + sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, chunk_idx, + &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + inline std::pair get_upper_neighbor_chunk_loc( + level_t level, node_id_t id) const { + // Shared lock: concurrent readers are fine, but must synchronize with + // add_upper_neighbor's exclusive lock to avoid data-race on + // slots_.size() inside HnswIndexHashMap (the emplace_back in alloc_slot + // is not atomic and concurrent find() may see a stale size value). + std::shared_lock lk(*upper_neighbor_rw_mutex_); + auto it = upper_neighbor_index_->find(id); + ailego_assert_abort(it != upper_neighbor_index_->end(), + "Get upper neighbor header failed"); + auto meta = reinterpret_cast(&it->second); + uint32_t chunk_idx = (meta->bits.index) >> upper_neighbor_mask_bits_; + uint32_t offset = + (((meta->bits.index) & upper_neighbor_mask_) + level - 1) * + upper_neighbor_size_; + sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, chunk_idx, + &upper_neighbor_chunks_); + ailego_assert_abort(chunk_idx < upper_neighbor_chunks_.size(), + "invalid chunk idx"); + ailego_assert_abort(offset < upper_neighbor_chunks_[chunk_idx]->data_size(), + "invalid chunk offset"); + return std::make_pair(chunk_idx, offset); + } + + //! return pair: chunk + chunk offset + inline std::pair get_neighbor_chunk_loc(level_t level, + node_id_t id) const { + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + + sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, chunk_idx, + &node_chunks_); + ailego_assert_abort(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_abort(offset < node_chunks_[chunk_idx]->data_size(), + "invalid chunk offset"); + return std::make_pair(node_chunks_[chunk_idx].get(), offset); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + return std::make_pair(upper_neighbor_chunks_[p.first].get(), p.second); + } + } + + //! Chunk hnsw index valid + int check_hnsw_index(const HNSWHeader *hd) const; + + size_t get_total_upper_neighbors_size(level_t level) const { + return level * upper_neighbor_size_; + } + + //! Add upper neighbor header and reserve space for upper neighbor + int add_upper_neighbor(level_t level, node_id_t id) { + if (level == 0) { + return 0; + } + // Exclusive lock: protects upper_neighbor_chunks_.emplace_back() and + // upper_neighbor_index_->insert() from racing with concurrent find() + // calls in get_upper_neighbor_chunk_loc(). + std::unique_lock lk(*upper_neighbor_rw_mutex_); + Chunk::Pointer chunk; + uint64_t chunk_offset = -1UL; + size_t neighbors_size = get_total_upper_neighbors_size(level); + uint64_t chunk_index = upper_neighbor_chunks_.size() - 1UL; + if (chunk_index == -1UL || + (upper_neighbor_chunks_[chunk_index]->padding_size() < + neighbors_size)) { // no space left and need to alloc + chunk_index++; + if (ailego_unlikely(upper_neighbor_chunks_.capacity() == + upper_neighbor_chunks_.size())) { + LOG_ERROR("add upper neighbor failed for no memory quota"); + return IndexError_IndexFull; + } + auto p = + broker_->alloc_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, + chunk_index, upper_neighbor_chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + chunk = p.second; + chunk_offset = 0UL; + upper_neighbor_chunks_.emplace_back(chunk); + } else { + chunk = upper_neighbor_chunks_[chunk_index]; + chunk_offset = chunk->data_size(); + } + ailego_assert_with((size_t)level < kMaxGraphLayers, "invalid level"); + ailego_assert_with(chunk_offset % upper_neighbor_size_ == 0, + "invalid offset"); + ailego_assert_with((chunk_offset / upper_neighbor_size_) < + (1U << upper_neighbor_mask_bits_), + "invalid offset"); + ailego_assert_with(chunk_index < (1U << (28 - upper_neighbor_mask_bits_)), + "invalid chunk index"); + UpperNeighborIndexMeta meta; + meta.bits.level = level; + meta.bits.index = (chunk_index << upper_neighbor_mask_bits_) | + (chunk_offset / upper_neighbor_size_); + size_t zero_start = chunk_offset; + chunk_offset += upper_neighbor_size_ * level; + + // IMPORTANT: order matters here. + // 1) resize so the chunk's data_size covers the new region. + // 2) zero-fill the new region: storage backends like BufferStorage do + // NOT zero on resize -- only metadata is updated, and the underlying + // page may contain stale content from a previously-evicted page. + // Without this step, NeighborsHeader::neighbor_cnt is garbage and + // select_entry_point()/search_neighbors() iterate over garbage + // node_ids, eventually triggering find()'s assertion in + // get_upper_neighbor_chunk_loc() at line 291. + // 3) ONLY THEN publish the entry to upper_neighbor_index_, so that any + // concurrent reader that finds this id already sees a properly + // zeroed upper-neighbor slot. + if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", (size_t)chunk_offset); + return IndexError_Runtime; + } + + // Use std::vector instead of a VLA: VLAs are a GNU extension and may + // produce different codegen / be rejected under clang/MSVC. + std::vector zeros(neighbors_size, 0); + if (ailego_unlikely(chunk->write(zero_start, zeros.data(), + neighbors_size) != neighbors_size)) { + LOG_ERROR("Chunk write zeros failed"); + return IndexError_Runtime; + } + + if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) { + LOG_ERROR("HashMap insert value failed"); + return IndexError_Runtime; + } + + return 0; + } + + size_t estimate_doc_capacity() const { + return node_chunks_.capacity() * node_cnt_per_chunk_; + } + + int init_chunk_params(size_t max_index_size, bool huge_page) { + node_cnt_per_chunk_ = std::max(1, chunk_size_ / node_size()); + //! align node cnt per chunk to pow of 2 + node_index_mask_bits_ = std::ceil(std::log2(node_cnt_per_chunk_)); + node_cnt_per_chunk_ = 1UL << node_index_mask_bits_; + if (huge_page) { + chunk_size_ = AlignHugePageSize(node_cnt_per_chunk_ * node_size()); + } else { + chunk_size_ = AlignPageSize(node_cnt_per_chunk_ * node_size()); + } + node_index_mask_ = node_cnt_per_chunk_ - 1; + + if (max_index_size == 0UL) { + max_index_size_ = chunk_size_ * kDefaultMaxChunkCnt; + } else { + max_index_size_ = max_index_size; + } + + //! To get a balanced upper neighbor chunk size. + //! If the upper chunk size is equal to node chunk size, it may waste + //! upper neighbor chunk space; if the upper neighbor chunk size is too + //! small, the will need large upper neighbor chunks index space. So to + //! get a balanced ratio be sqrt of the node/neighbor size ratio + float ratio = + std::sqrt(node_size() * scaling_factor() * 1.0f / upper_neighbor_size_); + if (huge_page) { + upper_neighbor_chunk_size_ = AlignHugePageSize( + std::max(get_total_upper_neighbors_size(kMaxGraphLayers), + static_cast(chunk_size_ / ratio))); + } else { + upper_neighbor_chunk_size_ = AlignPageSize( + std::max(get_total_upper_neighbors_size(kMaxGraphLayers), + static_cast(chunk_size_ / ratio))); + } + upper_neighbor_mask_bits_ = + std::ceil(std::log2(upper_neighbor_chunk_size_ / upper_neighbor_size_)); + upper_neighbor_mask_ = (1 << upper_neighbor_mask_bits_) - 1; + + size_t max_node_chunk_cnt = std::ceil(max_index_size_ / chunk_size_); + size_t max_upper_chunk_cnt = std::ceil( + (max_node_chunk_cnt * node_cnt_per_chunk_ * 1.0f / scaling_factor()) / + (upper_neighbor_chunk_size_ / upper_neighbor_size_)); + max_upper_chunk_cnt = + max_upper_chunk_cnt + std::ceil(max_upper_chunk_cnt / scaling_factor()); + + //! reserve space to avoid memmove in chunks vector emplace chunk, so + //! as to lock-free in reading chunk + node_chunks_.reserve(max_node_chunk_cnt); + upper_neighbor_chunks_.reserve(max_upper_chunk_cnt); + + LOG_DEBUG( + "Settings: nodeSize=%zu chunkSize=%u upperNeighborSize=%u " + "upperNeighborChunkSize=%u " + "nodeCntPerChunk=%u maxChunkCnt=%zu maxNeighborChunkCnt=%zu " + "maxIndexSize=%zu ratio=%.3f", + node_size(), chunk_size_, upper_neighbor_size_, + upper_neighbor_chunk_size_, node_cnt_per_chunk_, max_node_chunk_cnt, + max_upper_chunk_cnt, max_index_size_, ratio); + + return 0; + } + + //! Init node chunk and neighbor chunks + int init_chunks(const Chunk::Pointer &header_chunk); + + int flush_header(void) { + if (!broker_->dirty()) { + // do not need to flush + return 0; + } + auto header_chunk = + broker_->get_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_HEADER, + HnswRabitqChunkBroker::kDefaultChunkSeqId); + if (ailego_unlikely(!header_chunk)) { + LOG_ERROR("get header chunk failed"); + return IndexError_Runtime; + } + size_t size = header_chunk->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + + return 0; + } + + private: + HnswRabitqStreamerEntity(const HnswRabitqStreamerEntity &) = delete; + HnswRabitqStreamerEntity &operator=(const HnswRabitqStreamerEntity &) = + delete; + static constexpr uint64_t kUpperHashMemoryInflateRatio = 2.0f; + + private: + IndexStreamer::Stats &stats_; + HNSWHeader header_{}; + std::mutex mutex_{}; + size_t max_index_size_{0UL}; + uint32_t chunk_size_{kDefaultChunkSize}; + uint32_t upper_neighbor_chunk_size_{kDefaultChunkSize}; + uint32_t node_index_mask_bits_{0U}; + uint32_t node_cnt_per_chunk_{0U}; + uint32_t node_index_mask_{0U}; + uint32_t neighbor_size_{0U}; + uint32_t upper_neighbor_size_{0U}; + //! UpperNeighborIndex.index composite chunkIdx and offset in chunk by the + //! following mask + uint32_t upper_neighbor_mask_bits_{0U}; + uint32_t upper_neighbor_mask_{0U}; + bool filter_same_key_{false}; + bool get_vector_enabled_{false}; + bool use_key_info_map_{true}; + + // Shared via shared_ptr so that all cloned entities synchronize against + // the SAME mutex instance. A plain std::shared_mutex member would be + // independent per clone and provide no real protection for the shared + // upper_neighbor_index_ hashmap. + mutable std::shared_ptr upper_neighbor_rw_mutex_{}; + NIHashMapPointer upper_neighbor_index_{}; + + mutable std::shared_ptr keys_map_lock_{}; + HashMapPointer keys_map_{}; + + //! the chunks will be changed in searcher, so need mutable + //! data chunk include: vector, key, level 0 neighbors + mutable std::vector node_chunks_{}; + + //! upper neighbor chunk inlude: UpperNeighborHeader + (1~level) neighbors + mutable std::vector upper_neighbor_chunks_{}; + + HnswRabitqChunkBroker::Pointer broker_{}; // chunk broker +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_converter.cc b/src/core/algorithm/hnsw_rabitq/rabitq_converter.cc new file mode 100644 index 0000000..32c3204 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/rabitq_converter.cc @@ -0,0 +1,315 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "rabitq_converter.h" +#include +#include +#include +#include +#include +#include +#include "ailego/pattern/defer.h" +#include "algorithm/hnsw_rabitq/rabitq_reformer.h" +#include "zvec/core/framework/index_cluster.h" +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_factory.h" +#include "zvec/core/framework/index_features.h" +#include "zvec/core/framework/index_holder.h" +#include "zvec/core/framework/index_memory.h" +#include "zvec/core/framework/index_meta.h" +#include "rabitq_params.h" +#include "rabitq_utils.h" + +#ifdef _MSC_VER +#define strncasecmp _strnicmp +#endif + +namespace zvec { +namespace core { + +RabitqConverter::~RabitqConverter() { + this->cleanup(); +} + +int RabitqConverter::init(const IndexMeta &meta, const ailego::Params ¶ms) { + // Copy meta and ensure it has metric information + meta_ = meta; + dimension_ = meta.dimension(); + + if (meta_.metric_name().empty()) { + LOG_ERROR("Meta metric is empty"); + return IndexError_InvalidArgument; + } + + // Round up dimension to multiple of 64 + padded_dim_ = ((dimension_ + 63) / 64) * 64; + + // Get RaBitQ parameters with defaults + uint32_t total_bits = 0; + params.get(PARAM_RABITQ_TOTAL_BITS, &total_bits); + if (total_bits == 0) { + total_bits = kDefaultRabitqTotalBits; + } + if (total_bits < 1 || total_bits > 9) { + LOG_ERROR("Invalid total_bits: %zu, must be in [1, 9]", (size_t)total_bits); + return IndexError_InvalidArgument; + } + ex_bits_ = total_bits - 1; + + params.get(PARAM_RABITQ_NUM_CLUSTERS, &num_clusters_); + if (num_clusters_ == 0) { + num_clusters_ = kDefaultNumClusters; + } + + if (ex_bits_ > 8) { + LOG_ERROR("Invalid ex_bits: %zu, must be <= 8", ex_bits_); + return IndexError_InvalidArgument; + } + + if (meta.data_type() != IndexMeta::DataType::DT_FP32) { + LOG_ERROR("RaBitQ only supports FP32 data type"); + return IndexError_Unsupported; + } + params.get(PARAM_RABITQ_SAMPLE_COUNT, &sample_count_); + + std::string rotator_type_str; + params.get(PARAM_RABITQ_ROTATOR_TYPE, &rotator_type_str); + if (rotator_type_str.empty()) { + rotator_type_ = rabitqlib::RotatorType::FhtKacRotator; + } else if (strncasecmp(rotator_type_str.c_str(), "fht", 3) == 0) { + rotator_type_ = rabitqlib::RotatorType::FhtKacRotator; + } else if (strncasecmp(rotator_type_str.c_str(), "matrix", 6) == 0) { + rotator_type_ = rabitqlib::RotatorType::MatrixRotator; + } else { + LOG_ERROR("Invalid rotator_type: %s", rotator_type_str.c_str()); + return IndexError_InvalidArgument; + } + + // Create rotator + rotator_.reset( + rabitqlib::choose_rotator(dimension_, rotator_type_, padded_dim_)); + + LOG_INFO( + "RabitqConverter initialized: dim=%zu, padded_dim=%zu, " + "num_clusters=%zu, ex_bits=%zu, rotator_type=%d[%s] sample_count[%zu]", + dimension_, padded_dim_, num_clusters_, ex_bits_, (int)rotator_type_, + rotator_type_str.c_str(), sample_count_); + + return 0; +} + +int RabitqConverter::cleanup() { + centroids_.clear(); + rotated_centroids_.clear(); + result_holder_.reset(); + rotator_.reset(); + return 0; +} + +int RabitqConverter::train(IndexHolder::Pointer holder) { + if (!holder) { + LOG_ERROR("Null holder for training"); + return IndexError_InvalidArgument; + } + + ailego::ElapsedTime timer; + + size_t vector_count = holder->count(); + if (vector_count == 0) { + LOG_ERROR("No vectors for training"); + return IndexError_InvalidArgument; + } + + // do sampling from all data + size_t sample_count = vector_count; + if (sample_count_ > 0) { + sample_count = std::min(sample_count_, vector_count); + } + LOG_INFO("Training with %zu vectors from %zu of holder", sample_count, + vector_count); + auto sampler = std::make_shared>( + meta_, sample_count); + auto iter = holder->create_iterator(); + if (!iter) { + LOG_ERROR("Create iterator error"); + return IndexError_Runtime; + } + for (; iter->is_valid(); iter->next()) { + sampler->emplace(iter->data()); + } + + // Holder is not needed, cleanup it. + holder.reset(); + + if (sampler->count() == 0) { + LOG_ERROR("Load training data error"); + return IndexError_InvalidLength; + } + + + // Create KmeansCluster for training centroids + auto cluster = IndexFactory::CreateCluster("OptKmeansCluster"); + if (!cluster) { + LOG_ERROR("Failed to create OptKmeansCluster"); + return IndexError_NoExist; + } + + // Initialize cluster + LOG_INFO( + "Initializing KmeansCluster with meta: dim=%u, data_type=%d, metric=%s", + meta_.dimension(), (int)meta_.data_type(), meta_.metric_name().c_str()); + ailego::Params cluster_params; + int ret = cluster->init(meta_, cluster_params); + if (ret != 0) { + LOG_ERROR("Failed to initialize KmeansCluster: %d", ret); + return ret; + } + + ret = cluster->mount(sampler); + if (ret != 0) { + LOG_ERROR("Failed to mount training data: %d", ret); + return ret; + } + cluster->suggest(num_clusters_); + + // Perform clustering + IndexCluster::CentroidList cents; + // TODO: support specify threads with argument + auto threads = std::make_shared(0, false); + ret = cluster->cluster(threads, cents); + if (ret != 0) { + LOG_ERROR("Failed to perform clustering: %d", ret); + return ret; + } + + if (cents.size() != num_clusters_) { + LOG_WARN("Expected %zu clusters, got %zu", num_clusters_, cents.size()); + num_clusters_ = cents.size(); + } + // Extract original centroids (for LinearSeeker query) + centroids_.resize(num_clusters_ * dimension_); + // Extract rotated centroids (for quantization) + rotated_centroids_.resize(num_clusters_ * padded_dim_); + for (uint32_t i = 0; i < num_clusters_; ++i) { + const float *cent_data = static_cast(cents[i].feature()); + // Save original centroids + std::memcpy(¢roids_[i * dimension_], cent_data, + dimension_ * sizeof(float)); + // Save rotated centroids + this->rotator_->rotate(cent_data, &rotated_centroids_[i * padded_dim_]); + } + + stats_.set_trained_count(sampler->count()); + stats_.set_trained_costtime(timer.milli_seconds()); + + LOG_INFO("Training completed: %zu centroids, cost %zu ms", num_clusters_, + static_cast(timer.milli_seconds())); + + return 0; +} + + +int RabitqConverter::transform(IndexHolder::Pointer holder) { + if (!holder) { + LOG_ERROR("Null holder for transformation"); + return IndexError_InvalidArgument; + } + + if (rotated_centroids_.empty()) { + LOG_ERROR("Centroids not trained yet"); + return IndexError_NoReady; + } + + LOG_ERROR("Not implemented"); + return IndexError_NotImplemented; +} + +int RabitqConverter::dump(const IndexDumper::Pointer &dumper) { + if (!dumper) { + LOG_ERROR("Null dumper"); + return IndexError_InvalidArgument; + } + + if (rotated_centroids_.empty() || centroids_.empty()) { + LOG_ERROR("No centroids to dump"); + return IndexError_NoReady; + } + + ailego::ElapsedTime timer; + size_t dumped_size = 0; + + int ret = dump_rabitq_centroids( + dumper, dimension_, padded_dim_, ex_bits_, num_clusters_, rotator_type_, + rotated_centroids_, centroids_, rotator_, &dumped_size); + if (ret != 0) { + return ret; + } + + stats_.set_dumped_size(dumped_size); + stats_.set_dumped_costtime(timer.milli_seconds()); + + LOG_INFO("Dump completed: %zu bytes, cost %zu ms", stats_.dumped_size(), + static_cast(timer.milli_seconds())); + return 0; +} + +int RabitqConverter::to_reformer(IndexReformer::Pointer *reformer) { + auto memory_dumper = IndexFactory::CreateDumper("MemoryDumper"); + memory_dumper->init(ailego::Params()); + std::string file_id = ailego::StringHelper::Concat( + "rabitq_converter_", ailego::Monotime::MilliSeconds(), rand()); + int ret = memory_dumper->create(file_id); + if (ret != 0) { + LOG_ERROR("Failed to create memory dumper: %d", ret); + return ret; + } + // Release memory + AILEGO_DEFER([&file_id]() { IndexMemory::Instance()->remove(file_id); }); + ret = this->dump(memory_dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump RabitqConverter: %d", ret); + return ret; + } + ret = memory_dumper->close(); + if (ret != 0) { + LOG_ERROR("Failed to close memory dumper: %d", ret); + return ret; + } + + auto res = std::make_shared(); + ailego::Params reformer_params; + reformer_params.set(PARAM_RABITQ_METRIC_NAME, meta_.metric_name()); + ret = res->init(reformer_params); + if (ret != 0) { + LOG_ERROR("Failed to initialize RabitqReformer: %d", ret); + return ret; + } + auto memory_storage = IndexFactory::CreateStorage("MemoryReadStorage"); + ret = memory_storage->open(file_id, false); + if (ret != 0) { + LOG_ERROR("Failed to open memory storage: %d", ret); + return ret; + } + ret = res->load(memory_storage); + if (ret != 0) { + LOG_ERROR("Failed to load RabitqReformer: %d", ret); + return ret; + } + *reformer = std::move(res); + return 0; +} + + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_converter.h b/src/core/algorithm/hnsw_rabitq/rabitq_converter.h new file mode 100644 index 0000000..d7e52a7 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/rabitq_converter.h @@ -0,0 +1,101 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "zvec/core/framework/index_cluster.h" +#include "zvec/core/framework/index_converter.h" +#include "zvec/core/framework/index_reformer.h" +#include "zvec/core/framework/index_threads.h" +#include "rabitq_params.h" + +namespace zvec { +namespace core { + +class RabitqReformer; + +/*! RaBitQ Converter + * Trains KMeans centroids and quantizes vectors using RaBitQ + */ +class RabitqConverter : public IndexConverter { + public: + //! Constructor + RabitqConverter() = default; + + //! Destructor + ~RabitqConverter() override; + + //! Initialize Converter + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup Converter + int cleanup(void) override; + + //! Train the data - perform KMeans clustering + int train(IndexHolder::Pointer holder) override; + + //! Transform the data - quantize vectors using RaBitQ + int transform(IndexHolder::Pointer holder) override; + + //! Dump centroids and config into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve a holder as result + IndexHolder::Pointer result(void) const override { + return result_holder_; + } + + //! Retrieve Index Meta + const IndexMeta &meta(void) const override { + return meta_; + } + + int to_reformer(IndexReformer::Pointer *reformer) override; + + private: + static inline size_t AlignSize(size_t size) { + return (size + 0x1F) & (~0x1F); + } + + private: + IndexMeta meta_; + IndexHolder::Pointer result_holder_; + Stats stats_; + size_t sample_count_{0}; + + // RaBitQ parameters + size_t num_clusters_{0}; + size_t ex_bits_{0}; + size_t dimension_{0}; + size_t padded_dim_{0}; + + // Original centroids: num_clusters * dimension (for LinearSeeker query) + std::vector centroids_; + // Rotated centroids: num_clusters * padded_dim (for quantization) + std::vector rotated_centroids_; + + // Rotator for vector transformation + rabitqlib::RotatorType rotator_type_{rabitqlib::RotatorType::FhtKacRotator}; + std::unique_ptr> rotator_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_params.h b/src/core/algorithm/hnsw_rabitq/rabitq_params.h new file mode 100644 index 0000000..c30669f --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/rabitq_params.h @@ -0,0 +1,49 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +// Local metric type enum that mirrors rabitqlib::MetricType, +// without exposing rabitqlib headers to consumers of this file. +enum class RabitqMetricType { + kL2 = 0, + kIP = 1, +}; + +// RaBitQ Converter parameters +static const std::string PARAM_RABITQ_NUM_CLUSTERS( + "proxima.rabitq.num_clusters"); +static const std::string PARAM_RABITQ_TOTAL_BITS("proxima.rabitq.total_bits"); +static const std::string PARAM_RABITQ_METRIC_NAME("proxima.rabitq.metric_name"); +static const std::string PARAM_RABITQ_ROTATOR_TYPE( + "proxima.rabitq.rotator.type"); +static const std::string PARAM_RABITQ_SAMPLE_COUNT( + "proxima.rabitq.sample_count"); + +// Default values +constexpr size_t kDefaultNumClusters = 16; +// 4-bit, 5-bit, and 7-bit quantization typically achieve 90%, 95%, and 99% +// recall, respectively—without accessing raw vectors for reranking +constexpr size_t kDefaultRabitqTotalBits = 7; + +constexpr int kMinRabitqDimSize = 64; +constexpr int kMaxRabitqDimSize = 4095; + + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_reformer.cc b/src/core/algorithm/hnsw_rabitq/rabitq_reformer.cc new file mode 100644 index 0000000..4a4c0a5 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/rabitq_reformer.cc @@ -0,0 +1,471 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "rabitq_reformer.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include "core/algorithm/cluster/linear_seeker.h" +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_factory.h" +#include "zvec/core/framework/index_features.h" +#include "zvec/core/framework/index_meta.h" +#include "zvec/core/framework/index_storage.h" +#include "hnsw_rabitq_query_entity.h" +#include "rabitq_converter.h" +#include "rabitq_utils.h" + +namespace zvec { +namespace core { + +// All rabitqlib types are confined to this translation unit via pimpl. +struct RabitqReformer::Impl { + // RaBitQ parameters + size_t num_clusters{0}; + size_t ex_bits{0}; + size_t dimension{0}; + size_t padded_dim{0}; + size_t size_bin_data{0}; + size_t size_ex_data{0}; + bool loaded{false}; + + // Original centroids: num_clusters * dimension (for LinearSeeker query) + std::vector centroids; + // Rotated centroids: num_clusters * padded_dim (for quantization) + std::vector rotated_centroids; + + rabitqlib::RotatorType rotator_type{rabitqlib::RotatorType::FhtKacRotator}; + std::unique_ptr> rotator; + rabitqlib::quant::RabitqConfig query_config; + rabitqlib::quant::RabitqConfig config; + rabitqlib::MetricType metric_type{rabitqlib::METRIC_L2}; + + LinearSeeker::Pointer centroid_seeker; + CoherentIndexFeatures::Pointer centroid_features; + + // Translate local enum to rabitqlib enum (used only inside this .cc). + static rabitqlib::MetricType to_rabitq(RabitqMetricType m) { + return m == RabitqMetricType::kIP ? rabitqlib::METRIC_IP + : rabitqlib::METRIC_L2; + } + + // Translate rabitqlib enum to local enum. + static RabitqMetricType from_rabitq(rabitqlib::MetricType m) { + return m == rabitqlib::METRIC_IP ? RabitqMetricType::kIP + : RabitqMetricType::kL2; + } + + int quantize_vector(const float *raw_vector, uint32_t cluster_id, + std::string *quantized_data) const; +}; + +RabitqReformer::RabitqReformer() : impl_(std::make_unique()) {} + +RabitqReformer::~RabitqReformer() { + this->cleanup(); +} + +size_t RabitqReformer::num_clusters() const { + return impl_->num_clusters; +} + +size_t RabitqReformer::ex_bits() const { + return impl_->ex_bits; +} + +RabitqMetricType RabitqReformer::rabitq_metric_type() const { + return Impl::from_rabitq(impl_->metric_type); +} + +int RabitqReformer::init(const ailego::Params ¶ms) { + std::string metric_name = params.get_as_string(PARAM_RABITQ_METRIC_NAME); + if (metric_name == "SquaredEuclidean") { + impl_->metric_type = rabitqlib::METRIC_L2; + } else if (metric_name == "InnerProduct") { + impl_->metric_type = rabitqlib::METRIC_IP; + } else if (metric_name == "Cosine") { + impl_->metric_type = rabitqlib::METRIC_IP; + } else { + LOG_ERROR("Unsupported metric name: %s", metric_name.c_str()); + return IndexError_InvalidArgument; + } + LOG_DEBUG("Rabitq reformer init done. metric_name=%s metric_type=%d", + metric_name.c_str(), static_cast(impl_->metric_type)); + return 0; +} + +int RabitqReformer::cleanup() { + impl_->centroids.clear(); + impl_->rotated_centroids.clear(); + impl_->centroid_seeker.reset(); + impl_->centroid_features.reset(); + impl_->loaded = false; + impl_->rotator.reset(); + return 0; +} + +int RabitqReformer::unload() { + return this->cleanup(); +} + +int RabitqReformer::load(IndexStorage::Pointer storage) { + if (!storage) { + LOG_ERROR("Invalid storage for load"); + return IndexError_InvalidArgument; + } + + auto segment = storage->get(RABITQ_CONVERTER_SEG_ID); + if (!segment) { + LOG_ERROR("Failed to get segment %s", RABITQ_CONVERTER_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + + size_t offset = 0; + RabitqConverterHeader header; + IndexStorage::MemoryBlock block; + size_t size = segment->read(offset, block, sizeof(header)); + if (size != sizeof(header)) { + LOG_ERROR("Failed to read header"); + return IndexError_InvalidFormat; + } + memcpy(static_cast(&header), block.data(), sizeof(header)); + impl_->dimension = header.dim; + impl_->padded_dim = header.padded_dim; + impl_->ex_bits = header.ex_bits; + impl_->num_clusters = header.num_clusters; + impl_->rotator_type = + static_cast(header.rotator_type); + offset += sizeof(header); + + // Read rotated centroids + size_t rotated_centroids_size = + sizeof(float) * header.num_clusters * header.padded_dim; + size = segment->read(offset, block, rotated_centroids_size); + if (size != rotated_centroids_size) { + LOG_ERROR("Failed to read rotated centroids"); + return IndexError_InvalidFormat; + } + impl_->rotated_centroids.resize(header.num_clusters * header.padded_dim); + memcpy(impl_->rotated_centroids.data(), block.data(), rotated_centroids_size); + offset += size; + + // Read original centroids (for LinearSeeker query) + size_t centroids_size = sizeof(float) * header.num_clusters * header.dim; + size = segment->read(offset, block, centroids_size); + if (size != centroids_size) { + LOG_ERROR("Failed to read centroids"); + return IndexError_InvalidFormat; + } + impl_->centroids.resize(header.num_clusters * header.dim); + memcpy(impl_->centroids.data(), block.data(), centroids_size); + offset += size; + + // Read rotator + size_t rotator_size = header.rotator_size; + size = segment->read(offset, block, rotator_size); + if (size != rotator_size) { + LOG_ERROR("Failed to read rotator"); + return IndexError_InvalidFormat; + } + impl_->rotator.reset(rabitqlib::choose_rotator( + impl_->dimension, impl_->rotator_type, impl_->padded_dim)); + impl_->rotator->load(reinterpret_cast(block.data())); + offset += size; + + impl_->query_config = rabitqlib::quant::faster_config( + impl_->padded_dim, rabitqlib::SplitSingleQuery::kNumBits); + impl_->config = + rabitqlib::quant::faster_config(impl_->padded_dim, impl_->ex_bits + 1); + + impl_->size_bin_data = + rabitqlib::BinDataMap::data_bytes(impl_->padded_dim); + impl_->size_ex_data = rabitqlib::ExDataMap::data_bytes( + impl_->padded_dim, impl_->ex_bits); + + // Initialize LinearSeeker for centroid search + IndexMeta centroid_meta; + centroid_meta.set_data_type(IndexMeta::DataType::DT_FP32); + centroid_meta.set_dimension(static_cast(impl_->dimension)); + // Note: + // 1. spherical kmeans is used for InnerProduct and Cosine, so centroids are + // normalized. + // 2. for Cosine metric, `transform_to_entity` input is normalized, need to + // use InnerProduct metric as Cosine metric requires extra dimension which is + // unsuitable for centroids. + centroid_meta.set_metric(impl_->metric_type == rabitqlib::METRIC_L2 + ? "SquaredEuclidean" + : "InnerProduct", + 0, ailego::Params()); + + impl_->centroid_features = std::make_shared(); + impl_->centroid_features->mount(centroid_meta, impl_->centroids.data(), + impl_->centroids.size() * sizeof(float)); + + impl_->centroid_seeker = std::make_shared(); + int ret = impl_->centroid_seeker->init(centroid_meta); + if (ret != 0) { + LOG_ERROR("Failed to init centroid seeker. ret[%d]", ret); + return ret; + } + ret = impl_->centroid_seeker->mount(impl_->centroid_features); + if (ret != 0) { + LOG_ERROR("Failed to mount centroid features. ret[%d]", ret); + return ret; + } + + LOG_INFO( + "Rabitq reformer load done. dimension=%zu, padded_dim=%zu, " + "ex_bits=%zu, num_clusters=%zu, size_bin_data=%zu, size_ex_data=%zu " + "rotator_type=%d", + impl_->dimension, impl_->padded_dim, impl_->ex_bits, impl_->num_clusters, + impl_->size_bin_data, impl_->size_ex_data, (int)impl_->rotator_type); + impl_->loaded = true; + return 0; +} + +int RabitqReformer::convert(const void *record, const IndexQueryMeta &rmeta, + std::string *out, IndexQueryMeta *ometa) const { + if (!impl_->loaded) { + LOG_ERROR("Centroids not loaded yet"); + return IndexError_NoReady; + } + + if (!record || !out) { + LOG_ERROR("Invalid arguments for convert"); + return IndexError_InvalidArgument; + } + + // input may be transformed, require rmeta.dimension >= dimension + if (rmeta.dimension() < impl_->dimension || + rmeta.data_type() != IndexMeta::DataType::DT_FP32) { + LOG_ERROR("Invalid record meta: dimension=%zu, data_type=%d", + static_cast(rmeta.dimension()), (int)rmeta.data_type()); + return IndexError_InvalidArgument; + } + + // Find nearest centroid using LinearSeeker + Seeker::Document doc; + int ret = impl_->centroid_seeker->seek( + record, impl_->dimension * sizeof(float), &doc); + if (ret != 0) { + LOG_ERROR("Failed to seek centroid. ret[%d]", ret); + return ret; + } + uint32_t cluster_id = doc.index; + + const float *vector = static_cast(record); + ret = impl_->quantize_vector(vector, cluster_id, out); + if (ret != 0) { + LOG_ERROR("Failed to quantize vector"); + return ret; + } + + ometa->set_meta(IndexMeta::DataType::DT_INT8, (uint32_t)out->size()); + return 0; +} + +int RabitqReformer::transform(const void *, const IndexQueryMeta &, + std::string *, IndexQueryMeta *) const { + return IndexError_NotImplemented; +} + +int RabitqReformer::transform_to_entity(const void *query, + HnswRabitqQueryEntity *entity) const { + if (!impl_->loaded) { + LOG_ERROR("Centroids not loaded yet"); + return IndexError_NoReady; + } + + if (!query) { + LOG_ERROR("Invalid arguments for transform"); + return IndexError_InvalidArgument; + } + + const float *query_vector = static_cast(query); + + // Apply rotator + entity->rotated_query.resize(impl_->padded_dim); + impl_->rotator->rotate(query_vector, entity->rotated_query.data()); + + // Quantize query to 4-bit representation + entity->query_wrapper = std::make_unique>( + entity->rotated_query.data(), impl_->padded_dim, impl_->ex_bits, + impl_->query_config, impl_->metric_type); + + // Preprocess - get the distance from query to all centroids + entity->q_to_centroids.resize(impl_->num_clusters); + + if (impl_->metric_type == rabitqlib::METRIC_L2) { + for (size_t i = 0; i < impl_->num_clusters; i++) { + entity->q_to_centroids[i] = std::sqrt(rabitqlib::euclidean_sqr( + entity->rotated_query.data(), + impl_->rotated_centroids.data() + (i * impl_->padded_dim), + impl_->padded_dim)); + } + } else if (impl_->metric_type == rabitqlib::METRIC_IP) { + entity->q_to_centroids.resize(impl_->num_clusters * 2); + // first half as g_add, second half as g_error + for (size_t i = 0; i < impl_->num_clusters; i++) { + entity->q_to_centroids[i] = rabitqlib::dot_product( + entity->rotated_query.data(), + impl_->rotated_centroids.data() + (i * impl_->padded_dim), + impl_->padded_dim); + entity->q_to_centroids[i + impl_->num_clusters] = + std::sqrt(rabitqlib::euclidean_sqr( + entity->rotated_query.data(), + impl_->rotated_centroids.data() + (i * impl_->padded_dim), + impl_->padded_dim)); + } + } + + return 0; +} + +int RabitqReformer::Impl::quantize_vector(const float *raw_vector, + uint32_t cluster_id, + std::string *quantized_data) const { + std::vector rotated_data(padded_dim); + rotator->rotate(raw_vector, rotated_data.data()); + + // quantized format: cluster_id + bin_data + ex_data + quantized_data->resize(sizeof(cluster_id) + size_bin_data + size_ex_data); + memcpy(&(*quantized_data)[0], &cluster_id, sizeof(cluster_id)); + int bin_data_offset = sizeof(cluster_id); + int ex_data_offset = bin_data_offset + size_bin_data; + rabitqlib::quant::quantize_split_single( + rotated_data.data(), rotated_centroids.data() + (cluster_id * padded_dim), + padded_dim, ex_bits, &(*quantized_data)[bin_data_offset], + &(*quantized_data)[ex_data_offset], metric_type, config); + + return 0; +} + +int RabitqReformer::dump(const IndexDumper::Pointer &dumper) { + if (!dumper) { + LOG_ERROR("Null dumper"); + return IndexError_InvalidArgument; + } + + if (!impl_->loaded || impl_->rotated_centroids.empty() || + impl_->centroids.empty()) { + LOG_ERROR("No centroids to dump"); + return IndexError_NoReady; + } + + size_t dumped_size = 0; + int ret = dump_rabitq_centroids( + dumper, impl_->dimension, impl_->padded_dim, impl_->ex_bits, + impl_->num_clusters, impl_->rotator_type, impl_->rotated_centroids, + impl_->centroids, impl_->rotator, &dumped_size); + if (ret != 0) { + return ret; + } + + LOG_INFO("RabitqReformer dump completed: %zu bytes", dumped_size); + return 0; +} + +int RabitqReformer::dump(const IndexStorage::Pointer &storage) { + if (!storage) { + LOG_ERROR("Null storage"); + return IndexError_InvalidArgument; + } + + if (!impl_->loaded || impl_->rotated_centroids.empty() || + impl_->centroids.empty()) { + LOG_ERROR("No centroids to dump"); + return IndexError_NoReady; + } + + auto align_size = [](size_t size) -> size_t { + return (size + 0x1F) & (~0x1F); + }; + + size_t header_size = sizeof(RabitqConverterHeader); + size_t rotated_centroids_size = + impl_->rotated_centroids.size() * sizeof(float); + size_t centroids_size = impl_->centroids.size() * sizeof(float); + size_t rotator_size = impl_->rotator->dump_bytes(); + size_t data_size = + header_size + rotated_centroids_size + centroids_size + rotator_size; + size_t total_size = align_size(data_size); + + int ret = storage->append(RABITQ_CONVERTER_SEG_ID, total_size); + if (ret != 0) { + LOG_ERROR("Failed to append segment %s, ret=%d", + RABITQ_CONVERTER_SEG_ID.c_str(), ret); + return ret; + } + + auto segment = storage->get(RABITQ_CONVERTER_SEG_ID); + if (!segment) { + LOG_ERROR("Failed to get segment %s", RABITQ_CONVERTER_SEG_ID.c_str()); + return IndexError_ReadData; + } + + size_t offset = 0; + + RabitqConverterHeader header; + header.dim = static_cast(impl_->dimension); + header.padded_dim = static_cast(impl_->padded_dim); + header.num_clusters = static_cast(impl_->num_clusters); + header.ex_bits = static_cast(impl_->ex_bits); + header.rotator_type = static_cast(impl_->rotator_type); + header.rotator_size = static_cast(rotator_size); + size_t written = segment->write(offset, &header, header_size); + if (written != header_size) { + LOG_ERROR("Failed to write header: written=%zu, expected=%zu", written, + header_size); + return IndexError_WriteData; + } + offset += header_size; + + written = segment->write(offset, impl_->rotated_centroids.data(), + rotated_centroids_size); + if (written != rotated_centroids_size) { + LOG_ERROR("Failed to write rotated centroids: written=%zu, expected=%zu", + written, rotated_centroids_size); + return IndexError_WriteData; + } + offset += rotated_centroids_size; + + written = segment->write(offset, impl_->centroids.data(), centroids_size); + if (written != centroids_size) { + LOG_ERROR("Failed to write centroids: written=%zu, expected=%zu", written, + centroids_size); + return IndexError_WriteData; + } + offset += centroids_size; + + std::vector buffer(rotator_size); + impl_->rotator->save(buffer.data()); + written = segment->write(offset, buffer.data(), rotator_size); + if (written != rotator_size) { + LOG_ERROR("Failed to write rotator data: written=%zu, expected=%zu", + written, rotator_size); + return IndexError_WriteData; + } + + LOG_INFO("RabitqReformer dump to storage completed: %zu bytes", data_size); + return 0; +} + + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_reformer.h b/src/core/algorithm/hnsw_rabitq/rabitq_reformer.h new file mode 100644 index 0000000..e3ee9d2 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/rabitq_reformer.h @@ -0,0 +1,73 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +#pragma once + +#include +#include "zvec/core/framework/index_dumper.h" +#include "zvec/core/framework/index_reformer.h" +#include "zvec/core/framework/index_storage.h" +#include "rabitq_params.h" + +namespace zvec { +namespace core { + +struct HnswRabitqQueryEntity; + +/*! RaBitQ Reformer + * Loads centroids and performs query transformation and vector quantization. + * + * All rabitqlib types are hidden behind a pimpl to avoid leaking rabitqlib + * headers to consumers of this class. + */ +class RabitqReformer : public IndexReformer { + public: + typedef std::shared_ptr Pointer; + + RabitqReformer(); + ~RabitqReformer() override; + + // Non-copyable + RabitqReformer(const RabitqReformer &) = delete; + RabitqReformer &operator=(const RabitqReformer &) = delete; + + int init(const ailego::Params ¶ms) override; + int cleanup(void) override; + int load(IndexStorage::Pointer storage) override; + int unload(void) override; + + // transform() is not implemented for RabitqReformer; use transform_to_entity. + int transform(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const override; + + int convert(const void *record, const IndexQueryMeta &rmeta, std::string *out, + IndexQueryMeta *ometa) const override; + + int dump(const IndexDumper::Pointer &dumper); + int dump(const IndexStorage::Pointer &storage); + + int transform_to_entity(const void *query, + HnswRabitqQueryEntity *entity) const; + + size_t num_clusters() const; + size_t ex_bits() const; + RabitqMetricType rabitq_metric_type() const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_utils.cc b/src/core/algorithm/hnsw_rabitq/rabitq_utils.cc new file mode 100644 index 0000000..9945ad7 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/rabitq_utils.cc @@ -0,0 +1,115 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "rabitq_utils.h" +#include +#include +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_logger.h" + +namespace zvec { +namespace core { + +int dump_rabitq_centroids( + const IndexDumper::Pointer &dumper, size_t dimension, size_t padded_dim, + size_t ex_bits, size_t num_clusters, rabitqlib::RotatorType rotator_type, + const std::vector &rotated_centroids, + const std::vector ¢roids, + const std::unique_ptr> &rotator, + size_t *out_dumped_size) { + auto align_size = [](size_t size) -> size_t { + return (size + 0x1F) & (~0x1F); + }; + + uint32_t crc = 0; + size_t dumped_size = 0; + + // Write header + RabitqConverterHeader header; + header.dim = static_cast(dimension); + header.padded_dim = static_cast(padded_dim); + header.num_clusters = static_cast(num_clusters); + header.ex_bits = static_cast(ex_bits); + header.rotator_type = static_cast(rotator_type); + header.rotator_size = static_cast(rotator->dump_bytes()); + size_t size = dumper->write(&header, sizeof(header)); + if (size != sizeof(header)) { + LOG_ERROR("Failed to write header: written=%zu, expected=%zu", size, + sizeof(header)); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(&header, sizeof(header), crc); + dumped_size += size; + + // Write rotated centroids + size = dumper->write(rotated_centroids.data(), + rotated_centroids.size() * sizeof(float)); + if (size != rotated_centroids.size() * sizeof(float)) { + LOG_ERROR("Failed to write rotated centroids: written=%zu, expected=%zu", + size, rotated_centroids.size() * sizeof(float)); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(rotated_centroids.data(), + rotated_centroids.size() * sizeof(float), crc); + dumped_size += size; + + // Write original centroids + size = dumper->write(centroids.data(), centroids.size() * sizeof(float)); + if (size != centroids.size() * sizeof(float)) { + LOG_ERROR("Failed to write centroids: written=%zu, expected=%zu", size, + centroids.size() * sizeof(float)); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(centroids.data(), centroids.size() * sizeof(float), + crc); + dumped_size += size; + + // Write rotator data + std::vector buffer(rotator->dump_bytes()); + rotator->save(buffer.data()); + size = dumper->write(buffer.data(), buffer.size()); + if (size != buffer.size()) { + LOG_ERROR("Failed to write rotator data: written=%zu, expected=%zu", size, + buffer.size()); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(buffer.data(), buffer.size(), crc); + dumped_size += size; + + // Write padding + size_t padding_size = align_size(dumped_size) - dumped_size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + int ret = + dumper->append(RABITQ_CONVERTER_SEG_ID, dumped_size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", + RABITQ_CONVERTER_SEG_ID.c_str(), ret); + return ret; + } + + if (out_dumped_size) { + *out_dumped_size = dumped_size; + } + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_utils.h b/src/core/algorithm/hnsw_rabitq/rabitq_utils.h new file mode 100644 index 0000000..567c768 --- /dev/null +++ b/src/core/algorithm/hnsw_rabitq/rabitq_utils.h @@ -0,0 +1,53 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "zvec/core/framework/index_dumper.h" + +namespace zvec { +namespace core { + +inline const std::string RABITQ_CONVERTER_SEG_ID{"rabitq.converter"}; + +struct RabitqConverterHeader { + uint32_t num_clusters; + uint32_t dim; + uint32_t padded_dim; + uint32_t rotator_size; + uint8_t ex_bits; + uint8_t rotator_type; + uint8_t padding[2]; + uint32_t reserve[3]; + + RabitqConverterHeader() { + memset(static_cast(this), 0, sizeof(RabitqConverterHeader)); + } +}; +static_assert(sizeof(RabitqConverterHeader) % 32 == 0, + "RabitqConverterHeader must be aligned with 32 bytes"); + +// Common dump implementation for RabitqConverter and RabitqReformer +int dump_rabitq_centroids( + const IndexDumper::Pointer &dumper, size_t dimension, size_t padded_dim, + size_t ex_bits, size_t num_clusters, rabitqlib::RotatorType rotator_type, + const std::vector &rotated_centroids, + const std::vector ¢roids, + const std::unique_ptr> &rotator, + size_t *out_dumped_size = nullptr); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/CMakeLists.txt b/src/core/algorithm/hnsw_sparse/CMakeLists.txt new file mode 100644 index 0000000..15295b4 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/CMakeLists.txt @@ -0,0 +1,17 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(NOT APPLE) + set(CORE_KNN_HNSW_SPARSE_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_hnsw_sparse + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework sparsehash + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_HNSW_SPARSE_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.cc new file mode 100644 index 0000000..7538dfc --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.cc @@ -0,0 +1,524 @@ + +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_algorithm.h" +#include +#include +#include +#include + +namespace zvec { +namespace core { + +HnswSparseAlgorithm::HnswSparseAlgorithm(HnswSparseEntity &entity) + : entity_(entity), + mt_(std::chrono::system_clock::now().time_since_epoch().count()), + lock_pool_(kLockCnt) {} + +int HnswSparseAlgorithm::cleanup() { + return 0; +} + +int HnswSparseAlgorithm::add_node(node_id_t id, level_t level, + HnswSparseContext *ctx) { + spin_lock_.lock(); + + // std::cout << "id: " << id << ", level: " << level << std::endl; + + auto cur_max_level = entity_.cur_max_level(); + auto entry_point = entity_.entry_point(); + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + entity_.update_ep_and_level(id, level); + spin_lock_.unlock(); + return 0; + } + spin_lock_.unlock(); + + if (ailego_unlikely(level > cur_max_level)) { + mutex_.lock(); + // re-check max level + cur_max_level = entity_.cur_max_level(); + entry_point = entity_.entry_point(); + if (level <= cur_max_level) { + mutex_.unlock(); + } + } + + level_t cur_level = cur_max_level; + dist_t dist = ctx->dist_calculator()(entry_point); + for (; cur_level > level; --cur_level) { + select_entry_point(cur_level, &entry_point, &dist, ctx); + } + + for (; cur_level >= 0; --cur_level) { + search_neighbors(cur_level, &entry_point, &dist, ctx->level_topk(cur_level), + ctx); + } + + // add neighbors from down level to top level, to avoid upper level visible + // to knn_search but the under layer level not ready + for (cur_level = 0; cur_level <= level; ++cur_level) { + add_neighbors(id, cur_level, ctx->level_topk(cur_level), ctx); + ctx->level_topk(cur_level).clear(); + } + + if (ailego_unlikely(level > cur_max_level)) { + spin_lock_.lock(); + entity_.update_ep_and_level(id, level); + spin_lock_.unlock(); + mutex_.unlock(); + } + + return 0; +} + +int HnswSparseAlgorithm::search(HnswSparseContext *ctx) const { + spin_lock_.lock(); + auto maxLevel = entity_.cur_max_level(); + auto entry_point = entity_.entry_point(); + spin_lock_.unlock(); + + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + return 0; + } + + dist_t dist = ctx->dist_calculator().dist(entry_point); + for (level_t cur_level = maxLevel; cur_level >= 1; --cur_level) { + select_entry_point(cur_level, &entry_point, &dist, ctx); + } + + auto &topk_heap = ctx->topk_heap(); + topk_heap.clear(); + search_neighbors(0, &entry_point, &dist, topk_heap, ctx); + + if (ctx->group_by_search()) { + expand_neighbors_by_group(topk_heap, ctx); + } + + return 0; +} + +//! select_entry_point on hnsw level, ef = 1 +void HnswSparseAlgorithm::select_entry_point(level_t level, + node_id_t *entry_point, + dist_t *dist, + HnswSparseContext *ctx) const { + auto &entity = ctx->get_entity(); + HnswSparseDistCalculator &dc = ctx->dist_calculator(); + while (true) { + const Neighbors neighbors = entity.get_neighbors(level, *entry_point); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + uint32_t size = neighbors.size(); + if (size == 0) { + break; + } + + std::vector neighbor_block_vecs; + int ret = entity.get_vector_metas(&neighbors[0], size, neighbor_block_vecs); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } + if (ailego_unlikely(ret != 0)) { + break; + } + bool find_closer = false; + for (uint32_t i = 0; i < size; ++i) { + dist_t cur_dist = dc.dist(neighbor_block_vecs[i].data()); + if (cur_dist < *dist) { + *entry_point = neighbors[i]; + *dist = cur_dist; + find_closer = true; + } + } + + if (!find_closer) { + break; + } + } + + return; +} + +void HnswSparseAlgorithm::add_neighbors(node_id_t id, level_t level, + TopkHeap &topk_heap, + HnswSparseContext *ctx) { + if (ailego_unlikely(topk_heap.size() == 0)) { + return; + } + + HnswSparseDistCalculator &dc = ctx->dist_calculator(); + + update_neighbors(dc, id, level, topk_heap); + + // reverse update neighbors + for (size_t i = 0; i < topk_heap.size(); ++i) { + reverse_update_neighbors(dc, topk_heap[i].first, level, id, + topk_heap[i].second, ctx->update_heap()); + } + + return; +} + +void HnswSparseAlgorithm::search_neighbors(level_t level, + node_id_t *entry_point, dist_t *dist, + TopkHeap &topk, + HnswSparseContext *ctx) const { + const auto &entity = ctx->get_entity(); + HnswSparseDistCalculator &dc = ctx->dist_calculator(); + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + std::function filter = [](node_id_t) { return false; }; + if (ctx->filter().is_valid()) { + filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); }; + } + + candidates.clear(); + visit.clear(); + visit.set_visited(*entry_point); + if (!filter(*entry_point)) { + topk.emplace(*entry_point, *dist); + } + + candidates.emplace(*entry_point, *dist); + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t main_node = top->first; + dist_t main_dist = top->second; + + if (topk.full() && main_dist > topk[0].second) { + break; + } + + candidates.pop(); + const Neighbors neighbors = entity.get_neighbors(level, main_node); + ailego_prefetch(neighbors.data); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + + std::vector neighbor_ids(neighbors.size()); + uint32_t size = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) { + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_visit_dup_cnt())++; + } + continue; + } + visit.set_visited(node); + neighbor_ids[size++] = node; + } + if (size == 0) { + continue; + } + + std::vector neighbor_block_vecs; + int ret = + entity.get_vector_metas(neighbor_ids.data(), size, neighbor_block_vecs); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } + if (ailego_unlikely(ret != 0)) { + break; + } + + static constexpr node_id_t PREFETCH_STEP = 2; + static constexpr node_id_t SPARSE_PREFETCH_STEP = 1; + + for (uint32_t i = 0; i < std::min(PREFETCH_STEP, size); ++i) { + ailego_prefetch(neighbor_block_vecs[i].data()); + } + for (uint32_t i = 0; i < size; ++i) { + node_id_t node = neighbor_ids[i]; + node_id_t prefetch_id = i + PREFETCH_STEP; + if (prefetch_id < size) { + ailego_prefetch(neighbor_block_vecs[prefetch_id].data()); + } + + node_id_t sparse_prefetch_id = i + SPARSE_PREFETCH_STEP; + if (sparse_prefetch_id < size) { + IndexStorage::MemoryBlock sparse_block; + int sparse_length = 0; + entity.get_sparse_data_from_vector( + neighbor_block_vecs[sparse_prefetch_id].data(), sparse_block, + sparse_length); + auto sparse_data = std::make_pair(sparse_block.data(), sparse_length); + if (sparse_data.first != nullptr) { + ailego_prefetch(sparse_data.first); + } + } + + dist_t cur_dist = dc.dist(neighbor_block_vecs[i].data()); + if ((!topk.full()) || cur_dist < topk[0].second) { + candidates.emplace(node, cur_dist); + // update entry_point for next level scan + if (cur_dist < *dist) { + *entry_point = node; + *dist = cur_dist; + } + if (!filter(node)) { + topk.emplace(node, cur_dist); + } + } // end if + } // end for + } // while + + return; +} + +void HnswSparseAlgorithm::expand_neighbors_by_group( + TopkHeap &topk, HnswSparseContext *ctx) const { + if (!ctx->group_by().is_valid()) { + return; + } + + const auto &entity = ctx->get_entity(); + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity.get_key(id)); + }; + + // devide into groups + std::map &group_topk_heaps = ctx->group_topk_heaps(); + for (uint32_t i = 0; i < topk.size(); ++i) { + node_id_t id = topk[i].first; + auto score = topk[i].second; + + std::string group_id = group_by(id); + + auto &topk_heap = group_topk_heaps[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, score); + } + + // stage 2, expand to reach group num as possible + if (group_topk_heaps.size() < ctx->group_num()) { + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + HnswSparseDistCalculator &dc = ctx->dist_calculator(); + + std::function filter = [](node_id_t) { return false; }; + if (ctx->filter().is_valid()) { + filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); }; + } + + // refill to get enough groups + candidates.clear(); + visit.clear(); + for (uint32_t i = 0; i < topk.size(); ++i) { + node_id_t id = topk[i].first; + float score = topk[i].second; + + visit.set_visited(id); + candidates.emplace_back(id, score); + } + + // do expand + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t main_node = top->first; + + candidates.pop(); + const Neighbors neighbors = entity.get_neighbors(0, main_node); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_neighbors())++; + } + + std::vector neighbor_ids(neighbors.size()); + uint32_t size = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) { + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_visit_dup_cnt())++; + } + continue; + } + visit.set_visited(node); + neighbor_ids[size++] = node; + } + if (size == 0) { + continue; + } + + std::vector neighbor_block_vecs; + int ret = entity.get_vector_metas(neighbor_ids.data(), size, + neighbor_block_vecs); + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } + if (ailego_unlikely(ret != 0)) { + break; + } + + static constexpr node_id_t PREFETCH_STEP = 2; + for (uint32_t i = 0; i < size; ++i) { + node_id_t node = neighbor_ids[i]; + node_id_t prefetch_id = i + PREFETCH_STEP; + if (prefetch_id < size) { + ailego_prefetch(neighbor_block_vecs[prefetch_id].data()); + } + dist_t cur_dist = dc.dist(neighbor_block_vecs[i].data()); + + if (!filter(node)) { + std::string group_id = group_by(node); + + auto &topk_heap = group_topk_heaps[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(node, cur_dist); + + if (group_topk_heaps.size() >= ctx->group_num()) { + break; + } + } + + candidates.emplace(node, cur_dist); + } // end for + } // end while + } // end if +} + +void HnswSparseAlgorithm::update_neighbors(HnswSparseDistCalculator &dc, + node_id_t id, level_t level, + TopkHeap &topk_heap) { + topk_heap.sort(); + + uint32_t max_neighbor_cnt = entity_.neighbor_cnt(level); + if (topk_heap.size() <= static_cast(entity_.prune_cnt())) { + if (topk_heap.size() <= static_cast(max_neighbor_cnt)) { + entity_.update_neighbors(level, id, topk_heap); + return; + } + } + + uint32_t cur_size = 0; + for (size_t i = 0; i < topk_heap.size(); ++i) { + node_id_t cur_node = topk_heap[i].first; + dist_t cur_node_dist = topk_heap[i].second; + bool good = true; + for (uint32_t j = 0; j < cur_size; ++j) { + dist_t tmp_dist = dc.dist(cur_node, topk_heap[j].first); + if (tmp_dist <= cur_node_dist) { + good = false; + break; + } + } + + if (good) { + topk_heap[cur_size].first = cur_node; + topk_heap[cur_size].second = cur_node_dist; + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } + } + } + + // when after-prune neighbor count is too seldom, + // we use this strategy to make-up enough edges + // not only just make-up out-degrees + // we also make-up enough in-degrees + uint32_t min_neighbors = entity_.min_neighbor_cnt(); + for (size_t k = cur_size; cur_size < min_neighbors && k < topk_heap.size(); + ++k) { + bool exist = false; + for (size_t j = 0; j < cur_size; ++j) { + if (topk_heap[j].first == topk_heap[k].first) { + exist = true; + break; + } + } + if (!exist) { + topk_heap[cur_size].first = topk_heap[k].first; + topk_heap[cur_size].second = topk_heap[k].second; + cur_size++; + } + } + + topk_heap.resize(cur_size); + entity_.update_neighbors(level, id, topk_heap); + + return; +} + +void HnswSparseAlgorithm::reverse_update_neighbors(HnswSparseDistCalculator &dc, + node_id_t id, level_t level, + node_id_t link_id, + dist_t dist, + TopkHeap &update_heap) { + const size_t max_neighbor_cnt = entity_.neighbor_cnt(level); + + uint32_t lock_idx = id & kLockMask; + lock_pool_[lock_idx].lock(); + const Neighbors neighbors = entity_.get_neighbors(level, id); + size_t size = neighbors.size(); + ailego_assert_with(size <= max_neighbor_cnt, "invalid neighbor size"); + if (size < max_neighbor_cnt) { + entity_.add_neighbor(level, id, size, link_id); + lock_pool_[lock_idx].unlock(); + return; + } + + update_heap.emplace(link_id, dist); + + for (size_t i = 0; i < size; ++i) { + node_id_t node = neighbors[i]; + dist_t cur_dist = dc.dist(id, node); + update_heap.emplace(node, cur_dist); + } + + //! TODO: optimize prune + //! prune edges + update_heap.sort(); + size_t cur_size = 0; + for (size_t i = 0; i < update_heap.size(); ++i) { + node_id_t cur_node = update_heap[i].first; + dist_t cur_node_dist = update_heap[i].second; + bool good = true; + for (size_t j = 0; j < cur_size; ++j) { + dist_t tmp_dist = dc.dist(cur_node, update_heap[j].first); + if (tmp_dist <= cur_node_dist) { + good = false; + break; + } + } + + if (good) { + update_heap[cur_size].first = cur_node; + update_heap[cur_size].second = cur_node_dist; + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } + } + } + + update_heap.resize(cur_size); + entity_.update_neighbors(level, id, update_heap); + + lock_pool_[lock_idx].unlock(); + + update_heap.clear(); + + return; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.h new file mode 100644 index 0000000..5a7093f --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.h @@ -0,0 +1,127 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "hnsw_sparse_context.h" +#include "hnsw_sparse_dist_calculator.h" +#include "hnsw_sparse_entity.h" + +namespace zvec { +namespace core { + +//! hnsw graph algorithm implement +class HnswSparseAlgorithm { + public: + typedef std::unique_ptr UPointer; + + public: + //! Constructor + explicit HnswSparseAlgorithm(HnswSparseEntity &entity); + + //! Cleanup HnswSparseAlgorithm + int cleanup(); + + //! Add a node to hnsw graph + //! @id: the node unique id + //! @level: a node will be add to graph in each level [0, level] + //! return 0 on success, or errCode in failure + int add_node(node_id_t id, level_t level, HnswSparseContext *ctx); + + //! do knn search in graph + //! return 0 on success, or errCode in failure. results saved in ctx + int search(HnswSparseContext *ctx) const; + + //! Initiate HnswAlgorithm + int init() { + level_probas_.clear(); + double level_mult = + 1 / std::log(static_cast(entity_.scaling_factor())); + for (int level = 0;; level++) { + // refers faiss get_random_level alg + double proba = + std::exp(-level / level_mult) * (1 - std::exp(-1 / level_mult)); + if (proba < 1e-9) { + break; + } + level_probas_.push_back(proba); + } + + return 0; + } + + //! Generate a random level + //! return graph level + uint32_t get_random_level() const { + // gen rand float (0, 1) + double f = mt_() / static_cast(mt_.max()); + for (size_t level = 0; level < level_probas_.size(); level++) { + if (f < level_probas_[level]) { + return level; + } + f -= level_probas_[level]; + } + return level_probas_.size() - 1; + } + + private: + //! Select in upper layer to get entry point for next layer search + void select_entry_point(level_t level, node_id_t *entry_point, dist_t *dist, + HnswSparseContext *ctx) const; + + //! update node id neighbors from topkHeap, and reverse link is also updated + void add_neighbors(node_id_t id, level_t level, TopkHeap &topk_heap, + HnswSparseContext *ctx); + + //! Given a node id and level, search the nearest neighbors in graph + //! Note: the nearest neighbors result keeps in topk, and entry_point and + //! dist will be updated to current level nearest node id and distance + void search_neighbors(level_t level, node_id_t *entry_point, dist_t *dist, + TopkHeap &topk, HnswSparseContext *ctx) const; + + //! Update the node's neighbors + void update_neighbors(HnswSparseDistCalculator &dc, node_id_t id, + level_t level, TopkHeap &topk_heap); + + //! Checking linkId could be id's new neighbor, and add as neighbor if true + //! @dc distance calculator + //! @updateHeap temporary heap in updating neighbors + void reverse_update_neighbors(HnswSparseDistCalculator &dc, node_id_t id, + level_t level, node_id_t link_id, dist_t dist, + TopkHeap &update_heap); + + //! expand neighbors until group nums are reached + void expand_neighbors_by_group(TopkHeap &topk, HnswSparseContext *ctx) const; + + private: + HnswSparseAlgorithm(const HnswSparseAlgorithm &) = delete; + HnswSparseAlgorithm &operator=(const HnswSparseAlgorithm &) = delete; + + private: + static constexpr uint32_t kLockCnt{1U << 8}; + static constexpr uint32_t kLockMask{kLockCnt - 1U}; + + HnswSparseEntity &entity_; + mutable std::mt19937 mt_{}; + std::vector level_probas_{}; + + mutable ailego::SpinMutex spin_lock_{}; // global spin lock + std::mutex mutex_{}; // global mutex + // TODO: spin lock? + std::vector lock_pool_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc new file mode 100644 index 0000000..4ca8f89 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc @@ -0,0 +1,528 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_builder.h" +#include +#include +#include +#include +#include +#include +#include "hnsw_sparse_algorithm.h" +#include "hnsw_sparse_params.h" + +namespace zvec { +namespace core { + +HnswSparseBuilder::HnswSparseBuilder() {} + +int HnswSparseBuilder::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + LOG_INFO("Begin HnswSparseBuilder::init"); + + meta_ = meta; + auto params_copy = params; + meta_.set_builder("HnswSparseBuilder", HnswSparseEntity::kRevision, + std::move(params_copy)); + + size_t memory_quota = 0UL; + params.get(PARAM_HNSW_SPARSE_BUILDER_MEMORY_QUOTA, &memory_quota); + params.get(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, &thread_cnt_); + params.get(PARAM_HNSW_SPARSE_BUILDER_EFCONSTRUCTION, &ef_construction_); + params.get(PARAM_HNSW_SPARSE_BUILDER_CHECK_INTERVAL_SECS, + &check_interval_secs_); + + params.get(PARAM_HNSW_SPARSE_BUILDER_MAX_NEIGHBOR_COUNT, + &upper_max_neighbor_cnt_); + float multiplier = HnswSparseEntity::kDefaultL0MaxNeighborCntMultiplier; + params.get(PARAM_HNSW_SPARSE_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, + &multiplier); + l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_; + scaling_factor_ = upper_max_neighbor_cnt_; + params.get(PARAM_HNSW_SPARSE_BUILDER_SCALING_FACTOR, &scaling_factor_); + + multiplier = HnswSparseEntity::kDefaultNeighborPruneMultiplier; + params.get(PARAM_HNSW_SPARSE_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier); + size_t prune_cnt = multiplier * upper_max_neighbor_cnt_; + + if (ef_construction_ == 0) { + ef_construction_ = HnswSparseEntity::kDefaultEfConstruction; + } + if (upper_max_neighbor_cnt_ == 0) { + upper_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt; + } + if (upper_max_neighbor_cnt_ > kMaxNeighborCnt) { + LOG_ERROR("[%s] must be in range (0,%d]", + PARAM_HNSW_SPARSE_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), + kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) { + LOG_ERROR("[%s]-[%d] must be <= [%s]-[%d]", + PARAM_HNSW_SPARSE_BUILDER_MIN_NEIGHBOR_COUNT.c_str(), + min_neighbor_cnt_, + PARAM_HNSW_SPARSE_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), + upper_max_neighbor_cnt_); + return IndexError_InvalidArgument; + } + if (l0_max_neighbor_cnt_ == 0) { + l0_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt; + } + if (l0_max_neighbor_cnt_ > HnswSparseEntity::kMaxNeighborCnt) { + LOG_ERROR("L0MaxNeighborCnt must be in range (0,%d)", + HnswSparseEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (scaling_factor_ == 0U) { + scaling_factor_ = HnswSparseEntity::kDefaultScalingFactor; + } + if (scaling_factor_ < 5 || scaling_factor_ > 1000) { + LOG_ERROR("[%s] must be in range [5,1000]", + PARAM_HNSW_SPARSE_BUILDER_SCALING_FACTOR.c_str()); + return IndexError_InvalidArgument; + } + if (thread_cnt_ == 0) { + thread_cnt_ = std::thread::hardware_concurrency(); + } + if (thread_cnt_ > std::thread::hardware_concurrency()) { + LOG_WARN("[%s] greater than cpu cores %u", + PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT.c_str(), + std::thread::hardware_concurrency()); + } + if (prune_cnt == 0UL) { + prune_cnt = upper_max_neighbor_cnt_; + } + + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("CreateMeasure failed, name: %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + int ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMeasure init failed, ret=%d", ret); + return ret; + } + + entity_.set_ef_construction(ef_construction_); + entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_); + entity_.set_min_neighbor_cnt(min_neighbor_cnt_); + entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_); + entity_.set_scaling_factor(scaling_factor_); + entity_.set_memory_quota(memory_quota); + entity_.set_prune_cnt(prune_cnt); + + entity_.set_sparse_meta_size(HnswSparseEntity::kSparseMetaSize); + entity_.set_sparse_unit_size(meta.unit_size()); + + ret = entity_.init(); + if (ret != 0) { + return ret; + } + + alg_ = HnswSparseAlgorithm::UPointer(new HnswSparseAlgorithm(entity_)); + + ret = alg_->init(); + if (ret != 0) { + return ret; + } + + state_ = BUILD_STATE_INITED; + LOG_INFO( + "End HnswSparseBuilder::init, params: efConstruction=%u " + "l0NeighborCnt=%u upperNeighborCnt=%u scalingFactor=%u " + "memoryQuota=%zu neighborPruneCnt=%zu measureName=%s ", + ef_construction_, l0_max_neighbor_cnt_, upper_max_neighbor_cnt_, + scaling_factor_, memory_quota, prune_cnt, meta_.metric_name().c_str()); + + return 0; +} + +int HnswSparseBuilder::cleanup(void) { + LOG_INFO("Begin HnswSparseBuilder::cleanup"); + + l0_max_neighbor_cnt_ = HnswSparseEntity::kDefaultL0MaxNeighborCnt; + min_neighbor_cnt_ = 0; + upper_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt; + ef_construction_ = HnswSparseEntity::kDefaultEfConstruction; + scaling_factor_ = HnswSparseEntity::kDefaultScalingFactor; + check_interval_secs_ = kDefaultLogIntervalSecs; + errcode_ = 0; + error_ = false; + entity_.cleanup(); + alg_->cleanup(); + meta_.clear(); + metric_.reset(); + stats_.clear_attributes(); + stats_.set_trained_count(0UL); + stats_.set_built_count(0UL); + stats_.set_dumped_count(0UL); + stats_.set_discarded_count(0UL); + stats_.set_trained_costtime(0UL); + stats_.set_built_costtime(0UL); + stats_.set_dumped_costtime(0UL); + state_ = BUILD_STATE_INIT; + + LOG_INFO("End HnswSparseBuilder::cleanup"); + + return 0; +} + +int HnswSparseBuilder::train(IndexThreads::Pointer, + IndexSparseHolder::Pointer /*holder*/) { + if (state_ != BUILD_STATE_INITED) { + LOG_ERROR("Init the builder before HnswSparseBuilder::train"); + return IndexError_NoReady; + } + + stats_.set_trained_count(0UL); + stats_.set_trained_costtime(0UL); + state_ = BUILD_STATE_TRAINED; + + LOG_INFO("End HnswSparseBuilder::train"); + + return 0; +} + +int HnswSparseBuilder::train(const IndexTrainer::Pointer & /*trainer*/) { + if (state_ != BUILD_STATE_INITED) { + LOG_ERROR("Init the builder before HnswSparseBuilder::train"); + return IndexError_NoReady; + } + + LOG_INFO("Begin HnswSparseBuilder::train by trainer"); + + stats_.set_trained_count(0UL); + stats_.set_trained_costtime(0UL); + state_ = BUILD_STATE_TRAINED; + + LOG_INFO("End HnswSparseBuilder::train by trainer"); + + return 0; +} + +int HnswSparseBuilder::build(IndexThreads::Pointer threads, + IndexSparseHolder::Pointer holder) { + if (!holder) { + LOG_ERROR("Input holder is nullptr while building index"); + return IndexError_InvalidArgument; + } + + if (!holder->is_matched(meta_)) { + LOG_ERROR("Input holder doesn't match index meta while building index"); + return IndexError_Mismatch; + } + if (!threads) { + threads = std::make_shared(thread_cnt_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + + auto start_time = ailego::Monotime::MilliSeconds(); + + LOG_INFO("Begin HnswSparseBuilder::build sparse"); + + // holder should be hybrid holder + auto sparse_holder = std::dynamic_pointer_cast(holder); + + if (sparse_holder == nullptr) { + LOG_ERROR("HnswSparseBuilder failed to cast holder"); + return IndexError_Runtime; + } + + if (sparse_holder->count() != static_cast(-1)) { + LOG_DEBUG("HnswSparseBuilder holder documents count %lu", + sparse_holder->count()); + + int ret = entity_.reserve_space(sparse_holder->count(), + sparse_holder->total_sparse_count()); + if (ret != 0) { + LOG_ERROR("HnswBuilde reserver space failed"); + return ret; + } + } + auto iter = sparse_holder->create_iterator(); + if (!iter) { + LOG_ERROR("Create iterator for holder failed"); + return IndexError_Runtime; + } + + int ret; + error_ = false; + while (iter->is_valid()) { + level_t level = alg_->get_random_level(); + node_id_t id; + + ret = entity_.add_vector(level, iter->key(), iter->sparse_count(), + iter->sparse_indices(), iter->sparse_data(), &id); + + if (ailego_unlikely(ret != 0) && ret != IndexError_InvalidValue) { + return ret; + } + + iter->next(); + } + // Holder is not needed, cleanup it. + sparse_holder.reset(); + + LOG_INFO("Finished save vector, start build graph..."); + + std::atomic finished{0}; + + ret = build_graph(threads, finished); + if (ret != 0) { + LOG_ERROR("Failed to build graph"); + return ret; + } + + stats_.set_built_count(finished.load()); + stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time); + state_ = BUILD_STATE_BUILT; + + LOG_INFO("End HnswSparseBuilder::build"); + return 0; +} + +int HnswSparseBuilder::build_graph(IndexThreads::Pointer threads, + std::atomic &finished) { + auto task_group = threads->make_group(); + if (!task_group) { + LOG_ERROR("Failed to create task group"); + return IndexError_Runtime; + } + + for (size_t i = 0; i < threads->count(); ++i) { + task_group->submit(ailego::Closure ::New(this, &HnswSparseBuilder::do_build, + i, threads->count(), &finished)); + } + + { + std::unique_lock lk(mutex_); + while (finished.load() < entity_.doc_cnt()) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to build index while waiting finish"); + return errcode_; + } + LOG_INFO("Built cnt %u, finished percent %.3f%%", finished.load(), + finished.load() * 100.0f / entity_.doc_cnt()); + } + } + + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to build index while waiting finish"); + return errcode_; + } + task_group->wait_finish(); + + return 0; +} + +void HnswSparseBuilder::do_build(node_id_t idx, size_t step_size, + std::atomic *finished) { + AILEGO_DEFER([&]() { + std::lock_guard latch(mutex_); + cond_.notify_one(); + }); + + HnswSparseContext *ctx = new (std::nothrow) HnswSparseContext( + metric_, + std::shared_ptr(&entity_, [](HnswSparseEntity *) {})); + if (ailego_unlikely(ctx == nullptr)) { + if (!error_.exchange(true)) { + LOG_ERROR("Failed to create context"); + errcode_ = IndexError_NoMemory; + } + return; + } + HnswSparseContext::Pointer auto_ptr(ctx); + ctx->set_max_scan_num(entity_.doc_cnt()); + int ret = ctx->init(HnswSparseContext::kSparseBuilderContext); + if (ret != 0) { + if (!error_.exchange(true)) { + LOG_ERROR("Failed to init context"); + errcode_ = IndexError_Runtime; + } + return; + } + + IndexQueryMeta qmeta(meta_.data_type()); + for (node_id_t id = idx; id < entity_.doc_cnt(); id += step_size) { + const void *vec = entity_.get_vector_meta(id); + + auto sparse_data = entity_.get_sparse_data_from_vector(vec); + + ctx->reset_query(sparse_data.first); + + ret = alg_->add_node(id, entity_.get_level(id), ctx); + if (ailego_unlikely(ret != 0)) { + if (!error_.exchange(true)) { + LOG_ERROR("Hnsw graph add node failed"); + errcode_ = ret; + } + return; + } + ctx->clear(); + (*finished)++; + } +} + +int HnswSparseBuilder::dump(const IndexDumper::Pointer &dumper) { + if (state_ != BUILD_STATE_BUILT) { + LOG_INFO("Build the index before HnswSparseBuilder::dump"); + return IndexError_NoReady; + } + + LOG_INFO("Begin HnswSparseBuilder::dump"); + + meta_.set_searcher("HnswSparseSearcher", HnswSparseEntity::kRevision, + ailego::Params()); + auto start_time = ailego::Monotime::MilliSeconds(); + + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + + ret = entity_.dump(dumper); + if (ret != 0) { + LOG_ERROR("HnswSparseBuilder dump index failed"); + return ret; + } + + stats_.set_dumped_count(entity_.doc_cnt()); + stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time); + + LOG_INFO("EndHnswSparseBuilder::dump"); + return 0; +} + +int HnswSparseBuilder::build(IndexThreads::Pointer threads, size_t count, + const uint64_t *keys, + const uint64_t *sparse_indptr, + const uint32_t *sparse_indices, + const void *sparse_data) { + IndexQueryMeta qmeta(meta_.data_type()); + + return build(threads, qmeta, count, keys, sparse_indptr, sparse_indices, + sparse_data); +} + +int HnswSparseBuilder::build(IndexThreads::Pointer threads, + const IndexQueryMeta &qmeta, size_t count, + const uint64_t *keys, + const uint64_t *sparse_indptr, + const uint32_t *sparse_indices, + const void *sparse_data) { + if (!threads) { + threads = std::make_shared(thread_cnt_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + + auto start_time = ailego::Monotime::MilliSeconds(); + + LOG_INFO("Begin HnswSparseBuilder::build sparse, documents count %lu", count); + + size_t total_sparse_count = sparse_indptr[count]; + + int ret = entity_.reserve_space(count, total_sparse_count); + if (ret != 0) { + LOG_ERROR("HnswBuilde reserver space failed"); + return ret; + } + + if (qmeta.data_type() == meta_.data_type()) { + for (size_t i = 0; i < count; i++) { + level_t level = alg_->get_random_level(); + node_id_t id; + + uint32_t sparse_count = sparse_indptr[i + 1] - sparse_indptr[i]; + const uint32_t *sparse_indices_temp = sparse_indices + sparse_indptr[i]; + + const void *sparse_data_temp = static_cast(sparse_data) + + sparse_indptr[i] * qmeta.unit_size(); + + ret = entity_.add_vector(level, keys[i], sparse_count, + sparse_indices_temp, sparse_data_temp, &id); + if (ailego_unlikely(ret != 0) && ret != IndexError_InvalidValue) { + return ret; + } + } + } else if (meta_.data_type() == IndexMeta::DataType::DT_FP16 && + qmeta.data_type() == IndexMeta::DataType::DT_FP32) { + // transform from float 32 to float 16 + auto reformer = IndexFactory::CreateReformer("HalfFloatSparseReformer"); + if (!reformer) { + LOG_ERROR("Sparse reformer not existed."); + + return IndexError_NoExist; + } + + meta_.set_converter("HalfFloatSparseConverter", 0, ailego::Params()); + meta_.set_reformer("HalfFloatSparseReformer", 0, ailego::Params()); + + for (size_t i = 0; i < count; i++) { + level_t level = alg_->get_random_level(); + node_id_t id; + + uint32_t sparse_count = sparse_indptr[i + 1] - sparse_indptr[i]; + const uint32_t *sparse_indices_temp = sparse_indices + sparse_indptr[i]; + + const void *sparse_data_temp = static_cast(sparse_data) + + sparse_indptr[i] * qmeta.unit_size(); + + std::string query_fp16; + IndexQueryMeta ometa; + + reformer->transform(sparse_count, sparse_indices_temp, sparse_data_temp, + qmeta, &query_fp16, &ometa); + + ret = entity_.add_vector(level, keys[i], sparse_count, + sparse_indices_temp, query_fp16.data(), &id); + if (ailego_unlikely(ret != 0) && ret != IndexError_InvalidValue) { + return ret; + } + } + } else { + LOG_ERROR("Format not supported."); + + return IndexError_Unsupported; + } + + LOG_INFO("Finished save vector, start build graph..."); + + std::atomic finished{0}; + + ret = build_graph(threads, finished); + if (ret != 0) { + LOG_ERROR("Failed to build graph"); + return ret; + } + + stats_.set_built_count(finished.load()); + stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time); + state_ = BUILD_STATE_BUILT; + + LOG_INFO("End HnswSparseBuilder::build"); + return 0; +} + +INDEX_FACTORY_REGISTER_BUILDER(HnswSparseBuilder); + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.h new file mode 100644 index 0000000..35c45cc --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.h @@ -0,0 +1,103 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "hnsw_sparse_algorithm.h" +#include "hnsw_sparse_builder_entity.h" + +namespace zvec { +namespace core { + +class HnswSparseBuilder : public IndexBuilder { + public: + //! Constructor + HnswSparseBuilder(); + + //! Initialize the builder + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup the builder + int cleanup(void) override; + + //! Train the data + int train(IndexThreads::Pointer, IndexSparseHolder::Pointer holder) override; + + //! Train the data + int train(const IndexTrainer::Pointer &trainer) override; + + //! Build the index + int build(IndexThreads::Pointer threads, + IndexSparseHolder::Pointer holder) override; + + //! Build the index with indptr format + int build(IndexThreads::Pointer threads, const IndexQueryMeta &qmeta, + size_t count, const uint64_t *keys, const uint64_t *sparse_indptr, + const uint32_t *sparse_indices, const void *sparse_data) override; + + //! Build the index with indptr format + int build(IndexThreads::Pointer threads, size_t count, const uint64_t *keys, + const uint64_t *sparse_indptr, const uint32_t *sparse_indices, + const void *sparse_data) override; + + //! Dump index into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + private: + int build_graph(IndexThreads::Pointer threads, + std::atomic &finished); + void do_build(node_id_t idx, size_t step_size, + std::atomic *finished); + + constexpr static uint32_t kDefaultLogIntervalSecs = 15U; + constexpr static uint32_t kMaxNeighborCnt = 65535; + + private: + enum BUILD_STATE { + BUILD_STATE_INIT = 0, + BUILD_STATE_INITED = 1, + BUILD_STATE_TRAINED = 2, + BUILD_STATE_BUILT = 3 + }; + + HnswSparseBuilderEntity entity_{}; + HnswSparseAlgorithm::UPointer alg_; // impl graph algorithm + uint32_t thread_cnt_{0}; + uint32_t l0_max_neighbor_cnt_{HnswSparseEntity::kDefaultL0MaxNeighborCnt}; + uint32_t min_neighbor_cnt_{0}; + uint32_t upper_max_neighbor_cnt_{ + HnswSparseEntity::kDefaultUpperMaxNeighborCnt}; + uint32_t ef_construction_{HnswSparseEntity::kDefaultEfConstruction}; + uint32_t scaling_factor_{HnswSparseEntity::kDefaultScalingFactor}; + uint32_t check_interval_secs_{kDefaultLogIntervalSecs}; + + int errcode_{0}; + std::atomic_bool error_{false}; + IndexMeta meta_{}; + IndexMetric::Pointer metric_{}; + std::mutex mutex_{}; + std::condition_variable cond_{}; + Stats stats_{}; + + BUILD_STATE state_{BUILD_STATE_INIT}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.cc new file mode 100644 index 0000000..3c580fa --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.cc @@ -0,0 +1,291 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_builder_entity.h" +#include +#include "utility/sparse_utility.h" + +namespace zvec { +namespace core { + +HnswSparseBuilderEntity::HnswSparseBuilderEntity() { + update_ep_and_level(kInvalidNodeId, 0U); +} + +int HnswSparseBuilderEntity::cleanup() { + memory_quota_ = 0UL; + neighbors_size_ = 0U; + upper_neighbors_size_ = 0U; + padding_size_ = 0U; + vectors_buffer_.clear(); + keys_buffer_.clear(); + neighbors_buffer_.clear(); + upper_neighbors_buffer_.clear(); + neighbors_index_.clear(); + + vectors_buffer_.shrink_to_fit(); + keys_buffer_.shrink_to_fit(); + neighbors_buffer_.shrink_to_fit(); + upper_neighbors_buffer_.shrink_to_fit(); + neighbors_index_.shrink_to_fit(); + + this->HnswSparseEntity::cleanup(); + + return 0; +} + +int HnswSparseBuilderEntity::init() { + size_t size = vector_size(); + + size += sparse_meta_size(); + + //! aligned size to 32 + set_node_size(AlignSize(size)); + //! if node size is aligned to 1k, the build performance will downgrade + if (node_size() % 1024 == 0) { + set_node_size(AlignSize(node_size() + 1)); + } + + padding_size_ = node_size() - size; + + neighbors_size_ = neighbors_size(); + upper_neighbors_size_ = upper_neighbors_size(); + + return 0; +} + +int HnswSparseBuilderEntity::reserve_space(size_t docs, + size_t total_sparse_count) { + if (memory_quota_ > 0 && (node_size() * docs + neighbors_size_ * docs + + sizeof(SparseNeighborIndex) * docs > + memory_quota_)) { + return IndexError_NoMemory; + } + + vectors_buffer_.reserve(node_size() * docs); + keys_buffer_.reserve(sizeof(key_t) * docs); + neighbors_buffer_.reserve(neighbors_size_ * docs); + neighbors_index_.reserve(docs); + + sparse_data_buffer_.reserve(sizeof(uint32_t) * docs + + (sizeof(uint32_t)) * total_sparse_count + + sparse_unit_size() * total_sparse_count); + + return 0; +} + +int HnswSparseBuilderEntity::add_vector(level_t level, key_t key, + const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_vec, node_id_t *id) { + if (ailego_unlikely(sparse_count > HnswSparseEntity::kSparseMaxDimSize)) { + LOG_WARN( + "Failed to add sparse vector: number of non-zero elements (%u) exceeds " + "maximum allowed (%u), key=%zu", + sparse_count, HnswSparseEntity::kSparseMaxDimSize, (size_t)key); + return IndexError_InvalidValue; + } + + std::string sparse_buffer; + SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_vec, + sparse_unit_size(), sparse_buffer); + + uint32_t sparse_len = sparse_buffer.size(); + + if (memory_quota_ > 0 && + (vectors_buffer_.capacity() + keys_buffer_.capacity() + + neighbors_buffer_.capacity() + upper_neighbors_buffer_.capacity() + + neighbors_index_.capacity() * sizeof(SparseNeighborIndex) + + sparse_len > + memory_quota_)) { + LOG_ERROR("Add vector failed, used memory exceed quota, cur_doc=%u", + doc_cnt()); + return IndexError_NoMemory; + } + + vectors_buffer_.append(reinterpret_cast(&sparse_data_offset_), + sizeof(uint64_t)); + vectors_buffer_.append(reinterpret_cast(&sparse_len), + sizeof(uint32_t)); + vectors_buffer_.append(sizeof(uint32_t), + '\0'); // reserve to make it up to meta size + vectors_buffer_.append(padding_size_, '\0'); + + keys_buffer_.append(reinterpret_cast(&key), sizeof(key)); + + sparse_data_buffer_.append(sparse_buffer.data(), sparse_len); + sparse_data_offset_ += sparse_len; + + // init level 0 neighbors + neighbors_buffer_.append(neighbors_size_, '\0'); + + neighbors_index_.emplace_back(upper_neighbors_buffer_.size(), level); + + // init upper layer neighbors + for (level_t cur_level = 1; cur_level <= level; ++cur_level) { + upper_neighbors_buffer_.append(upper_neighbors_size_, '\0'); + } + + *id = (*mutable_doc_cnt())++; + + return 0; +} + +key_t HnswSparseBuilderEntity::get_key(node_id_t id) const { + return *(reinterpret_cast(keys_buffer_.data() + + id * sizeof(key_t))); +} + +const void *HnswSparseBuilderEntity::get_vector_meta(node_id_t id) const { + return vectors_buffer_.data() + id * node_size(); +} + +int HnswSparseBuilderEntity::get_vector_meta( + const node_id_t id, IndexStorage::MemoryBlock &block) const { + const void *vec = get_vector_meta(id); + block.reset((void *)vec); + return 0; +} + +int HnswSparseBuilderEntity::get_vector_metas(const node_id_t *ids, + uint32_t count, + const void **vecs) const { + for (uint32_t i = 0; i < count; ++i) { + vecs[i] = vectors_buffer_.data() + ids[i] * node_size(); + } + + return 0; +} + +int HnswSparseBuilderEntity::get_vector_metas( + const node_id_t *ids, uint32_t count, + std::vector &block_vecs) const { + std::vector vecs(count); + get_vector_metas(ids, count, vecs.data()); + for (uint32_t i = 0; i < count; ++i) { + block_vecs.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i])); + } + return 0; +} + +//! Get vector feature data by key +const void *HnswSparseBuilderEntity::get_sparse_data(uint64_t offset, + uint32_t /*len*/) const { + return reinterpret_cast(sparse_data_buffer_.data()) + offset; +} + +int HnswSparseBuilderEntity::get_sparse_data( + uint64_t offset, uint32_t len, IndexStorage::MemoryBlock &block) const { + const void *vec = get_sparse_data(offset, len); + block.reset((void *)vec); + return 0; +} + +//! Get sparse data from id +const void *HnswSparseBuilderEntity::get_sparse_data(node_id_t id) const { + auto sparse_data = get_sparse_data_from_vector(get_vector_meta(id)); + + return sparse_data.first; +} + +int HnswSparseBuilderEntity::get_sparse_data( + const node_id_t id, IndexStorage::MemoryBlock &block) const { + const void *vec = get_sparse_data(id); + block.reset((void *)vec); + return 0; +} + +//! Get sparse data from vector +std::pair +HnswSparseBuilderEntity::get_sparse_data_from_vector(const void *vec) const { + uint32_t vec_size = vector_size(); + const char *vec_ptr = reinterpret_cast(vec); + + uint64_t offset = *((uint64_t *)(vec_ptr + vec_size)); + uint32_t sparse_vector_len = + *((uint32_t *)(vec_ptr + vec_size + sizeof(uint64_t))); + + const void *sparse_data = get_sparse_data(offset, sparse_vector_len); + if (ailego_unlikely(sparse_data == nullptr)) { + LOG_ERROR("Get nullptr sparse, offset=%zu, len=%u", (size_t)offset, + sparse_vector_len); + + return std::make_pair(nullptr, 0); + } + + return std::make_pair(sparse_data, sparse_vector_len); +} + +int HnswSparseBuilderEntity::get_sparse_data_from_vector( + const void *vec, IndexStorage::MemoryBlock &block, + int &sparse_length) const { + std::pair sparse_data = + get_sparse_data_from_vector(vec); + block.reset((void *)sparse_data.first); + sparse_length = sparse_data.second; + return 0; +} + +const Neighbors HnswSparseBuilderEntity::get_neighbors(level_t level, + node_id_t id) const { + const NeighborsHeader *hd = get_neighbor_header(level, id); + return {hd->neighbor_cnt, hd->neighbors}; +} + +int HnswSparseBuilderEntity::update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) { + NeighborsHeader *hd = + const_cast(get_neighbor_header(level, id)); + for (size_t i = 0; i < neighbors.size(); ++i) { + hd->neighbors[i] = neighbors[i].first; + } + hd->neighbor_cnt = neighbors.size(); + + // std::cout << "id: " << id << ", neighbour, id: "; + // for (size_t i = 0; i < neighbors.size(); ++i) { + // if (i == neighbors.size()-1) + // std::cout << neighbors[i].first << ", score:" << neighbors[i].second << + // std::endl; + // else + // std::cout << neighbors[i].first << ", score:" << neighbors[i].second << + // ", id: "; + // } + + return 0; +} + +void HnswSparseBuilderEntity::add_neighbor(level_t level, node_id_t id, + uint32_t /*size*/, + node_id_t neighbor_id) { + NeighborsHeader *hd = + const_cast(get_neighbor_header(level, id)); + hd->neighbors[hd->neighbor_cnt++] = neighbor_id; + + return; +} + +int HnswSparseBuilderEntity::dump(const IndexDumper::Pointer &dumper) { + key_t *keys = + reinterpret_cast(const_cast(keys_buffer_.data())); + auto ret = + dump_segments(dumper, keys, [&](node_id_t id) { return get_level(id); }); + if (ailego_unlikely(ret < 0)) { + return ret; + } + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.h new file mode 100644 index 0000000..1fddc6b --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.h @@ -0,0 +1,159 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "hnsw_sparse_entity.h" + +namespace zvec { +namespace core { + +class HnswSparseBuilderEntity : public HnswSparseEntity { + public: + //! Add vector and key to hnsw entity, and local id will be saved in id + int add_vector(level_t level, key_t key, const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_vec, + node_id_t *id) override; + + //! Get primary key of the node id + key_t get_key(node_id_t id) const override; + + //! Get vector feature data by key + const void *get_vector_meta(node_id_t id) const override; + + int get_vector_meta(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Batch get vectors feature data by keys + int get_vector_metas(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + int get_vector_metas( + const node_id_t *ids, uint32_t count, + std::vector &block_vecs) const override; + + //! Get the node id's neighbors on graph level + const NeighborsHeader *get_neighbor_header(level_t level, + node_id_t id) const { + if (level == 0) { + return reinterpret_cast( + neighbors_buffer_.data() + neighbors_size_ * id); + } else { + size_t offset = neighbors_index_[id].offset; + return reinterpret_cast( + upper_neighbors_buffer_.data() + offset + + (level - 1) * upper_neighbors_size_); + } + } + + //! Get the node id's neighbors on graph level + const Neighbors get_neighbors(level_t level, node_id_t id) const override; + + //! Replace node id in level's neighbors + int update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) override; + + //! add a neighbor to id in graph level + void add_neighbor(level_t level, node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + + //! Get vector sparse feature data by chunk index and offset + const void *get_sparse_data(uint64_t offset, uint32_t len) const override; + //! Get sparse data from id + const void *get_sparse_data(node_id_t id) const override; + + int get_sparse_data(uint64_t offset, uint32_t len, + IndexStorage::MemoryBlock &block) const override; + + int get_sparse_data(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get sparse data from vector + std::pair get_sparse_data_from_vector( + const void *vec) const override; + + int get_sparse_data_from_vector(const void *vec, + IndexStorage::MemoryBlock &block, + int &sparse_length) const override; + + //! Dump the hnsw graph to dumper + int dump(const IndexDumper::Pointer &dumper) override; + + //! Cleanup the entity + int cleanup(void) override; + + public: + //! Constructor + HnswSparseBuilderEntity(); + + //! Get the node graph level by id + level_t get_level(node_id_t id) const { + return neighbors_index_[id].level; + } + + //! Init builerEntity + int init(); + + //! reserve buffer space for documents + //! @param docs number of documents + //! @param total_sparse_count total dim of sparse count + int reserve_space(size_t docs, size_t total_sparse_count); + + //! Set memory quota params + inline void set_memory_quota(size_t memory_quota) { + memory_quota_ = memory_quota; + } + + //! Get neighbors size + inline size_t neighbors_size() const { + return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); + } + + //! Get upper neighbors size + inline size_t upper_neighbors_size() const { + return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); + } + + public: + HnswSparseBuilderEntity(const HnswSparseBuilderEntity &) = delete; + HnswSparseBuilderEntity &operator=(const HnswSparseBuilderEntity &) = delete; + + private: + friend class HnswSparseSearcherEntity; + + //! class internal used only + struct SparseNeighborIndex { + SparseNeighborIndex(size_t off, level_t l) : offset(off), level(l) {} + uint64_t offset : 48; + uint64_t level : 16; + }; + + std::string vectors_buffer_{}; // aligned vectors + std::string keys_buffer_{}; // aligned vectors + std::string neighbors_buffer_{}; // level 0 neighbors buffer + std::string upper_neighbors_buffer_{}; // upper layer neighbors buffer + + std::string sparse_data_buffer_{}; // aligned spase data buffer + size_t sparse_data_offset_{0}; // + + // upper layer offset + level in upper_neighbors_buffer_ + std::vector neighbors_index_{}; + size_t memory_quota_{0UL}; + size_t neighbors_size_{0U}; // level 0 neighbors size + size_t upper_neighbors_size_{0U}; // level 0 neighbors size + size_t padding_size_{}; // padding size for each vector element +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.cc new file mode 100644 index 0000000..b3e6f49 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.cc @@ -0,0 +1,217 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_chunk.h" +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +int SparseChunkBroker::init_storage(size_t chunk_size) { + chunk_meta_.clear(); + chunk_meta_.chunk_size = chunk_size; + chunk_meta_.create_time = ailego::Realtime::Seconds(); + stats_.set_create_time(chunk_meta_.create_time); + chunk_meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(chunk_meta_.update_time); + + //! alloc meta chunk + size_t size = sizeof(HnswSparseChunkMeta); + size = (size + page_mask_) & (~page_mask_); + const std::string segment_id = + make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId); + int ret = stg_->append(segment_id, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Storage append segment failed for %s", IndexError::What(ret)); + return ret; + } + chunk_meta_segment_ = get_chunk(CHUNK_TYPE_META, kDefaultChunkSeqId); + if (ailego_unlikely(!chunk_meta_segment_)) { + LOG_ERROR("Get meta segment failed"); + return IndexError_Runtime; + } + + //! update meta info and write to storage + chunk_meta_.chunk_cnts[CHUNK_TYPE_META] += 1; + chunk_meta_.total_size += size; + (*stats_.mutable_index_size()) += size; + size = chunk_meta_segment_->write(0UL, &chunk_meta_, + sizeof(HnswSparseChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswSparseChunkMeta))) { + LOG_ERROR("Storage write data failed, wsize=%zu", size); + return IndexError_WriteData; + } + + return 0; +} + +int SparseChunkBroker::load_storage(size_t chunk_size) { + IndexStorage::MemoryBlock data_block; + size_t size = chunk_meta_segment_->read(0UL, data_block, + chunk_meta_segment_->data_size()); + if (size != sizeof(HnswSparseChunkMeta)) { + LOG_ERROR("Invalid hnsw meta chunk, read size=%zu chunk size=%zu", size, + chunk_meta_segment_->data_size()); + return IndexError_InvalidFormat; + } + std::memcpy(static_cast(&chunk_meta_), data_block.data(), size); + if (chunk_meta_.chunk_size != chunk_size) { + LOG_ERROR( + "Params hnsw chunk size=%zu mismatch from previous %zu " + "in index", + chunk_size, (size_t)chunk_meta_.chunk_size); + return IndexError_Mismatch; + } + + *stats_.mutable_check_point() = stg_->check_point(); + stats_.set_revision_id(chunk_meta_.revision_id); + stats_.set_update_time(chunk_meta_.update_time); + stats_.set_create_time(chunk_meta_.create_time); + + char create_time[32]; + char update_time[32]; + ailego::Realtime::Gmtime(chunk_meta_.create_time, "%Y-%m-%d %H:%M:%S", + create_time, sizeof(create_time)); + ailego::Realtime::Gmtime(chunk_meta_.update_time, "%Y-%m-%d %H:%M:%S", + update_time, sizeof(update_time)); + LOG_DEBUG( + "Load index, indexSize=%zu chunkSize=%zu nodeChunks=%zu " + "upperNeighborChunks=%zu revisionId=%zu " + "createTime=%s updateTime=%s", + (size_t)chunk_meta_.total_size, (size_t)chunk_meta_.chunk_size, + (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_NODE], + (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_UPPER_NEIGHBOR], + (size_t)chunk_meta_.revision_id, create_time, update_time); + + return 0; +} + +int SparseChunkBroker::open(IndexStorage::Pointer stg, size_t max_index_size, + size_t chunk_size, bool check_crc) { + if (ailego_unlikely(stg_)) { + LOG_ERROR("An storage instance is already opened"); + return IndexError_Duplicate; + } + stg_ = std::move(stg); + check_crc_ = check_crc; + max_chunks_size_ = max_index_size; + dirty_ = false; + + const std::string segment_id = + make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId); + chunk_meta_segment_ = stg_->get(segment_id); + if (!chunk_meta_segment_) { + LOG_DEBUG("Create new index"); + return init_storage(chunk_size); + } + + return load_storage(chunk_size); +} + +int SparseChunkBroker::close(void) { + flush(0UL); + + stg_.reset(); + check_crc_ = false; + dirty_ = false; + + return 0; +} + +int SparseChunkBroker::flush(uint64_t checkpoint) { + ailego_assert_with(chunk_meta_segment_, "invalid meta segment"); + + chunk_meta_.update_time = ailego::Realtime::Seconds(); + stats_.set_update_time(chunk_meta_.update_time); + + size_t size = chunk_meta_segment_->write(0UL, &chunk_meta_, + sizeof(HnswSparseChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswSparseChunkMeta))) { + LOG_ERROR("Storage write data failed, wsize=%zu", size); + } + + stg_->refresh(checkpoint); + int ret = stg_->flush(); + if (ret == 0) { + (*stats_.mutable_check_point()) = checkpoint; + } else { + LOG_ERROR("Storage flush failed for %s", IndexError::What(ret)); + } + return ret; +} + +std::pair SparseChunkBroker::alloc_chunk( + int type, uint64_t seq_id, size_t size) { + ailego_assert_with(type < CHUNK_TYPE_MAX, "chunk type overflow"); + + SparseChunk::Pointer chunk; + if (ailego_unlikely(!stg_)) { + LOG_ERROR("Init storage first"); + return std::make_pair(IndexError_Uninitialized, chunk); + } + + //! check exist a empty chunk with the same name + chunk = get_chunk(type, seq_id); + if (chunk) { + if (ailego_unlikely(chunk->capacity() == size && + chunk->data_size() == 0UL)) { + LOG_ERROR("Exist invalid chunk size %zu, expect size %zu", + chunk->capacity(), size); + chunk.reset(); + return std::make_pair(IndexError_Runtime, chunk); + } + return std::make_pair(0, chunk); + } + //! align to page size + size = (size + page_mask_) & (~page_mask_); + if (ailego_unlikely(chunk_meta_.total_size + size >= max_chunks_size_)) { + LOG_ERROR("No space to new a chunk, curIndexSize=%zu allocSize=%zu", + (size_t)chunk_meta_.total_size, size); + return std::make_pair(IndexError_IndexFull, chunk); + } + + std::string segment_id = make_segment_id(type, seq_id); + int ret = stg_->append(segment_id, size); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Storage append segment failed for %s", IndexError::What(ret)); + return std::make_pair(ret, chunk); + } + chunk_meta_.chunk_cnts[type] += 1; + chunk_meta_.total_size += size; + (*stats_.mutable_index_size()) += size; + + size = chunk_meta_segment_->write(0UL, &chunk_meta_, + sizeof(HnswSparseChunkMeta)); + if (ailego_unlikely(size != sizeof(HnswSparseChunkMeta))) { + LOG_ERROR("Storage append segment failed, wsize=%zu", size); + } + + chunk = get_chunk(type, seq_id); + return std::make_pair(chunk ? 0 : IndexError_NoMemory, chunk); +} + +SparseChunk::Pointer SparseChunkBroker::get_chunk(int type, + uint64_t seq_id) const { + std::string segment_id = make_segment_id(type, seq_id); + return stg_->get(segment_id); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.h new file mode 100644 index 0000000..88000ce --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.h @@ -0,0 +1,142 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +using SparseChunk = IndexStorage::Segment; + +class SparseChunkBroker { + public: + typedef std::shared_ptr Pointer; + + enum CHUNK_TYPE { + CHUNK_TYPE_HEADER = 1, + CHUNK_TYPE_META = 2, + CHUNK_TYPE_NODE = 3, + CHUNK_TYPE_UPPER_NEIGHBOR = 4, + CHUNK_TYPE_NEIGHBOR_INDEX = 5, + CHUNK_TYPE_SPARSE_NODE = 6, + CHUNK_TYPE_MAX = 8 + }; + static constexpr size_t kDefaultChunkSeqId = 0UL; + + SparseChunkBroker(IndexStreamer::Stats &stats) : stats_(stats) { + page_mask_ = ailego::MemoryHelper::PageSize() - 1; + } + + //! Open storage + int open(IndexStorage::Pointer stg, size_t max_index_size, size_t chunk_size, + bool check_crc); + + int close(void); + + int flush(uint64_t checkpoint); + + //! alloc a new chunk with size, not thread-safe + std::pair alloc_chunk(int type, uint64_t seq_id, + size_t size); + + //! alloc a new chunk with chunk size + inline std::pair alloc_chunk(int type, + uint64_t seq_id) { + return alloc_chunk(type, seq_id, chunk_meta_.chunk_size); + } + + SparseChunk::Pointer get_chunk(int type, uint64_t seq_id) const; + + inline size_t get_chunk_cnt(int type) const { + ailego_assert_with(type < CHUNK_TYPE_MAX, "chunk type overflow"); + return chunk_meta_.chunk_cnts[type]; + } + + inline bool dirty(void) const { + return dirty_; + } + + inline void mark_dirty(void) { + if (!dirty_) { + dirty_ = true; + chunk_meta_.revision_id += 1; + stats_.set_revision_id(chunk_meta_.revision_id); + } + } + + const IndexStorage::Pointer storage(void) const { + return stg_; + } + + private: + SparseChunkBroker(const SparseChunkBroker &) = delete; + SparseChunkBroker &operator=(const SparseChunkBroker &) = delete; + + struct HnswSparseChunkMeta { + HnswSparseChunkMeta(void) { + memset(static_cast(this), 0, sizeof(HnswSparseChunkMeta)); + } + void clear() { + memset(static_cast(this), 0, sizeof(HnswSparseChunkMeta)); + } + + uint64_t chunk_cnts[CHUNK_TYPE_MAX]; + uint64_t chunk_size; // size of per chunk + uint64_t total_size; // total size of allocated chunk + uint64_t revision_id; // index revision + uint64_t create_time; + uint64_t update_time; + uint64_t reserved[3]; + }; + + static_assert(sizeof(HnswSparseChunkMeta) % 32 == 0, + "HnswSparseChunkMeta must be aligned with 32 bytes"); + + //! Init the storage after open an empty index + int init_storage(size_t chunk_size); + + //! Load index from storage + int load_storage(size_t chunk_size); + + static inline const std::string make_segment_id(int type, uint64_t seq_id) { + return "HnswT" + ailego::StringHelper::ToString(type) + "S" + + ailego::StringHelper::ToString(seq_id); + } + + private: + IndexStreamer::Stats &stats_; + HnswSparseChunkMeta chunk_meta_{}; + size_t page_mask_{0UL}; + size_t max_chunks_size_{0UL}; + IndexStorage::Pointer stg_{}; + IndexStorage::Segment::Pointer chunk_meta_segment_{}; + bool check_crc_{false}; + bool dirty_{false}; // set as true if index is modified , the flag + // will not be cleared even if flushed +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.cc new file mode 100644 index 0000000..946d1e9 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.cc @@ -0,0 +1,294 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_context.h" +#include +#include "hnsw_sparse_params.h" + +namespace zvec { +namespace core { + +HnswSparseContext::HnswSparseContext(const IndexMetric::Pointer &metric, + const HnswSparseEntity::Pointer &entity) + : IndexContext(metric), entity_(entity), dc_(entity_.get(), metric) {} + +HnswSparseContext::~HnswSparseContext() { + visit_filter_.destroy(); +} + +int HnswSparseContext::init(ContextType type) { + int ret; + uint32_t doc_cnt; + + type_ = type; + + switch (type) { + case kSparseBuilderContext: + ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(), + max_scan_num_, negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + update_heap_.limit(entity_->l0_neighbor_cnt() + 1); + break; + + case kSparseSearcherContext: + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_, + negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + break; + + case kSparseStreamerContext: + // maxScanNum is unknown if inited from streamer, so the docCnt may + // change. we need to compute maxScanNum by scan ratio, and preserve + // max_doc_cnt space from visit filter + doc_cnt = entity_->doc_cnt(); + max_scan_num_ = compute_max_scan_num(doc_cnt); + reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt); + ret = visit_filter_.init(filter_mode_, reserve_max_doc_cnt_, + max_scan_num_, negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + + update_heap_.limit(entity_->l0_neighbor_cnt() + 1); + candidates_.limit(max_scan_num_); + + check_need_adjuct_ctx(); + break; + + default: + LOG_ERROR("Init context failed"); + return IndexError_Runtime; + } + + return 0; +} + +int HnswSparseContext::update(const ailego::Params ¶ms) { + LOG_DEBUG("Update hnsw context params"); + + auto update_visit_filter_param = [&]() { + bool need_update = false; + std::string p; + switch (type_) { + case kSparseSearcherContext: + p = PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_ENABLE; + break; + case kSparseStreamerContext: + p = PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_ENABLE; + break; + } + + if (params.has(p)) { + bool bf_enabled = false; + params.get(p, &bf_enabled); + if (bf_enabled ^ (filter_mode_ == VisitFilter::BloomFilter)) { + need_update = true; + filter_mode_ = + bf_enabled ? VisitFilter::BloomFilter : VisitFilter::ByteMap; + } + } + + float prob = negative_probability_; + p.clear(); + switch (type_) { + case kSparseSearcherContext: + p = PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB; + break; + case kSparseStreamerContext: + p = PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB; + break; + } + params.get(p, &prob); + if (filter_mode_ == VisitFilter::BloomFilter && + std::abs(prob - negative_probability_) > 1e-6) { + need_update = true; + } + if (need_update) { + visit_filter_.destroy(); + int max_doc_cnt = 0; + if (type_ == kSparseSearcherContext) { + max_doc_cnt = entity_->doc_cnt(); + } else { + max_doc_cnt = reserve_max_doc_cnt_; + } + int ret = visit_filter_.init(filter_mode_, max_doc_cnt, max_scan_num_, + negative_probability_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + } + return 0; + }; + + switch (type_) { + case kSparseSearcherContext: + if (params.has(PARAM_HNSW_SPARSE_SEARCHER_EF)) { + params.get(PARAM_HNSW_SPARSE_SEARCHER_EF, &ef_); + topk_heap_.limit(std::max(topk_, ef_)); + } + + if (params.has(PARAM_HNSW_SPARSE_SEARCHER_MAX_SCAN_RATIO)) { + params.get(PARAM_HNSW_SPARSE_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_); + max_scan_num_ = + static_cast(max_scan_ratio_ * entity_->doc_cnt()); + max_scan_num_ = std::max(10000U, max_scan_num_); + } + + if (params.has(PARAM_HNSW_SPARSE_SEARCHER_BRUTE_FORCE_THRESHOLD)) { + params.get(PARAM_HNSW_SPARSE_SEARCHER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + } + + return update_visit_filter_param(); + + case kSparseStreamerContext: + if (params.has(PARAM_HNSW_SPARSE_STREAMER_EF)) { + params.get(PARAM_HNSW_SPARSE_STREAMER_EF, &ef_); + topk_heap_.limit(std::max(topk_, ef_)); + } + params.get(PARAM_HNSW_SPARSE_STREAMER_EF, &ef_); + params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_); + params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_); + params.get(PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_); + if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) { + LOG_ERROR("[%s] must be in range (0.0f,1.0f]", + PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO.c_str()); + return IndexError_InvalidArgument; + } + if (max_scan_limit_ < min_scan_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT.c_str(), + PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT.c_str()); + return IndexError_InvalidArgument; + } + + if (params.has(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD)) { + params.get(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + } + + return update_visit_filter_param(); + + default: + LOG_ERROR("update context failed, type=%u", type_); + return IndexError_Runtime; + } +} + + +int HnswSparseContext::update_context(ContextType type, + const IndexMeta & /*meta*/, + const IndexMetric::Pointer &metric, + const HnswSparseEntity::Pointer &entity, + uint32_t magic_num) { + uint32_t doc_cnt; + + if (ailego_unlikely(static_cast(type) != type_)) { + LOG_ERROR( + "HnswSparseContext doesn't support shared by different type, " + "src=%u dst=%u", + type_, type); + return IndexError_Unsupported; + } + + magic_ = kInvalidMgic; + + // TODO: support change filter mode? + switch (type) { + case kSparseBuilderContext: + LOG_ERROR("BuildContext doesn't support update"); + return IndexError_NotImplemented; + + case kSparseSearcherContext: + if (!visit_filter_.reset(entity->doc_cnt(), max_scan_num_)) { + LOG_ERROR("Reset filter failed, mode %d", visit_filter_.get_mode()); + return IndexError_Runtime; + } + + candidates_.limit(max_scan_num_); + topk_heap_.limit(std::max(topk_, ef_)); + break; + + case kSparseStreamerContext: + doc_cnt = entity->doc_cnt(); + max_scan_num_ = compute_max_scan_num(doc_cnt); + reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt); + if (!visit_filter_.reset(reserve_max_doc_cnt_, max_scan_num_)) { + LOG_ERROR("Reset filter failed, mode %d", visit_filter_.get_mode()); + return IndexError_Runtime; + } + + update_heap_.limit(entity->l0_neighbor_cnt() + 1); + candidates_.limit(max_scan_num_); + topk_heap_.limit(std::max(topk_, ef_)); + break; + + default: + LOG_ERROR("update context failed"); + return IndexError_Runtime; + } + + entity_ = entity; + dc_.update(entity_.get(), metric); + magic_ = magic_num; + level_topks_.clear(); + + return 0; +} + +void HnswSparseContext::fill_random_to_topk_full(void) { + static std::mt19937 mt( + std::chrono::system_clock::now().time_since_epoch().count()); + std::uniform_int_distribution dt(0, entity_->doc_cnt() - 1); + std::function gen; + node_id_t seqid; + std::function myfilter = [](node_id_t) { return false; }; + if (this->filter().is_valid()) { + myfilter = [&](node_id_t id) { + return this->filter()(entity_->get_key(id)); + }; + } + + if (topk_heap_.limit() < entity_->doc_cnt() / 2) { + gen = [&](void) { return dt(mt); }; + } else { + // If topk limit is big value, gen sequential id from an random initial + seqid = dt(mt); + gen = [&](void) { + seqid = seqid == (entity_->doc_cnt() - 1) ? 0 : (seqid + 1); + return seqid; + }; + } + + for (size_t i = 0; !topk_heap_.full() && i < entity_->doc_cnt(); ++i) { + const auto id = gen(); + if (!visit_filter_.visited(id) && !myfilter(id)) { + visit_filter_.set_visited(id); + topk_heap_.emplace(id, dc_.dist(id)); + } + } + return; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h new file mode 100644 index 0000000..56ae1cb --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h @@ -0,0 +1,538 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "utility/sparse_utility.h" +#include "utility/visit_filter.h" +#include "hnsw_sparse_dist_calculator.h" + +namespace zvec { +namespace core { + +class HnswSparseContext : public IndexContext { + public: + //! Index Context Pointer + typedef std::unique_ptr Pointer; + + enum ContextType { + kUnknownContext = 0, + kSparseSearcherContext = 1, + kSparseBuilderContext = 2, + kSparseStreamerContext = 3, + }; + + //! Construct + HnswSparseContext(const IndexMetric::Pointer &metric, + const HnswSparseEntity::Pointer &entity); + + //! Destructor + ~HnswSparseContext() override; + + public: + //! Set topk of search result + void set_topk(uint32_t val) override { + topk_ = val; + topk_heap_.limit(std::max(val, ef_)); + } + + //! Retrieve search result + const IndexDocumentList &result(void) const override { + return results_[0]; + } + + //! Retrieve search result + const IndexDocumentList &result(size_t idx) const override { + return results_[idx]; + } + + //! Retrieve result object for output + IndexDocumentList *mutable_result(size_t idx) override { + ailego_assert_with(idx < results_.size(), "invalid idx"); + return &results_[idx]; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(void) const override { + return group_results_[0]; + } + + //! Retrieve search group result with index + const IndexGroupDocumentList &group_result(size_t idx) const override { + return group_results_[idx]; + } + + IndexGroupDocumentList *mutable_group_result(void) override { + return &group_results_[0]; + } + + IndexGroupDocumentList *mutable_group_result(size_t idx) override { + return &group_results_[idx]; + } + + uint32_t magic(void) const override { + return magic_; + } + + //! Set mode of debug + void set_debug_mode(bool enable) override { + debug_mode_ = enable; + } + + //! Retrieve mode of debug + bool debug_mode(void) const override { + return this->debugging(); + } + + //! Retrieve string of debug + std::string debug_string(void) const override { + char buf[4096]; + size_t size = snprintf( + buf, sizeof(buf), + "scan_cnt=%zu,get_vector_cnt=%u,get_neighbors_cnt=%u,dup_node=%u", + get_scan_num(), stats_get_vector_cnt_, stats_get_neighbors_cnt_, + stats_visit_dup_cnt_); + return std::string(buf, size); + } + + //! Update the parameters of context + int update(const ailego::Params ¶ms) override; + + public: + //! Init context + int init(ContextType type); + + //! Update context, the context may be shared by different searcher/streamer + int update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const HnswSparseEntity::Pointer &entity, + uint32_t magic_num); + + inline const HnswSparseEntity &get_entity() const { + return *entity_; + } + + inline void resize_results(size_t size) { + if (group_by_search()) { + group_results_.resize(size); + } else { + results_.resize(size); + } + } + + inline void topk_to_result() { + return topk_to_result(0); + } + + //! Construct result from topk heap, result will be normalized + inline void topk_to_result(uint32_t idx) { + if (group_by_search()) { + topk_to_group_result(idx); + } else { + topk_to_single_result(idx); + } + } + + inline void recal_topk_dist() { + TopkHeap heap(topk_heap_); + topk_heap_.clear(); + + for (size_t i = 0; i < heap.size(); ++i) { + node_id_t id = heap[i].first; + dist_t dist = dc_.dist(id); + topk_heap_.emplace_back(id, dist); + } + } + + inline void topk_to_single_result(uint32_t idx) { + if (force_padding_topk_ && !topk_heap_.full() && + topk_heap_.size() < entity_->doc_cnt()) { + this->fill_random_to_topk_full(); + } + if (ailego_unlikely(topk_heap_.size() == 0)) { + return; + } + + ailego_assert_with(idx < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(topk_heap_.size())); + topk_heap_.sort(); + results_[idx].clear(); + + for (int i = 0; i < size; ++i) { + auto score = topk_heap_[i].second; + if (score > this->threshold()) { + break; + } + + node_id_t id = topk_heap_[i].first; + if (fetch_vector_) { + IndexSparseDocument sparse_doc; + IndexStorage::MemoryBlock vec_block; + entity_->get_sparse_data(id, vec_block); + const void *sparse_data = vec_block.data(); + if (sparse_data != nullptr) { + SparseUtility::ReverseSparseFormat(sparse_data, sparse_doc, + entity_->sparse_unit_size()); + } + + results_[idx].emplace_back(entity_->get_key(id), score, id, + entity_->get_vector_meta(id), sparse_doc); + } else { + results_[idx].emplace_back(entity_->get_key(id), score, id); + } + } + + return; + } + + //! Construct result from topk heap, result will be normalized + inline void topk_to_group_result(uint32_t idx) { + ailego_assert_with(idx < group_results_.size(), "invalid idx"); + + group_results_[idx].clear(); + + std::vector> group_topk_list; + std::vector> best_score_in_groups; + for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end(); + itr++) { + const std::string &group_id = (*itr).first; + auto &heap = (*itr).second; + heap.sort(); + + if (heap.size() > 0) { + float best_score = heap[0].second; + best_score_in_groups.push_back(std::make_pair(group_id, best_score)); + } + } + + std::sort(best_score_in_groups.begin(), best_score_in_groups.end(), + [](const std::pair &a, + const std::pair &b) -> int { + return a.second < b.second; + }); + + // truncate to group num + for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size(); + ++i) { + const std::string &group_id = best_score_in_groups[i].first; + + group_topk_list.emplace_back( + std::make_pair(group_id, group_topk_heaps_[group_id])); + } + + group_results_[idx].resize(group_topk_list.size()); + + for (uint32_t i = 0; i < group_topk_list.size(); ++i) { + const std::string &group_id = group_topk_list[i].first; + group_results_[idx][i].set_group_id(group_id); + + uint32_t size = std::min( + group_topk_, static_cast(group_topk_list[i].second.size())); + + for (uint32_t j = 0; j < size; ++j) { + auto score = group_topk_list[i].second[j].second; + if (score > this->threshold()) { + break; + } + + node_id_t id = group_topk_list[i].second[j].first; + + if (fetch_vector_) { + IndexSparseDocument sparse_doc; + IndexStorage::MemoryBlock vec_block; + entity_->get_sparse_data(id, vec_block); + const void *sparse_data = vec_block.data(); + if (sparse_data != nullptr) { + SparseUtility::ReverseSparseFormat(sparse_data, sparse_doc, + entity_->sparse_unit_size()); + } + group_results_[idx][i].mutable_docs()->emplace_back( + entity_->get_key(id), score, id, entity_->get_vector_meta(id), + sparse_doc); + } else { + group_results_[idx][i].mutable_docs()->emplace_back( + entity_->get_key(id), score, id); + } + } + } + } + + inline void reset_query(const void *query) { + dc_.reset_query(query); + dc_.clear_compare_cnt(); + } + + inline HnswSparseDistCalculator &dist_calculator() { + return dc_; + } + + inline TopkHeap &topk_heap() { + return topk_heap_; + } + + inline TopkHeap &update_heap() { + return update_heap_; + } + + inline VisitFilter &visit_filter() { + return visit_filter_; + } + + inline CandidateHeap &candidates() { + return candidates_; + } + + inline void set_max_scan_num(uint32_t max_scan_num) { + max_scan_num_ = max_scan_num; + } + + inline void set_max_scan_limit(uint32_t max_scan_limit) { + max_scan_limit_ = max_scan_limit; + } + + inline void set_min_scan_limit(uint32_t min_scan_limit) { + min_scan_limit_ = min_scan_limit; + } + + inline void set_ef(uint32_t v) { + ef_ = v; + } + + inline void set_filter_mode(uint32_t v) { + filter_mode_ = v; + } + + inline void set_filter_negative_probability(float v) { + negative_probability_ = v; + } + + inline void set_max_scan_ratio(float v) { + max_scan_ratio_ = v; + } + + virtual void set_magic(uint32_t v) { + magic_ = v; + } + + virtual void set_force_padding_topk(bool v) { + force_padding_topk_ = v; + } + + void set_bruteforce_threshold(uint32_t v) override { + bruteforce_threshold_ = v; + } + + inline uint32_t get_bruteforce_threshold() const { + return bruteforce_threshold_; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + + bool fetch_vector() const override { + return fetch_vector_; + } + + //! Reset context + void reset(void) override { + set_filter(nullptr); + reset_threshold(); + set_fetch_vector(false); + set_group_params(0, 0); + reset_group_by(); + } + + inline std::map &group_topk_heaps() { + return group_topk_heaps_; + } + + inline TopkHeap &level_topk(int level) { + if (ailego_unlikely(level_topks_.size() <= static_cast(level))) { + int cur_level = level_topks_.size(); + level_topks_.resize(level + 1); + for (; cur_level <= level; ++cur_level) { + size_t heap_size = std::max(entity_->neighbor_cnt(cur_level), + entity_->ef_construction()); + level_topks_[cur_level].clear(); + level_topks_[cur_level].limit(heap_size); + } + } + + return level_topks_[level]; + } + + inline void check_need_adjuct_ctx(void) { + check_need_adjuct_ctx(entity_->doc_cnt()); + } + + inline size_t compute_reserve_cnt(uint32_t cur_doc) const { + if (cur_doc > kMaxReserveDocCnt) { + return kMaxReserveDocCnt; + } else if (cur_doc < kMinReserveDocCnt) { + return kMinReserveDocCnt; + } + return cur_doc; + } + + //! candidates heap and visitfilter need to resize as doc cnt growing up + inline void check_need_adjuct_ctx(uint32_t doc_cnt) { + if (ailego_unlikely(doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_)) { + while (doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_) { + reserve_max_doc_cnt_ = + reserve_max_doc_cnt_ + compute_reserve_cnt(reserve_max_doc_cnt_); + } + uint32_t max_scan_cnt = compute_max_scan_num(reserve_max_doc_cnt_); + max_scan_num_ = max_scan_cnt; + visit_filter_.reset(reserve_max_doc_cnt_, max_scan_cnt); + candidates_.clear(); + candidates_.limit(max_scan_num_); + } + } + + inline uint32_t compute_max_scan_num(uint32_t max_doc_cnt) const { + uint32_t max_scan = max_doc_cnt * max_scan_ratio_; + if (max_scan < min_scan_limit_) { + max_scan = min_scan_limit_; + } else if (max_scan > max_scan_limit_) { + max_scan = max_scan_limit_; + } + return max_scan; + } + + inline size_t get_scan_num() const { + return dc_.compare_cnt(); + } + + inline uint64_t reach_scan_limit() const { + return dc_.compare_cnt() >= max_scan_num_; + } + + inline bool error() const { + return dc_.error(); + } + + inline void clear() { + dc_.clear(); + if (ailego_unlikely(this->debugging())) { + stats_get_neighbors_cnt_ = 0u; + stats_get_vector_cnt_ = 0u; + stats_visit_dup_cnt_ = 0u; + } + // do not clear results_ for the next query will need it + for (auto &it : results_) { + it.clear(); + } + } + + uint32_t *mutable_stats_get_neighbors() { + return &stats_get_neighbors_cnt_; + } + + uint32_t *mutable_stats_get_vector() { + return &stats_get_vector_cnt_; + } + + uint32_t *mutable_stats_visit_dup_cnt() { + return &stats_visit_dup_cnt_; + } + + inline bool debugging(void) const { + return debug_mode_; + } + + inline void update_dist_caculator_distance( + const IndexMetric::MatrixSparseDistance &distance) { + dc_.update_distance(distance); + } + + //! Get topk + inline uint32_t topk() const override { + return topk_; + } + + //! Get group topk + inline uint32_t group_topk() const { + return group_topk_; + } + + //! Get group num + inline uint32_t group_num() const { + return group_num_; + } + + //! Get if group by search + inline bool group_by_search() { + return group_num_ > 0; + } + + //! Set group params + void set_group_params(uint32_t group_num, uint32_t group_topk) override { + group_num_ = group_num; + group_topk_ = group_topk; + + topk_ = group_topk_ * group_num_; + + topk_heap_.limit(std::max(topk_, ef_)); + + group_topk_heaps_.clear(); + } + + private: + // Filling random nodes if topk not full + void fill_random_to_topk_full(void); + + constexpr static uint32_t kTriggerReserveCnt = 4096UL; + constexpr static uint32_t kMinReserveDocCnt = 4096UL; + constexpr static uint32_t kMaxReserveDocCnt = 128 * 1024UL; + constexpr static uint32_t kInvalidMgic = -1U; + + private: + HnswSparseEntity::Pointer entity_; + HnswSparseDistCalculator dc_; + bool debug_mode_{false}; + bool force_padding_topk_{false}; + uint32_t max_scan_num_{0}; + uint32_t max_scan_limit_{0}; + uint32_t min_scan_limit_{0}; + uint32_t reserve_max_doc_cnt_{kMinReserveDocCnt}; + uint32_t topk_{0}; + uint32_t group_topk_{0}; + uint32_t filter_mode_{VisitFilter::ByteMap}; + float negative_probability_{HnswSparseEntity::kDefaultBFNegativeProbability}; + uint32_t ef_{HnswSparseEntity::kDefaultEf}; + float max_scan_ratio_{HnswSparseEntity::kDefaultScanRatio}; + uint32_t magic_{0U}; + std::vector results_{}; + std::vector group_results_{}; + TopkHeap topk_heap_{}; + TopkHeap update_heap_{}; + std::vector level_topks_{}; + CandidateHeap candidates_{}; + VisitFilter visit_filter_{}; + uint32_t bruteforce_threshold_{}; + bool fetch_vector_{false}; + + uint32_t group_num_{0}; + std::map group_topk_heaps_{}; + + uint32_t type_{kUnknownContext}; + //! debug stats info + uint32_t stats_get_neighbors_cnt_{0u}; + uint32_t stats_get_vector_cnt_{0u}; + uint32_t stats_visit_dup_cnt_{0u}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_dist_calculator.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_dist_calculator.h new file mode 100644 index 0000000..5bbfc1b --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_dist_calculator.h @@ -0,0 +1,184 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "hnsw_sparse_entity.h" + +namespace zvec { +namespace core { + +class HnswSparseDistCalculator { + public: + typedef std::shared_ptr Pointer; + + public: + //! Constructor + HnswSparseDistCalculator(const HnswSparseEntity *entity, + const IndexMetric::Pointer &metric) + : entity_(entity), + distance_(metric->sparse_distance()), + query_{nullptr}, + compare_cnt_(0) {} + + //! Constructor + HnswSparseDistCalculator(const HnswSparseEntity *entity, + const IndexMetric::Pointer &metric, + const void *query) + : entity_(entity), + distance_(metric->sparse_distance()), + query_(query), + compare_cnt_(0) {} + + void update(const HnswSparseEntity *entity, + const IndexMetric::Pointer &metric) { + entity_ = entity; + distance_ = metric->sparse_distance(); + } + + inline void update_distance( + const IndexMetric::MatrixSparseDistance &distance) { + distance_ = distance; + } + + //! Reset query vector data + inline void reset_query(const void *query) { + error_ = false; + query_ = query; + } + + //! Returns distance + inline dist_t dist(const void *sparse_data_lhs, const void *sparse_data_rhs) { + float score{0.0f}; + + if (ailego_unlikely(sparse_data_lhs == nullptr || + sparse_data_rhs == nullptr)) { + // LOG_WARN("Nullptr of sparse vector. Return dense score only"); + // error_ = true; + return score; + } + + distance_(sparse_data_lhs, sparse_data_rhs, &score); + + return score; + } + + //! Returns distance between query and vec. + inline dist_t dist(const void *vec) { + compare_cnt_++; + + auto sparse_data = entity_->get_sparse_data_from_vector(vec); + if (sparse_data.first == nullptr) { + error_ = true; + return 0.0f; + } + + return dist(sparse_data.first, query_); + } + + //! Return distance between query and node id. + inline dist_t dist(node_id_t id) { + compare_cnt_++; + + const void *feat = entity_->get_vector_meta(id); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + + auto sparse_data = entity_->get_sparse_data_from_vector(feat); + if (sparse_data.first == nullptr) { + error_ = true; + return 0.0f; + } + + return dist(sparse_data.first, query_); + } + + //! Return dist node lhs between node rhs + inline dist_t dist(node_id_t lhs, node_id_t rhs) { + compare_cnt_++; + + const void *feat = entity_->get_vector_meta(lhs); + const void *query = entity_->get_vector_meta(rhs); + if (ailego_unlikely(feat == nullptr || query == nullptr)) { + LOG_ERROR("Get nullptr vector"); + error_ = true; + return 0.0f; + } + + auto feat_sparse_data = entity_->get_sparse_data_from_vector(feat); + if (feat_sparse_data.first == nullptr) { + error_ = true; + return 0.0f; + } + + auto query_sparse_data = entity_->get_sparse_data_from_vector(query); + if (query_sparse_data.first == nullptr) { + error_ = true; + return 0.0f; + } + + return dist(feat_sparse_data.first, query_sparse_data.first); + } + + dist_t operator()(const void *vec) { + return dist(vec); + } + + dist_t operator()(node_id_t i) { + return dist(i); + } + + dist_t operator()(node_id_t lhs, node_id_t rhs) { + return dist(lhs, rhs); + } + + inline void clear() { + compare_cnt_ = 0; + error_ = false; + } + + inline void clear_compare_cnt() { + compare_cnt_ = 0; + } + + inline bool error() const { + return error_; + } + + //! Get distances compute times + inline uint32_t compare_cnt() const { + return compare_cnt_; + } + + private: + HnswSparseDistCalculator(const HnswSparseDistCalculator &) = delete; + HnswSparseDistCalculator &operator=(const HnswSparseDistCalculator &) = + delete; + + private: + const HnswSparseEntity *entity_; + + IndexMetric::MatrixSparseDistance distance_; + + const void *query_; + + uint32_t compare_cnt_; // record distance compute times + bool error_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.cc new file mode 100644 index 0000000..29ec39b --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.cc @@ -0,0 +1,443 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "hnsw_sparse_entity.h" +#include + +namespace zvec { +namespace core { + +const std::string HnswSparseEntity::kSparseGraphHeaderSegmentId = + "sparse_graph.header"; +const std::string HnswSparseEntity::kSparseGraphFeaturesSegmentId = + "sparse_graph.features"; +const std::string HnswSparseEntity::kSparseGraphKeysSegmentId = + "sparse_graph.keys"; +const std::string HnswSparseEntity::kSparseGraphNeighborsSegmentId = + "sparse_graph.neighbors"; +const std::string HnswSparseEntity::kSparseGraphOffsetsSegmentId = + "sparse_graph.offsets"; +const std::string HnswSparseEntity::kSparseGraphMappingSegmentId = + "sparse_graph.mapping"; +const std::string HnswSparseEntity::kSparseHnswHeaderSegmentId = + "sparse_hnsw.header"; +const std::string HnswSparseEntity::kSparseHnswNeighborsSegmentId = + "sparse_hnsw.neighbors"; +const std::string HnswSparseEntity::kSparseHnswOffsetsSegmentId = + "sparse_hnsw.offsets"; +const std::string HnswSparseEntity::kSparseGraphVectorsSegmentId = + "sparse_graph.vectors"; +const std::string HnswSparseEntity::kSparseGraphVectorMetaSegmentId = + "sparse_graph.vector_meta"; + +int HnswSparseEntity::CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, + size_t *padding_size) { + *padding_size = AlignSize(data_size) - data_size; + if (*padding_size == 0) { + return 0; + } + + std::string padding(*padding_size, '\0'); + if (dumper->write(padding.data(), *padding_size) != *padding_size) { + LOG_ERROR("Append padding failed, size %lu", *padding_size); + return IndexError_WriteData; + } + return 0; +} + +int64_t HnswSparseEntity::dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, + const void *data, size_t size) const { + size_t len = dumper->write(data, size); + if (len != size) { + LOG_ERROR("Dump segment %s data failed, expect: %lu, actual: %lu", + segment_id.c_str(), size, len); + return IndexError_WriteData; + } + + size_t padding_size = AlignSize(size) - size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + uint32_t crc = ailego::Crc32c::Hash(data, size); + int ret = dumper->append(segment_id, size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", segment_id.c_str(), ret); + return ret; + } + + return len + padding_size; +} + +int64_t HnswSparseEntity::dump_header(const IndexDumper::Pointer &dumper, + const HNSWSparseHeader &hd) const { + //! dump basic graph header. header is aligned and does not need padding + int64_t graph_hd_size = dump_segment(dumper, kSparseGraphHeaderSegmentId, + &hd.graph, hd.graph.size); + if (graph_hd_size < 0) { + return graph_hd_size; + } + + //! dump basic graph header. header is aligned and does not need padding + int64_t hnsw_hd_size = + dump_segment(dumper, kSparseHnswHeaderSegmentId, &hd.hnsw, hd.hnsw.size); + if (hnsw_hd_size < 0) { + return hnsw_hd_size; + } + + return graph_hd_size + hnsw_hd_size; +} + +void HnswSparseEntity::reshuffle_vectors( + const std::function & /*get_level*/, + std::vector * /*n2o_mapping*/, + std::vector * /*o2n_mapping*/, key_t * /*keys*/) const { + // TODO + return; +} + +int64_t HnswSparseEntity::dump_mapping_segment( + const IndexDumper::Pointer &dumper, const key_t *keys) const { + std::vector mapping(doc_cnt()); + + std::iota(mapping.begin(), mapping.end(), 0U); + std::sort(mapping.begin(), mapping.end(), + [&](node_id_t i, node_id_t j) { return keys[i] < keys[j]; }); + + size_t size = mapping.size() * sizeof(node_id_t); + return dump_segment(dumper, kSparseGraphMappingSegmentId, mapping.data(), + size); +} + +int64_t HnswSparseEntity::dump_segments( + const IndexDumper::Pointer &dumper, key_t *keys, + const std::function &get_level) const { + HNSWSparseHeader dump_hd(header()); + + dump_hd.graph.node_size = sparse_meta_size(); + + std::vector n2o_mapping; // map new id to origin id + std::vector o2n_mapping; // map origin id to new id + reshuffle_vectors(get_level, &n2o_mapping, &o2n_mapping, keys); + if (!o2n_mapping.empty()) { + dump_hd.hnsw.entry_point = o2n_mapping[entry_point()]; + } + + //! Dump header + int64_t hd_size = dump_header(dumper, dump_hd); + if (hd_size < 0) { + return hd_size; + } + + //! Dump vectors + int64_t sparse_vector_meta_size = + dump_sparse_vector_meta(dumper, n2o_mapping); + if (sparse_vector_meta_size < 0) { + return sparse_vector_meta_size; + } + + int64_t sparse_vecs_size = dump_sparse_vector(dumper, n2o_mapping); + if (sparse_vecs_size < 0) { + return sparse_vecs_size; + } + + //! Dump neighbors + auto neighbors_size = + dump_neighbors(dumper, get_level, n2o_mapping, o2n_mapping); + if (neighbors_size < 0) { + return neighbors_size; + } + //! free memory + n2o_mapping = std::vector(); + o2n_mapping = std::vector(); + + //! Dump keys + size_t key_segment_size = doc_cnt() * sizeof(key_t); + int64_t keys_size = + dump_segment(dumper, kSparseGraphKeysSegmentId, keys, key_segment_size); + if (keys_size < 0) { + return keys_size; + } + + //! Dump mapping + int64_t mapping_size = dump_mapping_segment(dumper, keys); + if (mapping_size < 0) { + return mapping_size; + } + + return hd_size + keys_size + sparse_vector_meta_size + sparse_vecs_size + + neighbors_size + mapping_size; +} + + +int64_t HnswSparseEntity::dump_sparse_vector_meta( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + const void *data = nullptr; + uint32_t crc = 0U; + size_t dump_size = 0UL; + + uint64_t sparse_data_offset = 0UL; + uint64_t sparse_data_len = 0UL; + + //! dump vectors + for (node_id_t id = 0; id < doc_cnt(); ++id) { + data = get_vector_meta(reorder_mapping.empty() ? id : reorder_mapping[id]); + if (ailego_unlikely(!data)) { + return IndexError_ReadData; + } + + const char *data_ptr = reinterpret_cast(data); + sparse_data_len = *((uint32_t *)(data_ptr + sizeof(uint64_t))); + + size_t len = dumper->write(&sparse_data_offset, sizeof(uint64_t)); + if (len != sizeof(uint64_t)) { + LOG_ERROR("Dump sparse data offset failed, write=%zu expect=%zu", len, + sizeof(uint64_t)); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(&sparse_data_offset, sizeof(uint64_t), crc); + dump_size += sizeof(uint64_t); + + len = dumper->write(&sparse_data_len, sizeof(uint64_t)); + if (len != sizeof(uint64_t)) { + LOG_ERROR("Dump sparse data len failed, write=%zu expect=%zu", len, + sizeof(uint64_t)); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(&sparse_data_len, sizeof(uint64_t), crc); + dump_size += sizeof(uint64_t); + + sparse_data_offset += sparse_data_len; + } + + int ret = + dumper->append(kSparseGraphVectorMetaSegmentId, dump_size, 0UL, crc); + if (ret != 0) { + LOG_ERROR("Dump vectors segment meta failed, ret %d", ret); + return ret; + } + + return dump_size; +} + +int64_t HnswSparseEntity::dump_sparse_vector( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + uint32_t crc = 0U; + size_t data_size = 0UL; + const void *data = nullptr; + + uint64_t sparse_data_len = 0UL; + uint32_t sparse_chunk_index = 0U; + uint32_t sparse_chunk_offset = 0U; + + //! dump vectors + for (node_id_t id = 0; id < doc_cnt(); ++id) { + data = get_vector_meta(reorder_mapping.empty() ? id : reorder_mapping[id]); + if (ailego_unlikely(!data)) { + return IndexError_ReadData; + } + + const char *data_ptr = reinterpret_cast(data); + + sparse_data_len = *((uint32_t *)(data_ptr + sizeof(uint64_t))); + + uint64_t sparse_offset = *((uint64_t *)(data_ptr)); + + const void *sparse = get_sparse_data(sparse_offset, sparse_data_len); + if (ailego_unlikely(sparse == nullptr)) { + LOG_ERROR("Get nullptr sparse, chunk index=%u, chunk offset=%u, len=%zu", + sparse_chunk_index, sparse_chunk_offset, + (size_t)sparse_data_len); + return IndexError_ReadData; + } + + size_t len = dumper->write(sparse, sparse_data_len); + if (len != sparse_data_len) { + LOG_ERROR("Dump sparse data failed, write=%zu expect=%zu", len, + (size_t)sparse_data_len); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(sparse, sparse_data_len, crc); + data_size += sparse_data_len; + } + + int ret = dumper->append(kSparseGraphVectorsSegmentId, data_size, 0UL, crc); + if (ret != 0) { + LOG_ERROR("Dump vectors segment meta failed, ret %d", ret); + return ret; + } + + return data_size; +} + +int64_t HnswSparseEntity::dump_graph_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + std::vector graph_meta; + graph_meta.reserve(doc_cnt()); + size_t offset = 0; + uint32_t crc = 0; + std::vector mapping(l0_neighbor_cnt()); + + uint32_t min_neighbor_count = 10000; + uint32_t max_neighbor_count = 0; + size_t sum_neighbor_count = 0; + + for (node_id_t id = 0; id < doc_cnt(); ++id) { + const Neighbors neighbors = + get_neighbors(0, reorder_mapping.empty() ? id : reorder_mapping[id]); + ailego_assert_with(!!neighbors.data, "invalid neighbors"); + ailego_assert_with(neighbors.size() <= l0_neighbor_cnt(), + "invalid neighbors"); + + uint32_t neighbor_count = neighbors.size(); + if (neighbor_count < min_neighbor_count) { + min_neighbor_count = neighbor_count; + } + if (neighbor_count > max_neighbor_count) { + max_neighbor_count = neighbor_count; + } + sum_neighbor_count += neighbor_count; + + graph_meta.emplace_back(offset, neighbor_count); + size_t size = neighbors.size() * sizeof(node_id_t); + const node_id_t *data = &neighbors[0]; + if (!neighbor_mapping.empty()) { + for (node_id_t i = 0; i < neighbors.size(); ++i) { + mapping[i] = neighbor_mapping[neighbors[i]]; + } + data = mapping.data(); + } + if (dumper->write(data, size) != size) { + LOG_ERROR("Dump graph neighbor id=%u failed, size %lu", id, size); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(data, size, crc); + offset += size; + } + + uint32_t average_neighbor_count = 0; + if (doc_cnt() > 0) { + average_neighbor_count = sum_neighbor_count / doc_cnt(); + } + LOG_INFO( + "Dump hnsw graph: min_neighbor_count[%u] max_neighbor_count[%u] " + "average_neighbor_count[%u]", + min_neighbor_count, max_neighbor_count, average_neighbor_count); + + size_t padding_size = 0; + int ret = CalcAndAddPadding(dumper, offset, &padding_size); + if (ret != 0) { + return ret; + } + ret = + dumper->append(kSparseGraphNeighborsSegmentId, offset, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s failed, ret %d", + kSparseGraphNeighborsSegmentId.c_str(), ret); + return ret; + } + + //! dump level 0 neighbors meta + auto len = + dump_segment(dumper, kSparseGraphOffsetsSegmentId, graph_meta.data(), + graph_meta.size() * sizeof(SparseGraphNeighborMeta)); + if (len < 0) { + return len; + } + + return len + offset + padding_size; +} + +int64_t HnswSparseEntity::dump_upper_neighbors( + const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + std::vector hnsw_meta; + hnsw_meta.reserve(doc_cnt()); + size_t offset = 0; + uint32_t crc = 0; + std::vector buffer(upper_neighbor_cnt() + 1); + size_t buffer_bytes = buffer.size() * sizeof(node_id_t); + for (node_id_t id = 0; id < doc_cnt(); ++id) { + node_id_t new_id = reorder_mapping.empty() ? id : reorder_mapping[id]; + auto level = get_level(new_id); + if (level == 0) { + hnsw_meta.emplace_back(0U, 0U); + continue; + } + hnsw_meta.emplace_back(offset, level); + ailego_assert_with((size_t)level < kMaxGraphLayers, "invalid level"); + for (level_t cur_level = 1; cur_level <= level; ++cur_level) { + const Neighbors neighbors = get_neighbors(cur_level, new_id); + ailego_assert_with(!!neighbors.data, "invalid neighbors"); + ailego_assert_with(neighbors.size() <= neighbor_cnt(cur_level), + "invalid neighbors"); + memset(buffer.data(), 0, buffer_bytes); + buffer[0] = neighbors.size(); + if (neighbor_mapping.empty()) { + memcpy(&buffer[1], &neighbors[0], neighbors.size() * sizeof(node_id_t)); + } else { + for (node_id_t i = 0; i < neighbors.size(); ++i) { + buffer[i + 1] = neighbor_mapping[neighbors[i]]; + } + } + if (dumper->write(buffer.data(), buffer_bytes) != buffer_bytes) { + LOG_ERROR("Dump graph neighbor id=%u failed, size %lu", id, + buffer_bytes); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(buffer.data(), buffer_bytes, crc); + offset += buffer_bytes; + } + } + size_t padding_size = 0; + int ret = CalcAndAddPadding(dumper, offset, &padding_size); + if (ret != 0) { + return ret; + } + + ret = + dumper->append(kSparseHnswNeighborsSegmentId, offset, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s failed, ret %d", + kSparseHnswNeighborsSegmentId.c_str(), ret); + return ret; + } + + //! dump level 0 neighbors meta + auto len = dump_segment(dumper, kSparseHnswOffsetsSegmentId, hnsw_meta.data(), + hnsw_meta.size() * sizeof(HnswSparseNeighborMeta)); + if (len < 0) { + return len; + } + + return len + offset + padding_size; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h new file mode 100644 index 0000000..7e6f814 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h @@ -0,0 +1,625 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +using node_id_t = uint32_t; +using key_t = uint64_t; +using level_t = int32_t; +using dist_t = float; +using TopkHeap = ailego::KeyValueHeap; +using CandidateHeap = + ailego::KeyValueHeap>; +constexpr node_id_t kInvalidNodeId = static_cast(-1); +constexpr key_t kInvalidKey = static_cast(-1); +class HnswSparseDistCalculator; + +struct SparseGraphHeader { + uint32_t size; + uint32_t version; + uint32_t graph_type; + uint32_t doc_count; + uint32_t vector_size; + uint32_t node_size; + uint32_t l0_neighbor_count; + uint32_t prune_type; + uint32_t prune_neighbor_count; + uint32_t ef_construction; + uint32_t options; + uint32_t min_neighbor_count; + uint32_t sparse_meta_size; + uint32_t sparse_unit_size; + uint32_t total_sparse_count; + uint8_t reserved[868]; +}; + +static_assert(sizeof(SparseGraphHeader) % 32 == 0, + "SparseGraphHeader must be aligned with 32 bytes"); + +//! Hnsw upper neighbor header +struct HnswSparseHeader { + uint32_t size; // header size + uint32_t revision; // current total docs of the graph + uint32_t upper_neighbor_count; + uint32_t ef_construction; + uint32_t scaling_factor; + uint32_t max_level; + uint32_t entry_point; + uint32_t options; + uint8_t reserved[30]; +}; + +struct SparseData { + public: + SparseData() {}; + + SparseData(uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_vec) + : count(sparse_count), indices(sparse_indices), vec(sparse_vec) {} + + uint32_t count{0}; + const uint32_t *indices{nullptr}; + const void *vec{nullptr}; +}; + +static_assert(sizeof(HnswSparseHeader) % 32 == 0, + "SparseGraphHeader must be aligned with 32 bytes"); + +//! Hnsw common header and upper neighbor header +struct HNSWSparseHeader { + HNSWSparseHeader() { + clear(); + } + + HNSWSparseHeader(const HNSWSparseHeader &header) { + memcpy(static_cast(this), &header, sizeof(header)); + } + + HNSWSparseHeader &operator=(const HNSWSparseHeader &header) { + memcpy(static_cast(this), &header, sizeof(header)); + return *this; + } + + //! Reset state to zero, and the params is untouched + void inline reset() { + graph.doc_count = 0U; + hnsw.entry_point = kInvalidNodeId; + hnsw.max_level = 0; + graph.total_sparse_count = 0U; + } + + //! Clear all fields to init value + void inline clear() { + memset(static_cast(this), 0, sizeof(HNSWSparseHeader)); + hnsw.entry_point = kInvalidNodeId; + graph.size = sizeof(SparseGraphHeader); + hnsw.size = sizeof(HnswSparseHeader); + graph.total_sparse_count = 0U; + } + + size_t neighbor_cnt() const { + return graph.l0_neighbor_count; + } + + size_t upper_neighbor_cnt() const { + return hnsw.upper_neighbor_count; + } + + size_t vector_size() const { + return graph.vector_size; + } + + size_t ef_construction() const { + return graph.ef_construction; + } + + size_t scaling_factor() const { + return hnsw.scaling_factor; + } + + size_t neighbor_prune_cnt() const { + return graph.prune_neighbor_count; + } + + node_id_t entry_point() const { + return hnsw.entry_point; + } + + node_id_t doc_cnt() const { + return graph.doc_count; + } + + uint32_t total_sparse_count() const { + return graph.total_sparse_count; + } + + SparseGraphHeader graph; + HnswSparseHeader hnsw; +}; + +struct NeighborsHeader { + uint32_t neighbor_cnt; +#ifdef _MSC_VER + node_id_t neighbors[]; +#else + node_id_t neighbors[0]; +#endif +}; + +struct Neighbors { + Neighbors() : cnt{0}, data{nullptr} {} + + Neighbors(uint32_t cnt_in, const node_id_t *data_in) + : cnt{cnt_in}, data{data_in} {} + + Neighbors(IndexStorage::MemoryBlock &&mem_block) + : neighbor_block{std::move(mem_block)} { + auto hd = reinterpret_cast(neighbor_block.data()); + cnt = hd->neighbor_cnt; + data = hd->neighbors; + } + + size_t size(void) const { + return cnt; + } + + const node_id_t &operator[](size_t idx) const { + return data[idx]; + } + + uint32_t cnt; + const node_id_t *data; + IndexStorage::MemoryBlock neighbor_block; +}; + +//! level 0 neighbors offset +struct SparseGraphNeighborMeta { + SparseGraphNeighborMeta(size_t o, size_t cnt) + : offset(o), neighbor_cnt(cnt) {} + + uint64_t offset : 48; + uint64_t neighbor_cnt : 16; +}; + +//! hnsw upper neighbors meta +struct HnswSparseNeighborMeta { + HnswSparseNeighborMeta(size_t o, size_t l) : offset(o), level(l) {} + + uint64_t offset : 48; // offset = idx * upper neighors size + uint64_t level : 16; +}; + +class HnswSparseEntity { + public: + //! Constructor + HnswSparseEntity() {} + + //! Constructor + HnswSparseEntity(const HNSWSparseHeader &hd) { + header_ = hd; + } + + //! Destructor + virtual ~HnswSparseEntity() {} + + //! HnswSparseEntity Pointerd; + typedef std::shared_ptr Pointer; + + //! Get max neighbor size of graph level + inline size_t neighbor_cnt(level_t level) const { + return level == 0 ? header_.graph.l0_neighbor_count + : header_.hnsw.upper_neighbor_count; + } + + //! get max neighbor size of graph level 0 + inline size_t l0_neighbor_cnt() const { + return header_.graph.l0_neighbor_count; + } + + //! get min neighbor size of graph + inline size_t min_neighbor_cnt() const { + return header_.graph.min_neighbor_count; + } + + //! get upper neighbor size of graph level other than 0 + inline size_t upper_neighbor_cnt() const { + return header_.hnsw.upper_neighbor_count; + } + + //! Get current total doc of the hnsw graph + inline node_id_t *mutable_doc_cnt() { + return &header_.graph.doc_count; + } + + inline node_id_t doc_cnt() const { + return header_.graph.doc_count; + } + + inline uint32_t *mutable_total_sparse_count() { + return &header_.graph.total_sparse_count; + } + + uint32_t total_sparse_count() const { + return header_.graph.total_sparse_count; + } + + //! Get hnsw graph scaling params + inline size_t scaling_factor() const { + return header_.hnsw.scaling_factor; + } + + //! Get prune_size + inline size_t prune_cnt() const { + return header_.graph.prune_neighbor_count; + } + + //! Current entity of top level graph + inline node_id_t entry_point() const { + return header_.hnsw.entry_point; + } + + //! Current max graph level + inline level_t cur_max_level() const { + return header_.hnsw.max_level; + } + + //! Retrieve index vector size + size_t vector_size() const { + return header_.graph.vector_size; + } + + //! Retrieve node size + size_t node_size() const { + return header_.graph.node_size; + } + + //! Retrieve ef constuction + size_t ef_construction() const { + return header_.graph.ef_construction; + } + + //! Retrieve sparse meta size + size_t sparse_meta_size() const { + return header_.graph.sparse_meta_size; + } + + //! Retrieve sparse unit size + size_t sparse_unit_size() const { + return header_.graph.sparse_unit_size; + } + + void set_vector_size(size_t size) { + header_.graph.vector_size = size; + } + + void set_prune_cnt(size_t v) { + header_.graph.prune_neighbor_count = v; + } + + void set_scaling_factor(size_t val) { + header_.hnsw.scaling_factor = val; + } + + void set_l0_neighbor_cnt(size_t cnt) { + header_.graph.l0_neighbor_count = cnt; + } + + void set_min_neighbor_cnt(size_t cnt) { + header_.graph.min_neighbor_count = cnt; + } + + void set_upper_neighbor_cnt(size_t cnt) { + header_.hnsw.upper_neighbor_count = cnt; + } + + void set_ef_construction(size_t ef) { + header_.graph.ef_construction = ef; + } + + void set_sparse_meta_size(size_t size) { + header_.graph.sparse_meta_size = size; + } + + void set_sparse_unit_size(size_t size) { + header_.graph.sparse_unit_size = size; + } + + protected: + inline const HNSWSparseHeader &header() const { + return header_; + } + + inline HNSWSparseHeader *mutable_header() { + return &header_; + } + + inline size_t header_size() const { + return sizeof(header_); + } + + void set_node_size(size_t size) { + header_.graph.node_size = size; + } + + //! Dump all segment by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_segments( + const IndexDumper::Pointer &dumper, key_t *keys, + const std::function &get_level) const; + + private: + //! dump mapping segment, for get_vector_by_key in provider + int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const; + + //! dump hnsw head by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_header(const IndexDumper::Pointer &dumper, + const HNSWSparseHeader &hd) const; + + //! dump vectors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_sparse_vector_meta( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + //! dump sparse vectors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_sparse_vector( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + //! dump hnsw neighbors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_neighbors(const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + auto len1 = dump_graph_neighbors(dumper, reorder_mapping, neighbor_mapping); + if (len1 < 0) { + return len1; + } + auto len2 = dump_upper_neighbors(dumper, get_level, reorder_mapping, + neighbor_mapping); + if (len2 < 0) { + return len2; + } + + return len1 + len2; + } + + //! dump segment by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, const void *data, + size_t size) const; + + //! Dump level 0 neighbors + //! Return dump size if success, errno(<0) in failure + int64_t dump_graph_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + //! Dump upper level neighbors + //! Return dump size if success, errno(<0) in failure + int64_t dump_upper_neighbors( + const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + public: + //! Cleanup the entity + virtual int cleanup(void) { + header_.clear(); + return 0; + } + + //! Make a copy of searcher entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + virtual const HnswSparseEntity::Pointer clone() const { + LOG_ERROR("Update neighbors not implemented"); + return HnswSparseEntity::Pointer(); + } + + //! Get primary key of the node id + virtual key_t get_key(node_id_t id) const = 0; + + //! Get vector feature data by key + virtual const void *get_vector_meta(node_id_t id) const = 0; + + virtual int get_vector_meta(const node_id_t id, + IndexStorage::MemoryBlock &block) const = 0; + + //! Get vectors feature data by keys + virtual int get_vector_metas(const node_id_t *ids, uint32_t count, + const void **vecs) const = 0; + virtual int get_vector_metas( + const node_id_t *ids, uint32_t count, + std::vector &block_vecs) const = 0; + + //! Retrieve a sparse vector using a primary key + virtual int get_sparse_vector_by_key( + uint64_t /*key*/, uint32_t * /*sparse_count*/, + std::string * /*sparse_indices_buffer*/, + std::string * /*sparse_values_buffer*/) const { + LOG_ERROR("get sparse vector not implemented"); + return IndexError_NotImplemented; + } + + //! Retrieve a sparse vector using a primary key + virtual int get_sparse_vector_by_id( + node_id_t /*id*/, uint32_t * /*sparse_count*/, + std::string * /*sparse_indices_buffer*/, + std::string * /*sparse_values_buffer*/) const { + LOG_ERROR("get sparse vector not implemented"); + return IndexError_NotImplemented; + } + + //! Get vector sparse feature data by chunk index and offset + virtual const void *get_sparse_data(uint64_t offset, uint32_t len) const = 0; + + //! Get sparse data from id + virtual const void *get_sparse_data(node_id_t id) const = 0; + + virtual int get_sparse_data(uint64_t offset, uint32_t len, + IndexStorage::MemoryBlock &block) const = 0; + + virtual int get_sparse_data(const node_id_t id, + IndexStorage::MemoryBlock &block) const = 0; + + //! Get sparse data from vector + virtual std::pair get_sparse_data_from_vector( + const void *vec) const = 0; + virtual int get_sparse_data_from_vector(const void *vec, + IndexStorage::MemoryBlock &block, + int &sparse_length) const = 0; + + //! Get the node id's neighbors on graph level + //! Note: the neighbors cannot be modified, using the following + //! method to get WritableNeighbors if want to + virtual const Neighbors get_neighbors(level_t level, node_id_t id) const = 0; + + //! Add vector and key to hnsw entity, and local id will be saved in id + virtual int add_vector(level_t /*level*/, key_t /*key*/, + const std::string & /*vec*/, uint32_t /*sparse_count*/, + node_id_t * /*id*/) { + return IndexError_NotImplemented; + } + + virtual int add_vector(level_t /*level*/, key_t /*key*/, + const uint32_t /*sparse_count*/, + const uint32_t * /*sparse_indices*/, + const void * /*sparse_vec*/, node_id_t * /*id*/) { + return IndexError_NotImplemented; + } + + //! Add vector and id + virtual int add_vector_with_id(level_t /*level*/, node_id_t /*id*/, + const std::string & /*vec*/, + uint32_t /*sparse_count*/) { + return IndexError_NotImplemented; + } + + virtual int update_neighbors( + level_t /*level*/, node_id_t /*id*/, + const std::vector> & /*neighbors*/) { + LOG_ERROR("Update neighbors dense not implemented"); + + return 0; + } + + //! Append neighbor_id to node id neighbors on level, size is the current + //! neighbors size. Notice: the caller must be ensure the neighbors not full + virtual void add_neighbor(level_t /*level*/, node_id_t /*id*/, + uint32_t /*size*/, node_id_t /*neighbor_id*/) { + LOG_ERROR("Add neighbor not implemented"); + } + + //! Update entry point and max level + virtual void update_ep_and_level(node_id_t ep, level_t level) { + header_.hnsw.entry_point = ep; + header_.hnsw.max_level = level; + } + + virtual int load(const IndexStorage::Pointer & /*container*/, + bool /*check_crc*/) { + LOG_ERROR("Load not implemented"); + return IndexError_NotImplemented; + } + + virtual int dump(const IndexDumper::Pointer & /*dumper*/) { + LOG_ERROR("Dump not implemented"); + return IndexError_NotImplemented; + } + + static int CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size); + + protected: + static inline size_t AlignSize(size_t size) { + return (size + 0x1F) & (~0x1F); + } + + static inline size_t AlignPageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::PageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + //! rearrange vectors to improve cache locality + void reshuffle_vectors(const std::function &get_level, + std::vector *n2o_mapping, + std::vector *o2n_mapping, + key_t *keys) const; + + public: + const static std::string kSparseGraphHeaderSegmentId; + const static std::string kSparseGraphFeaturesSegmentId; + const static std::string kSparseGraphKeysSegmentId; + const static std::string kSparseGraphNeighborsSegmentId; + const static std::string kSparseGraphOffsetsSegmentId; + const static std::string kSparseGraphMappingSegmentId; + const static std::string kSparseHnswHeaderSegmentId; + const static std::string kSparseHnswNeighborsSegmentId; + const static std::string kSparseHnswOffsetsSegmentId; + const static std::string kSparseGraphVectorsSegmentId; + const static std::string kSparseGraphVectorMetaSegmentId; + + constexpr static uint32_t kRevision = 0U; + constexpr static size_t kMaxGraphLayers = 15; + constexpr static uint32_t kDefaultEfConstruction = 500; + constexpr static uint32_t kDefaultEf = 500; + constexpr static uint32_t kDefaultUpperMaxNeighborCnt = 50; // M of HNSW + constexpr static uint32_t kDefaultL0MaxNeighborCnt = 100; + constexpr static uint32_t kMaxNeighborCnt = 65535; + constexpr static float kDefaultScanRatio = 0.1f; + constexpr static uint32_t kDefaultMinScanLimit = 10000; + constexpr static uint32_t kDefaultMaxScanLimit = + std::numeric_limits::max(); + constexpr static float kDefaultBFNegativeProbability = 0.001f; + constexpr static uint32_t kDefaultScalingFactor = 50U; + constexpr static uint32_t kDefaultBruteForceThreshold = 1000U; + constexpr static uint32_t kDefaultDocsHardLimit = 1 << 30U; // 1 billion + constexpr static float kDefaultDocsSoftLimitRatio = 0.9f; + constexpr static size_t kMaxChunkSize = 0xFFFFFFFF; + constexpr static size_t kDefaultChunkSize = 2UL * 1024UL * 1024UL; + constexpr static size_t kDefaultMaxChunkCnt = 50000UL; + constexpr static float kDefaultNeighborPruneMultiplier = + 1.0f; // prune_cnt = upper_max_neighbor_cnt * multiplier + constexpr static float kDefaultL0MaxNeighborCntMultiplier = + 2.0f; // l0_max_neighbor_cnt = upper_max_neighbor_cnt * multiplier + + constexpr static uint32_t kSparseMetaSize = 2u * sizeof(uint64_t); + constexpr static float kDefaultSparseNeighborRatio = 0.5f; + constexpr static uint32_t kSparseMaxDimSize = 16384; + constexpr static float kDefaultQueryFilteringRatio = 0.0f; // turn off + + protected: + HNSWSparseHeader header_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_index_hash.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_index_hash.h new file mode 100644 index 0000000..fd8e53e --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_index_hash.h @@ -0,0 +1,232 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "hnsw_sparse_chunk.h" + +namespace zvec { +namespace core { + +//! Persistent hashmap implement through open addressing algorithm +template ::value>::type> +class HnswSparseIndexHashMap { + using key_type = Key; + using val_type = Val; + + struct Iterator { + key_type first; + val_type second; + }; + typedef Iterator *iterator; + typedef Iterator Item; + typedef const Iterator *const_iterator; + + class Slot { + public: + Slot(SparseChunk::Pointer &&chunk, const void *data) + : chunk_(std::move(chunk)), + items_(reinterpret_cast(data)) {} + //! Return a empty loc or the key item loc + + Slot(SparseChunk::Pointer &&chunk, IndexStorage::MemoryBlock &&mem_block) + : chunk_(std::move(chunk)), items_block_(std::move(mem_block)) { + items_ = reinterpret_cast(items_block_.data()); + } + const_iterator find(key_type key, uint32_t max_items, uint32_t mask) const { + auto it = &items_[key & mask]; + for (auto i = 0U; i < max_items; ++i) { + if (it->first == key || it->second == EmptyVal) { + // LOG_DEBUG("i=%u", i); + return it; + } + ++it; + if (it == &items_[max_items]) { + it = &items_[0]; + } + } + return nullptr; + } + + bool update(const_iterator it) { + uint32_t offset = reinterpret_cast(it) - + reinterpret_cast(&items_[0]); + if (ailego_unlikely(chunk_->write(offset, it, sizeof(Item)) != + sizeof(Item))) { + LOG_ERROR("Chunk write failed"); + return false; + } + return true; + } + + private: + SparseChunk::Pointer chunk_{}; + const Item *items_{nullptr}; // point to chunk data + IndexStorage::MemoryBlock items_block_{}; + }; + + public: + //! Init the hash + //! broker the index allocator + //! chunk_size the size of per chunk allocated, actual size may greater + //! factor factor = 1/ratio, ratio is the probability of a squence + //! number inserted to this container + //! max the max number key can be inserted + //! expansion_ratio memory expansion ratio + int init(SparseChunkBroker::Pointer &broker, uint32_t chunk_size, + uint32_t factor, size_t max, float expansion_ratio) { + ailego_assert_with(expansion_ratio > 1.0f, "ratio must > 1.0f"); + broker_ = broker; + + size_t items = std::ceil(chunk_size * 1.0f / sizeof(Item)); + slot_items_ = 1UL << static_cast((std::ceil(std::log2(items)))); + size_t range = slot_items_ * factor / expansion_ratio; + mask_bits_ = std::floor(std::log2(range)); + range = 1UL << mask_bits_; + size_t max_slots = std::ceil(max * 1.0f / range); + slots_.reserve(max_slots); + slot_loc_mask_ = slot_items_ - 1U; + + int ret = load(); + if (ret != 0) { + return ret; + } + + LOG_DEBUG( + "HnswIndexHash init, chunkSize=%u factor=%u max=%zu " + "ratio=%f slotItems=%u maxSlots=%zu maskBits=%u " + "range=%zu", + chunk_size, factor, max, expansion_ratio, slot_items_, max_slots, + mask_bits_, range); + + return 0; + } + + int cleanup(void) { + broker_.reset(); + slots_.clear(); + slots_.shrink_to_fit(); + mask_bits_ = 0U; + slot_items_ = 0U; + slot_loc_mask_ = 0U; + + return 0; + } + + const_iterator end(void) const { + return nullptr; + } + + const_iterator find(const key_type key) const { + auto idx = key >> mask_bits_; + if (idx >= slots_.size()) { + return end(); + } + auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_); + return it && it->second != EmptyVal ? it : nullptr; + } + + bool insert(key_type key, val_type val) { + auto idx = key >> mask_bits_; + if (idx >= slots_.size()) { + if (ailego_unlikely(idx >= slots_.capacity())) { + LOG_ERROR("no space to insert"); + return false; + } + for (auto i = slots_.size(); i <= idx; ++i) { + if (ailego_unlikely(!alloc_slot(i))) { + return false; + } + } + } + auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_); + if (ailego_unlikely(it == nullptr)) { + LOG_ERROR("no space to insert"); + return false; + } + + //! TODO: write memory is ok? + const_cast(it)->first = key; + const_cast(it)->second = val; + + return slots_[idx].update(it); + } + + private: + bool alloc_slot(size_t idx) { + ailego_assert_with(idx == slots_.size(), "invalid idx"); + + size_t size = slot_items_ * sizeof(Item); + auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, + idx, size); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return false; + } + SparseChunk::Pointer chunk = p.second; + if (ailego_unlikely(chunk->resize(size) != size)) { + LOG_ERROR("Chunk resize failed, size=%zu", size); + return false; + } + //! Read the whole data to memory + IndexStorage::MemoryBlock data_block; + if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) { + LOG_ERROR("Chunk read failed, size=%zu", size); + return false; + } + + slots_.emplace_back(std::move(chunk), std::move(data_block)); + return true; + } + + int load(void) { + size_t slots_cnt = + broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX); + for (size_t i = 0UL; i < slots_cnt; ++i) { + auto chunk = + broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, i); + if (!chunk) { + LOG_ERROR("Get chunk failed, seq=%zu", i); + return IndexError_InvalidFormat; + } + size_t size = sizeof(Item) * slot_items_; + if (chunk->data_size() < size) { + LOG_ERROR( + "Hash params may be mismatch, seq=%zu, data_size=%zu " + "expect=%zu", + i, chunk->data_size(), size); + return IndexError_InvalidFormat; + } + //! Read the whole data to memory + IndexStorage::MemoryBlock data_block; + if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) { + LOG_ERROR("Chunk read failed, size=%zu", size); + return false; + } + slots_.emplace_back(std::move(chunk), std::move(data_block)); + } + return 0; + } + + private: + SparseChunkBroker::Pointer broker_{}; // chunk broker + std::vector slots_{}; + uint32_t mask_bits_{0U}; + uint32_t slot_items_{}; // must be a power of 2 + uint32_t slot_loc_mask_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_index_provider.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_index_provider.h new file mode 100644 index 0000000..d1a694c --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_index_provider.h @@ -0,0 +1,168 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "hnsw_sparse_entity.h" + +namespace zvec { +namespace core { + +class HnswSparseIndexProvider : public IndexSparseProvider { + public: + HnswSparseIndexProvider(const IndexMeta &meta, + const HnswSparseEntity::Pointer &entity, + const std::string &owner) + : meta_(meta), entity_(entity), owner_class_(owner) {} + + HnswSparseIndexProvider(const HnswSparseIndexProvider &) = delete; + HnswSparseIndexProvider &operator=(const HnswSparseIndexProvider &) = delete; + + public: + //! Create a new iterator + IndexSparseProvider::Iterator::Pointer create_iterator(void) override { + return IndexSparseProvider::Iterator::Pointer(new (std::nothrow) + Iterator(entity_)); + } + + //! Retrieve count of vectors + size_t count(void) const override { + return entity_->doc_cnt(); + } + + size_t total_sparse_count(void) const override { + return entity_->total_sparse_count(); + } + + //! Retrieve type of vector + IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + //! Retrieve a vector using a primary key + int get_sparse_vector(uint64_t key, uint32_t *sparse_count, + std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override { + return entity_->get_sparse_vector_by_key( + key, sparse_count, sparse_indices_buffer, sparse_values_buffer); + } + + //! Retrieve the owner class + const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class Iterator : public IndexSparseProvider::Iterator { + public: + Iterator(const HnswSparseEntity::Pointer &entity) + : entity_(entity), cur_id_(0U), valid_(false) { + const void *sparse_data = entity_->get_sparse_data(cur_id_); + if (sparse_data != nullptr) { + valid_ = true; + + sparse_indices_buffer_.clear(); + sparse_data_buffer_.clear(); + + SparseUtility::ReverseSparseFormat( + sparse_data, &sparse_count_, &sparse_indices_buffer_, + &sparse_data_buffer_, entity_->sparse_unit_size()); + } + } + + //! Retrieve sparse count + uint32_t sparse_count() const override { + return sparse_count_; + } + + //! Retrieve sparse indices + const uint32_t *sparse_indices() const override { + return reinterpret_cast(sparse_indices_buffer_.data()); + } + + //! Retrieve sparse data + const void *sparse_data() const override { + return reinterpret_cast(sparse_data_buffer_.data()); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt() && valid_; + } + + //! Retrieve primary key + uint64_t key(void) const override { + return entity_->get_key(cur_id_); + } + + //! Next iterator + void next(void) override { + cur_id_ = get_next_valid_id(cur_id_ + 1); + + if (cur_id_ < entity_->doc_cnt()) { + const void *sparse_data = entity_->get_sparse_data(cur_id_); + if (sparse_data != nullptr) { + valid_ = true; + + sparse_indices_buffer_.clear(); + sparse_data_buffer_.clear(); + + SparseUtility::ReverseSparseFormat( + sparse_data, &sparse_count_, &sparse_indices_buffer_, + &sparse_data_buffer_, entity_->sparse_unit_size()); + } else { + valid_ = false; + } + } + } + + //! Reset the iterator + void reset(void) { + cur_id_ = get_next_valid_id(0); + const void *sparse_data = entity_->get_sparse_data(cur_id_); + if (sparse_data != nullptr) { + valid_ = true; + + SparseUtility::ReverseSparseFormat( + sparse_data, &sparse_count_, &sparse_indices_buffer_, + &sparse_data_buffer_, entity_->sparse_unit_size()); + } + } + + private: + node_id_t get_next_valid_id(node_id_t start_id) { + for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) { + if (entity_->get_key(i) != kInvalidNodeId) { + return i; + } + } + return kInvalidNodeId; + } + + private: + const HnswSparseEntity::Pointer entity_; + node_id_t cur_id_; + uint32_t sparse_count_{0}; + std::string sparse_indices_buffer_; + std::string sparse_data_buffer_; + bool valid_{false}; + }; + + private: + const IndexMeta &meta_; + const HnswSparseEntity::Pointer entity_; + const std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_params.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_params.h new file mode 100644 index 0000000..23db15a --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_params.h @@ -0,0 +1,118 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +static const std::string PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT( + "proxima.hnsw.sparse_builder.thread_count"); +static const std::string PARAM_HNSW_SPARSE_BUILDER_MEMORY_QUOTA( + "proxima.hnsw.sparse_builder.memory_quota"); +static const std::string PARAM_HNSW_SPARSE_BUILDER_EFCONSTRUCTION( + "proxima.hnsw.sparse_builder.efconstruction"); +static const std::string PARAM_HNSW_SPARSE_BUILDER_SCALING_FACTOR( + "proxima.hnsw.sparse_builder.scaling_factor"); +static const std::string PARAM_HNSW_SPARSE_BUILDER_CHECK_INTERVAL_SECS( + "proxima.hnsw.sparse_builder.check_interval_secs"); +static const std::string PARAM_HNSW_SPARSE_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER( + "proxima.hnsw.sparse_builder.neighbor_prune_multiplier"); +static const std::string PARAM_HNSW_SPARSE_BUILDER_MIN_NEIGHBOR_COUNT( + "proxima.hnsw.sparse_builder.min_neighbor_count"); +static const std::string PARAM_HNSW_SPARSE_BUILDER_MAX_NEIGHBOR_COUNT( + "proxima.hnsw.sparse_builder.max_neighbor_count"); +static const std::string + PARAM_HNSW_SPARSE_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER( + "proxima.hnsw.sparse_builder.l0_max_neighbor_count_multiplier"); + +static const std::string PARAM_HNSW_SPARSE_SEARCHER_EF( + "proxima.hnsw.sparse_searcher.ef"); +static const std::string PARAM_HNSW_SPARSE_SEARCHER_BRUTE_FORCE_THRESHOLD( + "proxima.hnsw.sparse_searcher.brute_force_threshold"); +static const std::string PARAM_HNSW_SPARSE_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE( + "proxima.hnsw.sparse_searcher.neighbors_in_memory_enable"); +static const std::string PARAM_HNSW_SPARSE_SEARCHER_MAX_SCAN_RATIO( + "proxima.hnsw.sparse_searcher.max_scan_ratio"); +static const std::string PARAM_HNSW_SPARSE_SEARCHER_CHECK_CRC_ENABLE( + "proxima.hnsw.sparse_searcher.check_crc_enable"); +static const std::string PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_ENABLE( + "proxima.hnsw.sparse_searcher.visit_bloomfilter_enable"); +static const std::string + PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB( + "proxima.hnsw.sparse_searcher.visit_bloomfilter_negative_prob"); +static const std::string PARAM_HNSW_SPARSE_SEARCHER_FORCE_PADDING_RESULT_ENABLE( + "proxima.hnsw.sparse_searcher.force_padding_result_enable"); +static const std::string PARAM_HNSW_SPARSE_SEARCHER_QUERY_FILTERING_RATIO( + "proxima.hnsw.sparse_searcher.query_filtering_ratio"); + +static const std::string PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO( + "proxima.hnsw.sparse_streamer.max_scan_ratio"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT( + "proxima.hnsw.sparse_streamer.min_scan_limit"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT( + "proxima.hnsw.sparse_streamer.max_scan_limit"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_EF( + "proxima.hnsw.sparse_streamer.ef"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION( + "proxima.hnsw.sparse_streamer.efconstruction"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT( + "proxima.hnsw.sparse_streamer.max_neighbor_count"); +static const std::string + PARAM_HNSW_SPARSE_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER( + "proxima.hnsw.sparse_streamer.l0_max_neighbor_count_multiplier"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR( + "proxima.hnsw.sparse_streamer.scaling_factor"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD( + "proxima.hnsw.sparse_streamer.brute_force_threshold"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT( + "proxima.hnsw.sparse_streamer.docs_hard_limit"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT( + "proxima.hnsw.sparse_streamer.docs_soft_limit"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_MAX_INDEX_SIZE( + "proxima.hnsw.sparse_streamer.max_index_size"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_ENABLE( + "proxima.hnsw.sparse_streamer.visit_bloomfilter_enable"); +static const std::string + PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB( + "proxima.hnsw.sparse_streamer.visit_bloomfilter_negative_prob"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_CHECK_CRC_ENABLE( + "proxima.hnsw.sparse_streamer.check_crc_enable"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER( + "proxima.hnsw.sparse_streamer.neighbor_prune_multiplier"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_CHUNK_SIZE( + "proxima.hnsw.sparse_streamer.chunk_size"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_FILTER_SAME_KEY( + "proxima.hnsw.sparse_streamer.filter_same_key"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_GET_VECTOR_ENABLE( + "proxima.hnsw.sparse_streamer.get_vector_enable"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_MIN_NEIGHBOR_COUNT( + "proxima.hnsw.sparse_streamer.min_neighbor_count"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_FORCE_PADDING_RESULT_ENABLE( + "proxima.hnsw.sparse_streamer.force_padding_result_enable"); +static const std::string PARAM_HNSW_SPARSE_STREAMER_QUERY_FILTERING_RATIO( + "proxima.hnsw.sparse_streamer.query_filtering_ratio"); + +static const std::string PARAM_HNSW_SPARSE_REDUCER_WORKING_PATH( + "proxima.hnsw.sparse_reducer.working_path"); +static const std::string PARAM_HNSW_SPARSE_REDUCER_NUM_OF_ADD_THREADS( + "proxima.hnsw.sparse_reducer.num_of_add_threads"); +static const std::string PARAM_HNSW_SPARSE_REDUCER_INDEX_NAME( + "proxima.hnsw.sparse_reducer.index_name"); +static const std::string PARAM_HNSW_SPARSE_REDUCER_EFCONSTRUCTION( + "proxima.hnsw.sparse_reducer.efconstruction"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.cc new file mode 100644 index 0000000..51cbacb --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.cc @@ -0,0 +1,564 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_searcher.h" +#include "hnsw_sparse_algorithm.h" +#include "hnsw_sparse_index_provider.h" +#include "hnsw_sparse_params.h" + +namespace zvec { +namespace core { + +HnswSparseSearcher::HnswSparseSearcher() {} + +HnswSparseSearcher::~HnswSparseSearcher() {} + +int HnswSparseSearcher::init(const ailego::Params &search_params) { + params_ = search_params; + params_.get(PARAM_HNSW_SPARSE_SEARCHER_EF, &ef_); + params_.get(PARAM_HNSW_SPARSE_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_); + params_.get(PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_ENABLE, + &bf_enabled_); + params_.get(PARAM_HNSW_SPARSE_SEARCHER_CHECK_CRC_ENABLE, &check_crc_enabled_); + params_.get(PARAM_HNSW_SPARSE_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE, + &neighbors_in_memory_enabled_); + params_.get(PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB, + &bf_negative_probability_); + params_.get(PARAM_HNSW_SPARSE_SEARCHER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + params_.get(PARAM_HNSW_SPARSE_SEARCHER_FORCE_PADDING_RESULT_ENABLE, + &force_padding_topk_enabled_); + + query_filtering_enabled_ = + params_.get(PARAM_HNSW_SPARSE_SEARCHER_QUERY_FILTERING_RATIO, + &query_filtering_ratio_); + + if (ef_ == 0) { + ef_ = HnswSparseEntity::kDefaultEf; + } + if (bf_negative_probability_ <= 0.0f || bf_negative_probability_ >= 1.0f) { + LOG_ERROR( + "[%s] must be in range (0,1)", + PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str()); + return IndexError_InvalidArgument; + } + + if (query_filtering_enabled_ && + (query_filtering_ratio_ <= 0.0f || query_filtering_ratio_ >= 1.0f)) { + LOG_ERROR("[%s] must be in range (0, 1)", + PARAM_HNSW_SPARSE_SEARCHER_QUERY_FILTERING_RATIO.c_str()); + return IndexError_InvalidArgument; + } + + entity_.set_neighbors_in_memory(neighbors_in_memory_enabled_); + + state_ = STATE_INITED; + + LOG_DEBUG( + "Init params: ef=%u maxScanRatio=%f bfEnabled=%u checkCrcEnabled=%u " + "neighborsInMemoryEnabled=%u bfNagtiveProb=%f bruteForceThreshold=%u " + "forcePadding=%u filteringRatio=%f", + ef_, max_scan_ratio_, bf_enabled_, check_crc_enabled_, + neighbors_in_memory_enabled_, bf_negative_probability_, + bruteforce_threshold_, force_padding_topk_enabled_, + query_filtering_ratio_); + + return 0; +} + +void HnswSparseSearcher::print_debug_info() { + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + Neighbors neighbours = entity_.get_neighbors(0, id); + std::cout << "node: " << id << "; "; + for (uint32_t i = 0; i < neighbours.size(); ++i) { + std::cout << neighbours[i]; + + if (i == neighbours.size() - 1) { + std::cout << std::endl; + } else { + std::cout << ", "; + } + } + } +} + +int HnswSparseSearcher::cleanup() { + LOG_INFO("Begin HnswSparseSearcher:cleanup"); + + metric_.reset(); + meta_.clear(); + stats_.clear_attributes(); + stats_.set_loaded_count(0UL); + stats_.set_loaded_costtime(0UL); + max_scan_ratio_ = HnswSparseEntity::kDefaultScanRatio; + max_scan_num_ = 0U; + ef_ = HnswSparseEntity::kDefaultEf; + bf_enabled_ = false; + bf_negative_probability_ = HnswSparseEntity::kDefaultBFNegativeProbability; + bruteforce_threshold_ = HnswSparseEntity::kDefaultBruteForceThreshold; + check_crc_enabled_ = false; + neighbors_in_memory_enabled_ = false; + entity_.cleanup(); + state_ = STATE_INIT; + + LOG_INFO("End HnswSparseSearcher:cleanup"); + + return 0; +} + +int HnswSparseSearcher::load(IndexStorage::Pointer container, + IndexMetric::Pointer metric) { + if (state_ != STATE_INITED) { + LOG_ERROR("Init the searcher first before load index"); + return IndexError_Runtime; + } + + LOG_INFO("Begin HnswSparseSearcher:load"); + + auto start_time = ailego::Monotime::MilliSeconds(); + + int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize meta from container"); + return ret; + } + + ret = entity_.load(container, check_crc_enabled_); + if (ret != 0) { + LOG_ERROR("HnswSparseSearcher load index failed"); + return ret; + } + + alg_ = HnswSparseAlgorithm::UPointer(new HnswSparseAlgorithm(entity_)); + + if (metric) { + metric_ = metric; + } else { + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("CreateMeasure failed, name: %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMetric init failed, ret=%d", ret); + return ret; + } + if (metric_->query_metric()) { + metric_ = metric_->query_metric(); + } + } + + // if (!metric_->is_matched(meta_)) { + // LOG_ERROR("IndexMeasure not match index meta"); + // return IndexError_Mismatch; + // } + + max_scan_num_ = static_cast(max_scan_ratio_ * entity_.doc_cnt()); + max_scan_num_ = std::max(4096U, max_scan_num_); + + stats_.set_loaded_count(entity_.doc_cnt()); + stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time); + state_ = STATE_LOADED; + magic_ = IndexContext::GenerateMagic(); + + LOG_INFO("End HnswSparseSearcher::load"); + + return 0; +} + +int HnswSparseSearcher::unload() { + LOG_INFO("HnswSparseSearcher unload index"); + + meta_.clear(); + entity_.cleanup(); + metric_.reset(); + max_scan_num_ = 0; + stats_.set_loaded_count(0UL); + stats_.set_loaded_costtime(0UL); + state_ = STATE_INITED; + + return 0; +} + +int HnswSparseSearcher::update_context(HnswSparseContext *ctx) const { + const HnswSparseEntity::Pointer entity = entity_.clone(); + if (!entity) { + LOG_ERROR("Failed to clone search context entity"); + return IndexError_Runtime; + } + ctx->set_max_scan_num(max_scan_num_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + + return ctx->update_context(HnswSparseContext::kSparseSearcherContext, meta_, + metric_, entity, magic_); +} + +//! Similarity search with sparse inputs +int HnswSparseSearcher::search_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (ailego_unlikely(!context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + HnswSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswSparseContext failed"); + return IndexError_Cast; + } + + if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) { + return search_bf_impl(sparse_count, sparse_indices, sparse_query, qmeta, + count, context); + } + + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + int ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->resize_results(count); + + const uint32_t *sparse_indices_tmp = sparse_indices; + const void *sparse_query_tmp = sparse_query; + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + std::string sparse_query_filtered_buffer; + + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + if (query_filtering_enabled_) { + if (!SparseUtility::FilterSparseQuery( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + qmeta.data_type(), entity_.sparse_unit_size(), + query_filtering_ratio_, &sparse_query_filtered_buffer)) { + LOG_ERROR("Hnsw filtering failed"); + return IndexError_Runtime; + } + + ctx->reset_query(sparse_query_filtered_buffer.data()); + } else { + ctx->reset_query(sparse_query_buffer.data()); + } + + int ret = alg_->search(ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw searcher fast search failed"); + return ret; + } + + if (query_filtering_enabled_) { + ctx->reset_query(sparse_query_buffer.data()); + ctx->recal_topk_dist(); + } + + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +//! Similarity search with sparse inputs +int HnswSparseSearcher::search_bf_impl( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, uint32_t count, + IndexStreamer::Context::Pointer &context) const { + if (ailego_unlikely(!context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + HnswSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswSparseContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + int ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->resize_results(count); + + const uint32_t *sparse_indices_tmp = sparse_indices; + const void *sparse_query_tmp = sparse_query; + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_Runtime; + } + + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity_.get_key(id)); + }; + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + + ctx->group_topk_heaps().clear(); + + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + if (entity_.get_key(id) == kInvalidKey) { + continue; + } + + if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { + dist_t dist = ctx->dist_calculator().dist(id); + + std::string group_id = group_by(id); + + auto &topk_heap = ctx->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, dist); + } + } + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + } else { + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + + ctx->topk_heap().clear(); + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + if (entity_.get_key(id) == kInvalidKey) { + continue; + } + + if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { + dist_t dist = ctx->dist_calculator().dist(id); + ctx->topk_heap().emplace(id, dist); + } + } + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +//! Similarity search with sparse inputs +int HnswSparseSearcher::search_bf_by_p_keys_impl( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (ailego_unlikely(!context)) { + LOG_ERROR("The context is not created by this searcher"); + return IndexError_Mismatch; + } + + if (ailego_unlikely(p_keys.size() != count)) { + LOG_ERROR("The size of p_keys is not equal to count"); + return IndexError_InvalidArgument; + } + + HnswSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswSparseContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + int ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->resize_results(count); + + const uint32_t *sparse_indices_tmp = sparse_indices; + const void *sparse_query_tmp = sparse_query; + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_Runtime; + } + + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity_.get_key(id)); + }; + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + ctx->group_topk_heaps().clear(); + + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t pk = p_keys[q][idx]; + if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { + node_id_t id = entity_.get_id(pk); + if (id != kInvalidNodeId) { + dist_t dist = ctx->dist_calculator().dist(id); + + std::string group_id = group_by(id); + + auto &topk_heap = ctx->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, dist); + } + } + } + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + } else { + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + ctx->topk_heap().clear(); + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t pk = p_keys[q][idx]; + if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { + node_id_t id = entity_.get_id(pk); + if (id != kInvalidNodeId) { + dist_t dist = ctx->dist_calculator().dist(id); + ctx->topk_heap().emplace(id, dist); + } + } + } + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +IndexSearcher::Context::Pointer HnswSparseSearcher::create_context() const { + if (ailego_unlikely(state_ != STATE_LOADED)) { + LOG_ERROR("Load the index first before create context"); + return Context::Pointer(); + } + const HnswSparseEntity::Pointer search_ctx_entity = entity_.clone(); + if (!search_ctx_entity) { + LOG_ERROR("Failed to create search context entity"); + return Context::Pointer(); + } + HnswSparseContext *ctx = + new (std::nothrow) HnswSparseContext(metric_, search_ctx_entity); + if (ailego_unlikely(ctx == nullptr)) { + LOG_ERROR("Failed to new HnswSparseContext"); + return Context::Pointer(); + } + ctx->set_ef(ef_); + ctx->set_max_scan_num(max_scan_num_); + uint32_t filter_mode = + bf_enabled_ ? VisitFilter::BloomFilter : VisitFilter::ByteMap; + ctx->set_filter_mode(filter_mode); + ctx->set_filter_negative_probability(bf_negative_probability_); + ctx->set_magic(magic_); + ctx->set_force_padding_topk(force_padding_topk_enabled_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + if (ailego_unlikely(ctx->init(HnswSparseContext::kSparseSearcherContext)) != + 0) { + LOG_ERROR("Init HnswSparseContext failed"); + delete ctx; + return Context::Pointer(); + } + + return Context::Pointer(ctx); +} + +IndexSearcher::SparseProvider::Pointer +HnswSparseSearcher::create_sparse_provider(void) const { + LOG_DEBUG("HnswSparseSearcher create sparse provider"); + + auto entity = entity_.clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone HnswSparseEntity failed"); + return SparseProvider::Pointer(); + } + return SparseProvider::Pointer(new (std::nothrow) HnswSparseIndexProvider( + meta_, entity, "HnswSparseSearcher")); +} + +int HnswSparseSearcher::get_sparse_vector( + uint64_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const { + return entity_.get_sparse_vector_by_key( + key, sparse_count, sparse_indices_buffer, sparse_values_buffer); +} + +INDEX_FACTORY_REGISTER_SEARCHER(HnswSparseSearcher); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.h new file mode 100644 index 0000000..377f340 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.h @@ -0,0 +1,159 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "hnsw_sparse_searcher_entity.h" +#include "hnsw_sparse_streamer.h" + +namespace zvec { +namespace core { + +class HnswSparseSearcher : public IndexSearcher { + public: + using ContextPointer = IndexSearcher::Context::Pointer; + + public: + HnswSparseSearcher(void); + ~HnswSparseSearcher(void) override; + + HnswSparseSearcher(const HnswSparseSearcher &) = delete; + HnswSparseSearcher &operator=(const HnswSparseSearcher &) = delete; + + protected: + //! Initialize Searcher + int init(const ailego::Params ¶ms) override; + + //! Cleanup Searcher + int cleanup(void) override; + + //! Load Index from storage + int load(IndexStorage::Pointer container, + IndexMetric::Pointer measure) override; + + //! Unload index from storage + int unload(void) override; + + //! Similarity search with sparse inputs + int search_impl(const uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override { + return search_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1, + context); + } + + //! Similarity search with sparse inputs + int search_impl(const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Similarity brute force search with sparse inputs + int search_bf_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const override { + return search_bf_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1, + context); + } + + //! Similarity brute force search with sparse inputs + int search_bf_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(&sparse_count, sparse_indices, sparse_query, + p_keys, qmeta, 1, context); + } + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + //! Fetch sparser vector by key + int get_sparse_vector(uint64_t key, uint32_t *sparse_count, + std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override; + + //! Create a searcher context + ContextPointer create_context() const override; + + //! Create a new iterator + IndexSearcher::SparseProvider::Pointer create_sparse_provider( + void) const override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve params of index + const ailego::Params ¶ms(void) const override { + return params_; + } + + void print_debug_info() override; + + private: + //! To share ctx across streamer/searcher, we need to update the context for + //! current streamer/searcher + int update_context(HnswSparseContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; + + HnswSparseSearcherEntity entity_{}; + HnswSparseAlgorithm::UPointer alg_; // impl graph algorithm + + IndexMetric::Pointer metric_{}; + IndexMeta meta_{}; + ailego::Params params_{}; + Stats stats_; + uint32_t ef_{HnswSparseEntity::kDefaultEf}; + uint32_t max_scan_num_{0U}; + uint32_t bruteforce_threshold_{HnswSparseEntity::kDefaultBruteForceThreshold}; + float max_scan_ratio_{HnswSparseEntity::kDefaultScanRatio}; + bool bf_enabled_{false}; + bool check_crc_enabled_{false}; + bool neighbors_in_memory_enabled_{false}; + bool force_padding_topk_enabled_{false}; + float bf_negative_probability_{ + HnswSparseEntity::kDefaultBFNegativeProbability}; + + bool query_filtering_enabled_{false}; + float query_filtering_ratio_{HnswSparseEntity::kDefaultQueryFilteringRatio}; + + uint32_t magic_{0U}; + + State state_{STATE_INIT}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.cc new file mode 100644 index 0000000..4b23ab9 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.cc @@ -0,0 +1,630 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_searcher_entity.h" +#include +#include "utility/sparse_utility.h" + +namespace zvec { +namespace core { + +HnswSparseSearcherEntity::HnswSparseSearcherEntity() {} + +int HnswSparseSearcherEntity::cleanup(void) { + container_.reset(); + sparse_vector_meta_.reset(); + keys_.reset(); + neighbors_.reset(); + neighbors_meta_.reset(); + sparse_vectors_.reset(); + neighbors_in_memory_enabled_ = false; + loaded_ = false; + + this->HnswSparseEntity::cleanup(); + + return 0; +} + +key_t HnswSparseSearcherEntity::get_key(node_id_t id) const { + const void *key; + if (ailego_unlikely(keys_->read(id * sizeof(key_t), &key, sizeof(key_t)) != + sizeof(key_t))) { + LOG_ERROR("Read key from segment failed"); + return kInvalidKey; + } + return *(reinterpret_cast(key)); +} + +//! Get vector local id by key +node_id_t HnswSparseSearcherEntity::get_id(key_t key) const { + if (ailego_unlikely(!mapping_)) { + LOG_ERROR("Index missing mapping segment"); + return kInvalidNodeId; + } + + //! Do binary search + node_id_t start = 0UL; + node_id_t end = doc_cnt(); + const void *data; + node_id_t idx = 0u; + while (start < end) { + idx = start + (end - start) / 2; + if (ailego_unlikely( + mapping_->read(idx * sizeof(node_id_t), &data, sizeof(node_id_t)) != + sizeof(node_id_t))) { + LOG_ERROR("Read key from segment failed"); + return kInvalidNodeId; + } + const key_t *mkey; + node_id_t local_id = *reinterpret_cast(data); + if (ailego_unlikely(keys_->read(local_id * sizeof(key_t), + (const void **)(&mkey), + sizeof(key_t)) != sizeof(key_t))) { + LOG_ERROR("Read key from segment failed"); + return kInvalidNodeId; + } + if (*mkey < key) { + start = idx + 1; + } else if (*mkey > key) { + end = idx; + } else { + return local_id; + } + } + return kInvalidNodeId; +} + +int HnswSparseSearcherEntity::get_sparse_vector_by_key( + key_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const { + *sparse_count = 0; + + auto id = get_id(key); + if (id == kInvalidNodeId) { + return IndexError_NoExist; + } + + const void *sparse_data = get_sparse_data(id); + if (sparse_data == nullptr) { + return IndexError_InvalidValue; + } + + SparseUtility::ReverseSparseFormat(sparse_data, sparse_count, + sparse_indices_buffer, + sparse_values_buffer, sparse_unit_size()); + + return 0; +} + +const void *HnswSparseSearcherEntity::get_vector_meta(node_id_t id) const { + size_t read_size = sparse_meta_size(); + size_t offset = sparse_meta_size() * id; + + const void *vec; + if (ailego_unlikely(sparse_vector_meta_->read(offset, &vec, read_size) != + read_size)) { + LOG_ERROR("Read vector from segment failed"); + return nullptr; + } + return vec; +} + +int HnswSparseSearcherEntity::get_vector_meta( + const node_id_t id, IndexStorage::MemoryBlock &block) const { + const void *vec = get_vector_meta(id); + block.reset((void *)vec); + return 0; +} + +int HnswSparseSearcherEntity::get_vector_metas(const node_id_t *ids, + uint32_t count, + const void **vecs) const { + ailego_assert_with(count <= segment_datas_.size(), "invalid count"); + + size_t read_size = sparse_meta_size(); + + for (uint32_t i = 0; i < count; ++i) { + segment_datas_[i].offset = sparse_meta_size() * ids[i]; + segment_datas_[i].length = read_size; + + ailego_assert_with( + segment_datas_[i].offset < sparse_vector_meta_->data_size(), + "invalid offset"); + } + if (ailego_unlikely(!sparse_vector_meta_->read(&segment_datas_[0], count))) { + LOG_ERROR("Read vectors from segment failed"); + return IndexError_ReadData; + } + for (uint32_t i = 0; i < count; ++i) { + vecs[i] = segment_datas_[i].data; + } + + return 0; +} + +int HnswSparseSearcherEntity::get_vector_metas( + const node_id_t *ids, uint32_t count, + std::vector &block_vecs) const { + std::vector vecs(count); + get_vector_metas(ids, count, vecs.data()); + for (uint32_t i = 0; i < count; ++i) { + block_vecs.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i])); + } + return 0; +} + +const Neighbors HnswSparseSearcherEntity::get_neighbors(level_t level, + node_id_t id) const { + if (level == 0) { + if (neighbors_in_memory_enabled_) { + auto hd = reinterpret_cast( + fixed_neighbors_.get() + neighbors_size() * id); + return {hd->neighbor_cnt, hd->neighbors}; + } + + const SparseGraphNeighborMeta *m; + if (ailego_unlikely( + neighbors_meta_->read(id * sizeof(SparseGraphNeighborMeta), + (const void **)(&m), + sizeof(SparseGraphNeighborMeta)) != + sizeof(SparseGraphNeighborMeta))) { + LOG_ERROR("Read neighbors meta from segment failed"); + return {0, nullptr}; + } + + const void *data; + if (ailego_unlikely(neighbors_->read(m->offset, &data, + m->neighbor_cnt * sizeof(node_id_t)) != + m->neighbor_cnt * sizeof(node_id_t))) { + LOG_ERROR("Read neighbors from segment failed"); + return {0, nullptr}; + } + return {static_cast(m->neighbor_cnt), + reinterpret_cast(data)}; + } + + //! Read level > 0 neighbors + const HnswSparseNeighborMeta *m; + if (ailego_unlikely( + upper_neighbors_meta_->read(id * sizeof(HnswSparseNeighborMeta), + (const void **)(&m), + sizeof(HnswSparseNeighborMeta)) != + sizeof(HnswSparseNeighborMeta))) { + LOG_ERROR("Read neighbors meta from segment failed"); + return {0, nullptr}; + } + + ailego_assert_with(level <= m->level, "invalid level"); + size_t offset = m->offset + (level - 1) * upper_neighbors_size(); + ailego_assert_with(offset <= upper_neighbors_->data_size(), "invalid offset"); + const void *data; + if (ailego_unlikely( + upper_neighbors_->read(offset, &data, upper_neighbors_size()) != + upper_neighbors_size())) { + LOG_ERROR("Read neighbors from segment failed"); + return {0, nullptr}; + } + + auto hd = reinterpret_cast(data); + return {hd->neighbor_cnt, hd->neighbors}; +} + +int HnswSparseSearcherEntity::load(const IndexStorage::Pointer &container, + bool check_crc) { + container_ = container; + + int ret = load_segments(check_crc); + if (ret != 0) { + return ret; + } + + loaded_ = true; + + LOG_INFO( + "Index info: docCnt=%u entryPoint=%u maxLevel=%d efConstruct=%zu " + "l0NeighborCnt=%zu upperNeighborCnt=%zu scalingFactor=%zu " + "nodeSize=%zu sparesMetaSegmentSize=%zu keySegmentSize=%zu " + "neighborsSegmentSize=%zu neighborsMetaSegmentSize=%zu " + "sparseVectorSegmentSize=%zu", + doc_cnt(), entry_point(), cur_max_level(), ef_construction(), + l0_neighbor_cnt(), upper_neighbor_cnt(), scaling_factor(), node_size(), + sparse_vector_meta_->data_size(), keys_->data_size(), + neighbors_->data_size(), neighbors_meta_->data_size(), + sparse_vectors_->data_size()); + + return 0; +} + +int HnswSparseSearcherEntity::load_segments(bool check_crc) { + //! load header + const void *data = nullptr; + HNSWSparseHeader hd; + auto graph_hd_segment = container_->get(kSparseGraphHeaderSegmentId); + if (!graph_hd_segment || graph_hd_segment->data_size() < sizeof(hd.graph)) { + LOG_ERROR("Miss or invalid segment %s", + kSparseGraphHeaderSegmentId.c_str()); + return IndexError_InvalidFormat; + } + if (graph_hd_segment->read(0, reinterpret_cast(&data), + sizeof(hd.graph)) != sizeof(hd.graph)) { + LOG_ERROR("Read segment %s failed", kSparseGraphHeaderSegmentId.c_str()); + return IndexError_ReadData; + } + memcpy(&hd.graph, data, sizeof(hd.graph)); + + auto hnsw_hd_segment = container_->get(kSparseHnswHeaderSegmentId); + if (!hnsw_hd_segment || hnsw_hd_segment->data_size() < sizeof(hd.hnsw)) { + LOG_ERROR("Miss or invalid segment %s", kSparseHnswHeaderSegmentId.c_str()); + return IndexError_InvalidFormat; + } + if (hnsw_hd_segment->read(0, reinterpret_cast(&data), + sizeof(hd.hnsw)) != sizeof(hd.hnsw)) { + LOG_ERROR("Read segment %s failed", kSparseHnswHeaderSegmentId.c_str()); + return IndexError_ReadData; + } + memcpy(&hd.hnsw, data, sizeof(hd.hnsw)); + *mutable_header() = hd; + segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt())); + + sparse_vector_meta_ = container_->get(kSparseGraphVectorMetaSegmentId); + if (!sparse_vector_meta_) { + LOG_ERROR("IndexStorage get segment %s failed", + kSparseGraphVectorMetaSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + keys_ = container_->get(kSparseGraphKeysSegmentId); + if (!keys_) { + LOG_ERROR("IndexStorage get segment %s failed", + kSparseGraphKeysSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + sparse_vectors_ = container_->get(kSparseGraphVectorsSegmentId); + if (!sparse_vectors_) { + LOG_ERROR("IndexStorage get segment %s failed", + kSparseGraphVectorsSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + neighbors_ = container_->get(kSparseGraphNeighborsSegmentId); + if (!neighbors_ || (neighbors_->data_size() == 0 && doc_cnt() > 1)) { + LOG_ERROR("IndexStorage get segment %s failed or empty", + kSparseGraphNeighborsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + neighbors_meta_ = container_->get(kSparseGraphOffsetsSegmentId); + if (!neighbors_meta_ || neighbors_meta_->data_size() < + sizeof(SparseGraphNeighborMeta) * doc_cnt()) { + LOG_ERROR("IndexStorage get segment %s failed or invalid size", + kSparseGraphOffsetsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + upper_neighbors_ = container_->get(kSparseHnswNeighborsSegmentId); + if (!upper_neighbors_ || + (upper_neighbors_->data_size() == 0 && cur_max_level() > 0)) { + LOG_ERROR("IndexStorage get segment %s failed or empty", + kSparseHnswNeighborsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + upper_neighbors_meta_ = container_->get(kSparseHnswOffsetsSegmentId); + if (!upper_neighbors_meta_ || + upper_neighbors_meta_->data_size() < + sizeof(HnswSparseNeighborMeta) * doc_cnt()) { + LOG_ERROR("IndexStorage get segment %s failed or invalid size", + kSparseHnswOffsetsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + mapping_ = container_->get(kSparseGraphMappingSegmentId); + if (!mapping_ || mapping_->data_size() < sizeof(node_id_t) * doc_cnt()) { + LOG_ERROR("IndexStorage get segment %s failed or invalid size", + kSparseGraphMappingSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + if (check_crc) { + std::vector segments; + segments.emplace_back(graph_hd_segment); + segments.emplace_back(hnsw_hd_segment); + segments.emplace_back(sparse_vector_meta_); + segments.emplace_back(keys_); + + segments.emplace_back(sparse_vectors_); + + segments.emplace_back(neighbors_); + segments.emplace_back(neighbors_meta_); + segments.emplace_back(upper_neighbors_); + segments.emplace_back(upper_neighbors_meta_); + + if (!do_crc_check(segments)) { + LOG_ERROR("Check index crc failed, the index may broken"); + return IndexError_Runtime; + } + } + + if (neighbors_in_memory_enabled_) { + int ret = load_and_flat_neighbors(); + if (ret != 0) { + return ret; + } + } + + return 0; +} + +int HnswSparseSearcherEntity::load_and_flat_neighbors() { + fixed_neighbors_.reset( + new (std::nothrow) char[neighbors_size() * doc_cnt()]{}, + std::default_delete()); + if (!fixed_neighbors_) { + LOG_ERROR("Malloc memory failed"); + return IndexError_NoMemory; + } + + //! Get a new segemnt to release the buffer after loading neighbors + auto neighbors_meta = container_->get(kSparseGraphOffsetsSegmentId); + if (!neighbors_meta) { + LOG_ERROR("IndexStorage get segment graph.offsets failed"); + return IndexError_InvalidArgument; + } + + const SparseGraphNeighborMeta *neighbors_index = nullptr; + if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index), + neighbors_meta->data_size()) != + neighbors_meta->data_size()) { + LOG_ERROR("Read segment %s data failed", + kSparseGraphOffsetsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + const char *neighbor_data; + for (node_id_t id = 0; id < doc_cnt(); ++id) { + size_t rd_size = neighbors_index[id].neighbor_cnt * sizeof(node_id_t); + if (ailego_unlikely( + neighbors_->read(neighbors_index[id].offset, + reinterpret_cast(&neighbor_data), + rd_size) != rd_size)) { + LOG_ERROR("Read neighbors from segment failed"); + return IndexError_ReadData; + } + // copy level 0 neighbors to fixed size neighbors memory + char *dst = fixed_neighbors_.get() + neighbors_size() * id; + *reinterpret_cast(dst) = neighbors_index[id].neighbor_cnt; + memcpy(dst + sizeof(uint32_t), neighbor_data, rd_size); + } + + return 0; +} + +int HnswSparseSearcherEntity::get_fixed_neighbors( + std::vector *fixed_neighbors) const { + //! Get a new segemnt to release the buffer after loading neighbors + auto neighbors_meta = container_->get(kSparseGraphOffsetsSegmentId); + if (!neighbors_meta) { + LOG_ERROR("IndexStorage get segment graph.offsets failed"); + return IndexError_InvalidArgument; + } + + const SparseGraphNeighborMeta *neighbors_index = nullptr; + size_t meta_size = neighbors_meta->data_size(); + if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index), + meta_size) != meta_size) { + LOG_ERROR("Read segment %s data failed", + kSparseGraphOffsetsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + size_t fixed_neighbor_cnt = l0_neighbor_cnt(); + fixed_neighbors->resize((fixed_neighbor_cnt + 1) * doc_cnt(), kInvalidNodeId); + + size_t neighbors_cnt_offset = fixed_neighbor_cnt * doc_cnt(); + size_t total_neighbor_cnt = 0; + for (node_id_t id = 0; id < doc_cnt(); ++id) { + size_t cur_neighbor_cnt = neighbors_index[id].neighbor_cnt; + if (cur_neighbor_cnt == 0) { + (*fixed_neighbors)[neighbors_cnt_offset + id] = 0; + continue; + } + size_t rd_size = cur_neighbor_cnt * sizeof(node_id_t); + const uint32_t *neighbors; + if (neighbors_->read(neighbors_index[id].offset, + reinterpret_cast(&neighbors), + rd_size) != rd_size) { + LOG_ERROR("Read neighbors from segment failed"); + return IndexError_ReadData; + } + + // copy level 0 neighbors to fixed size neighbors memory + auto it = fixed_neighbors->begin() + id * fixed_neighbor_cnt; + std::copy(neighbors, neighbors + cur_neighbor_cnt, it); + + (*fixed_neighbors)[neighbors_cnt_offset + id] = cur_neighbor_cnt; + total_neighbor_cnt += cur_neighbor_cnt; + } + LOG_INFO("total neighbor cnt: %zu, average neighbor cnt: %zu", + total_neighbor_cnt, total_neighbor_cnt / doc_cnt()); + + return 0; +} + +bool HnswSparseSearcherEntity::do_crc_check( + std::vector &segments) const { + constexpr size_t blk_size = 4096; + const void *data; + for (auto &segment : segments) { + size_t offset = 0; + size_t rd_size; + uint32_t crc = 0; + while (offset < segment->data_size()) { + size_t size = std::min(blk_size, segment->data_size() - offset); + if ((rd_size = segment->read(offset, &data, size)) <= 0) { + break; + } + offset += rd_size; + crc = ailego::Crc32c::Hash(data, rd_size, crc); + } + if (crc != segment->data_crc()) { + return false; + } + } + return true; +} + +const HnswSparseEntity::Pointer HnswSparseSearcherEntity::clone() const { + auto keys = keys_->clone(); + if (ailego_unlikely(!keys)) { + LOG_ERROR("clone segment %s failed", kSparseGraphKeysSegmentId.c_str()); + return HnswSparseEntity::Pointer(); + } + + auto mapping = mapping_->clone(); + if (ailego_unlikely(!mapping)) { + LOG_ERROR("clone segment %s failed", kSparseGraphMappingSegmentId.c_str()); + return HnswSparseEntity::Pointer(); + } + + auto sparse_vector_meta = sparse_vector_meta_->clone(); + if (ailego_unlikely(!sparse_vector_meta)) { + LOG_ERROR("clone segment %s failed", + kSparseGraphVectorMetaSegmentId.c_str()); + return HnswSparseEntity::Pointer(); + } + + auto sparse_vectors = sparse_vectors_->clone(); + if (ailego_unlikely(!sparse_vectors)) { + LOG_ERROR("clone segment %s failed", kSparseGraphVectorsSegmentId.c_str()); + return HnswSparseEntity::Pointer(); + } + + auto neighbors = neighbors_->clone(); + if (ailego_unlikely(!neighbors)) { + LOG_ERROR("clone segment %s failed", + kSparseGraphNeighborsSegmentId.c_str()); + return HnswSparseEntity::Pointer(); + } + auto upper_neighbors = upper_neighbors_->clone(); + if (ailego_unlikely(!neighbors)) { + LOG_ERROR("clone segment %s failed", kSparseHnswNeighborsSegmentId.c_str()); + return HnswSparseEntity::Pointer(); + } + auto neighbors_meta = neighbors_meta_->clone(); + if (ailego_unlikely(!neighbors_meta)) { + LOG_ERROR("clone segment %s failed", kSparseGraphOffsetsSegmentId.c_str()); + return HnswSparseEntity::Pointer(); + } + auto upper_neighbors_meta = upper_neighbors_meta_->clone(); + if (ailego_unlikely(!upper_neighbors_meta)) { + LOG_ERROR("clone segment %s failed", kSparseHnswOffsetsSegmentId.c_str()); + return HnswSparseEntity::Pointer(); + } + + SegmentGroupParam neighbor_group{neighbors, neighbors_meta, upper_neighbors, + upper_neighbors_meta}; + SegmentGroupParam dense_neighbor_group{nullptr, nullptr, nullptr, nullptr}; + SegmentGroupParam sparse_neighbor_group{nullptr, nullptr, nullptr, nullptr}; + + HnswSparseSearcherEntity *entity = new (std::nothrow) + HnswSparseSearcherEntity(header(), keys, mapping, neighbor_group, + sparse_vector_meta, sparse_vectors, + fixed_neighbors_, neighbors_in_memory_enabled_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswSparseSearcherEntity new failed"); + } + + return HnswSparseEntity::Pointer(entity); +} + +//! Get vector sparse feature data by chunk index and offset +const void *HnswSparseSearcherEntity::get_sparse_data(uint64_t offset, + uint32_t len) const { + const void *sparse_data = nullptr; + + uint32_t real_length = sparse_vectors_->read(offset, &sparse_data, len); + + if (ailego_unlikely(real_length != len)) { + LOG_ERROR("Read sparse data from segment failed, %u vs %u", real_length, + len); + return nullptr; + } + + return sparse_data; +} + +int HnswSparseSearcherEntity::get_sparse_data( + uint64_t offset, uint32_t len, IndexStorage::MemoryBlock &block) const { + const void *vec = get_sparse_data(offset, len); + block.reset((void *)vec); + return 0; +} + +//! Get sparse data from id +const void *HnswSparseSearcherEntity::get_sparse_data(node_id_t id) const { + const void *vec = get_vector_meta(id); + if (vec == nullptr) { + LOG_ERROR("get vector failed, id: %u", id); + + return nullptr; + } + + auto sparse_data = get_sparse_data_from_vector(vec); + + return sparse_data.first; +} + +int HnswSparseSearcherEntity::get_sparse_data( + const node_id_t id, IndexStorage::MemoryBlock &block) const { + const void *vec = get_sparse_data(id); + block.reset((void *)vec); + return 0; +} + +//! Get sparse data from vector +std::pair +HnswSparseSearcherEntity::get_sparse_data_from_vector(const void *vec) const { + if (vec == nullptr) { + LOG_ERROR("vec is nullptr"); + + return std::make_pair(nullptr, 0); + } + + const char *vec_ptr = reinterpret_cast(vec); + + uint64_t offset = *((uint64_t *)(vec_ptr)); + uint32_t sparse_vector_len = *((uint32_t *)(vec_ptr + sizeof(uint64_t))); + + const void *sparse_data = get_sparse_data(offset, sparse_vector_len); + if (ailego_unlikely(sparse_data == nullptr)) { + LOG_ERROR("Get nullptr sparse, offset=%zu, len=%u", (size_t)offset, + sparse_vector_len); + + return std::make_pair(nullptr, 0); + } + + return std::make_pair(sparse_data, sparse_vector_len); +} + +int HnswSparseSearcherEntity::get_sparse_data_from_vector( + const void *vec, IndexStorage::MemoryBlock &block, + int &sparse_length) const { + std::pair sparse_data = + get_sparse_data_from_vector(vec); + block.reset((void *)sparse_data.first); + sparse_length = sparse_data.second; + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.h new file mode 100644 index 0000000..4860f16 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.h @@ -0,0 +1,181 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include "hnsw_sparse_builder_entity.h" +#include "hnsw_sparse_entity.h" + +namespace zvec { +namespace core { + +class HnswSparseSearcherEntity : public HnswSparseEntity { + public: + using Pointer = std::shared_ptr; + using SegmentPointer = IndexStorage::Segment::Pointer; + + public: + struct SegmentGroupParam { + SegmentGroupParam(SegmentPointer neighbors_in, + SegmentPointer neighbors_meta_in, + SegmentPointer upper_neighbors_in, + SegmentPointer upper_neighbors_meta_in) + : neighbors{neighbors_in}, + neighbors_meta{neighbors_meta_in}, + upper_neighbors{upper_neighbors_in}, + upper_neighbors_meta{upper_neighbors_meta_in} {} + + SegmentPointer neighbors{nullptr}; + SegmentPointer neighbors_meta{nullptr}; + SegmentPointer upper_neighbors{nullptr}; + SegmentPointer upper_neighbors_meta{nullptr}; + }; + + //! Constructor + HnswSparseSearcherEntity(); + + //! Make a copy of searcher entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + const HnswSparseEntity::Pointer clone() const override; + + //! Get primary key of the node id + key_t get_key(node_id_t id) const override; + + //! Get vector local id by key + node_id_t get_id(key_t key) const; + + //! Get sparse vector feature data by key + int get_sparse_vector_by_key( + key_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override; + + //! Get vector feature data by id + const void *get_vector_meta(node_id_t id) const override; + + int get_vector_meta(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get vector feature data by id + int get_vector_metas(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + + int get_vector_metas( + const node_id_t *ids, uint32_t count, + std::vector &block_vecs) const override; + + //! Get vector sparse feature data by chunk index and offset + const void *get_sparse_data(uint64_t offset, uint32_t len) const override; + + //! Get sparse data from id + const void *get_sparse_data(node_id_t id) const override; + + int get_sparse_data(uint64_t offset, uint32_t len, + IndexStorage::MemoryBlock &block) const override; + + int get_sparse_data(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get sparse data from vector + std::pair get_sparse_data_from_vector( + const void *vec) const override; + + int get_sparse_data_from_vector(const void *vec, + IndexStorage::MemoryBlock &block, + int &sparse_length) const override; + + //! Get the node id's neighbors on graph level + const Neighbors get_neighbors(level_t level, node_id_t id) const override; + + int load(const IndexStorage::Pointer &container, bool check_crc) override; + + int load_segments(bool check_crc); + + int cleanup(void) override; + + public: + bool is_loaded() const { + return loaded_; + } + + void set_neighbors_in_memory(bool enabled) { + neighbors_in_memory_enabled_ = enabled; + } + + //! get fixed length neighbors data + int get_fixed_neighbors(std::vector *fixed_neighbors) const; + + private: + //! Constructor + HnswSparseSearcherEntity(const HNSWSparseHeader &hd, + const SegmentPointer &keys, + const SegmentPointer &mapping, + const SegmentGroupParam &neighbor_group, + const SegmentPointer &sparse_vector_meta, + const SegmentPointer &sparse_vectors, + const std::shared_ptr &fixed_neighbors, + bool neighbors_in_memory_enabled) + : HnswSparseEntity(hd), + keys_(keys), + mapping_(mapping), + neighbors_(neighbor_group.neighbors), + neighbors_meta_(neighbor_group.neighbors_meta), + upper_neighbors_(neighbor_group.upper_neighbors), + upper_neighbors_meta_(neighbor_group.upper_neighbors_meta), + sparse_vector_meta_(sparse_vector_meta), + sparse_vectors_(sparse_vectors), + neighbors_in_memory_enabled_(neighbors_in_memory_enabled) { + segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()), + IndexStorage::SegmentData(0U, 0U)); + fixed_neighbors_ = fixed_neighbors; + } + + bool do_crc_check(std::vector &segments) const; + + inline size_t neighbors_size() const { + return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); + } + + inline size_t upper_neighbors_size() const { + return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); + } + + //! If neighbors_in_memory_enabled, load the level0 neighbors to memory + int load_and_flat_neighbors(void); + + public: + HnswSparseSearcherEntity(const HnswSparseSearcherEntity &) = delete; + HnswSparseSearcherEntity &operator=(const HnswSparseSearcherEntity &) = + delete; + + private: + IndexStorage::Pointer container_{}; + + SegmentPointer keys_{}; + SegmentPointer mapping_{}; + + SegmentPointer neighbors_{}; + SegmentPointer neighbors_meta_{}; + SegmentPointer upper_neighbors_{}; + SegmentPointer upper_neighbors_meta_{}; + + SegmentPointer sparse_vector_meta_{}; + SegmentPointer sparse_vectors_{}; + + mutable std::vector segment_datas_{}; + std::shared_ptr fixed_neighbors_{}; // level 0 fixed size neighbors + bool neighbors_in_memory_enabled_{false}; + bool loaded_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.cc new file mode 100644 index 0000000..20c2152 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.cc @@ -0,0 +1,960 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_streamer.h" +#include +#include +#include +#include +#include "hnsw_sparse_algorithm.h" +#include "hnsw_sparse_context.h" +#include "hnsw_sparse_dist_calculator.h" +#include "hnsw_sparse_index_provider.h" + +namespace zvec { +namespace core { + +HnswSparseStreamer::HnswSparseStreamer() : entity_(stats_) {} + +HnswSparseStreamer::~HnswSparseStreamer() { + if (state_ == STATE_INITED || state_ == STATE_OPENED) { + this->cleanup(); + } +} + +int HnswSparseStreamer::init(const IndexMeta &imeta, + const ailego::Params ¶ms) { + meta_ = imeta; + meta_.set_streamer("HnswSparseStreamer", HnswSparseEntity::kRevision, params); + + params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_INDEX_SIZE, &max_index_size_); + params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT, + &upper_max_neighbor_cnt_); + float multiplier = HnswSparseEntity::kDefaultL0MaxNeighborCntMultiplier; + params.get(PARAM_HNSW_SPARSE_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, + &multiplier); + l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_; + + multiplier = HnswSparseEntity::kDefaultNeighborPruneMultiplier; + params.get(PARAM_HNSW_SPARSE_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier); + size_t prune_cnt = multiplier * upper_max_neighbor_cnt_; + scaling_factor_ = upper_max_neighbor_cnt_; + params.get(PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR, &scaling_factor_); + + params.get(PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_); + params.get(PARAM_HNSW_SPARSE_STREAMER_EF, &ef_); + params.get(PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION, &ef_construction_); + params.get(PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_); + params.get(PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB, + &bf_negative_prob_); + params.get(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_); + params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_); + params.get(PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_); + params.get(PARAM_HNSW_SPARSE_STREAMER_CHECK_CRC_ENABLE, &check_crc_enabled_); + + params.get(PARAM_HNSW_SPARSE_STREAMER_CHUNK_SIZE, &chunk_size_); + params.get(PARAM_HNSW_SPARSE_STREAMER_FILTER_SAME_KEY, &filter_same_key_); + params.get(PARAM_HNSW_SPARSE_STREAMER_GET_VECTOR_ENABLE, + &get_vector_enabled_); + params.get(PARAM_HNSW_SPARSE_STREAMER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_); + params.get(PARAM_HNSW_SPARSE_STREAMER_FORCE_PADDING_RESULT_ENABLE, + &force_padding_topk_enabled_); + + query_filtering_enabled_ = + params.get(PARAM_HNSW_SPARSE_STREAMER_QUERY_FILTERING_RATIO, + &query_filtering_ratio_); + + params.get(PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit_); + if (docs_soft_limit_ > 0 && docs_soft_limit_ > docs_hard_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT.c_str(), + PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT.c_str()); + return IndexError_InvalidArgument; + } else if (docs_soft_limit_ == 0UL) { + docs_soft_limit_ = + docs_hard_limit_ * HnswSparseEntity::kDefaultDocsSoftLimitRatio; + } + + if (ef_ == 0U) { + ef_ = HnswSparseEntity::kDefaultEf; + } + if (ef_construction_ == 0U) { + ef_construction_ = HnswSparseEntity::kDefaultEfConstruction; + } + if (upper_max_neighbor_cnt_ == 0U) { + upper_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt; + } + if (upper_max_neighbor_cnt_ > HnswSparseEntity::kMaxNeighborCnt) { + LOG_ERROR("[%s] must be in range (0,%d)", + PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT.c_str(), + HnswSparseEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (l0_max_neighbor_cnt_ == 0U) { + l0_max_neighbor_cnt_ = HnswSparseEntity::kDefaultL0MaxNeighborCnt; + } + if (l0_max_neighbor_cnt_ > HnswSparseEntity::kMaxNeighborCnt) { + LOG_ERROR("UpperNeighborCnt must be in range (0,%d)", + HnswSparseEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) { + LOG_ERROR("[%s]-[%u] must be <= [%s]-[%u]", + PARAM_HNSW_SPARSE_STREAMER_MIN_NEIGHBOR_COUNT.c_str(), + min_neighbor_cnt_, + PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT.c_str(), + upper_max_neighbor_cnt_); + return IndexError_InvalidArgument; + } + + if (bf_negative_prob_ <= 0.0f || bf_negative_prob_ >= 1.0f) { + LOG_ERROR( + "[%s] must be in range (0,1)", + PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str()); + return IndexError_InvalidArgument; + } + + if (scaling_factor_ == 0U) { + scaling_factor_ = HnswSparseEntity::kDefaultScalingFactor; + } + if (scaling_factor_ < 5 || scaling_factor_ > 1000) { + LOG_ERROR("[%s] must be in range [5,1000]", + PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR.c_str()); + return IndexError_InvalidArgument; + } + + if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) { + LOG_ERROR("[%s] must be in range (0.0f,1.0f]", + PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO.c_str()); + return IndexError_InvalidArgument; + } + + if (max_scan_limit_ < min_scan_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT.c_str(), + PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT.c_str()); + return IndexError_InvalidArgument; + } + + if (prune_cnt == 0UL) { + prune_cnt = upper_max_neighbor_cnt_; + } + if (chunk_size_ == 0UL) { + chunk_size_ = HnswSparseEntity::kDefaultChunkSize; + } + if (chunk_size_ > HnswSparseEntity::kMaxChunkSize) { + LOG_ERROR("[%s] must be < %zu", + PARAM_HNSW_SPARSE_STREAMER_CHUNK_SIZE.c_str(), + HnswSparseEntity::kMaxChunkSize); + return IndexError_InvalidArgument; + } + + if (query_filtering_enabled_ && + (query_filtering_ratio_ <= 0.0f || query_filtering_ratio_ >= 1.0f)) { + LOG_ERROR("[%s] must be in range (0, 1)", + PARAM_HNSW_SPARSE_SEARCHER_QUERY_FILTERING_RATIO.c_str()); + return IndexError_InvalidArgument; + } + + entity_.set_ef_construction(ef_construction_); + entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_); + entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_); + entity_.set_scaling_factor(scaling_factor_); + entity_.set_prune_cnt(prune_cnt); + + entity_.set_chunk_size(chunk_size_); + entity_.set_filter_same_key(filter_same_key_); + entity_.set_get_vector(get_vector_enabled_); + entity_.set_min_neighbor_cnt(min_neighbor_cnt_); + + entity_.set_sparse_meta_size(HnswSparseEntity::kSparseMetaSize); + entity_.set_sparse_unit_size(meta_.unit_size()); + + int ret = entity_.init(max_index_size_, docs_hard_limit_); + if (ret != 0) { + LOG_ERROR("Hnsw entity init failed for %s", IndexError::What(ret)); + return ret; + } + LOG_DEBUG( + "Init params: maxIndexSize=%zu docsHardLimit=%zu docsSoftLimit=%zu " + "efConstruction=%u ef=%u l0NeighborCnt=%u upperNeighborCnt=%u " + "scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu " + "bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbability=%.5f " + "checkCrcEnabled=%d pruneSize=%zu chunkSize=%zu " + "filterSameKey=%u getVectorEnabled=%u " + "minNeighborCount=%u forcePadding=%u filteringRatio=%f", + max_index_size_, docs_hard_limit_, docs_soft_limit_, ef_construction_, + ef_, l0_max_neighbor_cnt_, upper_max_neighbor_cnt_, scaling_factor_, + max_scan_ratio_, min_scan_limit_, max_scan_limit_, bf_enabled_, + bruteforce_threshold_, bf_negative_prob_, check_crc_enabled_, prune_cnt, + chunk_size_, filter_same_key_, get_vector_enabled_, min_neighbor_cnt_, + force_padding_topk_enabled_, query_filtering_ratio_); + + alg_ = HnswSparseAlgorithm::UPointer(new HnswSparseAlgorithm(entity_)); + + ret = alg_->init(); + if (ret != 0) { + return ret; + } + + state_ = STATE_INITED; + + return 0; +} + +int HnswSparseStreamer::cleanup(void) { + if (state_ == STATE_OPENED) { + this->close(); + } + + LOG_INFO("HnswSparseStreamer cleanup"); + + meta_.clear(); + metric_.reset(); + stats_.clear(); + entity_.cleanup(); + + if (alg_) { + alg_->cleanup(); + } + + max_index_size_ = 0UL; + docs_hard_limit_ = HnswSparseEntity::kDefaultDocsHardLimit; + docs_soft_limit_ = 0UL; + upper_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt; + ef_ = HnswSparseEntity::kDefaultEf; + ef_construction_ = HnswSparseEntity::kDefaultEfConstruction; + bf_enabled_ = false; + scaling_factor_ = HnswSparseEntity::kDefaultScalingFactor; + bruteforce_threshold_ = HnswSparseEntity::kDefaultBruteForceThreshold; + max_scan_limit_ = HnswSparseEntity::kDefaultMaxScanLimit; + min_scan_limit_ = HnswSparseEntity::kDefaultMinScanLimit; + chunk_size_ = HnswSparseEntity::kDefaultChunkSize; + bf_negative_prob_ = HnswSparseEntity::kDefaultBFNegativeProbability; + max_scan_ratio_ = HnswSparseEntity::kDefaultScanRatio; + state_ = STATE_INIT; + check_crc_enabled_ = false; + filter_same_key_ = false; + get_vector_enabled_ = false; + + sparse_neighbor_ratio_ = HnswSparseEntity::kDefaultSparseNeighborRatio; + sparse_neighbor_cnt_ = 0UL; + sparse_min_neighbor_cnt_ = 0UL; + upper_sparse_neighbor_cnt_ = 0UL; + + return 0; +} + +int HnswSparseStreamer::open(IndexStorage::Pointer stg) { + LOG_INFO("HnswSparseStreamer open"); + + if (ailego_unlikely(state_ != STATE_INITED)) { + LOG_ERROR("Open storage failed, init streamer first!"); + return IndexError_NoReady; + } + int ret = entity_.open(std::move(stg), check_crc_enabled_); + if (ret != 0) { + return ret; + } + IndexMeta index_meta; + ret = entity_.get_index_meta(&index_meta); + if (ret == IndexError_NoExist) { + // Set IndexMeta for the new index + ret = entity_.set_index_meta(meta_); + if (ret != 0) { + LOG_ERROR("Failed to set index meta for %s", IndexError::What(ret)); + return ret; + } + } else if (ret != 0) { + LOG_ERROR("Failed to get index meta for %s", IndexError::What(ret)); + return ret; + } else { + if (index_meta.metric_name() != meta_.metric_name() || + index_meta.data_type() != meta_.data_type()) { + LOG_ERROR("IndexMeta mismatch from the previous in index"); + return IndexError_Mismatch; + } + // The IndexMetric Params may be updated like MipsSquaredEuclidean + auto metric_params = index_meta.metric_params(); + metric_params.merge(meta_.metric_params()); + meta_.set_metric(index_meta.metric_name(), 0, metric_params); + } + + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Failed to create metric %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("Failled to init metric, ret=%d", ret); + return ret; + } + + if (!metric_->sparse_distance()) { + LOG_ERROR("Invalid metric distance"); + return IndexError_InvalidArgument; + } + + add_distance_ = metric_->sparse_distance(); + search_distance_ = add_distance_; + + if (metric_->query_metric() && metric_->query_metric()->distance()) { + search_distance_ = metric_->query_metric()->sparse_distance(); + } + + state_ = STATE_OPENED; + magic_ = IndexContext::GenerateMagic(); + + return 0; +} + +int HnswSparseStreamer::close(void) { + LOG_INFO("HnswSparseStreamer close"); + + stats_.clear(); + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_.set_index_meta(meta_); + int ret = entity_.close(); + if (ret != 0) { + return ret; + } + state_ = STATE_INITED; + + return 0; +} + +int HnswSparseStreamer::flush(uint64_t checkpoint) { + LOG_INFO("HnswSparseStreamer flush checkpoint=%zu", (size_t)checkpoint); + + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_.set_index_meta(meta_); + return entity_.flush(checkpoint); +} + +int HnswSparseStreamer::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("HnswSparseStreamer dump"); + + shared_mutex_.lock(); + AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); + + meta_.set_searcher("HnswSparseSearcher", HnswSparseEntity::kRevision, + ailego::Params()); + + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + return entity_.dump(dumper); +} + +IndexStreamer::Context::Pointer HnswSparseStreamer::create_context(void) const { + if (ailego_unlikely(state_ != STATE_OPENED)) { + LOG_ERROR("Create context failed, open storage first!"); + return Context::Pointer(); + } + + HnswSparseEntity::Pointer entity = entity_.clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("CreateContext clone init failed"); + return Context::Pointer(); + } + HnswSparseContext *ctx = + new (std::nothrow) HnswSparseContext(metric_, entity); + if (ailego_unlikely(ctx == nullptr)) { + LOG_ERROR("Failed to new HnswSparseContext"); + return Context::Pointer(); + } + ctx->set_ef(ef_); + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_filter_mode(bf_enabled_ ? VisitFilter::BloomFilter + : VisitFilter::ByteMap); + ctx->set_filter_negative_probability(bf_negative_prob_); + ctx->set_magic(magic_); + ctx->set_force_padding_topk(force_padding_topk_enabled_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + + if (ailego_unlikely(ctx->init(HnswSparseContext::kSparseStreamerContext)) != + 0) { + LOG_ERROR("Init HnswSparseContext failed"); + delete ctx; + return Context::Pointer(); + } + + return Context::Pointer(ctx); +} + +IndexStreamer::SparseProvider::Pointer +HnswSparseStreamer::create_sparse_provider(void) const { + LOG_DEBUG("HnswSparseStreamer create sparse provider"); + + auto entity = entity_.clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone HnswSparseEntity failed"); + return SparseProvider::Pointer(); + } + return SparseProvider::Pointer( + new HnswSparseIndexProvider(meta_, entity, "HnswSparseStreamer")); +} + +int HnswSparseStreamer::update_context(HnswSparseContext *ctx) const { + const HnswSparseEntity::Pointer entity = entity_.clone(); + if (!entity) { + LOG_ERROR("Failed to clone search context entity"); + return IndexError_Runtime; + } + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + return ctx->update_context(HnswSparseContext::kSparseStreamerContext, meta_, + metric_, entity, magic_); +} + +//! Add a vector with id into index with sparse inputs +int HnswSparseStreamer::add_with_id_impl(uint32_t id, + const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + if (ailego_unlikely(sparse_count > HnswSparseEntity::kSparseMaxDimSize)) { + LOG_WARN( + "Failed to add sparse vector: number of non-zero elements (%u) exceeds " + "maximum allowed (%u), id=%u", + sparse_count, HnswSparseEntity::kSparseMaxDimSize, id); + return IndexError_InvalidValue; + } + + HnswSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswSparseContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) { + if (entity_.doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed [%s]", entity_.doc_cnt(), + PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT.c_str()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %u exceed [%s]", entity_.doc_cnt(), + PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT.c_str()); + } + } + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->update_dist_caculator_distance(add_distance_); + + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_query, + entity_.sparse_unit_size(), + sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + + level_t level = alg_->get_random_level(); + ret = + entity_.add_vector_with_id(level, id, sparse_query_buffer, sparse_count); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, level, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw stramer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + +//! Add a vector into index with sparse inputs +int HnswSparseStreamer::add_impl(uint64_t pkey, const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + if (ailego_unlikely(sparse_count > HnswSparseEntity::kSparseMaxDimSize)) { + LOG_WARN( + "Failed to add sparse vector: number of non-zero elements (%u) exceeds " + "maximum allowed (%u), key=%zu", + sparse_count, HnswSparseEntity::kSparseMaxDimSize, (size_t)pkey); + return IndexError_InvalidValue; + } + + HnswSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswSparseContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) { + if (entity_.doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed [%s]", entity_.doc_cnt(), + PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT.c_str()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %u exceed [%s]", entity_.doc_cnt(), + PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT.c_str()); + } + } + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->update_dist_caculator_distance(add_distance_); + + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_query, + entity_.sparse_unit_size(), + sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + + level_t level = alg_->get_random_level(); + node_id_t id; + ret = entity_.add_vector(level, pkey, sparse_query_buffer, sparse_count, &id); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, level, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw stramer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + +//! Similarity search with sparse inputs +int HnswSparseStreamer::search_impl( + const uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) const { + return search_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1, + context); +} + +//! Similarity search with sparse inputs +int HnswSparseStreamer::search_impl( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, uint32_t count, + IndexStreamer::Context::Pointer &context) const { + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + HnswSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswSparseContext failed"); + return IndexError_Cast; + } + + if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) { + return search_bf_impl(sparse_count, sparse_indices, sparse_query, qmeta, + count, context); + } + + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_); + ctx->resize_results(count); + ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + + const uint32_t *sparse_indices_tmp = sparse_indices; + const void *sparse_query_tmp = sparse_query; + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + std::string sparse_query_filtered_buffer; + + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + if (query_filtering_enabled_) { + if (!SparseUtility::FilterSparseQuery( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + qmeta.data_type(), entity_.sparse_unit_size(), + query_filtering_ratio_, &sparse_query_filtered_buffer)) { + LOG_ERROR("Hnsw filtering failed"); + + return IndexError_Runtime; + } + + ctx->reset_query(sparse_query_filtered_buffer.data()); + } else { + ctx->reset_query(sparse_query_buffer.data()); + } + + ret = alg_->search(ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw searcher fast search failed"); + return ret; + } + + if (query_filtering_enabled_) { + ctx->reset_query(sparse_query_buffer.data()); + ctx->recal_topk_dist(); + } + + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +//! Similarity search with sparse inputs +int HnswSparseStreamer::search_bf_impl( + const uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + IndexStreamer::Context::Pointer &context) const { + return search_bf_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1, + context); +} + +//! Similarity search with sparse inputs +int HnswSparseStreamer::search_bf_impl( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, uint32_t count, + IndexStreamer::Context::Pointer &context) const { + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + HnswSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswSparseContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_); + ctx->resize_results(count); + + const uint32_t *sparse_indices_tmp = sparse_indices; + const void *sparse_query_tmp = sparse_query; + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_Runtime; + } + + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity_.get_key(id)); + }; + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + ctx->group_topk_heaps().clear(); + + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + if (entity_.get_key(id) == kInvalidKey) { + continue; + } + + if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { + dist_t dist = ctx->dist_calculator().dist(id); + + std::string group_id = group_by(id); + + auto &topk_heap = ctx->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, dist); + } + } + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + } else { + auto &filter = ctx->filter(); + auto &topk = ctx->topk_heap(); + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + topk.clear(); + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + if (entity_.get_key(id) == kInvalidKey) { + continue; + } + + if (!filter.is_valid() || !filter(entity_.get_key(id))) { + dist_t dist = ctx->dist_calculator().dist(id); + topk.emplace(id, dist); + } + } + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + } + + return 0; +} + +//! Linear search by primary keys +int HnswSparseStreamer::search_bf_by_p_keys_impl( + const uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, ContextPointer &context) const { + return search_bf_by_p_keys_impl(&sparse_count, sparse_indices, sparse_query, + p_keys, qmeta, 1, context); +} + +//! Linear search by primary keys with sparse inputs +int HnswSparseStreamer::search_bf_by_p_keys_impl( + const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + int ret = check_params(qmeta); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + if (ailego_unlikely(p_keys.size() != count)) { + LOG_ERROR("The size of p_keys is not equal to count"); + return IndexError_InvalidArgument; + } + + HnswSparseContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to HnswSparseContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher or streamer + ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_); + ctx->resize_results(count); + + const uint32_t *sparse_indices_tmp = sparse_indices; + const void *sparse_query_tmp = sparse_query; + + if (ctx->group_by_search()) { + if (!ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by function"); + return IndexError_Runtime; + } + + std::function group_by = [&](node_id_t id) { + return ctx->group_by()(entity_.get_key(id)); + }; + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + ctx->group_topk_heaps().clear(); + + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + uint64_t pk = p_keys[q][idx]; + if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { + node_id_t id = entity_.get_id(pk); + if (id != kInvalidNodeId) { + dist_t dist = ctx->dist_calculator().dist(id); + + std::string group_id = group_by(id); + + auto &topk_heap = ctx->group_topk_heaps()[group_id]; + if (topk_heap.empty()) { + topk_heap.limit(ctx->group_topk()); + } + topk_heap.emplace_back(id, dist); + } + } + } + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + } else { + auto &filter = ctx->filter(); + auto &topk = ctx->topk_heap(); + + for (size_t q = 0; q < count; ++q) { + std::string sparse_query_buffer; + SparseUtility::TransSparseFormat( + sparse_count[q], sparse_indices_tmp, sparse_query_tmp, + entity_.sparse_unit_size(), sparse_query_buffer); + + ctx->reset_query(sparse_query_buffer.data()); + topk.clear(); + for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { + key_t pk = p_keys[q][idx]; + if (!filter.is_valid() || !filter(pk)) { + node_id_t id = entity_.get_id(pk); + if (id != kInvalidNodeId) { + dist_t dist = ctx->dist_calculator().dist(id); + topk.emplace(id, dist); + } + } + } + ctx->topk_to_result(q); + + sparse_indices_tmp += sparse_count[q]; + sparse_query_tmp = reinterpret_cast(sparse_query_tmp) + + sparse_count[q] * qmeta.unit_size(); + } + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +void HnswSparseStreamer::print_debug_info() { + for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { + Neighbors neighbours = entity_.get_neighbors(0, id); + std::cout << "node: " << id << "; "; + for (uint32_t i = 0; i < neighbours.size(); ++i) { + std::cout << neighbours[i]; + + if (i == neighbours.size() - 1) { + std::cout << std::endl; + } else { + std::cout << ", "; + } + } + } + + // entity_.print_key_map(); +} + +INDEX_FACTORY_REGISTER_STREAMER(HnswSparseStreamer); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.h new file mode 100644 index 0000000..da80d09 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.h @@ -0,0 +1,226 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "hnsw_sparse_algorithm.h" +#include "hnsw_sparse_streamer_entity.h" + +namespace zvec { +namespace core { + +class HnswSparseStreamer : public IndexStreamer { + public: + using ContextPointer = IndexStreamer::Context::Pointer; + + HnswSparseStreamer(void); + ~HnswSparseStreamer(void) override; + + HnswSparseStreamer(const HnswSparseStreamer &streamer) = delete; + HnswSparseStreamer &operator=(const HnswSparseStreamer &streamer) = delete; + + protected: + //! Initialize Streamer + int init(const IndexMeta &imeta, const ailego::Params ¶ms) override; + + //! Cleanup Streamer + int cleanup(void) override; + + //! Create a context + Context::Pointer create_context(void) const override; + + //! Create a new sparse iterator + IndexStreamer::SparseProvider::Pointer create_sparse_provider( + void) const override; + + int add_impl(uint64_t pkey, const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, Context::Pointer &context) override; + + int add_with_id_impl(uint32_t id, const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Similarity search with sparse inputs + int search_impl(const uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search with sparse inputs + int search_impl(const uint32_t *sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Similarity brute force search with sparse inputs + int search_bf_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search with sparse inputs + int search_bf_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const override; + + //! Linear search by primary keys + int search_bf_by_p_keys_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override; + + //! Linear search by primary keys with sparse inputs + int search_bf_by_p_keys_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + //! Fetch sparse vector by key + int get_sparse_vector(uint64_t key, uint32_t *sparse_count, + std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override { + return entity_.get_sparse_vector_by_key( + key, sparse_count, sparse_indices_buffer, sparse_values_buffer); + } + + //! Fetch vector by id + int get_sparse_vector_by_id( + uint32_t id, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override { + return entity_.get_sparse_vector_by_id( + id, sparse_count, sparse_indices_buffer, sparse_values_buffer); + } + + //! Open index from file path + int open(IndexStorage::Pointer stg) override; + + //! Close file + int close(void) override; + + //! flush file + int flush(uint64_t checkpoint) override; + + //! Dump index into storage + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve sparse meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + void print_debug_info() override; + + private: + inline int check_params(const IndexQueryMeta &qmeta) const { + if (ailego_unlikely(qmeta.data_type() != meta_.data_type())) { + LOG_ERROR("Unsupported query meta"); + return IndexError_Mismatch; + } + return 0; + } + + inline int check_sparse_count_is_zero(const uint32_t *sparse_count, + uint32_t count) const { + for (uint32_t i = 0; i < count; ++i) { + if (sparse_count[i] != 0) + LOG_ERROR("Sparse cout is not empty. Index: %u, Sparse Count: %u", i, + sparse_count[i]); + return IndexError_InvalidArgument; + } + + return 0; + } + + private: + //! To share ctx across streamer/searcher, we need to update the context for + //! current streamer/searcher + int update_context(HnswSparseContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 }; + class Stats : public IndexStreamer::Stats { + public: + void clear(void) { + set_revision_id(0u); + set_loaded_count(0u); + set_added_count(0u); + set_discarded_count(0u); + set_index_size(0u); + set_dumped_size(0u); + set_check_point(0u); + set_create_time(0u); + set_update_time(0u); + clear_attributes(); + } + }; + + Stats stats_{}; + HnswSparseStreamerEntity entity_; + HnswSparseAlgorithm::UPointer alg_; + IndexMeta meta_{}; + IndexMetric::Pointer metric_{}; + + IndexMetric::MatrixSparseDistance add_distance_{}; + IndexMetric::MatrixSparseDistance search_distance_{}; + std::mutex mutex_{}; + + size_t max_index_size_{0UL}; + size_t chunk_size_{HnswSparseEntity::kDefaultChunkSize}; + size_t docs_hard_limit_{HnswSparseEntity::kDefaultDocsHardLimit}; + size_t docs_soft_limit_{0UL}; + uint32_t min_neighbor_cnt_{0u}; + uint32_t upper_max_neighbor_cnt_{ + HnswSparseEntity::kDefaultUpperMaxNeighborCnt}; + uint32_t l0_max_neighbor_cnt_{HnswSparseEntity::kDefaultL0MaxNeighborCnt}; + uint32_t ef_{HnswSparseEntity::kDefaultEf}; + uint32_t ef_construction_{HnswSparseEntity::kDefaultEfConstruction}; + uint32_t scaling_factor_{HnswSparseEntity::kDefaultScalingFactor}; + size_t bruteforce_threshold_{HnswSparseEntity::kDefaultBruteForceThreshold}; + size_t max_scan_limit_{HnswSparseEntity::kDefaultMaxScanLimit}; + size_t min_scan_limit_{HnswSparseEntity::kDefaultMinScanLimit}; + float bf_negative_prob_{HnswSparseEntity::kDefaultBFNegativeProbability}; + float max_scan_ratio_{HnswSparseEntity::kDefaultScanRatio}; + float sparse_neighbor_ratio_{HnswSparseEntity::kDefaultSparseNeighborRatio}; + uint32_t sparse_neighbor_cnt_{0UL}; + uint32_t sparse_min_neighbor_cnt_{0UL}; + uint32_t upper_sparse_neighbor_cnt_{0UL}; + + bool query_filtering_enabled_{false}; + float query_filtering_ratio_{HnswSparseEntity::kDefaultQueryFilteringRatio}; + + uint32_t magic_{0U}; + State state_{STATE_INIT}; + bool bf_enabled_{false}; + bool check_crc_enabled_{false}; + bool filter_same_key_{false}; + bool get_vector_enabled_{false}; + bool force_padding_topk_enabled_{false}; + + //! avoid add vector while dumping index + ailego::SharedMutex shared_mutex_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.cc new file mode 100644 index 0000000..90d6b3a --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.cc @@ -0,0 +1,1001 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "hnsw_sparse_streamer_entity.h" +#include +#include +#include +#include +#include "utility/sparse_utility.h" +#include "hnsw_sparse_dist_calculator.h" + +namespace zvec { +namespace core { + +HnswSparseStreamerEntity::HnswSparseStreamerEntity(IndexStreamer::Stats &stats) + : stats_(stats) {} + +HnswSparseStreamerEntity::~HnswSparseStreamerEntity() {} + +int HnswSparseStreamerEntity::init(uint64_t max_index_size, + size_t max_doc_cnt) { + if (std::pow(scaling_factor(), kMaxGraphLayers) < max_doc_cnt) { + LOG_ERROR("scalingFactor=%zu is too small", scaling_factor()); + return IndexError_InvalidArgument; + } + + std::lock_guard lock(mutex_); + broker_ = std::make_shared(stats_); + upper_neighbor_index_ = std::make_shared(); + keys_map_lock_ = std::make_shared(); + keys_map_ = std::make_shared>(); + if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) { + LOG_ERROR("HnswSparseStreamerEntity new object failed"); + return IndexError_NoMemory; + } + keys_map_->set_empty_key(kInvalidKey); + + neighbor_size_ = neighbors_size(); + upper_neighbor_size_ = upper_neighbors_size(); + + //! vector + key + level 0 neighbors + size_t size = sizeof(key_t) + neighbor_size_ + sparse_meta_size(); + + size = AlignSize(size); + set_node_size(size); + + return init_chunk_params(max_index_size); +} + +int HnswSparseStreamerEntity::cleanup() { + std::lock_guard lock(mutex_); + mutable_header()->clear(); + chunk_size_ = kDefaultChunkSize; + node_index_mask_bits_ = 0U; + node_index_mask_ = 0U; + node_cnt_per_chunk_ = 0U; + neighbor_size_ = 0U; + upper_neighbor_size_ = 0U; + if (upper_neighbor_index_) { + upper_neighbor_index_->cleanup(); + } + if (keys_map_) { + keys_map_->clear(); + } + node_chunks_.clear(); + upper_neighbor_chunks_.clear(); + filter_same_key_ = false; + get_vector_enabled_ = false; + broker_.reset(); + + return 0; +} + +int HnswSparseStreamerEntity::update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) { + std::vector buffer(neighbor_size_); + NeighborsHeader *hd = reinterpret_cast(buffer.data()); + hd->neighbor_cnt = neighbors.size(); + size_t i = 0; + for (; i < neighbors.size(); ++i) { + hd->neighbors[i] = neighbors[i].first; + } + + auto loc = get_neighbor_chunk_loc(level, id); + size_t size = reinterpret_cast(&hd->neighbors[i]) - buffer.data(); + size_t ret = loc.first->write(loc.second, hd, size); + if (ailego_unlikely(ret != size)) { + LOG_ERROR("Write neighbor header failed, ret=%zu", ret); + + return IndexError_Runtime; + } + + return 0; +} + +const Neighbors HnswSparseStreamerEntity::get_neighbors(level_t level, + node_id_t id) const { + SparseChunk *chunk = nullptr; + size_t offset = 0UL; + size_t neighbor_size = neighbor_size_; + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + offset = (id & node_index_mask_) * node_size() + sizeof(key_t) + + sparse_meta_size(); + + sync_chunks(SparseChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + chunk = node_chunks_[chunk_idx].get(); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + chunk = upper_neighbor_chunks_[p.first].get(); + offset = p.second; + neighbor_size = upper_neighbor_size_; + } + + ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); + IndexStorage::MemoryBlock neighbor_block; + size_t size = chunk->read(offset, neighbor_block, neighbor_size); + if (ailego_unlikely(size != neighbor_size)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", size); + return Neighbors(); + } + return Neighbors(std::move(neighbor_block)); +} + +//! Get vector feature data by key +const void *HnswSparseStreamerEntity::get_vector_meta(node_id_t id) const { + auto loc = get_vector_chunk_loc(id); + const void *vec = nullptr; + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t read_size = sparse_meta_size(); + + size_t ret = node_chunks_[loc.first]->read(loc.second, &vec, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + } + + return vec; +} + +int HnswSparseStreamerEntity::get_vector_meta( + const node_id_t id, IndexStorage::MemoryBlock &block) const { + auto loc = get_vector_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t read_size = sparse_meta_size(); + + size_t ret = node_chunks_[loc.first]->read(loc.second, block, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + + return 0; +} + +int HnswSparseStreamerEntity::get_vector_metas(const node_id_t *ids, + uint32_t count, + const void **vecs) const { + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t read_size = sparse_meta_size(); + + size_t ret = node_chunks_[loc.first]->read(loc.second, &vecs[i], read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + } + + return 0; +} + +int HnswSparseStreamerEntity::get_vector_metas( + const node_id_t *ids, uint32_t count, + std::vector &block_vecs) const { + block_vecs.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t read_size = sparse_meta_size(); + + size_t ret = + node_chunks_[loc.first]->read(loc.second, block_vecs[i], read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + } + + return 0; +} + +//! Get vector feature data by key +const void *HnswSparseStreamerEntity::get_sparse_data(uint64_t offset, + uint32_t len) const { + uint32_t chunk_index = offset >> 32; + uint32_t chunk_offset = offset & 0xFFFFFFFF; + + auto loc = get_sparse_chunk_loc(chunk_index, chunk_offset); + const void *data = nullptr; + + ailego_assert_with(loc.first < sparse_node_chunks_.size(), + "invalid chunk idx"); + ailego_assert_with(loc.second < sparse_node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t ret = sparse_node_chunks_[loc.first]->read(loc.second, &data, len); + if (ailego_unlikely(ret != len)) { + LOG_ERROR("Read sparse vector failed, offset=%zu, read size=%u, ret=%zu", + (size_t)offset, len, ret); + } + return data; +} + +int HnswSparseStreamerEntity::get_sparse_data( + uint64_t offset, uint32_t len, IndexStorage::MemoryBlock &block) const { + uint32_t chunk_index = offset >> 32; + uint32_t chunk_offset = offset & 0xFFFFFFFF; + + auto loc = get_sparse_chunk_loc(chunk_index, chunk_offset); + ailego_assert_with(loc.first < sparse_node_chunks_.size(), + "invalid chunk idx"); + ailego_assert_with(loc.second < sparse_node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + + size_t ret = sparse_node_chunks_[loc.first]->read(loc.second, block, len); + if (ailego_unlikely(ret != len)) { + LOG_ERROR("Read sparse vector failed, offset=%zu, read size=%u, ret=%zu", + (size_t)offset, len, ret); + return IndexError_ReadData; + } + return 0; +} + +//! Get sparse data from id +const void *HnswSparseStreamerEntity::get_sparse_data(node_id_t id) const { + auto sparse_data = get_sparse_data_from_vector(get_vector_meta(id)); + + return sparse_data.first; +} + +int HnswSparseStreamerEntity::get_sparse_data( + node_id_t id, IndexStorage::MemoryBlock &block) const { + IndexStorage::MemoryBlock meta_block; + get_vector_meta(id, meta_block); + int sparse_length = 0; + return get_sparse_data_from_vector(meta_block.data(), block, sparse_length); +} + +//! Get sparse data from vector +std::pair +HnswSparseStreamerEntity::get_sparse_data_from_vector(const void *vec) const { + const char *vec_ptr = reinterpret_cast(vec); + + uint64_t offset = *((uint64_t *)(vec_ptr)); + uint32_t sparse_vector_len = *((uint32_t *)(vec_ptr + sizeof(uint64_t))); + + if (sparse_vector_len > 0) { + const void *sparse_data = get_sparse_data(offset, sparse_vector_len); + if (ailego_unlikely(sparse_data == nullptr)) { + LOG_ERROR("Get nullptr sparse, offset=%zu, len=%u", (size_t)offset, + sparse_vector_len); + + return std::make_pair(nullptr, 0); + } + + return std::make_pair(sparse_data, sparse_vector_len); + } + + return std::make_pair(nullptr, 0); +} + +int HnswSparseStreamerEntity::get_sparse_data_from_vector( + const void *vec, IndexStorage::MemoryBlock &block, + int &sparse_length) const { + const char *vec_ptr = reinterpret_cast(vec); + + uint64_t offset = *((uint64_t *)(vec_ptr)); + uint32_t sparse_vector_len = *((uint32_t *)(vec_ptr + sizeof(uint64_t))); + + if (sparse_vector_len > 0) { + int ret = get_sparse_data(offset, sparse_vector_len, block); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Get nullptr sparse, offset=%zu, len=%u", (size_t)offset, + sparse_vector_len); + return IndexError_ReadData; + } + sparse_length = sparse_vector_len; + } + return 0; +} + +key_t HnswSparseStreamerEntity::get_key(node_id_t id) const { + auto loc = get_key_chunk_loc(id); + IndexStorage::MemoryBlock key_block; + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t ret = + node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return kInvalidKey; + } + + return *reinterpret_cast(key_block.data()); +} + +void HnswSparseStreamerEntity::add_neighbor(level_t level, node_id_t id, + uint32_t size, + node_id_t neighbor_id) { + auto loc = get_neighbor_chunk_loc(level, id); + size_t offset = + loc.second + sizeof(NeighborsHeader) + size * sizeof(node_id_t); + ailego_assert_with(size < neighbor_cnt(level), "invalid neighbor size"); + ailego_assert_with(offset < loc.first->data_size(), "invalid chunk offset"); + size_t ret = loc.first->write(offset, &neighbor_id, sizeof(node_id_t)); + if (ailego_unlikely(ret != sizeof(node_id_t))) { + LOG_ERROR("Write neighbor id failed, ret=%zu", ret); + return; + } + + uint32_t neighbors = size + 1; + ret = loc.first->write(loc.second, &neighbors, sizeof(uint32_t)); + if (ailego_unlikely(ret != sizeof(uint32_t))) { + LOG_ERROR("Write neighbor cnt failed, ret=%zu", ret); + } + + return; +} + +int HnswSparseStreamerEntity::init_chunks( + const SparseChunk::Pointer &header_chunk) { + if (header_chunk->data_size() < header_size()) { + LOG_ERROR("Invalid header chunk size"); + return IndexError_InvalidFormat; + } + IndexStorage::MemoryBlock data_block; + size_t size = header_chunk->read(0UL, data_block, header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Read header chunk failed"); + return IndexError_ReadData; + } + *mutable_header() = + *reinterpret_cast(data_block.data()); + + int ret = check_hnsw_index(&header()); + if (ret != 0) { + broker_->close(); + return ret; + } + + node_chunks_.resize( + broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_NODE)); + for (auto seq = 0UL; seq < node_chunks_.size(); ++seq) { + node_chunks_[seq] = + broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_NODE, seq); + if (!node_chunks_[seq]) { + LOG_ERROR("Missing hnsw streamer data chunk %zu th of %zu", seq, + node_chunks_.size()); + return IndexError_InvalidFormat; + } + } + + upper_neighbor_chunks_.resize( + broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR)); + for (auto seq = 0UL; seq < upper_neighbor_chunks_.size(); ++seq) { + upper_neighbor_chunks_[seq] = + broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, seq); + if (!upper_neighbor_chunks_[seq]) { + LOG_ERROR("Missing hnsw streamer index chunk %zu th of %zu", seq, + upper_neighbor_chunks_.size()); + return IndexError_InvalidFormat; + } + } + + sparse_node_chunks_.resize( + broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE)); + for (auto seq = 0UL; seq < sparse_node_chunks_.size(); ++seq) { + sparse_node_chunks_[seq] = + broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE, seq); + if (!sparse_node_chunks_[seq]) { + LOG_ERROR("Missing hnsw streamer sparse data chunk %zu th of %zu", seq, + sparse_node_chunks_.size()); + return IndexError_InvalidFormat; + } + } + + return 0; +} + +int HnswSparseStreamerEntity::open(IndexStorage::Pointer stg, bool check_crc) { + std::lock_guard lock(mutex_); + int ret = + broker_->open(std::move(stg), max_index_size_, chunk_size_, check_crc); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Open index failed for %s", IndexError::What(ret)); + return ret; + } + ret = upper_neighbor_index_->init(broker_, upper_neighbor_chunk_size_, + scaling_factor(), estimate_doc_capacity(), + kUpperHashMemoryInflateRatio); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Init neighbor hash map failed"); + return ret; + } + + //! init header + auto header_chunk = broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_HEADER, + SparseChunkBroker::kDefaultChunkSeqId); + if (!header_chunk) { // open empty index, create one + auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_HEADER, + SparseChunkBroker::kDefaultChunkSeqId, + header_size()); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc header chunk failed"); + return p.first; + } + size_t size = p.second->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + return 0; + } + + //! Open an exist hnsw index + ret = init_chunks(header_chunk); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + //! total docs including features wrote in index but neighbors may not ready + node_id_t total_vecs = 0; + if (node_chunks_.size() > 0) { + size_t last_idx = node_chunks_.size() - 1; + auto last_chunk = node_chunks_[last_idx]; + if (last_chunk->data_size() % node_size()) { + LOG_WARN("The index may broken"); + return IndexError_InvalidFormat; + } + total_vecs = last_idx * node_cnt_per_chunk_ + + node_chunks_[last_idx]->data_size() / node_size(); + } + + LOG_INFO( + "Open index, l0NeighborCnt=%zu upperneighborCnt=%zu " + "efConstruction=%zu curDocCnt=%u totalVecs=%u maxLevel=%u", + l0_neighbor_cnt(), upper_neighbor_cnt(), ef_construction(), doc_cnt(), + total_vecs, cur_max_level()); + //! try to correct the docCnt if index not fully flushed + if (doc_cnt() != total_vecs) { + LOG_WARN("Index closed abnormally, using totalVecs as curDocCnt"); + *mutable_doc_cnt() = total_vecs; + } + if (filter_same_key_ || get_vector_enabled_) { + for (node_id_t id = 0U; id < doc_cnt(); ++id) { + (*keys_map_)[get_key(id)] = id; + } + } + + stats_.set_loaded_count(doc_cnt()); + + return 0; +} + +int HnswSparseStreamerEntity::close() { + LOG_DEBUG("close index"); + + std::lock_guard lock(mutex_); + flush_header(); + mutable_header()->reset(); + upper_neighbor_index_->cleanup(); + keys_map_->clear(); + header_.clear(); + node_chunks_.clear(); + upper_neighbor_chunks_.clear(); + + sparse_node_chunks_.clear(); + + return broker_->close(); +} + +int HnswSparseStreamerEntity::flush(uint64_t checkpoint) { + LOG_INFO("Flush index, curDocs=%u", doc_cnt()); + + std::lock_guard lock(mutex_); + flush_header(); + int ret = broker_->flush(checkpoint); + if (ret != 0) { + return ret; + } + + return 0; +} + +int HnswSparseStreamerEntity::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("Dump index, curDocs=%u", doc_cnt()); + + //! sort by keys, to support get_vector by key in searcher + std::vector keys(doc_cnt()); + for (node_id_t i = 0; i < doc_cnt(); ++i) { + keys[i] = get_key(i); + } + + //! dump neighbors + auto get_level = [&](node_id_t id) { + auto it = upper_neighbor_index_->find(id); + if (it == upper_neighbor_index_->end()) { + return 0U; + }; + auto meta = reinterpret_cast(&it->second); + return meta->bits.level; + }; + auto ret = dump_segments(dumper, keys.data(), get_level); + if (ailego_unlikely(ret < 0)) { + return ret; + } + *stats_.mutable_dumped_size() += ret; + + return 0; +} + +int HnswSparseStreamerEntity::check_hnsw_index( + const HNSWSparseHeader *hd) const { + if (l0_neighbor_cnt() != hd->neighbor_cnt() || + upper_neighbor_cnt() != hd->upper_neighbor_cnt()) { + LOG_ERROR("Param neighbors:%zu:%zu mismatch index previous %zu:%zu", + l0_neighbor_cnt(), upper_neighbor_cnt(), hd->neighbor_cnt(), + hd->upper_neighbor_cnt()); + return IndexError_Mismatch; + } + if (ef_construction() != hd->ef_construction()) { + LOG_WARN("Param efConstruction %zu mismatch index previous %zu", + ef_construction(), hd->ef_construction()); + } + if (scaling_factor() != hd->scaling_factor()) { + LOG_WARN("Param scalingFactor %zu mismatch index previous %zu", + scaling_factor(), hd->scaling_factor()); + return IndexError_Mismatch; + } + if (prune_cnt() != hd->neighbor_prune_cnt()) { + LOG_WARN("Param pruneCnt %zu mismatch index previous %zu", prune_cnt(), + hd->neighbor_prune_cnt()); + return IndexError_Mismatch; + } + if ((hd->entry_point() != kInvalidNodeId && + hd->entry_point() >= hd->doc_cnt()) || + (hd->entry_point() == kInvalidNodeId && hd->doc_cnt() > 0U)) { + LOG_WARN("Invalid entryPoint %u, docCnt %u", hd->entry_point(), + hd->doc_cnt()); + return IndexError_InvalidFormat; + } + if (hd->entry_point() == kInvalidNodeId && + broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_NODE) > 0) { + LOG_WARN("The index is broken, maybe it haven't flush"); + return IndexError_InvalidFormat; + } + + return 0; +} + +int HnswSparseStreamerEntity::add_vector(level_t level, key_t key, + const std::string &sparse_vec, + uint32_t sparse_count, node_id_t *id) { + // allocat sparse chunk + uint32_t sparse_vector_len = sparse_vec.size(); + + sparse_vector_len = AlignSize(sparse_vector_len); + + if (sparse_vector_len > sparse_chunk_size_) { + LOG_ERROR( + "Sparse Vector Length exceed the chunk size, sparse vec len: %u, chunk " + "size: %u", + sparse_vector_len, sparse_chunk_size_); + return IndexError_InvalidArgument; + } + + SparseChunk::Pointer node_chunk; + SparseChunk::Pointer sparse_node_chunk; + + size_t chunk_offset = static_cast(-1); + size_t sparse_chunk_offset = static_cast(-1); + + std::lock_guard lock(mutex_); + // duplicate check + if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) { + LOG_WARN("Try to add duplicate key, ignore it"); + return IndexError_Duplicate; + } + + node_id_t local_id = static_cast(doc_cnt()); + + uint32_t chunk_index = node_chunks_.size() - 1U; + if (chunk_index == -1U || + (node_chunks_[chunk_index]->data_size() >= + node_cnt_per_chunk_ * node_size())) { // no space left and need to alloc + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + chunk_index++; + auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_NODE, + chunk_index, chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + chunk_offset = 0UL; + node_chunks_.emplace_back(node_chunk); + } else { + node_chunk = node_chunks_[chunk_index]; + chunk_offset = node_chunk->data_size(); + } + + uint32_t sparse_chunk_index = sparse_node_chunks_.size() - 1U; + if (sparse_chunk_index == -1U || + sparse_node_chunks_[sparse_chunk_index]->data_size() + sparse_vector_len > + sparse_chunk_size_) { + if (ailego_unlikely(sparse_node_chunks_.capacity() == + sparse_node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + sparse_chunk_index++; + auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE, + sparse_chunk_index, sparse_chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + sparse_node_chunk = p.second; + + sparse_node_chunks_.emplace_back(sparse_node_chunk); + + sparse_chunk_offset = 0UL; + } else { + sparse_node_chunk = sparse_node_chunks_[sparse_chunk_index]; + sparse_chunk_offset = sparse_node_chunk->data_size(); + } + + // write sparse vector + if (sparse_vec.size() > 0) { + size_t size = sparse_node_chunk->write( + sparse_chunk_offset, sparse_vec.data(), sparse_vec.size()); + if (ailego_unlikely(size != sparse_vec.size())) { + LOG_ERROR("SparseChunk write sparse vec failed, ret=%zu", size); + return IndexError_WriteData; + } + } + + uint64_t sparse_offset = sparse_chunk_index; + sparse_offset = (sparse_offset << 32) + sparse_chunk_offset; + + size_t size = + node_chunk->write(chunk_offset, &sparse_offset, sizeof(uint64_t)); + if (ailego_unlikely(size != sizeof(uint64_t))) { + LOG_ERROR("SparseChunk write sparse vec index failed, ret=%zu", size); + return IndexError_WriteData; + } + + size = node_chunk->write(chunk_offset + sizeof(uint64_t), &sparse_vector_len, + sizeof(uint32_t)); + if (ailego_unlikely(size != sizeof(uint32_t))) { + LOG_ERROR("SparseChunk write sparse vec len failed, ret=%zu", size); + return IndexError_WriteData; + } + + size = + node_chunk->write(chunk_offset + sparse_meta_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("SparseChunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + + //! level 0 neighbors is inited to zero by default + int ret = add_upper_neighbor(level, local_id); + if (ret != 0) { + return ret; + } + + if (sparse_vector_len > 0) { + sparse_chunk_offset += sparse_vector_len; + if (ailego_unlikely(sparse_node_chunk->resize(sparse_chunk_offset) != + sparse_chunk_offset)) { + LOG_ERROR("SparseChunk resize to %zu failed", sparse_chunk_offset); + return IndexError_Runtime; + } + } + + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("SparseChunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + + if (filter_same_key_ || get_vector_enabled_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = local_id; + keys_map_lock_->unlock(); + } + + *mutable_doc_cnt() += 1; + *mutable_total_sparse_count() += sparse_count; + + broker_->mark_dirty(); + *id = local_id; + + return 0; +} + +int HnswSparseStreamerEntity::add_vector_with_id(level_t level, node_id_t id, + const std::string &sparse_vec, + uint32_t sparse_count) { + key_t key = id; + SparseChunk::Pointer node_chunk; + SparseChunk::Pointer sparse_node_chunk; + size_t chunk_offset = static_cast(-1); + size_t sparse_chunk_offset = static_cast(-1); + + // allocat sparse chunk + uint32_t sparse_vector_len = sparse_vec.size(); + + sparse_vector_len = AlignSize(sparse_vector_len); + + if (sparse_vector_len > sparse_chunk_size_) { + LOG_ERROR( + "Sparse Vector Length exceed the chunk size, sparse vec len: %u, chunk " + "size: %u", + sparse_vector_len, sparse_chunk_size_); + return IndexError_InvalidArgument; + } + + + std::lock_guard lock(mutex_); + + // duplicate check + if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) { + LOG_WARN("Try to add duplicate key, ignore it"); + return IndexError_Duplicate; + } + + auto func_get_sparse_node_chunk_and_offset = [&](node_id_t node_id) -> int { + uint32_t chunk_index = node_id >> node_index_mask_bits_; + ailego_assert_with(chunk_index <= node_chunks_.size(), "invalid chunk idx"); + // belongs to next chunk + if (chunk_index == node_chunks_.size()) { + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_NODE, + chunk_index, chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + node_chunks_.emplace_back(node_chunk); + } + + node_chunk = node_chunks_[chunk_index]; + chunk_offset = (node_id & node_index_mask_) * node_size(); + return 0; + }; + + for (size_t start_id = doc_cnt(); start_id < id; ++start_id) { + if (auto ret = func_get_sparse_node_chunk_and_offset(start_id); ret != 0) { + LOG_ERROR("func_get_sparse_node_chunk_and_offset failed"); + return ret; + } + size_t size = node_chunk->write(chunk_offset + sparse_meta_size(), + &kInvalidKey, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("SparseChunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("SparseChunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (auto ret = func_get_sparse_node_chunk_and_offset(id); ret != 0) { + LOG_ERROR("func_get_sparse_node_chunk_and_offset failed"); + return ret; + } + + uint32_t sparse_chunk_index = sparse_node_chunks_.size() - 1U; + if (sparse_chunk_index == -1U || + sparse_node_chunks_[sparse_chunk_index]->data_size() + sparse_vector_len > + sparse_chunk_size_) { + if (ailego_unlikely(sparse_node_chunks_.capacity() == + sparse_node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + sparse_chunk_index++; + auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE, + sparse_chunk_index, sparse_chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + sparse_node_chunk = p.second; + + sparse_node_chunks_.emplace_back(sparse_node_chunk); + + sparse_chunk_offset = 0UL; + } else { + sparse_node_chunk = sparse_node_chunks_[sparse_chunk_index]; + sparse_chunk_offset = sparse_node_chunk->data_size(); + } + + // write sparse vector + if (sparse_vec.size() > 0) { + size_t size = sparse_node_chunk->write( + sparse_chunk_offset, sparse_vec.data(), sparse_vec.size()); + if (ailego_unlikely(size != sparse_vec.size())) { + LOG_ERROR("SparseChunk write sparse vec failed, ret=%zu", size); + return IndexError_WriteData; + } + } + + uint64_t sparse_offset = sparse_chunk_index; + sparse_offset = (sparse_offset << 32) + sparse_chunk_offset; + + size_t size = + node_chunk->write(chunk_offset, &sparse_offset, sizeof(uint64_t)); + if (ailego_unlikely(size != sizeof(uint64_t))) { + LOG_ERROR("SparseChunk write sparse vec index failed, ret=%zu", size); + return IndexError_WriteData; + } + + size = node_chunk->write(chunk_offset + sizeof(uint64_t), &sparse_vector_len, + sizeof(uint32_t)); + if (ailego_unlikely(size != sizeof(uint32_t))) { + LOG_ERROR("SparseChunk write sparse vec len failed, ret=%zu", size); + return IndexError_WriteData; + } + + size = + node_chunk->write(chunk_offset + sparse_meta_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("SparseChunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + + //! level 0 neighbors is inited to zero by default + int ret = add_upper_neighbor(level, id); + if (ret != 0) { + return ret; + } + + if (sparse_vector_len > 0) { + sparse_chunk_offset += sparse_vector_len; + if (ailego_unlikely(sparse_node_chunk->resize(sparse_chunk_offset) != + sparse_chunk_offset)) { + LOG_ERROR("SparseChunk resize to %zu failed", sparse_chunk_offset); + return IndexError_Runtime; + } + } + + if (*mutable_doc_cnt() <= id) { + *mutable_doc_cnt() = id + 1; + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + *mutable_total_sparse_count() += sparse_count; + + if (filter_same_key_ || get_vector_enabled_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = id; + keys_map_lock_->unlock(); + } + + broker_->mark_dirty(); + + return 0; +} + +void HnswSparseStreamerEntity::update_ep_and_level(node_id_t ep, + level_t level) { + HnswSparseEntity::update_ep_and_level(ep, level); + flush_header(); + + return; +} + +const HnswSparseEntity::Pointer HnswSparseStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswSparseStreamerEntity get chunk failed in clone"); + return HnswSparseEntity::Pointer(); + } + } + + std::vector sparse_node_chunks; + sparse_node_chunks.reserve(sparse_node_chunks_.size()); + for (size_t i = 0UL; i < sparse_node_chunks_.size(); ++i) { + sparse_node_chunks.emplace_back(sparse_node_chunks_[i]->clone()); + if (ailego_unlikely(!sparse_node_chunks[i])) { + LOG_ERROR("HnswSparseStreamerEntity get sparse chunk failed in clone"); + return HnswSparseEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswSparseStreamerEntity get chunk failed in clone"); + return HnswSparseEntity::Pointer(); + } + } + + HnswSparseStreamerEntity *entity = + new (std::nothrow) HnswSparseStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + sparse_chunk_size_, upper_neighbor_index_, keys_map_lock_, keys_map_, + std::move(node_chunks), std::move(upper_neighbor_chunks), + std::move(sparse_node_chunks), broker_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswSparseStreamerEntity new failed"); + } + return HnswSparseEntity::Pointer(entity); +} + +//! Get sparse vector feature data by key +int HnswSparseStreamerEntity::get_sparse_vector_by_key( + key_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const { + *sparse_count = 0; + + auto id = get_id(key); + if (id == kInvalidNodeId) { + return IndexError_NoExist; + } + + return get_sparse_vector_by_id(id, sparse_count, sparse_indices_buffer, + sparse_values_buffer); +} + +int HnswSparseStreamerEntity::get_sparse_vector_by_id( + node_id_t id, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const { + IndexStorage::MemoryBlock block; + get_sparse_data(id, block); + const void *sparse_data = block.data(); + if (sparse_data == nullptr) { + return IndexError_InvalidValue; + } + + SparseUtility::ReverseSparseFormat(sparse_data, sparse_count, + sparse_indices_buffer, + sparse_values_buffer, sparse_unit_size()); + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.h new file mode 100644 index 0000000..a474561 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.h @@ -0,0 +1,540 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include "hnsw_sparse_chunk.h" +#include "hnsw_sparse_entity.h" +#include "hnsw_sparse_index_hash.h" +#include "hnsw_sparse_params.h" + +namespace zvec { +namespace core { + +//! HnswSparseStreamerEntity manage vector data, pkey, and node's neighbors +class HnswSparseStreamerEntity : public HnswSparseEntity { + public: + //! Cleanup + //! return 0 on success, or errCode in failure + int cleanup() override; + + //! Make a copy of streamer entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + const HnswSparseEntity::Pointer clone() const override; + + //! Get primary key of the node id + key_t get_key(node_id_t id) const override; + + //! Get vector feature data by key + const void *get_vector_meta(node_id_t id) const override; + + int get_vector_meta(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get vectors feature data by local ids + int get_vector_metas(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + int get_vector_metas( + const node_id_t *ids, uint32_t count, + std::vector &block_vecs) const override; + + //! Get vector sparse feature data by chunk index and offset + const void *get_sparse_data(uint64_t offset, uint32_t len) const override; + + //! Get sparse data from id + const void *get_sparse_data(node_id_t id) const override; + + int get_sparse_data(uint64_t offset, uint32_t len, + IndexStorage::MemoryBlock &block) const override; + int get_sparse_data(node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get sparse data from vector + std::pair get_sparse_data_from_vector( + const void *vec) const override; + int get_sparse_data_from_vector(const void *vec, + IndexStorage::MemoryBlock &block, + int &sparse_length) const override; + + //! Get sparse vector feature data by key + int get_sparse_vector_by_key( + key_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override; + + //! Get sparse vector feature data by id + int get_sparse_vector_by_id(node_id_t id, uint32_t *sparse_count, + std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const override; + + //! Get the node id's neighbors on graph level + //! Note: the neighbors cannot be modified, using the following + //! method to get WritableNeighbors if want to + const Neighbors get_neighbors(level_t level, node_id_t id) const override; + + + //! Add vector and key to hnsw entity, and local id will be saved in id + int add_vector(level_t level, key_t key, const std::string &sparse_vec_buffer, + uint32_t sparse_count, node_id_t *id) override; + + //! Add vector and id to hnsw entity + int add_vector_with_id(level_t level, node_id_t id, + const std::string &sparse_vec, + uint32_t sparse_count) override; + + int update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) override; + + //! Replace node id in level's neighbors + int update_neighbors_dense( + level_t level, node_id_t id, + const std::vector> &neighbors); + + //! Replace node id in level's neighbors + int update_neighbors_sparse( + level_t level, node_id_t id, + const std::vector> &neighbors); + + //! Append neighbor_id to node id neighbors on level + //! Notice: the caller must be ensure the neighbors not full + void add_neighbor(level_t level, node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + + //! Dump index by dumper + int dump(const IndexDumper::Pointer &dumper) override; + + void update_ep_and_level(node_id_t ep, level_t level) override; + + public: + //! Constructor + HnswSparseStreamerEntity(IndexStreamer::Stats &stats); + + //! Destructor + ~HnswSparseStreamerEntity(); + + //! Init entity + int init(uint64_t max_index_size, size_t max_doc_cnt); + + //! Flush graph entity to disk + //! return 0 on success, or errCode in failure + int flush(uint64_t checkpoint); + + //! Open entity from storage + //! return 0 on success, or errCode in failure + int open(IndexStorage::Pointer stg, bool check_crc); + + //! Close entity + //! return 0 on success, or errCode in failure + int close(); + + //! Set meta information from entity + int set_index_meta(const IndexMeta &meta) const { + return IndexHelper::SerializeToStorage(meta, broker_->storage().get()); + } + + //! Get meta information from entity + int get_index_meta(IndexMeta *meta) const { + return IndexHelper::DeserializeFromStorage(broker_->storage().get(), meta); + } + + //! Set params: chunk size + inline void set_chunk_size(size_t val) { + chunk_size_ = val; + } + + //! Set params + inline void set_filter_same_key(bool val) { + filter_same_key_ = val; + } + + //! Set params + inline void set_get_vector(bool val) { + get_vector_enabled_ = val; + } + + //! Get vector local id by key + inline node_id_t get_id(key_t key) const { + keys_map_lock_->lock_shared(); + auto it = keys_map_->find(key); + keys_map_lock_->unlock_shared(); + return it == keys_map_->end() ? kInvalidNodeId : it->second; + } + + void print_key_map() { + std::cout << "key map begins" << std::endl; + + auto iter = keys_map_->begin(); + while (iter != keys_map_->end()) { + std::cout << "key: " << iter->first << ", id: " << iter->second + << std::endl; + ; + iter++; + } + + std::cout << "key map ends" << std::endl; + } + + //! Get neighbors size + inline size_t neighbors_size() const { + return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); + } + + //! Get upper neighbors size + inline size_t upper_neighbors_size() const { + return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); + } + + private: + union UpperNeighborIndexMeta { + struct { + uint32_t level : 4; + uint32_t index : 28; // index is composite type: chunk idx, and the + // N th neighbors in chunk, they two composite + // the 28 bits location + } bits; + uint32_t data; + }; + + template + using HashMap = google::dense_hash_map>; + template + using HashMapPointer = std::shared_ptr>; + + template + using HashSet = google::dense_hash_set>; + template + using HashSetPointer = std::shared_ptr>; + + //! upper neighbor index hashmap + using NIHashMap = HnswSparseIndexHashMap; + using NIHashMapPointer = std::shared_ptr; + + //! Private construct, only be called by clone method + HnswSparseStreamerEntity( + IndexStreamer::Stats &stats, const HNSWSparseHeader &hd, + size_t chunk_size, uint32_t node_index_mask_bits, + uint32_t upper_neighbor_mask_bits, bool filter_same_key, + bool get_vector_enabled, uint32_t sparse_chunk_size, + const NIHashMapPointer &upper_neighbor_index, + std::shared_ptr &keys_map_lock, + const HashMapPointer &keys_map, + std::vector &&node_chunks, + std::vector &&upper_neighbor_chunks, + std::vector &&sparse_node_chunks, + const SparseChunkBroker::Pointer &broker) + : stats_(stats), + chunk_size_(chunk_size), + node_index_mask_bits_(node_index_mask_bits), + node_cnt_per_chunk_(1UL << node_index_mask_bits_), + node_index_mask_(node_cnt_per_chunk_ - 1), + upper_neighbor_mask_bits_(upper_neighbor_mask_bits), + upper_neighbor_mask_((1U << upper_neighbor_mask_bits_) - 1), + filter_same_key_(filter_same_key), + get_vector_enabled_(get_vector_enabled), + sparse_chunk_size_(sparse_chunk_size), + upper_neighbor_index_(upper_neighbor_index), + keys_map_lock_(keys_map_lock), + keys_map_(keys_map), + node_chunks_(std::move(node_chunks)), + upper_neighbor_chunks_(std::move(upper_neighbor_chunks)), + sparse_node_chunks_(std::move(sparse_node_chunks)), + broker_(broker) { + *mutable_header() = hd; + + neighbor_size_ = neighbors_size(); + upper_neighbor_size_ = upper_neighbors_size(); + } + + //! Called only in searching procedure per context, so no need to lock + void sync_chunks(SparseChunkBroker::CHUNK_TYPE type, size_t idx, + std::vector *chunks) const { + if (ailego_likely(idx < chunks->size())) { + return; + } + for (size_t i = chunks->size(); i <= idx; ++i) { + auto chunk = broker_->get_chunk(type, i); + // the storage can ensure get chunk will success after the first get + ailego_assert_with(!!chunk, "get chunk failed"); + chunks->emplace_back(std::move(chunk)); + } + } + + //! return pair: chunk index + chunk offset + inline std::pair get_vector_chunk_loc( + node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size(); + + sync_chunks(SparseChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + //! return pair: chunk index + chunk offset + inline std::pair get_key_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + + offset += sparse_meta_size(); + + sync_chunks(SparseChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + //! return pair: chunk index + chunk offset + inline std::pair get_sparse_chunk_loc( + uint32_t chunk_index, uint32_t chunk_offset) const { + sync_chunks(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE, chunk_index, + &sparse_node_chunks_); + + return std::make_pair(chunk_index, chunk_offset); + } + + inline std::pair get_upper_neighbor_chunk_loc( + level_t level, node_id_t id) const { + auto it = upper_neighbor_index_->find(id); + ailego_assert_abort(it != upper_neighbor_index_->end(), + "Get upper neighbor header failed"); + auto meta = reinterpret_cast(&it->second); + uint32_t chunk_idx = (meta->bits.index) >> upper_neighbor_mask_bits_; + uint32_t offset = + (((meta->bits.index) & upper_neighbor_mask_) + level - 1) * + upper_neighbor_size_; + sync_chunks(SparseChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, chunk_idx, + &upper_neighbor_chunks_); + ailego_assert_abort(chunk_idx < upper_neighbor_chunks_.size(), + "invalid chunk idx"); + ailego_assert_abort(offset < upper_neighbor_chunks_[chunk_idx]->data_size(), + "invalid chunk offset"); + return std::make_pair(chunk_idx, offset); + } + + //! return pair: chunk + chunk offset + inline std::pair get_neighbor_chunk_loc( + level_t level, node_id_t id) const { + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + + offset += sparse_meta_size(); + + sync_chunks(SparseChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_abort(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_abort(offset < node_chunks_[chunk_idx]->data_size(), + "invalid chunk offset"); + return std::make_pair(node_chunks_[chunk_idx].get(), offset); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + return std::make_pair(upper_neighbor_chunks_[p.first].get(), p.second); + } + } + + //! Chunk hnsw index valid + int check_hnsw_index(const HNSWSparseHeader *hd) const; + + size_t get_total_upper_neighbors_size(level_t level) const { + return level * upper_neighbor_size_; + } + + //! Add upper neighbor header and reserve space for upper neighbor + int add_upper_neighbor(level_t level, node_id_t id) { + if (level == 0) { + return 0; + } + SparseChunk::Pointer chunk; + uint64_t chunk_offset = UINT64_MAX; + size_t neighbors_size = get_total_upper_neighbors_size(level); + uint64_t chunk_index = upper_neighbor_chunks_.size() - 1ULL; + if (chunk_index == UINT64_MAX || + (upper_neighbor_chunks_[chunk_index]->padding_size() < + neighbors_size)) { // no space left and need to alloc + chunk_index++; + if (ailego_unlikely(upper_neighbor_chunks_.capacity() == + upper_neighbor_chunks_.size())) { + LOG_ERROR("add upper neighbor failed for no memory quota"); + return IndexError_IndexFull; + } + auto p = + broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, + chunk_index, upper_neighbor_chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + chunk = p.second; + chunk_offset = 0UL; + upper_neighbor_chunks_.emplace_back(chunk); + } else { + chunk = upper_neighbor_chunks_[chunk_index]; + chunk_offset = chunk->data_size(); + } + ailego_assert_with((size_t)level < kMaxGraphLayers, "invalid level"); + ailego_assert_with(chunk_offset % upper_neighbor_size_ == 0, + "invalid offset"); + ailego_assert_with((chunk_offset / upper_neighbor_size_) < + (1U << upper_neighbor_mask_bits_), + "invalid offset"); + ailego_assert_with(chunk_index < (1U << (28 - upper_neighbor_mask_bits_)), + "invalid chunk index"); + UpperNeighborIndexMeta meta; + meta.bits.level = level; + meta.bits.index = (chunk_index << upper_neighbor_mask_bits_) | + (chunk_offset / upper_neighbor_size_); + chunk_offset += upper_neighbor_size_ * level; + if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) { + LOG_ERROR("HashMap insert value failed"); + return IndexError_Runtime; + } + + if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("SparseChunk resize to %zu failed", (size_t)chunk_offset); + return IndexError_Runtime; + } + + return 0; + } + + size_t estimate_doc_capacity() const { + return node_chunks_.capacity() * node_cnt_per_chunk_; + } + + int init_chunk_params(size_t max_index_size) { + sparse_chunk_size_ = AlignPageSize(chunk_size_); + + node_cnt_per_chunk_ = std::max(1, chunk_size_ / node_size()); + //! align node cnt per chunk to pow of 2 + node_index_mask_bits_ = std::ceil(std::log2(node_cnt_per_chunk_)); + node_cnt_per_chunk_ = 1UL << node_index_mask_bits_; + chunk_size_ = AlignPageSize(node_cnt_per_chunk_ * node_size()); + node_index_mask_ = node_cnt_per_chunk_ - 1; + + if (max_index_size == 0UL) { + max_index_size_ = chunk_size_ * kDefaultMaxChunkCnt; + } else { + max_index_size_ = max_index_size; + } + + //! To get a balanced upper neighbor chunk size. + //! If the upper chunk size is equal to node chunk size, it may waste + //! upper neighbor chunk space; if the upper neighbor chunk size is too + //! small, the will need large upper neighbor chunks index space. So to + //! get a balanced ratio be sqrt of the node/neighbor size ratio + float ratio = + std::sqrt(node_size() * scaling_factor() * 1.0f / upper_neighbor_size_); + upper_neighbor_chunk_size_ = + AlignPageSize(std::max(get_total_upper_neighbors_size(kMaxGraphLayers), + static_cast(chunk_size_ / ratio))); + upper_neighbor_mask_bits_ = + std::ceil(std::log2(upper_neighbor_chunk_size_ / upper_neighbor_size_)); + upper_neighbor_mask_ = (1 << upper_neighbor_mask_bits_) - 1; + + size_t max_node_chunk_cnt = std::ceil(max_index_size_ / chunk_size_); + size_t max_upper_chunk_cnt = std::ceil( + (max_node_chunk_cnt * node_cnt_per_chunk_ * 1.0f / scaling_factor()) / + (upper_neighbor_chunk_size_ / upper_neighbor_size_)); + max_upper_chunk_cnt = + max_upper_chunk_cnt + std::ceil(max_upper_chunk_cnt / scaling_factor()); + + //! reserve space to avoid memmove in chunks vector emplace chunk, so + //! as to lock-free in reading chunk + node_chunks_.reserve(max_node_chunk_cnt); + sparse_node_chunks_.reserve(max_node_chunk_cnt); + upper_neighbor_chunks_.reserve(max_upper_chunk_cnt); + + LOG_DEBUG( + "Settings: nodeSize=%zu chunkSize=%u upperNeighborSize=%u " + "upperNeighborChunkSize=%u " + "nodeCntPerChunk=%u maxChunkCnt=%zu maxNeighborChunkCnt=%zu " + "maxIndexSize=%zu ratio=%.3f", + node_size(), chunk_size_, upper_neighbor_size_, + upper_neighbor_chunk_size_, node_cnt_per_chunk_, max_node_chunk_cnt, + max_upper_chunk_cnt, max_index_size_, ratio); + + return 0; + } + + //! Init node chunk and neighbor chunks + int init_chunks(const SparseChunk::Pointer &header_chunk); + + int flush_header(void) { + if (!broker_->dirty()) { + // do not need to flush + return 0; + } + auto header_chunk = + broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_HEADER, + SparseChunkBroker::kDefaultChunkSeqId); + if (ailego_unlikely(!header_chunk)) { + LOG_ERROR("get header chunk failed"); + return IndexError_Runtime; + } + size_t size = header_chunk->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + + return 0; + } + + private: + HnswSparseStreamerEntity(const HnswSparseStreamerEntity &) = delete; + HnswSparseStreamerEntity &operator=(const HnswSparseStreamerEntity &) = + delete; + static constexpr uint64_t kUpperHashMemoryInflateRatio = 2.0f; + + private: + IndexStreamer::Stats &stats_; + HNSWSparseHeader header_{}; + std::mutex mutex_{}; + size_t max_index_size_{0UL}; + uint32_t chunk_size_{kDefaultChunkSize}; + uint32_t upper_neighbor_chunk_size_{kDefaultChunkSize}; + uint32_t node_index_mask_bits_{0U}; + uint32_t node_cnt_per_chunk_{0U}; + uint32_t node_index_mask_{0U}; + uint32_t neighbor_size_{0U}; + uint32_t upper_neighbor_size_{0U}; + //! UpperNeighborIndex.index composite chunkIdx and offset in chunk by the + //! following mask + uint32_t upper_neighbor_mask_bits_{0U}; + uint32_t upper_neighbor_mask_{0U}; + bool filter_same_key_{false}; + bool get_vector_enabled_{false}; + + uint32_t sparse_chunk_size_{kDefaultChunkSize}; + + NIHashMapPointer upper_neighbor_index_{}; + + mutable std::shared_ptr keys_map_lock_{}; + HashMapPointer keys_map_{}; + + //! the chunks will be changed in searcher, so need mutable + //! data chunk include: vector, key, level 0 neighbors + mutable std::vector node_chunks_{}; + //! upper neighbor chunk inlude: UpperNeighborHeader + (1~level) neighbors + mutable std::vector upper_neighbor_chunks_{}; + //! chunk that holds up sparse part + mutable std::vector sparse_node_chunks_{}; + + SparseChunkBroker::Pointer broker_{}; // chunk broker +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/ivf/CMakeLists.txt b/src/core/algorithm/ivf/CMakeLists.txt new file mode 100644 index 0000000..8e3872f --- /dev/null +++ b/src/core/algorithm/ivf/CMakeLists.txt @@ -0,0 +1,16 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(NOT APPLE) + set(CORE_KNN_IVF_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_ivf STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS zvec_ailego core_framework core_knn_cluster + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_IVF_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/ivf/ivf_builder.cc b/src/core/algorithm/ivf/ivf_builder.cc new file mode 100644 index 0000000..16b3cbf --- /dev/null +++ b/src/core/algorithm/ivf/ivf_builder.cc @@ -0,0 +1,829 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "ivf_builder.h" +#include +#include +#include "algorithm/cluster/cluster_params.h" +#include "ivf_dumper.h" + +namespace zvec { +namespace core { + +/*! IndexHolder support filtered by vector labels + */ +class LabelFilteredIndexHolder : public IndexHolder { + public: + /*! Index Holder Iterator + */ + class Iterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(const IVFBuilder::RandomAccessIndexHolder::Pointer &holder, + const std::vector *elems) + : holder_(holder), elems_(elems) {} + + //! Destructor + ~Iterator(void) override {} + + //! Retrieve pointer of data + const void *data(void) const override { + return holder_->element((*elems_)[index_]); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return index_ < elems_->size(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return (*elems_)[index_]; + } + + //! Next iterator + void next(void) override { + ++index_; + } + + private: + //! Members + const IVFBuilder::RandomAccessIndexHolder::Pointer holder_{nullptr}; + const std::vector *elems_{nullptr}; + size_t index_{0}; + }; + + //! Constructor + LabelFilteredIndexHolder( + const IVFBuilder::RandomAccessIndexHolder::Pointer &holder, + const std::vector &items) + : holder_(holder), elems_(&items) {} + + //! Retrieve count of elements in holder (-1 indicates unknown) + size_t count(void) const override { + return elems_->size(); + } + + //! Retrieve dimension + size_t dimension(void) const override { + return holder_->dimension(); + } + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return holder_->data_type(); + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return holder_->element_size(); + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return true; + } + + //! Create a new iterator + IndexHolder::Iterator::Pointer create_iterator(void) override { + return IndexHolder::Iterator::Pointer( + new LabelFilteredIndexHolder::Iterator(holder_, elems_)); + } + + private: + //! Members + const IVFBuilder::RandomAccessIndexHolder::Pointer holder_{}; + const std::vector *elems_{}; +}; + +IVFBuilder::IVFBuilder() {} + +IVFBuilder::~IVFBuilder() { + this->cleanup(); +} + +int IVFBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { + LOG_INFO("Begin IVFBuilder::init!"); + + if (state_ != INIT) { + LOG_ERROR("IVFBuilder state wrong. state=%d", state_); + return IndexError_Logic; + } + + meta_ = meta; + converted_meta_ = meta; + quantized_meta_ = meta; + // Clear the converter/reformer params for external transforms + converted_meta_.set_reformer(std::string(), 0, ailego::Params()); + converted_meta_.set_converter(std::string(), 0, ailego::Params()); + quantized_meta_.set_reformer(std::string(), 0, ailego::Params()); + quantized_meta_.set_converter(std::string(), 0, ailego::Params()); + params_ = params; + + if (!IndexFactory::HasMetric(meta_.metric_name())) { + LOG_ERROR("Metric %s not exist", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + + int ret = parse_centroids_num(params); + ivf_check_with_msg(ret, "Failed to parse centroids, ret=%d", ret); + + ret = parse_clustering_params(params); + ivf_check_with_msg(ret, "Failed to parse clustering params, ret=%d", ret); + + ret = parse_general_params(params); + ivf_check_with_msg(ret, "Failed to parse general params, ret=%d", ret); + + LOG_INFO("End IVFBuilder::init!"); + + LOG_DEBUG( + "Converter=%s Quantizer=%s Optimizer=%s " + "OptimizerQuantizer=%s QuantizeByCentroid=%u StoreFeatures=%u " + "ClusterClass=%s TrainSamplesCount=%u TrainSampleRatio=%f " + "BlockVectorCount=%u", + params.get_as_string(PARAM_IVF_BUILDER_CONVERTER_CLASS).c_str(), + params.get_as_string(PARAM_IVF_BUILDER_QUANTIZER_CLASS).c_str(), + params.get_as_string(PARAM_IVF_BUILDER_OPTIMIZER_CLASS).c_str(), + params.get_as_string(PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_CLASS).c_str(), + params.get_as_bool(PARAM_IVF_BUILDER_QUANTIZE_BY_CENTROID), + params.get_as_bool(PARAM_IVF_BUILDER_STORE_ORIGINAL_FEATURES), + params.get_as_string(PARAM_IVF_BUILDER_CLUSTER_CLASS).c_str(), + params.get_as_uint32(PARAM_IVF_BUILDER_TRAIN_SAMPLE_COUNT), + params.get_as_float(PARAM_IVF_BUILDER_TRAIN_SAMPLE_RATIO), + block_vector_count_); + + state_ = INITED; + return 0; +} + +int IVFBuilder::cleanup(void) { + LOG_INFO("Begin IVFBuilder::cleanup"); + + state_ = INIT; + stats_.clear_attributes(); + stats_.set_built_costtime(0u); + stats_.set_built_count(0u); + stats_.set_discarded_count(0u); + stats_.set_dumped_costtime(0u); + stats_.set_dumped_count(0u); + stats_.set_trained_costtime(0u); + stats_.set_trained_count(0u); + + centroid_num_vec_.clear(); + cluster_class_.clear(); + converter_class_.clear(); + cluster_params_.clear(); + + labels_.clear(); + centroid_index_.reset(); + holder_.reset(); + converted_meta_ = meta_; + converter_.reset(); + quantized_meta_ = meta_; + quantizers_.clear(); + + error_ = false; + err_code_ = 0; + + thread_count_ = 0; + sample_count_ = 0; + cluster_auto_tuning_ = false; + store_original_features_ = false; + quantize_by_centroid_ = false; + + LOG_INFO("End IVFBuilder::cleanup"); + + return 0; +} + +int IVFBuilder::train(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) { + LOG_INFO("Begin IVFBuilder::train with holder"); + if (state_ != INITED) { + LOG_ERROR("IVFBuilder train failed, wrong state=%d", state_); + return IndexError_Runtime; + } + + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + ailego::ElapsedTime timer; + if (!holder || holder->count() == 0) { + LOG_ERROR("Input holder is nullptr or empty while train index"); + return IndexError_InvalidArgument; + } + if (!holder->is_matched(meta_)) { + LOG_ERROR("Input holder doesn't match index meta while train index"); + return IndexError_Mismatch; + } + + if (converter_) { + int ret = IndexConverter::TrainAndTransform(converter_, std::move(holder)); + ivf_check_with_msg(ret, "Failed to train or transform by converter %s", + converter_->name().c_str()); + converted_meta_ = converter_->meta(); + holder = converter_->result(); + } + + ailego::Params train_params; + int ret = prepare_trainer_params(train_params); + ivf_check_with_msg(ret, "Failed to prepare trainer params, ret=%d", ret); + + IndexTrainer::Pointer trainer = + IndexFactory::CreateTrainer("StratifiedClusterTrainer"); + ivf_assert_with_msg(trainer, IndexError_NoExist, "Failed to create trainer"); + + ret = trainer->init(converted_meta_, train_params); + ivf_check_with_msg(ret, "Trainer init failed with ret %d", ret); + + ret = trainer->train(std::move(threads), std::move(holder)); + ivf_check_with_msg(ret, "Trainer train failed with ret %d", ret); + + ret = this->train(trainer); + ivf_check_error_code(ret); + + stats_.set_trained_costtime(timer.milli_seconds()); + + LOG_INFO("End IVFBuilder::train with holder"); + + state_ = TRAINED; + return 0; +} + +int IVFBuilder::train(const IndexTrainer::Pointer &trainer) { + LOG_DEBUG("Begin IVFBuilder::train by trainer"); + ailego::ElapsedTime timer; + + if (state_ != INITED) { + LOG_ERROR("IVFBuilder train failed, wrong state=%d", state_); + return IndexError_Runtime; + } + + if (!trainer) { + LOG_ERROR("Input trainer is nullptr while train index"); + return IndexError_InvalidArgument; + } + + IndexCluster::CentroidList centroid_list; + IndexBundle::Pointer boundle = trainer->indexes(); + int ret = IndexCluster::Deserialize(trainer->meta(), boundle, ¢roid_list); + ivf_check_with_msg(ret, "Failed to deserialize index"); + + const IndexMeta &meta = trainer->meta(); + if (meta.data_type() != converted_meta_.data_type() || + meta.metric_name().compare(converted_meta_.metric_name()) != 0 || + meta.element_size() != converted_meta_.element_size()) { + if (meta.converter_name() != converter_class_) { + LOG_ERROR("Input trainer doesn't match index meta while train index"); + return IndexError_Mismatch; + } + //! Create converter from trainer params + LOG_INFO("Train IVFBuilder by trainer with converter"); + converter_ = CreateAndInitConverter(meta_, meta.converter_name(), + meta.converter_params()); + ivf_assert(converter_, IndexError_Runtime); + converted_meta_ = meta; + } + + centroid_index_ = std::make_shared(); + if (!centroid_index_) { + return IndexError_NoMemory; + } + ret = centroid_index_->init(converted_meta_, params_); + ivf_check_error_code(ret); + + ret = centroid_index_->build(centroid_list); + ivf_check_with_msg(ret, "Failed to build centroid index"); + + if (params_.has(PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_CLASS)) { + //! Quantize the centroids for searcher + searcher_centroid_index_ = std::make_shared(); + if (!searcher_centroid_index_) { + return IndexError_NoMemory; + } + ailego::Params params; + params_.get(PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_PARAMS, ¶ms); + searcher_centroid_index_->set_quantizer( + params_.get_as_string(PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_CLASS), + params); + ret = searcher_centroid_index_->init(converted_meta_, params_); + ivf_check_error_code(ret); + + ret = searcher_centroid_index_->build(centroid_list); + ivf_check_with_msg(ret, "Failed to build centroid index"); + } + + stats_.set_trained_costtime(timer.milli_seconds()); + + LOG_DEBUG("End IVFBuilder::train by trainer"); + + state_ = TRAINED; + return 0; +} + +int IVFBuilder::build(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) { + LOG_INFO("Begin IVFBuilder::build!"); + + if (state_ != TRAINED) { + LOG_ERROR("Train the index first before build"); + return IndexError_Runtime; + } + + ailego::ElapsedTime timer; + if (!holder || holder->count() == 0) { + LOG_ERROR("Input holder is nullptr or empty while building index"); + return IndexError_InvalidArgument; + } + + if (!holder->is_matched(meta_)) { + LOG_ERROR("Input holder doesn't match index meta while building index"); + return IndexError_Mismatch; + } + if (!threads) { + threads = std::make_shared(thread_count_, false); + if (!threads) { + return IndexError_NoMemory; + } + } + + holder_ = std::make_shared(meta_); + if (!holder_) { + return IndexError_NoMemory; + } + if (holder->count() > 0) { + holder_->reserve(holder->count()); + } + for (auto iter = holder->create_iterator(); iter && iter->is_valid(); + iter->next()) { + holder_->emplace(iter->key(), iter->data()); + } + + // Holder is not needed, cleanup it. + holder.reset(); + + IndexHolder::Pointer converted_holder = holder_; + if (converter_) { + int ret = converter_->transform(holder_); + ivf_check_with_msg(ret, "Failed to transform by converter %s", + converter_->name().c_str()); + converted_holder = converter_->result(); + } + + labels_.resize(centroid_index_->centroids_count()); + int ret = this->build_label_index(threads.get(), converted_holder); + ivf_check_with_msg(ret, "Failed to build index for %s", + IndexError::What(ret)); + + ret = this->prepare_quantizer(threads.get()); + ivf_check_error_code(ret); + + stats_.set_built_costtime(timer.milli_seconds()); + + LOG_INFO("End IVFBuilder::build"); + + state_ = BUILT; + return 0; +} + +int IVFBuilder::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("Begin IVFBuilder::dump"); + + if (state_ != BUILT) { + LOG_ERROR("Build the index before dump QC Index"); + return IndexError_Runtime; + } + + ailego::ElapsedTime timer; + int ret = this->dump_index(dumper); + ivf_check_with_msg(ret, "Failed to dump index with ret=%d", ret); + + // the fitting function for the follow points: 1000000(0.02) 10000000(0.01) + // 50000000(0.005) 100000000(0.001) + float scan_ratio = -0.004 * std::log(holder_->count()) + 0.0751; + scan_ratio = std::max(scan_ratio, 0.0001f); + + // Set Searcher Params + ailego::Params params; + params.set(PARAM_IVF_SEARCHER_SCAN_RATIO, scan_ratio); + meta_.set_searcher("IVFSearcher", 0, std::move(params)); + meta_.set_builder("IVFBuilder", 0, std::move(params_)); + + ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + + stats_.set_discarded_count(stats_.built_count() - stats_.dumped_count()); + stats_.set_dumped_costtime(timer.milli_seconds()); + + LOG_INFO("End IVFBuilder::dump"); + + return 0; +} + +int IVFBuilder::CheckAndUpdateMajorOrder(IndexMeta &meta) { + const std::string &metric_name = meta.metric_name(); + auto metric = IndexFactory::CreateMetric(metric_name); + if (!metric) { + LOG_ERROR("CreateMetric %s failed", metric_name.c_str()); + return IndexError_InvalidArgument; + } + int ret = metric->init(meta, meta.metric_params()); + ivf_check_with_msg(ret, "IndexMetric %s init failed", metric_name.c_str()); + + bool support_column_major = true; + for (size_t m = 32; m != 0; m /= 2) { + for (size_t n = m; n != 0; n /= 2) { + if (metric->distance_matrix(m, n) == nullptr) { + support_column_major = false; + break; + } + } + if (!support_column_major) { + break; + } + } + support_column_major &= + meta.element_size() % IndexMeta::AlignSizeof(meta.data_type()) == 0; + + if (meta.major_order() == IndexMeta::MO_UNDEFINED) { + if (support_column_major && meta.dimension() <= 512) { + meta.set_major_order(IndexMeta::MO_COLUMN); + } else { + meta.set_major_order(IndexMeta::MO_ROW); + } + } else { + if (!support_column_major && meta.major_order() == IndexMeta::MO_COLUMN) { + LOG_WARN( + "Index Metric %s Unsupported " + "Column Major Order", + metric_name.c_str()); + return IndexError_Unsupported; + } + } + + if (block_vector_count_ * quantized_meta_.element_size() % 32 != 0) { + LOG_ERROR( + "block_vector_count * quantized_element_size not align with 32 bytes."); + return IndexError_InvalidArgument; + } + + return 0; +} + +int IVFBuilder::parse_centroids_num(const ailego::Params ¶ms) { + std::string centroids_num = + params.get_as_string(PARAM_IVF_BUILDER_CENTROID_COUNT); + if (centroids_num.empty()) { + LOG_ERROR("Param %s is required", PARAM_IVF_BUILDER_CENTROID_COUNT.c_str()); + return IndexError_InvalidArgument; + } + + std::vector centroid_str_vec; + ailego::StringHelper::Split(centroids_num, CENTROID_SEPERATOR, + ¢roid_str_vec); + size_t level_cnt = centroid_str_vec.size(); + if ((level_cnt <= 0) || (level_cnt > 2)) { + LOG_ERROR("Centroids level count must be [1,2]"); + return IndexError_InvalidArgument; + } + + for (size_t idx = 0; idx < level_cnt; ++idx) { + uint32_t centroid_cnt = 0; + if (!ailego::StringHelper::ToUint32(centroid_str_vec[idx], ¢roid_cnt)) { + LOG_ERROR("Invalid centroids count %s", centroid_str_vec[idx].c_str()); + return IndexError_InvalidArgument; + } + centroid_num_vec_.push_back(centroid_cnt); + } + + return 0; +} + +int IVFBuilder::parse_clustering_params(const ailego::Params ¶ms) { + params.get(PARAM_IVF_BUILDER_CLUSTER_AUTO_TUNING, &cluster_auto_tuning_); + + cluster_class_ = params.get_as_string(PARAM_IVF_BUILDER_CLUSTER_CLASS); + if (cluster_class_.empty()) { + // OptKmeansCluster does not support custom metric + cluster_class_ = meta_.metric_name() == kMipsMetricName + ? "KmeansCluster" + : "OptKmeansCluster"; + LOG_INFO("Using [%s] as default cluster class", cluster_class_.c_str()); + } + for (size_t i = 1; i <= centroid_num_vec_.size(); ++i) { + std::string level_params_key = + PARAM_IVF_BUILDER_CLUSTER_PARAMS_IN_LEVEL_PREFIX + std::to_string(i); + ailego::Params level_params; + params.get(level_params_key, &level_params); + cluster_params_.push_back(level_params); + } + + return 0; +} + +int IVFBuilder::parse_general_params(const ailego::Params ¶ms) { + thread_count_ = params.get_as_uint32(PARAM_IVF_BUILDER_THREAD_COUNT); + sample_count_ = params.get_as_uint32(PARAM_IVF_BUILDER_TRAIN_SAMPLE_COUNT); + sample_ratio_ = params.get_as_float(PARAM_IVF_BUILDER_TRAIN_SAMPLE_RATIO); + + params.get(PARAM_IVF_BUILDER_QUANTIZE_BY_CENTROID, &quantize_by_centroid_); + params.get(PARAM_IVF_BUILDER_STORE_ORIGINAL_FEATURES, + &store_original_features_); + + //! Prepare Converter for training + if (meta_.metric_name() == kIPMetricName) { + converter_class_ = kMipsConverterName; + } + params.get(PARAM_IVF_BUILDER_CONVERTER_CLASS, &converter_class_); + if (!converter_class_.empty()) { + ailego::Params converter_params; + params_.get(PARAM_IVF_BUILDER_CONVERTER_PARAMS, &converter_params); + converter_ = + CreateAndInitConverter(meta_, converter_class_, converter_params); + ivf_assert(converter_, IndexError_NoExist); + } + + params_.get(PARAM_IVF_BUILDER_BLOCK_VECTOR_COUNT, &block_vector_count_); + if (block_vector_count_ == 0) { + block_vector_count_ = kDefaultBlockCount; + } + if (block_vector_count_ > kDefaultBlockCount || + block_vector_count_ & (block_vector_count_ - 1)) { + LOG_ERROR("block_vector_count only can be [1|2|4|8|16|32]."); + return IndexError_InvalidArgument; + } + if (block_vector_count_ * meta_.element_size() % 32 != 0) { + LOG_ERROR("block_vector_count * element_size not align with 32 bytes."); + return IndexError_InvalidArgument; + } + return 0; +} + +int IVFBuilder::prepare_trainer_params(ailego::Params ¶ms) { + params.set(STRATIFIED_TRAINER_SAMPLE_COUNT, sample_count_); + params.set(STRATIFIED_TRAINER_SAMPLE_RATIO, sample_ratio_); + params.set(STRATIFIED_TRAINER_THREAD_COUNT, thread_count_); + params.set(STRATIFIED_TRAINER_AUTOAUNE, cluster_auto_tuning_); + if (centroid_num_vec_.empty()) { + LOG_ERROR("Centroids no specified."); + return IndexError_InvalidArgument; + } + std::string cluster_count = std::to_string(centroid_num_vec_[0]); + if (centroid_num_vec_.size() > 1) { + cluster_count += + (CENTROID_SEPERATOR + std::to_string(centroid_num_vec_[1])); + } + params.set(STRATIFIED_TRAINER_CLUSTER_COUNT, cluster_count); + + for (size_t i = 1; i <= cluster_params_.size(); ++i) { + std::string level_params_key = + STRATIFIED_TRAINER_PARAMS_IN_LEVEL_PREFIX + std::to_string(i); + params.set(level_params_key, cluster_params_[i - 1]); + } + params.set(STRATIFIED_TRAINER_CLASS_NAME, cluster_class_); + + return 0; +} + +int IVFBuilder::build_label_index(IndexThreads *threads, + const IndexHolder::Pointer &holder) { + auto iter = holder->create_iterator(); + if (!iter) { + LOG_ERROR("Create iterator for holder failed"); + return IndexError_Runtime; + } + + auto task_group = threads->make_group(); + if (!task_group) { + LOG_ERROR("Failed to create task group"); + return IndexError_Runtime; + } + + size_t id = 0UL; + AILEGO_DEFER([&]() { + task_group->wait_finish(); + stats_.set_built_count(id); + LOG_INFO("Finished building, total=%zu", id); + }); + + size_t elem_size = holder->element_size(); + std::shared_ptr vectors = std::make_shared(); + ivf_assert(vectors, IndexError_NoMemory); + for (; iter && iter->is_valid(); iter->next()) { + ivf_assert(!error_, err_code_); + vectors->emplace_back(iter->data(), elem_size, id); + id++; + if (vectors->size() == kBatchSize || id == holder_->count()) { + auto task = ailego::Closure ::New(const_cast(this), + &IVFBuilder::label, vectors); + task_group->submit(std::move(task)); + vectors = std::make_shared(); + ivf_assert(vectors, IndexError_NoMemory); + vectors->reserve(kBatchSize); + } + if (!(id & 0xFFFFF)) { + LOG_INFO("Current built count:%zu", id); + } + } + ailego_assert_with(vectors->size() == 0, "invalid size"); + + return err_code_; +} + +int IVFBuilder::dump_index(const IndexDumper::Pointer &dumper) { + int ret = CheckAndUpdateMajorOrder(quantized_meta_); + ivf_check_error_code(ret); + + IVFDumper::Pointer ivf_dumper = std::make_shared( + quantized_meta_, dumper, centroid_index_->centroids_count(), + block_vector_count_); + if (!ivf_dumper) { + LOG_ERROR("Alloc IVFDumper failed"); + return IndexError_NoMemory; + } + + //! Dump inverted vectors + std::vector dumped_ids; + std::function record_dumped_id = [&](uint32_t) {}; + if (store_original_features_) { + dumped_ids.reserve(holder_->count()); + record_dumped_id = [&](uint32_t id) { dumped_ids.emplace_back(id); }; + } + if (quantizers_.size() == 0) { + //! No quantizer for inverted vectors + for (size_t i = 0; i < centroid_index_->centroids_count(); ++i) { + ailego_assert_with(i < labels_.size(), "Index Overflow"); + for (size_t j = 0; j < labels_[i].size(); ++j) { + auto id = labels_[i][j]; + record_dumped_id(id); + ret = ivf_dumper->dump_inverted_vector(i, holder_->key(id), + holder_->element(id)); + ivf_check_error_code(ret); + } + } + } else { + for (size_t i = 0; i < centroid_index_->centroids_count(); ++i) { + ailego_assert_with(i < labels_.size(), "Index Overflow"); + auto holder = + std::make_shared(holder_, labels_[i]); + if (!holder) { + return IndexError_NoMemory; + } + auto quantizer = quantize_by_centroid_ ? quantizers_[i] : quantizers_[0]; + ret = quantizer->transform(holder); + ivf_check_error_code(ret); + + auto iter = quantizer->result()->create_iterator(); + for (; iter->is_valid(); iter->next()) { + uint32_t id = iter->key(); + record_dumped_id(id); + ret = + ivf_dumper->dump_inverted_vector(i, holder_->key(id), iter->data()); + ivf_check_error_code(ret); + } + } + } + + ret = ivf_dumper->dump_inverted_vector_finished(); + ivf_check_error_code(ret); + + ret = ivf_dumper->dump_quantizer_params(quantizers_); + ivf_check_error_code(ret); + + auto centroid_index = + searcher_centroid_index_ ? searcher_centroid_index_ : centroid_index_; + ret = ivf_dumper->dump_centroid_index(centroid_index->data(), + centroid_index->size()); + ivf_check_with_msg(ret, "Failed to dump CentroidIndex"); + + if (store_original_features_) { + for (size_t i = 0; i < dumped_ids.size(); ++i) { + ret = ivf_dumper->dump_original_vector(holder_->element(dumped_ids[i]), + holder_->element_size()); + ivf_check_error_code(ret); + } + } + + stats_.set_dumped_count(stats_.dumped_count() + ivf_dumper->dumped_count()); + + return 0; +} + +int IVFBuilder::prepare_quantizer(IndexThreads *threads) { + std::string quantizer_name; + params_.get(PARAM_IVF_BUILDER_QUANTIZER_CLASS, &quantizer_name); + if (quantizer_name.empty()) { + return 0; + } + + //! Prepare Quantizers for inverted index + ailego::Params quantizer_params; + params_.get(PARAM_IVF_BUILDER_QUANTIZER_PARAMS, &quantizer_params); + if (((quantizer_name != kInt8QuantizerName && + quantizer_name != kInt4QuantizerName) || + meta_.metric_name() != kIPMetricName) && + quantize_by_centroid_) { + LOG_WARN("%s is supported in InnerProduct only", + PARAM_IVF_BUILDER_QUANTIZE_BY_CENTROID.c_str()); + quantize_by_centroid_ = false; + } + if (quantizer_name == kInt4QuantizerName && meta_.dimension() & 0x1) { + LOG_ERROR("Unsupport quantizer=%s for dim=%u", kInt4QuantizerName, + meta_.dimension()); + return IndexError_Unsupported; + } + + int ret = 0; + auto create_and_init_quantizer = [&]() { + auto quantizer = IndexFactory::CreateConverter(quantizer_name); + if (!quantizer) { + LOG_ERROR("Failed to create converter %s", quantizer_name.c_str()); + ret = IndexError_NoExist; + return IndexConverter::Pointer(); + } + ret = quantizer->init(meta_, quantizer_params); + if (ret != 0) { + LOG_ERROR("Failed to initialize converter %s for %s", + quantizer_name.c_str(), IndexError::What(ret)); + return IndexConverter::Pointer(); + } + return quantizer; + }; + for (size_t i = 0; i < centroid_index_->centroids_count(); ++i) { + quantizers_.emplace_back(create_and_init_quantizer()); + ivf_check_error_code(ret); + if (!quantize_by_centroid_) { + break; + } + } + + //! Train the quantizers + auto train_data = [&](size_t i) { + IndexHolder::Pointer holder = holder_; + size_t idx = 0; + if (quantize_by_centroid_) { + holder = std::make_shared(holder_, labels_[i]); + if (!holder && !error_.exchange(true)) { + err_code_ = IndexError_NoMemory; + return; + } + idx = i; + } + if (holder->count() == 0) { + return; + } + ret = quantizers_[idx]->train(holder); + if (ret != 0) { + LOG_ERROR("Failed to train converter %s for %s", quantizer_name.c_str(), + IndexError::What(ret)); + if (!error_.exchange(true)) { + err_code_ = IndexError_Runtime; + } + } + }; + + auto task_group = threads->make_group(); + if (!task_group) { + LOG_ERROR("Failed to create task group"); + return IndexError_Runtime; + } + + for (size_t i = 0; i < quantizers_.size(); ++i) { + if (error_) { + task_group->wait_finish(); + return err_code_; + } + task_group->submit(ailego::Closure ::New(train_data, i)); + } + + task_group->wait_finish(); + if (quantizers_.size() > 0) { + quantized_meta_ = quantizers_[0]->meta(); + } + + return 0; +} + +INDEX_FACTORY_REGISTER_BUILDER(IVFBuilder); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_builder.h b/src/core/algorithm/ivf/ivf_builder.h new file mode 100644 index 0000000..5140ff2 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_builder.h @@ -0,0 +1,315 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "ivf_centroid_index.h" + +namespace zvec { +namespace core { + +/*! IVF Builder + */ +class IVFBuilder : public IndexBuilder { + public: + //! Constructor + IVFBuilder(); + + //! Destructor + ~IVFBuilder() override; + + //! Disable them + IVFBuilder(const IVFBuilder &) = delete; + IVFBuilder &operator=(const IVFBuilder &) = delete; + + public: + //! Initialize the builder + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup the builder + int cleanup(void) override; + + //! Train the data + int train(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) override; + + //! Train the data + int train(const IndexTrainer::Pointer &trainer) override; + + //! Build the index + int build(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) override; + + //! Dump index into file system + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + IVFCentroidIndex::Pointer centroid_index() const { + return centroid_index_; + } + + public: + /*! Random Access Index Holder + */ + class RandomAccessIndexHolder : public IndexHolder { + public: + //! Index Holder Iterator Pointer + typedef std::shared_ptr Pointer; + + /*! Random Access Index Holder Iterator + */ + class Iterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(RandomAccessIndexHolder *owner) : holder_(owner) {} + + //! Destructor + ~Iterator(void) override {} + + //! Retrieve pointer of data + const void *data(void) const override { + return holder_->element(id_); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return id_ < holder_->count(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return holder_->key(id_); + } + + //! Next iterator + void next(void) override { + ++id_; + } + + private: + //! Members + RandomAccessIndexHolder *holder_{nullptr}; + uint32_t id_{0}; + }; + + //! Constructor + RandomAccessIndexHolder(const IndexMeta &meta) + : features_(std::make_shared(meta)) {} + + //! Retrieve count of elements in holder (-1 indicates unknown) + size_t count(void) const override { + return features_->count(); + } + + //! Retrieve dimension + size_t dimension(void) const override { + return features_->dimension(); + } + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return features_->data_type(); + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return features_->element_size(); + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return true; + } + + //! Create a new iterator + IndexHolder::Iterator::Pointer create_iterator(void) override { + return IndexHolder::Iterator::Pointer( + new RandomAccessIndexHolder::Iterator(this)); + } + + void reserve(size_t elems) { + features_->reserve(elems); + keys_.reserve(elems); + } + + //! Append an element into holder + void emplace(uint64_t pkey, const void *vec) { + features_->emplace(vec); + keys_.emplace_back(pkey); + } + + //! Retrieve feature via local id + const void *element(size_t id) const { + return features_->element(id); + } + + //! Retrieve key via local id + uint64_t key(size_t id) const { + ailego_assert_with(id < keys_.size(), "Index Overflow"); + return keys_[id]; + } + + private: + //! Disable them + RandomAccessIndexHolder(void) = delete; + + //! Members + CompactIndexFeatures::Pointer features_{}; + std::vector keys_{}; + }; + + private: + /*! Wrapper of feature + */ + class Vector { + public: + typedef std::shared_ptr Pointer; + + Vector(const void *vec, size_t len, uint32_t idx) + : vec_(reinterpret_cast(vec), len), id_{idx} {} + + const void *data() const { + return vec_.data(); + } + + size_t size() const { + return vec_.size(); + } + + uint32_t id(void) const { + return id_; + } + + private: + std::string vec_{}; + uint32_t id_{0u}; + }; + + using VectorList = std::vector; + + //! Check MajorOrder in meta, and update the major order if needed + int CheckAndUpdateMajorOrder(IndexMeta &meta); + + //! Parse params + int parse_centroids_num(const ailego::Params ¶ms); + int parse_clustering_params(const ailego::Params ¶ms); + int parse_general_params(const ailego::Params ¶ms); + + //! Prepare params for trainer + int prepare_trainer_params(ailego::Params ¶ms); + + //! Build the index + int build_label_index(IndexThreads *threads, + const IndexHolder::Pointer &holder); + + //! Dump the index to dumper + int dump_index(const IndexDumper::Pointer &dumper); + + //! Prepare the quantizer for inverted index + int prepare_quantizer(IndexThreads *threads); + + //! Quantize the centrods list + int quantize_centroids(); + + //! Create converter and init with params + static IndexConverter::Pointer CreateAndInitConverter( + const IndexMeta &meta, const std::string &name, + const ailego::Params ¶ms) { + auto converter = IndexFactory::CreateConverter(name); + if (!converter) { + LOG_ERROR("Failed to create converter %s", name.c_str()); + return IndexConverter::Pointer(); + } + int ret = converter->init(meta, params); + if (ret != 0) { + LOG_ERROR("Failed to initialize converter %s for %s", name.c_str(), + IndexError::What(ret)); + return IndexConverter::Pointer(); + } + return converter; + } + + //! Select the nearest centroid id for the vector + void label(const std::shared_ptr &vecs) { + for (size_t i = 0; i < vecs->size(); ++i) { + auto &vec = (*vecs)[i]; + + uint32_t centroid_idx = + centroid_index_->search_nearest_centroid(vec.data(), vec.size()); + if (centroid_idx == IVFCentroidIndex::kInvalidID) { + LOG_ERROR("Failed to search nearest centroid in CentroidIndex"); + if (!error_.exchange(true)) { + err_code_ = IndexError_Runtime; + } + return; + } + ailego_assert_with(centroid_idx < labels_.size(), "Index Overflow"); + mutex_.lock(); + labels_[centroid_idx].emplace_back(vec.id()); + mutex_.unlock(); + } + } + + + private: + //! Constants + static constexpr size_t kThreadPoolQueueSize = 300u; + static constexpr size_t kBatchSize = 10u; + static constexpr size_t kDefaultBlockCount = 32u; + + enum BuilderState { INIT = 0, INITED = 1, TRAINED = 2, BUILT = 3 }; + + //! Members + BuilderState state_{INIT}; + Stats stats_{}; + ailego::Params params_{}; + IndexMeta meta_{}; + + std::vector centroid_num_vec_{}; + std::string cluster_class_{}; + std::string converter_class_{}; + std::vector cluster_params_{}; + + std::vector> labels_{}; + std::mutex mutex_{}; + IVFCentroidIndex::Pointer centroid_index_{}; + IVFCentroidIndex::Pointer searcher_centroid_index_{}; + RandomAccessIndexHolder::Pointer holder_{}; + IndexMeta converted_meta_{}; + IndexConverter::Pointer converter_{}; + IndexMeta quantized_meta_{}; + std::vector quantizers_{}; + + std::atomic_bool error_{false}; + int err_code_{0}; + + uint32_t thread_count_{0}; + uint32_t sample_count_{0}; + float sample_ratio_{0.0}; + uint32_t block_vector_count_{kDefaultBlockCount}; + bool cluster_auto_tuning_{false}; + bool store_original_features_{false}; + bool quantize_by_centroid_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_centroid_index.cc b/src/core/algorithm/ivf/ivf_centroid_index.cc new file mode 100644 index 0000000..daa1f66 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_centroid_index.cc @@ -0,0 +1,617 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "ivf_centroid_index.h" +#include +#include +#include "metric/metric_params.h" + +namespace zvec { +namespace core { + +/*! Fake Trainer to supply centroids in bundle + */ +class FakeClusterTrainer : public IndexTrainer { + public: + //! Constructor + FakeClusterTrainer(const IndexMeta &imeta, const IndexBundle::Pointer &bundle) + : meta_(imeta), bundle_(bundle) {} + + //! Destructor + ~FakeClusterTrainer(void) override {} + + protected: + //! Initialize Trainer + int init(const IndexMeta &, const ailego::Params &) override { + return 0; + } + + //! Cleanup Trainer + int cleanup(void) override { + return 0; + } + + //! Train the data + int train(IndexHolder::Pointer) override { + return 0; + } + + //! Train the data + int train(IndexThreads::Pointer, IndexHolder::Pointer) override { + return 0; + } + + //! Load index from file path or dir + int load(IndexStorage::Pointer) override { + return 0; + } + + //! Dump index into file path or dir + int dump(const IndexDumper::Pointer &) override { + return 0; + } + + //! Retrieve Index Meta + const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve statistics + const IndexTrainer::Stats &stats(void) const override { + return stats_; + } + + //! Retrieve the output indexes + IndexBundle::Pointer indexes(void) const override { + return bundle_; + } + + private: + //! Members + IndexMeta meta_{}; + Stats stats_{}; + IndexBundle::Pointer bundle_{}; +}; + +/*! Int8QuantizerReformer for InnerProduct Measure + */ +class Int8QuantizerReformer4IP : public IndexReformer { + public: + //! Initialize Reformer + int init(const ailego::Params &) override { + return 0; + } + + //! Cleanup Reformer + int cleanup(void) override { + return 0; + } + + //! Load index from container + int load(IndexStorage::Pointer) override { + return 0; + } + + //! Unload index + int unload(void) override { + return 0; + } + + //! Transform query + int transform(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const override { +#if 0 + size_t dim = qmeta.dimension(); + out->resize(IndexMeta::ElementSizeof( + IndexMeta::DataType::DT_INT8, dim)); + ometa->set_meta(IndexMeta::DataType::DT_INT8, dim); + const float *ivec = reinterpret_cast(query); + int8_t *ovec = reinterpret_cast(&(*out)[0]); + float abs_max = 0.0f; + for (size_t i = 0; i < dim; ++i) { + auto abs = std::abs(ivec[i]); + if (abs > abs_max) { + abs_max = abs; + } + } + if (abs_max > 0.0f) { + float scale = 127 / abs_max; + for (size_t i = 0; i < dim; ++i) { + ovec[i] = static_cast(std::round(ivec[i] * scale)); + } + } else { + std::fill(ovec, ovec + dim, static_cast(1)); + } + return 0; +#else + return IndexError_NotImplemented; +#endif + } + + //! Transform queries + int transform(const void *query, const IndexQueryMeta &qmeta, uint32_t count, + std::string *oquery, IndexQueryMeta *ometa) const override { + size_t dim = qmeta.dimension(); + oquery->resize(count * + IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8, dim)); + ometa->set_meta(IndexMeta::DataType::DT_INT8, dim); + const float *ivec = reinterpret_cast(query); + int8_t *ovec = reinterpret_cast(&(*oquery)[0]); + for (size_t q = 0; q < count; ++q) { + float abs_max = 0.0f; + const float *in = &ivec[q * dim]; + int8_t *out = &ovec[q * dim]; + for (size_t i = 0; i < dim; ++i) { + auto abs = std::abs(in[i]); + if (abs > abs_max) { + abs_max = abs; + } + } + if (abs_max > 0.0f) { + float scale = 127 / abs_max; + for (size_t i = 0; i < dim; ++i) { + out[i] = static_cast(std::round(in[i] * scale)); + } + } else { + std::fill(out, out + dim, static_cast(1)); + } + } + return 0; + } + + //! Normalize results + int normalize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + IndexDocumentList & /*result*/) const override { + return 0; + } +}; + +/*! Int4QuantizerReformer for InnerProduct Metric + */ +class Int4QuantizerReformer4IP : public IndexReformer { + public: + //! Initialize Reformer + int init(const ailego::Params &) override { + return 0; + } + + //! Cleanup Reformer + int cleanup(void) override { + return 0; + } + + //! Load index from container + int load(IndexStorage::Pointer) override { + return 0; + } + + //! Unload index + int unload(void) override { + return 0; + } + + //! Transform query + int transform(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const override { + return IndexError_NotImplemented; + } + + //! Transform queries + int transform(const void *query, const IndexQueryMeta &qmeta, uint32_t count, + std::string *oquery, IndexQueryMeta *ometa) const override { + if (qmeta.dimension() & 0x1) { + LOG_ERROR("Unsuuport dim=%u for transform", qmeta.dimension()); + return IndexError_Unsupported; + } + + size_t dim = qmeta.dimension(); + oquery->resize(count * + IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT4, dim)); + ometa->set_meta(IndexMeta::DataType::DT_INT4, dim); + const float *ivec = reinterpret_cast(query); + uint8_t *ovec = reinterpret_cast(&(*oquery)[0]); + for (size_t q = 0; q < count; ++q) { + float abs_max = 0.0f; + float max = -std::numeric_limits::max(); + const float *in = &ivec[q * dim]; + uint8_t *out = &ovec[q * dim / 2]; + for (size_t i = 0; i < dim; ++i) { + float abs = std::abs(in[i]); + abs_max = std::max(abs_max, abs); + max = std::max(max, in[i]); + } + if (abs_max > 0.0f) { + float scale = ((7 * abs_max > 8 * max) ? 8 : 7) / abs_max; + for (size_t i = 0; i < dim; i += 2) { + auto v1 = static_cast(std::round(in[i] * scale)); + auto v2 = static_cast(std::round(in[i + 1] * scale)); + out[i / 2] = (static_cast(v1) << 4) | + (static_cast(v2) & 0xF); + } + } else { + std::fill(out, out + dim / 2, static_cast(9)); + } + } + return 0; + } + + //! Normalize results + int normalize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + IndexDocumentList & /*result*/) const override { + return 0; + } +}; + +int IVFCentroidIndex::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + meta_ = meta; + + params.get(PARAM_IVF_BUILDER_OPTIMIZER_CLASS, &builder_class_); + params.get(PARAM_IVF_BUILDER_OPTIMIZER_PARAMS, &builder_params_); + params.get(PARAM_IVF_SEARCHER_OPTIMIZER, &searcher_class_); + params.get(PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS, &searcher_params_); + + return 0; +} + +int IVFCentroidIndex::search(const void *query, const IndexQueryMeta &qmeta, + size_t count, + IndexSearcher::Context::Pointer &ctx) { + int ret = 0; + + if (reformer_) { + std::string buffer; + IndexQueryMeta ometa; + ret = reformer_->transform(query, qmeta, count, &buffer, &ometa); + if (ret != 0) { + LOG_ERROR("Failed to transform querys by reformer"); + return ret; + } + ret = searcher_->search_impl(buffer.data(), ometa, count, ctx); + } else { + ret = searcher_->search_impl(query, qmeta, count, ctx); + } + + ivf_check_with_msg(ret, "Failed to search in centroid index for %s", + IndexError::What(ret)); + + return 0; +} + +uint32_t IVFCentroidIndex::search_nearest_centroid(const void *query, + size_t len) { + //! Called in building index precedure, so transform the query is needless + if (len != meta_.element_size()) { + LOG_ERROR("Invalid query size actual: %zu, expected: %u", len, + meta_.element_size()); + return kInvalidID; + } + + thread_local IndexSearcher::Context::Pointer context( + searcher_->create_context()); + context->set_topk(1); + + IndexQueryMeta qmeta(meta_.data_type(), meta_.dimension()); + int ret = searcher_->search_impl(query, qmeta, context); + if (ret != 0 || context->result().empty()) { + LOG_ERROR("Failed to search nearest centroid, with ret %d", ret); + return kInvalidID; + } + + return static_cast(context->result()[0].key()); +} + +uint32_t IVFCentroidIndex::transform_and_search_nearest_centroid( + const void *record, const IndexQueryMeta &rmeta, + IndexSearcher::Context::Pointer &ctx) const { + int ret = 0; + if (reformer_) { + std::string buffer; + IndexQueryMeta ometa; + ret = reformer_->convert(record, rmeta, &buffer, &ometa); + if (ret != 0) { + LOG_ERROR("Failed to transform querys by reformer"); + return kInvalidID; + } + ret = searcher_->search_impl(buffer.data(), ometa, ctx); + } else { + ret = searcher_->search_impl(record, rmeta, ctx); + } + if (ret != 0 || ctx->result().empty()) { + LOG_ERROR("Failed to search in centroid index for %s", + IndexError::What(ret)); + return kInvalidID; + } + + return static_cast(ctx->result()[0].key()); +} + +IndexHolder::Pointer IVFCentroidIndex::quantize_holder( + const IndexHolder::Pointer &holder) { + auto input = holder; + if (meta_.reformer_name() == kMipsReformerName && + meta_.metric_name() == kL2MetricName && + (quantizer_class_ == kInt8QuantizerName || + quantizer_class_ == kInt4QuantizerName)) { + //! Reverse for Mips if do convert by integer quantizer + auto reverse = IndexFactory::CreateConverter(kMipsRevConverterName); + if (!reverse) { + LOG_ERROR("Failed to create converter %s", kMipsRevConverterName); + return nullptr; + } + ailego::Params params; + auto p = meta_.reformer_params(); + params.set(MIPS_REVERSE_CONVERTER_M_VALUE, + p.get_as_uint32(MIPS_REFORMER_M_VALUE)); + params.set(MIPS_REVERSE_CONVERTER_U_VALUE, + p.get_as_float(MIPS_REFORMER_U_VALUE)); + params.set(MIPS_REVERSE_CONVERTER_L2_NORM, + p.get_as_uint32(MIPS_REFORMER_L2_NORM)); + params.set(MIPS_REVERSE_CONVERTER_FORCED_SINGLE_FLOAT, + p.get_as_float(MIPS_REFORMER_FORCED_HALF_FLOAT)); + int ret = reverse->init(meta_, params); + if (ret != 0) { + LOG_ERROR("Fail to init converter %s", kMipsRevConverterName); + return nullptr; + } + ret = IndexConverter::TrainAndTransform(reverse, holder); + if (ret != 0) { + LOG_ERROR("Fail to transform converter %s", kMipsRevConverterName); + return nullptr; + } + input = reverse->result(); + meta_ = reverse->meta(); + meta_.set_metric(kIPMetricName, 0, ailego::Params()); + meta_.set_reformer("", 0, ailego::Params()); + } + + auto converter = IndexFactory::CreateConverter(quantizer_class_); + if (!converter) { + LOG_ERROR("Failed to create converter %s", quantizer_class_.c_str()); + return nullptr; + } + int ret = converter->init(meta_, quantizer_params_); + if (ret != 0) { + LOG_ERROR("Fail to init converter %s", quantizer_class_.c_str()); + return nullptr; + } + + ret = IndexConverter::TrainAndTransform(converter, input); + if (ret != 0) { + LOG_ERROR("Fail to tranform converter %s", quantizer_class_.c_str()); + return nullptr; + } + + meta_ = converter->meta(); + return converter->result(); +} + +int IVFCentroidIndex::build_index( + const IndexCluster::CentroidList ¢roid_list, + const IndexDumper::Pointer &dumper) { + IndexBuilder::Pointer builder = IndexFactory::CreateBuilder(builder_class_); + if (!builder) { + LOG_ERROR("Failed to create builder %s", builder_class_.c_str()); + return IndexError_NoExist; + } + + IndexHolder::Pointer holder = + std::make_shared(meta_, centroid_list); + if (!holder) { + return IndexError_NoMemory; + } + if (holder->count() == 0) { + LOG_ERROR("No centroids to build"); + return IndexError_InvalidArgument; + } + centroids_count_ = holder->count(); + + //! Set default params if not given + auto count = std::to_string( + static_cast(std::ceil(std::sqrt(centroids_count_ / 10.0)))); + // if (IsHcBuilder(builder_class_) && + // !builder_params_.has(hc::PARAM_HC_BUILDER_CENTROID_COUNT)) { + // builder_params_.set(hc::PARAM_HC_BUILDER_CENTROID_COUNT, count); + // } else if (builder_class_ == "GcBuilder" && + // !builder_params_.has(hc::PARAM_GC_BUILDER_CENTROID_COUNT)) { + // builder_params_.set(hc::PARAM_GC_BUILDER_CENTROID_COUNT, count); + // } + if (!quantizer_class_.empty()) { + holder = this->quantize_holder(holder); + if (!holder) { + return IndexError_Runtime; + } + } + + const auto name = builder_class_.c_str(); + int ret = builder->init(meta_, builder_params_); + ivf_check_with_msg(ret, "%s init failed, ret=%d", name, ret); + + // if (IsHcBuilder(builder_class_) && quantizer_class_.empty()) { + // auto trainer = this->prepare_trainer(centroid_list); + // ret = trainer ? builder->train(trainer) : builder->train(holder); + // } else { + // ret = builder->train(holder); + // } + + ret = builder->train(holder); + ivf_check_with_msg(ret, "%s train failed, ret=%d", name, ret); + + ret = builder->build(holder); + ivf_check_with_msg(ret, "%s build failed, ret=%d", name, ret); + + ret = builder->dump(dumper); + ivf_check_with_msg(ret, "%s dump failed, ret=%d", name, ret); + + ret = dumper->close(); + ivf_check_error_code(ret); + + return 0; +} + +int IVFCentroidIndex::build(const IndexCluster::CentroidList ¢roid_list) { + index_building_ = true; + //! Build and dump the index + IndexDumper::Pointer dumper = IndexFactory::CreateDumper("MemoryDumper"); + if (!dumper) { + LOG_ERROR("Failed to create MemoryDumper"); + return IndexError_NoExist; + } + path_ = IVFUtility::GenerateRandomPath(kTempralPathPrefix); + int ret = dumper->create(path_); + if (ret != 0) { + LOG_ERROR("IndexDumper create path %s failed", path_.c_str()); + return ret; + } + ret = this->build_index(centroid_list, dumper); + ivf_check_error_code(ret); + + auto rope = IndexMemory::Instance()->open(path_); + if (!rope) { + LOG_ERROR("Open memory path %s failed.", path_.c_str()); + return ret; + } + if (rope->count() != 1) { + LOG_ERROR("Graph Rope block count not equal with 1."); + return ret; + } + (*rope)[0].read(0, &data_, 0); + size_ = (*rope)[0].size(); + + //! Load the index + IndexStorage::Pointer container = + IndexFactory::CreateStorage("MemoryReadStorage"); + if (!container) { + LOG_ERROR("Failed to create MemoryReadStorage"); + return IndexError_NoExist; + } + ret = container->init(ailego::Params()); + ivf_check_with_msg(ret, "Failed to initialize MemoryReadStorage for %s", + IndexError::What(ret)); + ret = container->open(path_, false); + ivf_check_with_msg(ret, "Failed to load path in MemoryReadStorage for %s", + IndexError::What(ret)); + + ailego::Params searcher_params; + if (!searcher_class_.empty()) { + searcher_params.set(PARAM_IVF_SEARCHER_OPTIMIZER, searcher_class_); + } + if (!searcher_params_.empty()) { + searcher_params.set(PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS, searcher_params_); + } + ret = this->load(container, searcher_params); + ivf_check_with_msg(ret, "IVFCentroidIndex load failed with %s", + IndexError::What(ret)); + + return 0; +} + +int IVFCentroidIndex::load(const IndexStorage::Pointer &container, + const ailego::Params params) { + if (!container) { + LOG_ERROR("Invalid container"); + return IndexError_InvalidArgument; + } + + int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize meta from container"); + return ret; + } + + auto reformer_name = meta_.reformer_name(); + if (!reformer_name.empty()) { + LOG_DEBUG("Load CentroidIndex with reformer %s, metric %s", + reformer_name.c_str(), meta_.metric_name().c_str()); + if ((reformer_name == kInt8ReformerName || + reformer_name == kInt4ReformerName) && + meta_.metric_name() == kIPMetricName) { + if (reformer_name == kInt8ReformerName) { + reformer_ = std::make_shared(); + } else { + reformer_ = std::make_shared(); + } + if (!reformer_) { + return IndexError_NoMemory; + } + } else { + reformer_ = IndexFactory::CreateReformer(reformer_name); + if (!reformer_) { + LOG_ERROR("Failed to create reformer %s", reformer_name.c_str()); + return IndexError_NoExist; + } + } + ret = reformer_->init(meta_.reformer_params()); + ivf_check_with_msg(ret, "Failed to initialize reformer %s", + reformer_name.c_str()); + } + + searcher_class_ = meta_.searcher_name(); + params.get(PARAM_IVF_SEARCHER_OPTIMIZER, &searcher_class_); + params.get(PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS, &searcher_params_); + searcher_ = IndexFactory::CreateSearcher(searcher_class_); + if (!searcher_) { + LOG_ERROR("Failed to create searcher %s", searcher_class_.c_str()); + return IndexError_Runtime; + } + + auto searcher_params = meta_.searcher_params(); + searcher_params.merge(searcher_params_); + ret = searcher_->init(searcher_params); + ivf_check_with_msg(ret, "Failed to initialize searcher %s", + searcher_class_.c_str()); + + IndexMetric::Pointer metric; + if (index_building_) { + // The searcher index metric should specified in building process, + // otherwise the query_metric will be used in searching + metric = IndexFactory::CreateMetric(meta_.metric_name()); + ivf_assert_with_msg(metric, IndexError_NoExist, + "Failed to create metric %s", + meta_.metric_name().c_str()); + ret = metric->init(meta_, meta_.metric_params()); + ivf_check_with_msg(ret, "Failed to initialize metric"); + } + ret = searcher_->load(container, metric); + ivf_check_with_msg(ret, "Failed to load searcher %s", + searcher_class_.c_str()); + + return 0; +} + +IndexTrainer::Pointer IVFCentroidIndex::prepare_trainer( + const IndexCluster::CentroidList ¢roid_list) { + IndexCluster::CentroidList level1_centroids; + bool two_level = false; + for (auto &it : centroid_list) { + auto centroid = it; + if (!centroid.subitems().empty()) { + two_level = true; + } + centroid.mutable_subitems()->clear(); + centroid.mutable_similars()->clear(); + level1_centroids.emplace_back(centroid); + } + if (!two_level) { + return IndexTrainer::Pointer(); + } + + IndexBundle::Pointer bundle; + IndexCluster::Serialize(meta_, level1_centroids, &bundle); + return std::make_shared(meta_, bundle); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_centroid_index.h b/src/core/algorithm/ivf/ivf_centroid_index.h new file mode 100644 index 0000000..ca3ef4a --- /dev/null +++ b/src/core/algorithm/ivf/ivf_centroid_index.h @@ -0,0 +1,247 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "ivf_params.h" +#include "ivf_utility.h" + +namespace zvec { +namespace core { + +/*! IVF Centroid Index + */ +class IVFCentroidIndex { + public: + typedef std::shared_ptr Pointer; + + //! Constructor + IVFCentroidIndex(void) {} + + //! Destructor + ~IVFCentroidIndex(void) { + IndexMemory *instance = IndexMemory::Instance(); + if (instance) { + if (instance->has(path_)) { + instance->remove(path_); + } + } + } + + //! Initialize + int init(const IndexMeta &meta, const ailego::Params ¶ms); + + //! Set Quantizer for the index + void set_quantizer(const std::string &quantizer_name, + ailego::Params &quantizer_params) { + quantizer_class_ = quantizer_name; + quantizer_params_ = quantizer_params; + } + + //! Retrieve data address of the index + const void *data(void) const { + return data_; + } + + //! Retrieve size of the index + size_t size(void) const { + return size_; + } + + //! Create searcher context for centroid index + IndexSearcher::Context::Pointer create_context() const { + return searcher_ ? searcher_->create_context() : nullptr; + } + + //! Similarity search + int search(const void *query, const IndexQueryMeta &qmeta, size_t count, + IndexSearcher::Context::Pointer &ctx); + + //! Search the nearest point, must be called in local thread pool + uint32_t search_nearest_centroid(const void *query, size_t len); + + //! Transform Data and Search the nearest point, called while adding record + uint32_t transform_and_search_nearest_centroid( + const void *record, const IndexQueryMeta &rmeta, + IndexSearcher::Context::Pointer &ctx) const; + + //! Build Centroid Index From Centroid List + int build(const IndexCluster::CentroidList ¢roid_list); + + //! Load Centroid Index From container + int load(const IndexStorage::Pointer &container, const ailego::Params params); + + //! Retrieve centroid count of the index + size_t centroids_count(void) const { + return centroids_count_; + } + + //! Retrieve meta + const IndexMeta &meta() const { + return meta_; + } + + //! Retrieve reformer of the index + const IndexReformer::Pointer reformer(void) const { + return reformer_; + } + + static constexpr uint32_t kInvalidID = std::numeric_limits::max(); + + private: + /*! Centroids IndexHolder + */ + class CentroidsIndexHolder : public IndexHolder { + public: + class Iterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(std::vector *features) : features_(features) {} + + //! Destructor + ~Iterator(void) override {} + + //! Retrieve pointer of data + const void *data(void) const override { + return (*features_)[id_]; + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return id_ < features_->size(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return id_; + } + + //! Next iterator + void next(void) override { + ++id_; + } + + private: + //! Members + std::vector *features_{nullptr}; + uint32_t id_{0}; + }; + + //! Constructor + CentroidsIndexHolder(const IndexMeta &meta, + const IndexCluster::CentroidList ¢roid_list) + : dimension_(meta.dimension()), + element_size_(meta.element_size()), + data_type_(meta.data_type()) { + using CentroidList = IndexCluster::CentroidList; + + std::function get_leaf_features = + [&](const CentroidList ¢s) { + if (cents.empty()) { + return; + } + for (const auto &it : cents) { + if (it.subitems().empty()) { + features_.emplace_back(it.feature()); + } else { + get_leaf_features(it.subitems()); + } + } + }; + + get_leaf_features(centroid_list); + } + + //! Retrieve count of elements in holder (-1 indicates unknown) + size_t count(void) const override { + return features_.size(); + } + + //! Retrieve dimension + size_t dimension(void) const override { + return dimension_; + } + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return data_type_; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return element_size_; + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return true; + } + + //! Create a new iterator + IndexHolder::Iterator::Pointer create_iterator(void) override { + return IndexHolder::Iterator::Pointer( + new CentroidsIndexHolder::Iterator(&features_)); + } + + private: + //! Members + std::vector features_{}; + size_t dimension_{0}; + size_t element_size_{0}; + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_UNDEFINED}; + }; + + int build_index(const IndexCluster::CentroidList ¢roid_list, + const IndexDumper::Pointer &dumper); + + //! Prepare trainer for clustering index + IndexTrainer::Pointer prepare_trainer( + const IndexCluster::CentroidList ¢roid_list); + + //! Quantize the centroid vectors in holder + IndexHolder::Pointer quantize_holder(const IndexHolder::Pointer &holder); + + + private: + //! Constants + constexpr static const char *kDefaultBuilder = "FlatBuilder"; + constexpr static const char *kTempralPathPrefix = "IVF"; + + //! Members + IndexMeta meta_{}; + + IndexSearcher::Pointer searcher_{}; + IndexReformer::Pointer reformer_{}; + std::string builder_class_{kDefaultBuilder}; + std::string searcher_class_{}; + std::string quantizer_class_{}; + + std::string path_{}; + + ailego::Params builder_params_{}; + ailego::Params searcher_params_{}; + ailego::Params quantizer_params_{}; + + const void *data_{}; + size_t size_{}; + size_t centroids_count_{0}; + bool index_building_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_distance_calculator.cc b/src/core/algorithm/ivf/ivf_distance_calculator.cc new file mode 100644 index 0000000..f0c731c --- /dev/null +++ b/src/core/algorithm/ivf/ivf_distance_calculator.cc @@ -0,0 +1,46 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "ivf_distance_calculator.h" +#include + +namespace zvec { +namespace core { + +IVFDistanceCalculator::IVFDistanceCalculator(const IndexMeta &meta, + const IndexMetric::Pointer &metric, + uint32_t block_vec_cnt) + : metric_ptr_(metric), block_vec_cnt_(block_vec_cnt) { + row_distance_ = metric->distance(); + distanceXx1_ = metric->distance_matrix(block_vec_cnt, 1); + distances_.resize(33); + for (size_t b = 32; b != 0; b /= 2) { + distances_[b] = metric->distance_matrix(block_vec_cnt, b); + } + element_size_ = meta.element_size(); + dimension_ = meta.dimension(); + if (meta.major_order() == IndexMeta::MajorOrder::MO_COLUMN) { + column_major_order_ = true; + } else { + column_major_order_ = false; + } +} + +IVFDistanceCalculator::~IVFDistanceCalculator() { + row_distance_ = nullptr; + distanceXx1_ = nullptr; + distances_.clear(); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_distance_calculator.h b/src/core/algorithm/ivf/ivf_distance_calculator.h new file mode 100644 index 0000000..17a075d --- /dev/null +++ b/src/core/algorithm/ivf/ivf_distance_calculator.h @@ -0,0 +1,185 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include + +namespace zvec { +namespace core { + +class IVFDistanceCalculator { + public: + typedef std::shared_ptr Pointer; + + //! Constructor + IVFDistanceCalculator(const IndexMeta &meta, + const IndexMetric::Pointer &metric, + uint32_t block_vec_cnt); + + virtual ~IVFDistanceCalculator(); + + public: + inline void query_centroids_distance(const void *query, size_t qnum, + const void *feature, size_t fnum, + float *distances); + + inline void query_centroids_distance(const void *query, const void *feature, + size_t fnum, float *distances); + + inline void query_features_distance(const void *query, const void *feature, + size_t fnum, float *distances); + + inline void query_features_distance(const void *query, const void *feature, + bool column_major, size_t fnum, + float *distances); + + protected: + //! Row Major Distances -> Online + inline void row_major_distance(const void *query, size_t qnum, + const void *feature, size_t fnum, float *out); + + inline void row_major_distance(const void *query, const void *feature, + size_t fnum, float *out); + + template + inline void batch_query_centroids_distance(const void *query, + const void *feature, size_t fnum, + float *distances); + + protected: + IndexMetric::Pointer metric_ptr_{}; + IndexMetric::MatrixDistance row_distance_{nullptr}; + IndexMetric::MatrixDistance distanceXx1_{nullptr}; + std::vector distances_{}; + + size_t element_size_{0}; + size_t dimension_{0}; + uint32_t block_vec_cnt_{0}; + bool column_major_order_{false}; +}; + +void IVFDistanceCalculator::query_centroids_distance(const void *query, + size_t qnum, + const void *feature, + size_t fnum, + float *distances) { + if (column_major_order_) { + switch (qnum) { + case 1: + batch_query_centroids_distance<1>(query, feature, fnum, distances); + break; + case 16: + batch_query_centroids_distance<16>(query, feature, fnum, distances); + break; + case 8: + batch_query_centroids_distance<8>(query, feature, fnum, distances); + break; + case 4: + batch_query_centroids_distance<4>(query, feature, fnum, distances); + break; + case 2: + batch_query_centroids_distance<2>(query, feature, fnum, distances); + break; + case 32: + batch_query_centroids_distance<32>(query, feature, fnum, distances); + break; + default: + LOG_ERROR("Unsupported query num %zu.", qnum); + break; + } + } else { + const uint8_t *cur_query = reinterpret_cast(query); + for (size_t q = 0; q < qnum; ++q) { + this->row_major_distance(cur_query, feature, fnum, distances); + cur_query += element_size_; + distances += block_vec_cnt_; + } + } +} + +void IVFDistanceCalculator::query_centroids_distance(const void *query, + const void *feature, + size_t fnum, + float *distances) { + this->query_features_distance(query, feature, fnum, distances); +} + +void IVFDistanceCalculator::query_features_distance(const void *query, + const void *feature, + size_t fnum, + float *distances) { + if (column_major_order_) { + if (fnum == block_vec_cnt_) { + distanceXx1_(feature, query, dimension_, distances); + } else { + this->row_major_distance(query, feature, fnum, distances); + } + } else { + this->row_major_distance(query, feature, fnum, distances); + } +} + +void IVFDistanceCalculator::query_features_distance(const void *query, + const void *feature, + bool column_major, + size_t fnum, + float *distances) { + if (column_major) { + ailego_assert_with(fnum == block_vec_cnt_, "Invalid Block"); + distanceXx1_(feature, query, dimension_, distances); + } else { + this->row_major_distance(query, feature, fnum, distances); + } +} + +template +void IVFDistanceCalculator::batch_query_centroids_distance(const void *query, + const void *feature, + size_t fnum, + float *distances) { + if (fnum == block_vec_cnt_) { + distances_[Q](feature, query, dimension_, distances); + } else { + row_major_distance(query, Q, feature, fnum, distances); + } +} + +void IVFDistanceCalculator::row_major_distance(const void *query, size_t qnum, + const void *feature, size_t fnum, + float *out) { + const uint8_t *cur_query = reinterpret_cast(query); + for (size_t q = 0; q < qnum; ++q) { + const uint8_t *tmp_feature = reinterpret_cast(feature); + float *cur_out = out + q * fnum; + for (size_t f = 0; f < fnum; ++f) { + row_distance_(cur_query, tmp_feature, dimension_, cur_out + f); + tmp_feature += element_size_; + } + cur_query += element_size_; + } +} + +void IVFDistanceCalculator::row_major_distance(const void *query, + const void *feature, size_t fnum, + float *out) { + const uint8_t *cur_feature = reinterpret_cast(feature); + for (size_t f = 0; f < fnum; ++f) { + row_distance_(query, cur_feature, dimension_, out + f); + cur_feature += element_size_; + } +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_dumper.cc b/src/core/algorithm/ivf/ivf_dumper.cc new file mode 100644 index 0000000..761adaa --- /dev/null +++ b/src/core/algorithm/ivf/ivf_dumper.cc @@ -0,0 +1,409 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "ivf_dumper.h" + +namespace zvec { +namespace core { + +int IVFDumper::dump_inverted_vector(uint32_t inverted_list_id, uint64_t key, + const void *vec) { + int ret = this->check_dump_inverted_list(inverted_list_id); + ivf_check_error_code(ret); + + ++inverted_lists_meta_[cur_list_id_].vector_count; + ++header_.total_vector_count; + block_.emplace(key, vec, IndexMeta::MajorOrder::MO_ROW); + if (block_.full()) { + ret = this->dump_block(); + ivf_check_error_code(ret); + } + return 0; +} + +int IVFDumper::dump_inverted_block(uint32_t inverted_list_id, + const uint64_t *keys, const void *vecs, + uint32_t vector_count, bool column_major) { + int ret = this->check_dump_inverted_list(inverted_list_id); + ivf_check_error_code(ret); + + if (block_.match_order(column_major ? IndexMeta::MajorOrder::MO_COLUMN + : IndexMeta::MajorOrder::MO_ROW) && + vector_count == block_.capacity()) { + // Dump the block directly + size_t size = vector_count * meta_.element_size(); + size_t pd_size = ailego_align(size, 32) - size; + if (dumper_->write(vecs, size) != size) { + LOG_ERROR("Failed to write data into dumper %s", dumper_->name().c_str()); + return IndexError_WriteData; + } + if (pd_size > 0) { + std::string padding(pd_size, '\0'); + if (dumper_->write(padding.data(), pd_size) != pd_size) { + return IndexError_WriteData; + } + } + std::copy(keys, keys + vector_count, std::back_inserter(keys_)); + ++inverted_lists_meta_[cur_list_id_].block_count; + ++header_.block_count; + header_.inverted_body_size += size; + } else { + size_t step_size = meta_.element_size(); + if (column_major) { + step_size = IndexMeta::AlignSizeof(meta_.data_type()); + } + for (size_t i = 0; i < vector_count; ++i) { + auto v = reinterpret_cast(vecs) + i * step_size; + block_.emplace(keys[i], v, + column_major ? IndexMeta::MajorOrder::MO_COLUMN + : IndexMeta::MajorOrder::MO_ROW); + if (block_.full()) { + ret = this->dump_block(); + ivf_check_error_code(ret); + } + } + } + + inverted_lists_meta_[cur_list_id_].vector_count += vector_count; + header_.total_vector_count += vector_count; + + return 0; +} + +int IVFDumper::dump_container_segment(const IndexStorage::Pointer &container, + const std::string &segmemt_id) { + auto seg = container->get(segmemt_id, 2); + if (!seg) { + LOG_ERROR("Failed to fetch segment %s from %s", segmemt_id.c_str(), + container->name().c_str()); + return IndexError_InvalidFormat; + } + + const size_t batch_size = 32 * 1024; + const size_t total_size = seg->data_size() + seg->padding_size(); + size_t off = 0; + while (off < total_size) { + const void *data = nullptr; + size_t rd_size = std::min(batch_size, total_size - off); + if (seg->read(off, &data, rd_size) != rd_size) { + LOG_ERROR("Failed to read data, off=%zu size=%zu", off, rd_size); + return IndexError_ReadData; + } + if (dumper_->write(data, rd_size) != rd_size) { + LOG_ERROR("Failed to write data, size=%zu", rd_size); + return IndexError_WriteData; + } + off += rd_size; + } + + int ret = dumper_->append(segmemt_id, seg->data_size(), seg->padding_size(), + seg->data_crc()); + ivf_check_with_msg(ret, "Failed to append %s", segmemt_id.c_str()); + + dumped_size_ += total_size; + + return 0; +} + +int IVFDumper::dump_inverted_vector_finished(void) { + //! Dump Inverted Index Segment + if (!block_.empty()) { + int ret = this->dump_block(); + ivf_check_error_code(ret); + } + header_.block_size = block_.block_size(); + size_t segment_size = header_.inverted_body_size; + int ret = dumper_->append(IVF_INVERTED_BODY_SEG_ID, segment_size, 0, 0); + if (ret != 0) { + LOG_ERROR("Failed to append to segment %s, ret=%d", + IVF_INVERTED_BODY_SEG_ID.c_str(), ret); + return ret; + } + dumped_size_ += segment_size; + + //! Dump Inverted Index Header Segment + std::string str; + meta_.serialize(&str); + header_.header_size = sizeof(header_) + str.size(); + header_.index_meta_size = str.size(); + header_.inverted_list_count = inverted_lists_meta_.size(); + if (dumper_->write(&header_, sizeof(header_)) != sizeof(header_)) { + LOG_ERROR("Failed to write data, size %zu", sizeof(header_)); + return IndexError_WriteData; + } + if (dumper_->write(str.data(), str.size()) != str.size()) { + LOG_ERROR("Failed to write data, size %zu", str.size()); + return IndexError_WriteData; + } + size_t padding_size = 0; + ret = this->dump_padding(header_.header_size, &padding_size); + ivf_check_error_code(ret); + ret = dumper_->append(IVF_INVERTED_HEADER_SEG_ID, header_.header_size, + padding_size, 0); + if (ret != 0) { + LOG_ERROR("Failed to append to segment %s, ret:%d", + IVF_INVERTED_HEADER_SEG_ID.c_str(), ret); + return ret; + } + dumped_size_ += header_.header_size + padding_size; + + LOG_DEBUG( + "Dump header info: blocks=%u block_size=%u block_vec_count=%u " + "inverted_list_count=%u total_vecs=%u inverted_size=%zu", + header_.block_count, header_.block_size, header_.block_vector_count, + header_.inverted_list_count, header_.total_vector_count, + static_cast(header_.inverted_body_size)); + + //! Dump Inverted Lists Meta Segment + segment_size = inverted_lists_meta_.size() * sizeof(InvertedListMeta); + ret = this->dump_segment(IVF_INVERTED_META_SEG_ID, + inverted_lists_meta_.data(), segment_size); + ivf_check_error_code(ret); + + //! Dump Keys Segment + ret = this->dump_segment(IVF_KEYS_SEG_ID, keys_.data(), + keys_.size() * sizeof(keys_[0])); + ivf_check_error_code(ret); + + //! Dump Mapping Segment + auto mapping = std::make_shared>(); + IVFUtility::Sort(keys_.data(), mapping.get(), keys_.size()); + ret = this->dump_segment(IVF_MAPPING_SEG_ID, mapping->data(), + mapping->size() * sizeof(uint32_t)); + ivf_check_error_code(ret); + mapping.reset(); + + //! Dump the Offsets Segment + return this->dump_offsets_segment(); +} + +int IVFDumper::dump_centroid_index(const void *data, size_t size) { + int ret = this->dump_segment(IVF_CENTROID_SEG_ID, data, size); + ivf_check_error_code(ret); + + return 0; +} + +int IVFDumper::dump_quantizer_params( + const std::vector &quantizers) { + if (meta_.reformer_name() != kInt8ReformerName && + meta_.reformer_name() != kInt4ReformerName) { + // IntegerQuantizer params is support only + return 0; + } + if (quantizers.size() == 1) { + //! Donot dump, using reformer params in IndexMeta + return 0; + } + + if (quantizers.size() != header_.inverted_list_count) { + LOG_ERROR("Mismatch size, quantizers=%zu, inverted_list_count=%u", + quantizers.size(), header_.inverted_list_count); + return IndexError_Logic; + } + bool int8_quantizer = meta_.reformer_name() == kInt8ReformerName; + std::vector params; + params.resize(header_.inverted_list_count); + for (size_t i = 0; i < quantizers.size(); ++i) { + auto &p = quantizers[i]->meta().reformer_params(); + auto &scale_key = int8_quantizer ? INT8_QUANTIZER_REFORMER_SCALE + : INT4_QUANTIZER_REFORMER_SCALE; + auto &bias_key = int8_quantizer ? INT8_QUANTIZER_REFORMER_BIAS + : INT4_QUANTIZER_REFORMER_BIAS; + if (inverted_lists_meta_[i].vector_count > 0 && + (!p.has(scale_key) || !p.has(bias_key))) { + LOG_ERROR("Miss reformer params %s or %s", bias_key.c_str(), + scale_key.c_str()); + return IndexError_Logic; + } + + params[i].bias = p.get_as_float(bias_key); + params[i].scale = p.get_as_float(scale_key); + } + + return this->dump_segment( + int8_quantizer ? IVF_INT8_QUANTIZED_PARAMS_SEG_ID + : IVF_INT4_QUANTIZED_PARAMS_SEG_ID, + params.data(), params.size() * sizeof(InvertedIntegerQuantizerParams)); +} + +int IVFDumper::dump_original_vector(const void *data, size_t size) { + if (dumped_feature_count_ >= header_.total_vector_count) { + LOG_ERROR("Dump too much orignal features, expect=%u", + header_.total_vector_count); + return IndexError_Logic; + } + + if (dumper_->write(data, size) != size) { + LOG_ERROR("Dumper write features failed"); + return IndexError_WriteData; + } + dumped_features_size_ += size; + ++dumped_feature_count_; + if (dumped_feature_count_ == header_.total_vector_count) { + //! Dump features finished, dump the meta + size_t padding_size = 0; + int ret = this->dump_padding(size, &padding_size); + ivf_check_error_code(ret); + + ret = dumper_->append(IVF_FEATURES_SEG_ID, dumped_features_size_, + padding_size, 0); + if (ret != 0) { + LOG_ERROR("Dumper append segment %s failed, ret:%d", + IVF_FEATURES_SEG_ID.c_str(), ret); + return ret; + } + dumped_size_ += dumped_features_size_; + } + + return 0; +} + +int IVFDumper::check_dump_inverted_list(uint32_t inverted_list_id) { + if (inverted_list_id < cur_list_id_) { + LOG_ERROR("Invalid backward vector dumping, want=%u cur=%u", + inverted_list_id, cur_list_id_); + return IndexError_Logic; + } + if (inverted_list_id >= inverted_lists_meta_.size()) { + LOG_ERROR("Invalid inverted_list_id=%u, lists_size=%zu", inverted_list_id, + inverted_lists_meta_.size()); + return IndexError_Logic; + } + if (inverted_list_id != cur_list_id_) { + //! flush previous inverted_list block + int ret = this->dump_block(); + ivf_check_error_code(ret); + for (auto idx = cur_list_id_ + 1; idx <= inverted_list_id; ++idx) { + inverted_lists_meta_[idx].offset = header_.inverted_body_size; + inverted_lists_meta_[idx].id_offset = header_.total_vector_count; + } + cur_list_id_ = inverted_list_id; + } + + return 0; +} + +int IVFDumper::dump_offsets_segment(void) const { + bool col_pri = meta_.major_order() == IndexMeta::MajorOrder::MO_COLUMN; + size_t total_size = 0; + for (size_t i = 0; i < inverted_lists_meta_.size(); ++i) { + std::vector offsets; + const auto &m = inverted_lists_meta_[i]; + size_t vec_cnt = m.vector_count; + size_t idx = 0; + uint64_t off = m.offset; + size_t align_idx = vec_cnt - vec_cnt % block_vector_count_; + for (size_t j = 0; j < vec_cnt; ++j) { + if (col_pri && j < align_idx) { + offsets.emplace_back(off + idx * block_.align_size(), true); + } else { + offsets.emplace_back(off + idx * block_.element_size(), false); + } + ++idx; + if (idx == block_vector_count_) { + off += header_.block_size; + idx = 0; + } + } + if (idx != 0) { + off += (vec_cnt - align_idx) * meta_.element_size(); + } + + size_t len = offsets.size() * sizeof(offsets[0]); + size_t actual_len = dumper_->write(offsets.data(), len); + if (actual_len != len) { + LOG_ERROR("Write offsets failed expect %zu, actual: %zu.", len, + actual_len); + return IndexError_WriteData; + } + total_size += len; + } + + size_t padding_size = 0; + int ret = this->dump_padding(total_size, &padding_size); + ivf_check_error_code(ret); + + ret = dumper_->append(IVF_OFFSETS_SEG_ID, total_size, padding_size, 0); + if (ret != 0) { + LOG_ERROR("Dumper append segment %s failed, ret:%d", + IVF_OFFSETS_SEG_ID.c_str(), ret); + return ret; + } + + dumped_size_ += total_size + padding_size; + + return 0; +} + +int IVFDumper::dump_segment(const std::string &segment_id, const void *data, + size_t size) const { + size_t len = dumper_->write(data, size); + if (len != size) { + LOG_ERROR("Dump segment %s data failed, expect=%zu, actual=%zu", + segment_id.c_str(), size, len); + return IndexError_WriteData; + } + + size_t padding_size = 0; + int ret = this->dump_padding(size, &padding_size); + ivf_check_error_code(ret); + + uint32_t crc = ailego::Crc32c::Hash(data, size); + ret = dumper_->append(segment_id, size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", segment_id.c_str(), ret); + return ret; + } + dumped_size_ += size + padding_size; + + return 0; +} + +int IVFDumper::dump_padding(size_t data_size, size_t *padding_size) const { + *padding_size = IVFUtility::AlignedSize(data_size) - data_size; + if (*padding_size == 0) { + return 0; + } + + std::string padding(*padding_size, '\0'); + if (dumper_->write(padding.data(), *padding_size) != *padding_size) { + LOG_ERROR("Append padding failed, size %lu", *padding_size); + return IndexError_WriteData; + } + + return 0; +} + +int IVFDumper::dump_block(void) { + if (block_.empty()) { + return 0; + } + + size_t size = ailego_align(block_.bytes(), 32); + if (dumper_->write(block_.data(), size) != size) { + LOG_ERROR("Failed to write data into dumper %s", dumper_->name().c_str()); + return IndexError_WriteData; + } + auto &keys = block_.keys(); + std::copy(keys.begin(), keys.end(), std::back_inserter(keys_)); + ++inverted_lists_meta_[cur_list_id_].block_count; + ++header_.block_count; + header_.inverted_body_size += size; + block_.clear(); + + return 0; +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/ivf/ivf_dumper.h b/src/core/algorithm/ivf/ivf_dumper.h new file mode 100644 index 0000000..81c03c8 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_dumper.h @@ -0,0 +1,260 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "metric/metric_params.h" +#include "ivf_index_format.h" +#include "ivf_params.h" +#include "ivf_utility.h" + +namespace zvec { +namespace core { + +/*! Quantized Clustering Dumper + */ +class IVFDumper { + public: + typedef std::shared_ptr Pointer; + + //! Vectors block + class Block { + public: + //! Initialize block + void init(const IndexMeta &meta, uint32_t max_vec_count) { + element_size_ = meta.element_size(); + auto bsize = IVFUtility::AlignedSize(max_vec_count, element_size_); + data_.resize(bsize); + count_ = 0u; + major_order_ = meta.major_order(); + align_size_ = IndexMeta::AlignSizeof(meta.data_type()); + units_ = element_size_ / align_size_; + max_vec_count_ = max_vec_count; + keys_.reserve(max_vec_count_); + } + + //! Add a vector to the block in row major order + //! If the block is full and the block order is column, make a + //! transpose + void emplace(uint64_t key, const void *vec, IndexMeta::MajorOrder order) { + switch (align_size_) { + case 2: + do_emplace(vec, order); + break; + case 4: + do_emplace(vec, order); + break; + case 8: + do_emplace(vec, order); + break; + default: + ailego_check_with(false, "Unsupport Aligned Size"); + } + keys_.emplace_back(key); + } + + bool full(void) const { + return count_ == max_vec_count_; + } + + const void *data(void) const { + return data_.data(); + } + + void clear(void) { + count_ = 0u; + keys_.clear(); + } + + bool empty(void) const { + return count_ == 0u; + } + + size_t size(void) const { + return count_; + } + + size_t capacity(void) const { + return max_vec_count_; + } + + size_t align_size(void) const { + return align_size_; + } + + size_t element_size(void) const { + return element_size_; + } + + //! Retrieve block data size + size_t bytes(void) const { + return element_size_ * count_; + } + + //! Retrieve max block size in bytes + size_t block_size(void) const { + return data_.size(); + } + + IndexMeta::MajorOrder major_order(void) const { + return major_order_; + } + + const std::vector &keys(void) const { + return keys_; + } + + bool match_order(IndexMeta::MajorOrder column_major) const { + return major_order_ == column_major; + } + + private: + //! Transpose the block vectors + void transpose() { + std::vector buf(data_.size()); + IVFUtility::Transpose(align_size_, data_.data(), count_, units_, + buf.data()); + data_.swap(buf); + } + + template + void do_emplace(const void *vec, IndexMeta::MajorOrder order) { + ailego_assert_with(count_ < max_vec_count_, "emplace a full block"); + + T *dst = reinterpret_cast(data_.data() + element_size_ * count_); + const T *src = reinterpret_cast(vec); + size_t step = order == IndexMeta::MO_ROW ? 1 : max_vec_count_; + for (auto i = 0u; i < units_; ++i) { + *dst = *src; + dst++; + src += step; + } + + count_++; + if (full() && major_order_ == IndexMeta::MO_COLUMN) { + transpose(); + } + } + + private: + //! Members + std::vector data_{}; + std::vector keys_{}; + uint32_t count_{0u}; + uint32_t units_{0u}; + uint32_t align_size_{0u}; + uint32_t element_size_{0u}; + uint32_t max_vec_count_{0u}; + IndexMeta::MajorOrder major_order_{}; + }; + + //! Constructor + IVFDumper(const IndexMeta &meta, const IndexDumper::Pointer &dumper, + size_t inverted_list_count, size_t block_vector_count) + : meta_(meta), + dumper_(dumper), + block_vector_count_(block_vector_count), + inverted_lists_meta_(inverted_list_count) { + block_.init(meta, block_vector_count_); + header_.block_vector_count = block_vector_count_; + } + + //! Constructor + IVFDumper(const IndexMeta &meta, const IndexDumper::Pointer &dumper, + size_t inverted_list_count) + : IVFDumper(meta, dumper, inverted_list_count, kDefaultBlockCount) {} + + //! Destructor + ~IVFDumper() { + // Check the dumper status + if (dumped_feature_count_ > 0 && + dumped_feature_count_ != header_.total_vector_count) { + LOG_ERROR("Dumped features=%u mismatch from invertedVecs=%u", + dumped_feature_count_, header_.total_vector_count); + ailego_assert_with(false, "invalid status"); + } + } + + //! Dump a vector in row major order + int dump_inverted_vector(uint32_t inverted_list_id, uint64_t key, + const void *vec); + + int dump_inverted_block(uint32_t inverted_list_id, const uint64_t *keys, + const void *vecs, uint32_t vector_count, + bool column_major); + + //! Finish dump the inverted vectors + int dump_inverted_vector_finished(void); + + //! Dump the centroids index + int dump_centroid_index(const void *data, size_t size); + + //! Dump params for each inverted list quantizer + int dump_quantizer_params( + const std::vector &quantizers); + + //! Dump the original vector, which doesnot been quantized + int dump_original_vector(const void *data, size_t size); + + //! Retrieve total dumped size + size_t dumped_size(void) const { + return dumped_size_; + } + + //! Retrieve total dumped vector count + size_t dumped_count(void) const { + return header_.total_vector_count; + } + + //! Dump the segment from container + int dump_container_segment(const IndexStorage::Pointer &container, + const std::string &segmemt_id); + + private: + int check_dump_inverted_list(uint32_t inverted_list_id); + + //! Dump offsets segment + int dump_offsets_segment(void) const; + + //! Dump a segment + int dump_segment(const std::string &segment_id, const void *data, + size_t size) const; + + //! Dump segment padding + int dump_padding(size_t data_size, size_t *padding_size) const; + + //! Dump a vector block + int dump_block(void); + + private: + //! Constants + static constexpr size_t kDefaultBlockCount = 32u; + + //! Members + Block block_{}; // vectors grouped in block + const IndexMeta meta_{}; // IndexMeta of the inverted index + const IndexDumper::Pointer dumper_{}; + size_t block_vector_count_{kDefaultBlockCount}; + std::vector inverted_lists_meta_{}; + std::vector keys_{}; + uint32_t cur_list_id_{0}; + uint32_t dumped_feature_count_{0}; + size_t dumped_features_size_{0}; + mutable size_t dumped_size_{0}; + InvertedIndexHeader header_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_entity.cc b/src/core/algorithm/ivf/ivf_entity.cc new file mode 100644 index 0000000..decc86d --- /dev/null +++ b/src/core/algorithm/ivf/ivf_entity.cc @@ -0,0 +1,989 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "ivf_entity.h" +#include +#include "ivf_utility.h" +namespace zvec { +namespace core { + +//! Initialize +int IVFEntity::IVFReformerWrapper::init(const IndexMeta &imeta) { + auto &name = imeta.reformer_name(); + + if (name.empty()) { + type_ = kReformerTpNone; + return 0; + } + + auto reformer = IndexFactory::CreateReformer(name); + if (!reformer) { + LOG_ERROR("Failed to create reformer %s", name.c_str()); + return IndexError_NoExist; + } + int ret = reformer->init(imeta.reformer_params()); + ivf_check_with_msg(ret, "Failed to init reformer %s", name.c_str()); + + reformer_ = std::move(reformer); + + if (name == kInt8ReformerName) { + if (imeta.metric_name() == kIPMetricName) { + type_ = kReformerTpInnerProductInt8; + return 0; + } + auto &key = INT8_QUANTIZER_REFORMER_SCALE; + if (!imeta.reformer_params().has(key)) { + LOG_ERROR("Missing param %s in reformer %s", key.c_str(), name.c_str()); + return IndexError_InvalidArgument; + }; + float scale = imeta.reformer_params().get_as_float(key); + reciprocal_ = scale == 0.0 ? 1.0 : (1.0 / scale); + type_ = kReformerTpInt8; + } else if (name == kInt4ReformerName) { + if (imeta.metric_name() == kIPMetricName) { + type_ = kReformerTpInnerProductInt4; + return 0; + } + auto &key = INT4_QUANTIZER_REFORMER_SCALE; + if (!imeta.reformer_params().has(key)) { + LOG_ERROR("Missing param %s in reformer %s", key.c_str(), name.c_str()); + return IndexError_InvalidArgument; + }; + float scale = imeta.reformer_params().get_as_float(key); + reciprocal_ = scale == 0.0 ? 1.0 : (1.0 / scale); + type_ = kReformerTpInt4; + } else { + type_ = kReformerTpDefault; + } + + LOG_DEBUG("Init QcReformer with %s, type=%u", name.c_str(), type_); + + return 0; +} + +//! Load reformer state (e.g. rotation matrix) from storage +int IVFEntity::IVFReformerWrapper::load(const IndexStorage::Pointer &storage) { + if (!reformer_) { + return 0; + } + int ret = reformer_->load(storage); + ivf_check_with_msg(ret, "Failed to load reformer state"); + return 0; +} + +//! Update the params, Called by gpu searcher only +int IVFEntity::IVFReformerWrapper::update(const IndexMeta &meta) { + auto &name = meta.reformer_name(); + if (name == kInt4ReformerName && meta.metric_name() == kL2MetricName) { + auto &key = INT4_QUANTIZER_REFORMER_SCALE; + if (!meta.reformer_params().has(key)) { + LOG_ERROR("Missing param %s in reformer %s", key.c_str(), name.c_str()); + return IndexError_InvalidArgument; + }; + float scale = meta.reformer_params().get_as_float(key); + reciprocal_ = scale == 0.0 ? 1.0 : (1.0 / scale / kNormalizeScaleFactor); + type_ = kReformerTpInt8; + + ailego::Params params; + float int8_scale = scale * kNormalizeScaleFactor; + params.set(INT8_QUANTIZER_REFORMER_SCALE, int8_scale); + float bias = + meta.reformer_params().get_as_float(INT4_QUANTIZER_REFORMER_BIAS); + params.set(INT8_QUANTIZER_REFORMER_BIAS, bias); + params.set( + INT4_QUANTIZER_REFORMER_METRIC, + meta.reformer_params().get_as_string(INT4_QUANTIZER_REFORMER_METRIC)); + + auto reformer = IndexFactory::CreateReformer(kInt8ReformerName); + if (!reformer) { + LOG_ERROR("Failed to create reformer %s", name.c_str()); + return IndexError_NoExist; + } + int ret = reformer->init(params); + ivf_check_with_msg(ret, "Failed to init reformer %s", name.c_str()); + + reformer_ = reformer; + + LOG_DEBUG("Init QcReformer with %s, type=%u", name.c_str(), type_); + } + + return 0; +} + +//! Transform a query +int IVFEntity::IVFReformerWrapper::transform(const void *query, + const IndexQueryMeta &qmeta, + const void **out, + IndexQueryMeta *ometa) { + int ret = 0; + + switch (type_) { + case kReformerTpNone: + *out = query; + *ometa = qmeta; + break; + + case kReformerTpInnerProductInt8: + if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + scales_.resize(1); + buffer_.resize(IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8, + qmeta.dimension())); + this->transform(0, static_cast(query), qmeta.dimension(), + reinterpret_cast(&buffer_[0])); + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_INT8, qmeta.dimension()); + *out = buffer_.data(); + break; + + case kReformerTpInnerProductInt4: + if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + scales_.resize(1); + buffer_.resize(IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT4, + qmeta.dimension())); + this->transform(0, static_cast(query), qmeta.dimension(), + reinterpret_cast(&buffer_[0])); + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_INT4, qmeta.dimension()); + *out = buffer_.data(); + break; + + case kReformerTpInt8: + case kReformerTpInt4: + /* FALLTHRU */ + case kReformerTpDefault: + ret = reformer_->transform(query, qmeta, &buffer_, ometa); + *out = buffer_.data(); + break; + + default: + ret = IndexError_Unsupported; + break; + } + + return ret; +} + +//! Transform querys +int IVFEntity::IVFReformerWrapper::transform(const void *query, + const IndexQueryMeta &qmeta, + uint32_t count, const void **out, + IndexQueryMeta *ometa) { + int ret = 0; + + switch (type_) { + case kReformerTpNone: + *out = query; + *ometa = qmeta; + break; + + case kReformerTpInnerProductInt8: + if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + scales_.resize(count); + buffer_.resize(count * + IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8, + qmeta.dimension())); + { + const float *ivec = reinterpret_cast(query); + int8_t *ovec = reinterpret_cast(&buffer_[0]); + for (size_t i = 0; i < count; ++i) { + this->transform(i, &ivec[i * qmeta.dimension()], qmeta.dimension(), + &ovec[i * qmeta.dimension()]); + } + } + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_INT8, qmeta.dimension()); + *out = buffer_.data(); + break; + + case kReformerTpInnerProductInt4: + if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + scales_.resize(count); + buffer_.resize(count * + IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT4, + qmeta.dimension())); + { + const float *ivec = reinterpret_cast(query); + uint8_t *ovec = reinterpret_cast(&buffer_[0]); + for (size_t i = 0; i < count; ++i) { + this->transform(i, &ivec[i * qmeta.dimension()], qmeta.dimension(), + &ovec[i * qmeta.dimension() / 2]); + } + } + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_INT4, qmeta.dimension()); + *out = buffer_.data(); + break; + + case kReformerTpInt8: + case kReformerTpInt4: + /* FALLTHRU */ + case kReformerTpDefault: + ret = reformer_->transform(query, qmeta, count, &buffer_, ometa); + *out = buffer_.data(); + break; + + default: + ret = IndexError_Unsupported; + break; + } + + return ret; +} + +//! Transform querys +int IVFEntity::IVFReformerWrapper::transform_gpu(const void *query, + const IndexQueryMeta &qmeta, + uint32_t count, + const void **out, + IndexQueryMeta *ometa) { + int ret = 0; + + switch (type_) { + case kReformerTpNone: + case kReformerTpDefault: + *out = query; + *ometa = qmeta; + break; + + case kReformerTpInnerProductInt4: + case kReformerTpInnerProductInt8: + if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + scales_.resize(count); + buffer_.resize(count * + IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8, + qmeta.dimension())); + { + const float *ivec = reinterpret_cast(query); + int8_t *ovec = reinterpret_cast(&buffer_[0]); + for (size_t i = 0; i < count; ++i) { + this->transform(i, &ivec[i * qmeta.dimension()], qmeta.dimension(), + &ovec[i * qmeta.dimension()]); + } + } + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_INT8, qmeta.dimension()); + *out = buffer_.data(); + break; + + case kReformerTpInt8: + case kReformerTpInt4: + ret = reformer_->transform(query, qmeta, count, &buffer_, ometa); + *out = buffer_.data(); + break; + + default: + ret = IndexError_Unsupported; + break; + } + + return ret; +} + + +//! Convert a record +int IVFEntity::IVFReformerWrapper::convert(const void *record, + const IndexQueryMeta &rmeta, + const void **out, + IndexQueryMeta *ometa) { + if (type_ == kReformerTpNone) { + *out = record; + *ometa = rmeta; + return 0; + } + + int ret = reformer_->convert(record, rmeta, &buffer_, ometa); + *out = buffer_.data(); + return ret; +} + +//! Convert records +int IVFEntity::IVFReformerWrapper::convert(const void *records, + const IndexQueryMeta &rmeta, + uint32_t count, const void **out, + IndexQueryMeta *ometa) { + if (type_ == kReformerTpNone) { + *out = records; + *ometa = rmeta; + return 0; + } + int ret = reformer_->convert(records, rmeta, count, &buffer_, ometa); + *out = buffer_.data(); + return ret; +} + +//! Normalize score +void IVFEntity::IVFReformerWrapper::normalize(size_t qidx, + IndexDocumentHeap *heap) const { + switch (type_) { + case kReformerTpNone: + return; + + case kReformerTpInnerProductInt8: + case kReformerTpInnerProductInt4: + ailego_assert_with(qidx < scales_.size(), "invalid index"); + { + auto reciprocal = 1.0f / scales_[qidx]; + for (auto &it : *heap) { + *it.mutable_score() *= reciprocal; + } + } + break; + + case kReformerTpInt8: + case kReformerTpInt4: + for (auto &it : *heap) { + *it.mutable_score() *= reciprocal_; + } + break; + + default: + // Not support + break; + } +} + +//! Normalize score +void IVFEntity::IVFReformerWrapper::normalize(size_t qidx, const void *query, + const IndexQueryMeta &qmeta, + IndexDocumentHeap *heap) const { + switch (type_) { + case kReformerTpNone: + return; + + case kReformerTpInnerProductInt8: + case kReformerTpInnerProductInt4: + ailego_assert_with(qidx < scales_.size(), "invalid index"); + { + auto reciprocal = 1.0f / scales_[qidx]; + for (auto &it : *heap) { + *it.mutable_score() *= reciprocal; + } + } + break; + + case kReformerTpInt8: + case kReformerTpInt4: + for (auto &it : *heap) { + *it.mutable_score() *= reciprocal_; + } + break; + + case kReformerTpDefault: + reformer_->normalize(query, qmeta, *heap); + break; + + default: + // Not support + LOG_ERROR("Not a supported type in QC reformer, type: %u", type_); + break; + } +} + +void IVFEntity::IVFReformerWrapper::transform(size_t qidx, const float *in, + size_t dim, int8_t *out) { + ailego_assert_with(qidx < scales_.size(), "invalid index"); + + float abs_max = 0.0f; + for (size_t i = 0; i < dim; ++i) { + auto abs = std::abs(in[i]); + if (abs > abs_max) { + abs_max = abs; + } + } + + if (abs_max > 0.0f) { + float scale = 127 / abs_max; + for (size_t i = 0; i < dim; ++i) { + out[i] = static_cast(std::round(in[i] * scale)); + } + scales_[qidx] = scale; + } else { + std::fill(out, out + dim, static_cast(1)); + scales_[qidx] = std::numeric_limits::max(); + } +} + +void IVFEntity::IVFReformerWrapper::transform(size_t qidx, const float *in, + size_t dim, uint8_t *out) { + ailego_assert_with(qidx < scales_.size(), "invalid index"); + ailego_assert_with(dim % 2 == 0, "invalid dim"); + + float abs_max = 0.0f; + float max = -std::numeric_limits::max(); + for (size_t i = 0; i < dim; ++i) { + float abs = std::abs(in[i]); + abs_max = std::max(abs_max, abs); + max = std::max(max, in[i]); + } + if (abs_max > 0.0f) { + float scale = ((7 * abs_max > 8 * max) ? 8 : 7) / abs_max; + for (size_t i = 0; i < dim; i += 2) { + auto v1 = static_cast(std::round(in[i] * scale)); + auto v2 = static_cast(std::round(in[i + 1] * scale)); + out[i / 2] = + (static_cast(v1) & 0xF) | (static_cast(v2) << 4); + } + scales_[qidx] = scale; + } else { + std::fill(out, out + dim / 2, static_cast(9)); + scales_[qidx] = std::numeric_limits::max(); + } +} + +int IVFEntity::load_header(const IndexStorage::Pointer &container) { + //! Load the Header Segment + auto header = container->get(IVF_INVERTED_HEADER_SEG_ID); + if (!header) { + LOG_ERROR("Failed to get segment %s", IVF_INVERTED_HEADER_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + if (header->data_size() < sizeof(header_)) { + LOG_ERROR("Invalid format for segment %s", + IVF_INVERTED_HEADER_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + const void *data = nullptr; + if (header->read(0, &data, header->data_size()) != header->data_size()) { + LOG_ERROR("Failed to read data, segment %s", + IVF_INVERTED_HEADER_SEG_ID.c_str()); + return IndexError_ReadData; + } + std::memcpy(&header_, data, sizeof(header_)); + if (header_.header_size < sizeof(header_) + header_.index_meta_size || + header_.header_size > header->data_size()) { + LOG_ERROR("Invalid header size %u", header_.header_size); + return IndexError_InvalidFormat; + } + + //! Load the index meta + if (!meta_.deserialize( + reinterpret_cast(data) + sizeof(header_), + header_.index_meta_size)) { + LOG_ERROR("Failed to deserialize index meta"); + return IndexError_InvalidFormat; + } + + int ret = reformer_.init(meta_); + ivf_check_error_code(ret); + + //! Create the distance calculator + auto metric = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric) { + LOG_ERROR("Failed to create metric %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = metric->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("Failed to initialize metric %s", meta_.metric_name().c_str()); + return ret; + } + calculator_ = std::make_shared( + meta_, metric->query_metric() ? metric->query_metric() : metric, + header_.block_vector_count); + if (!calculator_) { + return IndexError_NoMemory; + } + + return 0; +} + +int IVFEntity::load(const IndexStorage::Pointer &container) { + int ret = this->load_header(container); + ivf_check_error_code(ret); + + //! Load the remaining segments + container_ = container; + + //! Load reformer state (e.g. rotation matrix) from the main container, + //! which holds the rotator segment dumped at build time. + ret = reformer_.load(container); + ivf_check_error_code(ret); + + size_t expect_size = header_.inverted_body_size; + inverted_ = load_segment(IVF_INVERTED_BODY_SEG_ID, expect_size); + if (!inverted_) { + LOG_ERROR("Failed to load segment, inverted_size=%zu block_count=%u", + static_cast(header_.inverted_body_size), + header_.block_count); + return IndexError_InvalidFormat; + } + + expect_size = header_.inverted_list_count * sizeof(InvertedListMeta); + inverted_meta_ = load_segment(IVF_INVERTED_META_SEG_ID, expect_size); + if (!inverted_meta_) { + LOG_ERROR("Failed to load segment, inverted_lists=%u", + header_.inverted_list_count); + return IndexError_InvalidFormat; + } + + expect_size = header_.total_vector_count * sizeof(uint64_t); + keys_ = load_segment(IVF_KEYS_SEG_ID, expect_size); + if (!keys_) { + return IndexError_InvalidFormat; + } + + expect_size = header_.total_vector_count * sizeof(InvertedVecLocation); + offsets_ = load_segment(IVF_OFFSETS_SEG_ID, expect_size); + if (!offsets_) { + return IndexError_InvalidFormat; + } + + expect_size = header_.total_vector_count * sizeof(uint32_t); + mapping_ = load_segment(IVF_MAPPING_SEG_ID, expect_size); + if (!mapping_) { + return IndexError_InvalidFormat; + } + + norm_value_sqrt_ = + meta_.metric_name() == "Euclidean" || meta_.metric_name() == "Manhattan"; + if (container_->get(IVF_INT8_QUANTIZED_PARAMS_SEG_ID) || + container->get(IVF_INT4_QUANTIZED_PARAMS_SEG_ID)) { + expect_size = + header_.inverted_list_count * sizeof(InvertedIntegerQuantizerParams); + auto &seg_id = meta_.reformer_name() == kInt8ReformerName + ? IVF_INT8_QUANTIZED_PARAMS_SEG_ID + : IVF_INT4_QUANTIZED_PARAMS_SEG_ID; + integer_quantizer_params_ = load_segment(seg_id, expect_size); + if (!integer_quantizer_params_) { + return IndexError_InvalidFormat; + } + norm_value_ = 0.0f; + } else if (meta_.reformer_name() == kInt8ReformerName || + meta_.reformer_name() == kInt4ReformerName) { + auto &scale_key = meta_.reformer_name() == kInt8ReformerName + ? INT8_QUANTIZER_REFORMER_SCALE + : INT4_QUANTIZER_REFORMER_SCALE; + auto scale = meta_.reformer_params().get_as_float(scale_key); + norm_value_ = this->convert_to_normalize_value(scale); + } else { + norm_value_ = 1.0f; + } + + if (container_->get(IVF_FEATURES_SEG_ID)) { + features_ = load_segment(IVF_FEATURES_SEG_ID, 0); + if (!features_) { + return IndexError_InvalidFormat; + } + if (features_->data_size() % vector_count() != 0) { + LOG_ERROR("Invalid featureSegment size=%zu, totalVecs=%zu", + features_->data_size(), vector_count()); + return IndexError_InvalidFormat; + } + } + + LOG_DEBUG( + "Load inverted index done, docs=%u invertedListCnt=%u " + "elementSize=%u metric=%s reformer=%s", + header_.total_vector_count, header_.inverted_list_count, + meta_.element_size(), meta_.metric_name().c_str(), + meta_.reformer_name().c_str()); + return 0; +} + +int IVFEntity::search(size_t inverted_list_id, const void *query, + const IndexFilter &filter, uint32_t *scan_count, + IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const { + ailego_assert_with(inverted_list_id < header_.inverted_list_count, + "invalid id"); + auto list_meta = this->inverted_list_meta(inverted_list_id); + ivf_assert(list_meta, IndexError_ReadData); + + const void *data = nullptr; + const size_t block_vecs = header_.block_vector_count; + std::vector distances(block_vecs); + const size_t batch_size = kBatchBlocks; + const size_t block_size = header_.block_size; + const auto norm_val = this->inverted_list_normalize_value(inverted_list_id); + for (size_t i = 0; i < list_meta->block_count; i += batch_size) { + //! Read vecs + const size_t off = list_meta->offset + i * block_size; + const size_t blocks = std::min(batch_size, list_meta->block_count - i); + const size_t size = + std::min(blocks * block_size, + static_cast(header_.inverted_body_size - off)); + if (inverted_->read(off, &data, size) != size) { + LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); + return IndexError_ReadData; + } + + //! Read keys + size_t items = std::min(blocks * block_vecs, + list_meta->vector_count - (i * block_vecs)); + auto keys = get_keys(list_meta->id_offset + i * block_vecs, items); + if (!keys) { + return IndexError_ReadData; + } + + //! Compute distances for each block + for (size_t b = 0; b < blocks; ++b) { + const size_t vecs_count = + std::min(block_vecs, list_meta->vector_count - (i + b) * block_vecs); + auto block_keys = keys + b * block_vecs; + size_t keeps = 0; + ailego_assert_with(block_vecs < sizeof(keeps) * 8, "bits overflow"); + for (size_t k = 0; k < vecs_count; ++k) { + if (!filter(block_keys[k])) { + keeps |= (1 << k); + } else { + ++(*context_stats->mutable_filtered_count()); + } + } + if (keeps == 0) { + continue; + } + + const void *block_data = static_cast(data) + b * block_size; + calculator_->query_features_distance(query, block_data, vecs_count, + distances.data()); + + *(context_stats->mutable_dist_calced_count()) += vecs_count; + + uint32_t id_off = list_meta->id_offset + (i + b) * block_vecs; + for (size_t k = 0; k < vecs_count; ++k) { + if (keeps & (1 << k)) { + if (block_keys[k] != kInvalidKey) { + heap->emplace(block_keys[k], distances[k] * norm_val, id_off + k); + } + } + } + } + } + + *scan_count = list_meta->vector_count; + return 0; +} + +//! search in inverted list without filter +int IVFEntity::search(size_t inverted_list_id, const void *query, + uint32_t *scan_count, IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const { + ailego_assert_with(inverted_list_id < header_.inverted_list_count, + "invalid id"); + auto list_meta = inverted_list_meta(inverted_list_id); + ivf_assert(list_meta, IndexError_ReadData); + + const void *data = nullptr; + const size_t block_vecs = header_.block_vector_count; + std::vector distances(block_vecs); + const size_t batch_size = kBatchBlocks; + const size_t block_size = header_.block_size; + const auto norm_val = this->inverted_list_normalize_value(inverted_list_id); + for (size_t i = 0; i < list_meta->block_count; i += batch_size) { + //! Read vecs + const size_t off = list_meta->offset + i * block_size; + const size_t blocks = std::min(batch_size, list_meta->block_count - i); + const size_t size = + std::min(blocks * block_size, + static_cast(header_.inverted_body_size - off)); + if (inverted_->read(off, &data, size) != size) { + LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); + return IndexError_ReadData; + } + + //! Read keys + size_t items = std::min(blocks * block_vecs, + list_meta->vector_count - (i * block_vecs)); + auto keys = get_keys(list_meta->id_offset + i * block_vecs, items); + if (!keys) { + return IndexError_ReadData; + } + + //! Compute distances for each block + for (size_t b = 0; b < blocks; ++b) { + const size_t vecs_count = + std::min(block_vecs, list_meta->vector_count - (i + b) * block_vecs); + auto block_keys = keys + b * block_vecs; + const void *block_data = static_cast(data) + b * block_size; + calculator_->query_features_distance(query, block_data, vecs_count, + distances.data()); + for (size_t k = 0; k < vecs_count; ++k) { + if (block_keys[k] != kInvalidKey) { + uint32_t id = list_meta->id_offset + (i + b) * block_vecs + k; + heap->emplace(block_keys[k], distances[k] * norm_val, id); + } + } + *(context_stats->mutable_dist_calced_count()) += vecs_count; + } + } + + *scan_count = list_meta->vector_count; + return 0; +} + +//! search all inverted list with filter +int IVFEntity::search(const void *query, const IndexFilter &filter, + IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const { + for (size_t i = 0; i < header_.inverted_list_count; ++i) { + uint32_t scan_count; + int ret = this->search(i, query, filter, &scan_count, heap, context_stats); + if (ret != 0) { + return ret; + } + } + + return 0; +} + +//! search all inverted list without filter +int IVFEntity::search(const void *query, IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const { + for (size_t i = 0; i < header_.inverted_list_count; ++i) { + uint32_t scan_count; + int ret = this->search(i, query, &scan_count, heap, context_stats); + if (ret != 0) { + return ret; + } + } + + return 0; +} + +const void *IVFEntity::get_vector(size_t id) const { + if (features_) { + const void *data = nullptr; + size_t element_size = features_->data_size() / vector_count(); + size_t off = id * element_size; + if (features_->read(off, &data, element_size) != element_size) { + LOG_ERROR("Failed to read segment, off=%zu size=%zu", off, element_size); + return nullptr; + } + return data; + } + + const void *data = nullptr; + size_t size = sizeof(InvertedVecLocation); + if (offsets_->read(id * size, &data, size) != size) { + LOG_ERROR("Failed to read offsets segment, id=%zu", id); + return nullptr; + } + auto &loc = *reinterpret_cast(data); + if (loc.column_major) { + vector_.resize(meta_.element_size()); + auto unit_size = IndexMeta::AlignSizeof(meta_.data_type()); + size_t cols = meta_.element_size() / unit_size; + size_t step = block_vector_count() * unit_size; + size_t rd_size = step * (cols - 1) + unit_size; + if (inverted_->read(loc.offset, &data, rd_size) != rd_size) { + LOG_ERROR("Failed to read data, off=%zu size=%zu", + static_cast(loc.offset), rd_size); + return nullptr; + } + for (size_t c = 0; c < cols; ++c) { + vector_.replace(c * unit_size, unit_size, + reinterpret_cast(data) + c * step, + unit_size); + } + return vector_.data(); + } else { + if (inverted_->read(loc.offset, &data, meta_.element_size()) != + meta_.element_size()) { + LOG_ERROR("Failed to read data, off=%zu size=%u", + static_cast(loc.offset), meta_.element_size()); + return nullptr; + } + return data; + } +} + +int IVFEntity::get_vector(size_t id, IndexStorage::MemoryBlock &block) const { + if (features_) { + size_t element_size = features_->data_size() / vector_count(); + size_t off = id * element_size; + if (features_->read(off, block, element_size) != element_size) { + LOG_ERROR("Failed to read segment, off=%zu size=%zu", off, element_size); + return IndexError_Runtime; + } + return 0; + } + + + IndexStorage::MemoryBlock data_block; + size_t size = sizeof(InvertedVecLocation); + if (offsets_->read(id * size, data_block, size) != size) { + LOG_ERROR("Failed to read offsets segment, id=%zu", id); + return IndexError_Runtime; + } + const void *data = data_block.data(); + auto &loc = *reinterpret_cast(data); + if (loc.column_major) { + vector_.resize(meta_.element_size()); + auto unit_size = IndexMeta::AlignSizeof(meta_.data_type()); + size_t cols = meta_.element_size() / unit_size; + size_t step = block_vector_count() * unit_size; + size_t rd_size = step * (cols - 1) + unit_size; + if (inverted_->read(loc.offset, &data, rd_size) != rd_size) { + LOG_ERROR("Failed to read data, off=%zu size=%zu", + static_cast(loc.offset), rd_size); + return IndexError_Runtime; + } + for (size_t c = 0; c < cols; ++c) { + vector_.replace(c * unit_size, unit_size, + reinterpret_cast(data) + c * step, + unit_size); + } + block.reset(vector_.data()); + return 0; + } else { + if (inverted_->read(loc.offset, block, meta_.element_size()) != + meta_.element_size()) { + LOG_ERROR("Failed to read data, off=%zu size=%u", + static_cast(loc.offset), meta_.element_size()); + return IndexError_Runtime; + } + return 0; + } +} + +uint32_t IVFEntity::key_to_id(uint64_t key) const { + //! Do binary search + uint32_t start = 0UL; + uint32_t end = vector_count(); + const void *data = nullptr; + uint32_t idx = 0u; + while (start < end) { + idx = start + (end - start) / 2; + if (ailego_unlikely(mapping_->read(idx * sizeof(uint32_t), &data, + sizeof(uint32_t)) != sizeof(uint32_t))) { + LOG_ERROR("Failed to read mapping segment, idx=%u", idx); + return std::numeric_limits::max(); + } + const uint64_t *mkey; + uint32_t local_id = *reinterpret_cast(data); + if (ailego_unlikely(keys_->read(local_id * sizeof(uint64_t), + (const void **)(&mkey), + sizeof(uint64_t)) != sizeof(uint64_t))) { + LOG_ERROR("Read key from segment failed"); + return std::numeric_limits::max(); + } + if (*mkey < key) { + start = idx + 1; + } else if (*mkey > key) { + end = idx; + } else { + return local_id; + } + } + return std::numeric_limits::max(); +} + +const void *IVFEntity::get_vector_by_key(uint64_t key) const { + uint32_t id = this->key_to_id(key); + if (id != std::numeric_limits::max()) { + return get_vector(id); + } else { + return nullptr; + } +} + +int IVFEntity::get_vector_by_key(uint64_t key, + IndexStorage::MemoryBlock &block) const { + uint32_t id = this->key_to_id(key); + if (id != std::numeric_limits::max()) { + return get_vector(id, block); + } else { + return IndexError_Runtime; + } +} + +IVFEntity::Pointer IVFEntity::clone(void) const { + auto entity = std::make_shared(); + return clone(entity); +} + +IVFEntity::Pointer IVFEntity::clone(const IVFEntity::Pointer &entity) const { + if (!entity) { + LOG_ERROR("Failed to alloc IVFEntity"); + return nullptr; + } + + auto inverted = inverted_->clone(); + ivf_assert_with_msg(inverted, nullptr, "Failed to clone inverted segment"); + + auto inverted_meta = inverted_meta_->clone(); + ivf_assert_with_msg(inverted_meta, nullptr, + "Failed to clone inverted meta segment"); + + auto keys = keys_->clone(); + ivf_assert_with_msg(keys, nullptr, "Failed to clone keys segment"); + + auto offsets = offsets_->clone(); + ivf_assert_with_msg(offsets, nullptr, "Failed to clone offsets segment"); + + auto mapping = mapping_->clone(); + ivf_assert_with_msg(mapping, nullptr, "Failed to clone mapping segment"); + + IndexStorage::Segment::Pointer integer_quantizer_params; + if (integer_quantizer_params_) { + integer_quantizer_params = integer_quantizer_params_->clone(); + if (!integer_quantizer_params) { + LOG_ERROR("Failed to clone integer quantizer params segment"); + return nullptr; + } + } + IndexStorage::Segment::Pointer features; + if (features_) { + features = features_->clone(); + if (!features) { + LOG_ERROR("Failed to clone features segment"); + return nullptr; + } + } + + entity->meta_ = this->meta_; + entity->reformer_ = this->reformer_; + entity->calculator_ = this->calculator_; + entity->header_ = this->header_; + entity->container_ = this->container_; + + entity->inverted_ = inverted; + entity->inverted_meta_ = inverted_meta; + entity->keys_ = keys; + entity->offsets_ = offsets; + entity->mapping_ = mapping; + entity->integer_quantizer_params_ = integer_quantizer_params; + entity->features_ = features; + entity->norm_value_ = this->norm_value_; + entity->norm_value_sqrt_ = this->norm_value_sqrt_; + + return entity; +} + +IndexStorage::Segment::Pointer IVFEntity::load_segment( + const std::string &seg_id, size_t expect_size) const { + auto segment = container_->get(seg_id); + if (!segment) { + LOG_ERROR("Failed to get segment %s", seg_id.c_str()); + return nullptr; + } + if (expect_size && segment->data_size() != expect_size) { + LOG_ERROR("Invalid segment %s size=%zu, total_vecs=%u", seg_id.c_str(), + segment->data_size(), header_.total_vector_count); + return nullptr; + } + return segment; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_entity.h b/src/core/algorithm/ivf/ivf_entity.h new file mode 100644 index 0000000..e0265b6 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_entity.h @@ -0,0 +1,367 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "metric/metric_params.h" +#include "ivf_distance_calculator.h" +#include "ivf_index_format.h" +#include "ivf_params.h" + +namespace zvec { +namespace core { + +/*! IVF Entity + */ +class IVFEntity { + public: + typedef std::shared_ptr Pointer; + + class IVFReformerWrapper; + + //! Constructor + IVFEntity() {} + + //! Destructor + virtual ~IVFEntity() {} + + //! Disable them + IVFEntity(const IVFEntity &) = delete; + IVFEntity &operator=(const IVFEntity &) = delete; + + //! load the index from container + virtual int load(const IndexStorage::Pointer &container); + + //! search in inverted list with filter + int search(size_t inverted_list_id, const void *query, + const IndexFilter &filter, uint32_t *scan_count, + IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const; + + //! search in inverted list without filter + int search(size_t inverted_list_id, const void *query, uint32_t *scan_count, + IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const; + + //! search all inverted list with filter + int search(const void *query, const IndexFilter &filter, + IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const; + + //! search all inverted list without filter + int search(const void *query, IndexDocumentHeap *heap, + IndexContext::Stats *context_stats) const; + + //! Clone the entity + virtual IVFEntity::Pointer clone(void) const; + + //! Clone the entity + IVFEntity::Pointer clone(const IVFEntity::Pointer &entity) const; + + //! Retrieve the primary keys by local id in heap + int retrieve_keys(IndexDocumentHeap *heap) const { + for (auto &it : (*heap)) { + uint64_t key = this->get_key(it.index()); + if (key == kInvalidKey) { + return IndexError_ReadData; + } + it.set_key(key); + } + + return 0; + } + + //! Retrieve the total vectors in the index + size_t vector_count(void) const { + return header_.total_vector_count; + } + + //! Retrieve the inverted list count + size_t inverted_list_count(void) const { + return header_.inverted_list_count; + } + + //! Retrieve block size of the inverted vector + size_t inverted_block_size(void) const { + return header_.block_size; + } + + //! Retrieve the vectors count in one block + size_t block_vector_count(void) const { + return header_.block_vector_count; + } + + //! Retrieve IndexMeta of the inverted index + const IndexMeta &meta(void) const { + return meta_; + } + + //! Retrieve a block of vectors + const void *read_block(size_t inverted_list_id, size_t local_block_id, + size_t *vecs_count) const { + auto iv_meta = this->inverted_list_meta(inverted_list_id); + if (!iv_meta || local_block_id >= iv_meta->block_count) { + LOG_ERROR("Failed to read inverted list, listId=%zu blockIdx=%zu", + inverted_list_id, local_block_id); + return nullptr; + } + + size_t block_vecs = header_.block_vector_count; + *vecs_count = std::min(block_vecs, + iv_meta->vector_count - local_block_id * block_vecs); + ailego_assert_with(*vecs_count <= header_.block_vector_count, + "invalid vecs"); + const size_t off = iv_meta->offset + local_block_id * header_.block_size; + const size_t size = *vecs_count * meta_.element_size(); + const void *data = nullptr; + if (inverted_->read(off, &data, size) != size) { + LOG_ERROR("Failed to read block off=%zu size=%zu", off, size); + return nullptr; + } + + return data; + } + + //! Retrieve the inverted list meta + const InvertedListMeta *inverted_list_meta(size_t inverted_list_id) const { + const void *data = nullptr; + const size_t size = sizeof(InvertedListMeta); + const size_t offset = inverted_list_id * size; + if (inverted_meta_->read(offset, &data, size) != size) { + LOG_ERROR("Failed to read inverted meta, id=%zu, size=%zu", + inverted_list_id, size); + return nullptr; + } + + return static_cast(data); + } + + //! Retrieve the keys by consecutive local ids + const uint64_t *get_keys(size_t id, size_t count) const { + const void *data = nullptr; + const size_t offset = id * sizeof(uint64_t); + const size_t size = count * sizeof(uint64_t); + if (keys_->read(offset, &data, size) != size) { + LOG_ERROR("Failed to read keys, id=%zu, size=%zu", id, size); + return nullptr; + } + + return static_cast(data); + } + + //! Retrieve the key by local id + uint64_t get_key(size_t id) const { + const void *data = nullptr; + const size_t offset = id * sizeof(uint64_t); + const size_t size = sizeof(uint64_t); + if (keys_->read(offset, &data, size) != size) { + LOG_ERROR("Failed to read key, id=%zu", id); + return kInvalidKey; + } + + return *static_cast(data); + } + + //! Retrieve the key-order mapping (sorted rank -> local_id). + //! mapping[rank] is the local_id of the vector with the rank-th smallest + //! key. Returns nullptr if mapping segment is unavailable. + const uint32_t *get_key_order_mapping() const { + if (!mapping_) return nullptr; + const void *data = nullptr; + const size_t size = vector_count() * sizeof(uint32_t); + if (mapping_->read(0, &data, size) != size) { + return nullptr; + } + return static_cast(data); + } + + //! Retrieve vector by local id + const void *get_vector(size_t id) const; + + //! Retrieve vector by local id + const void *get_vector_by_key(uint64_t key) const; + + int get_vector(size_t id, IndexStorage::MemoryBlock &block) const; + + int get_vector_by_key(uint64_t key, IndexStorage::MemoryBlock &block) const; + + uint32_t key_to_id(uint64_t key) const; + + //! Transform a query + int transform(const void *query, const IndexQueryMeta &qmeta, + const void **out, IndexQueryMeta *ometa) const { + return reformer_.transform(query, qmeta, out, ometa); + } + + //! Transform queries + int transform(const void *query, const IndexQueryMeta &qmeta, uint32_t count, + const void **out, IndexQueryMeta *ometa) const { + return reformer_.transform(query, qmeta, count, out, ometa); + } + + //! Normalize the score in query part + void normalize(size_t qidx, IndexDocumentHeap *heap) const { + return reformer_.normalize(qidx, heap); + } + + //! Retrieve the value for each inverted list to multiply for normalizing + float inverted_list_normalize_value(size_t inverted_list_id) const { + if (norm_value_ != 0.0f) { + return norm_value_; + } + + // ailego_assert_with(integer_quantizer_params_, "nullptr"); + if (integer_quantizer_params_ != nullptr) { + const void *data = nullptr; + size_t size = sizeof(InvertedIntegerQuantizerParams); + size_t off = inverted_list_id * size; + if (integer_quantizer_params_->read(off, &data, size) != size) { + LOG_ERROR("Failed to read data from segment, off=%zu", off); + return 1.0f; + } + auto scale = + static_cast(data)->scale; + return this->convert_to_normalize_value(scale); + } + + return norm_value_; + } + + //! Check whether the feature segment exist + bool has_orignal_feature() const { + return !!features_; + } + + //! Retrieve reformer + const IVFReformerWrapper &reformer(void) const { + return reformer_; + } + + /*! Index Reformer Wrapper + * To transform query in inverted index searching, and normalize the score + */ + class IVFReformerWrapper { + public: + //! Constructor + IVFReformerWrapper() {} + + //! Assignment + IVFReformerWrapper &operator=(const IVFReformerWrapper &wrapper) { + reformer_ = wrapper.reformer_; + type_ = wrapper.type_; + buffer_.clear(); + buffer_.shrink_to_fit(); + reciprocal_ = wrapper.reciprocal_; + return *this; + } + + //! Initialize + int init(const IndexMeta &imeta); + + //! Load reformer state (e.g. rotation matrix) from storage + int load(const IndexStorage::Pointer &storage); + + //! Update + int update(const IndexMeta &meta); + + //! Transform a query + int transform(const void *query, const IndexQueryMeta &qmeta, + const void **out, IndexQueryMeta *ometa); + + //! Transform queries + int transform(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, const void **out, IndexQueryMeta *ometa); + + //! Convert a record + virtual int convert(const void *record, const IndexQueryMeta &rmeta, + const void **out, IndexQueryMeta *ometa); + + //! Convert records + virtual int convert(const void *records, const IndexQueryMeta &rmeta, + uint32_t count, const void **out, + IndexQueryMeta *ometa); + + //! Transform queries + int transform_gpu(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, const void **out, IndexQueryMeta *ometa); + + //! Normalize the score in query part + void normalize(size_t qidx, IndexDocumentHeap *heap) const; + + //! Normalize the score in query part + void normalize(size_t qidx, const void *query, const IndexQueryMeta &qmeta, + IndexDocumentHeap *heap) const; + + private: + //! Transform query from fp32 to int8 + void transform(size_t qidx, const float *in, size_t dim, int8_t *out); + + //! Transform query from fp32 to int4 + void transform(size_t qidx, const float *in, size_t dim, uint8_t *out); + + private: + //! Constants + enum Type { + kReformerTpNone = 0, + kReformerTpInnerProductInt8 = 1, + kReformerTpInnerProductInt4 = 2, + kReformerTpInt8 = 3, + kReformerTpInt4 = 4, + kReformerTpDefault = 7, + }; + + //! Members + Type type_{kReformerTpNone}; + IndexReformer::Pointer reformer_{}; + std::string buffer_{}; + float reciprocal_{0.0}; // for int8 + std::vector scales_{}; // for int8 IP + }; + + private: + //! Load the segment by seg_id in expect_size segment size + IndexStorage::Segment::Pointer load_segment(const std::string &seg_id, + size_t expect_size) const; + + //! Load the header segment + int load_header(const IndexStorage::Pointer &container); + + //! Convert the int8 quantizer scale to normalize value + float convert_to_normalize_value(float scale) const { + auto v = scale == 0.0 ? 1.0 : (1.0 / scale); + return !norm_value_sqrt_ ? v : std::sqrt(v); + } + + protected: + //! Constants + static constexpr size_t kBatchBlocks = 10u; + + //! Members + IndexMeta meta_{}; + mutable IVFReformerWrapper reformer_{}; + IVFDistanceCalculator::Pointer calculator_{}; + IndexStorage::Pointer container_{}; + IndexStorage::Segment::Pointer inverted_{}; + IndexStorage::Segment::Pointer inverted_meta_{}; + IndexStorage::Segment::Pointer keys_{}; + IndexStorage::Segment::Pointer offsets_{}; + IndexStorage::Segment::Pointer mapping_{}; + IndexStorage::Segment::Pointer features_{}; + IndexStorage::Segment::Pointer integer_quantizer_params_{}; + mutable std::string vector_{}; // temporary buffer for colomn major order + float norm_value_{0.0f}; // normalize the inverted vector to orignal score + bool norm_value_sqrt_{false}; // does the norm value need to sqrt + InvertedIndexHeader header_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_index_format.h b/src/core/algorithm/ivf/ivf_index_format.h new file mode 100644 index 0000000..e7df0cc --- /dev/null +++ b/src/core/algorithm/ivf/ivf_index_format.h @@ -0,0 +1,174 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include + +namespace zvec { +namespace core { + +using node_id_t = uint32_t; +using key_t = uint64_t; + +static constexpr uint64_t kInvalidKey = std::numeric_limits::max(); + +/*! Index Format of Inverted Index Header + */ +struct InvertedIndexHeader { + uint32_t header_size{0}; + uint32_t total_vector_count{0}; + uint64_t inverted_body_size{0}; + uint32_t inverted_list_count{0}; + uint32_t block_vector_count{0}; + uint32_t block_size{0}; + uint32_t block_count{0}; + uint32_t index_meta_size{0}; + char reserved_[28]; +#ifdef _MSC_VER + char index_meta[]; +#else + char index_meta[0]; +#endif +}; + +/*! Index Format of Inverted Index Meta for each Inverted list + */ +struct InvertedListMeta { + uint64_t offset{0}; + uint32_t block_count{0}; + uint32_t vector_count{0}; + uint32_t id_offset{0}; + char reserved_[16]; +}; + +/*! Index Format of Location in Inverted Index for each vector + */ +struct InvertedVecLocation { + InvertedVecLocation(size_t off, bool col) + : offset(off), column_major(col), reserved(0u) {} + + uint64_t offset : 48; // feature offset in posting block segment + uint64_t column_major : 1; // coloum major if true + uint64_t reserved : 15; +}; + +/*! Index Format of Integer Quantizer params for each inverted list + */ +struct InvertedIntegerQuantizerParams { + float scale{1.0}; + float bias{0.0}; +}; + +/*! Location of Vectors Block in Storage Segment + */ +struct BlockLocation { + uint16_t segment_id; + uint16_t block_index; +}; + +/*! The Header of a Block in Storage Segment + */ +struct BlockHeader { + BlockLocation next; + uint16_t vector_count; + uint16_t column_major : 1; + uint16_t reserved_ : 15; +}; + +struct DeletionMap { + void set(uint32_t index) { + bitset.set(index); + } + + void reset(uint32_t index) { + bitset.reset(index); + } + + bool test(uint32_t index) const { + return bitset.test(index); + } + + bool is_dirty() const { + return bitset.test_any(); + } + + ailego::FixedBitset<32> bitset{}; +}; + +static_assert(sizeof(DeletionMap) == 4, "DeletionMap must be 4 bytes"); + +/*! Meta Information of Streamer Entity + */ +struct StreamerInvertedMeta { + uint64_t create_time{0}; + uint64_t update_time{0}; + uint64_t revision_id{0}; + uint32_t segment_count{0}; + uint32_t segment_size{0}; + uint8_t reserved_[32]; + InvertedIndexHeader header; +}; + +/*! Location of Vector in Storage Segment + */ +struct VectorLocation { + //! Constructor + VectorLocation(void) {} + + //! Constructor + VectorLocation(uint16_t id, bool col, uint32_t off) + : segment_id(id), column_major(col), offset(off) {} + + uint16_t segment_id; + uint16_t column_major : 1; + uint16_t reserved_ : 15; + uint32_t offset; + + public: + bool operator==(const VectorLocation &other) const { + return segment_id == other.segment_id && + column_major == other.column_major && offset == other.offset; + } +}; + +static_assert(sizeof(VectorLocation) == sizeof(uint64_t), + "VectorLocation must be size of 8 bytes"); + +struct KeyInfo { + KeyInfo(void) {} + KeyInfo(uint32_t idx, const VectorLocation &loc) + : centroid_idx(idx), location(loc) {} + KeyInfo(VectorLocation loc) : location(loc) {} + uint32_t centroid_idx; + VectorLocation location; +}; + +// Segments ID +const std::string IVF_CENTROID_SEG_ID("ivf.centroid"); +const std::string IVF_INVERTED_BODY_SEG_ID("ivf.inverted_body"); +const std::string IVF_INVERTED_HEADER_SEG_ID("ivf.inverted_header"); +const std::string IVF_INVERTED_META_SEG_ID("ivf.inverted_meta"); +const std::string IVF_KEYS_SEG_ID("hc.keys"); +const std::string IVF_OFFSETS_SEG_ID("ivf.offsets"); +const std::string IVF_MAPPING_SEG_ID("ivf.mapping"); +const std::string IVF_FEATURES_SEG_ID("ivf.features"); +const std::string IVF_INT8_QUANTIZED_PARAMS_SEG_ID("ivf.int8_quantized_params"); +const std::string IVF_INT4_QUANTIZED_PARAMS_SEG_ID("ivf.int4_quantized_params"); + +const std::string IVF_INVERTED_LIST_HEAD_SEG_ID("ivf.inverted_list_head"); +const std::string IVF_STORAGE_SEGMENT_ID("ivf.S"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_index_provider.h b/src/core/algorithm/ivf/ivf_index_provider.h new file mode 100644 index 0000000..1b9e9fe --- /dev/null +++ b/src/core/algorithm/ivf/ivf_index_provider.h @@ -0,0 +1,162 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include "ivf_entity.h" + +namespace zvec { +namespace core { + +/*! IVF IndexProvider + */ +class IVFIndexProvider : public IndexProvider { + public: + IVFIndexProvider(const IndexMeta &meta, const IVFEntity::Pointer &entity, + const std::string &owner) + : meta_(meta), entity_(entity), owner_class_(owner) {} + + IVFIndexProvider(const IVFIndexProvider &) = delete; + IVFIndexProvider &operator=(const IVFIndexProvider &) = delete; + + public: + //! Create a new iterator + Iterator::Pointer create_iterator(void) override { + return Iterator::Pointer(new (std::nothrow) SortedIterator(entity_)); + } + + //! Retrieve count of vectors + size_t count(void) const override { + return entity_->vector_count(); + } + + //! Retrieve dimension of vector + size_t dimension(void) const override { + return meta_.dimension(); + } + + //! Retrieve type of vector + IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + //! Retrieve vector size in bytes + size_t element_size(void) const override { + return meta_.element_size(); + } + + //! Retrieve a vector using a primary key + const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + //! Retrieve the owner class + const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class SortedIterator : public IndexProvider::Iterator { + public: + SortedIterator(const IVFEntity::Pointer &entity) : entity_(entity) { + count_ = entity_->vector_count(); + mapping_ = entity_->get_key_order_mapping(); + if (!mapping_) { + // Fallback: compute sorting if mapping segment is unavailable + fallback_.resize(count_); + std::iota(fallback_.begin(), fallback_.end(), size_t(0)); + std::sort(fallback_.begin(), fallback_.end(), [&](size_t a, size_t b) { + return entity_->get_key(a) < entity_->get_key(b); + }); + } + } + + //! Retrieve pointer of data + //! NOTICE: the vec feature will be changed after iterating to next, so + //! the caller need to keep a copy of it before iterator to next vector + const void *data(void) const override { + return entity_->get_vector(current_local_id()); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return pos_ < count_; + } + + //! Retrieve primary key + uint64_t key(void) const override { + return entity_->get_key(current_local_id()); + } + + //! Next iterator + void next(void) override { + ++pos_; + } + + private: + size_t current_local_id() const { + return mapping_ ? static_cast(mapping_[pos_]) : fallback_[pos_]; + } + + //! Members + IVFEntity::Pointer entity_; + const uint32_t *mapping_{nullptr}; // points into mapping_ segment data + std::vector fallback_; // used only if mapping_ unavailable + size_t count_{0}; + size_t pos_{0}; + }; + + //! Original sequential iterator (kept for potential internal use) + class Iterator : public IndexProvider::Iterator { + public: + Iterator(const IVFEntity::Pointer &entity) : entity_(entity) {} + + //! Retrieve pointer of data + const void *data(void) const override { + return entity_->get_vector(index_); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return index_ < entity_->vector_count(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return entity_->get_key(index_); + } + + //! Next iterator + void next(void) override { + ++index_; + } + + private: + //! Members + IVFEntity::Pointer entity_; + size_t index_{0}; + }; + + private: + //! Members + const IndexMeta &meta_; + IVFEntity::Pointer entity_; + std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_params.h b/src/core/algorithm/ivf/ivf_params.h new file mode 100644 index 0000000..ac8a70f --- /dev/null +++ b/src/core/algorithm/ivf/ivf_params.h @@ -0,0 +1,90 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +static const std::string SEPARATOR("/"); +static const std::string CENTROID_SEPERATOR = "*"; + +// builder params +static const std::string PARAM_IVF_BUILDER_CENTROID_COUNT( + "proxima.ivf.builder.centroid_count"); +static const std::string PARAM_IVF_BUILDER_CLUSTER_CLASS( + "proxima.ivf.builder.cluster_class"); +static const std::string PARAM_IVF_BUILDER_THREAD_COUNT( + "proxima.ivf.builder.thread_count"); +static const std::string PARAM_IVF_BUILDER_CLUSTER_AUTO_TUNING( + "proxima.ivf.builder.cluster_auto_tuning"); +static const std::string PARAM_IVF_BUILDER_TRAIN_SAMPLE_COUNT( + "proxima.ivf.builder.train_sample_count"); +static const std::string PARAM_IVF_BUILDER_TRAIN_SAMPLE_RATIO( + "proxima.ivf.builder.train_sample_ratio"); +static const std::string PARAM_IVF_BUILDER_CONVERTER_PARAMS( + "proxima.ivf.builder.converter_params"); +static const std::string PARAM_IVF_BUILDER_CONVERTER_CLASS( + "proxima.ivf.builder.converter_class"); +static const std::string PARAM_IVF_BUILDER_STORE_ORIGINAL_FEATURES( + "proxima.ivf.builder.store_original_features"); +static const std::string PARAM_IVF_BUILDER_QUANTIZER_CLASS( + "proxima.ivf.builder.quantizer_class"); +static const std::string PARAM_IVF_BUILDER_QUANTIZE_BY_CENTROID( + "proxima.ivf.builder.quantize_by_centroid"); +static const std::string PARAM_IVF_BUILDER_QUANTIZER_PARAMS( + "proxima.ivf.builder.quantizer_params"); +static const std::string PARAM_IVF_BUILDER_CLUSTER_PARAMS_IN_LEVEL_PREFIX( + "proxima.ivf.builder.cluster_params_in_level_"); +static const std::string PARAM_IVF_BUILDER_OPTIMIZER_CLASS( + "proxima.ivf.builder.optimizer_class"); +static const std::string PARAM_IVF_BUILDER_OPTIMIZER_PARAMS( + "proxima.ivf.builder.optimizer_params"); +static const std::string PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_CLASS( + "proxima.ivf.builder.optimizer_quantizer_class"); +static const std::string PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_PARAMS( + "proxima.ivf.builder.optimizer_quantizer_params"); +static const std::string PARAM_IVF_BUILDER_BLOCK_VECTOR_COUNT( + "proxima.ivf.builder.block_vector_count"); + +// searcher params +static const std::string PARAM_IVF_SEARCHER_SCAN_RATIO( + "proxima.ivf.searcher.scan_ratio"); +static const std::string PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD( + "proxima.ivf.searcher.brute_force_threshold"); +static const std::string PARAM_IVF_SEARCHER_OPTIMIZER( + "proxima.ivf.searcher.optimizer"); +static const std::string PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS( + "proxima.ivf.searcher.optimizer_params"); +static const std::string PARAM_IVF_SEARCHER_CONVERTER_REFORMER( + "proxima.ivf.searcher.converter_reformer"); +static const std::string PARAM_IVF_SEARCHER_NPROBE( + "proxima.ivf.searcher.nprobe"); + +// Constants +static constexpr char const *kIPMetricName = "InnerProduct"; +static constexpr char const *kMipsMetricName = "MipsSquaredEuclidean"; +static constexpr char const *kL2MetricName = "SquaredEuclidean"; +static constexpr char const *kMipsConverterName = "MipsConverter"; +static constexpr char const *kMipsRevConverterName = "MipsReverseConverter"; +static constexpr char const *kMipsReformerName = "MipsReformer"; +static constexpr char const *kInt8QuantizerName = "Int8QuantizerConverter"; +static constexpr char const *kInt4QuantizerName = "Int4QuantizerConverter"; +static constexpr char const *kInt8ReformerName = "Int8QuantizerReformer"; +static constexpr char const *kInt4ReformerName = "Int4QuantizerReformer"; +static constexpr float kNormalizeScaleFactor = 16.0f; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_searcher.cc b/src/core/algorithm/ivf/ivf_searcher.cc new file mode 100644 index 0000000..e8cda56 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_searcher.cc @@ -0,0 +1,324 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "ivf_searcher.h" +#include +#include +#include "ivf_centroid_index.h" +#include "ivf_index_provider.h" +#include "ivf_params.h" + +namespace zvec { +namespace core { + +int IVFSearcher::init(const ailego::Params ¶meters) { + params_ = parameters; + + params_.get(PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD, &bruteforce_threshold_); + + searcher_state_ = STATE_INITED; + + return 0; +} + +int IVFSearcher::cleanup(void) { + this->unload(); + + params_.clear(); + bruteforce_threshold_ = kDefaultBfThreshold; + + searcher_state_ = STATE_INIT; + return 0; +} + +int IVFSearcher::load(IndexStorage::Pointer container, + IndexMetric::Pointer /*metric*/) { + if (!container) { + LOG_ERROR("Invalid container"); + return IndexError_InvalidArgument; + } + if (searcher_state_ != STATE_INITED) { + LOG_ERROR("Initalize the searcher first before load index"); + return IndexError_Runtime; + } + + ailego::ElapsedTime timer; + int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize meta from container"); + return ret; + } + + //! Load centroid index + centroid_index_ = std::make_shared(); + if (!centroid_index_) { + return IndexError_NoMemory; + } + auto seg = container->get(IVF_CENTROID_SEG_ID, 0); + if (!seg) { + LOG_ERROR("Failed to get segment %s", IVF_CENTROID_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + IndexStorage::Pointer seg_container = + std::make_shared(seg); + if (!seg_container) { + return IndexError_NoMemory; + } + ret = seg_container->open(std::string(), false); + if (ret != 0) { + LOG_ERROR("IndexSegmentStorage load failed for %s", IndexError::What(ret)); + return ret; + } + ret = centroid_index_->load(seg_container, params_); + if (ret != 0) { + LOG_ERROR("Failed to load index for %s", IndexError::What(ret)); + return ret; + } + + auto reformer = centroid_index_->reformer(); + if (reformer) { + //! The centroid index is loaded from the centroid sub-segment which does + //! not contain the rotator segment. Load the reformer state (e.g. rotation + //! matrix) from the main container instead. + ret = reformer->load(container); + ivf_check_error_code(ret); + } + params_.set(PARAM_IVF_SEARCHER_CONVERTER_REFORMER, reformer); + + //! load iverted index + entity_ = std::make_shared(); + if (!entity_) { + return IndexError_NoMemory; + } + ret = entity_->load(container); + ivf_check_error_code(ret); + + magic_ = IndexContext::GenerateMagic(); + + stats_.set_loaded_count(entity_->vector_count()); + stats_.set_loaded_costtime(timer.milli_seconds()); + + searcher_state_ = STATE_LOADED; + return 0; +} + +int IVFSearcher::unload(void) { + magic_ = 0; + centroid_index_.reset(); + entity_.reset(); + stats_.set_loaded_count(0UL); + stats_.set_loaded_costtime(0UL); + stats_.clear_attributes(); + searcher_state_ = STATE_INITED; + + return 0; +} + +int IVFSearcher::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_bf_impl(query, qmeta, 1, context); +} + +int IVFSearcher::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + if (!query || qmeta.element_size() != meta_.element_size()) { + LOG_ERROR("Null query or invalid qmeta"); + return IndexError_InvalidArgument; + } + IVFSearcherContext *ctx = dynamic_cast(context.get()); + if (!ctx || ctx->topk() == 0) { + LOG_ERROR("Invalid context or topk not set yet"); + return IndexError_InvalidArgument; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher + int ret = this->update_context(ctx); + ivf_check_error_code(ret); + } + + ctx->reset_results(count); + auto &entity = ctx->entity(); + auto &filter = ctx->filter(); + + //! Transform the querys for querying in inverted vector index later + IndexQueryMeta iv_qmeta; + int ret = entity->transform(query, qmeta, count, &query, &iv_qmeta); + ivf_check_with_msg(ret, "Failed to transform querys"); + + // TODO: do batch search in matrix + for (size_t q = 0; q < count; ++q) { + auto &context_stats = ctx->mutable_stats(q); + auto &heap = ctx->mutable_result_heap(); + heap.clear(); + if (!filter.is_valid()) { + ret = entity->search(query, &heap, &context_stats); + } else { + ret = entity->search(query, filter, &heap, &context_stats); + } + ivf_check_with_msg(ret, "Failed to search in entity for %s", + IndexError::What(ret)); + heap.sort(); // sort the results + if (!filter.is_valid()) { + // mapping the local id to key if query without filter + ret = entity->retrieve_keys(&heap); + ivf_check_error_code(ret); + } + entity->normalize(q, &heap); + ctx->topk_to_result(q); + query = static_cast(query) + iv_qmeta.element_size(); + } + + return 0; +} + +int IVFSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return this->search_impl(query, qmeta, 1, context); +} + +int IVFSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const { + if (entity_->vector_count() <= bruteforce_threshold_) { + return this->search_bf_impl(query, qmeta, count, context); + } + if (!query || qmeta.element_size() != meta_.element_size()) { + LOG_ERROR("Null query or invalid qmeta"); + return IndexError_InvalidArgument; + } + + IVFSearcherContext *ctx = dynamic_cast(context.get()); + if (!ctx || ctx->topk() == 0) { + LOG_ERROR("Invalid context or topk not set yet"); + return IndexError_InvalidArgument; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher + int ret = update_context(ctx); + ivf_check_error_code(ret); + } + + ctx->reset_results(count); + auto &entity = ctx->entity(); + auto &filter = ctx->filter(); + + auto ¢roid_index_ctx = ctx->centroid_searcher_ctx(); + int ret = centroid_index_->search(query, qmeta, count, centroid_index_ctx); + ivf_check_error_code(ret); + + //! Transform the querys for querying in inverted vector index later + IndexQueryMeta iv_qmeta; + ret = entity->transform(query, qmeta, count, &query, &iv_qmeta); + ivf_check_with_msg(ret, "Failed to transform querys"); + + for (size_t q = 0; q < count; ++q) { + auto ¢roids = centroid_index_ctx->result(q); + auto &context_stats = ctx->mutable_stats(q); + auto &heap = ctx->mutable_result_heap(); + heap.clear(); + size_t total_scan_count = 0; + for (size_t i = 0; + i < centroids.size() && total_scan_count < ctx->max_scan_count(); + ++i) { + auto cid = centroids[i].key(); + uint32_t scan_count = 0; + if (!filter.is_valid()) { + ret = entity->search(cid, query, &scan_count, &heap, &context_stats); + } else { + ret = entity->search(cid, query, filter, &scan_count, &heap, + &context_stats); + } + ivf_check_with_msg(ret, "Failed to search in entity for %s", + IndexError::What(ret)); + total_scan_count += scan_count; + } + heap.sort(); // sort the results + if (!filter.is_valid()) { + // mapping the local id to key if query without filter + ret = entity->retrieve_keys(&heap); + ivf_check_error_code(ret); + } + entity->normalize(q, &heap); + ctx->topk_to_result(q); + query = static_cast(query) + iv_qmeta.element_size(); + } + + return 0; +} + +const IndexSearcher::Stats &IVFSearcher::stats(void) const { + return stats_; +} + +IndexSearcher::Context::Pointer IVFSearcher::create_context() const { + if (searcher_state_ != STATE_LOADED) { + LOG_ERROR("Load the index first before create context"); + return nullptr; + } + + auto entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone IVFEntity"); + return nullptr; + } + + auto centroid_index_ctx = centroid_index_->create_context(); + if (!centroid_index_ctx) { + LOG_ERROR("Failed to create centroid index context"); + return nullptr; + } + + auto context = + new (std::nothrow) IVFSearcherContext(entity, centroid_index_ctx); + if (!context) { + LOG_ERROR("Failed to alloc IVFSearcherContext"); + return nullptr; + } + int ret = context->init(params_); + if (ret != 0) { + delete context; + return nullptr; + } + + context->set_magic(magic_); + + return Context::Pointer(context); +} + +IndexProvider::Pointer IVFSearcher::create_provider(void) const { + if (searcher_state_ != STATE_LOADED) { + LOG_ERROR("Load the index first before create provider"); + return nullptr; + } + + auto entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone IVFEntity"); + return Provider::Pointer(); + } + + auto *provider = new (std::nothrow) + IVFIndexProvider(entity->has_orignal_feature() ? meta_ : entity->meta(), + entity, "IVFSearcher"); + if (!provider) { + LOG_ERROR("Failed to alloc IVFIndexProvider"); + return Provider::Pointer(); + } + + return Provider::Pointer(provider); +} + +INDEX_FACTORY_REGISTER_SEARCHER(IVFSearcher); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_searcher.h b/src/core/algorithm/ivf/ivf_searcher.h new file mode 100644 index 0000000..dbe0d79 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_searcher.h @@ -0,0 +1,112 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "ivf_centroid_index.h" +#include "ivf_entity.h" +#include "ivf_searcher_context.h" + +namespace zvec { +namespace core { + +/*! IVF Searcher + */ +class IVFSearcher : public IndexSearcher { + public: + //! Initialize Searcher + int init(const ailego::Params ¶meters) override; + + //! Cleanup Searcher + int cleanup(void) override; + + //! Load index from container + int load(IndexStorage::Pointer container, + IndexMetric::Pointer metric) override; + + //! Unload index + int unload(void) override; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Retrieve statistics + const Stats &stats(void) const override; + + //! Create a searcher context + Context::Pointer create_context(void) const override; + + //! Create a new iterator + IndexProvider::Pointer create_provider(void) const override; + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve params of index + const ailego::Params ¶ms(void) const override { + return params_; + } + + protected: + int update_context(IVFSearcherContext *ctx) const { + auto entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone QcEntity"); + return IndexError_Runtime; + } + + //! The centroid index searcher may be different, so need to create one + auto centroid_ctx = centroid_index_->create_context(); + if (!centroid_ctx) { + LOG_ERROR("Failed to create centroid index searcher context"); + return IndexError_Runtime; + } + + return ctx->update_context(entity, centroid_ctx, params_, magic_); + } + + private: + //! Constants + static constexpr uint32_t kDefaultBfThreshold = 1000u; + + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; + + //! Members + IndexMeta meta_{}; + ailego::Params params_{}; + IVFCentroidIndex::Pointer centroid_index_{}; + IVFEntity::Pointer entity_{}; + uint32_t bruteforce_threshold_{kDefaultBfThreshold}; + uint32_t magic_{0}; + Stats stats_{}; + State searcher_state_{STATE_INIT}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_searcher_context.h b/src/core/algorithm/ivf/ivf_searcher_context.h new file mode 100644 index 0000000..594e045 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_searcher_context.h @@ -0,0 +1,239 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "ivf_entity.h" +#include "ivf_utility.h" + +namespace zvec { +namespace core { + +/*! IVF Searcher Context + */ +class IVFSearcherContext : public IndexSearcher::Context { + public: + IVFSearcherContext(const IVFEntity::Pointer &ivf_entity, + IndexSearcher::Context::Pointer ¢roid_ctx) + : entity_(ivf_entity), centroid_searcher_ctx_(std::move(centroid_ctx)) {} + + public: + //! Set topk of search result + void set_topk(uint32_t k) override { + topk_ = k; + result_heap_.limit(topk_); + result_heap_.set_threshold(this->threshold()); + } + + //! Retrieve search result + const IndexDocumentList &result(void) const override { + return results_[0]; + } + + //! Retrieve search result with index + const IndexDocumentList &result(size_t idx) const override { + ailego_assert_with(results_.size() > idx, "invalid index"); + return results_[idx]; + } + + //! Retrieve mutable result with index + IndexDocumentList *mutable_result(size_t idx) override { + ailego_assert_with(idx < results_.size(), "invalid idx"); + return &results_[idx]; + } + + inline IndexDocumentHeap *result_heap() { + return &result_heap_; + } + + //! Update the parameters of context + int update(const ailego::Params ¶ms) override { + params.get(PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + params.get(PARAM_IVF_SEARCHER_SCAN_RATIO, &scan_ratio_); + if (scan_ratio_ <= 0.0) { + LOG_ERROR("Invalid params %s=%f", PARAM_IVF_SEARCHER_SCAN_RATIO.c_str(), + scan_ratio_); + return IndexError_InvalidArgument; + } + size_t topk_val = + std::max(static_cast( + std::round(entity_->inverted_list_count() * scan_ratio_)), + 1u); + + uint32_t nprobe = 0; + params.get(PARAM_IVF_SEARCHER_NPROBE, &nprobe); + if (nprobe > 0) { + nprobe = std::min(nprobe, + static_cast(entity_->inverted_list_count())); + topk_val = nprobe; + } + + centroid_searcher_ctx_->set_topk(topk_val); + + // When nprobe is explicitly set, lift max_scan_count to the total vector + // count to ensure all selected clusters can be fully scanned. + if (nprobe > 0 && entity_->inverted_list_count() > 0) { + max_scan_count_ = static_cast(entity_->vector_count()); + } else { + max_scan_count_ = static_cast( + std::ceil(entity_->vector_count() * scan_ratio_)); + } + max_scan_count_ = std::max(bruteforce_threshold_, max_scan_count_); + return 0; + } + + //! Retrieve magic number + uint32_t magic(void) const override { + return magic_; + } + + public: + //! Initialize the context + int init(const ailego::Params ¶ms) { + return this->update(params); + } + + //! Update the magic number + void set_magic(uint32_t mag) { + magic_ = mag; + } + + //! Get Topk Value + uint32_t topk() const override { + return topk_; + } + + //! Retrieve scan ratio + float scan_ratio(void) const { + return scan_ratio_; + } + + //! Retrieve max scan count + uint32_t max_scan_count(void) const { + return max_scan_count_; + } + + uint32_t bruteforce_threshold() const { + return bruteforce_threshold_; + } + + //! Retrieve magic number + const IVFEntity::Pointer &entity() const { + return entity_; + } + + //! Retrieve Mutable Query Result By Query Index + IndexDocumentHeap &mutable_result_heap() { + return result_heap_; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + + bool fetch_vector(void) const override { + return fetch_vector_; + } + + //! Reset all the query results + void reset_results(size_t qnum) { + results_.resize(qnum); + stats_vec_.resize(qnum); + for (size_t i = 0; i < qnum; ++i) { + results_[i].clear(); + stats_vec_[i].clear(); + } + result_heap_.clear(); + result_heap_.limit(topk_); + result_heap_.set_threshold(this->threshold()); + } + + //! Update context, the context may be shared by different searcher + int update_context(IVFEntity::Pointer &new_entity, + IndexSearcher::Context::Pointer ¢roid_ctx, + const ailego::Params ¶ms, uint32_t magic_num) { + entity_ = new_entity; + centroid_searcher_ctx_ = std::move(centroid_ctx); + int ret = this->update(params); + ivf_check_error_code(ret); + + magic_ = magic_num; + + return 0; + } + + //! Retrieve the centroid index context + IndexSearcher::Context::Pointer ¢roid_searcher_ctx(void) { + return centroid_searcher_ctx_; + } + + const Stats &stats(size_t idx = 0) const { + ailego_assert_with(stats_vec_.size() > idx, "invalid index"); + return stats_vec_[idx]; + } + + Stats &mutable_stats(size_t idx = 0) { + ailego_assert_with(stats_vec_.size() > idx, "invalid index"); + return stats_vec_[idx]; + } + + void topk_to_result(uint32_t idx) { + if (ailego_unlikely(result_heap_.size() == 0)) { + return; + } + + ailego_assert_with(idx < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(result_heap_.size())); + result_heap_.sort(); + results_[idx].clear(); + for (int i = 0; i < size; ++i) { + auto score = result_heap_[i].score(); + if (score > this->threshold()) { + break; + } + + key_t key = result_heap_[i].key(); + if (fetch_vector_) { + IndexStorage::MemoryBlock block; + entity_->get_vector_by_key(key, block); + results_[idx].emplace_back(key, score, key, block); + } else { + results_[idx].emplace_back(key, score); + } + } + } + + private: + //! Constants + static constexpr float kDefaultScanRatio = 0.1f; + static constexpr uint32_t kDefaultBfThreshold = 1000u; + + //! Members + IVFEntity::Pointer entity_{}; + IndexSearcher::Context::Pointer centroid_searcher_ctx_{}; + IndexDocumentHeap result_heap_; + std::vector results_{}; + std::vector stats_vec_{}; + + bool fetch_vector_{false}; + uint32_t topk_{0}; + uint32_t magic_{0}; + float scan_ratio_{kDefaultScanRatio}; + uint32_t max_scan_count_{0}; + uint32_t bruteforce_threshold_{kDefaultBfThreshold}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_streamer.cc b/src/core/algorithm/ivf/ivf_streamer.cc new file mode 100644 index 0000000..6bb9654 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_streamer.cc @@ -0,0 +1,324 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "ivf_streamer.h" +#include +#include +#include "ivf_centroid_index.h" +#include "ivf_index_provider.h" +#include "ivf_params.h" + +namespace zvec { +namespace core { + +int IVFStreamer::init(const IndexMeta &meta, const ailego::Params ¶meters) { + meta_ = meta; + params_ = parameters; + + params_.get(PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD, &bruteforce_threshold_); + + searcher_state_ = STATE_INITED; + + return 0; +} + +int IVFStreamer::cleanup(void) { + this->unload(); + + params_.clear(); + bruteforce_threshold_ = kDefaultBfThreshold; + + searcher_state_ = STATE_INIT; + return 0; +} + +int IVFStreamer::open(IndexStorage::Pointer storage) { + if (!storage) { + LOG_ERROR("Invalid storage"); + return IndexError_InvalidArgument; + } + if (searcher_state_ != STATE_INITED) { + LOG_ERROR("Initalize the searcher first before load index"); + return IndexError_Runtime; + } + + ailego::ElapsedTime timer; + int ret = IndexHelper::DeserializeFromStorage(storage.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize meta from storage"); + return ret; + } + + //! Load centroid index + centroid_index_ = std::make_shared(); + if (!centroid_index_) { + return IndexError_NoMemory; + } + auto seg = storage->get(IVF_CENTROID_SEG_ID, 0); + if (!seg) { + LOG_ERROR("Failed to get segment %s", IVF_CENTROID_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + IndexStorage::Pointer seg_container = + std::make_shared(seg); + if (!seg_container) { + return IndexError_NoMemory; + } + ret = seg_container->open(std::string(), false); + if (ret != 0) { + LOG_ERROR("IndexSegmentStorage load failed for %s", IndexError::What(ret)); + return ret; + } + ret = centroid_index_->load(seg_container, params_); + if (ret != 0) { + LOG_ERROR("Failed to load index for %s", IndexError::What(ret)); + return ret; + } + + auto reformer = centroid_index_->reformer(); + if (reformer) { + //! The centroid index is loaded from the centroid sub-segment which does + //! not contain the rotator segment. Load the reformer state (e.g. rotation + //! matrix) from the main storage instead. + ret = reformer->load(storage); + ivf_check_error_code(ret); + } + params_.set(PARAM_IVF_SEARCHER_CONVERTER_REFORMER, reformer); + + //! load iverted index + entity_ = std::make_shared(); + if (!entity_) { + return IndexError_NoMemory; + } + ret = entity_->load(storage); + ivf_check_error_code(ret); + + magic_ = IndexContext::GenerateMagic(); + + stats_.set_loaded_count(entity_->vector_count()); + stats_.set_loaded_costtime(timer.milli_seconds()); + + searcher_state_ = STATE_LOADED; + return 0; +} + +int IVFStreamer::unload(void) { + magic_ = 0; + centroid_index_.reset(); + entity_.reset(); + stats_.set_loaded_count(0UL); + stats_.set_loaded_costtime(0UL); + stats_.clear_attributes(); + searcher_state_ = STATE_INITED; + + return 0; +} + +int IVFStreamer::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_bf_impl(query, qmeta, 1, context); +} + +int IVFStreamer::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + if (!query || qmeta.element_size() != meta_.element_size()) { + LOG_ERROR("Null query or invalid qmeta"); + return IndexError_InvalidArgument; + } + IVFSearcherContext *ctx = dynamic_cast(context.get()); + if (!ctx || ctx->topk() == 0) { + LOG_ERROR("Invalid context or topk not set yet"); + return IndexError_InvalidArgument; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher + int ret = this->update_context(ctx); + ivf_check_error_code(ret); + } + + ctx->reset_results(count); + auto &entity = ctx->entity(); + auto &filter = ctx->filter(); + + //! Transform the querys for querying in inverted vector index later + IndexQueryMeta iv_qmeta; + int ret = entity->transform(query, qmeta, count, &query, &iv_qmeta); + ivf_check_with_msg(ret, "Failed to transform querys"); + + // TODO: do batch search in matrix + for (size_t q = 0; q < count; ++q) { + auto &context_stats = ctx->mutable_stats(q); + auto &heap = ctx->mutable_result_heap(); + heap.clear(); + if (!filter.is_valid()) { + ret = entity->search(query, &heap, &context_stats); + } else { + ret = entity->search(query, filter, &heap, &context_stats); + } + ivf_check_with_msg(ret, "Failed to search in entity for %s", + IndexError::What(ret)); + heap.sort(); // sort the results + if (!filter.is_valid()) { + // mapping the local id to key if query without filter + ret = entity->retrieve_keys(&heap); + ivf_check_error_code(ret); + } + entity->normalize(q, &heap); + ctx->topk_to_result(q); + query = static_cast(query) + iv_qmeta.element_size(); + } + + return 0; +} + +int IVFStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return this->search_impl(query, qmeta, 1, context); +} + +int IVFStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const { + if (entity_->vector_count() <= bruteforce_threshold_) { + return this->search_bf_impl(query, qmeta, count, context); + } + if (!query || qmeta.element_size() != meta_.element_size()) { + LOG_ERROR("Null query or invalid qmeta"); + return IndexError_InvalidArgument; + } + + IVFSearcherContext *ctx = dynamic_cast(context.get()); + if (!ctx || ctx->topk() == 0) { + LOG_ERROR("Invalid context or topk not set yet"); + return IndexError_InvalidArgument; + } + if (ctx->magic() != magic_) { + //! context is created by another searcher + int ret = update_context(ctx); + ivf_check_error_code(ret); + } + + ctx->reset_results(count); + auto &entity = ctx->entity(); + auto &filter = ctx->filter(); + + auto ¢roid_index_ctx = ctx->centroid_searcher_ctx(); + int ret = centroid_index_->search(query, qmeta, count, centroid_index_ctx); + ivf_check_error_code(ret); + + //! Transform the querys for querying in inverted vector index later + IndexQueryMeta iv_qmeta; + ret = entity->transform(query, qmeta, count, &query, &iv_qmeta); + ivf_check_with_msg(ret, "Failed to transform querys"); + + for (size_t q = 0; q < count; ++q) { + auto ¢roids = centroid_index_ctx->result(q); + auto &context_stats = ctx->mutable_stats(q); + auto &heap = ctx->mutable_result_heap(); + heap.clear(); + size_t total_scan_count = 0; + for (size_t i = 0; + i < centroids.size() && total_scan_count < ctx->max_scan_count(); + ++i) { + auto cid = centroids[i].key(); + uint32_t scan_count = 0; + if (!filter.is_valid()) { + ret = entity->search(cid, query, &scan_count, &heap, &context_stats); + } else { + ret = entity->search(cid, query, filter, &scan_count, &heap, + &context_stats); + } + ivf_check_with_msg(ret, "Failed to search in entity for %s", + IndexError::What(ret)); + total_scan_count += scan_count; + } + heap.sort(); // sort the results + if (!filter.is_valid()) { + // mapping the local id to key if query without filter + ret = entity->retrieve_keys(&heap); + ivf_check_error_code(ret); + } + entity->normalize(q, &heap); + ctx->topk_to_result(q); + query = static_cast(query) + iv_qmeta.element_size(); + } + + return 0; +} + +const IndexSearcher::Stats &IVFStreamer::stats(void) const { + return stats_; +} + +IndexSearcher::Context::Pointer IVFStreamer::create_context() const { + if (searcher_state_ != STATE_LOADED) { + LOG_ERROR("Load the index first before create context"); + return nullptr; + } + + auto entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone IVFEntity"); + return nullptr; + } + + auto centroid_index_ctx = centroid_index_->create_context(); + if (!centroid_index_ctx) { + LOG_ERROR("Failed to create centroid index context"); + return nullptr; + } + + auto context = + new (std::nothrow) IVFSearcherContext(entity, centroid_index_ctx); + if (!context) { + LOG_ERROR("Failed to alloc IVFSearcherContext"); + return nullptr; + } + int ret = context->init(params_); + if (ret != 0) { + delete context; + return nullptr; + } + + context->set_magic(magic_); + + return Context::Pointer(context); +} + +IndexProvider::Pointer IVFStreamer::create_provider(void) const { + if (searcher_state_ != STATE_LOADED) { + LOG_ERROR("Load the index first before create provider"); + return nullptr; + } + + auto entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone IVFEntity"); + return Provider::Pointer(); + } + + auto *provider = new (std::nothrow) + IVFIndexProvider(entity->has_orignal_feature() ? meta_ : entity->meta(), + entity, "IVFStreamer"); + if (!provider) { + LOG_ERROR("Failed to alloc IVFIndexProvider"); + return Provider::Pointer(); + } + + return Provider::Pointer(provider); +} + +INDEX_FACTORY_REGISTER_STREAMER(IVFStreamer); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf/ivf_streamer.h b/src/core/algorithm/ivf/ivf_streamer.h new file mode 100644 index 0000000..79c9d56 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_streamer.h @@ -0,0 +1,123 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#ifndef __IVF_STREAMER_H__ +#define __IVF_STREAMER_H__ + +#include +#include "ivf_centroid_index.h" +#include "ivf_entity.h" +#include "ivf_searcher_context.h" + +namespace zvec { +namespace core { + +/*! IVF Searcher + */ +class IVFStreamer : public IndexStreamer { + public: + //! Initialize Searcher + int init(const IndexMeta & /*meta*/, + const ailego::Params & /*params*/) override; + + //! Cleanup Searcher + int cleanup(void) override; + + //! Load index from container + int open(IndexStorage::Pointer storage) override; + + int flush(uint64_t /*check_point*/) override { + return 0; + } + int close(void) override { + return this->unload(); + } + + //! Unload index + int unload(void) override; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Retrieve statistics + const Stats &stats(void) const override; + + //! Create a searcher context + Context::Pointer create_context(void) const override; + + //! Create a new iterator + IndexProvider::Pointer create_provider(void) const override; + + //! Retrieve meta of index + const IndexMeta &meta(void) const override { + return meta_; + } + + int get_vector_by_id(const uint32_t id, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector_by_key(id, block); + } + + protected: + int update_context(IVFSearcherContext *ctx) const { + auto entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone QcEntity"); + return IndexError_Runtime; + } + + //! The centroid index searcher may be different, so need to create one + auto centroid_ctx = centroid_index_->create_context(); + if (!centroid_ctx) { + LOG_ERROR("Failed to create centroid index searcher context"); + return IndexError_Runtime; + } + + return ctx->update_context(entity, centroid_ctx, params_, magic_); + } + + private: + //! Constants + static constexpr uint32_t kDefaultBfThreshold = 1000u; + + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; + + //! Members + IndexMeta meta_{}; + ailego::Params params_{}; + IndexBuilder::Pointer builder_; + IVFCentroidIndex::Pointer centroid_index_{}; + IVFEntity::Pointer entity_{}; + uint32_t bruteforce_threshold_{kDefaultBfThreshold}; + uint32_t magic_{0}; + Stats stats_{}; + State searcher_state_{STATE_INIT}; +}; + +} // namespace core +} // namespace zvec + +#endif //__IVF_STREAMER_H__ diff --git a/src/core/algorithm/ivf/ivf_utility.h b/src/core/algorithm/ivf/ivf_utility.h new file mode 100644 index 0000000..e7d2131 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_utility.h @@ -0,0 +1,177 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +#ifndef ivf_check_error_code +#define ivf_check_error_code(code) \ + if (ailego_unlikely((code) != 0)) return code +#endif + +#ifndef ivf_assert +#define ivf_assert(cond, code) \ + if (ailego_unlikely(!(cond))) return code +#endif + +#ifndef ivf_check_with_msg +#define ivf_check_with_msg(code, fmt, ...) \ + do { \ + if (ailego_unlikely((code) != 0)) { \ + LOG_ERROR(fmt, ##__VA_ARGS__); \ + return code; \ + } \ + } while (0) +#endif + +#ifndef ivf_assert_with_msg +#define ivf_assert_with_msg(cond, err, fmt, ...) \ + do { \ + if (ailego_unlikely(!(cond))) { \ + LOG_ERROR(fmt, ##__VA_ARGS__); \ + return err; \ + } \ + } while (0) +#endif + +/*! Quantized Clustering Utility + */ +class IVFUtility { + public: + //! Generator a random path with specificed prefix + static inline std::string GenerateRandomPath(const std::string &prefix) { + uint64_t timestamp = ailego::Monotime::MicroSeconds(); + return prefix + std::to_string(timestamp); + } + + //! Compute the default scan ratio for total vectors + static inline float ComputeScanRatio(size_t vector_count) { + // the fitting function for the follow points: 1000000(0.02) + // 10000000(0.01) 50000000(0.005) 100000000(0.001) + float scan_ratio = -0.004 * std::log(vector_count) + 0.0751; + scan_ratio = std::max(scan_ratio, 0.0001f); + return scan_ratio; + } + + //! Transpose the vectors in row major order to column major order + static inline void Transpose(size_t align_size, const void *src, size_t m, + size_t dim, void *dst); + + //! Transpose the vectors in column major order to row major order + static inline void ReverseTranspose(size_t align_size, const void *src, + size_t m, size_t dim, void *dst); + + //! Aligned size of a block vectors buffer + static inline size_t AlignedSize(size_t fnum, size_t element_size); + + //! Aligned size of one vector buffer + static inline size_t AlignedSize(size_t element_size); + + //! Sort arr with size in ascending order, and keep the index postion + //! n2o keep the mapping: new position => origin postion + //! For example, the input arr = [5, 3, 9, 6, 7], size = 5, after sort + // arr = [3, 5, 6, 7, 9] + // n2o = [1, 0, 3, 4, 2] + //! To save memory, no extra memory is allocated + template + static void Sort(T *arr, std::vector *n2o, size_t size) { + std::vector o2n; + o2n.resize(size); + n2o->resize(size); + + std::iota(n2o->begin(), n2o->end(), 0U); + std::sort(n2o->begin(), n2o->end(), + [&](I i, I j) { return arr[i] < arr[j]; }); + for (I i = 0U; i < size; ++i) { + o2n[(*n2o)[i]] = i; + } + //! reorder arr in place, according to given n2o index + for (I i = 0; i < size; ++i) { + if (i != (*n2o)[i]) { + T tmp = arr[i]; + I j = i, k; + while (i != (k = (*n2o)[j])) { + arr[j] = arr[k]; + (*n2o)[j] = j; + j = k; + } + arr[j] = tmp; + (*n2o)[j] = j; + } + } + + for (I i = 0U; i < size; ++i) { + (*n2o)[o2n[i]] = i; + } + } + + //! Transpose one vector in block + template + static inline void TransposeOne(const void *src, size_t M, size_t N, + void *dst) { + for (size_t i = 0; i < N; ++i) { + reinterpret_cast(dst)[i] = reinterpret_cast(src)[i * M]; + } + } +}; + +void IVFUtility::Transpose(size_t align_size, const void *src, size_t m, + size_t dim, void *dst) { + switch (align_size) { + case 2: + ailego::MatrixHelper::Transpose(src, m, dim, dst); + break; + case 4: + ailego::MatrixHelper::Transpose(src, m, dim, dst); + break; + case 8: + ailego::MatrixHelper::Transpose(src, m, dim, dst); + break; + } +} + +void IVFUtility::ReverseTranspose(size_t align_size, const void *src, size_t m, + size_t dim, void *dst) { + switch (align_size) { + case 2: + ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst); + break; + case 4: + ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst); + break; + case 8: + ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst); + break; + } +} + +size_t IVFUtility::AlignedSize(size_t fnum, size_t element_size) { + return ailego_align(fnum * element_size, 32); +} + +size_t IVFUtility::AlignedSize(size_t element_size) { + return ailego_align(element_size, 32); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/CMakeLists.txt b/src/core/algorithm/vamana/CMakeLists.txt new file mode 100644 index 0000000..b2feaf9 --- /dev/null +++ b/src/core/algorithm/vamana/CMakeLists.txt @@ -0,0 +1,17 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(NOT APPLE) + set(CORE_KNN_VAMANA_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_vamana + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework core_knn_hnsw sparsehash + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_VAMANA_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/vamana/vamana_algorithm.cc b/src/core/algorithm/vamana/vamana_algorithm.cc new file mode 100644 index 0000000..bf0fe5e --- /dev/null +++ b/src/core/algorithm/vamana/vamana_algorithm.cc @@ -0,0 +1,620 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "vamana_algorithm.h" +#include +#include + +namespace zvec { +namespace core { + +// ============================================================================ +// add_node: Insert a new node into the Vamana graph. +// +// Algorithm (from DiskANN / Vamana paper): +// 1. GreedySearch from entry_point to find candidate neighbors +// 2. RobustPrune to select diverse neighbors for the new node +// 3. Update the new node's neighbor list +// 4. For each new neighbor, add reverse link; if over-degree, RobustPrune +// 5. If this is the first node, set it as entry point +// ============================================================================ +template +int VamanaAlgorithm::add_node(node_id_t id, VamanaContext *ctx) { + // Lazily initialize distance storage on first insert + entity_.ensure_dist_storage(); + + spin_lock_.lock(); + auto entry_point = entity_.entry_point(); + + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + entity_.update_entry_point(id); + spin_lock_.unlock(); + return 0; + } + spin_lock_.unlock(); + + // Step 1: GreedySearch to find candidate neighbors + uint32_t search_list_size = entity_.search_list_size(); + ctx->topk_heap().clear(); + ctx->topk_heap().limit(search_list_size); + ctx->dist_calculator().clear_compare_cnt(); + + // Set query to the new node's vector. Use reset_query (same as search path) + // so that greedy_search works with the search-optimized distance kernel. + const void *query_vec = entity_.get_vector(id); + if (ailego_unlikely(query_vec == nullptr)) { + LOG_ERROR("Failed to get vector for node %u", id); + return IndexError_ReadData; + } + ctx->reset_query(query_vec); + + greedy_search(entry_point, ctx, /*use_pool=*/false); + + auto &topk_heap = ctx->topk_heap(); + + // Step 2: RobustPrune to select diverse neighbors + robust_prune(id, topk_heap, entity_.alpha(), entity_.max_degree(), ctx); + // Copy result before reverse updates (which also call robust_prune) + auto pruned_neighbors = ctx->prune_result(); + + // Step 3: Update the new node's neighbor list and distances + entity_.update_neighbors(id, pruned_neighbors); + entity_.update_neighbor_dists(id, pruned_neighbors); + + // Step 4: Reverse-link updates + update_neighbors_and_reverse_links(id, pruned_neighbors, ctx); + + return 0; +} + +// ============================================================================ +// search: Greedy search for approximate nearest neighbors. +// ============================================================================ +template +int VamanaAlgorithm::search(VamanaContext *ctx) const { + spin_lock_.lock(); + auto entry_point = entity_.entry_point(); + spin_lock_.unlock(); + + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + return 0; + } + + auto &topk_heap = ctx->topk_heap(); + topk_heap.clear(); + + // Use ef (query-time parameter) instead of entity.search_list_size() + // (build-time L parameter). search_list_size controls construction; + // ef controls search quality and is user-configurable at query time. + uint32_t ef_search = std::max(static_cast(ctx->topk()), ctx->ef()); + topk_heap.limit(ef_search); + + greedy_search(entry_point, ctx, /*use_pool=*/true); + + return 0; +} + +// ============================================================================ +// greedy_search helper templates +// +// Two specialized inner loops, dispatched from greedy_search(): +// +// fast_greedy_search: mmap/contiguous with direct vector pointers. +// Uses batch_dist on a pointer array. +// slow_greedy_search: BufferPool-backed storage: must fetch MemBlock +// wrappers via get_vector_typed to pin pages. +// +// Both accept either BlockHeap or LinearPool as `HeapType` because the +// two expose the same reset(n, ef, block_size) / push_block(dists, ids, n) +// surface (LinearPool adapts via push_block and ignores the block_size hint). +// ============================================================================ + +// mmap/contiguous variant: resolve vectors via get_vector_ptr +// and dispatch to the classic pointer-array batch_dist. +template +void fast_greedy_search(const EntityType &entity, HeapType &pool, + VisitFilter &visit, VamanaDistCalculator &dc, + uint32_t topk, uint32_t ef, node_id_t entry_point, + uint32_t prefetch_lines, uint32_t prefetch_offset) { + const uint32_t max_deg = entity.max_degree(); + const uint32_t cap = std::max(topk, ef); + pool.reset(static_cast(cap), static_cast(max_deg)); + visit.clear(); + + dist_t ep_dist = dc.batch_dist(entry_point); + visit.set_visited(entry_point); + pool.push_block(&ep_dist, &entry_point, 1); + + uint32_t buf_capacity = max_deg; + std::vector neighbor_ids(buf_capacity); + std::vector dists(buf_capacity); + std::vector neighbor_vecs(buf_capacity); + + while (pool.has_next()) { + auto current_node = pool.pop(); + + const auto neighbors = entity.get_neighbors_typed(current_node); + ailego_prefetch(neighbors.data); + + if (neighbors.size() > buf_capacity) { + buf_capacity = neighbors.size(); + neighbor_ids.resize(buf_capacity); + dists.resize(buf_capacity); + neighbor_vecs.resize(buf_capacity); + } + + const uint32_t po = + std::min(static_cast(neighbors.size()), prefetch_offset); + uint32_t unvisited_count = 0; + uint32_t i = 0; + + for (; i < po; ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) continue; + visit.set_visited(node); + const void *vec_ptr = entity.get_vector_ptr(node); + const char *p = reinterpret_cast(vec_ptr); + for (uint32_t cl = 0; cl < prefetch_lines; ++cl) { + ailego_prefetch(p + cl * 64); + } + neighbor_ids[unvisited_count] = node; + neighbor_vecs[unvisited_count] = vec_ptr; + unvisited_count++; + } + for (; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) continue; + visit.set_visited(node); + neighbor_ids[unvisited_count] = node; + neighbor_vecs[unvisited_count] = entity.get_vector_ptr(node); + unvisited_count++; + } + + if (unvisited_count == 0) continue; + dc.batch_dist(neighbor_vecs.data(), unvisited_count, dists.data()); + pool.push_block(dists.data(), neighbor_ids.data(), + static_cast(unvisited_count)); + } +} + +// ============================================================================ +// dual_heap_greedy_search: shared core for the fallback dual-heap path. +// +// Maintains a candidate min-heap + topk heap + VisitFilter. Uses plain +// batch_dist. +// ============================================================================ +template +void dual_heap_greedy_search(const EntityType &entity, VamanaContext *ctx, + VamanaDistCalculator &dc, node_id_t entry_point, + FilterFn &&filter) { + const uint32_t prefetch_offset = ctx->po(); + const uint32_t prefetch_lines = + ctx->pl() > 0 ? ctx->pl() : (entity.vector_size() + 63) / 64; + + uint32_t buf_capacity = entity.max_degree(); + std::vector neighbor_ids(buf_capacity); + std::vector neighbor_vec_blocks; + neighbor_vec_blocks.reserve(buf_capacity); + std::vector dists(buf_capacity); + std::vector neighbor_vecs(buf_capacity); + + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + auto &topk_heap = ctx->topk_heap(); + candidates.clear(); + visit.clear(); + + // Initialize with entry point using batch_dist (single-element batch). + // We must NOT use dc.dist(entry_point) here because dist() calls + // distance_() which is squared_euclidean_int8_distance (sign/abs trick, + // expects two raw int8 inputs), but query_ has been preprocessed by + // reset_query (+128 shift to uint8). batch_dist() correctly calls + // batch_distance_() which expects the preprocessed uint8 query. + dist_t entry_dist = dc.batch_dist(entry_point); + if (ailego_unlikely(dc.error())) { + return; + } + visit.set_visited(entry_point); + if (!filter(entry_point)) { + topk_heap.emplace(entry_point, entry_dist); + } + candidates.emplace(entry_point, entry_dist); + + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t current_node = top->first; + dist_t current_dist = top->second; + + // Early termination: if the closest candidate is worse than the worst + // result in a full topk heap, we won't find anything better. + if (topk_heap.full() && current_dist > topk_heap[0].second) { + break; + } + + candidates.pop(); + + // Expand neighbors using typed access (no virtual dispatch) + const auto neighbors = entity.get_neighbors_typed(current_node); + ailego_prefetch(neighbors.data); + + // Resize buffers if this node has more neighbors than expected + if (neighbors.size() > buf_capacity) { + buf_capacity = neighbors.size(); + neighbor_ids.resize(buf_capacity); + dists.resize(buf_capacity); + neighbor_vecs.resize(buf_capacity); + } + + // Collect unvisited neighbors (reuse pre-allocated buffer) + uint32_t unvisited_count = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) continue; + visit.set_visited(node); + neighbor_ids[unvisited_count++] = node; + } + if (unvisited_count == 0) continue; + + // Batch fetch vectors using typed access (reuse pre-allocated buffer) + neighbor_vec_blocks.clear(); + int ret = entity.get_vector_typed(neighbor_ids.data(), unvisited_count, + neighbor_vec_blocks); + if (ailego_unlikely(ret != 0)) break; + + for (uint32_t i = 0; i < std::min(prefetch_offset, unvisited_count); ++i) { + const char *base = + static_cast(neighbor_vec_blocks[i].data()); + for (uint32_t cl = 0; cl < prefetch_lines; ++cl) { + ailego_prefetch(base + cl * 64); + } + } + + // Batch distance computation (reuse pre-allocated buffers). + for (uint32_t i = 0; i < unvisited_count; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + dc.batch_dist(neighbor_vecs.data(), unvisited_count, dists.data()); + + for (uint32_t i = 0; i < unvisited_count; ++i) { + node_id_t node = neighbor_ids[i]; + dist_t node_dist = dists[i]; + if ((!topk_heap.full()) || node_dist < topk_heap[0].second) { + candidates.emplace(node, node_dist); + if (!filter(node)) { + topk_heap.emplace(node, node_dist); + } + } + } + } +} + +// ============================================================================ +// greedy_search: Beam search from entry_point. +// +// Maintains a candidate min-heap (ordered by distance) and a visited set. +// At each step, pops the closest unvisited candidate, expands its neighbors, +// and adds unvisited neighbors to both the candidate heap and the topk heap. +// Stops when the closest candidate is farther than the worst in topk, or +// when the scan limit is reached. +// ============================================================================ +template +void VamanaAlgorithm::greedy_search(node_id_t entry_point, + VamanaContext *ctx, + bool use_pool) const { + const auto &entity = static_cast(ctx->get_entity()); + VamanaDistCalculator &dc = ctx->dist_calculator(); + + const IndexFilter &index_filter = + static_cast(ctx)->filter(); + + const uint32_t prefetch_lines = + ctx->pl() > 0 ? ctx->pl() : (entity.vector_size() + 63) / 64; + + if (!use_pool || index_filter.is_valid()) { + // Fallback path used by add_node (use_pool=false) and filtered search. + // Dispatched to dual_heap_greedy_search (plain batch_dist). + auto run_with_filter = [&](auto &&filter) { + dual_heap_greedy_search( + entity, ctx, dc, entry_point, std::forward(filter)); + }; + + if (index_filter.is_valid()) { + auto filter = [&](node_id_t id) { + return index_filter(entity.get_key_typed(id)); + }; + run_with_filter(filter); + } else { + auto filter = [](node_id_t) { return false; }; + run_with_filter(filter); + } + } else { + // Fast pool-based path for mmap/contiguous entities that support + // direct pointer access. BlockHeap (AVX2) or LinearPool (scalar) + // are used for top-k tracking. BufferPool entities fall back to + // dual_heap_greedy_search since they lack direct pointer access. + if constexpr (std::is_same_v) { + const uint32_t topk_v = static_cast(ctx->topk()); + const uint32_t ef_v = ctx->ef(); + const bool avx2_ok = + zvec::ailego::internal::CpuFeatures::static_flags_.AVX2; + auto &topk_heap = ctx->topk_heap(); + + auto &visit = ctx->visit_filter(); + + if (avx2_ok) { + auto &bpool = ctx->block_pool(); + fast_greedy_search(entity, bpool, visit, dc, topk_v, ef_v, entry_point, + prefetch_lines, ctx->po()); + copy_pool_to_topk(bpool, topk_heap); + } else { + auto &lpool = ctx->pool(); + fast_greedy_search(entity, lpool, visit, dc, topk_v, ef_v, entry_point, + prefetch_lines, ctx->po()); + copy_pool_to_topk(lpool, topk_heap); + } + } else { + // BufferPool entities: fallback to dual-heap path. + auto filter = [](node_id_t) { return false; }; + dual_heap_greedy_search(entity, ctx, dc, + entry_point, filter); + } + } +} + +// ============================================================================ +// robust_prune: Select up to max_degree diverse neighbors from candidates. +// +// Faithfully follows DiskANN's occlude_list algorithm: +// 1. Sort candidates by distance (ascending) +// 2. Truncate to max_occlusion_size (DiskANN's maxc parameter) +// 3. Multi-round alpha sweep: cur_alpha starts at 1.0, increments by *1.2 +// each round until reaching alpha. This progressively relaxes the +// occlusion criterion. +// 4. For each candidate, compute occlude_factor as: +// max over all selected neighbors p of: dist(query, candidate) / dist(p, +// candidate) +// If occlude_factor > cur_alpha, the candidate is occluded in this round. +// 5. After all rounds, if _saturate_graph and alpha > 1, fill remaining +// slots with any un-selected candidates. +// ============================================================================ +template +void VamanaAlgorithm::robust_prune(node_id_t id, + TopkHeap &candidates, + float alpha, uint32_t max_degree, + VamanaContext *ctx) const { + auto &result = ctx->prune_result(); + result.clear(); + + if (candidates.size() == 0) return; + + // Sort candidates by distance (ascending — closest first) + candidates.sort(); + + VamanaDistCalculator &dc = ctx->dist_calculator(); + size_t n = candidates.size(); + + // Build-time distance offset: shifts the internal distance into a + // non-negative range so the ratio-based occlude_factor below is + // geometrically meaningful. Zero for metrics whose internal distance is + // already non-negative (SquaredEuclidean, etc.); 1.0 for the quantized int8 + // Cosine / NormalizedCosine path (internal distance is -cos, so + // offset=1 yields 1-cos which matches DiskANN's normalized-L2 semantics). + const float dist_offset = ctx->build_distance_offset(); + + // Truncate to max_occlusion_size (DiskANN's maxc parameter) + size_t maxc = entity_.max_occlusion_size(); + if (maxc > 0 && n > maxc) { + n = maxc; + } + + // Pre-cache all candidate vectors at once + auto &vec_cache = ctx->prune_vec_cache(); + vec_cache.resize(n); + for (size_t i = 0; i < n; ++i) { + vec_cache[i] = entity_.get_vector(candidates[i].first); + } + + // occlude_factor: tracks the maximum occlusion ratio for each candidate + // (DiskANN: occlude_factor[t] = max over selected p of + // dist_to_query/dist_to_p) + auto &occlude_factor = ctx->prune_occlude_factor(); + occlude_factor.assign(n, 0.0f); + + // Pre-allocated buffers for batch distance computation + auto &batch_vecs = ctx->batch_vecs_buf(); + auto &batch_dists = ctx->batch_dists_buf(); + auto &batch_indices = ctx->batch_indices_buf(); + batch_vecs.resize(n); + batch_dists.resize(n); + batch_indices.resize(n); + + // Multi-round alpha sweep (DiskANN: cur_alpha starts at 1, increments *1.2) + float cur_alpha = 1.0f; + while (cur_alpha <= alpha + 1e-6f && result.size() < max_degree) { + for (size_t i = 0; i < n && result.size() < max_degree; ++i) { + if (occlude_factor[i] > cur_alpha) { + continue; + } + + // Mark as consumed so it won't be reconsidered + occlude_factor[i] = std::numeric_limits::max(); + + // Skip self-loops + if (candidates[i].first == id) continue; + + const void *selected_vec = vec_cache[i]; + if (ailego_unlikely(selected_vec == nullptr)) continue; + + // Add this candidate as a neighbor + node_id_t candidate_id = candidates[i].first; + dist_t candidate_dist = candidates[i].second; + result.emplace_back(candidate_id, candidate_dist); + + // Update occlude_factor for remaining candidates + // Collect candidates that haven't been consumed yet + uint32_t batch_count = 0; + for (size_t j = i + 1; j < n; ++j) { + if (occlude_factor[j] > alpha) continue; // already fully occluded + if (ailego_unlikely(vec_cache[j] == nullptr)) continue; + batch_vecs[batch_count] = vec_cache[j]; + batch_indices[batch_count] = static_cast(j); + batch_count++; + } + + if (batch_count > 0) { + // Compute distances from selected candidate to remaining candidates. + // distance_ is the symmetric data-to-data kernel (no pairwise split). + for (uint32_t k = 0; k < batch_count; ++k) { + batch_dists[k] = dc.dist(selected_vec, batch_vecs[k]); + } + + // DiskANN (L2/Cosine): + // occlude_factor[t] = max(occlude_factor[t], dist_to_query / + // dist_between) + // where dist_to_query = candidates[j].second (distance from query to j) + // dist_between = batch_dists[k] (distance from selected to j) + // + // `dist_offset` shifts both distances into a non-negative range for + // metrics whose internal distance can be negative (e.g. quantized + // int8 cosine stores -cos(m,q) in [-1, 1]). Without this shift the + // ratio loses its geometric meaning and RobustPrune produces a + // poor-quality graph (low recall at low ef_search). + for (uint32_t k = 0; k < batch_count; ++k) { + uint32_t j = batch_indices[k]; + float dist_selected_to_candidate = batch_dists[k] + dist_offset; + if (dist_selected_to_candidate <= 0.0f) { + occlude_factor[j] = std::numeric_limits::max(); + } else { + occlude_factor[j] = std::max(occlude_factor[j], + (candidates[j].second + dist_offset) / + dist_selected_to_candidate); + } + } + } + } + cur_alpha *= 1.2f; + } + + // Saturate graph: if enabled and alpha > 1, fill remaining slots with any + // un-selected candidates. This improves graph connectivity (better recall) + // at the cost of slightly more distance computations during search. + // Configurable via proxima.vamana.streamer.saturate_graph (default: false, + // matching DiskANN's default behavior). + if (entity_.saturate_graph() && alpha > 1.0f) { + for (size_t i = 0; i < n && result.size() < max_degree; ++i) { + if (candidates[i].first == id) continue; + bool already_selected = false; + for (const auto &r : result) { + if (r.first == candidates[i].first) { + already_selected = true; + break; + } + } + if (!already_selected) { + result.emplace_back(candidates[i].first, candidates[i].second); + } + } + } +} + +// ============================================================================ +// update_neighbors_and_reverse_links: For each new neighbor of `id`, +// add a reverse link from neighbor back to `id`. If the neighbor's degree +// exceeds max_degree, prune it using RobustPrune. +// ============================================================================ +template +void VamanaAlgorithm::update_neighbors_and_reverse_links( + node_id_t id, + const std::vector> &new_neighbors, + VamanaContext *ctx) { + for (const auto &[neighbor_id, dist] : new_neighbors) { + reverse_update_neighbor(id, neighbor_id, dist, ctx); + } +} + +// ============================================================================ +// reverse_update_neighbor: Add `id` as a neighbor of `neighbor_id`. +// If neighbor_id already has max_degree neighbors, collect all neighbors +// + the new one into a candidate set and RobustPrune. +// ============================================================================ +template +void VamanaAlgorithm::reverse_update_neighbor(node_id_t id, + node_id_t neighbor_id, + dist_t dist, + VamanaContext *ctx) { + std::lock_guard lock(lock_pool_[neighbor_id & kLockMask]); + + const Neighbors current_neighbors = entity_.get_neighbors(neighbor_id); + uint32_t current_size = current_neighbors.size(); + uint32_t max_deg = entity_.max_degree(); + + // Check if `id` is already a neighbor + for (uint32_t i = 0; i < current_size; ++i) { + if (current_neighbors[i] == id) return; + } + + if (current_size < max_deg) { + // Simply append and record distance + entity_.add_neighbor(neighbor_id, current_size, id); + entity_.set_neighbor_dist(neighbor_id, current_size, dist); + return; + } + + // Need to prune: collect current neighbors + new node into candidates + VamanaDistCalculator &dc = ctx->dist_calculator(); + + // Reuse update_heap from context instead of creating a new TopkHeap each time + TopkHeap &prune_candidates = ctx->update_heap(); + prune_candidates.clear(); + prune_candidates.limit(max_deg + 1); + + // Add existing neighbors — use cached distances when available + const dist_t *cached_dists = entity_.get_neighbor_dists(neighbor_id); + if (cached_dists != nullptr) { + // Fast path: read distances from storage, no recomputation needed + for (uint32_t i = 0; i < current_size; ++i) { + prune_candidates.emplace(current_neighbors[i], cached_dists[i]); + } + } else { + // Fallback: compute distances (first time or dist storage not loaded) + const void *neighbor_vec = entity_.get_vector(neighbor_id); + if (ailego_unlikely(neighbor_vec == nullptr)) return; + for (uint32_t i = 0; i < current_size; ++i) { + node_id_t nbr = current_neighbors[i]; + const void *nbr_vec = entity_.get_vector(nbr); + if (ailego_unlikely(nbr_vec == nullptr)) continue; + dist_t nbr_dist = dc.dist(neighbor_vec, nbr_vec); + prune_candidates.emplace(nbr, nbr_dist); + } + } + + // Add the new reverse link + prune_candidates.emplace(id, dist); + + // RobustPrune from neighbor_id's perspective + robust_prune(neighbor_id, prune_candidates, entity_.alpha(), max_deg, ctx); + + // Update neighbor_id's neighbor list and distances + const auto &prune_result = ctx->prune_result(); + entity_.update_neighbors(neighbor_id, prune_result); + entity_.update_neighbor_dists(neighbor_id, prune_result); +} + +// Explicit template instantiation for all entity types +template class VamanaAlgorithm; +template class VamanaAlgorithm; +template class VamanaAlgorithm; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_algorithm.h b/src/core/algorithm/vamana/vamana_algorithm.h new file mode 100644 index 0000000..d8e0c93 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_algorithm.h @@ -0,0 +1,116 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include "vamana_context.h" +#include "vamana_dist_calculator.h" +#include "vamana_streamer_entity.h" + +namespace zvec { +namespace core { + +// Non-template base class providing a type-erased interface so that +// VamanaStreamer can hold a pointer without knowing the EntityType. +class VamanaAlgorithmBase { + public: + typedef std::unique_ptr UPointer; + + virtual ~VamanaAlgorithmBase() = default; + + virtual int cleanup() = 0; + + // Insert a new node into the Vamana graph. + // The node's vector must already be stored in the entity. + virtual int add_node(node_id_t id, VamanaContext *ctx) = 0; + + // Greedy search: find approximate nearest neighbors. + virtual int search(VamanaContext *ctx) const = 0; + + virtual int init() = 0; +}; + +// Vamana graph algorithm, templated on EntityType for hot-path optimization. +// EntityType should be VamanaMmapStreamerEntity, +// VamanaBufferPoolStreamerEntity, or VamanaContiguousStreamerEntity. +// +// Core operations: +// - GreedySearch: beam search from entry point, expanding best candidates +// - RobustPrune: select diverse neighbors using alpha-based pruning +// - add_node: insert + prune + reverse-link update +template +class VamanaAlgorithm : public VamanaAlgorithmBase { + public: + using MemBlockType = typename EntityType::MemoryBlock; + + explicit VamanaAlgorithm(EntityType &entity) + : entity_(entity), lock_pool_(kLockCnt) {} + + ~VamanaAlgorithm() override = default; + + int cleanup() override { + return 0; + } + + int init() override { + return 0; + } + + // Insert node `id` into the graph. Its vector must already be in the entity. + int add_node(node_id_t id, VamanaContext *ctx) override; + + // Greedy search from entry point. Results are stored in ctx->topk_heap(). + int search(VamanaContext *ctx) const override; + + private: + // GreedySearch: starting from entry_point, greedily expand the closest + // unvisited candidate until the search list is exhausted or scan limit + // is reached. Results accumulate in topk_heap. + void greedy_search(node_id_t entry_point, VamanaContext *ctx, + bool use_pool) const; + + // RobustPrune: given a candidate set (topk_heap), select up to max_degree + // diverse neighbors using alpha-based distance comparison. + // Result is stored in ctx->prune_result(). + void robust_prune(node_id_t id, TopkHeap &candidates, float alpha, + uint32_t max_degree, VamanaContext *ctx) const; + + // Update node's neighbors and handle reverse links. + void update_neighbors_and_reverse_links( + node_id_t id, + const std::vector> &new_neighbors, + VamanaContext *ctx); + + // Check if adding `id` as a reverse neighbor of `neighbor_id` requires + // pruning, and if so, prune neighbor_id's neighbor list. + void reverse_update_neighbor(node_id_t id, node_id_t neighbor_id, dist_t dist, + VamanaContext *ctx); + + private: + VamanaAlgorithm(const VamanaAlgorithm &) = delete; + VamanaAlgorithm &operator=(const VamanaAlgorithm &) = delete; + + static constexpr uint32_t kLockCnt{1U << 8}; + static constexpr uint32_t kLockMask{kLockCnt - 1U}; + + EntityType &entity_; + mutable ailego::SpinMutex spin_lock_{}; + std::vector lock_pool_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_context.cc b/src/core/algorithm/vamana/vamana_context.cc new file mode 100644 index 0000000..ee52c51 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_context.cc @@ -0,0 +1,176 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "vamana_context.h" +#include +#include "vamana_params.h" + +namespace zvec { +namespace core { + +VamanaContext::VamanaContext(size_t dimension, + const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity) + : IndexContext(metric), + entity_(entity), + dc_(entity.get(), metric, dimension), + metric_(metric) { + if (metric) { + build_distance_offset_ = metric->build_distance_offset(); + } +} + +VamanaContext::VamanaContext(const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity) + : IndexContext(metric), + entity_(entity), + dc_(entity.get(), metric), + metric_(metric) { + if (metric) { + build_distance_offset_ = metric->build_distance_offset(); + } +} + +VamanaContext::~VamanaContext() { + visit_filter_.destroy(); +} + +int VamanaContext::init(ContextType type) { + int ret; + uint32_t doc_cnt; + + type_ = type; + results_.resize(1); + topk_heap_.limit(std::max(topk_, ef_)); + update_heap_.limit(entity_->max_degree()); + + switch (type) { + case kBuilderContext: + ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(), + max_scan_num_, filter_negative_prob_); + if (ret != 0) { + LOG_ERROR("Create visit filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + break; + + case kSearcherContext: + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_, + filter_negative_prob_); + if (ret != 0) { + LOG_ERROR("Create visit filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + break; + + case kStreamerContext: + doc_cnt = entity_->doc_cnt(); + max_scan_num_ = compute_max_scan_num(doc_cnt); + reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt); + ret = visit_filter_.init(filter_mode_, reserve_max_doc_cnt_, + max_scan_num_, filter_negative_prob_); + if (ret != 0) { + LOG_ERROR("Create visit filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + check_need_adjuct_ctx(); + break; + + default: + break; + } + + return 0; +} + +int VamanaContext::update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity, + uint32_t magic_num) { + if (magic_ == magic_num) { + return 0; + } + type_ = type; + entity_ = entity; + metric_ = metric; + magic_ = magic_num; + if (metric) { + build_distance_offset_ = metric->build_distance_offset(); + } + dc_.update(entity.get(), metric, meta.dimension()); + return 0; +} + +int VamanaContext::update(const ailego::Params ¶ms) { + uint32_t ef = ef_; + params.get(PARAM_VAMANA_STREAMER_EF, &ef); + ef_ = ef; + topk_heap_.limit(std::max(topk_, ef_)); + uint32_t po = po_; + params.get(PARAM_VAMANA_STREAMER_PO, &po); + po_ = po; + uint32_t pl = pl_; + params.get(PARAM_VAMANA_STREAMER_PL, &pl); + pl_ = pl; + return 0; +} + +void VamanaContext::topk_to_result(uint32_t idx) { + if (force_padding_topk_ && !topk_heap_.full() && + topk_heap_.size() < entity_->doc_cnt()) { + this->fill_random_to_topk_full(); + } + if (ailego_unlikely(topk_heap_.size() == 0)) { + return; + } + + ailego_assert_with(idx < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(topk_heap_.size())); + topk_heap_.sort(); + results_[idx].clear(); + + for (int i = 0; i < size; ++i) { + auto score = topk_heap_[i].second; + if (score > this->threshold()) { + break; + } + node_id_t id = topk_heap_[i].first; + if (fetch_vector_) { + results_[idx].emplace_back(entity_->get_key(id), score, id, + entity_->get_vector(id)); + } else { + results_[idx].emplace_back(entity_->get_key(id), score, id); + } + } +} + +void VamanaContext::fill_random_to_topk_full() { + std::mt19937 rng(42); + uint32_t doc_cnt = entity_->doc_cnt(); + uint32_t max_attempts = doc_cnt * 2; + uint32_t attempts = 0; + while (!topk_heap_.full() && doc_cnt > 0 && attempts < max_attempts) { + node_id_t random_id = rng() % doc_cnt; + if (entity_->get_key(random_id) != kInvalidKey) { + dist_t random_dist = dc_.dist(random_id); + topk_heap_.emplace_back(random_id, random_dist); + } + ++attempts; + } +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_context.h b/src/core/algorithm/vamana/vamana_context.h new file mode 100644 index 0000000..a576353 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_context.h @@ -0,0 +1,350 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "utility/block_heap.h" +#include "utility/linear_pool.h" +#include "utility/visit_filter.h" +#include "vamana_dist_calculator.h" +#include "vamana_entity.h" + +namespace zvec { +namespace core { + +class VamanaContext : public IndexContext { + public: + typedef std::unique_ptr Pointer; + + enum ContextType { + kUnknownContext = 0, + kSearcherContext = 1, + kBuilderContext = 2, + kStreamerContext = 3 + }; + + VamanaContext(size_t dimension, const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity); + + VamanaContext(const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity); + + ~VamanaContext() override; + + void set_topk(uint32_t val) override { + topk_ = val; + topk_heap_.limit(std::max(val, ef_)); + } + + const IndexDocumentList &result(void) const override { + return results_[0]; + } + + const IndexDocumentList &result(size_t idx) const override { + return results_[idx]; + } + + IndexDocumentList *mutable_result(size_t idx) override { + ailego_assert_with(idx < results_.size(), "invalid idx"); + return &results_[idx]; + } + + uint32_t magic(void) const override { + return magic_; + } + + void set_debug_mode(bool enable) override { + debug_mode_ = enable; + } + bool debug_mode(void) const override { + return debug_mode_; + } + + std::string debug_string(void) const override { + char buf[4096]; + size_t size = snprintf(buf, sizeof(buf), "scan_cnt=%zu", get_scan_num()); + return std::string(buf, size); + } + + int update(const ailego::Params ¶ms) override; + + int init(ContextType type); + + int update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity, uint32_t magic_num); + + inline const VamanaEntity &get_entity() const { + return *entity_; + } + + inline void resize_results(size_t size) { + results_.resize(size); + } + + inline void topk_to_result() { + topk_to_result(0); + } + + void topk_to_result(uint32_t idx); + + inline void reset_query(const void *query) { + if (auto query_preprocess_func = index_metric_->get_query_preprocess_func(); + query_preprocess_func != nullptr) { + size_t dim = dc_.dimension(); + preprocess_buffer_.resize(dim); + memcpy(preprocess_buffer_.data(), query, dim); + query_preprocess_func(preprocess_buffer_.data(), dim); + query = preprocess_buffer_.data(); + } + dc_.reset_query(query); + dc_.clear_compare_cnt(); + } + + inline VamanaDistCalculator &dist_calculator() { + return dc_; + } + inline TopkHeap &topk_heap() { + return topk_heap_; + } + inline TopkHeap &update_heap() { + return update_heap_; + } + inline LinearPool &pool() { + return pool_; + } + // Block-insert pool used by the AVX2-gated greedy_search fast path. + // Only accessed under a runtime CpuFeatures::AVX2 guard at call sites. + inline BlockHeap &block_pool() { + return block_pool_; + } + inline VisitFilter &visit_filter() { + return visit_filter_; + } + inline CandidateHeap &candidates() { + return candidates_; + } + + // Pre-allocated buffers for robust_prune optimization + inline std::vector &prune_vec_cache() { + return prune_vec_cache_; + } + inline std::vector &prune_active() { + return prune_active_; + } + inline std::vector &prune_occlude_factor() { + return prune_occlude_factor_; + } + inline std::vector> &prune_result() { + return prune_result_; + } + inline std::vector &batch_vecs_buf() { + return batch_vecs_buf_; + } + inline std::vector &batch_dists_buf() { + return batch_dists_buf_; + } + inline std::vector &batch_indices_buf() { + return batch_indices_buf_; + } + + //! Build-time distance offset cached from the metric. Used by RobustPrune + //! to shift the internal distance to a non-negative range before computing + //! the ratio-based occlude_factor. Zero for metrics whose internal distance + //! is already non-negative (e.g. SquaredEuclidean). + inline float build_distance_offset() const { + return build_distance_offset_; + } + + inline void set_max_scan_num(uint32_t max_scan_num) { + max_scan_num_ = max_scan_num; + } + inline void set_ef(uint32_t v) { + ef_ = v; + } + + inline uint32_t ef() const { + return ef_; + } + inline void set_po(uint32_t v) { + po_ = v; + } + + inline uint32_t po() const { + return po_; + } + + inline void set_pl(uint32_t v) { + pl_ = v; + } + + inline uint32_t pl() const { + return pl_; + } + inline void set_max_scan_ratio(float v) { + max_scan_ratio_ = v; + } + virtual void set_magic(uint32_t v) { + magic_ = v; + } + virtual void set_force_padding_topk(bool v) { + force_padding_topk_ = v; + } + + void set_bruteforce_threshold(uint32_t v) override { + bruteforce_threshold_ = v; + } + inline uint32_t get_bruteforce_threshold() const { + return bruteforce_threshold_; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + bool fetch_vector() const override { + return fetch_vector_; + } + + void set_max_scan_limit(size_t v) { + max_scan_limit_ = v; + } + void set_min_scan_limit(size_t v) { + min_scan_limit_ = v; + } + + void set_filter_mode(VisitFilter::Mode mode) { + filter_mode_ = mode; + } + void set_filter_negative_probability(float prob) { + filter_negative_prob_ = prob; + } + + void reset(void) override { + dc_.clear(); + for (auto &it : results_) { + it.clear(); + } + IndexContext::reset_filter(); + IndexContext::reset_threshold(); + IndexContext::set_fetch_vector(false); + } + + inline void check_need_adjuct_ctx(void) { + check_need_adjuct_ctx(entity_->doc_cnt()); + } + + inline void check_need_adjuct_ctx(uint32_t doc_cnt) { + if (ailego_unlikely(doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_)) { + while (doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_) { + reserve_max_doc_cnt_ = + reserve_max_doc_cnt_ + compute_reserve_cnt(reserve_max_doc_cnt_); + } + uint32_t max_scan_cnt = compute_max_scan_num(reserve_max_doc_cnt_); + max_scan_num_ = max_scan_cnt; + visit_filter_.reset(reserve_max_doc_cnt_, max_scan_cnt); + candidates_.clear(); + candidates_.limit(max_scan_num_); + } + } + + inline size_t get_scan_num() const { + return dc_.compare_cnt(); + } + + inline uint64_t reach_scan_limit() const { + return dc_.compare_cnt() >= max_scan_num_; + } + + inline bool error() const { + return dc_.error(); + } + + inline void clear() { + dc_.clear(); + for (auto &it : results_) { + it.clear(); + } + } + + inline uint32_t topk() const override { + return topk_; + } + + private: + void fill_random_to_topk_full(void); + + inline size_t compute_reserve_cnt(uint32_t cur_doc) const { + if (cur_doc > kMaxReserveDocCnt) return kMaxReserveDocCnt; + if (cur_doc < kMinReserveDocCnt) return kMinReserveDocCnt; + return cur_doc; + } + + inline uint32_t compute_max_scan_num(uint32_t max_doc_cnt) const { + uint32_t max_scan = max_doc_cnt * max_scan_ratio_; + if (max_scan < min_scan_limit_) max_scan = min_scan_limit_; + if (max_scan > max_scan_limit_) max_scan = max_scan_limit_; + return max_scan; + } + + constexpr static uint32_t kTriggerReserveCnt = 4096UL; + constexpr static uint32_t kMinReserveDocCnt = 4096UL; + constexpr static uint32_t kMaxReserveDocCnt = 128 * 1024UL; + + VamanaEntity::Pointer entity_; + VamanaDistCalculator dc_; + IndexMetric::Pointer metric_; + + bool debug_mode_{false}; + bool force_padding_topk_{false}; + uint32_t max_scan_num_{0}; + uint32_t reserve_max_doc_cnt_{kMinReserveDocCnt}; + uint32_t topk_{0}; + uint32_t ef_{VamanaEntity::kDefaultEf}; + uint32_t po_{8}; + uint32_t pl_{0}; + float max_scan_ratio_{VamanaEntity::kDefaultScanRatio}; + size_t max_scan_limit_{VamanaEntity::kDefaultMaxScanLimit}; + size_t min_scan_limit_{VamanaEntity::kDefaultMinScanLimit}; + uint32_t magic_{0U}; + std::vector results_{}; + TopkHeap topk_heap_{}; + TopkHeap update_heap_{}; + CandidateHeap candidates_{}; + VisitFilter visit_filter_{}; + uint32_t bruteforce_threshold_{}; + bool fetch_vector_{false}; + uint32_t type_{kUnknownContext}; + std::string preprocess_buffer_; + + // Pre-allocated buffers for robust_prune optimization + std::vector prune_vec_cache_; + std::vector prune_active_; + std::vector prune_occlude_factor_; + std::vector> prune_result_; + std::vector batch_vecs_buf_; + std::vector batch_dists_buf_; + std::vector batch_indices_buf_; + + //! Cached build-time distance offset (see build_distance_offset()). + float build_distance_offset_{0.0f}; + + VisitFilter::Mode filter_mode_{VisitFilter::ByteMap}; + float filter_negative_prob_{VamanaEntity::kDefaultBFNegativeProbability}; + + LinearPool pool_; + BlockHeap block_pool_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_dist_calculator.h b/src/core/algorithm/vamana/vamana_dist_calculator.h new file mode 100644 index 0000000..ed15b0b --- /dev/null +++ b/src/core/algorithm/vamana/vamana_dist_calculator.h @@ -0,0 +1,174 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include "vamana_entity.h" + +namespace zvec { +namespace core { + +class VamanaDistCalculator { + public: + typedef std::shared_ptr Pointer; + + VamanaDistCalculator(const VamanaEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(nullptr), + dim_(dim), + compare_cnt_(0) {} + + VamanaDistCalculator(const VamanaEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim, + const void *query) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(query), + dim_(dim), + compare_cnt_(0) {} + + VamanaDistCalculator(const VamanaEntity *entity, + const IndexMetric::Pointer &metric) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(nullptr), + dim_(0), + compare_cnt_(0) {} + + void update(const VamanaEntity *entity, const IndexMetric::Pointer &metric) { + entity_ = entity; + distance_ = metric->distance(); + batch_distance_ = metric->batch_distance(); + } + + void update(const VamanaEntity *entity, const IndexMetric::Pointer &metric, + uint32_t dim) { + entity_ = entity; + distance_ = metric->distance(); + batch_distance_ = metric->batch_distance(); + dim_ = dim; + } + + inline void reset_query(const void *query) { + error_ = false; + query_ = query; + } + + inline dist_t dist(const void *vec_lhs, const void *vec_rhs) { + if (ailego_unlikely(vec_lhs == nullptr || vec_rhs == nullptr)) { + LOG_ERROR("Nullptr of dense vector"); + error_ = true; + return 0.0f; + } + float score{0.0f}; + distance_(vec_lhs, vec_rhs, dim_, &score); + return score; + } + + inline dist_t dist(const void *vec) { + compare_cnt_++; + return dist(vec, query_); + } + + inline dist_t dist(node_id_t id) { + compare_cnt_++; + const void *feat = entity_->get_vector(id); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + return dist(feat, query_); + } + + inline dist_t dist(node_id_t lhs, node_id_t rhs) { + compare_cnt_++; + const void *feat = entity_->get_vector(lhs); + const void *query = entity_->get_vector(rhs); + if (ailego_unlikely(feat == nullptr || query == nullptr)) { + LOG_ERROR("Get nullptr vector"); + error_ = true; + return 0.0f; + } + return dist(feat, query); + } + + inline void batch_dist(const void **vecs, uint32_t count, float *dists) { + compare_cnt_ += count; + batch_distance_(vecs, query_, count, dim_, dists); + } + + // Single-node batch distance: compute distance between query and a stored + // node using batch_distance_. Consistent with HnswDistCalculator::batch_dist. + inline dist_t batch_dist(node_id_t id) { + compare_cnt_++; + const void *feat = entity_->get_vector(id); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + dist_t score = 0; + batch_distance_(&feat, query_, 1, dim_, &score); + return score; + } + + dist_t operator()(const void *vec) { + return dist(vec); + } + dist_t operator()(node_id_t i) { + return dist(i); + } + dist_t operator()(node_id_t lhs, node_id_t rhs) { + return dist(lhs, rhs); + } + + inline void clear() { + compare_cnt_ = 0; + error_ = false; + } + + inline void clear_compare_cnt() { + compare_cnt_ = 0; + } + inline bool error() const { + return error_; + } + inline uint32_t compare_cnt() const { + return compare_cnt_; + } + inline uint32_t dimension() const { + return dim_; + } + + private: + VamanaDistCalculator(const VamanaDistCalculator &) = delete; + VamanaDistCalculator &operator=(const VamanaDistCalculator &) = delete; + + const VamanaEntity *entity_; + IndexMetric::MatrixDistance distance_; + IndexMetric::MatrixBatchDistance batch_distance_; + const void *query_; + uint32_t dim_; + uint32_t compare_cnt_; + bool error_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_entity.cc b/src/core/algorithm/vamana/vamana_entity.cc new file mode 100644 index 0000000..4114dc6 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_entity.cc @@ -0,0 +1,203 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "vamana_entity.h" +#include + +namespace zvec { +namespace core { + +const std::string VamanaEntity::kGraphHeaderSegmentId = "vamana.graph.header"; +const std::string VamanaEntity::kGraphFeaturesSegmentId = + "vamana.graph.features"; +const std::string VamanaEntity::kGraphKeysSegmentId = "vamana.graph.keys"; +const std::string VamanaEntity::kGraphNeighborsSegmentId = + "vamana.graph.neighbors"; +const std::string VamanaEntity::kGraphOffsetsSegmentId = "vamana.graph.offsets"; +const std::string VamanaEntity::kGraphMappingSegmentId = "vamana.graph.mapping"; +const std::string VamanaEntity::kGraphNeighborDistsSegmentId = + "vamana.graph.neighbor_dists"; + +int VamanaEntity::CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size) { + *padding_size = AlignSize(data_size) - data_size; + if (*padding_size == 0) { + return 0; + } + std::string padding(*padding_size, '\0'); + if (dumper->write(padding.data(), *padding_size) != *padding_size) { + LOG_ERROR("Append padding failed, size %lu", *padding_size); + return IndexError_WriteData; + } + return 0; +} + +int64_t VamanaEntity::dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, + const void *data, size_t size) const { + size_t len = dumper->write(data, size); + if (len != size) { + LOG_ERROR("Dump segment %s data failed, expect: %lu, actual: %lu", + segment_id.c_str(), size, len); + return IndexError_WriteData; + } + + size_t padding_size = AlignSize(size) - size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + uint32_t crc = ailego::Crc32c::Hash(data, size); + int ret = dumper->append(segment_id, size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", segment_id.c_str(), ret); + return ret; + } + + return len + padding_size; +} + +int64_t VamanaEntity::dump_header(const IndexDumper::Pointer &dumper, + const VamanaHeader &hd) const { + return dump_segment(dumper, kGraphHeaderSegmentId, &hd.graph, hd.graph.size); +} + +int64_t VamanaEntity::dump_vectors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + size_t total_size = doc_cnt() * vector_size(); + std::vector buffer(total_size); + + for (node_id_t i = 0; i < doc_cnt(); ++i) { + node_id_t old_id = reorder_mapping[i]; + const void *vec = get_vector(old_id); + if (vec == nullptr) { + LOG_ERROR("Get vector failed for node %u", old_id); + return IndexError_ReadData; + } + memcpy(buffer.data() + static_cast(i) * vector_size(), vec, + vector_size()); + } + + return dump_segment(dumper, kGraphFeaturesSegmentId, buffer.data(), + total_size); +} + +int64_t VamanaEntity::dump_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + size_t nbr_size = neighbors_size(); + size_t total_size = doc_cnt() * nbr_size; + std::vector buffer(total_size, 0); + + for (node_id_t i = 0; i < doc_cnt(); ++i) { + node_id_t old_id = reorder_mapping[i]; + const Neighbors nbrs = get_neighbors(old_id); + + auto *hd = reinterpret_cast( + buffer.data() + static_cast(i) * nbr_size); + hd->neighbor_cnt = nbrs.size(); + for (uint32_t j = 0; j < nbrs.size(); ++j) { + hd->neighbors[j] = neighbor_mapping[nbrs[j]]; + } + } + + return dump_segment(dumper, kGraphNeighborsSegmentId, buffer.data(), + total_size); +} + +int64_t VamanaEntity::dump_neighbor_dists( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + if (!dist_storage_loaded()) { + // No distance data to dump — this is fine for read-only indices + return 0; + } + + uint32_t max_deg = static_cast(max_degree()); + size_t dist_entry = max_deg * sizeof(dist_t); + size_t total_size = doc_cnt() * dist_entry; + std::vector buffer(total_size, 0); + + for (node_id_t i = 0; i < doc_cnt(); ++i) { + node_id_t old_id = reorder_mapping[i]; + const dist_t *dists = get_neighbor_dists(old_id); + if (dists != nullptr) { + memcpy(buffer.data() + static_cast(i) * dist_entry, dists, + dist_entry); + } + } + + return dump_segment(dumper, kGraphNeighborDistsSegmentId, buffer.data(), + total_size); +} + +int64_t VamanaEntity::dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const { + size_t total_size = doc_cnt() * sizeof(key_t); + return dump_segment(dumper, kGraphMappingSegmentId, keys, total_size); +} + +void VamanaEntity::reshuffle_vectors(std::vector *n2o_mapping, + std::vector *o2n_mapping, + key_t *keys) const { + uint32_t count = doc_cnt(); + n2o_mapping->resize(count); + o2n_mapping->resize(count); + + // Simple identity mapping for now (can be optimized with BFS traversal) + for (uint32_t i = 0; i < count; ++i) { + (*n2o_mapping)[i] = i; + (*o2n_mapping)[i] = i; + keys[i] = get_key(i); + } +} + +int64_t VamanaEntity::dump_segments(const IndexDumper::Pointer &dumper, + key_t *keys) const { + std::vector n2o_mapping; + std::vector o2n_mapping; + reshuffle_vectors(&n2o_mapping, &o2n_mapping, keys); + + // Remap entry point + VamanaHeader dump_header_copy = header_; + if (dump_header_copy.graph.entry_point != kInvalidNodeId) { + dump_header_copy.graph.entry_point = + o2n_mapping[dump_header_copy.graph.entry_point]; + } + + int64_t hd_size = dump_header(dumper, dump_header_copy); + if (hd_size < 0) return hd_size; + + int64_t vec_size = dump_vectors(dumper, n2o_mapping); + if (vec_size < 0) return vec_size; + + int64_t nbr_size = dump_neighbors(dumper, n2o_mapping, o2n_mapping); + if (nbr_size < 0) return nbr_size; + + int64_t map_size = dump_mapping_segment(dumper, keys); + if (map_size < 0) return map_size; + + int64_t dist_size = dump_neighbor_dists(dumper, n2o_mapping); + if (dist_size < 0) return dist_size; + + return hd_size + vec_size + nbr_size + map_size + dist_size; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_entity.h b/src/core/algorithm/vamana/vamana_entity.h new file mode 100644 index 0000000..976d385 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_entity.h @@ -0,0 +1,366 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +// Reuse typed MemoryBlock and NeighborsT from hnsw_entity.h +#include "algorithm/hnsw/hnsw_entity.h" + +namespace zvec { +namespace core { + +// Vamana graph header — single-layer graph (no hierarchical levels) +struct VamanaGraphHeader { + uint32_t size; + uint32_t version; + uint32_t graph_type; + uint32_t doc_count; + uint32_t vector_size; + uint32_t node_size; + uint32_t max_degree; // R: maximum out-degree + uint32_t search_list_size; // L: search list size for construction + uint32_t max_occlusion_size; // C: max candidate size for RobustPrune + uint32_t entry_point; // medoid node id + uint32_t options; + uint32_t reserved_pad; + float alpha; // alpha parameter for RobustPrune + uint8_t reserved_[4076]; +}; + +static_assert(sizeof(VamanaGraphHeader) % 32 == 0, + "VamanaGraphHeader must be aligned with 32 bytes"); + +struct VamanaHeader { + VamanaHeader() { + clear(); + } + + VamanaHeader(const VamanaHeader &header) { + memcpy(static_cast(this), &header, sizeof(header)); + } + + VamanaHeader &operator=(const VamanaHeader &header) { + memcpy(static_cast(this), &header, sizeof(header)); + return *this; + } + + void inline reset() { + graph.doc_count = 0U; + graph.entry_point = kInvalidNodeId; + } + + void inline clear() { + memset(static_cast(this), 0, sizeof(VamanaHeader)); + graph.entry_point = kInvalidNodeId; + graph.size = sizeof(VamanaGraphHeader); + graph.alpha = 1.2f; + } + + size_t max_degree() const { + return graph.max_degree; + } + size_t vector_size() const { + return graph.vector_size; + } + size_t search_list_size() const { + return graph.search_list_size; + } + size_t max_occlusion_size() const { + return graph.max_occlusion_size; + } + float alpha() const { + return graph.alpha; + } + node_id_t entry_point() const { + return graph.entry_point; + } + node_id_t doc_cnt() const { + return graph.doc_count; + } + + VamanaGraphHeader graph; +}; + +// VamanaEntity: base class for Vamana graph data management +class VamanaEntity { + public: + VamanaEntity() {} + VamanaEntity(const VamanaHeader &hd) { + header_ = hd; + } + virtual ~VamanaEntity() {} + + typedef std::shared_ptr Pointer; + + // Options bit flags (stored in VamanaGraphHeader::options) + static constexpr uint32_t kOptionSaturateGraph = 1U << 0; + + // Default constants + static constexpr uint32_t kDefaultMaxDegree = 64; + static constexpr uint32_t kDefaultSearchListSize = 100; + static constexpr uint32_t kDefaultMaxOcclusionSize = 750; + static constexpr float kDefaultAlpha = 1.2f; + static constexpr bool kDefaultSaturateGraph = false; + static constexpr uint32_t kDefaultEf = 200; + static constexpr float kDefaultScanRatio = 0.1f; + static constexpr uint32_t kDefaultBruteForceThreshold = 1000U; + static constexpr uint32_t kDefaultDocsHardLimit = 1 << 30U; + static constexpr float kDefaultDocsSoftLimitRatio = 0.9f; + static constexpr size_t kMaxChunkSize = 0xFFFFFFFF; + static constexpr size_t kDefaultChunkSize = 2UL * 1024UL * 1024UL; + static constexpr size_t kDefaultMaxChunkCnt = 50000UL; + static constexpr uint32_t kDefaultMinScanLimit = 10000; + static constexpr uint32_t kDefaultMaxScanLimit = + std::numeric_limits::max(); + static constexpr float kDefaultBFNegativeProbability = 0.001f; + + inline size_t max_degree() const { + return header_.graph.max_degree; + } + inline node_id_t *mutable_doc_cnt() { + return &header_.graph.doc_count; + } + inline node_id_t doc_cnt() const { + return header_.graph.doc_count; + } + inline float alpha() const { + return header_.graph.alpha; + } + inline size_t search_list_size() const { + return header_.graph.search_list_size; + } + inline size_t max_occlusion_size() const { + return header_.graph.max_occlusion_size; + } + inline node_id_t entry_point() const { + return header_.graph.entry_point; + } + inline size_t vector_size() const { + return header_.graph.vector_size; + } + inline size_t node_size() const { + return header_.graph.node_size; + } + + void set_vector_size(size_t size) { + header_.graph.vector_size = size; + } + void set_max_degree(uint32_t val) { + header_.graph.max_degree = val; + } + void set_search_list_size(uint32_t val) { + header_.graph.search_list_size = val; + } + void set_max_occlusion_size(uint32_t val) { + header_.graph.max_occlusion_size = val; + } + void set_alpha(float val) { + header_.graph.alpha = val; + } + + inline bool saturate_graph() const { + return (header_.graph.options & kOptionSaturateGraph) != 0; + } + void set_saturate_graph(bool val) { + if (val) { + header_.graph.options |= kOptionSaturateGraph; + } else { + header_.graph.options &= ~kOptionSaturateGraph; + } + } + + // Neighbor size: NeighborsHeader + max_degree * sizeof(node_id_t) + inline size_t neighbors_size() const { + return sizeof(NeighborsHeader) + max_degree() * sizeof(node_id_t); + } + + virtual void update_entry_point(node_id_t ep) { + header_.graph.entry_point = ep; + } + + // Calculate medoid (entry point) as the data point closest to the centroid + // of all vectors, following DiskANN's standard approach. + // Parameters: + // dimension: vector dimension (number of elements per vector) + // data_type: IndexMeta::DataType value (e.g. DT_FP32=2, DT_INT8=4, + // DT_FP16=1) + // Returns the medoid node ID, or kInvalidNodeId if no valid data. + virtual node_id_t calculate_medoid(uint32_t /*dimension*/, + uint32_t /*data_type*/) { + return kInvalidNodeId; + } + + virtual int cleanup() { + header_.clear(); + return 0; + } + + virtual const VamanaEntity::Pointer clone() const { + return VamanaEntity::Pointer(); + } + + // Pure virtual interface + virtual key_t get_key(node_id_t id) const = 0; + virtual const void *get_vector(node_id_t id) const = 0; + virtual int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const = 0; + virtual int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const = 0; + virtual int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const = 0; + virtual const Neighbors get_neighbors(node_id_t id) const = 0; + + virtual int add_vector(key_t /*key*/, const void * /*vec*/, + node_id_t * /*id*/) { + return IndexError_NotImplemented; + } + virtual int add_vector_with_id(node_id_t /*id*/, const void * /*vec*/) { + return IndexError_NotImplemented; + } + virtual int update_neighbors( + node_id_t /*id*/, + const std::vector> & /*neighbors*/) { + return IndexError_NotImplemented; + } + virtual void add_neighbor(node_id_t /*id*/, uint32_t /*size*/, + node_id_t /*neighbor_id*/) {} + + // --- Neighbor distance storage (CSR-like, lazy-loaded) --- + // Each node has max_degree dist_t slots, the i-th slot stores the distance + // from this node to its i-th neighbor. Only allocated/loaded when needed + // (first write operation). Search-only paths never touch this data. + + // Ensure distance storage is allocated/loaded. Must be called before + // any get/set neighbor dist operations. Thread-safe (idempotent). + virtual int ensure_dist_storage() { + return 0; + } + + // Whether distance storage is currently loaded + virtual bool dist_storage_loaded() const { + return false; + } + + // Get pointer to the distance array for node `id`. + // Returns nullptr if dist storage is not loaded. + virtual const dist_t *get_neighbor_dists(node_id_t /*id*/) const { + return nullptr; + } + + // Update all neighbor distances for node `id` from a prune result. + virtual void update_neighbor_dists( + node_id_t /*id*/, + const std::vector> & /*neighbors*/) {} + + // Set the distance for the `idx`-th neighbor of node `id`. + virtual void set_neighbor_dist(node_id_t /*id*/, uint32_t /*idx*/, + dist_t /*dist*/) {} + + virtual int dump(const IndexDumper::Pointer & /*dumper*/) { + return IndexError_NotImplemented; + } + + virtual const void *get_vector_by_key(uint64_t /*key*/) const { + return nullptr; + } + virtual int get_vector_by_key(const key_t /*key*/, + IndexStorage::MemoryBlock & /*block*/) const { + return IndexError_NotImplemented; + } + + static int CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size); + + protected: + inline const VamanaHeader &header() const { + return header_; + } + inline VamanaHeader *mutable_header() { + return &header_; + } + inline size_t header_size() const { + return sizeof(header_); + } + + void set_node_size(size_t size) { + header_.graph.node_size = size; + } + + int64_t dump_segments(const IndexDumper::Pointer &dumper, key_t *keys) const; + + int64_t dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, const void *data, + size_t size) const; + + int64_t dump_header(const IndexDumper::Pointer &dumper, + const VamanaHeader &hd) const; + + int64_t dump_vectors(const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + int64_t dump_neighbors(const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + int64_t dump_neighbor_dists( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const; + + void reshuffle_vectors(std::vector *n2o_mapping, + std::vector *o2n_mapping, + key_t *keys) const; + + static inline size_t AlignSize(size_t size) { + return (size + 0x1F) & (~0x1F); + } + + static inline size_t AlignPageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::PageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + static inline size_t AlignHugePageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::HugePageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + public: + const static std::string kGraphHeaderSegmentId; + const static std::string kGraphFeaturesSegmentId; + const static std::string kGraphKeysSegmentId; + const static std::string kGraphNeighborsSegmentId; + const static std::string kGraphOffsetsSegmentId; + const static std::string kGraphMappingSegmentId; + const static std::string kGraphNeighborDistsSegmentId; + + static constexpr uint32_t kRevision = 0U; + + protected: + VamanaHeader header_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_index_provider.h b/src/core/algorithm/vamana/vamana_index_provider.h new file mode 100644 index 0000000..ab70ba5 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_index_provider.h @@ -0,0 +1,117 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "vamana_entity.h" + +namespace zvec { +namespace core { + +class VamanaIndexProvider : public IndexProvider { + public: + VamanaIndexProvider(const IndexMeta &meta, + const VamanaEntity::Pointer &entity, + const std::string &owner) + : meta_(meta), entity_(entity), owner_class_(owner) {} + + VamanaIndexProvider(const VamanaIndexProvider &) = delete; + VamanaIndexProvider &operator=(const VamanaIndexProvider &) = delete; + + public: + IndexProvider::Iterator::Pointer create_iterator() override { + return VamanaIndexProvider::Iterator::Pointer(new (std::nothrow) + Iterator(entity_)); + } + + size_t count(void) const override { + return entity_->doc_cnt(); + } + + size_t dimension(void) const override { + return meta_.dimension(); + } + + IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + size_t element_size(void) const override { + return meta_.element_size(); + } + + const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector_by_key(key, block); + } + + const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class Iterator : public IndexProvider::Iterator { + public: + Iterator(const VamanaEntity::Pointer &entity) + : entity_(entity), cur_id_(0U) { + cur_id_ = get_next_valid_id(0); + } + + const void *data(void) const override { + return entity_->get_vector(cur_id_); + } + + bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt(); + } + + uint64_t key(void) const override { + return entity_->get_key(cur_id_); + } + + void next(void) override { + cur_id_ = get_next_valid_id(cur_id_ + 1); + } + + void reset(void) { + cur_id_ = get_next_valid_id(0); + } + + private: + node_id_t get_next_valid_id(node_id_t start_id) { + for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) { + if (entity_->get_key(i) != kInvalidNodeId) { + return i; + } + } + return kInvalidNodeId; + } + + const VamanaEntity::Pointer entity_; + node_id_t cur_id_; + }; + + const IndexMeta &meta_; + const VamanaEntity::Pointer entity_; + const std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_params.h b/src/core/algorithm/vamana/vamana_params.h new file mode 100644 index 0000000..4d7df96 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_params.h @@ -0,0 +1,81 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +// Builder parameters +static const std::string PARAM_VAMANA_BUILDER_THREAD_COUNT( + "proxima.vamana.builder.thread_count"); +static const std::string PARAM_VAMANA_BUILDER_MEMORY_QUOTA( + "proxima.vamana.builder.memory_quota"); +static const std::string PARAM_VAMANA_BUILDER_MAX_DEGREE( + "proxima.vamana.builder.max_degree"); +static const std::string PARAM_VAMANA_BUILDER_SEARCH_LIST_SIZE( + "proxima.vamana.builder.search_list_size"); +static const std::string PARAM_VAMANA_BUILDER_ALPHA( + "proxima.vamana.builder.alpha"); +static const std::string PARAM_VAMANA_BUILDER_MAX_OCCLUSION_SIZE( + "proxima.vamana.builder.max_occlusion_size"); + +// Searcher parameters +static const std::string PARAM_VAMANA_SEARCHER_SEARCH_LIST_SIZE( + "proxima.vamana.searcher.search_list_size"); +static const std::string PARAM_VAMANA_SEARCHER_BRUTE_FORCE_THRESHOLD( + "proxima.vamana.searcher.brute_force_threshold"); + +// Streamer parameters +static const std::string PARAM_VAMANA_STREAMER_MAX_DEGREE( + "proxima.vamana.streamer.max_degree"); +static const std::string PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE( + "proxima.vamana.streamer.search_list_size"); +static const std::string PARAM_VAMANA_STREAMER_ALPHA( + "proxima.vamana.streamer.alpha"); +static const std::string PARAM_VAMANA_STREAMER_MAX_OCCLUSION_SIZE( + "proxima.vamana.streamer.max_occlusion_size"); +static const std::string PARAM_VAMANA_STREAMER_EF("proxima.vamana.streamer.ef"); +static const std::string PARAM_VAMANA_STREAMER_PO("proxima.vamana.streamer.po"); +static const std::string PARAM_VAMANA_STREAMER_PL("proxima.vamana.streamer.pl"); +static const std::string PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD( + "proxima.vamana.streamer.brute_force_threshold"); +static const std::string PARAM_VAMANA_STREAMER_MAX_SCAN_RATIO( + "proxima.vamana.streamer.max_scan_ratio"); +static const std::string PARAM_VAMANA_STREAMER_DOCS_HARD_LIMIT( + "proxima.vamana.streamer.docs_hard_limit"); +static const std::string PARAM_VAMANA_STREAMER_DOCS_SOFT_LIMIT( + "proxima.vamana.streamer.docs_soft_limit"); +static const std::string PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE( + "proxima.vamana.streamer.max_index_size"); +static const std::string PARAM_VAMANA_STREAMER_CHUNK_SIZE( + "proxima.vamana.streamer.chunk_size"); +static const std::string PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE( + "proxima.vamana.streamer.get_vector_enable"); +static const std::string PARAM_VAMANA_STREAMER_FORCE_PADDING_RESULT_ENABLE( + "proxima.vamana.streamer.force_padding_result_enable"); +static const std::string PARAM_VAMANA_STREAMER_USE_ID_MAP( + "proxima.vamana.streamer.use_id_map"); +static const std::string PARAM_VAMANA_STREAMER_MAX_SCAN_LIMIT( + "proxima.vamana.streamer.max_scan_limit"); +static const std::string PARAM_VAMANA_STREAMER_MIN_SCAN_LIMIT( + "proxima.vamana.streamer.min_scan_limit"); +static const std::string PARAM_VAMANA_STREAMER_SATURATE_GRAPH( + "proxima.vamana.streamer.saturate_graph"); +static const std::string PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY( + "proxima.vamana.streamer.use_contiguous_memory"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_streamer.cc b/src/core/algorithm/vamana/vamana_streamer.cc new file mode 100644 index 0000000..0c08a82 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_streamer.cc @@ -0,0 +1,715 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "vamana_streamer.h" +#include +#include +#include +#include "vamana_algorithm.h" +#include "vamana_context.h" +#include "vamana_dist_calculator.h" +#include "vamana_index_provider.h" + +namespace zvec { +namespace core { + +VamanaStreamer::VamanaStreamer() = default; + +VamanaStreamer::~VamanaStreamer() { + if (state_ == STATE_INITED || state_ == STATE_OPENED) { + this->cleanup(); + } +} + +int VamanaStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { + meta_ = imeta; + meta_.set_streamer("VamanaStreamer", VamanaEntity::kRevision, params); + + params.get(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, &max_index_size_); + params.get(PARAM_VAMANA_STREAMER_MAX_DEGREE, &max_degree_); + params.get(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, &search_list_size_); + params.get(PARAM_VAMANA_STREAMER_ALPHA, &alpha_); + params.get(PARAM_VAMANA_STREAMER_MAX_OCCLUSION_SIZE, &max_occlusion_size_); + params.get(PARAM_VAMANA_STREAMER_EF, &ef_); + params.get(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + params.get(PARAM_VAMANA_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_); + params.get(PARAM_VAMANA_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_); + params.get(PARAM_VAMANA_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_); + params.get(PARAM_VAMANA_STREAMER_CHUNK_SIZE, &chunk_size_); + params.get(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, &get_vector_enabled_); + params.get(PARAM_VAMANA_STREAMER_FORCE_PADDING_RESULT_ENABLE, + &force_padding_topk_enabled_); + params.get(PARAM_VAMANA_STREAMER_USE_ID_MAP, &use_id_map_); + params.get(PARAM_VAMANA_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_); + params.get(PARAM_VAMANA_STREAMER_SATURATE_GRAPH, &saturate_graph_); + params.get(PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY, + &use_contiguous_memory_); + + size_t docs_soft_limit = 0; + params.get(PARAM_VAMANA_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit); + if (docs_soft_limit > 0 && docs_soft_limit > docs_hard_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_VAMANA_STREAMER_DOCS_HARD_LIMIT.c_str(), + PARAM_VAMANA_STREAMER_DOCS_SOFT_LIMIT.c_str()); + return IndexError_InvalidArgument; + } else if (docs_soft_limit == 0UL) { + docs_soft_limit_ = + docs_hard_limit_ * VamanaEntity::kDefaultDocsSoftLimitRatio; + } else { + docs_soft_limit_ = docs_soft_limit; + } + + // Validate parameters + if (max_degree_ == 0U) max_degree_ = VamanaEntity::kDefaultMaxDegree; + if (search_list_size_ == 0U) + search_list_size_ = VamanaEntity::kDefaultSearchListSize; + if (max_occlusion_size_ == 0U) + max_occlusion_size_ = VamanaEntity::kDefaultMaxOcclusionSize; + if (alpha_ <= 0.0f) alpha_ = VamanaEntity::kDefaultAlpha; + if (ef_ == 0U) ef_ = VamanaEntity::kDefaultEf; + if (chunk_size_ == 0UL) chunk_size_ = VamanaEntity::kDefaultChunkSize; + if (chunk_size_ > VamanaEntity::kMaxChunkSize) { + LOG_ERROR("[%s] must be < %zu", PARAM_VAMANA_STREAMER_CHUNK_SIZE.c_str(), + VamanaEntity::kMaxChunkSize); + return IndexError_InvalidArgument; + } + if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) { + LOG_ERROR("[%s] must be in range (0.0f,1.0f]", + PARAM_VAMANA_STREAMER_MAX_SCAN_RATIO.c_str()); + return IndexError_InvalidArgument; + } + if (max_scan_limit_ < min_scan_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_VAMANA_STREAMER_MAX_SCAN_LIMIT.c_str(), + PARAM_VAMANA_STREAMER_MIN_SCAN_LIMIT.c_str()); + return IndexError_InvalidArgument; + } + + LOG_DEBUG( + "Vamana init params: maxIndexSize=%zu docsHardLimit=%zu " + "docsSoftLimit=%zu maxDegree=%u searchListSize=%u alpha=%.2f " + "maxOcclusionSize=%u ef=%u maxScanRatio=%.3f minScanLimit=%zu " + "maxScanLimit=%zu bruteForceThreshold=%zu chunkSize=%zu " + "getVectorEnabled=%u forcePadding=%u", + max_index_size_, docs_hard_limit_, docs_soft_limit_, max_degree_, + search_list_size_, alpha_, max_occlusion_size_, ef_, max_scan_ratio_, + min_scan_limit_, max_scan_limit_, bruteforce_threshold_, chunk_size_, + get_vector_enabled_, force_padding_topk_enabled_); + + state_ = STATE_INITED; + return 0; +} + +int VamanaStreamer::cleanup(void) { + if (state_ == STATE_OPENED) { + this->close(); + } + + LOG_INFO("VamanaStreamer cleanup"); + + meta_.clear(); + metric_.reset(); + stats_.clear(); + if (entity_) entity_->cleanup(); + if (alg_) alg_->cleanup(); + + max_index_size_ = 0UL; + docs_hard_limit_ = VamanaEntity::kDefaultDocsHardLimit; + docs_soft_limit_ = 0UL; + max_degree_ = VamanaEntity::kDefaultMaxDegree; + search_list_size_ = VamanaEntity::kDefaultSearchListSize; + max_occlusion_size_ = VamanaEntity::kDefaultMaxOcclusionSize; + alpha_ = VamanaEntity::kDefaultAlpha; + ef_ = VamanaEntity::kDefaultEf; + bruteforce_threshold_ = VamanaEntity::kDefaultBruteForceThreshold; + max_scan_limit_ = VamanaEntity::kDefaultMaxScanLimit; + min_scan_limit_ = VamanaEntity::kDefaultMinScanLimit; + chunk_size_ = VamanaEntity::kDefaultChunkSize; + max_scan_ratio_ = VamanaEntity::kDefaultScanRatio; + state_ = STATE_INIT; + check_crc_enabled_ = false; + get_vector_enabled_ = false; + + return 0; +} + +int VamanaStreamer::setup_entity() { + entity_->set_use_key_info_map(use_id_map_); + entity_->set_vector_size(meta_.element_size()); + entity_->set_chunk_size(chunk_size_); + entity_->set_get_vector(get_vector_enabled_); + + // Set Vamana-specific parameters via public setters + entity_->set_max_degree(max_degree_); + entity_->set_search_list_size(search_list_size_); + entity_->set_max_occlusion_size(max_occlusion_size_); + entity_->set_alpha(alpha_); + entity_->set_saturate_graph(saturate_graph_); + + int ret = entity_->init(docs_hard_limit_); + if (ret != 0) { + LOG_ERROR("Vamana entity init failed: %s", IndexError::What(ret)); + } + return ret; +} + +int VamanaStreamer::open(IndexStorage::Pointer stg) { + LOG_INFO("VamanaStreamer open"); + + if (ailego_unlikely(state_ != STATE_INITED)) { + LOG_ERROR("Open storage failed, init streamer first!"); + return IndexError_NoReady; + } + + // Create entity based on storage type + switch (stg->memory_block_type()) { + case IndexStorage::MemoryBlock::MBT_BUFFERPOOL: { + entity_ = std::make_unique(stats_); + break; + } + default: { + if (use_contiguous_memory_) { + entity_ = std::make_unique(stats_); + } else { + entity_ = std::make_unique(stats_); + } + break; + } + } + + auto cleanup_on_error = [this]() { + if (entity_) { + entity_->close(); + entity_.reset(); + } + alg_.reset(); + metric_.reset(); + }; + + int ret = setup_entity(); + if (ret != 0) { + cleanup_on_error(); + return ret; + } + + ret = entity_->open(std::move(stg), max_index_size_, check_crc_enabled_); + if (ret != 0) { + cleanup_on_error(); + return ret; + } + + // Handle IndexMeta + IndexMeta index_meta; + ret = entity_->get_index_meta(&index_meta); + if (ret == IndexError_NoExist) { + ret = entity_->set_index_meta(meta_); + if (ret != 0) { + LOG_ERROR("Failed to set index meta: %s", IndexError::What(ret)); + cleanup_on_error(); + return ret; + } + } else if (ret != 0) { + LOG_ERROR("Failed to get index meta: %s", IndexError::What(ret)); + cleanup_on_error(); + return ret; + } else { + if (index_meta.dimension() != meta_.dimension() || + index_meta.element_size() != meta_.element_size() || + index_meta.metric_name() != meta_.metric_name() || + index_meta.data_type() != meta_.data_type()) { + LOG_ERROR("IndexMeta mismatch from the previous in index"); + cleanup_on_error(); + return IndexError_Mismatch; + } + auto metric_params = index_meta.metric_params(); + metric_params.merge(meta_.metric_params()); + meta_.set_metric(index_meta.metric_name(), 0, metric_params); + // Propagate reformer info from stored meta (needed for quantizers + // whose reformer params are computed during training, e.g. UniformInt8) + if (!index_meta.reformer_name().empty()) { + meta_.set_reformer(index_meta.reformer_name(), 0, + index_meta.reformer_params()); + } + } + + // Create metric + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Failed to create metric %s", meta_.metric_name().c_str()); + cleanup_on_error(); + return IndexError_NoExist; + } + ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("Failed to init metric, ret=%d", ret); + cleanup_on_error(); + return ret; + } + if (!metric_->distance() || !metric_->batch_distance()) { + LOG_ERROR("Invalid metric distance functions"); + cleanup_on_error(); + return IndexError_InvalidArgument; + } + + // Create algorithm based on entity storage mode + switch (entity_->storage_mode()) { + case VamanaStorageMode::kBufferPool: + alg_ = VamanaAlgorithmBase::UPointer( + new VamanaAlgorithm( + static_cast(*entity_))); + break; + case VamanaStorageMode::kContiguous: { + auto &contiguous_entity = + static_cast(*entity_); + int build_ret = contiguous_entity.build_contiguous_memory(); + if (build_ret != 0) { + LOG_ERROR("Failed to build contiguous memory, ret=%d", build_ret); + cleanup_on_error(); + return build_ret; + } + alg_ = VamanaAlgorithmBase::UPointer( + new VamanaAlgorithm( + contiguous_entity)); + break; + } + default: + alg_ = VamanaAlgorithmBase::UPointer( + new VamanaAlgorithm( + static_cast(*entity_))); + break; + } + + ret = alg_->init(); + if (ret != 0) { + cleanup_on_error(); + return ret; + } + + state_ = STATE_OPENED; + magic_ = IndexContext::GenerateMagic(); + + return 0; +} + +int VamanaStreamer::close(void) { + LOG_INFO("VamanaStreamer close"); + + stats_.clear(); + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_->set_index_meta(meta_); + int ret = entity_->close(); + if (ret != 0) return ret; + state_ = STATE_INITED; + return 0; +} + +int VamanaStreamer::flush(uint64_t checkpoint) { + LOG_INFO("VamanaStreamer flush checkpoint=%zu", (size_t)checkpoint); + + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_->set_index_meta(meta_); + return entity_->flush(checkpoint); +} + +int VamanaStreamer::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("VamanaStreamer dump"); + + shared_mutex_.lock(); + AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); + + meta_.set_searcher("VamanaSearcher", VamanaEntity::kRevision, + ailego::Params()); + + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + + // Calculate medoid (DiskANN standard: entry point = closest to centroid). + // At dump time, data_type and dimension are fully known from meta_. + if (entity_->doc_cnt() > 0) { + node_id_t medoid = entity_->calculate_medoid( + meta_.dimension(), static_cast(meta_.data_type())); + if (medoid != kInvalidNodeId && medoid != entity_->entry_point()) { + LOG_INFO("Updating entry point from %u to medoid %u", + entity_->entry_point(), medoid); + entity_->update_entry_point(medoid); + } + } + + return entity_->dump(dumper); +} + +IndexStreamer::Context::Pointer VamanaStreamer::create_context(void) const { + if (ailego_unlikely(state_ != STATE_OPENED)) { + LOG_ERROR("Create context failed, open storage first!"); + return Context::Pointer(); + } + + VamanaEntity::Pointer entity = entity_->clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("CreateContext clone failed"); + return Context::Pointer(); + } + + VamanaContext *ctx = + new (std::nothrow) VamanaContext(meta_.dimension(), metric_, entity); + if (ailego_unlikely(ctx == nullptr)) { + LOG_ERROR("Failed to new VamanaContext"); + return Context::Pointer(); + } + + ctx->set_ef(ef_); + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_magic(magic_); + ctx->set_force_padding_topk(force_padding_topk_enabled_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + + if (ailego_unlikely(ctx->init(VamanaContext::kStreamerContext) != 0)) { + LOG_ERROR("Init VamanaContext failed"); + delete ctx; + return Context::Pointer(); + } + + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + return Context::Pointer(ctx); +} + +IndexProvider::Pointer VamanaStreamer::create_provider(void) const { + LOG_DEBUG("VamanaStreamer create provider"); + + auto entity = entity_->clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone VamanaEntity failed"); + return nullptr; + } + return IndexProvider::Pointer( + new VamanaIndexProvider(meta_, entity, "VamanaStreamer")); +} + +int VamanaStreamer::update_context(VamanaContext *ctx) const { + const VamanaEntity::Pointer entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone search context entity"); + return IndexError_Runtime; + } + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + return ctx->update_context(VamanaContext::kStreamerContext, meta_, metric_, + entity, magic_); +} + +int VamanaStreamer::add_impl(uint64_t pkey, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + if (ailego_unlikely(entity_->doc_cnt() >= docs_soft_limit_)) { + if (entity_->doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed hard limit", entity_->doc_cnt()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %u exceed soft limit", entity_->doc_cnt()); + } + } + + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + if (metric_->support_train()) { + const std::lock_guard lk(mutex_); + ret = metric_->train(query, meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer metric train failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + } + + node_id_t id; + ret = entity_->add_vector(pkey, query, &id); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + +int VamanaStreamer::add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + if (ailego_unlikely(entity_->doc_cnt() >= docs_soft_limit_)) { + if (entity_->doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed hard limit", entity_->doc_cnt()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %u exceed soft limit", entity_->doc_cnt()); + } + } + + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + if (metric_->support_train()) { + const std::lock_guard lk(mutex_); + ret = metric_->train(query, meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer metric train failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + } + + ret = entity_->add_vector_with_id(id, query); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + +int VamanaStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_impl(query, qmeta, 1, context); +} + +int VamanaStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + + if (entity_->doc_cnt() <= ctx->get_bruteforce_threshold()) { + return search_bf_impl(query, qmeta, count, context); + } + + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + ctx->clear(); + ctx->resize_results(count); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + ret = alg_->search(ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana search failed"); + return ret; + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) return IndexError_Runtime; + return 0; +} + +void VamanaStreamer::print_debug_info() { + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) continue; + Neighbors neighbours = entity_->get_neighbors(id); + std::cout << "node: " << id << "; "; + if (neighbours.size() == 0) { + std::cout << std::endl; + continue; + } + for (uint32_t i = 0; i < neighbours.size(); ++i) { + std::cout << neighbours[i]; + if (i == neighbours.size() - 1) { + std::cout << std::endl; + } else { + std::cout << ", "; + } + } + } +} + +int VamanaStreamer::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_bf_impl(query, qmeta, 1, context); +} + +int VamanaStreamer::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + ctx->clear(); + ctx->resize_results(count); + + const auto &filter = static_cast(ctx)->filter(); + auto &topk = ctx->topk_heap(); + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + topk.clear(); + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) continue; + if (!filter.is_valid() || !filter(entity_->get_key(id))) { + dist_t dist = ctx->dist_calculator().batch_dist(id); + topk.emplace(id, dist); + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) return IndexError_Runtime; + return 0; +} + +int VamanaStreamer::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + ctx->clear(); + ctx->resize_results(count); + + auto &topk = ctx->topk_heap(); + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + topk.clear(); + for (const auto &keys : p_keys) { + for (auto key : keys) { + node_id_t id = entity_->get_id(key); + if (id == kInvalidNodeId) continue; + dist_t dist = ctx->dist_calculator().batch_dist(id); + topk.emplace(id, dist); + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) return IndexError_Runtime; + return 0; +} + +INDEX_FACTORY_REGISTER_STREAMER(VamanaStreamer); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_streamer.h b/src/core/algorithm/vamana/vamana_streamer.h new file mode 100644 index 0000000..a1217c9 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_streamer.h @@ -0,0 +1,183 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include "vamana_algorithm.h" +#include "vamana_streamer_entity.h" + +namespace zvec { +namespace core { + +class VamanaStreamer : public IndexStreamer { + public: + using ContextPointer = IndexStreamer::Context::Pointer; + + VamanaStreamer(void); + ~VamanaStreamer(void) override; + + VamanaStreamer(const VamanaStreamer &) = delete; + VamanaStreamer &operator=(const VamanaStreamer &) = delete; + + protected: + int init(const IndexMeta &imeta, const ailego::Params ¶ms) override; + + int cleanup(void) override; + + Context::Pointer create_context(void) const override; + + IndexProvider::Pointer create_provider(void) const override; + + int add_impl(uint64_t pkey, const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + int add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector_by_key(key, block); + } + + const void *get_vector_by_id(uint32_t id) const override { + return entity_->get_vector(id); + } + + int get_vector_by_id(const uint32_t id, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector(id, block); + } + + int open(IndexStorage::Pointer stg) override; + + int close(void) override; + + int flush(uint64_t checkpoint) override; + + int dump(const IndexDumper::Pointer &dumper) override; + + const Stats &stats(void) const override { + return stats_; + } + + const IndexMeta &meta(void) const override { + return meta_; + } + + void print_debug_info() override; + + private: + inline int check_params(const void *query, + const IndexQueryMeta &qmeta) const { + if (ailego_unlikely(!query)) { + LOG_ERROR("null query"); + return IndexError_InvalidArgument; + } + if (ailego_unlikely(qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size())) { + LOG_ERROR("Unsupported query meta"); + return IndexError_Mismatch; + } + return 0; + } + + int setup_entity(); + int update_context(VamanaContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 }; + + class Stats : public IndexStreamer::Stats { + public: + void clear(void) { + set_revision_id(0u); + set_loaded_count(0u); + set_added_count(0u); + set_discarded_count(0u); + set_index_size(0u); + set_dumped_size(0u); + set_check_point(0u); + set_create_time(0u); + set_update_time(0u); + clear_attributes(); + } + }; + + std::unique_ptr entity_; + VamanaAlgorithmBase::UPointer alg_; + IndexMeta meta_{}; + IndexMetric::Pointer metric_{}; + + Stats stats_{}; + std::mutex mutex_{}; + + size_t max_index_size_{0UL}; + size_t chunk_size_{VamanaEntity::kDefaultChunkSize}; + size_t docs_hard_limit_{VamanaEntity::kDefaultDocsHardLimit}; + size_t docs_soft_limit_{0UL}; + uint32_t max_degree_{VamanaEntity::kDefaultMaxDegree}; + uint32_t search_list_size_{VamanaEntity::kDefaultSearchListSize}; + uint32_t max_occlusion_size_{VamanaEntity::kDefaultMaxOcclusionSize}; + float alpha_{VamanaEntity::kDefaultAlpha}; + uint32_t ef_{VamanaEntity::kDefaultEf}; + size_t bruteforce_threshold_{VamanaEntity::kDefaultBruteForceThreshold}; + size_t max_scan_limit_{VamanaEntity::kDefaultMaxScanLimit}; + size_t min_scan_limit_{VamanaEntity::kDefaultMinScanLimit}; + float max_scan_ratio_{VamanaEntity::kDefaultScanRatio}; + + uint32_t magic_{0U}; + State state_{STATE_INIT}; + bool check_crc_enabled_{false}; + bool get_vector_enabled_{false}; + bool force_padding_topk_enabled_{false}; + bool use_id_map_{true}; + bool saturate_graph_{VamanaEntity::kDefaultSaturateGraph}; + bool use_contiguous_memory_{false}; + + ailego::SharedMutex shared_mutex_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_streamer_entity.cc b/src/core/algorithm/vamana/vamana_streamer_entity.cc new file mode 100644 index 0000000..4324ea2 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_streamer_entity.cc @@ -0,0 +1,926 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "vamana_streamer_entity.h" +#include +#include +#include + +namespace zvec { +namespace core { + +VamanaStreamerEntity::VamanaStreamerEntity(IndexStreamer::Stats &stats) + : stats_(stats) { + keys_map_lock_ = std::make_shared(); + keys_map_ = std::make_shared>(); + keys_map_->set_empty_key(kInvalidKey); + broker_ = std::make_shared(stats); +} + +VamanaStreamerEntity::~VamanaStreamerEntity() {} + +int VamanaStreamerEntity::cleanup() { + node_chunks_.clear(); + if (keys_map_) { + keys_map_->clear(); + } + header_.clear(); + return 0; +} + +int VamanaStreamerEntity::init(size_t /*max_doc_cnt*/) { + // node_size = vector_size + key_size + neighbors_size + set_node_size(vector_size() + sizeof(key_t) + neighbors_size()); + neighbor_size_ = neighbors_size(); + return 0; +} + +key_t VamanaStreamerEntity::get_key(node_id_t id) const { + if (!use_key_info_map_) return id; + auto loc = get_key_chunk_loc(id); + if (ailego_unlikely(loc.first >= node_chunks_.size())) return kInvalidKey; + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(ptr); +} + +const void *VamanaStreamerEntity::get_vector(node_id_t id) const { + auto loc = get_vector_chunk_loc(id); + if (ailego_unlikely(loc.first >= node_chunks_.size())) return nullptr; + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, vector_size()); + if (ailego_unlikely(ret != vector_size())) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return nullptr; + } + return ptr; +} + +int VamanaStreamerEntity::get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const { + auto loc = get_vector_chunk_loc(id); + if (ailego_unlikely(loc.first >= node_chunks_.size())) + return IndexError_NoExist; + size_t ret = node_chunks_[loc.first]->read(loc.second, block, vector_size()); + if (ailego_unlikely(ret != vector_size())) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return IndexError_ReadData; + } + return 0; +} + +int VamanaStreamerEntity::get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const { + for (uint32_t i = 0; i < count; ++i) { + vecs[i] = get_vector(ids[i]); + if (ailego_unlikely(vecs[i] == nullptr)) { + return IndexError_NoExist; + } + } + return 0; +} + +int VamanaStreamerEntity::get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (uint32_t i = 0; i < count; ++i) { + int ret = get_vector(ids[i], vec_blocks[i]); + if (ailego_unlikely(ret != 0)) return ret; + } + return 0; +} + +const Neighbors VamanaStreamerEntity::get_neighbors(node_id_t id) const { + auto loc = get_neighbor_chunk_loc(id); + IndexStorage::MemoryBlock mem_block; + size_t ret = loc.first->read(loc.second, mem_block, neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return Neighbors(); + } + return Neighbors(mem_block); +} + +int VamanaStreamerEntity::add_vector(key_t key, const void *vec, + node_id_t *id) { + Chunk::Pointer node_chunk; + size_t chunk_offset = static_cast(-1); + + std::lock_guard lock(mutex_); + + node_id_t local_id = static_cast(doc_cnt()); + uint32_t chunk_index = node_chunks_.size() - 1U; + if (chunk_index == static_cast(-1) || + (node_chunks_[chunk_index]->data_size() >= + node_cnt_per_chunk_ * node_size())) { + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + chunk_index++; + if (auto dret = ensure_dist_chunk_for(chunk_index); dret != 0) { + return dret; + } + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NODE, chunk_index, + chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + chunk_offset = 0UL; + { + std::lock_guard chunks_lock(node_chunks_mutex_); + node_chunks_.emplace_back(node_chunk); + } + } else { + node_chunk = node_chunks_[chunk_index]; + chunk_offset = node_chunk->data_size(); + } + + // Write vector + size_t size = node_chunk->write(chunk_offset, vec, vector_size()); + if (ailego_unlikely(size != vector_size())) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + // Write key + size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + // Neighbors are initialized to zero by default (chunk is zero-filled) + + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + + if (use_key_info_map_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = local_id; + keys_map_lock_->unlock(); + } + + *mutable_doc_cnt() += 1; + broker_->mark_dirty(); + *id = local_id; + + return 0; +} + +int VamanaStreamerEntity::add_vector_with_id(node_id_t id, const void *vec) { + Chunk::Pointer node_chunk; + size_t chunk_offset = static_cast(-1); + key_t key = id; + + std::lock_guard lock(mutex_); + + auto func_get_node_chunk_and_offset = [&](node_id_t node_id) -> int { + uint32_t chunk_idx = node_id >> node_index_mask_bits_; + ailego_assert_with(chunk_idx <= node_chunks_.size(), "invalid chunk idx"); + if (chunk_idx == node_chunks_.size()) { + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + if (auto dret = ensure_dist_chunk_for(chunk_idx); dret != 0) { + return dret; + } + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, + chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + { + std::lock_guard chunks_lock(node_chunks_mutex_); + node_chunks_.emplace_back(node_chunk); + } + } + node_chunk = node_chunks_[chunk_idx]; + chunk_offset = (node_id & node_index_mask_) * node_size(); + return 0; + }; + + // Fill gaps with invalid keys + for (size_t start_id = doc_cnt(); start_id < id; ++start_id) { + if (auto ret = func_get_node_chunk_and_offset(start_id); ret != 0) { + return ret; + } + size_t size = node_chunk->write(chunk_offset + vector_size(), &kInvalidKey, + sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (auto ret = func_get_node_chunk_and_offset(id); ret != 0) { + return ret; + } + + // Write vector + size_t size = node_chunk->write(chunk_offset, vec, vector_size()); + if (ailego_unlikely(size != vector_size())) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + // Write key + size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + + if (*mutable_doc_cnt() <= id) { + *mutable_doc_cnt() = id + 1; + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (use_key_info_map_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = id; + keys_map_lock_->unlock(); + } + + broker_->mark_dirty(); + return 0; +} + +int VamanaStreamerEntity::update_neighbors( + node_id_t id, const std::vector> &neighbors) { + auto loc = get_neighbor_chunk_loc(id); + uint32_t count = std::min(static_cast(neighbors.size()), + static_cast(max_degree())); + + // Build neighbor data in a local buffer + size_t nbr_size = neighbors_size(); + std::vector buffer(nbr_size, 0); + auto *hd = reinterpret_cast(buffer.data()); + hd->neighbor_cnt = count; + for (uint32_t i = 0; i < count; ++i) { + hd->neighbors[i] = neighbors[i].first; + } + + size_t ret = loc.first->write(loc.second, buffer.data(), nbr_size); + if (ailego_unlikely(ret != nbr_size)) { + LOG_ERROR("Write neighbors failed, ret=%zu", ret); + return IndexError_WriteData; + } + return 0; +} + +void VamanaStreamerEntity::add_neighbor(node_id_t id, uint32_t size, + node_id_t neighbor_id) { + auto loc = get_neighbor_chunk_loc(id); + if (size >= max_degree()) return; + + // Read current neighbors + IndexStorage::MemoryBlock mem_block; + size_t ret = loc.first->read(loc.second, mem_block, neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return; + } + + // Copy to mutable buffer, update, and write back + std::vector buffer(neighbor_size_); + memcpy(buffer.data(), mem_block.data(), neighbor_size_); + auto *hd = reinterpret_cast(buffer.data()); + hd->neighbors[size] = neighbor_id; + hd->neighbor_cnt = size + 1; + + ret = loc.first->write(loc.second, buffer.data(), neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Write neighbor failed, ret=%zu", ret); + } +} + +void VamanaStreamerEntity::update_entry_point(node_id_t ep) { + VamanaEntity::update_entry_point(ep); + flush_header(); +} + +int VamanaStreamerEntity::open(IndexStorage::Pointer stg, + uint64_t max_index_size, bool check_crc) { + std::lock_guard lock(mutex_); + bool huge_page = stg->isHugePage(); + + int ret = broker_->open(std::move(stg), chunk_size_, check_crc); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Open index failed: %s", IndexError::What(ret)); + return ret; + } + + ret = init_chunk_params(max_index_size, huge_page); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("init_chunk_params failed: %s", IndexError::What(ret)); + return ret; + } + broker_->set_max_chunks_size(max_index_size_); + + // Init header + auto header_chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId); + if (!header_chunk) { + // Open empty index, create header + auto p = + broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId, header_size()); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc header chunk failed"); + return p.first; + } + size_t size = p.second->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + return 0; + } + + // Open existing index + ret = init_chunks(header_chunk); + if (ailego_unlikely(ret != 0)) return ret; + + // Verify total docs + node_id_t total_vecs = 0; + if (!node_chunks_.empty()) { + size_t last_idx = node_chunks_.size() - 1; + if (node_chunks_[last_idx]->data_size() % node_size()) { + LOG_WARN("The index may be broken"); + return IndexError_InvalidFormat; + } + total_vecs = last_idx * node_cnt_per_chunk_ + + node_chunks_[last_idx]->data_size() / node_size(); + } + + LOG_INFO("Open Vamana index, maxDegree=%zu docCnt=%u totalVecs=%u", + max_degree(), doc_cnt(), total_vecs); + + if (doc_cnt() != total_vecs) { + LOG_WARN("Index closed abnormally, using totalVecs as curDocCnt"); + *mutable_doc_cnt() = total_vecs; + } + + // Rebuild key map + if (use_key_info_map_) { + for (node_id_t i = 0; i < doc_cnt(); ++i) { + key_t k = get_key(i); + if (k != kInvalidKey) { + (*keys_map_)[k] = i; + } + } + } + + stats_.set_loaded_count(doc_cnt()); + return 0; +} + +int VamanaStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) { + // Read header from chunk + const void *hd_ptr = nullptr; + size_t ret = header_chunk->read(0UL, &hd_ptr, header_size()); + if (ailego_unlikely(ret != header_size())) { + LOG_ERROR("Read header chunk failed"); + return IndexError_ReadData; + } + auto *hd = reinterpret_cast(hd_ptr); + + // Validate + if (vector_size() != hd->vector_size()) { + LOG_ERROR("vector size %zu mismatch index previous %zu", vector_size(), + hd->vector_size()); + return IndexError_Mismatch; + } + if (max_degree() != hd->max_degree()) { + LOG_ERROR("max_degree %zu mismatch index previous %zu", max_degree(), + hd->max_degree()); + return IndexError_Mismatch; + } + + *mutable_header() = *hd; + + // Load node chunks + size_t chunk_cnt = broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NODE); + for (size_t i = 0; i < chunk_cnt; ++i) { + auto chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_NODE, i); + if (ailego_unlikely(!chunk)) { + LOG_ERROR("Get node chunk %zu failed", i); + return IndexError_ReadData; + } + node_chunks_.emplace_back(std::move(chunk)); + } + return 0; +} + +int VamanaStreamerEntity::close() { + LOG_DEBUG("close Vamana index"); + std::lock_guard lock(mutex_); + flush_header(); + mutable_header()->reset(); + keys_map_->clear(); + header_.clear(); + node_chunks_.clear(); + dist_chunks_.clear(); + dist_loaded_ = false; + return broker_->close(); +} + +int VamanaStreamerEntity::flush(uint64_t checkpoint) { + LOG_INFO("Flush Vamana index, curDocs=%u", doc_cnt()); + std::lock_guard lock(mutex_); + flush_header(); + return broker_->flush(checkpoint); +} + +int VamanaStreamerEntity::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("Dump Vamana index, curDocs=%u", doc_cnt()); + + std::vector keys(doc_cnt()); + auto ret = dump_segments(dumper, keys.data()); + if (ailego_unlikely(ret < 0)) { + return static_cast(ret); + } + *stats_.mutable_dumped_size() += ret; + return 0; +} + +const VamanaEntity::Pointer VamanaStreamerEntity::clone() const { + std::vector cloned_chunks; + cloned_chunks.reserve(node_chunks_.size()); + for (size_t i = 0; i < node_chunks_.size(); ++i) { + cloned_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!cloned_chunks[i])) { + LOG_ERROR("VamanaStreamerEntity get chunk failed in clone"); + return VamanaEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) VamanaStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, get_vector_enabled_, + use_key_info_map_, keys_map_lock_, keys_map_, std::move(cloned_chunks), + broker_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("VamanaStreamerEntity new failed"); + } + return VamanaEntity::Pointer(entity); +} + +const VamanaEntity::Pointer VamanaMmapStreamerEntity::clone() const { + std::vector cloned_chunks; + cloned_chunks.reserve(node_chunks_.size()); + for (size_t i = 0; i < node_chunks_.size(); ++i) { + cloned_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!cloned_chunks[i])) { + LOG_ERROR("VamanaMmapStreamerEntity get chunk failed in clone"); + return VamanaEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) VamanaMmapStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, get_vector_enabled_, + use_key_info_map_, keys_map_lock_, keys_map_, std::move(cloned_chunks), + broker_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("VamanaMmapStreamerEntity new failed"); + } + return VamanaEntity::Pointer(entity); +} + +const VamanaEntity::Pointer VamanaContiguousStreamerEntity::clone() const { + std::vector cloned_chunks; + cloned_chunks.reserve(node_chunks_.size()); + for (size_t i = 0; i < node_chunks_.size(); ++i) { + cloned_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!cloned_chunks[i])) { + LOG_ERROR("VamanaContiguousStreamerEntity get chunk failed in clone"); + return VamanaEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) VamanaContiguousStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, get_vector_enabled_, + use_key_info_map_, keys_map_lock_, keys_map_, std::move(cloned_chunks), + broker_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("VamanaContiguousStreamerEntity new failed"); + return VamanaEntity::Pointer(); + } + + // Share contiguous memory with the clone (zero-copy) + entity->vector_memory_ = vector_memory_; + entity->vector_base_ = vector_base_; + entity->vector_stride_ = vector_stride_; + entity->graph_memory_ = graph_memory_; + entity->graph_base_ = graph_base_; + entity->graph_stride_ = graph_stride_; + + return VamanaEntity::Pointer(entity); +} + +// ============================================================================ +// VamanaContiguousStreamerEntity implementation +// ============================================================================ + +char *VamanaContiguousStreamerEntity::allocate_contiguous(size_t size) { + if (size == 0) return nullptr; + void *ptr = ailego::MemoryHelper::AllocateHugePage(size); + if (!ptr) { + LOG_ERROR("AllocateHugePage failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +} + +int VamanaContiguousStreamerEntity::build_contiguous_memory() { + vector_memory_.reset(); + vector_base_ = nullptr; + vector_stride_ = 0; + graph_memory_.reset(); + graph_base_ = nullptr; + + const uint32_t total_docs = doc_cnt(); + if (total_docs == 0) return 0; + + const size_t per_node = node_size(); + const size_t vec_size = vector_size(); + + // Pad per-vector stride up to kVectorAlignment (64B) so every vector + // starts on a cache-line boundary. + vector_stride_ = + (vec_size + (kVectorAlignment - 1)) & ~(kVectorAlignment - 1); + // graph_stride = key + neighbors (everything except vector) + graph_stride_ = sizeof(key_t) + neighbors_size(); + + // Allocate flat vector array (stride = vector_stride_, padded for 64B) + const size_t total_vec_data = + static_cast(total_docs) * vector_stride_; + size_t vector_memory_size = AlignHugePageSize(total_vec_data); + char *raw_vec = allocate_contiguous(vector_memory_size); + if (!raw_vec) return IndexError_Runtime; + vector_memory_.reset(raw_vec, ContiguousDeleter{vector_memory_size}); + vector_base_ = raw_vec; + + // Allocate graph array (stride = sizeof(key_t) + neighbors_size) + const size_t total_graph_data = + static_cast(total_docs) * graph_stride_; + size_t graph_memory_size = AlignHugePageSize(total_graph_data); + char *raw_graph = allocate_contiguous(graph_memory_size); + if (!raw_graph) { + vector_memory_.reset(); + vector_base_ = nullptr; + vector_stride_ = 0; + return IndexError_Runtime; + } + graph_memory_.reset(raw_graph, ContiguousDeleter{graph_memory_size}); + graph_base_ = raw_graph; + + // Split node data from chunks into vector / graph arrays. + // Original node layout: [vector (vec_size) | key (8B) | neighbors] + // Padding bytes in vector_base_ are left zero (anon mmap is zero-filled). + const auto &chunks = node_chunks_; + const uint32_t nodes_per_chunk = 1U << node_index_mask_bits_; + for (size_t chunk_idx = 0; chunk_idx < chunks.size(); ++chunk_idx) { + const void *chunk_data = nullptr; + size_t data_size = chunks[chunk_idx]->data_size(); + chunks[chunk_idx]->read(0, &chunk_data, data_size); + + uint32_t base_id = chunk_idx * nodes_per_chunk; + uint32_t count_in_chunk = std::min(nodes_per_chunk, total_docs - base_id); + + const char *src = static_cast(chunk_data); + for (uint32_t i = 0; i < count_in_chunk; ++i) { + const char *node_src = src + static_cast(i) * per_node; + size_t global_id = static_cast(base_id + i); + + // Copy vector to flat vector array at padded stride + std::memcpy(vector_base_ + global_id * vector_stride_, node_src, + vec_size); + + // Copy key + neighbors to graph array + std::memcpy(graph_base_ + global_id * graph_stride_, node_src + vec_size, + graph_stride_); + } + } + + LOG_INFO( + "Built Vamana contiguous memory: " + "vector_mem=%zu graph_mem=%zu total_docs=%u " + "node_chunks=%zu vector_size=%zu vector_stride=%zu " + "(cache-line aligned to %zuB)", + vector_memory_size, graph_memory_size, total_docs, chunks.size(), + vec_size, vector_stride_, kVectorAlignment); + + return 0; +} + +// ============================================================================ +// Neighbor distance storage implementation (CSR-like, lazy-loaded) +// ============================================================================ + +int VamanaStreamerEntity::ensure_dist_storage() { + if (dist_loaded_) return 0; + + std::lock_guard lock(mutex_); + if (dist_loaded_) return 0; // double-check after lock + + dist_entry_size_ = static_cast(max_degree() * sizeof(dist_t)); + + // Pre-reserve dist_chunks_ to match node_chunks_ capacity so that + // subsequent emplace_back in ensure_dist_chunk_for never triggers + // reallocation while concurrent add_node threads read dist_chunks_. + dist_chunks_.reserve(node_chunks_.capacity()); + + // Calculate how many dist chunks we need for existing nodes + uint32_t total_docs = doc_cnt(); + if (total_docs == 0) { + dist_loaded_ = true; + return 0; + } + + // Check if dist chunks already exist in storage (reopened index) + size_t existing_dist_chunks = + broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST); + if (existing_dist_chunks > 0) { + // Load existing dist chunks + for (size_t i = 0; i < existing_dist_chunks; ++i) { + auto chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST, i); + if (ailego_unlikely(!chunk)) { + LOG_ERROR("Failed to load dist chunk %zu", i); + return IndexError_ReadData; + } + dist_chunks_.emplace_back(std::move(chunk)); + } + LOG_INFO("Loaded %zu existing dist chunks", existing_dist_chunks); + } else { + // Allocate new dist chunks for all existing nodes + int ret = alloc_dist_chunks_for_existing_nodes(); + if (ret != 0) return ret; + } + + dist_loaded_ = true; + return 0; +} + +int VamanaStreamerEntity::ensure_dist_chunk_for(uint32_t chunk_index) { + // No-op when dist storage is not active. + if (!dist_loaded_ || dist_entry_size_ == 0) return 0; + + // Idempotent: nothing to do if this dist chunk slot already exists and is + // populated. (Slots created by the placeholder loop below will hold + // nullptr and must still be (re-)allocated.) + if (chunk_index < dist_chunks_.size() && dist_chunks_[chunk_index]) { + return 0; + } + + uint32_t dist_chunk_data_size = node_cnt_per_chunk_ * dist_entry_size_; + uint32_t dist_chunk_size = AlignPageSize(dist_chunk_data_size); + auto dp = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST, + chunk_index, dist_chunk_size); + if (ailego_unlikely(dp.first != 0)) { + LOG_ERROR("Alloc dist chunk %u failed", chunk_index); + return dp.first; + } + dp.second->resize(dist_chunk_data_size); + { + // Protect dist_chunks_ modification against concurrent readers in + // get_neighbor_dists/update_neighbor_dists (called from add_node without + // mutex_). Uses node_chunks_mutex_ which is the same lock used by + // sync_chunks for CHUNK_TYPE_NEIGHBOR_DIST. + std::lock_guard chunks_lock(node_chunks_mutex_); + while (dist_chunks_.size() <= chunk_index) { + dist_chunks_.emplace_back(nullptr); + } + dist_chunks_[chunk_index] = std::move(dp.second); + } + return 0; +} + +int VamanaStreamerEntity::alloc_dist_chunks_for_existing_nodes() { + uint32_t total_docs = doc_cnt(); + if (total_docs == 0) return 0; + + // Calculate dist chunk size: same number of nodes per chunk as node chunks + uint32_t dist_chunk_data_size = node_cnt_per_chunk_ * dist_entry_size_; + uint32_t dist_chunk_size = AlignPageSize(dist_chunk_data_size); + + uint32_t num_chunks_needed = + (total_docs + node_cnt_per_chunk_ - 1) >> node_index_mask_bits_; + + for (uint32_t i = 0; i < num_chunks_needed; ++i) { + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST, i, + dist_chunk_size); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc dist chunk %u failed", i); + return p.first; + } + // Resize to cover all nodes in this chunk + uint32_t nodes_in_chunk = + std::min(node_cnt_per_chunk_, total_docs - i * node_cnt_per_chunk_); + size_t data_size = static_cast(nodes_in_chunk) * dist_entry_size_; + p.second->resize(data_size); + dist_chunks_.emplace_back(std::move(p.second)); + } + broker_->mark_dirty(); + LOG_INFO("Allocated %u dist chunks for %u existing nodes", num_chunks_needed, + total_docs); + return 0; +} + +const dist_t *VamanaStreamerEntity::get_neighbor_dists(node_id_t id) const { + if (!dist_loaded_) return nullptr; + + auto loc = get_dist_chunk_loc(id); + if (ailego_unlikely(loc.first >= dist_chunks_.size())) { + sync_dist_chunks(loc.first); + } + if (ailego_unlikely(loc.first >= dist_chunks_.size())) return nullptr; + + const void *ptr = nullptr; + size_t ret = + dist_chunks_[loc.first]->read(loc.second, &ptr, dist_entry_size_); + if (ailego_unlikely(ret != dist_entry_size_)) return nullptr; + return static_cast(ptr); +} + +void VamanaStreamerEntity::update_neighbor_dists( + node_id_t id, const std::vector> &neighbors) { + if (!dist_loaded_) return; + + auto loc = get_dist_chunk_loc(id); + // Dist chunk must have been pre-allocated by add_vector or + // ensure_dist_storage + if (ailego_unlikely(loc.first >= dist_chunks_.size() || + dist_chunks_[loc.first] == nullptr)) { + LOG_ERROR("Dist chunk %u not allocated for node %u", loc.first, id); + return; + } + + // Write distances: fill max_degree slots, zero-pad unused slots + uint32_t max_deg = static_cast(max_degree()); + std::vector dists(max_deg, 0.0f); + for (size_t i = 0; i < neighbors.size() && i < max_deg; ++i) { + dists[i] = neighbors[i].second; + } + dist_chunks_[loc.first]->write(loc.second, dists.data(), dist_entry_size_); +} + +void VamanaStreamerEntity::set_neighbor_dist(node_id_t id, uint32_t idx, + dist_t dist) { + if (!dist_loaded_) return; + + auto loc = get_dist_chunk_loc(id); + if (ailego_unlikely(loc.first >= dist_chunks_.size())) return; + + uint32_t offset = loc.second + idx * sizeof(dist_t); + dist_chunks_[loc.first]->write(offset, &dist, sizeof(dist_t)); +} + +// calculate_medoid: Compute the medoid (entry point) following DiskANN's +// standard approach: +// 1. Compute the centroid (component-wise mean) of all vectors in float +// space. +// 2. Find the data point closest to this centroid using squared L2 distance +// in float space. +// 3. Return that point's node ID as the medoid. +// +// Called at dump time to set the optimal entry point for the persisted index. +// data_type uses IndexMeta::DataType values: DT_FP16=1, DT_FP32=2, DT_INT8=4. +// ============================================================================ +node_id_t VamanaStreamerEntity::calculate_medoid(uint32_t dimension, + uint32_t data_type) { + uint32_t n = doc_cnt(); + if (n == 0) return kInvalidNodeId; + if (dimension == 0) return kInvalidNodeId; + + // data_type constants matching IndexMeta::DataType + constexpr uint32_t DT_FP16 = 1; + constexpr uint32_t DT_FP32 = 2; + constexpr uint32_t DT_INT8 = 4; + + if (data_type != DT_FP32 && data_type != DT_INT8 && data_type != DT_FP16) { + LOG_WARN("calculate_medoid: unsupported data_type=%u, skip", data_type); + return entry_point(); + } + + // Step 1: Compute centroid (mean) of all vectors in float space. + std::vector centroid(dimension, 0.0f); + uint32_t valid_count = 0; + + for (node_id_t i = 0; i < n; ++i) { + if (get_key(i) == kInvalidKey) continue; + const void *vec = get_vector(i); + if (vec == nullptr) continue; + + switch (data_type) { + case DT_FP32: { + const float *fv = static_cast(vec); + for (uint32_t d = 0; d < dimension; ++d) centroid[d] += fv[d]; + break; + } + case DT_INT8: { + const int8_t *iv = static_cast(vec); + for (uint32_t d = 0; d < dimension; ++d) + centroid[d] += static_cast(iv[d]); + break; + } + case DT_FP16: { + const uint16_t *hv = static_cast(vec); + for (uint32_t d = 0; d < dimension; ++d) + centroid[d] += ailego::FloatHelper::ToFP32(hv[d]); + break; + } + } + valid_count++; + } + + if (valid_count == 0) return kInvalidNodeId; + if (valid_count == 1) { + for (node_id_t i = 0; i < n; ++i) { + if (get_key(i) != kInvalidKey && get_vector(i) != nullptr) return i; + } + return kInvalidNodeId; + } + + float inv = 1.0f / static_cast(valid_count); + for (uint32_t d = 0; d < dimension; ++d) centroid[d] *= inv; + + // Step 2: Find the data point closest to the centroid (squared L2). + node_id_t medoid = kInvalidNodeId; + float min_dist = std::numeric_limits::max(); + + for (node_id_t i = 0; i < n; ++i) { + if (get_key(i) == kInvalidKey) continue; + const void *vec = get_vector(i); + if (vec == nullptr) continue; + + float dist = 0.0f; + switch (data_type) { + case DT_FP32: { + const float *fv = static_cast(vec); + for (uint32_t d = 0; d < dimension; ++d) { + float diff = fv[d] - centroid[d]; + dist += diff * diff; + } + break; + } + case DT_INT8: { + const int8_t *iv = static_cast(vec); + for (uint32_t d = 0; d < dimension; ++d) { + float diff = static_cast(iv[d]) - centroid[d]; + dist += diff * diff; + } + break; + } + case DT_FP16: { + const uint16_t *hv = static_cast(vec); + for (uint32_t d = 0; d < dimension; ++d) { + float diff = ailego::FloatHelper::ToFP32(hv[d]) - centroid[d]; + dist += diff * diff; + } + break; + } + } + + if (dist < min_dist) { + min_dist = dist; + medoid = i; + } + } + + LOG_INFO("Calculated medoid: node_id=%u, min_sq_dist=%.4f, valid_docs=%u", + medoid, min_dist, valid_count); + return medoid; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_streamer_entity.h b/src/core/algorithm/vamana/vamana_streamer_entity.h new file mode 100644 index 0000000..ecc3f70 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_streamer_entity.h @@ -0,0 +1,713 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include "algorithm/hnsw/hnsw_chunk.h" +#include "algorithm/hnsw/hnsw_entity.h" // MmapMemoryBlock, BufferPoolMemoryBlock, NeighborsT +#include "vamana_entity.h" +#include "vamana_params.h" + +namespace zvec { +namespace core { + + +// Storage mode for VamanaStreamerEntity +enum class VamanaStorageMode { kMmap = 0, kBufferPool = 1, kContiguous = 2 }; + +// VamanaStreamerEntity manages vector data, primary keys, and neighbors +// for a single-layer Vamana graph in streaming (incremental) mode. +// Unlike HNSW, Vamana has no upper-level neighbors — only a single +// neighbor list per node. Node layout in chunk: +// [vector_data (vector_size) | key (sizeof(key_t)) | NeighborsHeader + +// neighbors (neighbors_size)] +class VamanaStreamerEntity : public VamanaEntity { + public: + // Virtual interface implementation + int cleanup() override; + const VamanaEntity::Pointer clone() const override; + key_t get_key(node_id_t id) const override; + const void *get_vector(node_id_t id) const override; + int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const override; + const Neighbors get_neighbors(node_id_t id) const override; + + int add_vector(key_t key, const void *vec, node_id_t *id) override; + int add_vector_with_id(node_id_t id, const void *vec) override; + int update_neighbors( + node_id_t id, + const std::vector> &neighbors) override; + void add_neighbor(node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + int dump(const IndexDumper::Pointer &dumper) override; + void update_entry_point(node_id_t ep) override; + + // Calculate medoid: find the data point closest to the centroid + // of all vectors (DiskANN standard entry point selection). + node_id_t calculate_medoid(uint32_t dimension, uint32_t data_type) override; + + // --- Neighbor distance storage --- + int ensure_dist_storage() override; + bool dist_storage_loaded() const override { + return dist_loaded_; + } + const dist_t *get_neighbor_dists(node_id_t id) const override; + void update_neighbor_dists( + node_id_t id, + const std::vector> &neighbors) override; + void set_neighbor_dist(node_id_t id, uint32_t idx, dist_t dist) override; + + virtual VamanaStorageMode storage_mode() const { + return VamanaStorageMode::kMmap; + } + + void set_use_key_info_map(bool use_id_map) { + use_key_info_map_ = use_id_map; + } + + public: + VamanaStreamerEntity(IndexStreamer::Stats &stats); + ~VamanaStreamerEntity(); + + const void *get_vector_by_key(key_t key) const override { + auto id = get_id(key); + return id == kInvalidNodeId ? nullptr : get_vector(id); + } + + int get_vector_by_key(const key_t key, + IndexStorage::MemoryBlock &block) const override { + auto id = get_id(key); + if (id != kInvalidNodeId) { + return get_vector(id, block); + } + return IndexError_InvalidArgument; + } + + int init(size_t max_doc_cnt); + int flush(uint64_t checkpoint); + int open(IndexStorage::Pointer stg, uint64_t max_index_size, bool check_crc); + int close(); + + int set_index_meta(const IndexMeta &meta) const { + return IndexHelper::SerializeToStorage(meta, broker_->storage().get()); + } + + int get_index_meta(IndexMeta *meta) const { + return IndexHelper::DeserializeFromStorage(broker_->storage().get(), meta); + } + + inline void set_chunk_size(size_t val) { + chunk_size_ = val; + } + inline void set_get_vector(bool val) { + get_vector_enabled_ = val; + } + + inline node_id_t get_id(key_t key) const { + if (use_key_info_map_) { + keys_map_lock_->lock_shared(); + auto it = keys_map_->find(key); + keys_map_lock_->unlock_shared(); + return it == keys_map_->end() ? kInvalidNodeId : it->second; + } + return key; + } + + // --- Typed access methods for hot-path optimization --- + // These are templated on MemBlock type to avoid runtime branching. + + template + inline NeighborsT get_neighbors_typed(node_id_t id) const; + + template + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const; + + template + inline key_t get_key_typed(node_id_t id) const; + + protected: + inline void sync_node_chunks(size_t idx) const { + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, idx, &node_chunks_); + } + + protected: + template + using HashMap = google::dense_hash_map>; + template + using HashMapPointer = std::shared_ptr>; + + //! Clone constructor, used by clone method in subclasses + VamanaStreamerEntity(IndexStreamer::Stats &stats, const VamanaHeader &hd, + size_t chunk_size, uint32_t node_index_mask_bits, + bool get_vector_enabled, bool use_key_info_map, + std::shared_ptr &keys_map_lock, + const HashMapPointer &keys_map, + std::vector &&node_chunks, + const ChunkBroker::Pointer &broker) + : stats_(stats), + chunk_size_(chunk_size), + node_index_mask_bits_(node_index_mask_bits), + node_cnt_per_chunk_(1UL << node_index_mask_bits_), + node_index_mask_(node_cnt_per_chunk_ - 1), + get_vector_enabled_(get_vector_enabled), + use_key_info_map_(use_key_info_map), + keys_map_lock_(keys_map_lock), + keys_map_(keys_map), + broker_(broker), + node_chunks_(std::move(node_chunks)) { + *mutable_header() = hd; + neighbor_size_ = neighbors_size(); + } + + //! Lazy chunk synchronization: fetches chunks from broker when needed. + //! Protected by node_chunks_mutex_ to synchronize with add_vector's + //! emplace_back during concurrent build. + void sync_chunks(ChunkBroker::CHUNK_TYPE type, size_t idx, + std::vector *chunks) const { + if (ailego_likely(idx < chunks->size())) { + return; + } + std::lock_guard lock(node_chunks_mutex_); + // Double-check after acquiring lock + if (idx < chunks->size()) { + return; + } + for (size_t i = chunks->size(); i <= idx; ++i) { + auto chunk = broker_->get_chunk(type, i); + ailego_assert_with(!!chunk, "get chunk failed"); + chunks->emplace_back(std::move(chunk)); + } + } + + inline std::pair get_vector_chunk_loc( + node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size(); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + inline std::pair get_key_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + inline std::pair get_neighbor_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_abort(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + return std::make_pair(node_chunks_[chunk_idx].get(), offset); + } + + // Get chunk location for neighbor distance data. + // Uses the same chunk indexing as node chunks but with dist_entry_size_. + inline std::pair get_dist_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * dist_entry_size_; + return std::make_pair(chunk_idx, offset); + } + + void sync_dist_chunks(size_t idx) const { + sync_chunks(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST, idx, &dist_chunks_); + } + + int ensure_dist_chunk_for(uint32_t chunk_index); + + int alloc_dist_chunks_for_existing_nodes(); + + size_t estimate_doc_capacity() const { + return node_chunks_.capacity() * node_cnt_per_chunk_; + } + + int init_chunk_params(size_t max_index_size, bool huge_page) { + node_cnt_per_chunk_ = std::max(1, chunk_size_ / node_size()); + node_index_mask_bits_ = std::ceil(std::log2(node_cnt_per_chunk_)); + node_cnt_per_chunk_ = 1UL << node_index_mask_bits_; + if (huge_page) { + chunk_size_ = AlignHugePageSize(node_cnt_per_chunk_ * node_size()); + } else { + chunk_size_ = AlignPageSize(node_cnt_per_chunk_ * node_size()); + } + node_index_mask_ = node_cnt_per_chunk_ - 1; + + if (max_index_size == 0UL) { + max_index_size_ = chunk_size_ * kDefaultMaxChunkCnt; + } else { + max_index_size_ = max_index_size; + } + + size_t max_node_chunk_cnt = + std::ceil(static_cast(max_index_size_) / chunk_size_); + node_chunks_.reserve(max_node_chunk_cnt); + + LOG_DEBUG( + "VamanaSettings: nodeSize=%zu chunkSize=%u nodeCntPerChunk=%u " + "maxChunkCnt=%zu maxIndexSize=%zu", + node_size(), chunk_size_, node_cnt_per_chunk_, max_node_chunk_cnt, + max_index_size_); + + return 0; + } + + int init_chunks(const Chunk::Pointer &header_chunk); + + int flush_header(void) { + if (!broker_->dirty()) { + return 0; + } + auto header_chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId); + if (ailego_unlikely(!header_chunk)) { + LOG_ERROR("get header chunk failed"); + return IndexError_Runtime; + } + size_t size = header_chunk->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + return 0; + } + + protected: + IndexStreamer::Stats &stats_; + std::mutex mutex_{}; + size_t max_index_size_{0UL}; + uint32_t chunk_size_{kDefaultChunkSize}; + uint32_t node_index_mask_bits_{0U}; + uint32_t node_cnt_per_chunk_{0U}; + uint32_t node_index_mask_{0U}; + uint32_t neighbor_size_{0U}; + bool get_vector_enabled_{false}; + bool use_key_info_map_{true}; + + mutable std::shared_ptr keys_map_lock_; + HashMapPointer keys_map_; + + ChunkBroker::Pointer broker_; + + //! Protects node_chunks_ against concurrent emplace_back from add_vector + //! (writer) and sync_chunks from greedy_search (reader threads during build). + mutable std::mutex node_chunks_mutex_{}; + mutable std::vector node_chunks_{}; + + private: + mutable std::vector dist_chunks_{}; + bool dist_loaded_{false}; + uint32_t dist_entry_size_{0}; // max_degree * sizeof(dist_t) +}; + +// --- Template specializations for typed MemoryBlock access --- + +template <> +inline NeighborsT +VamanaStreamerEntity::get_neighbors_typed(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(offset < node_chunks_[chunk_idx]->data_size(), + "invalid chunk offset"); + const void *ptr = nullptr; + size_t ret = node_chunks_[chunk_idx]->read(offset, &ptr, neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return NeighborsT(); + } + MmapMemoryBlock block(const_cast(ptr)); + return NeighborsT(std::move(block)); +} + +template <> +inline NeighborsT +VamanaStreamerEntity::get_neighbors_typed( + node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + IndexStorage::MemoryBlock mem_block; + size_t ret = node_chunks_[chunk_idx]->read(offset, mem_block, neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return NeighborsT(); + } + BufferPoolMemoryBlock block; + if (mem_block.type_ == IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH) { + block = BufferPoolMemoryBlock::MakeOwned(mem_block.data_); + mem_block.data_ = nullptr; + mem_block.type_ = IndexStorage::MemoryBlock::MBT_UNKNOWN; + } else { + block = BufferPoolMemoryBlock(mem_block.buffer_pool_handle_, + mem_block.buffer_block_id_, mem_block.data_); + mem_block.buffer_pool_handle_ = nullptr; + } + return NeighborsT(std::move(block)); +} + +template <> +inline int VamanaStreamerEntity::get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, vector_size()); + if (ailego_unlikely(ret != vector_size())) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return IndexError_ReadData; + } + vec_blocks[i].reset(const_cast(ptr)); + } + return 0; +} + +template <> +inline int VamanaStreamerEntity::get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + IndexStorage::MemoryBlock mem_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, mem_block, vector_size()); + if (ailego_unlikely(ret != vector_size())) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return IndexError_ReadData; + } + vec_blocks[i] = [&]() { + if (mem_block.type_ == IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH) { + BufferPoolMemoryBlock b = + BufferPoolMemoryBlock::MakeOwned(mem_block.data_); + mem_block.data_ = nullptr; + mem_block.type_ = IndexStorage::MemoryBlock::MBT_UNKNOWN; + return b; + } + BufferPoolMemoryBlock b(mem_block.buffer_pool_handle_, + mem_block.buffer_block_id_, mem_block.data_); + mem_block.buffer_pool_handle_ = nullptr; + return b; + }(); + } + return 0; +} + +template <> +inline key_t VamanaStreamerEntity::get_key_typed( + node_id_t id) const { + if (!use_key_info_map_) return id; + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(ptr); +} + +template <> +inline key_t VamanaStreamerEntity::get_key_typed( + node_id_t id) const { + if (!use_key_info_map_) return id; + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + IndexStorage::MemoryBlock key_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(key_block.data()); +} + +// --- Typed entity subclass for mmap mode --- +// Caches chunk base addresses to eliminate virtual function calls on the +// search hot path. For mmap mode, chunk data is memory-mapped at init time, +// so we can directly compute pointers via base_addr + offset. +class VamanaMmapStreamerEntity : public VamanaStreamerEntity { + public: + using MemoryBlock = MmapMemoryBlock; + using TypedNeighbors = NeighborsT; + + using VamanaStreamerEntity::VamanaStreamerEntity; + + VamanaStorageMode storage_mode() const override { + return VamanaStorageMode::kMmap; + } + + //! Override clone to return correct subclass type, so that + //! static_cast in the algorithm is safe. + const VamanaEntity::Pointer clone() const override; + + ailego_force_inline TypedNeighbors get_neighbors_typed(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + const char *base = get_node_chunk_base(chunk_idx); + MmapMemoryBlock block(const_cast(base + offset)); + return TypedNeighbors(std::move(block)); + } + + ailego_force_inline int get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + uint32_t chunk_idx = ids[i] >> node_index_mask_bits_; + uint32_t offset = (ids[i] & node_index_mask_) * node_size(); + const char *base = get_node_chunk_base(chunk_idx); + vec_blocks[i].reset(const_cast(base + offset)); + } + return 0; + } + + ailego_force_inline key_t get_key_typed(node_id_t id) const { + if (!use_key_info_map_) return id; + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + const char *base = get_node_chunk_base(chunk_idx); + return *reinterpret_cast(base + offset); + } + + //! Direct vector pointer access (no MemoryBlock wrapper). + //! For use in the merged search loop to avoid intermediate allocations. + ailego_force_inline const void *get_vector_ptr(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size(); + return get_node_chunk_base(chunk_idx) + offset; + } + + private: + ailego_force_inline const char *get_node_chunk_base( + uint32_t chunk_idx) const { + if (ailego_unlikely(chunk_idx >= node_chunk_bases_.size())) { + sync_node_chunk_bases(chunk_idx); + } + return node_chunk_bases_[chunk_idx]; + } + + void sync_node_chunk_bases(uint32_t chunk_idx) const { + std::lock_guard lock(chunk_bases_mutex_); + // Double-check after acquiring lock to avoid redundant sync + if (chunk_idx < node_chunk_bases_.size()) { + return; + } + // Pre-reserve to match node_chunks_ capacity so that subsequent + // push_back never triggers reallocation — the lock-free fast path in + // get_node_chunk_base reads existing elements without holding the mutex. + if (node_chunk_bases_.capacity() < node_chunks_.capacity()) { + node_chunk_bases_.reserve(node_chunks_.capacity()); + } + sync_node_chunks(chunk_idx); + const auto &chunks = node_chunks_; + for (size_t i = node_chunk_bases_.size(); i <= chunk_idx; ++i) { + const void *ptr = nullptr; + chunks[i]->read(0, &ptr, 1); + node_chunk_bases_.push_back(static_cast(ptr)); + } + } + + mutable std::mutex chunk_bases_mutex_{}; + mutable std::vector node_chunk_bases_{}; +}; + +// --- Typed entity subclass for buffer pool mode --- +class VamanaBufferPoolStreamerEntity : public VamanaStreamerEntity { + public: + using MemoryBlock = BufferPoolMemoryBlock; + using TypedNeighbors = NeighborsT; + + using VamanaStreamerEntity::VamanaStreamerEntity; + + VamanaStorageMode storage_mode() const override { + return VamanaStorageMode::kBufferPool; + } + + inline TypedNeighbors get_neighbors_typed(node_id_t id) const { + return VamanaStreamerEntity::get_neighbors_typed(id); + } + + inline int get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + return VamanaStreamerEntity::get_vector_typed( + ids, count, vec_blocks); + } + + inline key_t get_key_typed(node_id_t id) const { + return VamanaStreamerEntity::get_key_typed(id); + } +}; + +// --- Typed entity subclass for contiguous memory mode --- +// Splits node data into two dense arrays during build: +// 1. vector_base_: flat vector array (stride = vector_size) +// 2. graph_base_: key + neighbors (stride = graph_stride_) +// Total memory = vector_size + graph_stride_ per node (same as original +// node_size), but each access pattern gets optimal cache locality. +class VamanaContiguousStreamerEntity : public VamanaMmapStreamerEntity { + public: + using VamanaMmapStreamerEntity::VamanaMmapStreamerEntity; + + VamanaStorageMode storage_mode() const override { + return VamanaStorageMode::kContiguous; + } + + //! Override clone to return correct subclass type. + //! Cloned entity shares contiguous memory via shared_ptr. + const VamanaEntity::Pointer clone() const override; + + ~VamanaContiguousStreamerEntity() = default; + + // Build contiguous memory from chunks after open. + int build_contiguous_memory(); + + //! Degrade to mmap mode by releasing contiguous memory and falling back + //! to chunk-based access. + void degrade_to_mmap() { + vector_memory_.reset(); + vector_base_ = nullptr; + vector_stride_ = 0; + graph_memory_.reset(); + graph_base_ = nullptr; + LOG_INFO("Vamana contiguous entity degraded to mmap mode for insertion"); + } + + bool is_contiguous() const { + return vector_base_ != nullptr; + } + + //! Per-entry stride of the flat vector array (0 if no contiguous build). + //! Padded up to kVectorAlignment (64B), so it is also the amount that + //! should be prefetched per vector. + size_t vector_stride() const { + return vector_stride_; + } + + int add_vector(key_t key, const void *vec, node_id_t *id) override { + if (ailego_unlikely(is_contiguous())) degrade_to_mmap(); + return VamanaMmapStreamerEntity::add_vector(key, vec, id); + } + + int add_vector_with_id(node_id_t id, const void *vec) override { + if (ailego_unlikely(is_contiguous())) degrade_to_mmap(); + return VamanaMmapStreamerEntity::add_vector_with_id(id, vec); + } + + ailego_force_inline TypedNeighbors get_neighbors_typed(node_id_t id) const { + if (ailego_likely(graph_base_ != nullptr)) { + // graph layout: [key (sizeof(key_t)) | NeighborsHeader + neighbors] + const char *ptr = + graph_base_ + static_cast(id) * graph_stride_ + sizeof(key_t); + MmapMemoryBlock block(const_cast(ptr)); + return TypedNeighbors(std::move(block)); + } + return VamanaMmapStreamerEntity::get_neighbors_typed(id); + } + + ailego_force_inline int get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + if (ailego_likely(vector_base_ != nullptr)) { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + const char *ptr = + vector_base_ + static_cast(ids[i]) * vector_stride_; + vec_blocks[i].reset(const_cast(ptr)); + } + return 0; + } + return VamanaMmapStreamerEntity::get_vector_typed(ids, count, vec_blocks); + } + + ailego_force_inline key_t get_key_typed(node_id_t id) const { + if (ailego_likely(graph_base_ != nullptr)) { + if (!use_key_info_map_) return id; + // key is at offset 0 within each graph node + const char *ptr = graph_base_ + static_cast(id) * graph_stride_; + return *reinterpret_cast(ptr); + } + return VamanaMmapStreamerEntity::get_key_typed(id); + } + + //! Direct vector pointer from flat vector array. + //! Stride is padded up to kVectorAlignment (64B) to preserve cache-line + //! alignment even when vector_size is not a multiple of 64. The padding is + //! purely in-memory and does NOT affect the on-disk index file layout. + ailego_force_inline const void *get_vector_ptr(node_id_t id) const { + if (ailego_likely(vector_base_ != nullptr)) { + return vector_base_ + static_cast(id) * vector_stride_; + } + return VamanaMmapStreamerEntity::get_vector_ptr(id); + } + + protected: + //! Custom deleter for contiguous memory allocated via + //! MemoryHelper::AllocateHugePage. `size` is the (already huge-page-aligned) + //! length passed at allocation time, required by the mmap/munmap path. + struct ContiguousDeleter { + size_t size; + void operator()(char *ptr) const { + ailego::MemoryHelper::FreeHugePage(ptr, size); + } + }; + + //! Flat vector array: vectors stored densely with per-vector stride + //! padded up to kVectorAlignment (64B) to keep each vector's starting + //! address cache-line aligned. Base is page-aligned by the allocator. + std::shared_ptr vector_memory_{}; + char *vector_base_{nullptr}; + //! Per-vector stride = AlignUp(vector_size(), kVectorAlignment). + size_t vector_stride_{0}; + + //! Graph array: [key | neighbors] stored densely (stride = graph_stride_). + std::shared_ptr graph_memory_{}; + char *graph_base_{nullptr}; + size_t graph_stride_{0}; // sizeof(key_t) + neighbors_size() + + //! Cache-line alignment used for per-vector stride in the flat array. + static constexpr size_t kVectorAlignment = 64; + + private: + static char *allocate_contiguous(size_t size); +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/CMakeLists.txt b/src/core/framework/CMakeLists.txt new file mode 100644 index 0000000..5091592 --- /dev/null +++ b/src/core/framework/CMakeLists.txt @@ -0,0 +1,11 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +cc_library( + NAME core_framework + STATIC STRICT ALWAYS_LINK + SRCS *.cc + LIBS zvec_ailego + INCS . ${PROJECT_ROOT_DIR}/src/core + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/framework/index_cluster.cc b/src/core/framework/index_cluster.cc new file mode 100644 index 0000000..5fbda92 --- /dev/null +++ b/src/core/framework/index_cluster.cc @@ -0,0 +1,159 @@ +// namespace aitheta2 +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +namespace zvec { +namespace core { + +static const std::string CLUSTER_CENTROIDS_FEATURES_NAME = + "IndexCluster.Centroids.Features"; +static const std::string CLUSTER_CENTROIDS_INDEXES_NAME = + "IndexCluster.Centroids.Indexes"; + +/*! Item Centroid Format + */ +struct ItemCentroidFormat { + uint32_t parent; + uint32_t reserved0_; + uint64_t follows; + double score; + uint64_t reserved1_; +}; + +static inline bool GatherSubitemsCount(const ItemCentroidFormat *format, + size_t count, + std::vector *out) { + out->resize(count + 1); + + for (const ItemCentroidFormat *it = format, *end = format + count; it != end; + ++it) { + uint32_t parent = it->parent + 1; + if (parent > count) { + return false; + } + (*out)[parent] += 1; + } + return (out->front() != 0); +} + +int IndexCluster::Deserialize(const IndexMeta &meta, + IndexBundle::Pointer bundle, + CentroidList *cents) { + if (!bundle || !cents) { + return IndexError_InvalidArgument; + } + + ailego::BlobWrap features = bundle->get(CLUSTER_CENTROIDS_FEATURES_NAME); + ailego::BlobWrap indexes = bundle->get(CLUSTER_CENTROIDS_INDEXES_NAME); + + if (!features.is_valid() || !indexes.is_valid()) { + return IndexError_InvalidArgument; + } + + if (features.size() % meta.element_size() != 0 || + indexes.size() % sizeof(ItemCentroidFormat) != 0) { + return IndexError_InvalidLength; + } + + size_t count = features.size() / meta.element_size(); + if (indexes.size() / sizeof(ItemCentroidFormat) != count) { + return IndexError_InvalidLength; + } + + const ItemCentroidFormat *format = + reinterpret_cast(indexes.buffer()); + std::vector subitems; + + if (!GatherSubitemsCount(format, count, &subitems)) { + return IndexError_InvalidFormat; + } + + std::vector items; + items.reserve(count); + cents->clear(); + cents->reserve(subitems.front()); + + const uint8_t *feat = reinterpret_cast(features.buffer()); + size_t feat_size = meta.element_size(); + + for (size_t i = 0; i < count; ++i, ++format, feat += feat_size) { + CentroidList *current = cents; + + if (format->parent != static_cast(-1)) { + if (format->parent >= items.size()) { + return IndexError_InvalidFormat; + } + current = items[format->parent]->mutable_subitems(); + } + current->emplace_back(feat, feat_size); + + // Update information + Centroid *last_one = &(current->back()); + last_one->set_follows(static_cast(format->follows)); + last_one->set_score(format->score); + last_one->mutable_subitems()->reserve(subitems[i + 1]); + items.push_back(last_one); + } + return 0; +} + +static void SerializeToBuffers(const IndexCluster::CentroidList ¢s, + std::string *features, std::string *indexes) { + uint32_t parent = + static_cast(indexes->size() / sizeof(ItemCentroidFormat)) - 1; + + for (const auto &it : cents) { + ItemCentroidFormat format{parent, 0, it.follows(), it.score(), 0}; + indexes->append(reinterpret_cast(&format), sizeof(format)); + features->append(reinterpret_cast(it.feature()), it.size()); + + if (!it.subitems().empty()) { + SerializeToBuffers(it.subitems(), features, indexes); + } + } +} + +int IndexCluster::Serialize(const IndexMeta &meta, const CentroidList ¢s, + IndexBundle::Pointer *out) { + size_t cents_total = cents.size(); + + // Check the centroids + for (const auto &it : cents) { + if (!it.is_matched(meta)) { + return IndexError_Mismatch; + } + cents_total += it.subcount(); + } + + std::string features, indexes; + features.reserve(cents_total * meta.element_size()); + indexes.reserve(cents_total * sizeof(ItemCentroidFormat)); + SerializeToBuffers(cents, &features, &indexes); + + std::shared_ptr bundle = + std::make_shared(); + + bundle->set(CLUSTER_CENTROIDS_FEATURES_NAME, std::move(features)); + bundle->set(CLUSTER_CENTROIDS_INDEXES_NAME, std::move(indexes)); + *out = std::move(bundle); + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_context.cc b/src/core/framework/index_context.cc new file mode 100644 index 0000000..08812d5 --- /dev/null +++ b/src/core/framework/index_context.cc @@ -0,0 +1,28 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +namespace zvec { +namespace core { + +uint32_t IndexContext::GenerateMagic(void) { + static std::atomic_uint32_t magic_number{std::random_device()()}; + return magic_number.fetch_add(1); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_converter.cc b/src/core/framework/index_converter.cc new file mode 100644 index 0000000..dcf8480 --- /dev/null +++ b/src/core/framework/index_converter.cc @@ -0,0 +1,43 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +namespace zvec { +namespace core { + +int IndexConverter::TrainAndTransform(const IndexConverter::Pointer &converter, + IndexHolder::Pointer holder) { + auto two_pass_holder = IndexHelper::MakeTwoPassHolder(std::move(holder)); + int ret = converter->train(two_pass_holder); + if (ret == 0) { + ret = converter->transform(std::move(two_pass_holder)); + } + return ret; +} + +int IndexConverter::TrainTransformAndDump( + const IndexConverter::Pointer &converter, IndexHolder::Pointer holder, + const IndexDumper::Pointer &dumper) { + int ret = IndexConverter::TrainAndTransform(converter, std::move(holder)); + if (ret == 0) { + ret = converter->dump(dumper); + } + return ret; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_error.cc b/src/core/framework/index_error.cc new file mode 100644 index 0000000..0a259ec --- /dev/null +++ b/src/core/framework/index_error.cc @@ -0,0 +1,73 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include + +namespace zvec { +namespace core { + +INDEX_ERROR_CODE_DEFINE(Success, 0, "Success"); +INDEX_ERROR_CODE_DEFINE(Runtime, 1, "Runtime error"); +INDEX_ERROR_CODE_DEFINE(Logic, 2, "Logic error"); +INDEX_ERROR_CODE_DEFINE(Type, 3, "Type error"); +INDEX_ERROR_CODE_DEFINE(System, 4, "System call error"); +INDEX_ERROR_CODE_DEFINE(Cast, 5, "Cast error"); +INDEX_ERROR_CODE_DEFINE(IO, 6, "IO error"); +INDEX_ERROR_CODE_DEFINE(AuthExpired, 7, "Auth expired error"); + +INDEX_ERROR_CODE_DEFINE(NotImplemented, 11, "Not implemented"); +INDEX_ERROR_CODE_DEFINE(Unsupported, 12, "Unsupported"); +INDEX_ERROR_CODE_DEFINE(Denied, 13, "Permission denied"); +INDEX_ERROR_CODE_DEFINE(Canceled, 14, "Operation canceled"); +INDEX_ERROR_CODE_DEFINE(Overflow, 15, "Overflow"); +INDEX_ERROR_CODE_DEFINE(Underflow, 16, "Underflow"); +INDEX_ERROR_CODE_DEFINE(OutOfRange, 17, "Out of range"); +INDEX_ERROR_CODE_DEFINE(NoBuffer, 18, "No buffer space available"); +INDEX_ERROR_CODE_DEFINE(NoMemory, 19, "Not enough space"); +INDEX_ERROR_CODE_DEFINE(NoParamFound, 20, "No parameter found"); +INDEX_ERROR_CODE_DEFINE(NoReady, 21, "No ready"); +INDEX_ERROR_CODE_DEFINE(NoExist, 22, "No exist"); +INDEX_ERROR_CODE_DEFINE(Exist, 23, "Already exist"); +INDEX_ERROR_CODE_DEFINE(Mismatch, 24, "Mismatch"); +INDEX_ERROR_CODE_DEFINE(Duplicate, 25, "Duplicate"); +INDEX_ERROR_CODE_DEFINE(Uninitialized, 26, "Uninitialized"); + +INDEX_ERROR_CODE_DEFINE(InvalidArgument, 31, "Invalid argument"); +INDEX_ERROR_CODE_DEFINE(InvalidFormat, 32, "Invalid format"); +INDEX_ERROR_CODE_DEFINE(InvalidLength, 33, "Invalid length"); +INDEX_ERROR_CODE_DEFINE(InvalidChecksum, 34, "Invalid checksum"); +INDEX_ERROR_CODE_DEFINE(InvalidValue, 35, "Invalid value"); + +INDEX_ERROR_CODE_DEFINE(CreateDirectory, 101, "Create directory error"); +INDEX_ERROR_CODE_DEFINE(OpenDirectory, 102, "Open directory error"); +INDEX_ERROR_CODE_DEFINE(Serialize, 105, "Serialize error"); +INDEX_ERROR_CODE_DEFINE(Deserialize, 106, "Deserialize error"); +INDEX_ERROR_CODE_DEFINE(CreateFile, 111, "Create file error"); +INDEX_ERROR_CODE_DEFINE(OpenFile, 112, "Open file error"); +INDEX_ERROR_CODE_DEFINE(SeekFile, 113, "Seek file error"); +INDEX_ERROR_CODE_DEFINE(CloseFile, 114, "Close file error"); +INDEX_ERROR_CODE_DEFINE(TruncateFile, 115, "TruncateFile file error"); +INDEX_ERROR_CODE_DEFINE(MMapFile, 116, "MMap file error"); +INDEX_ERROR_CODE_DEFINE(FlushFile, 117, "Flush file error"); +INDEX_ERROR_CODE_DEFINE(WriteData, 121, "Write data error"); +INDEX_ERROR_CODE_DEFINE(ReadData, 122, "Read data error"); + +INDEX_ERROR_CODE_DEFINE(PackIndex, 201, "Read data error"); +INDEX_ERROR_CODE_DEFINE(UnpackIndex, 202, "Read data error"); +INDEX_ERROR_CODE_DEFINE(IndexLoaded, 203, "Index loaded"); +INDEX_ERROR_CODE_DEFINE(NoIndexLoaded, 204, "No index loaded"); +INDEX_ERROR_CODE_DEFINE(NoTrained, 205, "No trained"); +INDEX_ERROR_CODE_DEFINE(IndexFull, 206, "Index full"); + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_factory.cc b/src/core/framework/index_factory.cc new file mode 100644 index 0000000..69fe0e9 --- /dev/null +++ b/src/core/framework/index_factory.cc @@ -0,0 +1,261 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +namespace zvec { +namespace core { + +IndexMetric::Pointer IndexFactory::CreateMetric(const std::string &name) { + IndexMetric::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasMetric(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllMetrics(void) { + return ailego::Factory::Classes(); +} + +IndexLogger::Pointer IndexFactory::CreateLogger(const std::string &name) { + IndexLogger::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasLogger(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllLoggers(void) { + return ailego::Factory::Classes(); +} + +IndexDumper::Pointer IndexFactory::CreateDumper(const std::string &name) { + IndexDumper::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasDumper(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllDumpers(void) { + return ailego::Factory::Classes(); +} + +IndexStorage::Pointer IndexFactory::CreateStorage(const std::string &name) { + IndexStorage::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasStorage(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllStorages(void) { + return ailego::Factory::Classes(); +} + +IndexConverter::Pointer IndexFactory::CreateConverter(const std::string &name) { + IndexConverter::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasConverter(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllConverters(void) { + return ailego::Factory::Classes(); +} + +IndexReformer::Pointer IndexFactory::CreateReformer(const std::string &name) { + IndexReformer::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasReformer(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllReformers(void) { + return ailego::Factory::Classes(); +} + +IndexTrainer::Pointer IndexFactory::CreateTrainer(const std::string &name) { + IndexTrainer::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasTrainer(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllTrainers(void) { + return ailego::Factory::Classes(); +} + +IndexBuilder::Pointer IndexFactory::CreateBuilder(const std::string &name) { + IndexBuilder::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasBuilder(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllBuilders(void) { + return ailego::Factory::Classes(); +} + +IndexSearcher::Pointer IndexFactory::CreateSearcher(const std::string &name) { + IndexSearcher::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasSearcher(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllSearchers(void) { + return ailego::Factory::Classes(); +} + +IndexStreamer::Pointer IndexFactory::CreateStreamer(const std::string &name) { + IndexStreamer::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasStreamer(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllStreamers(void) { + return ailego::Factory::Classes(); +} + +IndexReducer::Pointer IndexFactory::CreateReducer(const std::string &name) { + IndexReducer::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasReducer(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllReducers(void) { + return ailego::Factory::Classes(); +} + + +IndexCluster::Pointer IndexFactory::CreateCluster(const std::string &name) { + IndexCluster::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasCluster(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllClusters(void) { + return ailego::Factory::Classes(); +} + +IndexStreamerReducer::Pointer IndexFactory::CreateStreamerReducer( + const std::string &name) { + IndexStreamerReducer::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasStreamerReducer(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllStreamerReducers(void) { + return ailego::Factory::Classes(); +} + +IndexRefiner::Pointer IndexFactory::CreateRefiner(const std::string &name) { + IndexRefiner::Pointer obj = + ailego::Factory::MakeShared(name.c_str()); + if (obj) { + obj->set_name(name); + } + return obj; +} + +bool IndexFactory::HasRefiner(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllRefiners(void) { + return ailego::Factory::Classes(); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_flow.cc b/src/core/framework/index_flow.cc new file mode 100644 index 0000000..d375c73 --- /dev/null +++ b/src/core/framework/index_flow.cc @@ -0,0 +1,882 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include + +//! Default storage +#define INDEX_FLOW_STORAGE_DEFAULT "MMapFileReadStorage" + +namespace zvec { +namespace core { + +// Index Flow +int IndexFlow::set_storage(const std::string &name, + const ailego::Params ¶ms) { + storage_ = IndexFactory::CreateStorage(name); + if (!storage_) { + LOG_ERROR("Failed to create a index storage with name: %s", name.c_str()); + return IndexError_NoExist; + } + int ret = storage_->init(params); + if (ret < 0) { + storage_ = nullptr; + LOG_ERROR("Failed to initialize index storage %s", name.c_str()); + return ret; + } + return 0; +} + +int IndexFlow::set_searcher(IndexSearcher::Pointer searcher) { + user_searcher_ = searcher; + + return 0; +} + +int IndexFlow::set_searcher(const std::string &name, + const ailego::Params ¶ms) { + user_searcher_ = IndexFactory::CreateSearcher(name); + if (!user_searcher_) { + LOG_ERROR("Failed to create a index searcher with name: %s", name.c_str()); + return IndexError_NoExist; + } + int ret = user_searcher_->init(params); + if (ret < 0) { + user_searcher_ = nullptr; + LOG_ERROR("Failed to initialize index searcher %s", name.c_str()); + return ret; + } + return 0; +} + +int IndexFlow::set_reformer(const std::string &name, + const ailego::Params ¶ms) { + user_reformer_ = IndexFactory::CreateReformer(name); + if (!user_reformer_) { + LOG_ERROR("Failed to create a index reformer with name: %s", name.c_str()); + return IndexError_NoExist; + } + int ret = user_reformer_->init(params); + if (ret < 0) { + user_reformer_ = nullptr; + LOG_ERROR("Failed to initialize index reformer %s", name.c_str()); + return ret; + } + return 0; +} + +int IndexFlow::set_metric(const std::string &name, + const ailego::Params ¶ms) { + if (!IndexFactory::HasMetric(name)) { + LOG_ERROR("The index metric with name %s does not exist.", name.c_str()); + return IndexError_NoExist; + } + user_metric_name_ = name; + user_metric_params_ = params; + return 0; +} + +int IndexFlow::load(const std::string &path) { + // Prepare storage + if (!storage_) { + this->set_storage(INDEX_FLOW_STORAGE_DEFAULT, ailego::Params()); + } + + if (!storage_) { + LOG_ERROR("The index storage is uninitialized."); + return IndexError_Uninitialized; + } + + int ret = storage_->open(path, false); + if (ret != 0) { + LOG_ERROR("Failed to load index with storage %s", storage_->name().c_str()); + return ret; + } + + ret = IndexHelper::DeserializeFromStorage(storage_.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize index meta with storage %s", + storage_->name().c_str()); + return ret; + } + + ret = load_internal(); + if (ret != 0) { + LOG_ERROR("Failed to load index with storage %s", storage_->name().c_str()); + return ret; + } + + return 0; +} + +int IndexFlow::load_internal() { + // Prepare metric + const std::string &metric_name = + user_metric_name_.empty() ? meta_.metric_name() : user_metric_name_; + const ailego::Params &metric_params = + user_metric_name_.empty() ? meta_.metric_params() : user_metric_params_; + if (metric_name.empty()) { + LOG_ERROR("The metric name from index file is empty."); + return IndexError_NoExist; + } + metric_ = IndexFactory::CreateMetric(metric_name); + if (!metric_) { + LOG_ERROR("Failed to create a index metric with name: %s", + metric_name.c_str()); + return IndexError_NoExist; + } + int ret = metric_->init(meta_, metric_params); + if (ret < 0) { + LOG_ERROR("Failed to initialize index metric %s", metric_name.c_str()); + metric_ = nullptr; + return ret; + } + if (!metric_->is_matched(meta_)) { + LOG_ERROR("The index meta is unmatched for index metric %s", + metric_->name().c_str()); + return IndexError_Mismatch; + } + auto query_metric = metric_->query_metric(); + if (query_metric) { + metric_ = query_metric; + } + + // Prepare reformer + if (!user_reformer_) { + const std::string &reformer_name = meta_.reformer_name(); + if (!reformer_name.empty()) { + reformer_ = IndexFactory::CreateReformer(reformer_name); + if (!reformer_) { + LOG_ERROR("Failed to create a index reformer with name: %s", + reformer_name.c_str()); + return IndexError_NoExist; + } + ret = reformer_->init(meta_.reformer_params()); + if (ret < 0) { + LOG_ERROR("Failed to initialize index reformer %s", + reformer_name.c_str()); + reformer_ = nullptr; + return ret; + } + } + } else { + // Using user reformer + reformer_ = user_reformer_; + } + + if (reformer_) { + ret = reformer_->load(storage_); + if (ret < 0) { + LOG_ERROR("Failed to load index with reformer %s, storage %s", + reformer_->name().c_str(), storage_->name().c_str()); + return ret; + } + } + + // Prepare searcher + if (!user_searcher_) { + const std::string &name = meta_.searcher_name(); + if (name.empty()) { + LOG_ERROR("The searcher name from index file is empty."); + return IndexError_NoExist; + } + searcher_ = IndexFactory::CreateSearcher(name); + if (!searcher_) { + LOG_ERROR("Failed to create a index searcher with name: %s", + name.c_str()); + return IndexError_NoExist; + } + ret = searcher_->init(meta_.searcher_params()); + if (ret < 0) { + LOG_ERROR("Failed to initialize index searcher %s", name.c_str()); + searcher_ = nullptr; + return ret; + } + } else { + // Using user searcher + searcher_ = user_searcher_; + } + + ret = searcher_->load(storage_, metric_); + if (ret < 0) { + LOG_ERROR("Failed to load index with searcher %s, storage %s, metric %s", + searcher_->name().c_str(), storage_->name().c_str(), + metric_->name().c_str()); + return ret; + } + + // searcher_->print_all_neighbour(); + + return 0; +} + +int IndexFlow::unload(void) { + if (searcher_) { + int ret = searcher_->unload(); + if (ret < 0) { + LOG_WARN("Unload index searcher %s error, %d", searcher_->name().c_str(), + ret); + } + searcher_ = nullptr; + } + if (reformer_) { + int ret = reformer_->unload(); + if (ret < 0) { + LOG_WARN("Unload index reformer %s error, %d", reformer_->name().c_str(), + ret); + } + reformer_ = nullptr; + } + if (metric_) { + int ret = metric_->cleanup(); + if (ret < 0) { + LOG_WARN("Cleanup index metric %s error, %d", metric_->name().c_str(), + ret); + } + metric_ = nullptr; + } + if (storage_) { + int ret = storage_->cleanup(); + if (ret < 0) { + LOG_WARN("Unload index searcher %s error, %d", storage_->name().c_str(), + ret); + } + storage_ = nullptr; + } + return 0; +} + +int IndexFlow::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + if (ailego_unlikely(!query || !context)) { + return IndexError_InvalidArgument; + } + + int error_code = 0; + if (reformer_) { + IndexQueryMeta new_qmeta; + error_code = reformer_->transform(query, qmeta, context->mutable_features(), + &new_qmeta); + if (error_code == 0) { + if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) { + return IndexError_Mismatch; + } + error_code = searcher_->search_bf_impl( + reinterpret_cast(context->features().data()), new_qmeta, + context->searcher_context()); + } + } else { + if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) { + return IndexError_Mismatch; + } + error_code = + searcher_->search_bf_impl(query, qmeta, context->searcher_context()); + } + + if (error_code == 0) { + if (metric_->support_normalize()) { + for (auto &it : const_cast( + context->searcher_context()->result())) { + metric_->normalize(it.mutable_score()); + } + } + if (reformer_) { + error_code = + reformer_->normalize(query, qmeta, + const_cast( + context->searcher_context()->result())); + } + } + return error_code; +} + +int IndexFlow::search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + if (ailego_unlikely(!query || !context)) { + return IndexError_InvalidArgument; + } + + int error_code = 0; + if (reformer_) { + IndexQueryMeta new_qmeta; + error_code = reformer_->transform(query, qmeta, context->mutable_features(), + &new_qmeta); + if (error_code == 0) { + if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) { + return IndexError_Mismatch; + } + error_code = searcher_->search_impl( + reinterpret_cast(context->features().data()), new_qmeta, + context->searcher_context()); + } + } else { + if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) { + return IndexError_Mismatch; + } + error_code = + searcher_->search_impl(query, qmeta, context->searcher_context()); + } + + if (error_code == 0) { + if (metric_->support_normalize()) { + for (auto &it : const_cast( + context->searcher_context()->result())) { + metric_->normalize(it.mutable_score()); + } + } + if (reformer_) { + error_code = + reformer_->normalize(query, qmeta, + const_cast( + context->searcher_context()->result())); + } + } + return error_code; +} + +int IndexFlow::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const { + if (ailego_unlikely(!query || !count || !context)) { + return IndexError_InvalidArgument; + } + + int error_code = 0; + if (reformer_) { + IndexQueryMeta new_qmeta; + error_code = reformer_->transform(query, qmeta, count, + context->mutable_features(), &new_qmeta); + if (error_code == 0) { + if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) { + return IndexError_Mismatch; + } + error_code = searcher_->search_bf_impl( + reinterpret_cast(context->features().data()), new_qmeta, + count, context->searcher_context()); + } + } else { + if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) { + return IndexError_Mismatch; + } + error_code = searcher_->search_bf_impl(query, qmeta, count, + context->searcher_context()); + } + + if (error_code == 0) { + if (metric_->support_normalize()) { + for (uint32_t i = 0; i < count; ++i) { + IndexDocumentList &result = const_cast( + context->searcher_context()->result(i)); + + for (auto &it : result) { + metric_->normalize(it.mutable_score()); + } + } + } + if (reformer_) { + size_t offset = 0; + for (uint32_t i = 0; i < count; ++i) { + error_code = reformer_->normalize( + reinterpret_cast(query) + offset, qmeta, + const_cast( + context->searcher_context()->result(i))); + if (error_code != 0) { + break; + } + offset += qmeta.element_size(); + } + } + } + return error_code; +} + +int IndexFlow::search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const { + if (ailego_unlikely(!query || !count || !context)) { + return IndexError_InvalidArgument; + } + + int error_code = 0; + if (reformer_) { + IndexQueryMeta new_qmeta; + error_code = reformer_->transform(query, qmeta, count, + context->mutable_features(), &new_qmeta); + if (error_code == 0) { + if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) { + return IndexError_Mismatch; + } + error_code = searcher_->search_impl( + reinterpret_cast(context->features().data()), new_qmeta, + count, context->searcher_context()); + } + } else { + if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) { + return IndexError_Mismatch; + } + error_code = searcher_->search_impl(query, qmeta, count, + context->searcher_context()); + } + + if (error_code == 0) { + if (metric_->support_normalize()) { + for (uint32_t i = 0; i < count; ++i) { + IndexDocumentList &result = const_cast( + context->searcher_context()->result(i)); + + for (auto &it : result) { + metric_->normalize(it.mutable_score()); + } + } + } + if (reformer_) { + size_t offset = 0; + for (uint32_t i = 0; i < count; ++i) { + error_code = reformer_->normalize( + reinterpret_cast(query) + offset, qmeta, + const_cast( + context->searcher_context()->result(i))); + if (error_code != 0) { + break; + } + offset += qmeta.element_size(); + } + } + } + return error_code; +} + +// Index Sparse Flow +int IndexSparseFlow::set_storage(const std::string &name, + const ailego::Params ¶ms) { + storage_ = IndexFactory::CreateStorage(name); + if (!storage_) { + LOG_ERROR("Failed to create a index storage with name: %s", name.c_str()); + return IndexError_NoExist; + } + int ret = storage_->init(params); + if (ret < 0) { + storage_ = nullptr; + LOG_ERROR("Failed to initialize index storage %s", name.c_str()); + return ret; + } + return 0; +} + +int IndexSparseFlow::set_searcher(IndexSearcher::Pointer searcher) { + user_searcher_ = searcher; + + return 0; +} + +int IndexSparseFlow::set_searcher(const std::string &name, + const ailego::Params ¶ms) { + user_searcher_ = IndexFactory::CreateSearcher(name); + if (!user_searcher_) { + LOG_ERROR("Failed to create a index sparse searcher with name: %s", + name.c_str()); + return IndexError_NoExist; + } + int ret = user_searcher_->init(params); + if (ret < 0) { + user_searcher_ = nullptr; + LOG_ERROR("Failed to initialize index sparse searcher %s", name.c_str()); + return ret; + } + return 0; +} + +int IndexSparseFlow::set_reformer(const std::string &name, + const ailego::Params ¶ms) { + user_reformer_ = IndexFactory::CreateReformer(name); + if (!user_reformer_) { + LOG_ERROR("Failed to create a index sparse reformer with name: %s", + name.c_str()); + return IndexError_NoExist; + } + int ret = user_reformer_->init(params); + if (ret < 0) { + user_reformer_ = nullptr; + LOG_ERROR("Failed to initialize index sparse reformer %s", name.c_str()); + return ret; + } + return 0; +} + +int IndexSparseFlow::set_metric(const std::string &name, + const ailego::Params ¶ms) { + if (!IndexFactory::HasMetric(name)) { + LOG_ERROR("The index metric with name %s does not exist.", name.c_str()); + return IndexError_NoExist; + } + user_metric_name_ = name; + user_metric_params_ = params; + return 0; +} + +int IndexSparseFlow::load(const std::string &path) { + // Prepare storage + if (!storage_) { + this->set_storage(INDEX_FLOW_STORAGE_DEFAULT, ailego::Params()); + } + + if (!storage_) { + LOG_ERROR("The index storage is uninitialized."); + return IndexError_Uninitialized; + } + + int ret = storage_->open(path, false); + if (ret != 0) { + LOG_ERROR("Failed to load index with storage %s", storage_->name().c_str()); + return ret; + } + + ret = IndexHelper::DeserializeFromStorage(storage_.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize index meta with storage %s", + storage_->name().c_str()); + return ret; + } + + ret = load_internal(); + if (ret != 0) { + LOG_ERROR("Failed to load index with storage %s", storage_->name().c_str()); + return ret; + } + + return 0; +} + +int IndexSparseFlow::load_internal() { + // Prepare metric + const std::string &metric_name = + user_metric_name_.empty() ? meta_.metric_name() : user_metric_name_; + const ailego::Params &metric_params = + user_metric_name_.empty() ? meta_.metric_params() : user_metric_params_; + if (metric_name.empty()) { + LOG_ERROR("The metric name from index file is empty."); + return IndexError_NoExist; + } + metric_ = IndexFactory::CreateMetric(metric_name); + if (!metric_) { + LOG_ERROR("Failed to create a index metric with name: %s", + metric_name.c_str()); + return IndexError_NoExist; + } + int ret = metric_->init(meta_, metric_params); + if (ret < 0) { + LOG_ERROR("Failed to initialize index metric %s", metric_name.c_str()); + metric_ = nullptr; + return ret; + } + + auto query_metric = metric_->query_metric(); + if (query_metric) { + metric_ = query_metric; + } + + // Prepare reformer + if (!user_reformer_) { + const std::string &reformer_name = meta_.reformer_name(); + if (!reformer_name.empty()) { + reformer_ = IndexFactory::CreateReformer(reformer_name); + if (!reformer_) { + LOG_ERROR("Failed to create a index sparse reformer with name: %s", + reformer_name.c_str()); + return IndexError_NoExist; + } + ret = reformer_->init(meta_.reformer_params()); + if (ret < 0) { + LOG_ERROR("Failed to initialize index reformer %s", + reformer_name.c_str()); + reformer_ = nullptr; + return ret; + } + } + } else { + // Using user reformer + reformer_ = user_reformer_; + } + + if (reformer_) { + ret = reformer_->load(storage_); + if (ret < 0) { + LOG_ERROR("Failed to load index with reformer %s, storage %s", + reformer_->name().c_str(), storage_->name().c_str()); + return ret; + } + } + + // Prepare searcher + if (!user_searcher_) { + const std::string &name = meta_.searcher_name(); + if (name.empty()) { + LOG_ERROR("The searcher name from index file is empty."); + return IndexError_NoExist; + } + searcher_ = IndexFactory::CreateSearcher(name); + if (!searcher_) { + LOG_ERROR("Failed to create a index searcher with name: %s", + name.c_str()); + return IndexError_NoExist; + } + ret = searcher_->init(meta_.searcher_params()); + if (ret < 0) { + LOG_ERROR("Failed to initialize index searcher %s", name.c_str()); + searcher_ = nullptr; + return ret; + } + } else { + // Using user searcher + searcher_ = user_searcher_; + } + + ret = searcher_->load(storage_, metric_); + if (ret < 0) { + LOG_ERROR("Failed to load index with searcher %s, storage %s, metric %s", + searcher_->name().c_str(), storage_->name().c_str(), + metric_->name().c_str()); + return ret; + } + + // searcher_->print_all_neighbour(); + + return 0; +} + +int IndexSparseFlow::unload(void) { + if (searcher_) { + int ret = searcher_->unload(); + if (ret < 0) { + LOG_WARN("Unload index searcher %s error, %d", searcher_->name().c_str(), + ret); + } + searcher_ = nullptr; + } + if (reformer_) { + int ret = reformer_->unload(); + if (ret < 0) { + LOG_WARN("Unload index reformer %s error, %d", reformer_->name().c_str(), + ret); + } + reformer_ = nullptr; + } + if (metric_) { + int ret = metric_->cleanup(); + if (ret < 0) { + LOG_WARN("Cleanup index metric %s error, %d", metric_->name().c_str(), + ret); + } + metric_ = nullptr; + } + if (storage_) { + int ret = storage_->cleanup(); + if (ret < 0) { + LOG_WARN("Unload index searcher %s error, %d", storage_->name().c_str(), + ret); + } + storage_ = nullptr; + } + return 0; +} + +int IndexSparseFlow::search_bf_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + if (ailego_unlikely(!context)) { + return IndexError_InvalidArgument; + } + + int error_code = 0; + if (reformer_) { + std::string ovec; + IndexQueryMeta new_qmeta; + error_code = reformer_->transform(sparse_count, sparse_indices, + sparse_query, qmeta, &ovec, &new_qmeta); + if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) { + return IndexError_Mismatch; + } + if (error_code == 0) { + error_code = + searcher_->search_bf_impl(sparse_count, sparse_indices, ovec.data(), + new_qmeta, context->searcher_context()); + } + } else { + if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) { + return IndexError_Mismatch; + } + error_code = + searcher_->search_bf_impl(sparse_count, sparse_indices, sparse_query, + qmeta, context->searcher_context()); + } + + if (error_code == 0) { + if (metric_->support_normalize()) { + for (auto &it : const_cast( + context->searcher_context()->result())) { + metric_->normalize(it.mutable_score()); + } + } + } + return error_code; +} + +int IndexSparseFlow::search_impl(const uint32_t sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + if (ailego_unlikely(!context)) { + return IndexError_InvalidArgument; + } + + int error_code = 0; + if (reformer_) { + std::string ovec; + IndexQueryMeta new_qmeta; + error_code = reformer_->transform(sparse_count, sparse_indices, + sparse_query, qmeta, &ovec, &new_qmeta); + if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) { + return IndexError_Mismatch; + } + if (error_code == 0) { + error_code = + searcher_->search_impl(sparse_count, sparse_indices, ovec.data(), + new_qmeta, context->searcher_context()); + } + } else { + if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) { + return IndexError_Mismatch; + } + error_code = + searcher_->search_impl(sparse_count, sparse_indices, sparse_query, + qmeta, context->searcher_context()); + } + + if (error_code == 0) { + if (metric_->support_normalize()) { + for (auto &it : const_cast( + context->searcher_context()->result())) { + metric_->normalize(it.mutable_score()); + } + } + } + return error_code; +} + +int IndexSparseFlow::search_bf_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (ailego_unlikely(!count || !context)) { + return IndexError_InvalidArgument; + } + + int error_code = 0; + if (reformer_) { + std::string ovec; + IndexQueryMeta new_qmeta; + error_code = + reformer_->transform(sparse_count, sparse_indices, sparse_query, qmeta, + count, &ovec, &new_qmeta); + + if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) { + return IndexError_Mismatch; + } + + if (error_code == 0) { + error_code = searcher_->search_bf_impl(sparse_count, sparse_indices, + ovec.data(), new_qmeta, count, + context->searcher_context()); + } + } else { + if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) { + return IndexError_Mismatch; + } + error_code = + searcher_->search_bf_impl(sparse_count, sparse_indices, sparse_query, + qmeta, count, context->searcher_context()); + } + + if (error_code == 0) { + if (metric_->support_normalize()) { + for (uint32_t i = 0; i < count; ++i) { + IndexDocumentList &result = const_cast( + context->searcher_context()->result(i)); + + for (auto &it : result) { + metric_->normalize(it.mutable_score()); + } + } + } + } + return error_code; +} + +int IndexSparseFlow::search_impl(const uint32_t *sparse_count, + const uint32_t *sparse_indices, + const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (ailego_unlikely(!count || !context)) { + return IndexError_InvalidArgument; + } + + int error_code = 0; + if (reformer_) { + std::string ovec; + IndexQueryMeta new_qmeta; + error_code = + reformer_->transform(sparse_count, sparse_indices, sparse_query, qmeta, + count, &ovec, &new_qmeta); + + if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) { + return IndexError_Mismatch; + } + + if (error_code == 0) { + error_code = + searcher_->search_impl(sparse_count, sparse_indices, ovec.data(), + new_qmeta, count, context->searcher_context()); + } + } else { + if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) { + return IndexError_Mismatch; + } + error_code = + searcher_->search_impl(sparse_count, sparse_indices, sparse_query, + qmeta, count, context->searcher_context()); + } + + if (error_code == 0) { + if (metric_->support_normalize()) { + for (uint32_t i = 0; i < count; ++i) { + IndexDocumentList &result = const_cast( + context->searcher_context()->result(i)); + + for (auto &it : result) { + metric_->normalize(it.mutable_score()); + } + } + } + } + return error_code; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_helper.cc b/src/core/framework/index_helper.cc new file mode 100644 index 0000000..1401fe6 --- /dev/null +++ b/src/core/framework/index_helper.cc @@ -0,0 +1,255 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +namespace zvec { +namespace core { + +int IndexHelper::SerializeToDumper(const IndexMeta &mt, IndexDumper *dumper, + const std::string &key) { + std::string buffer; + mt.serialize(&buffer); + + size_t data_size = buffer.size(); + uint32_t data_crc = ailego::Crc32c::Hash(buffer.data(), buffer.size(), 0); + buffer.resize((data_size + 31u) & ~31u); + + if (dumper->write(buffer.data(), buffer.size()) != buffer.size()) { + return IndexError_WriteData; + } + if (dumper->append(key, data_size, buffer.size() - data_size, data_crc) != + 0) { + return IndexError_WriteData; + } + return IndexError_Success; +} + +int IndexHelper::SerializeToStorage(const IndexMeta &mt, IndexStorage *storage, + const std::string &key) { + std::string buffer; + mt.serialize(&buffer); + + auto segment = storage->get(key); + if (!segment) { + const size_t align_size = 4096 * 4; + size_t meta_size = + (buffer.size() + align_size - 1) / align_size * align_size; + + if (storage->append(key, meta_size) != 0) { + return IndexError_WriteData; + } + + segment = storage->get(key); + if (!segment) { + return IndexError_NoExist; + } + } + + if (segment->write(0, buffer.data(), buffer.size()) != buffer.size()) { + return IndexError_WriteData; + } + segment->resize(buffer.size()); + segment->update_data_crc( + ailego::Crc32c::Hash(buffer.data(), buffer.size(), 0)); + return IndexError_Success; +} + +int IndexHelper::DeserializeFromStorage(IndexStorage *storage, + const std::string &key, + IndexMeta *out) { + auto segment = storage->get(key); + if (!segment) { + return IndexError_NoExist; + } + + uint32_t crc = segment->data_crc(); + size_t len = segment->data_size(); + IndexStorage::MemoryBlock block; + if (segment->read(0, block, len) != len) { + return IndexError_ReadData; + } + const void *data = block.data(); + if (crc != 0u && ailego::Crc32c::Hash(data, len, 0u) != crc) { + return IndexError_InvalidChecksum; + } + if (!out->deserialize(data, len)) { + return IndexError_Deserialize; + } + return IndexError_Success; +} + +/*! Two Pass Index Holder + */ +class TwoPassIndexHolder : public IndexHolder { + private: + /*! First Pass Iterator + * store elements during iterating for second iterating. + */ + class FirstPassIterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + FirstPassIterator(TwoPassIndexHolder *owner, + IndexHolder::Iterator::Pointer &&iter) + : holder_(owner), front_iter_(std::move(iter)) {} + + //! Destructor + ~FirstPassIterator(void) override {} + + //! Retrieve pointer of data + const void *data(void) const override { + return front_iter_->data(); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return front_iter_->is_valid(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return front_iter_->key(); + } + + //! Next iterator + void next(void) override { + holder_->features_.emplace_back( + front_iter_->key(), std::string((const char *)front_iter_->data(), + holder_->front_->element_size())); + front_iter_->next(); + } + + private: + TwoPassIndexHolder *holder_{nullptr}; + IndexHolder::Iterator::Pointer front_iter_{}; + }; + + class SecondPassIterator : public IndexHolder::Iterator { + public: + //! Second Pass Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + SecondPassIterator(TwoPassIndexHolder *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + ~SecondPassIterator(void) override {} + + //! Retrieve pointer of data + const void *data(void) const override { + return features_iter_->second.data(); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return (features_iter_ != holder_->features_.end()); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return features_iter_->first; + } + + //! Next iterator + void next(void) override { + holder_->features_.erase(features_iter_++); + } + + private: + TwoPassIndexHolder *holder_{nullptr}; + typename std::list>::iterator + features_iter_{}; + }; + + public: + //! Constructor + TwoPassIndexHolder(IndexHolder::Pointer &&front) + : front_(std::move(front)), + data_type_(front_->data_type()), + dimension_(front_->dimension()), + element_size_(front_->element_size()), + count_(front_->count()) {} + + //! Retrieve count of elements in holder (-1 indicates unknown) + size_t count(void) const override { + return count_; + } + + //! Retrieve dimension + size_t dimension(void) const override { + return dimension_; + } + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return data_type_; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return element_size_; + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return false; + } + + //! Create a new iterator + IndexHolder::Iterator::Pointer create_iterator(void) override { + ++pass_; + if (pass_ == 1) { + IndexHolder::Iterator::Pointer iter = front_->create_iterator(); + return iter ? IndexHolder::Iterator::Pointer( + new TwoPassIndexHolder::FirstPassIterator( + this, std::move(iter))) + : IndexHolder::Iterator::Pointer(); + } else if (pass_ == 2) { + return IndexHolder::Iterator::Pointer( + new TwoPassIndexHolder::SecondPassIterator(this)); + } + return nullptr; + } + + private: + //! Disable them + TwoPassIndexHolder(void) = delete; + + //! Members + IndexHolder::Pointer front_{}; + std::list> features_{}; + size_t pass_{0}; + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_UNDEFINED}; + size_t dimension_; + size_t element_size_; + size_t count_; +}; + +IndexHolder::Pointer IndexHelper::MakeTwoPassHolder( + IndexHolder::Pointer holder) { + if (holder->multipass()) { + return holder; + } + return IndexHolder::Pointer(new TwoPassIndexHolder(std::move(holder))); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_logger.cc b/src/core/framework/index_logger.cc new file mode 100644 index 0000000..4ab6e2e --- /dev/null +++ b/src/core/framework/index_logger.cc @@ -0,0 +1,72 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +const int IndexLogger::LEVEL_DEBUG = 0; +const int IndexLogger::LEVEL_INFO = 1; +const int IndexLogger::LEVEL_WARN = 2; +const int IndexLogger::LEVEL_ERROR = 3; +const int IndexLogger::LEVEL_FATAL = 4; + +/*! Console Logger + */ +struct ConsoleLogger : public IndexLogger { + //! Initialize Logger + int init(const zvec::ailego::Params &) override { + return 0; + } + + //! Cleanup Logger + int cleanup(void) override { + return 0; + } + + //! Log Message + void log(int level, const char *file, int line, const char *format, + va_list args) override { + char buffer[8192]; + std::ostringstream stream; + + ailego::Realtime::Localtime(buffer, sizeof(buffer)); + stream << '[' << LevelString(level) << ' ' << buffer << ' ' + << std::this_thread::get_id() << ' ' << ailego::File::BaseName(file) + << ':' << line << "] "; + + vsnprintf(buffer, sizeof(buffer), format, args); + stream << buffer << '\n'; + + if (level <= LEVEL_INFO) { + std::cout << stream.str() << std::flush; + } else { + std::cerr << stream.str() << std::flush; + } + } +}; + +//! Logger Level +int IndexLoggerBroker::logger_level_ = 0; + +//! Logger +IndexLogger::Pointer IndexLoggerBroker::logger_(new ConsoleLogger); + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/framework/index_mapping.cc b/src/core/framework/index_mapping.cc new file mode 100644 index 0000000..7f5f5a1 --- /dev/null +++ b/src/core/framework/index_mapping.cc @@ -0,0 +1,617 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "ailego/utility/memory_helper.h" + +#if defined(_WIN32) || defined(_WIN64) +#include +#endif + +#ifdef __linux__ +#include +#include +#ifndef HUGETLBFS_MAGIC +#define HUGETLBFS_MAGIC 0x958458f6 +#endif +#endif + +namespace zvec { +namespace core { + +static inline size_t CalcPageAlignedSize(size_t size, bool huge_size) { + size_t page_size = ailego::MemoryHelper::PageSize(); + if (huge_size) { + page_size = ailego::MemoryHelper::HugePageSize(); + } + return (size + page_size - 1) / page_size * page_size; +} + +static inline bool WritePadding(ailego::File &file, size_t size) { + std::string padding(ailego::MemoryHelper::PageSize(), 0); + for (size_t i = 0, count = size / padding.size(); i < count; ++i) { + if (file.write(padding.data(), padding.size()) != padding.size()) { + return false; + } + } + padding.resize(size % padding.size()); + if (padding.size()) { + if (file.write(padding.data(), padding.size()) != padding.size()) { + return false; + } + } + return true; +} + +static inline int UnpackMappingSize(ailego::File &file, size_t *len) { + IndexFormat::MetaHeader header; + if (file.read(&header, sizeof(header)) != sizeof(header)) { + LOG_ERROR("Failed to read file, %s", + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_ReadData; + } + + if (header.meta_header_size != sizeof(IndexFormat::MetaHeader) || + header.meta_footer_size != sizeof(IndexFormat::MetaFooter)) { + return IndexError_InvalidValue; + } + + if (ailego::Crc32c::Hash(&header, sizeof(header), header.header_crc) != + header.header_crc) { + return IndexError_InvalidChecksum; + } + + if ((int32_t)header.meta_footer_offset < 0) { + return IndexError_Unsupported; + } + + *len = header.meta_footer_offset + header.meta_footer_size; + if (*len > file.size()) { + return IndexError_InvalidLength; + } + return 0; +} + +int IndexMapping::open(const std::string &path, bool cow, bool full_mode) { + path_ = path; + full_mode_ = full_mode; + copy_on_write_ = cow; + huge_page_ = Ishugetlbfs(path); + + bool read_only = copy_on_write_ && !full_mode_; + if (!file_.open(path.c_str(), read_only, false)) { + LOG_ERROR("Failed to open file %s, %s", path.c_str(), + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_OpenFile; + } + + size_t mapping_size = 0u; + int error_code = UnpackMappingSize(file_, &mapping_size); + if (error_code != 0) { + file_.close(); + return error_code; + } + + if (!file_.seek(0, ailego::File::Origin::End)) { + LOG_ERROR("Failed to seek file %s, %s", path.c_str(), + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_SeekFile; + } + return this->init_index_mapping(mapping_size); +} + +int IndexMapping::create(const std::string &path, size_t seg_meta_capacity) { + path_ = path; + seg_meta_capacity_ = seg_meta_capacity; + current_header_start_offset_ = 0; + + // write() & copying to mmap() will auto extend the file size + if (!file_.create(path.c_str(), 0)) { + LOG_ERROR("Failed to create file %s, %s", path.c_str(), + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_CreateFile; + } + huge_page_ = Ishugetlbfs(path); + if (huge_page_) { + return init_hugepage_meta_section(); + } + return init_meta_section(); +} + +int IndexMapping::init_meta_section() { + if (current_header_start_offset_ % ailego::MemoryHelper::PageSize() != 0) { + LOG_ERROR("File offset %zu is not a multiple of the page size: %zu", + (size_t)current_header_start_offset_, + ailego::MemoryHelper::PageSize()); + return IndexError_InvalidValue; + } + + auto &path = path_; + size_t len = + CalcPageAlignedSize(seg_meta_capacity_ + sizeof(IndexFormat::MetaHeader) + + sizeof(IndexFormat::MetaFooter), + false); + + IndexFormat::MetaHeader meta_header; + IndexFormat::MetaFooter meta_footer; + + // Write index header + IndexFormat::SetupMetaHeader(&meta_header, len - sizeof(meta_footer), len); + if (!file_.seek(current_header_start_offset_, ailego::File::Origin::Begin)) { + LOG_ERROR("Failed to seek file %s, %s", path.c_str(), + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_SeekFile; + } + if (file_.write(&meta_header, sizeof(meta_header)) != sizeof(meta_header)) { + LOG_ERROR("Failed to write file: %s, %s", path.c_str(), + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_WriteData; + } + + // Write padding data + uint32_t segments_meta_size = + static_cast(len - (sizeof(meta_header) + sizeof(meta_footer))); + if (!WritePadding(file_, segments_meta_size)) { + LOG_ERROR("Failed to write file: %s, %s", path.c_str(), + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_WriteData; + } + + // Write index footer + IndexFormat::SetupMetaFooter(&meta_footer); + meta_footer.segments_meta_size = segments_meta_size; + meta_footer.total_size = len; + IndexFormat::UpdateMetaFooter(&meta_footer, 0); + if (file_.write(&meta_footer, sizeof(meta_footer)) != sizeof(meta_footer)) { + LOG_ERROR("Failed to write file: %s, %s", path.c_str(), + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_WriteData; + } + return this->init_index_mapping(len); +} + +int IndexMapping::init_hugepage_meta_section() { + ssize_t file_offset = (ssize_t)current_header_start_offset_; + if (file_offset % ailego::MemoryHelper::HugePageSize() != 0) { + LOG_ERROR("File offset %zu is not a multiple of the page size: %zu", + file_offset, ailego::MemoryHelper::HugePageSize()); + return IndexError_InvalidValue; + } + + size_t len = + CalcPageAlignedSize(seg_meta_capacity_ + sizeof(IndexFormat::MetaHeader) + + sizeof(IndexFormat::MetaFooter), + true); + int opts = ailego::File::MMAP_SHARED | ailego::File::MMAP_HUGE_PAGE; + void *addr = + ailego::File::MemoryMap(file_.native_handle(), file_offset, len, opts); + + IndexFormat::MetaHeader meta_header; + IndexFormat::MetaFooter meta_footer; + + // Write index header + IndexFormat::SetupMetaHeader(&meta_header, len - sizeof(meta_footer), len); + memcpy((char *)addr + file_offset, &meta_header, sizeof(meta_header)); + file_offset += sizeof(meta_header); + + // Write padding data + uint32_t segments_meta_size = + static_cast(len - (sizeof(meta_header) + sizeof(meta_footer))); + std::string padding(ailego::MemoryHelper::HugePageSize(), 0); + for (size_t i = 0, count = segments_meta_size / padding.size(); i < count; + ++i) { + memcpy((char *)addr + file_offset, padding.data(), padding.size()); + file_offset += padding.size(); + } + padding.resize(segments_meta_size % padding.size()); + if (padding.size()) { + memcpy((char *)addr + file_offset, padding.data(), padding.size()); + file_offset += padding.size(); + } + + // Write index footer + IndexFormat::SetupMetaFooter(&meta_footer); + meta_footer.segments_meta_size = segments_meta_size; + meta_footer.total_size = len; + IndexFormat::UpdateMetaFooter(&meta_footer, 0); + memcpy((char *)addr + file_offset, &meta_footer, sizeof(meta_footer)); + file_offset += sizeof(meta_footer); + + return this->init_index_mapping(len); +} + +void IndexMapping::close(void) { + // Unmap all memory + this->unmap_all(); + if (header_) { + for (auto item : header_addr_map_) { + auto header = item.second; + ailego::File::MemoryUnmap(header, header->content_offset); + } + } + // Reset members + segment_ids_offset_ = 0; + segment_start_ = nullptr; + header_ = nullptr; + header_addr_map_.clear(); + footer_ = nullptr; + index_size_ = 0u; + segments_.clear(); + file_.close(); + copy_on_write_ = false; + full_mode_ = false; + header_dirty_ = false; + huge_page_ = false; +} + +void IndexMapping::refresh(uint64_t check_point) { + // support add_with_id + for (auto item : header_addr_map_) { + auto header = item.second; + auto footer = reinterpret_cast( + reinterpret_cast(header) + header->meta_footer_offset); + auto segment_start = reinterpret_cast( + reinterpret_cast(header) + + (header->meta_footer_offset - footer->segments_meta_size)); + footer->segments_meta_crc = + ailego::Crc32c::Hash(segment_start, footer->segments_meta_size, 0); + IndexFormat::UpdateMetaFooter(footer, check_point); + } + header_dirty_ = true; +} + +int IndexMapping::append(const std::string &id, size_t size) { + size = CalcPageAlignedSize(size, huge_page_); + if (size == 0) { + return IndexError_InvalidArgument; + } + + if (segments_.find(id) != segments_.end()) { + return IndexError_Duplicate; + } + + size_t id_size = std::strlen(id.c_str()) + 1; + size_t need_size = sizeof(IndexFormat::SegmentMeta) + id_size; + if (sizeof(IndexFormat::SegmentMeta) * footer_->segment_count + need_size > + segment_ids_offset_) { + LOG_DEBUG("segment meta section expanded: %s", path_.c_str()); + footer_->next_meta_header_offset = index_size_; + refresh(0); + flush(); + // mmap file storage write() will update segment's meta + // ailego::File::MemoryUnmap(header_, header_->content_offset); + header_ = nullptr; + footer_ = nullptr; + + current_header_start_offset_ = index_size_; + const int ret = + huge_page_ ? init_hugepage_meta_section() : init_meta_section(); + if (ret != 0) { + return ret; + } + } + + if (!copy_on_write_ && !file_.truncate(index_size_ + size)) { + LOG_ERROR("Failed to truncate file, %s", + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_TruncateFile; + } + + // Update segment table + segment_ids_offset_ -= static_cast(id_size); + IndexFormat::SegmentMeta *segment = segment_start_ + footer_->segment_count; + segment->segment_id_offset = segment_ids_offset_; + segment->data_index = + index_size_ - header_->content_offset - current_header_start_offset_; + segment->data_size = 0; + segment->data_crc = 0; + segment->padding_size = size; + memcpy((uint8_t *)segment_start_ + segment_ids_offset_, id.c_str(), id_size); + index_size_ += size; + + // Update index footer + footer_->segments_meta_crc = + ailego::Crc32c::Hash(segment_start_, footer_->segments_meta_size, 0); + footer_->segment_count += 1; + footer_->content_size += size; + footer_->total_size += size; + IndexFormat::UpdateMetaFooter(footer_, 0); + segments_.emplace( + id, SegmentInfo{Segment{segment}, current_header_start_offset_, header_}); + header_dirty_ = true; + return 0; +} + +IndexMapping::Segment *IndexMapping::map(const std::string &id, bool warmup, + bool locked) { + auto iter = segments_.find(id); + if (iter == segments_.end()) { + return nullptr; + } + SegmentInfo &segment_info = iter->second; + Segment *item = &segment_info.segment; + if (!item->data()) { + auto meta = item->meta(); + size_t mapping_size = meta->data_size + meta->padding_size; + size_t offset = segment_info.segment_header_start_offset + + segment_info.segment_header->content_offset + + meta->data_index; + + void *addr = nullptr; + if (!copy_on_write_) { + int opts = ailego::File::MMAP_SHARED; + if (huge_page_) { + opts |= ailego::File::MMAP_HUGE_PAGE; + } + addr = ailego::File::MemoryMap(file_.native_handle(), offset, + mapping_size, opts); + } else { + size_t file_size = file_.size(); + int opts = ailego::File::MMAP_POPULATE; + if (huge_page_) { + opts |= ailego::File::MMAP_HUGE_PAGE; + } + if (offset < file_size) { + ailego_assert(offset + mapping_size <= file_size); + addr = ailego::File::MemoryMap(file_.native_handle(), offset, + mapping_size, opts); + } else { + addr = ailego::File::MemoryMap(mapping_size, opts); + } + } + + if (!addr) { + LOG_ERROR("Map segment failed, segment id %s", id.c_str()); + return nullptr; + } + item->set_data(addr); + + // Lock memory + if (locked) { + ailego::File::MemoryLock(item->data(), mapping_size); + } + // Warmup memory + if (warmup && meta->data_size) { + ailego::File::MemoryWarmup(item->data(), meta->data_size); + } + } + return item; +} + +void IndexMapping::unmap(const std::string &id) { + auto iter = segments_.find(id); + if (iter != segments_.end()) { + SegmentInfo &segment_info = iter->second; + Segment *item = &segment_info.segment; + + if (item->data()) { + ailego::File::MemoryUnmap( + item->data(), item->meta()->data_size + item->meta()->padding_size); + item->set_data(nullptr); + } + } +} + +void IndexMapping::unmap_all(void) { + for (auto iter = segments_.begin(); iter != segments_.end(); ++iter) { + SegmentInfo &segment_info = iter->second; + Segment *item = &segment_info.segment; + + if (item->data()) { + ailego::File::MemoryUnmap( + item->data(), item->meta()->data_size + item->meta()->padding_size); + item->set_data(nullptr); + } + } +} + +int IndexMapping::flush(void) { + if ((file_.size() < index_size_) && !file_.truncate(index_size_)) { + LOG_ERROR("Failed to truncate file size %zu, %s", index_size_, + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_TruncateFile; + } + + for (auto iter = segments_.begin(); iter != segments_.end(); ++iter) { + SegmentInfo &segment_info = iter->second; + Segment *item = &segment_info.segment; + if (!item->data() || !item->dirty()) { + continue; + } + + size_t segment_size = item->meta()->data_size + item->meta()->padding_size; + if (full_mode_ && copy_on_write_) { + size_t off = segment_info.segment_header_start_offset + + segment_info.segment_header->content_offset + + item->meta()->data_index; + if (file_.write(off, item->data(), segment_size) != segment_size) { + LOG_ERROR("Failed to write segment, size %zu, %s", segment_size, + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_WriteData; + } + } else { + ailego::File::MemoryFlush(item->data(), segment_size); + } + item->reset_dirty(); + } + + if (!header_dirty_) { + return 0; + } + + header_dirty_ = false; + if (full_mode_ && copy_on_write_) { + for (auto item : header_addr_map_) { + auto header_start_offset = item.first; + auto header = item.second; + if (file_.write(header_start_offset, header, header->content_offset) != + header->content_offset) { + LOG_ERROR("Failed to write segment, size %zu, %s", + (size_t)header->content_offset, + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_WriteData; + } + } + } else { + for (auto item : header_addr_map_) { + auto header = item.second; + ailego::File::MemoryFlush(header, header->content_offset); + } + } + return 0; +} + +int IndexMapping::init_index_mapping(size_t len) { + int opts = + copy_on_write_ ? ailego::File::MMAP_POPULATE : ailego::File::MMAP_SHARED; + if (huge_page_) { + opts |= ailego::File::MMAP_HUGE_PAGE; + } + uint8_t *start = reinterpret_cast(ailego::File::MemoryMap( + file_.native_handle(), current_header_start_offset_, len, opts)); + if (!start) { + LOG_ERROR("Failed to map file, %s", + ailego::FileHelper::GetLastErrorString().c_str()); + return IndexError_MMapFile; + } + + // Unpack header + header_ = reinterpret_cast(start); + header_addr_map_.insert({current_header_start_offset_, header_}); + if (header_->meta_header_size != sizeof(IndexFormat::MetaHeader)) { + return IndexError_InvalidLength; + } + if (ailego::Crc32c::Hash(header_, sizeof(*header_), header_->header_crc) != + header_->header_crc) { + return IndexError_InvalidChecksum; + } + + switch (header_->version) { + case IndexFormat::FORMAT_VERSION: + break; + default: + LOG_ERROR("Unsupported index version: %u", header_->version); + return IndexError_Unsupported; + } + + // Unpack footer + if (header_->meta_footer_size != sizeof(IndexFormat::MetaFooter)) { + return IndexError_InvalidLength; + } + if ((int32_t)header_->meta_footer_offset < 0) { + return IndexError_Unsupported; + } + size_t footer_offset = header_->meta_footer_offset; + if (footer_offset + header_->meta_footer_size > len) { + return IndexError_InvalidLength; + } + + footer_ = reinterpret_cast(start + footer_offset); + if (footer_offset < footer_->segments_meta_size) { + return IndexError_InvalidLength; + } + + index_size_ = file_.size(); + if ((footer_->total_size > index_size_) || + (footer_->content_size + footer_->content_padding_size + + header_->content_offset > + index_size_)) { + return IndexError_InvalidLength; + } + if (ailego::Crc32c::Hash(footer_, sizeof(*footer_), footer_->footer_crc) != + footer_->footer_crc) { + return IndexError_InvalidChecksum; + } + + // Unpack segment table + if (sizeof(IndexFormat::SegmentMeta) * footer_->segment_count > + footer_->segments_meta_size) { + return IndexError_InvalidLength; + } + + segment_start_ = reinterpret_cast( + start + (footer_offset - footer_->segments_meta_size)); + if (ailego::Crc32c::Hash(segment_start_, footer_->segments_meta_size, 0u) != + footer_->segments_meta_crc) { + LOG_ERROR("Index segments meta checksum is invalid."); + return IndexError_InvalidChecksum; + } + + segment_ids_offset_ = footer_->segments_meta_size; + for (IndexFormat::SegmentMeta *iter = segment_start_, + *end = segment_start_ + footer_->segment_count; + iter != end; ++iter) { + if (iter->segment_id_offset > footer_->segments_meta_size) { + return IndexError_InvalidValue; + } + if (iter->data_index > footer_->content_size) { + return IndexError_InvalidValue; + } + if (iter->data_index + iter->data_size > footer_->content_size) { + return IndexError_InvalidLength; + } + + if (iter->segment_id_offset < segment_ids_offset_) { + segment_ids_offset_ = iter->segment_id_offset; + } + segments_.emplace( + std::string(reinterpret_cast(segment_start_) + + iter->segment_id_offset), + SegmentInfo{Segment{iter}, current_header_start_offset_, header_}); + } + if (sizeof(IndexFormat::SegmentMeta) * footer_->segment_count > + segment_ids_offset_) { + return IndexError_InvalidLength; + } + + // if (header_->version == IndexFormat::COMPATIBLE_FORMAT_VERSION_0X0002) { + // header_->version = IndexFormat::CURRENT_FORMAT_VERSION; + // LOG_INFO("Index file format upgraded"); + // IndexFormat::UpdateMetaHeader(header_); + // footer_->segments_meta_crc = + // ailego::Crc32c::Hash(segment_start_, footer_->segments_meta_size, 0); + // IndexFormat::UpdateMetaFooter(footer_, 0); + // header_dirty_ = true; + // } + + if (footer_->next_meta_header_offset > 0) { + current_header_start_offset_ = footer_->next_meta_header_offset; + // Meta sections have all the same size, so we can use the same size to map + // the next meta section + return this->init_index_mapping(len); + } + + return 0; +} + +bool IndexMapping::Ishugetlbfs(const std::string &path) const { +#ifdef __linux__ + struct statfs buf; + if (statfs(path.c_str(), &buf) != 0) { + perror("statfs"); + return false; + } + return static_cast(buf.f_type) == HUGETLBFS_MAGIC; +#else + static_cast(path); + return false; +#endif +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_meta.cc b/src/core/framework/index_meta.cc new file mode 100644 index 0000000..11d54cb --- /dev/null +++ b/src/core/framework/index_meta.cc @@ -0,0 +1,208 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +namespace zvec { +namespace core { + +/*! Index Meta Buffer Format + */ +struct IndexMetaFormatHeader { + uint32_t header_size; + uint32_t meta_type; + uint32_t major_order; + uint32_t data_type; + uint32_t dimension; + uint32_t unit_size; + uint32_t space_id; + uint32_t attachment_offset; + uint32_t attachment_size; + uint8_t reserved_[4092]; +}; + +static_assert(sizeof(IndexMetaFormatHeader) % 32 == 0, + "IndexMetaBufferFormat must be aligned with 32 bytes"); + +void IndexMeta::serialize(std::string *out) const { + ailego::Params attachment; + IndexMetaFormatHeader format; + memset(&format, 0, sizeof(format)); + format.header_size = sizeof(format); + format.meta_type = static_cast(meta_type_); + format.major_order = static_cast(major_order_); + format.data_type = static_cast(data_type_); + format.dimension = dimension_; + format.unit_size = unit_size_; + format.space_id = space_id_; + + if (!metric_name_.empty()) { + ailego::Params item; + item.set("name", metric_name_); + item.set("revision", metric_revision_); + item.set("params", metric_params_); + attachment.set("metric", std::move(item)); + } + + if (!converter_name_.empty()) { + ailego::Params item; + item.set("name", converter_name_); + item.set("revision", converter_revision_); + item.set("params", converter_params_); + attachment.set("converter", std::move(item)); + } + if (!reformer_name_.empty()) { + ailego::Params item; + item.set("name", reformer_name_); + item.set("revision", reformer_revision_); + item.set("params", reformer_params_); + attachment.set("reformer", std::move(item)); + } + if (!trainer_name_.empty()) { + ailego::Params item; + item.set("name", trainer_name_); + item.set("revision", trainer_revision_); + item.set("params", trainer_params_); + attachment.set("trainer", std::move(item)); + } + if (!builder_name_.empty()) { + ailego::Params item; + item.set("name", builder_name_); + item.set("revision", builder_revision_); + item.set("params", builder_params_); + attachment.set("builder", std::move(item)); + } + if (!reducer_name_.empty()) { + ailego::Params item; + item.set("name", reducer_name_); + item.set("revision", reducer_revision_); + item.set("params", reducer_params_); + attachment.set("reducer", std::move(item)); + } + if (!searcher_name_.empty()) { + ailego::Params item; + item.set("name", searcher_name_); + item.set("revision", searcher_revision_); + item.set("params", searcher_params_); + attachment.set("searcher", std::move(item)); + } + if (!streamer_name_.empty()) { + ailego::Params item; + item.set("name", streamer_name_); + item.set("revision", streamer_revision_); + item.set("params", streamer_params_); + attachment.set("streamer", std::move(item)); + } + + if (!attributes_.empty()) { + attachment.set("attributes", attributes_); + } + + out->assign(reinterpret_cast(&format), sizeof(format)); + size_t offset = static_cast(out->size()); + + if (!attachment.empty()) { + std::string buf; + ailego::Params::SerializeToBuffer(attachment, &buf); + out->append(buf.data(), buf.size()); + IndexMetaFormatHeader *header = (IndexMetaFormatHeader *)out->data(); + header->attachment_offset = static_cast(offset); + header->attachment_size = static_cast(buf.size()); + offset += buf.size(); + } +} + +bool IndexMeta::deserialize(const void *data, size_t len) { + const IndexMetaFormatHeader *format = + reinterpret_cast(data); + + this->clear(); + if (sizeof(IndexMetaFormatHeader) > len) { + return false; + } + if (sizeof(IndexMetaFormatHeader) > format->header_size) { + return false; + } + + meta_type_ = static_cast(format->meta_type); + major_order_ = static_cast(format->major_order); + data_type_ = static_cast(format->data_type); + dimension_ = format->dimension; + unit_size_ = format->unit_size; + element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_); + space_id_ = format->space_id; + + // Read attachment + ailego::Params attachment; + if (format->attachment_size) { + if (format->attachment_offset + format->attachment_size > len) { + return false; + } + std::string str( + reinterpret_cast(data) + format->attachment_offset, + format->attachment_size); + if (!ailego::Params::ParseFromBuffer(str, &attachment)) { + return false; + } + } + + ailego::Params item; + if (attachment.get("metric", &item)) { + item.get("name", &metric_name_); + item.get("revision", &metric_revision_); + item.get("params", &metric_params_); + } + if (attachment.get("converter", &item)) { + item.get("name", &converter_name_); + item.get("revision", &converter_revision_); + item.get("params", &converter_params_); + } + if (attachment.get("reformer", &item)) { + item.get("name", &reformer_name_); + item.get("revision", &reformer_revision_); + item.get("params", &reformer_params_); + } + if (attachment.get("trainer", &item)) { + item.get("name", &trainer_name_); + item.get("revision", &trainer_revision_); + item.get("params", &trainer_params_); + } + if (attachment.get("builder", &item)) { + item.get("name", &builder_name_); + item.get("revision", &builder_revision_); + item.get("params", &builder_params_); + } + if (attachment.get("reducer", &item)) { + item.get("name", &reducer_name_); + item.get("revision", &reducer_revision_); + item.get("params", &reducer_params_); + } + if (attachment.get("searcher", &item)) { + item.get("name", &searcher_name_); + item.get("revision", &searcher_revision_); + item.get("params", &searcher_params_); + } + if (attachment.get("streamer", &item)) { + item.get("name", &streamer_name_); + item.get("revision", &streamer_revision_); + item.get("params", &streamer_params_); + } + attachment.get("attributes", &attributes_); + + return true; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_plugin.cc b/src/core/framework/index_plugin.cc new file mode 100644 index 0000000..f9209a6 --- /dev/null +++ b/src/core/framework/index_plugin.cc @@ -0,0 +1,59 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include + +namespace zvec { +namespace core { + +bool IndexPlugin::load(const std::string &path) { + if (handle_) { + return false; + } + handle_ = ailego::DLHelper::Load(path, nullptr); + return (!!handle_); +} + +bool IndexPlugin::load(const std::string &path, std::string *err) { + if (handle_) { + *err = "plugin loaded"; + return false; + } + handle_ = ailego::DLHelper::Load(path, err); + return !!handle_; +} + +void IndexPlugin::unload(void) { + if (handle_) { + ailego::DLHelper::Unload(handle_); + handle_ = nullptr; + } +} + +bool IndexPluginBroker::emplace(IndexPlugin &&plugin) { + if (!plugin.is_valid()) { + return false; + } + for (auto iter = plugins_.begin(); iter != plugins_.end(); ++iter) { + if (iter->handle() == plugin.handle()) { + plugin.unload(); + return true; + } + } + plugins_.push_back(std::move(plugin)); + return true; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_version.cc b/src/core/framework/index_version.cc new file mode 100644 index 0000000..9cd33e7 --- /dev/null +++ b/src/core/framework/index_version.cc @@ -0,0 +1,34 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + + +namespace zvec { +namespace core { + +static const char AITHETA_VERSION_DETAILS[] = + AILEGO_VERSION_COMPILE_DETAILS("All rights reserved.\n"); + +const char *IndexVersion::String(void) { + return AITHETA_VERSION_DETAILS; +} + +const char *IndexVersion::Details(void) { + return AITHETA_VERSION_DETAILS; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/interface/CMakeLists.txt b/src/core/interface/CMakeLists.txt new file mode 100644 index 0000000..d5475a4 --- /dev/null +++ b/src/core/interface/CMakeLists.txt @@ -0,0 +1,10 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +cc_library( + NAME core_interface STATIC STRICT ALWAYS_LINK + SRCS *.cc indexes/*.cc + INCS . ${PROJECT_ROOT_DIR}/src/ ${PROJECT_ROOT_DIR}/src/core + LIBS zvec_ailego core_framework core_plugin sparsehash magic_enum rabitqlib + VERSION "${PROXIMA_ZVEC_VERSION}" +) diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc new file mode 100644 index 0000000..6c69530 --- /dev/null +++ b/src/core/interface/index.cc @@ -0,0 +1,1069 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include "mixed_reducer/mixed_reducer_params.h" +#include "utility/utility_params.h" + +namespace zvec::core_interface { + +namespace { + +bool has_group_by_search(const BaseIndexQueryParam::Pointer &search_param) { + return search_param->group_by_param && search_param->group_by_param->group_by; +} + +} // namespace + +// eliminate the pre-alloc of the context pool +thread_local static std::array() - 1) * 2> + _context_list; + + +bool Index::init_context() { + context_index_ = (magic_enum::enum_integer(param_.index_type) - 1) * 2 + + static_cast(is_sparse_); + if (_context_list[context_index_] == nullptr) { + if ((_context_list[context_index_] = streamer_->create_context()) == + nullptr) { + LOG_ERROR("Failed to create context"); + return false; + } + } + return true; +} + +core::IndexContext::Pointer &Index::acquire_context() { + init_context(); + return _context_list[context_index_]; +} + +int Index::ParseMetricName(const BaseIndexParam ¶m) { + std::string metric_name; + if (is_sparse_) { + // only inner product is supported for sparse index + switch (param.metric_type) { + case MetricType::kInnerProduct: + metric_name = "InnerProductSparse"; + break; + case MetricType::kMIPSL2sq: + metric_name = "MipsSquaredEuclideanSparse"; + break; + default: + LOG_ERROR("Unsupported metric type"); + return core::IndexError_Runtime; + } + } else { + switch (param.metric_type) { + case MetricType::kL2sq: + metric_name = "SquaredEuclidean"; + break; + case MetricType::kInnerProduct: + metric_name = "InnerProduct"; + break; + case MetricType::kCosine: + metric_name = "Cosine"; // This is already the normalizedCosine + break; + case MetricType::kMIPSL2sq: + metric_name = "MipsSquaredEuclidean"; + break; + default: + LOG_ERROR("Unsupported metric type"); + return core::IndexError_Runtime; + } + } + // TODO: MIPS need to set some param + // for streamer open() + proxima_index_meta_.set_metric(metric_name, 0, ailego::Params()); + return 0; +} + +int Index::CreateAndInitMetric(const BaseIndexParam & /*param*/) { + auto &metric_name = proxima_index_meta_.metric_name(); + + metric_ = core::IndexFactory::CreateMetric(metric_name); + if (!metric_) { + LOG_ERROR("Failed to create metric, name %s", metric_name.c_str()); + return core::IndexError_Runtime; + } + if (const auto ret = metric_->init(proxima_index_meta_, + proxima_index_meta_.metric_params()); + ret != 0) { + LOG_ERROR("Failed to create and init metric, name %s, code %d, desc: %s", + metric_name.c_str(), ret, core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (metric_->query_metric()) { + metric_ = metric_->query_metric(); + } + + return core::IndexError_Success; +} + +int Index::CreateAndInitConverterReformer(const QuantizerParam ¶m, + const BaseIndexParam &index_param) { + ailego::Params converter_params; + std::string converter_name; + if (is_sparse_) { + switch (param.type) { + case QuantizerType::kNone: + return core::IndexError_Success; + case QuantizerType::kFP16: + converter_name = "HalfFloatSparseConverter"; + break; + default: + LOG_ERROR("Unsupported quantizer type: "); + return core::IndexError_Unsupported; + } + } else { + if (index_param.metric_type == MetricType::kCosine) { + switch (param.type) { + case QuantizerType::kNone: + if (index_param.data_type == DataType::DT_FP16) { + converter_name = "CosineHalfFloatConverter"; + } else if (index_param.data_type == DataType::DT_FP32) { + converter_name = "CosineNormalizeConverter"; + } else { + LOG_ERROR("Unsupported data type: "); + return core::IndexError_Unsupported; + } + break; + case QuantizerType::kRabitq: + if (index_param.data_type == DataType::DT_FP32) { + converter_name = "CosineNormalizeConverter"; + } else { + LOG_ERROR("Unsupported data type: "); + return core::IndexError_Unsupported; + } + break; + case QuantizerType::kFP16: + converter_name = "CosineFp16Converter"; + break; + case QuantizerType::kInt8: + converter_name = "CosineInt8Converter"; + break; + case QuantizerType::kInt4: + converter_name = "CosineInt4Converter"; + break; + default: + LOG_ERROR("Unsupported quantizer type: "); + return core::IndexError_Unsupported; + } + } else { + switch (param.type) { + case QuantizerType::kNone: + return core::IndexError_Success; + case QuantizerType::kFP16: + converter_name = "HalfFloatConverter"; + break; + case QuantizerType::kInt8: + converter_name = "Int8StreamingConverter"; + break; + case QuantizerType::kInt4: + converter_name = "Int4StreamingConverter"; + break; + case QuantizerType::kRabitq: + // no converter here + return 0; + case QuantizerType::kUniformInt8: + converter_name = "UniformInt8StreamingConverter"; + break; + default: + LOG_ERROR("Unsupported quantizer type: "); + return core::IndexError_Unsupported; + } + } + } + + // Pass enable_rotate to converter_params (effective for INT8 and INT4) + if (param.enable_rotate) { + if (param.type == QuantizerType::kInt8 || + param.type == QuantizerType::kInt4) { + if (index_param.metric_type == MetricType::kCosine) { + converter_params.set("cosine.converter.enable_rotate", true); + } else { + converter_params.set("integer_streaming.converter.enable_rotate", true); + } + } else { + LOG_ERROR( + "enable_rotate is only supported for INT8/INT4 quantizer, " + "but got quantizer type: %d", + static_cast(param.type)); + return core::IndexError_Unsupported; + } + } + + proxima_index_meta_.set_converter(converter_name, 0, converter_params); + converter_ = core::IndexFactory::CreateConverter(converter_name); + if (converter_ == nullptr || + converter_->init(proxima_index_meta_, converter_params) != 0) { + LOG_ERROR("Failed to create and init converter"); + return core::IndexError_Runtime; + } + + proxima_index_meta_ = converter_->meta(); + + if (!proxima_index_meta_.reformer_name().empty()) { + reformer_ = + core::IndexFactory::CreateReformer(proxima_index_meta_.reformer_name()); + if (reformer_ == nullptr || + reformer_->init(proxima_index_meta_.reformer_params()) != 0) { + LOG_ERROR("Failed to create and init reformer"); + return core::IndexError_Runtime; + } + } + + streamer_vector_meta_.set_meta(proxima_index_meta_.data_type(), + proxima_index_meta_.dimension()); + streamer_vector_meta_.set_meta_type(proxima_index_meta_.meta_type()); + + return core::IndexError_Success; +} + +int Index::Init(const BaseIndexParam ¶m) { + param_ = param; // will lose the original type info + + is_sparse_ = param.is_sparse; + is_huge_page_ = param.is_huge_page; + + proxima_index_meta_.set_meta(param.data_type, param.dimension); + proxima_index_meta_.set_meta_type(is_sparse_ ? IndexMeta::MetaType::MT_SPARSE + : IndexMeta::MetaType::MT_DENSE); + + input_vector_meta_.set_meta(proxima_index_meta_.data_type(), + proxima_index_meta_.dimension()); + input_vector_meta_.set_meta_type(proxima_index_meta_.meta_type()); + streamer_vector_meta_ = input_vector_meta_; + + + // when quantizer=int8/int4, the converter.init() will change the metric to + // QuantizedInteger with params + if (ParseMetricName(param) != 0) { + LOG_ERROR("Failed to parse metric name"); + return core::IndexError_Runtime; + } + + if (CreateAndInitConverterReformer(param.quantizer_param, param) != 0) { + LOG_ERROR("Failed to create and init converter"); + return core::IndexError_Runtime; + } + + // must after quantizer handled. e.g., cosine doesn't support int8 quantizer + if (CreateAndInitMetric(param) != 0) { + LOG_ERROR("Failed to create and init metric"); + return core::IndexError_Runtime; + } + + if (CreateAndInitStreamer(param) != 0) { + LOG_ERROR("Failed to create and init streamer"); + return core::IndexError_Runtime; + } + return 0; +} + + +int Index::Open(const std::string &file_path, StorageOptions storage_options) { + ailego::Params storage_params; + // storage_params.set("proxima.mmap_file.storage.memory_warmup", true); + // storage_params.set("proxima.mmap_file.storage.segment_meta_capacity", + // 1024); + + storage_params.set(core::MMAPFILE_STORAGE_COPY_ON_WRITE, + storage_options.copy_on_write); + // force_flush must be enabled with copy_on_write so that + // IndexMapping::flush() actually persists data written via file_.write() to + // disk; without it, data would be lost. See IndexMapping::flush() in + // index_mapping.cc. + storage_params.set(core::MMAPFILE_STORAGE_FORCE_FLUSH, + storage_options.copy_on_write); + + switch (storage_options.type) { + case StorageOptions::StorageType::kMMAP: { + storage_ = core::IndexFactory::CreateStorage("MMapFileStorage"); + if (storage_ == nullptr) { + LOG_ERROR("Failed to create MMapFileStorage"); + return core::IndexError_Runtime; + } + int ret = storage_->init(storage_params); + if (ret != 0) { + LOG_ERROR("Failed to init MMapFileStorage, path: %s, err: %s", + file_path.c_str(), core::IndexError::What(ret)); + return ret; + } + break; + } + case StorageOptions::StorageType::kBufferPool: { + storage_ = core::IndexFactory::CreateStorage("BufferStorage"); + if (storage_ == nullptr) { + LOG_ERROR("Failed to create BufferStorage"); + return core::IndexError_Runtime; + } + int ret = storage_->init(storage_params); + if (ret != 0) { + LOG_ERROR("Failed to init BufferStorage, path: %s, err: %s", + file_path.c_str(), core::IndexError::What(ret)); + return ret; + } + break; + } + default: + LOG_ERROR("Unsupported storage type"); + return core::IndexError_Unsupported; + } + + // read_options.create_new + int ret = storage_->open(file_path, storage_options.create_new); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path.c_str()); + return core::IndexError_Runtime; + } + + // If a converter exists but reformer was not created during Init() + // (converters like UniformInt8 whose reformer params are only available + // after train()), create it now from the persisted meta that the streamer + // has loaded. When there is no converter (QuantizerType::kNone), reformer_ + // is nullptr by design — skip this block entirely. + if (converter_ != nullptr && reformer_ == nullptr) { + const auto &meta = streamer_->meta(); + if (meta.reformer_name().empty()) { + LOG_ERROR( + "Index::Open: converter exists but reformer not initialized and " + "no reformer in persisted meta"); + return core::IndexError_Runtime; + } + reformer_ = core::IndexFactory::CreateReformer(meta.reformer_name()); + if (!reformer_ || reformer_->init(meta.reformer_params()) != 0) { + LOG_ERROR("Failed to create reformer '%s' from persisted meta", + meta.reformer_name().c_str()); + return core::IndexError_Runtime; + } + } + + // converter/reformer/metric are created in IndexFactory::CreateIndex + // TODO: init + + // Load reformer data from storage (e.g., rotation matrix for + // IntegerStreaming) + if (reformer_ != nullptr) { + // When building a new index, dump converter state (e.g., rotator) to + // storage so the reformer can load it. This is needed for + // enable_rotate with INT8 quantization. + if (storage_options.create_new && converter_ != nullptr) { + if (converter_->dump_to_storage(storage_) != 0) { + LOG_ERROR("Failed to dump converter to storage, path: %s", + file_path.c_str()); + return core::IndexError_Runtime; + } + } + if (reformer_->load(storage_) != 0) { + LOG_ERROR("Failed to load reformer, path: %s", file_path.c_str()); + return core::IndexError_Runtime; + } + } + + // TODO: context pool + if (!init_context()) { // to validate if any error, will be overwritten + LOG_ERROR("Failed to init context"); + return core::IndexError_Runtime; + } + + is_open_ = true; + is_read_only_ = storage_options.read_only; + return 0; +} + +int Index::Close() { + if (!is_open_) { + LOG_ERROR("Index is not open"); + return core::IndexError_Runtime; + } + + if (!is_read_only_) { + if (ailego_unlikely(Flush() != 0)) { + LOG_ERROR("Failed to cleanup streamer"); + return core::IndexError_Runtime; + } + } + if (ailego_unlikely(streamer_->cleanup() != 0)) { + LOG_ERROR("Failed to cleanup streamer"); + return core::IndexError_Runtime; + } + if (ailego_unlikely(storage_->close() != 0)) { + LOG_ERROR("Failed to close storage"); + return core::IndexError_Runtime; + } + is_open_ = false; + return 0; +} + +int Index::Flush() { + if (!is_open_) { + LOG_ERROR("Index is not open"); + return core::IndexError_Runtime; + } + + if (is_read_only_) { + LOG_ERROR("Cannot flush read-only index"); + return core::IndexError_Runtime; + } + if (ailego_unlikely(streamer_->flush(0) != 0)) { + LOG_ERROR("Failed to flush streamer"); + return core::IndexError_Runtime; + } + if (ailego_unlikely(storage_->flush() != 0)) { + LOG_ERROR("Failed to flush storage"); + return core::IndexError_Runtime; + } + return 0; +} + +bool Index::IsDirty() const { + if (!storage_) { + return false; + } + return storage_->is_dirty(); +} + +int Index::Fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { + if (!is_open_) { + LOG_ERROR("Index is not open"); + return core::IndexError_Runtime; + } + if (is_sparse_) { + return _sparse_fetch(doc_id, vector_data_buffer); + } + return _dense_fetch(doc_id, vector_data_buffer); +} + +int Index::Add(const VectorData &vector_data, const uint32_t doc_id) { + if (!is_open_) { + LOG_ERROR("Index is not open"); + return core::IndexError_Runtime; + } + + if (is_read_only_) { + LOG_ERROR("Cannot add to read-only index"); + return core::IndexError_Runtime; + } + + auto &context = acquire_context(); + if (!context) { + LOG_ERROR("Failed to acquire context"); + return core::IndexError_Runtime; + } + + int ret = 0; + if (is_sparse_) { + ret = _sparse_add(vector_data, doc_id, context); + } else { + ret = _dense_add(vector_data, doc_id, context); + } + context->reset(); + return ret; +} + +int Index::AddWithSource(const VectorData & /*vector*/, uint32_t /*doc_id*/, + const core::VectorSource & /*src*/) { + LOG_ERROR("AddWithSource is not supported by this index type"); + return core::IndexError_Unsupported; +} + +int Index::SearchWithSource( + const VectorData & /*query*/, + const BaseIndexQueryParam::Pointer & /*search_param*/, + const core::VectorSource & /*src*/, SearchResult * /*result*/) { + LOG_ERROR("SearchWithSource is not supported by this index type"); + return core::IndexError_Unsupported; +} + +int Index::Search(const VectorData &vector_data, + const BaseIndexQueryParam::Pointer &search_param, + SearchResult *result) { + if (!is_open_) { + LOG_ERROR("Index is not open"); + return core::IndexError_Runtime; + } + + const bool has_group_by = has_group_by_search(search_param); + if (has_group_by && is_group_by_unsupported_index(param_.index_type)) { + LOG_ERROR("group_by search is not supported for this index type"); + return core::IndexError_Unsupported; + } + + if (search_param->refiner_param != nullptr && has_group_by) { + LOG_ERROR("group_by search is not supported with refiner"); + return core::IndexError_Unsupported; + } + + if (!is_trained_ && this->Train() != 0) { + LOG_ERROR("Failed to train index"); + return core::IndexError_Runtime; + } + + auto &context = acquire_context(); + if (!context) { + LOG_ERROR("Failed to acquire context"); + return core::IndexError_Runtime; + } + + if (_prepare_for_search(vector_data, search_param, context) != 0) { + LOG_ERROR("Failed to prepare for search"); + context->reset(); + return core::IndexError_Runtime; + } + + if (is_sparse_) { + int ret = _sparse_search(vector_data, search_param, result, context); + context->reset(); + return ret; + } + + // dense supports refiner, but sparse doesn't + int ret = 0; + if (search_param->refiner_param == nullptr) { + ret = _dense_search(vector_data, search_param, result, context); + context->reset(); + } else { + auto &reference_index = search_param->refiner_param->reference_index; + if (reference_index == nullptr) { + LOG_ERROR("Reference index is not set"); + context->reset(); + return core::IndexError_Runtime; + } + // TODO: tackle query_param's type info loss to loosen the constraint + if (reference_index->param_.index_type != IndexType::kFlat) { + LOG_ERROR("Reference index is not flat"); + context->reset(); + return core::IndexError_Runtime; + } + + context->set_topk(_get_coarse_search_topk(search_param)); + context->set_fetch_vector(false); // no need to fetch vector + if (_dense_search(vector_data, search_param, result, context) != 0) { + LOG_ERROR("Failed to search"); + context->reset(); + return core::IndexError_Runtime; + } + + auto &base_result = context->result(); + std::vector keys(base_result.size()); + for (size_t i = 0; i < base_result.size(); ++i) { + keys[i] = base_result[i].key(); + } + + FlatQueryParam::Pointer flat_search_param = + std::make_shared(); + flat_search_param->topk = search_param->topk; + flat_search_param->fetch_vector = search_param->fetch_vector; + flat_search_param->filter = search_param->filter; + // TODO: should copy other params? + flat_search_param->bf_pks = std::make_shared>(keys); + + ret = reference_index->Search(vector_data, flat_search_param, result); + context->reset(); + } + return ret; +} + + +int Index::_dense_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer) { + core::IndexStorage::MemoryBlock vector_block; + int ret = streamer_->get_vector_by_id(doc_id, vector_block); + if (ret != 0) { + LOG_ERROR("Failed to fetch vector, doc_id: %u", doc_id); + return core::IndexError_Runtime; + } + const void *vector = vector_block.data(); + + DenseVectorBuffer dense_vector_buffer; + std::string &out_vector_buffer = dense_vector_buffer.data; + // for int4, unit_size * dim != element_size + out_vector_buffer.resize(input_vector_meta_.element_size()); + + if (reformer_ != nullptr) { + if (reformer_->revert(vector, streamer_vector_meta_, &out_vector_buffer) != + 0) { + LOG_ERROR("Failed to convert vector"); + return core::IndexError_Runtime; + } + } else { + out_vector_buffer = std::string( + static_cast(vector), + input_vector_meta_.dimension() * input_vector_meta_.unit_size()); + } + vector_data_buffer->vector_buffer = std::move(dense_vector_buffer); + return 0; +} + + +int Index::_sparse_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer) { + SparseVectorBuffer sparse_vector_buffer; + + if (0 != streamer_->get_sparse_vector_by_id( + doc_id, &sparse_vector_buffer.count, + &sparse_vector_buffer.indices, &sparse_vector_buffer.values)) { + LOG_ERROR("Failed to fetch vector"); + return core::IndexError_Runtime; + } + + if (reformer_ != nullptr) { + std::string reverted_sparse_values_buffer; + if (reformer_->revert( + sparse_vector_buffer.count, sparse_vector_buffer.get_indices(), + sparse_vector_buffer.get_values(), streamer_vector_meta_, + &reverted_sparse_values_buffer) != 0) { + LOG_ERROR("Failed to convert vector"); + return core::IndexError_Runtime; + } + sparse_vector_buffer.values = std::move(reverted_sparse_values_buffer); + } + vector_data_buffer->vector_buffer = std::move(sparse_vector_buffer); + return 0; +} + +int Index::_dense_add(const VectorData &vector_data, const uint32_t doc_id, + core::IndexContext::Pointer &context) { + if (!std::holds_alternative(vector_data.vector)) { + LOG_ERROR("Invalid vector data"); + return core::IndexError_Runtime; + } + const DenseVector &dense_vector = std::get(vector_data.vector); + if (reformer_ != nullptr) { + core::IndexQueryMeta new_meta; + std::string new_vector; + int ret; + ret = reformer_->convert(dense_vector.data, input_vector_meta_, &new_vector, + &new_meta); + if (ret != 0) { + LOG_ERROR("Failed to convert vector"); + return core::IndexError_Runtime; + } + ret = streamer_->add_with_id_impl(doc_id, new_vector.data(), new_meta, + context); + if (ret != 0) { + LOG_ERROR("Failed to add vector"); + return core::IndexError_Runtime; + } + } else { + int ret = streamer_->add_with_id_impl(doc_id, dense_vector.data, + input_vector_meta_, context); + if (ret != 0) { + LOG_ERROR("Failed to add vector"); + return core::IndexError_Runtime; + } + } + return 0; +} + + +int Index::_sparse_add(const VectorData &vector_data, const uint32_t doc_id, + core::IndexContext::Pointer &context) { + if (!std::holds_alternative(vector_data.vector)) { + LOG_ERROR("Invalid vector data"); + return core::IndexError_Runtime; + } + const SparseVector &sparse_vector = + std::get(vector_data.vector); + + if (reformer_ != nullptr) { + std::string converted_sparse_values_buffer; + core::IndexQueryMeta new_meta; + int ret; + ret = reformer_->convert(sparse_vector.count, sparse_vector.get_indices(), + sparse_vector.get_values(), input_vector_meta_, + &converted_sparse_values_buffer, &new_meta); + if (ret != 0) { + LOG_ERROR("Failed to convert vector"); + return core::IndexError_Runtime; + } + ret = streamer_->add_with_id_impl( + doc_id, sparse_vector.count, sparse_vector.get_indices(), + converted_sparse_values_buffer.data(), new_meta, context); + if (ret != 0) { + LOG_ERROR("Failed to add vector"); + return core::IndexError_Runtime; + } + } else { + int ret = streamer_->add_with_id_impl( + doc_id, sparse_vector.count, sparse_vector.get_indices(), + sparse_vector.get_values(), input_vector_meta_, context); + if (ret != 0) { + LOG_ERROR("Failed to add vector"); + return core::IndexError_Runtime; + } + } + return 0; +} + + +int Index::_dense_search(const VectorData &vector_data, + const BaseIndexQueryParam::Pointer &search_param, + SearchResult *result, + core::IndexContext::Pointer &context) { + if (!std::holds_alternative(vector_data.vector)) { + LOG_ERROR("Invalid vector data"); + return core::IndexError_Runtime; + } + const DenseVector &dense_vector = std::get(vector_data.vector); + auto vector = dense_vector.data; + // Check if need to transform feature + std::string new_vector; + core::IndexQueryMeta new_meta = input_vector_meta_; + if (reformer_ != nullptr) { + if (reformer_->transform(dense_vector.data, input_vector_meta_, &new_vector, + &new_meta) != 0) { + LOG_ERROR("Failed to transform vector"); + return core::IndexError_Runtime; + } + vector = new_vector.data(); + } + if (search_param->bf_pks != nullptr) { + // should we eliminate the copy of bf_pks? + if (streamer_->search_bf_by_p_keys_impl( + vector, std::vector>{*search_param->bf_pks}, + new_meta, 1, context) != 0) { + LOG_ERROR("Failed to search_bf_by_p_keys_impl vector"); + return core::IndexError_Runtime; + } + } else if (search_param->is_linear) { + if (streamer_->search_bf_impl(vector, new_meta, 1, context) != 0) { + LOG_ERROR("Failed to search vector"); + return core::IndexError_Runtime; + } + } else { + if (streamer_->search_impl(vector, new_meta, 1, context) != 0) { + LOG_ERROR("Failed to search vector"); + return core::IndexError_Runtime; + } + } + + // Retrieve group_by results if applicable + bool has_group_by = + (search_param->group_by_param && search_param->group_by_param->group_by); + if (has_group_by) { + auto *group_result = context->mutable_group_result(); + if (group_result == nullptr) { + LOG_ERROR("Failed to retrieve group_by result"); + return core::IndexError_Runtime; + } + result->group_doc_list_ = std::move(*group_result); + } else { + result->doc_list_ = std::move(context->result()); + } + + if (metric_->support_normalize()) { + if (has_group_by) { + for (auto &group : result->group_doc_list_) { + for (auto &doc : *group.mutable_docs()) { + metric_->normalize(doc.mutable_score()); + } + } + } else { + for (auto &doc : result->doc_list_) { + metric_->normalize(doc.mutable_score()); + } + } + } + if (reformer_) { + if (has_group_by) { + for (auto &group : result->group_doc_list_) { + auto *docs = group.mutable_docs(); + if (reformer_->normalize(dense_vector.data, input_vector_meta_, + *docs) != 0) { + LOG_ERROR("Failed to normalize vector"); + return core::IndexError_Runtime; + } + } + } else { + if (reformer_->normalize(dense_vector.data, input_vector_meta_, + result->doc_list_) != 0) { + LOG_ERROR("Failed to normalize vector"); + return core::IndexError_Runtime; + } + } + if (context->fetch_vector() && reformer_->need_revert()) { + int revert_err = 0; + auto revert_one = [&](const void *vec, std::vector *out) { + if (revert_err) return; + std::string reverted_vector; + reverted_vector.resize(input_vector_meta_.dimension() * + input_vector_meta_.unit_size()); + if (reformer_->revert(vec, new_meta, &reverted_vector) != 0) { + LOG_ERROR("Failed to revert vector"); + revert_err = core::IndexError_Runtime; + return; + } + out->push_back(std::move(reverted_vector)); + }; + auto revert_docs = [&](auto &docs, std::vector &out) { + out.reserve(docs.size()); + for (auto &doc : docs) { + revert_one(doc.vector(), &out); + } + }; + if (has_group_by) { + result->group_reverted_vector_list_.reserve( + result->group_doc_list_.size()); + for (auto &group : result->group_doc_list_) { + std::vector group_vectors; + revert_docs(*group.mutable_docs(), group_vectors); + result->group_reverted_vector_list_.push_back( + std::move(group_vectors)); + } + } else { + revert_docs(result->doc_list_, result->reverted_vector_list_); + } + if (revert_err) return revert_err; + } + } + + return 0; +} + + +int Index::_sparse_search(const VectorData &vector_data, + const BaseIndexQueryParam::Pointer &search_param, + SearchResult *result, + core::IndexContext::Pointer &context) { + if (!std::holds_alternative(vector_data.vector)) { + LOG_ERROR("Invalid vector data"); + return core::IndexError_Runtime; + } + const SparseVector &sparse_vector = + std::get(vector_data.vector); + auto indices = sparse_vector.get_indices(); + auto values = sparse_vector.get_values(); + + std::string converted_sparse_values_buffer; + core::IndexQueryMeta new_meta = input_vector_meta_; + if (reformer_ != nullptr) { + if (reformer_->transform(sparse_vector.count, indices, values, + input_vector_meta_, + &converted_sparse_values_buffer, &new_meta) != 0) { + LOG_ERROR("Failed to transform vector"); + return core::IndexError_Runtime; + } + values = converted_sparse_values_buffer.data(); + } + + if (search_param->bf_pks != nullptr) { + if (streamer_->search_bf_by_p_keys_impl( + sparse_vector.count, indices, values, + std::vector>{*search_param->bf_pks}, new_meta, + context) != 0) { + LOG_ERROR("Failed to search_bf_by_p_keys_impl vector"); + return core::IndexError_Runtime; + } + } else if (search_param->is_linear) { + if (streamer_->search_bf_impl(sparse_vector.count, indices, values, + new_meta, context) != 0) { + LOG_ERROR("Failed to search vector"); + return core::IndexError_Runtime; + } + } else { + if (streamer_->search_impl(sparse_vector.count, indices, values, new_meta, + context) != 0) { + LOG_ERROR("Failed to search vector"); + return core::IndexError_Runtime; + } + } + // Retrieve group_by results if applicable + const bool has_group_by = has_group_by_search(search_param); + if (has_group_by) { + auto *group_result = context->mutable_group_result(); + if (group_result == nullptr) { + LOG_ERROR("Failed to retrieve group_by result"); + return core::IndexError_Runtime; + } + result->group_doc_list_ = std::move(*group_result); + } else { + result->doc_list_ = std::move(context->result()); + } + + if (metric_->support_normalize()) { + if (has_group_by) { + for (auto &group : result->group_doc_list_) { + for (auto &doc : *group.mutable_docs()) { + metric_->normalize(doc.mutable_score()); + } + } + } else { + for (auto &doc : result->doc_list_) { + metric_->normalize(doc.mutable_score()); + } + } + } + if (reformer_) { + // TODO: no need to call reformer_->normalize() when sparse? + if (context->fetch_vector() && reformer_->need_revert()) { + int revert_err = 0; + auto revert_one = [&](const core::IndexDocument &doc, + std::vector *out) { + if (revert_err) return; + auto &result_doc = doc.sparse_doc(); + std::string reverted_sparse_values; + reverted_sparse_values.resize(result_doc.sparse_count() * + input_vector_meta_.unit_size()); + if (reformer_->revert(result_doc.sparse_count(), + reinterpret_cast( + result_doc.sparse_indices().data()), + reinterpret_cast( + result_doc.sparse_values().data()), + new_meta, &reverted_sparse_values) != 0) { + LOG_ERROR("Failed to revert sparse vector"); + revert_err = core::IndexError_Runtime; + return; + } + out->push_back(std::move(reverted_sparse_values)); + }; + auto revert_docs = [&](auto &docs, std::vector &out) { + out.reserve(docs.size()); + for (auto &doc : docs) { + revert_one(doc, &out); + } + }; + if (has_group_by) { + result->group_reverted_sparse_values_list_.reserve( + result->group_doc_list_.size()); + for (auto &group : result->group_doc_list_) { + std::vector group_sparse_values; + revert_docs(*group.mutable_docs(), group_sparse_values); + result->group_reverted_sparse_values_list_.push_back( + std::move(group_sparse_values)); + } + } else { + revert_docs(result->doc_list_, result->reverted_sparse_values_list_); + } + if (revert_err) return revert_err; + } + } + return 0; +} + + +int Index::Merge(const std::vector &indexes, + const IndexFilter &filter, const MergeOptions &options) { + if (indexes.empty()) { + return core::IndexError_Success; + } + // ivf need builder + auto reducer = + core::IndexFactory::CreateStreamerReducer("MixedStreamerReducer"); + if (reducer == nullptr) { + LOG_ERROR("Failed to create reducer"); + return core::IndexError_Runtime; + } + + if (options.write_concurrency == 0) { + LOG_ERROR("Write concurrency must be greater than 0"); + return core::IndexError_InvalidArgument; + } + // must declare here to ensure its lifespan can cover reducer->reduce() + std::unique_ptr local_thread_pool = nullptr; + if (options.pool != nullptr) { + reducer->set_thread_pool(options.pool); + } else { + local_thread_pool = + std::make_unique(options.write_concurrency); + reducer->set_thread_pool(local_thread_pool.get()); + } + + ailego::Params reducer_params; + reducer_params.set(core::PARAM_MIXED_STREAMER_REDUCER_ENABLE_PK_REWRITE, + true); + reducer_params.set(core::PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, + options.write_concurrency); + if (reducer->init(reducer_params) != 0) { + LOG_ERROR("Failed to init reducer"); + return core::IndexError_Runtime; + } + if (reducer->set_target_streamer_wiht_info(builder_, streamer_, converter_, + reformer_, + input_vector_meta_) != 0) { + LOG_ERROR("Failed to set target streamer"); + return core::IndexError_Runtime; + } + + for (const auto &index : indexes) { + if (reducer->feed_streamer_with_reformer(index->streamer_, + index->reformer_) != 0) { + LOG_ERROR("Failed to feed streamer"); + return core::IndexError_Runtime; + } + } + if (reducer->reduce(filter) != 0) { + LOG_ERROR("Failed to reduce"); + return core::IndexError_Runtime; + } + is_trained_ = true; + return 0; +} + +int Index::_get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param) { + float scale_factor = search_param->refiner_param->scale_factor_; + if (scale_factor == 0) { + scale_factor = 1; + } + return floor(search_param->topk * scale_factor); +} + +void Index::_set_group_by_on_context( + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { + if (search_param->group_by_param && search_param->group_by_param->group_by) { + context->set_group_by(search_param->group_by_param->group_by); + context->set_group_params(search_param->group_by_param->group_count, + search_param->group_by_param->group_topk); + } +} + +std::string Index::get_metric_name(MetricType metric_type, bool is_sparse) { + if (is_sparse) { + switch (metric_type) { + case MetricType::kInnerProduct: + return "InnerProductSparse"; + case MetricType::kMIPSL2sq: + return "MipsSquaredEuclideanSparse"; + default: + return ""; + } + } else { + switch (metric_type) { + case MetricType::kL2sq: + return "SquaredEuclidean"; + case MetricType::kInnerProduct: + return "InnerProduct"; + case MetricType::kCosine: + return "Cosine"; + case MetricType::kMIPSL2sq: + return "MipsSquaredEuclidean"; + default: + return ""; + } + } +} + +} // namespace zvec::core_interface diff --git a/src/core/interface/index_factory.cc b/src/core/interface/index_factory.cc new file mode 100644 index 0000000..5d9dea1 --- /dev/null +++ b/src/core/interface/index_factory.cc @@ -0,0 +1,415 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "core/interface/utils/utils.h" + +namespace zvec::core_interface { + + +Index::Pointer IndexFactory::CreateAndInitIndex(const BaseIndexParam ¶m) { + Index::Pointer ptr = nullptr; + // if (param.index_type == IndexType::kIVF) { + // const IVFIndexParam *_param = dynamic_cast(¶m); ptr = std::make_shared(param); + + // if (_param->l1Index) { + // // TODO: create l1 index + // } + // if (_param->l2Index) { + // // TODO: create l2 index + // } + // } + // if (param.index_type == IndexType::kHNSW) { + // ptr = std::make_shared(param); + // } + if (param.index_type == IndexType::kFlat) { + // ptr = std::make_shared(param); + ptr = std::make_shared(); + } else if (param.index_type == IndexType::kHNSW) { + ptr = std::make_shared(); + } else if (param.index_type == IndexType::kIVF) { + ptr = std::make_shared(); + } else if (param.index_type == IndexType::kHNSWRabitq) { + ptr = std::make_shared(); +#if DISKANN_SUPPORTED + } else if (param.index_type == IndexType::kDiskAnn) { + ptr = std::make_shared(); +#endif + } else if (param.index_type == IndexType::kVamana) { + ptr = std::make_shared(); + } else { + LOG_ERROR("Unsupported index type: "); + return nullptr; + } + + if (!ptr) { + LOG_ERROR("Failed to create index"); + return nullptr; + } + if (0 != ptr->Init(param)) { + LOG_ERROR("Failed to init index"); + return nullptr; + } + return ptr; +} + +BaseIndexParam::Pointer IndexFactory::DeserializeIndexParamFromJson( + const std::string &json_str) { + ailego::JsonValue json_value; + if (!json_value.parse(json_str)) { + LOG_ERROR("Failed to parse json string: %s", json_str.c_str()); + return nullptr; + } + ailego::JsonObject json_obj = json_value.as_object(); + ailego::JsonValue tmp_json_value; + + IndexType index_type; + + if (!extract_enum_from_json(json_obj, "index_type", index_type, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize index type"); + return nullptr; + } + + switch (index_type) { + case IndexType::kFlat: { + FlatIndexParam::Pointer param = std::make_shared(); + if (!param->DeserializeFromJson(json_str)) { + LOG_ERROR("Failed to deserialize flat index param"); + return nullptr; + } + return param; + } + case IndexType::kHNSW: { + HNSWIndexParam::Pointer param = std::make_shared(); + if (!param->DeserializeFromJson(json_str)) { + LOG_ERROR("Failed to deserialize hnsw index param"); + return nullptr; + } + return param; + } + case IndexType::kIVF: { + IVFIndexParam::Pointer param = std::make_shared(); + if (!param->DeserializeFromJson(json_str)) { + LOG_ERROR("Failed to deserialize hnsw index param"); + return nullptr; + } + return param; + } + case IndexType::kHNSWRabitq: { + HNSWRabitqIndexParam::Pointer param = + std::make_shared(); + if (!param->DeserializeFromJson(json_str)) { + LOG_ERROR("Failed to deserialize hnsqrabitq index param"); + return nullptr; + } + return param; + } + case IndexType::kVamana: { + VamanaIndexParam::Pointer param = std::make_shared(); + if (!param->DeserializeFromJson(json_str)) { + LOG_ERROR("Failed to deserialize vamana index param"); + return nullptr; + } + return param; + } + default: + LOG_ERROR("Unsupported index type: %s", + magic_enum::enum_name(index_type).data()); + return nullptr; + } +} + +template , bool> > +std::string IndexFactory::QueryParamSerializeToJson(const QueryParamType ¶m, + bool omit_empty_value) { + ailego::JsonObject json_obj; + + // BaseIndexQueryParam + // omit filter & bf_pks + if (!omit_empty_value || param.topk != 0) { + json_obj.set("topk", ailego::JsonValue(param.topk)); + } + if (!omit_empty_value || param.fetch_vector) { + json_obj.set("fetch_vector", ailego::JsonValue(param.fetch_vector)); + } + if (!omit_empty_value || param.radius != 0.0f) { + json_obj.set("radius", ailego::JsonValue(param.radius)); + } + if (!omit_empty_value || param.is_linear) { + json_obj.set("is_linear", ailego::JsonValue(param.is_linear)); + } + + IndexType index_type{IndexType::kNone}; + if constexpr (std::is_same_v) { + // index_type + index_type = IndexType::kFlat; + } else if constexpr (std::is_same_v) { + if (!omit_empty_value || param.ef_search != 0) { + json_obj.set("ef_search", ailego::JsonValue(param.ef_search)); + } + if (!omit_empty_value || param.prefetch_offset != 0) { + json_obj.set("prefetch_offset", ailego::JsonValue(param.prefetch_offset)); + } + if (!omit_empty_value || param.prefetch_lines != 0) { + json_obj.set("prefetch_lines", ailego::JsonValue(param.prefetch_lines)); + } + index_type = IndexType::kHNSW; + } else if constexpr (std::is_same_v) { + if (!omit_empty_value || param.nprobe != 0) { + json_obj.set("nprobe", ailego::JsonValue(param.nprobe)); + } + index_type = IndexType::kIVF; + // json_obj.set("l1QueryParam", + // ailego::JsonValue(QueryParamSerializeToJson(param.l1QueryParam))); + // json_obj.set("l2QueryParam", + // ailego::JsonValue(QueryParamSerializeToJson(param.l2QueryParam))); + } else if constexpr (std::is_same_v) { + if (!omit_empty_value || param.ef_search != 0) { + json_obj.set("ef_search", ailego::JsonValue(param.ef_search)); + } + index_type = IndexType::kHNSWRabitq; + } else if constexpr (std::is_same_v) { + if (!omit_empty_value || param.ef_search != 0) { + json_obj.set("ef_search", ailego::JsonValue(param.ef_search)); + } + if (!omit_empty_value || param.prefetch_offset != 0) { + json_obj.set("prefetch_offset", ailego::JsonValue(param.prefetch_offset)); + } + if (!omit_empty_value || param.prefetch_lines != 0) { + json_obj.set("prefetch_lines", ailego::JsonValue(param.prefetch_lines)); + } + index_type = IndexType::kVamana; + } + + json_obj.set("index_type", + ailego::JsonValue(magic_enum::enum_name(index_type).data())); + + return ailego::JsonValue(json_obj).as_json_string().as_stl_string(); +} + +template std::string +IndexFactory::QueryParamSerializeToJson( + const BaseIndexQueryParam ¶m, bool omit_empty_value); +template std::string IndexFactory::QueryParamSerializeToJson( + const FlatQueryParam ¶m, bool omit_empty_value); +template std::string IndexFactory::QueryParamSerializeToJson( + const HNSWQueryParam ¶m, bool omit_empty_value); +template std::string IndexFactory::QueryParamSerializeToJson( + const IVFQueryParam ¶m, bool omit_empty_value); + +template , bool> > +typename QueryParamType::Pointer IndexFactory::QueryParamDeserializeFromJson( + const std::string &json_str) { + ailego::JsonValue tmp_json_value; + if (!tmp_json_value.parse(json_str)) { + LOG_ERROR("Failed to parse json string: %s", json_str.c_str()); + return nullptr; + } + ailego::JsonObject json_obj = tmp_json_value.as_object(); + + auto parse_common_fields = [&](auto ¶m) -> bool { + if (!extract_value_from_json(json_obj, "topk", param->topk, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize topk"); + return false; + } + + if (!extract_value_from_json(json_obj, "fetch_vector", param->fetch_vector, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize fetch_vector"); + return false; + } + + if (!extract_value_from_json(json_obj, "radius", param->radius, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize radius"); + return false; + } + + if (!extract_value_from_json(json_obj, "is_linear", param->is_linear, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize is_linear"); + return false; + } + return true; + }; + + IndexType index_type; + + if (!extract_enum_from_json(json_obj, "index_type", index_type, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize index type"); + return nullptr; + } + + if constexpr (std::is_same_v) { + if (index_type == IndexType::kFlat) { + auto param = std::make_shared(); + if (!parse_common_fields(param)) { + return nullptr; + } + return param; + } else if (index_type == IndexType::kHNSW) { + auto param = std::make_shared(); + if (!parse_common_fields(param)) { + return nullptr; + } + if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize ef_search"); + return nullptr; + } + if (!extract_value_from_json(json_obj, "prefetch_offset", + param->prefetch_offset, tmp_json_value)) { + LOG_ERROR("Failed to deserialize prefetch_offset"); + return nullptr; + } + if (!extract_value_from_json(json_obj, "prefetch_lines", + param->prefetch_lines, tmp_json_value)) { + LOG_ERROR("Failed to deserialize prefetch_lines"); + return nullptr; + } + return param; + } else if (index_type == IndexType::kIVF) { + auto param = std::make_shared(); + if (!parse_common_fields(param)) { + return nullptr; + } + if (!extract_value_from_json(json_obj, "nprobe", param->nprobe, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize nprobe"); + return nullptr; + } + return param; + } else if (index_type == IndexType::kHNSWRabitq) { + auto param = std::make_shared(); + if (!parse_common_fields(param)) { + return nullptr; + } + if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize ef_search"); + return nullptr; + } + return param; + } else if (index_type == IndexType::kVamana) { + auto param = std::make_shared(); + if (!parse_common_fields(param)) { + return nullptr; + } + if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize ef_search"); + return nullptr; + } + if (!extract_value_from_json(json_obj, "prefetch_offset", + param->prefetch_offset, tmp_json_value)) { + LOG_ERROR("Failed to deserialize prefetch_offset"); + return nullptr; + } + if (!extract_value_from_json(json_obj, "prefetch_lines", + param->prefetch_lines, tmp_json_value)) { + LOG_ERROR("Failed to deserialize prefetch_lines"); + return nullptr; + } + return param; + } else { + LOG_ERROR("Unsupported index type: %s", + magic_enum::enum_name(index_type).data()); + return nullptr; + } + } else { + auto param = std::make_shared(); + if (!parse_common_fields(param)) { + return nullptr; + } + if constexpr (std::is_same_v) { + } else if constexpr (std::is_same_v) { + if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize ef_search"); + return nullptr; + } + if (!extract_value_from_json(json_obj, "prefetch_offset", + param->prefetch_offset, tmp_json_value)) { + LOG_ERROR("Failed to deserialize prefetch_offset"); + return nullptr; + } + if (!extract_value_from_json(json_obj, "prefetch_lines", + param->prefetch_lines, tmp_json_value)) { + LOG_ERROR("Failed to deserialize prefetch_lines"); + return nullptr; + } + } else if constexpr (std::is_same_v) { + if (!extract_value_from_json(json_obj, "nprobe", param->nprobe, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize nprobe"); + return nullptr; + } + } else if constexpr (std::is_same_v) { + if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize ef_search"); + return nullptr; + } + } else if constexpr (std::is_same_v) { + if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize ef_search"); + return nullptr; + } + if (!extract_value_from_json(json_obj, "prefetch_offset", + param->prefetch_offset, tmp_json_value)) { + LOG_ERROR("Failed to deserialize prefetch_offset"); + return nullptr; + } + if (!extract_value_from_json(json_obj, "prefetch_lines", + param->prefetch_lines, tmp_json_value)) { + LOG_ERROR("Failed to deserialize prefetch_lines"); + return nullptr; + } + } else { + LOG_ERROR("Unsupported index type: %s", + magic_enum::enum_name(index_type).data()); + return nullptr; + } + return param; + } +} + +template BaseIndexQueryParam::Pointer +IndexFactory::QueryParamDeserializeFromJson( + const std::string &json_str); +template FlatQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< + FlatQueryParam>(const std::string &json_str); +template HNSWQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< + HNSWQueryParam>(const std::string &json_str); +template IVFQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< + IVFQueryParam>(const std::string &json_str); +template std::string IndexFactory::QueryParamSerializeToJson( + const VamanaQueryParam ¶m, bool omit_empty_value); +template VamanaQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< + VamanaQueryParam>(const std::string &json_str); + +} // namespace zvec::core_interface diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc new file mode 100644 index 0000000..5d75276 --- /dev/null +++ b/src/core/interface/index_param.cc @@ -0,0 +1,343 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include "core/interface/utils/utils.h" + +namespace zvec { +namespace core_interface { +ailego::JsonObject BaseIndexParam::SerializeToJsonObject( + bool omit_empty_value) const { + ailego::JsonObject json_obj; + + if (!omit_empty_value || index_type != IndexType::kNone) { + json_obj.set("index_type", + ailego::JsonValue(magic_enum::enum_name(index_type).data())); + } + if (!omit_empty_value || metric_type != MetricType::kNone) { + json_obj.set("metric_type", + ailego::JsonValue(magic_enum::enum_name(metric_type).data())); + } + if (!omit_empty_value || dimension != 0) { + json_obj.set("dimension", ailego::JsonValue(dimension)); + } + if (!omit_empty_value || version != 0) { + json_obj.set("version", ailego::JsonValue(version)); + } + if (!omit_empty_value || is_sparse) { + json_obj.set("is_sparse", ailego::JsonValue(is_sparse)); + } + if (!omit_empty_value || data_type != DataType::DT_UNDEFINED) { + json_obj.set("data_type", + ailego::JsonValue(magic_enum::enum_name(data_type).data())); + } + if (!omit_empty_value || use_id_map) { + json_obj.set("use_id_map", ailego::JsonValue(use_id_map)); + } + if (!omit_empty_value || is_huge_page) { + json_obj.set("is_huge_page", ailego::JsonValue(is_huge_page)); + } + if (!omit_empty_value || use_external_vector) { + json_obj.set("use_external_vector", ailego::JsonValue(use_external_vector)); + } + + // if (preprocess_param) { + // json.set("preprocess_param", preprocess_param->SerializeToJson()); + // } + if (!omit_empty_value || quantizer_param.type != QuantizerType::kNone) { + json_obj.set("quantizer_param", + quantizer_param.SerializeToJsonObject(omit_empty_value)); + } + // if (refiner_param) { + // json.set("refiner_param", refiner_param->SerializeToJson()); + // } + // if (default_query_param) { + // json.set("default_query_param", + // default_query_param->SerializeToJson()); + // } + return json_obj; +} + + +ailego::JsonObject FlatIndexParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); + if (!omit_empty_value || major_order != IndexMeta::MajorOrder::MO_UNDEFINED) { + json_obj.set("major_order", + ailego::JsonValue(magic_enum::enum_name(major_order).data())); + } + return json_obj; +} + +ailego::JsonObject HNSWIndexParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); + json_obj.set("m", ailego::JsonValue(m)); + json_obj.set("ef_construction", ailego::JsonValue(ef_construction)); + if (!omit_empty_value || use_contiguous_memory) { + json_obj.set("use_contiguous_memory", + ailego::JsonValue(use_contiguous_memory)); + } + return json_obj; +} + +bool BaseIndexParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + DESERIALIZE_ENUM_FIELD(json_obj, index_type, IndexType); + DESERIALIZE_ENUM_FIELD(json_obj, metric_type, MetricType); + DESERIALIZE_ENUM_FIELD(json_obj, data_type, DataType); + + DESERIALIZE_VALUE_FIELD(json_obj, dimension); + DESERIALIZE_VALUE_FIELD(json_obj, version); + DESERIALIZE_VALUE_FIELD(json_obj, is_sparse); + DESERIALIZE_VALUE_FIELD(json_obj, use_id_map); + DESERIALIZE_VALUE_FIELD(json_obj, is_huge_page); + DESERIALIZE_VALUE_FIELD(json_obj, use_external_vector); + + ailego::JsonValue tmp_json_value; + if (json_obj.has("quantizer_param")) { + if (json_obj.get("quantizer_param", &tmp_json_value); + tmp_json_value.is_object()) { + quantizer_param.DeserializeFromJsonObject(tmp_json_value.as_object()); + } + } + + return true; +} + +bool FlatIndexParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + + if (index_type != IndexType::kFlat) { + LOG_ERROR("index_type is not kFlat"); + return false; + } + + DESERIALIZE_ENUM_FIELD(json_obj, major_order, IndexMeta::MajorOrder); + return true; +} + +bool HNSWIndexParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + + if (index_type != IndexType::kHNSW) { + LOG_ERROR("index_type is not kHNSW"); + return false; + } + + DESERIALIZE_VALUE_FIELD(json_obj, m); + DESERIALIZE_VALUE_FIELD(json_obj, ef_construction); + DESERIALIZE_VALUE_FIELD(json_obj, use_contiguous_memory); + + return true; +} + +bool HNSWRabitqIndexParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + + if (index_type != IndexType::kHNSWRabitq) { + LOG_ERROR("index_type is not kHNSWRabitq"); + return false; + } + + DESERIALIZE_VALUE_FIELD(json_obj, m); + DESERIALIZE_VALUE_FIELD(json_obj, ef_construction); + DESERIALIZE_VALUE_FIELD(json_obj, total_bits); + DESERIALIZE_VALUE_FIELD(json_obj, num_clusters); + DESERIALIZE_VALUE_FIELD(json_obj, sample_count); + + return true; +} + +ailego::JsonObject HNSWRabitqIndexParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); + json_obj.set("m", ailego::JsonValue(m)); + json_obj.set("ef_construction", ailego::JsonValue(ef_construction)); + json_obj.set("total_bits", ailego::JsonValue(total_bits)); + json_obj.set("num_clusters", ailego::JsonValue(num_clusters)); + if (!omit_empty_value || sample_count != 0) { + json_obj.set("sample_count", ailego::JsonValue(sample_count)); + } + return json_obj; +} + +ailego::JsonObject VamanaIndexParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); + json_obj.set("max_degree", ailego::JsonValue(max_degree)); + json_obj.set("search_list_size", ailego::JsonValue(search_list_size)); + json_obj.set("alpha", ailego::JsonValue(alpha)); + if (!omit_empty_value || + max_occlusion_size != static_cast(kDefaultVamanaMaxOcclusionSize)) { + json_obj.set("max_occlusion_size", ailego::JsonValue(max_occlusion_size)); + } + if (!omit_empty_value || saturate_graph) { + json_obj.set("saturate_graph", ailego::JsonValue(saturate_graph)); + } + if (!omit_empty_value || use_contiguous_memory) { + json_obj.set("use_contiguous_memory", + ailego::JsonValue(use_contiguous_memory)); + } + return json_obj; +} + +bool DiskAnnIndexParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + + if (index_type != IndexType::kDiskAnn) { + LOG_ERROR("index_type is not DiskAnn"); + return false; + } + + return true; +} + +ailego::JsonObject DiskAnnIndexParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); + return json_obj; +} + +bool VamanaIndexParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + + if (index_type != IndexType::kVamana) { + LOG_ERROR("index_type is not kVamana"); + return false; + } + + DESERIALIZE_VALUE_FIELD(json_obj, max_degree); + DESERIALIZE_VALUE_FIELD(json_obj, search_list_size); + DESERIALIZE_VALUE_FIELD(json_obj, alpha); + DESERIALIZE_VALUE_FIELD(json_obj, max_occlusion_size); + DESERIALIZE_VALUE_FIELD(json_obj, saturate_graph); + DESERIALIZE_VALUE_FIELD(json_obj, use_contiguous_memory); + + return true; +} + +ailego::JsonObject QuantizerParam::SerializeToJsonObject( + bool omit_empty_value) const { + ailego::JsonObject json_obj; + if (!omit_empty_value || type != QuantizerType::kNone) { + json_obj.set("type", + zvec::ailego::JsonValue(magic_enum::enum_name(type).data())); + } + if (!omit_empty_value || enable_rotate) { + json_obj.set("enable_rotate", ailego::JsonValue(enable_rotate)); + } + return json_obj; +} + +bool QuantizerParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + DESERIALIZE_ENUM_FIELD(json_obj, type, QuantizerType); + DESERIALIZE_VALUE_FIELD(json_obj, enable_rotate); + return true; +} + +// bool BaseIndexQueryParam::DeserializeFromJsonObject( +// const ailego::JsonObject &json_obj) { +// DESERIALIZE_ENUM_FIELD(json_obj, index_type, IndexType); +// DESERIALIZE_VALUE_FIELD(json_obj, topk); +// DESERIALIZE_VALUE_FIELD(json_obj, fetch_vector); +// DESERIALIZE_VALUE_FIELD(json_obj, radius); +// DESERIALIZE_VALUE_FIELD(json_obj, is_linear); +// return true; +// } + +// ailego::JsonObject BaseIndexQueryParam::SerializeToJsonObject( +// bool omit_empty_value) const { +// ailego::JsonObject json_obj; +// if (!omit_empty_value || index_type != IndexType::kNone) { +// json_obj.set("index_type", +// ailego::JsonValue(magic_enum::enum_name(index_type).data())); +// } +// if (!omit_empty_value || topk != 0) { +// json_obj.set("topk", ailego::JsonValue(topk)); +// } +// if (!omit_empty_value || fetch_vector) { +// json_obj.set("fetch_vector", ailego::JsonValue(fetch_vector)); +// } +// if (!omit_empty_value || radius != 0.0f) { +// json_obj.set("radius", ailego::JsonValue(radius)); +// } +// if (!omit_empty_value || is_linear) { +// json_obj.set("is_linear", ailego::JsonValue(is_linear)); +// } +// return json_obj; +// } + +// bool FlatQueryParam::DeserializeFromJsonObject( +// const ailego::JsonObject &json_obj) { +// if (!BaseIndexQueryParam::DeserializeFromJsonObject(json_obj)) { +// return false; +// } +// if (index_type != IndexType::kFlat) { +// LOG_ERROR("index_type is not kFlat"); +// return false; +// } +// return true; +// } + +// ailego::JsonObject FlatQueryParam::SerializeToJsonObject( +// bool omit_empty_value) const { +// auto json_obj = +// BaseIndexQueryParam::SerializeToJsonObject(omit_empty_value); +// return json_obj; +// } + +// bool HNSWQueryParam::DeserializeFromJsonObject( +// const ailego::JsonObject &json_obj) { +// if (!BaseIndexQueryParam::DeserializeFromJsonObject(json_obj)) { +// return false; +// } +// if (index_type != IndexType::kHNSW) { +// LOG_ERROR("index_type is not kHNSW"); +// return false; +// } +// DESERIALIZE_VALUE_FIELD(json_obj, ef_search); +// return true; +// } + +// ailego::JsonObject HNSWQueryParam::SerializeToJsonObject( +// bool omit_empty_value) const { +// auto json_obj = +// BaseIndexQueryParam::SerializeToJsonObject(omit_empty_value); +// if (!omit_empty_value || ef_search != 0) { +// json_obj.set("ef_search", ailego::JsonValue(ef_search)); +// } +// return json_obj; +// } + + +} // namespace core_interface +} // namespace zvec \ No newline at end of file diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc new file mode 100644 index 0000000..bf1233a --- /dev/null +++ b/src/core/interface/indexes/diskann_index.cc @@ -0,0 +1,323 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "algorithm/diskann/diskann_params.h" +#include "holder_builder.h" + +namespace zvec::core_interface { + +namespace { + +// Implicitly bring the DiskAnn runtime online on first use. This keeps the +// DiskAnn index an ordinary public API (users just instantiate a +// DiskAnnIndexParam) while still letting the rest of the library — HNSW, +// IVF, Flat, Vamana — run on hosts that happen to lack libaio. On such +// hosts only DiskAnn fails, with a clear, actionable error message, and +// every other index type stays fully functional. +int EnsureDiskAnnRuntimeReady() { + static std::once_flag once; + static int cached_result = 0; + std::call_once(once, []() { + const int status = ::zvec::LoadDiskAnnPlugin(); + if (status == kDiskAnnPluginOk) { + cached_result = 0; + return; + } + switch (status) { + case kDiskAnnPluginLibAioMissing: + LOG_ERROR( + "DiskAnn requires libaio at runtime, but it was not found on this " + "host. Install it (e.g. 'apt-get install libaio1' on " + "Debian/Ubuntu, " + "or 'libaio1t64' on Ubuntu 24.04+) and retry."); + break; + case kDiskAnnPluginUnsupportedPlatform: + LOG_ERROR("DiskAnn is only supported on Linux x86_64."); + break; + case kDiskAnnPluginDlopenFailed: + default: + LOG_ERROR("Failed to initialize the DiskAnn runtime (status=%d).", + status); + break; + } + cached_result = core::IndexError_Runtime; + }); + return cached_result; +} + +} // namespace + +int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { + // Fail fast and cleanly if the DiskAnn runtime cannot be brought up on + // this host (most commonly: libaio is missing). The rest of zvec keeps + // running; only DiskAnn is unusable. + if (int rc = EnsureDiskAnnRuntimeReady(); rc != 0) { + return rc; + } + + if (is_sparse_) { + LOG_ERROR("Failed to create streamer. Sparse is not Supported."); + return core::IndexError_Unsupported; + } + + param_ = dynamic_cast(param); + param_.max_degree = std::min(100, param_.max_degree); + param_.list_size = std::min(100, param_.list_size); + param_.pq_chunk_num = std::min(1024, param_.pq_chunk_num); + + proxima_index_params_.set(core::PARAM_DISKANN_BUILDER_MAX_DEGREE, + param_.max_degree); + proxima_index_params_.set(core::PARAM_DISKANN_BUILDER_LIST_SIZE, + param_.list_size); + proxima_index_params_.set(core::PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM, + param_.pq_chunk_num); + + builder_ = core::IndexFactory::CreateBuilder("DiskAnnBuilder"); + streamer_ = core::IndexFactory::CreateStreamer("DiskAnnStreamer"); + + if (ailego_unlikely(!builder_ || !streamer_)) { + LOG_ERROR( + "Failed to create DiskAnnBuilder/DiskAnnStreamer: DiskAnn factory " + "entries are not registered. This usually means the DiskAnn shared " + "module could not be located next to the hosting binary."); + return core::IndexError_Runtime; + } + + IndexMeta real_meta; + if (converter_) { + real_meta = converter_->meta(); + } else { + real_meta = proxima_index_meta_; + } + + const int builder_ret = builder_->init(real_meta, proxima_index_params_); + const int streamer_ret = streamer_->init(real_meta, proxima_index_params_); + if (ailego_unlikely(builder_ret != 0 || streamer_ret != 0)) { + LOG_ERROR( + "Failed to init builder or streamer, builder_ret: %d, " + "streamer_ret: %d", + builder_ret, streamer_ret); + return core::IndexError_Runtime; + } + + return 0; +} + +int DiskAnnIndex::Open(const std::string &file_path, + StorageOptions storage_options) { + ailego::Params storage_params; + file_path_ = file_path; + is_read_only_ = storage_options.read_only; + switch (storage_options.type) { + case StorageOptions::StorageType::kMMAP: + case StorageOptions::StorageType::kBufferPool: { + // NOTE: DiskAnn index is dumped via FileDumper (plain binary file), which + // is not compatible with BufferStorage's IndexFormat layout. Fall back to + // FileReadStorage for both MMAP and BufferPool storage types. + storage_ = core::IndexFactory::CreateStorage("FileReadStorage"); + if (storage_ == nullptr) { + LOG_ERROR("Failed to create FileReadStorage"); + return core::IndexError_Runtime; + } + int ret = storage_->init(storage_params); + if (ret != 0) { + LOG_ERROR("Failed to init FileReadStorage, path: %s, err: %s", + file_path_.c_str(), core::IndexError::What(ret)); + return ret; + } + break; + } + default: { + LOG_ERROR("Unsupported storage type"); + return core::IndexError_Unsupported; + } + } + + if (!storage_options.create_new) { + int ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + is_trained_ = true; + } + is_open_ = true; + return 0; +} + +int DiskAnnIndex::GenerateHolder() { + return BuildMultiPassHolder(param_.data_type, param_.dimension, doc_cache_, + converter_, &holder_); +} + +int DiskAnnIndex::Add(const VectorData &vector, uint32_t doc_id) { + if (is_trained_) { + LOG_ERROR("this diskann index is trained"); + return core::IndexError_Runtime; + } + if (!std::holds_alternative(vector.vector)) { + LOG_ERROR("Invalid vector data"); + return core::IndexError_Runtime; + } + const DenseVector &dense_vector = std::get(vector.vector); + std::string out_vector_buffer = std::string( + static_cast(dense_vector.data), + input_vector_meta_.dimension() * input_vector_meta_.unit_size()); + + std::lock_guard lock(mutex_); + if (doc_cache_.size() <= doc_id) { + std::string fake_data( + input_vector_meta_.dimension() * input_vector_meta_.unit_size(), 0); + doc_cache_.resize(doc_id + 1, std::make_pair(kInvalidKey, fake_data)); + } + doc_cache_[doc_id] = std::make_pair(doc_id, out_vector_buffer); + return 0; +} + +int DiskAnnIndex::Train() { + int ret = GenerateHolder(); + if (ret != 0) { + LOG_ERROR("Failed to generate holder, err: %s", + core::IndexError::What(ret)); + return ret; + } + ret = builder_->train(holder_); + if (ret != 0) { + LOG_ERROR("Failed to train builder, err: %s", core::IndexError::What(ret)); + return ret; + } + ret = builder_->build(holder_); + if (ret != 0) { + LOG_ERROR("Failed to build index, err: %s", core::IndexError::What(ret)); + return ret; + } + auto dumper = core::IndexFactory::CreateDumper("FileDumper"); + if (dumper == nullptr) { + LOG_ERROR("Failed to create FileDumper"); + return core::IndexError_Runtime; + } + + ret = dumper->create(file_path_); + if (ret != 0) { + LOG_ERROR("Failed to create dumper, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + ret = builder_->dump(dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump index, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + dumper->close(); + ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + is_trained_ = true; + return 0; +} + +int DiskAnnIndex::_dense_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer) { + if (is_trained_) { + return Index::_dense_fetch(doc_id, vector_data_buffer); + } else { + DenseVectorBuffer dense_vector_buffer; + std::string &out_vector_buffer = dense_vector_buffer.data; + out_vector_buffer = doc_cache_[doc_id].second; + vector_data_buffer->vector_buffer = std::move(dense_vector_buffer); + return 0; + } +} + +int DiskAnnIndex::_prepare_for_search( + const VectorData & /*query*/, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { + const auto &diskann_search_param = + std::dynamic_pointer_cast(search_param); + if (diskann_search_param == nullptr) { + LOG_ERROR("Invalid search param: expected DiskAnnQueryParam"); + return core::IndexError_Runtime; + } + + if (search_param->group_by_param && search_param->group_by_param->group_by) { + LOG_ERROR("group_by search is not supported for DiskAnn index"); + return core::IndexError_Unsupported; + } + + context->set_topk(diskann_search_param->topk); + + // Propagate the query-time beam-search list size into the context. Must be + // at least topk to keep enough candidates for a correct result. + ailego::Params params; + params.set( + core::PARAM_DISKANN_SEARCHER_LIST_SIZE, + std::max(diskann_search_param->topk, diskann_search_param->list_size)); + context->update(params); + + return 0; +} + +int DiskAnnIndex::Merge(const std::vector &indexes, + const IndexFilter &filter, + const MergeOptions &options) { + int pre_ret = Index::Merge(indexes, filter, options); + if (pre_ret != 0) { + return pre_ret; + } + auto dumper = core::IndexFactory::CreateDumper("FileDumper"); + + dumper->create(file_path_); + int ret = builder_->dump(dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump index, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + + dumper->close(); + + ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + is_trained_ = true; + return 0; +} + +} // namespace zvec::core_interface \ No newline at end of file diff --git a/src/core/interface/indexes/flat_index.cc b/src/core/interface/indexes/flat_index.cc new file mode 100644 index 0000000..032f366 --- /dev/null +++ b/src/core/interface/indexes/flat_index.cc @@ -0,0 +1,72 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "algorithm/flat/flat_utility.h" + +namespace zvec::core_interface { + +int FlatIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { + param_ = dynamic_cast(param); + + proxima_index_params_.set(core::PARAM_FLAT_COLUMN_MAJOR_ORDER, + param_.major_order == IndexMeta::MO_COLUMN); + proxima_index_params_.set(core::PARAM_FLAT_USE_ID_MAP, param_.use_id_map); + if (is_sparse_) { + streamer_ = core::IndexFactory::CreateStreamer("FlatSparseStreamer"); + } else { + streamer_ = core::IndexFactory::CreateStreamer("FlatStreamer"); + } + + if (ailego_unlikely(!streamer_)) { + LOG_ERROR("Failed to create streamer"); + return core::IndexError_Runtime; + } + if (ailego_unlikely( + streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init streamer"); + return core::IndexError_Runtime; + } + return 0; +} + +int FlatIndex::_prepare_for_search( + const VectorData & /*vector_data*/, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { + auto flat_search_param = + std::dynamic_pointer_cast(search_param); + + if (ailego_unlikely(!flat_search_param)) { + LOG_ERROR("Invalid search param type, expected FlatQueryParam"); + return core::IndexError_Runtime; + } + + context->set_topk(flat_search_param->topk); + context->set_fetch_vector(flat_search_param->fetch_vector); + if (flat_search_param->filter) { + context->set_filter(std::move(*flat_search_param->filter)); + } + if (flat_search_param->radius > 0.0f) { + context->set_threshold(flat_search_param->radius); + } + _set_group_by_on_context(search_param, context); + + return 0; +} + + +} // namespace zvec::core_interface \ No newline at end of file diff --git a/src/core/interface/indexes/hnsw_index.cc b/src/core/interface/indexes/hnsw_index.cc new file mode 100644 index 0000000..9756289 --- /dev/null +++ b/src/core/interface/indexes/hnsw_index.cc @@ -0,0 +1,185 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "algorithm/hnsw/hnsw_context.h" +#include "algorithm/hnsw/hnsw_params.h" +#include "algorithm/hnsw/hnsw_streamer.h" +#include "algorithm/hnsw/hnsw_streamer_entity.h" +#include "algorithm/hnsw_sparse/hnsw_sparse_params.h" + +namespace zvec::core_interface { + +std::string HNSWIndex::storage_mode() const { + if (!streamer_) { + return ""; + } + auto *hnsw_streamer = dynamic_cast(streamer_.get()); + if (!hnsw_streamer) { + // e.g. sparse branch uses HnswSparseStreamer which is a different type + return ""; + } + switch (hnsw_streamer->storage_mode()) { + case core::HnswStorageMode::kMmap: + return "mmap"; + case core::HnswStorageMode::kBufferPool: + return "buffer_pool"; + case core::HnswStorageMode::kContiguous: + return "contiguous"; + case core::HnswStorageMode::kExternal: + return "external"; + } + return ""; +} + +int HNSWIndex::AddWithSource(const VectorData &vector_data, + const uint32_t doc_id, + const core::VectorSource &src) { + auto &context = acquire_context(); + if (!context) { + LOG_ERROR("Failed to acquire context for AddWithSource"); + return core::IndexError_Runtime; + } + if (auto *ctx = dynamic_cast(context.get())) { + ctx->set_vector_source(&src); + } + return Index::Add(vector_data, doc_id); +} + +int HNSWIndex::SearchWithSource( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + const core::VectorSource &src, SearchResult *result) { + auto &context = acquire_context(); + if (!context) { + LOG_ERROR("Failed to acquire context for SearchWithSource"); + return core::IndexError_Runtime; + } + if (auto *ctx = dynamic_cast(context.get())) { + ctx->set_vector_source(&src); + } + return Index::Search(query, search_param, result); +} + +int HNSWIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { + param_ = dynamic_cast(param); + + // valid + param_.ef_construction = std::max(1, std::min(2048, param_.ef_construction)); + param_.m = std::max(5, std::min(1024, param_.m)); + + if (is_sparse_) { + proxima_index_params_.set(core::PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION, + param_.ef_construction); + proxima_index_params_.set( + core::PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT, param_.m); + + // TODO: add_vector_with_id & fetch_by_id don't rely on this param + proxima_index_params_.set( + core::PARAM_HNSW_SPARSE_STREAMER_GET_VECTOR_ENABLE, true); + + // TODO: use index params' default query param here + proxima_index_params_.set(core::PARAM_HNSW_SPARSE_STREAMER_EF, + kDefaultHnswEfSearch); + streamer_ = core::IndexFactory::CreateStreamer("HnswSparseStreamer"); + + } else { + proxima_index_params_.set(core::PARAM_HNSW_STREAMER_EFCONSTRUCTION, + param_.ef_construction); + proxima_index_params_.set(core::PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, + param_.m); + + // TODO: add_vector_with_id & fetch_by_id don't rely on this param + proxima_index_params_.set(core::PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, + true); + + // TODO: use index params' default query param here + proxima_index_params_.set(core::PARAM_HNSW_STREAMER_EF, + kDefaultHnswEfSearch); + proxima_index_params_.set(core::PARAM_HNSW_STREAMER_USE_ID_MAP, + param_.use_id_map); + proxima_index_params_.set(core::PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY, + param_.use_contiguous_memory); + proxima_index_params_.set(core::PARAM_HNSW_STREAMER_USE_EXTERNAL_VECTOR, + param_.use_external_vector); + streamer_ = core::IndexFactory::CreateStreamer("HnswStreamer"); + } + + if (ailego_unlikely(!streamer_)) { + LOG_ERROR("Failed to create streamer"); + return core::IndexError_Runtime; + } + if (ailego_unlikely( + streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init streamer"); + return core::IndexError_Runtime; + } + return 0; +} + + +int HNSWIndex::_prepare_for_search( + const VectorData & /*vector_data*/, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { + const auto &hnsw_search_param = + std::dynamic_pointer_cast(search_param); + + if (ailego_unlikely(!hnsw_search_param)) { + LOG_ERROR("Invalid search param type, expected HNSWQueryParam"); + return core::IndexError_Runtime; + } + + if (0 >= hnsw_search_param->ef_search || + hnsw_search_param->ef_search > 2048) { + LOG_ERROR( + "ef_search must be greater than 0 and less than or equal to 2048."); + return core::IndexError_Runtime; + } + + context->set_topk(hnsw_search_param->topk); + context->set_fetch_vector(hnsw_search_param->fetch_vector); + if (hnsw_search_param->filter) { + context->set_filter(std::move(*hnsw_search_param->filter)); + } + if (hnsw_search_param->radius > 0.0f) { + context->set_threshold(hnsw_search_param->radius); + } + ailego::Params params; + const int real_search_ef = + std::max(1u, std::min(2048u, hnsw_search_param->ef_search)); + params.set(core::PARAM_HNSW_STREAMER_EF, real_search_ef); + const uint32_t real_search_po = + std::min(256u, hnsw_search_param->prefetch_offset); + params.set(core::PARAM_HNSW_STREAMER_PO, real_search_po); + const uint32_t real_search_pl = + std::min(256u, hnsw_search_param->prefetch_lines); + params.set(core::PARAM_HNSW_STREAMER_PL, real_search_pl); + context->update(params); + _set_group_by_on_context(search_param, context); + return 0; +} + +int HNSWIndex::_get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param) { + const auto &hnsw_search_param = + std::dynamic_pointer_cast(search_param); + + // scale_factor doesn't take effect for hnsw. + auto ret = std::max(search_param->topk, hnsw_search_param->ef_search); + return ret; +} + +} // namespace zvec::core_interface \ No newline at end of file diff --git a/src/core/interface/indexes/hnsw_rabitq_index.cc b/src/core/interface/indexes/hnsw_rabitq_index.cc new file mode 100644 index 0000000..c555389 --- /dev/null +++ b/src/core/interface/indexes/hnsw_rabitq_index.cc @@ -0,0 +1,145 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "zvec/core/framework/index_error.h" + +#if RABITQ_SUPPORTED +#include "algorithm/hnsw_rabitq/hnsw_rabitq_params.h" +#include "algorithm/hnsw_rabitq/hnsw_rabitq_streamer.h" +#include "algorithm/hnsw_rabitq/rabitq_params.h" +#endif + +namespace zvec::core_interface { + +int HNSWRabitqIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { +#if !RABITQ_SUPPORTED + (void)param; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + param_ = dynamic_cast(param); + + if (is_sparse_) { + LOG_ERROR("Sparse index is not supported"); + return core::IndexError_Runtime; + } + + if (param.dimension < core::kMinRabitqDimSize || + param.dimension > core::kMaxRabitqDimSize) { + LOG_ERROR("Unsupported dimension: %d", param.dimension); + return core::IndexError_Unsupported; + } + + // validate parameters + param_.ef_construction = std::max(1, std::min(2048, param_.ef_construction)); + param_.m = std::max(5, std::min(1024, param_.m)); + + proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_EFCONSTRUCTION, + param_.ef_construction); + proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT, + param_.m); + proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_GET_VECTOR_ENABLE, + true); + proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_EF, + kDefaultHnswEfSearch); + proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_USE_ID_MAP, + param_.use_id_map); + proxima_index_params_.set(core::PARAM_HNSW_RABITQ_GENERAL_DIMENSION, + input_vector_meta_.dimension()); + proxima_index_params_.set(core::PARAM_RABITQ_TOTAL_BITS, param_.total_bits); + // num_clusters, sample_count are parameters for rabitq converter + // proxima_index_params_.set(core::PARAM_RABITQ_NUM_CLUSTERS, + // param_.num_clusters); + + auto streamer = std::make_shared(); + streamer->set_provider(param_.provider); + streamer->set_reformer(param_.reformer); + streamer_ = streamer; + + if (ailego_unlikely(!streamer_)) { + LOG_ERROR("Failed to create HnswRabitqStreamer"); + return core::IndexError_Runtime; + } + if (ailego_unlikely( + streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init HnswRabitqStreamer"); + return core::IndexError_Runtime; + } + return 0; +#endif // RABITQ_SUPPORTED +} + +int HNSWRabitqIndex::_prepare_for_search( + const VectorData & /*vector_data*/, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { +#if !RABITQ_SUPPORTED + (void)search_param; + (void)context; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + const auto &hnsw_search_param = + std::dynamic_pointer_cast(search_param); + + if (ailego_unlikely(!hnsw_search_param)) { + LOG_ERROR("Invalid search param type, expected HNSWRabitqQueryParam"); + return core::IndexError_Runtime; + } + + if (0 >= hnsw_search_param->ef_search || + hnsw_search_param->ef_search > 2048) { + LOG_ERROR( + "ef_search must be greater than 0 and less than or equal to 2048."); + return core::IndexError_Runtime; + } + + context->set_topk(hnsw_search_param->topk); + context->set_fetch_vector(hnsw_search_param->fetch_vector); + if (hnsw_search_param->filter) { + context->set_filter(std::move(*hnsw_search_param->filter)); + } + if (hnsw_search_param->radius > 0.0f) { + context->set_threshold(hnsw_search_param->radius); + } + ailego::Params params; + const int real_search_ef = + std::max(1u, std::min(2048u, hnsw_search_param->ef_search)); + params.set(core::PARAM_HNSW_RABITQ_STREAMER_EF, real_search_ef); + context->update(params); + _set_group_by_on_context(search_param, context); + return 0; +#endif // RABITQ_SUPPORTED +} + +int HNSWRabitqIndex::_get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param) { +#if !RABITQ_SUPPORTED + (void)search_param; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return -1; +#else + const auto &hnsw_search_param = + std::dynamic_pointer_cast(search_param); + + auto ret = std::max(search_param->topk, hnsw_search_param->ef_search); + return ret; +#endif // RABITQ_SUPPORTED +} + + +} // namespace zvec::core_interface diff --git a/src/core/interface/indexes/holder_builder.h b/src/core/interface/indexes/holder_builder.h new file mode 100644 index 0000000..0714c2c --- /dev/null +++ b/src/core/interface/indexes/holder_builder.h @@ -0,0 +1,87 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec::core_interface { + +inline constexpr uint64_t kInvalidKey = std::numeric_limits::max(); + +template +inline int BuildMultiPassHolderImpl( + uint32_t dimension, + const std::vector> &doc_cache, + core::IndexHolder::Pointer *holder_out) { + auto holder = + std::make_shared>(dimension); + for (const auto &doc : doc_cache) { + if (doc.first == kInvalidKey) { + continue; + } + ailego::NumericalVector vec(doc.second); + if (!holder->emplace(doc.first, vec)) { + LOG_ERROR("Failed to add vector"); + return core::IndexError_Runtime; + } + } + *holder_out = holder; + return 0; +} + +inline int BuildMultiPassHolder( + DataType data_type, uint32_t dimension, + const std::vector> &doc_cache, + const core::IndexConverter::Pointer &converter, + core::IndexHolder::Pointer *holder) { + int ret = 0; + switch (data_type) { + case DataType::DT_FP16: + ret = BuildMultiPassHolderImpl( + dimension, doc_cache, holder); + break; + case DataType::DT_FP32: + ret = BuildMultiPassHolderImpl( + dimension, doc_cache, holder); + break; + case DataType::DT_INT8: + ret = BuildMultiPassHolderImpl( + dimension, doc_cache, holder); + break; + default: + LOG_ERROR("data_type is not support"); + return core::IndexError_Runtime; + } + if (ret != 0) { + return ret; + } + if (converter) { + core::IndexConverter::TrainAndTransform(converter, *holder); + *holder = converter->result(); + } + return 0; +} + +} // namespace zvec::core_interface diff --git a/src/core/interface/indexes/ivf_index.cc b/src/core/interface/indexes/ivf_index.cc new file mode 100644 index 0000000..9df540f --- /dev/null +++ b/src/core/interface/indexes/ivf_index.cc @@ -0,0 +1,273 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "algorithm/ivf/ivf_params.h" +#include "holder_builder.h" + +namespace zvec::core_interface { + +int IVFIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { + if (is_sparse_) { + LOG_ERROR("IVF Index not support sparse vector"); + return core::IndexError_InvalidArgument; + } + + param_ = dynamic_cast(param); + param_.nlist = std::max(1, std::min(1024, param_.nlist)); + param_.niters = std::max(1, std::min(1024, param_.niters)); + + proxima_index_params_.set(core::PARAM_IVF_BUILDER_CENTROID_COUNT, + param_.nlist); + + // TODO: add_vector_with_id & fetch_by_id don't rely on this param + builder_ = core::IndexFactory::CreateBuilder("IVFBuilder"); + streamer_ = core::IndexFactory::CreateStreamer("IVFStreamer"); + + if (ailego_unlikely(!builder_)) { + LOG_ERROR("Failed to create builder"); + return core::IndexError_Runtime; + } + if (ailego_unlikely(!streamer_)) { + LOG_ERROR("Failed to create streamer"); + return core::IndexError_Runtime; + } + IndexMeta real_meta; + if (converter_) { + real_meta = converter_->meta(); + } else { + real_meta = proxima_index_meta_; + } + if (ailego_unlikely(builder_->init(real_meta, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init builder"); + return core::IndexError_Runtime; + } + if (ailego_unlikely(streamer_->init(real_meta, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init streamer"); + return core::IndexError_Runtime; + } + return 0; +} + +int IVFIndex::Open(const std::string &file_path, + StorageOptions storage_options) { + ailego::Params storage_params; + file_path_ = file_path; + is_read_only_ = storage_options.read_only; + switch (storage_options.type) { + case StorageOptions::StorageType::kMMAP: { + storage_ = core::IndexFactory::CreateStorage("MMapFileReadStorage"); + if (storage_ == nullptr) { + LOG_ERROR("Failed to create MMapFileStorage"); + return core::IndexError_Runtime; + } + int ret = storage_->init(storage_params); + if (ret != 0) { + LOG_ERROR("Failed to init MMapFileStorage, path: %s, err: %s", + file_path_.c_str(), core::IndexError::What(ret)); + return ret; + } + break; + } + case StorageOptions::StorageType::kBufferPool: { + // NOTE: IVF index is dumped via FileDumper (plain binary file), which is + // not compatible with BufferStorage's IndexFormat layout (header/footer + // chain). Until IVF gains a BufferStorage-aware dump path, fall back to + // MMapFileReadStorage so the freshly-dumped file can be reopened. + storage_ = core::IndexFactory::CreateStorage("MMapFileReadStorage"); + if (storage_ == nullptr) { + LOG_ERROR( + "Failed to create MMapFileReadStorage (IVF buffer-pool fallback)"); + return core::IndexError_Runtime; + } + int ret = storage_->init(storage_params); + if (ret != 0) { + LOG_ERROR( + "Failed to init MMapFileReadStorage (IVF buffer-pool fallback), " + "path: %s, err: %s", + file_path_.c_str(), core::IndexError::What(ret)); + return ret; + } + break; + } + default: { + LOG_ERROR("Unsupported storage type"); + return core::IndexError_Unsupported; + } + } + + if (is_read_only_ || !storage_options.create_new) { + // read_options.create_new + int ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + // Load reformer data from storage (e.g., rotation matrix for INT8+rotate) + if (reformer_ != nullptr && reformer_->load(storage_) != 0) { + LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + is_trained_ = true; + } + is_open_ = true; + return 0; +} + +int IVFIndex::GenerateHolder() { + return BuildMultiPassHolder(param_.data_type, param_.dimension, doc_cache_, + converter_, &holder_); +} + +int IVFIndex::Add(const VectorData &vector, uint32_t doc_id) { + if (is_trained_) { + LOG_ERROR("this IVF index is trained"); + return core::IndexError_Runtime; + } + if (!std::holds_alternative(vector.vector)) { + LOG_ERROR("Invalid vector data"); + return core::IndexError_Runtime; + } + const DenseVector &dense_vector = std::get(vector.vector); + std::string out_vector_buffer = std::string( + static_cast(dense_vector.data), + input_vector_meta_.dimension() * input_vector_meta_.unit_size()); + + std::lock_guard lock(mutex_); + while (doc_cache_.size() <= doc_id) { + std::string fake_data( + input_vector_meta_.dimension() * input_vector_meta_.unit_size(), 0); + doc_cache_.push_back(std::make_pair(kInvalidKey, fake_data)); + } + doc_cache_[doc_id] = std::make_pair(doc_id, out_vector_buffer); + return 0; +} + +int IVFIndex::Train() { + GenerateHolder(); + builder_->train(holder_); + builder_->build(holder_); + auto dumper = core::IndexFactory::CreateDumper("FileDumper"); + + dumper->create(file_path_); + builder_->dump(dumper); + // Dump converter state (e.g., rotator for INT8+rotate) to dumper + if (converter_ && converter_->dump(dumper) != 0) { + LOG_ERROR("Failed to dump converter, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + dumper->close(); + int ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + // Load reformer data from storage (e.g., rotation matrix) + if (reformer_ != nullptr && reformer_->load(storage_) != 0) { + LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + is_trained_ = true; + return 0; +} + +int IVFIndex::_dense_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer) { + if (is_trained_) { + return Index::_dense_fetch(doc_id, vector_data_buffer); + } else { + DenseVectorBuffer dense_vector_buffer; + std::string &out_vector_buffer = dense_vector_buffer.data; + out_vector_buffer = doc_cache_[doc_id].second; + vector_data_buffer->vector_buffer = std::move(dense_vector_buffer); + return 0; + } +} + +int IVFIndex::_prepare_for_search( + const VectorData & /*query*/, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { + const auto &ivf_search_param = + std::dynamic_pointer_cast(search_param); + + if (search_param->group_by_param && search_param->group_by_param->group_by) { + LOG_ERROR("group_by search is not supported for IVF index"); + return core::IndexError_Unsupported; + } + + context->set_topk(ivf_search_param->topk); + context->set_fetch_vector(ivf_search_param->fetch_vector); + if (ivf_search_param->filter) { + context->set_filter(std::move(*ivf_search_param->filter)); + } + if (ivf_search_param->radius > 0.0f) { + context->set_threshold(ivf_search_param->radius); + } + + if (ivf_search_param->nprobe > 0) { + ailego::Params params; + params.set(core::PARAM_IVF_SEARCHER_NPROBE, ivf_search_param->nprobe); + context->update(params); + } + return 0; +} + +int IVFIndex::Merge(const std::vector &indexes, + const IndexFilter &filter, const MergeOptions &options) { + int pre_ret = Index::Merge(indexes, filter, options); + if (pre_ret != 0) { + return pre_ret; + } + auto dumper = core::IndexFactory::CreateDumper("FileDumper"); + + dumper->create(file_path_); + builder_->dump(dumper); + // Dump converter state (e.g., rotator for INT8+rotate) to dumper + if (converter_ && converter_->dump(dumper) != 0) { + LOG_ERROR("Failed to dump converter, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + dumper->close(); + int ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + // Load reformer data from storage (e.g., rotation matrix) + if (reformer_ != nullptr && reformer_->load(storage_) != 0) { + LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + is_trained_ = true; + return 0; +} +} // namespace zvec::core_interface \ No newline at end of file diff --git a/src/core/interface/indexes/vamana_index.cc b/src/core/interface/indexes/vamana_index.cc new file mode 100644 index 0000000..f40991b --- /dev/null +++ b/src/core/interface/indexes/vamana_index.cc @@ -0,0 +1,116 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "algorithm/vamana/vamana_params.h" + +namespace zvec::core_interface { + +int VamanaIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { + param_ = dynamic_cast(param); + + // Validate parameters + param_.max_degree = std::max(5, std::min(256, param_.max_degree)); + param_.search_list_size = + std::max(10, std::min(2048, param_.search_list_size)); + if (param_.alpha <= 0.0f) param_.alpha = kDefaultVamanaAlpha; + + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_MAX_DEGREE, + static_cast(param_.max_degree)); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, + static_cast(param_.search_list_size)); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_ALPHA, param_.alpha); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_MAX_OCCLUSION_SIZE, + static_cast(param_.max_occlusion_size)); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_SATURATE_GRAPH, + param_.saturate_graph); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, + true); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_EF, + kDefaultVamanaEfSearch); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_USE_ID_MAP, + param_.use_id_map); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY, + param_.use_contiguous_memory); + + streamer_ = core::IndexFactory::CreateStreamer("VamanaStreamer"); + + if (ailego_unlikely(!streamer_)) { + LOG_ERROR("Failed to create VamanaStreamer"); + return core::IndexError_Runtime; + } + if (ailego_unlikely( + streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init VamanaStreamer"); + return core::IndexError_Runtime; + } + return 0; +} + +int VamanaIndex::_prepare_for_search( + const VectorData & /*vector_data*/, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { + const auto &vamana_search_param = + std::dynamic_pointer_cast(search_param); + + if (ailego_unlikely(!vamana_search_param)) { + LOG_ERROR("Invalid search param type, expected VamanaQueryParam"); + return core::IndexError_Runtime; + } + + if (search_param->group_by_param && search_param->group_by_param->group_by) { + LOG_ERROR("group_by search is not supported for Vamana index"); + return core::IndexError_Unsupported; + } + + if (vamana_search_param->ef_search == 0 || + vamana_search_param->ef_search > 2048) { + LOG_ERROR( + "ef_search must be greater than 0 and less than or equal to 2048."); + return core::IndexError_Runtime; + } + + context->set_topk(vamana_search_param->topk); + context->set_fetch_vector(vamana_search_param->fetch_vector); + if (vamana_search_param->filter) { + context->set_filter(std::move(*vamana_search_param->filter)); + } + if (vamana_search_param->radius > 0.0f) { + context->set_threshold(vamana_search_param->radius); + } + ailego::Params params; + const uint32_t real_search_ef = + std::max(1u, std::min(2048u, vamana_search_param->ef_search)); + params.set(core::PARAM_VAMANA_STREAMER_EF, real_search_ef); + const uint32_t real_search_po = + std::min(256u, vamana_search_param->prefetch_offset); + params.set(core::PARAM_VAMANA_STREAMER_PO, real_search_po); + const uint32_t real_search_pl = + std::min(256u, vamana_search_param->prefetch_lines); + params.set(core::PARAM_VAMANA_STREAMER_PL, real_search_pl); + context->update(params); + return 0; +} + +int VamanaIndex::_get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param) { + const auto &vamana_search_param = + std::dynamic_pointer_cast(search_param); + return std::max(search_param->topk, vamana_search_param->ef_search); +} + +} // namespace zvec::core_interface diff --git a/src/core/interface/utils/utils.h b/src/core/interface/utils/utils.h new file mode 100644 index 0000000..5fc5365 --- /dev/null +++ b/src/core/interface/utils/utils.h @@ -0,0 +1,104 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include + +namespace zvec { +namespace core_interface { + +template +constexpr bool extract_enum_from_json(const ailego::JsonObject &json_obj, + const char *key, EnumType &enum_value, + ailego::JsonValue &tmp_json_value) { + if (json_obj.has(key)) { + if (json_obj.get(key, &tmp_json_value); tmp_json_value.is_string()) { + auto optional_enum_value = + magic_enum::enum_cast(tmp_json_value.as_stl_string()); + if (optional_enum_value.has_value()) { + enum_value = optional_enum_value.value(); + } else { + LOG_ERROR("Invalid enum value for key: %s, value: %s", key, + tmp_json_value.as_c_string()); + return false; + } + } else { + LOG_ERROR("Invalid json field type for key: %s", key); + return false; + } + } + return true; +} + +template +constexpr bool extract_value_from_json(const ailego::JsonObject &json_obj, + const char *key, T &value, + ailego::JsonValue &tmp_json_value) { + if (json_obj.has(key)) { + json_obj.get(key, &tmp_json_value); + if constexpr (std::is_same_v) { + if (tmp_json_value.is_boolean()) { + value = tmp_json_value.as_bool(); + } else { + LOG_ERROR("Invalid json field type for key: %s; expected: boolean", + key); + return false; + } + } else if constexpr (std::is_floating_point_v) { + if (tmp_json_value.is_float() || tmp_json_value.is_integer()) { + value = static_cast(tmp_json_value.as_float()); + } else { + LOG_ERROR("Invalid json field type for key: %s; expected: float", key); + return false; + } + } else if constexpr (std::is_integral_v) { + if (tmp_json_value.is_integer()) { + value = static_cast(tmp_json_value.as_integer()); + } else { + LOG_ERROR("Invalid json field type for key: %s; expected: integer", + key); + return false; + } + } else { + abort(); + } + } + return true; +} + +#define DESERIALIZE_ENUM_FIELD(json_obj, field_name, EnumType) \ + { \ + ailego::JsonValue tmp_json_value; \ + if (!extract_enum_from_json(json_obj, #field_name, field_name, \ + tmp_json_value)) { \ + LOG_ERROR("Error when deserialize json - field:%s", #field_name); \ + return false; \ + } \ + } + + +#define DESERIALIZE_VALUE_FIELD(json_obj, field_name) \ + { \ + ailego::JsonValue tmp_json_value; \ + if (!extract_value_from_json(json_obj, #field_name, field_name, \ + tmp_json_value)) { \ + LOG_ERROR("Error when deserialize json - field:%s", #field_name); \ + return false; \ + } \ + } +} // namespace core_interface +} // namespace zvec \ No newline at end of file diff --git a/src/core/metric/CMakeLists.txt b/src/core/metric/CMakeLists.txt new file mode 100644 index 0000000..2918b90 --- /dev/null +++ b/src/core/metric/CMakeLists.txt @@ -0,0 +1,17 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(NOT APPLE) + set(CORE_METRIC_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_metric + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS zvec_ailego zvec_turbo core_framework + INCS . ${PROJECT_ROOT_DIR}/src/core + LDFLAGS "${CORE_METRIC_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/metric/cosine_metric.cc b/src/core/metric/cosine_metric.cc new file mode 100644 index 0000000..1b0fb05 --- /dev/null +++ b/src/core/metric/cosine_metric.cc @@ -0,0 +1,235 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +//! Retrieve distance function for index features +inline IndexMetric::MatrixDistanceHandle CosineDistanceMatrixFp32(size_t m, + size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +//! Retrieve distance function for index features +inline IndexMetric::MatrixDistanceHandle CosineDistanceMatrixFp16(size_t m, + size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute), + reinterpret_cast( + ailego::CosineDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +/*! Cosine Metric + */ +class CosineMetric : public IndexMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + IndexMeta::DataType ft = meta.data_type(); + if (ft != IndexMeta::DataType::DT_FP16 && + ft != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + if (IndexMeta::UnitSizeof(ft) != meta.unit_size()) { + return IndexError_Unsupported; + } + data_type_ = ft; + params_ = index_params; + + return 0; + } + + //! Cleanup Metric + int cleanup(void) override { + return 0; + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta) const override { + return (meta.data_type() == data_type_ && + meta.unit_size() == IndexMeta::UnitSizeof(data_type_)); + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return (qmeta.data_type() == data_type_ && + qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) && + qmeta.dimension() == meta.dimension()); + } + + //! Retrieve distance function for query + MatrixDistance distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::CosineDistanceMatrix::Compute); + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::CosineDistanceMatrix::Compute); + + default: + return nullptr; + } + } + + //! Retrieve distance function for index features + MatrixDistance distance_matrix(size_t m, size_t n) const override { + if (m != 1 || n != 1) { + return nullptr; + } + return distance(); + } + + //! Retrieve distance function for query + MatrixBatchDistance batch_distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + default: + return nullptr; + } + } + + //! Retrieve params of Metric + const ailego::Params ¶ms(void) const override { + return params_; + } + + //! Retrieve query metric object of this index metric + Pointer query_metric(void) const override { + return nullptr; + } + + private: + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32}; + ailego::Params params_{}; +}; + +INDEX_FACTORY_REGISTER_METRIC_ALIAS(Cosine, CosineMetric); + +} // namespace core +} // namespace zvec diff --git a/src/core/metric/euclidean_metric.cc b/src/core/metric/euclidean_metric.cc new file mode 100644 index 0000000..565284f --- /dev/null +++ b/src/core/metric/euclidean_metric.cc @@ -0,0 +1,839 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle +SquaredEuclideanDistanceMatrixFp32(size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle +SquaredEuclideanDistanceMatrixFp16(size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +static inline IndexMetric::MatrixDistanceHandle +SquaredEuclideanDistanceMatrixInt8(size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +//! Retrieve distance function for index features in Int4 +static inline IndexMetric::MatrixDistanceHandle +SquaredEuclideanDistanceMatrixInt4(size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle EuclideanDistanceMatrixFp32( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle EuclideanDistanceMatrixFp16( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +static inline IndexMetric::MatrixDistanceHandle EuclideanDistanceMatrixInt8( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +//! Retrieve distance function for index features in Int4 +static inline IndexMetric::MatrixDistanceHandle EuclideanDistanceMatrixInt4( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute), + reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +/*! Squared Euclidean Distance Metric + */ +class SquaredEuclideanMetric : public IndexMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + IndexMeta::DataType dt = meta.data_type(); + if (dt != IndexMeta::DataType::DT_FP16 && + dt != IndexMeta::DataType::DT_FP32 && + dt != IndexMeta::DataType::DT_INT8 && + dt != IndexMeta::DataType::DT_INT4) { + return IndexError_Unsupported; + } + if (IndexMeta::UnitSizeof(dt) != meta.unit_size()) { + return IndexError_Unsupported; + } + data_type_ = dt; + params_ = index_params; + + return 0; + } + + //! Cleanup Metric + int cleanup(void) override { + return 0; + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta) const override { + return (meta.data_type() == data_type_ && + meta.unit_size() == IndexMeta::UnitSizeof(data_type_)); + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return (qmeta.data_type() == data_type_ && + qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) && + qmeta.dimension() == meta.dimension()); + } + + //! Retrieve distance function for query + MatrixDistance distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_INT8: + return reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_INT4: + return reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute); + + default: + return nullptr; + } + } + + //! Retrieve sparse distance function for query + MatrixSparseDistance sparse_distance(void) const override { + return reinterpret_cast( + ailego::SquaredEuclideanSparseDistanceMatrix::Compute); + } + + //! Retrieve distance function for index features + MatrixDistance distance_matrix(size_t m, size_t n) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return SquaredEuclideanDistanceMatrixFp16(m, n); + + case IndexMeta::DataType::DT_FP32: + return SquaredEuclideanDistanceMatrixFp32(m, n); + + case IndexMeta::DataType::DT_INT8: + return SquaredEuclideanDistanceMatrixInt8(m, n); + + case IndexMeta::DataType::DT_INT4: + return SquaredEuclideanDistanceMatrixInt4(m, n); + + default: + return nullptr; + } + } + + //! Retrieve distance function for query + MatrixBatchDistance batch_distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + + case IndexMeta::DataType::DT_INT8: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + + case IndexMeta::DataType::DT_INT4: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + + default: + return nullptr; + } + } + + //! Retrieve params of Metric + const ailego::Params ¶ms(void) const override { + return params_; + } + + //! Retrieve query metric object of this index metric + Pointer query_metric(void) const override { + return nullptr; + } + + private: + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32}; + ailego::Params params_{}; +}; + +/*! Euclidean Distance Metric + */ +class EuclideanMetric : public IndexMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + IndexMeta::DataType dt = meta.data_type(); + if (dt != IndexMeta::DataType::DT_FP16 && + dt != IndexMeta::DataType::DT_FP32 && + dt != IndexMeta::DataType::DT_INT8 && + dt != IndexMeta::DataType::DT_INT4) { + return IndexError_Unsupported; + } + if (IndexMeta::UnitSizeof(dt) != meta.unit_size()) { + return IndexError_Unsupported; + } + data_type_ = dt; + params_ = index_params; + return 0; + } + + //! Cleanup Metric + int cleanup(void) override { + return 0; + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta) const override { + return (meta.data_type() == data_type_ && + meta.unit_size() == IndexMeta::UnitSizeof(data_type_)); + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return (qmeta.data_type() == data_type_ && + qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) && + qmeta.dimension() == meta.dimension()); + } + + //! Retrieve distance function for query + MatrixDistance distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_INT8: + return reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_INT4: + return reinterpret_cast( + ailego::EuclideanDistanceMatrix::Compute); + + default: + return nullptr; + } + } + + //! Retrieve distance function for index features + MatrixDistance distance_matrix(size_t m, size_t n) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return EuclideanDistanceMatrixFp16(m, n); + + case IndexMeta::DataType::DT_FP32: + return EuclideanDistanceMatrixFp32(m, n); + + case IndexMeta::DataType::DT_INT8: + return EuclideanDistanceMatrixInt8(m, n); + + case IndexMeta::DataType::DT_INT4: + return EuclideanDistanceMatrixInt4(m, n); + + default: + return nullptr; + } + } + + //! Retrieve distance function for query + MatrixBatchDistance batch_distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + + case IndexMeta::DataType::DT_INT8: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + + case IndexMeta::DataType::DT_INT4: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + + default: + return nullptr; + } + } + + //! Retrieve params of Metric + const ailego::Params ¶ms(void) const override { + return params_; + } + + //! Retrieve query metric object of this index metric + Pointer query_metric(void) const override { + return nullptr; + } + + private: + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32}; + ailego::Params params_{}; +}; + +/*! Squared Euclidean Sparse Metric + */ +class SquaredEuclideanSparseMetric : public IndexMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + IndexMeta::DataType data_type = meta.data_type(); + if (data_type != IndexMeta::DataType::DT_FP16 && + data_type != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + + if (IndexMeta::UnitSizeof(data_type) != meta.unit_size()) { + return IndexError_Unsupported; + } + + data_type_ = data_type; + params_ = index_params; + + return 0; + } + + //! Cleanup Metric + int cleanup(void) override { + return 0; + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta) const override { + return (meta.data_type() == data_type_ && + meta.unit_size() == IndexMeta::UnitSizeof(data_type_)); + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return (qmeta.data_type() == data_type_ && + qmeta.data_type() == meta.data_type() && + qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) && + qmeta.unit_size() == meta.unit_size()); + } + + //! Retrieve sparse distance function for query + MatrixSparseDistance sparse_distance(void) const override { + return reinterpret_cast( + ailego::SquaredEuclideanSparseDistanceMatrix::Compute); + } + + //! Retrieve params of Metric + const ailego::Params ¶ms(void) const override { + return params_; + } + + //! Retrieve query metric object of this index metric + Pointer query_metric(void) const override { + return nullptr; + } + + private: + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32}; + + ailego::Params params_{}; +}; + +INDEX_FACTORY_REGISTER_METRIC_ALIAS(SquaredEuclidean, SquaredEuclideanMetric); +INDEX_FACTORY_REGISTER_METRIC_ALIAS(Euclidean, EuclideanMetric); + +INDEX_FACTORY_REGISTER_METRIC_ALIAS(SquaredEuclideanSparse, + SquaredEuclideanSparseMetric); + +} // namespace core +} // namespace zvec diff --git a/src/core/metric/inner_product_metric.cc b/src/core/metric/inner_product_metric.cc new file mode 100644 index 0000000..8ef0a11 --- /dev/null +++ b/src/core/metric/inner_product_metric.cc @@ -0,0 +1,534 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle MinusInnerProductMatrixFp32( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle MinusInnerProductMatrixFp16( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +static inline IndexMetric::MatrixDistanceHandle MinusInnerProductMatrixInt8( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +//! Retrieve distance function for index features in Int4 +static inline IndexMetric::MatrixDistanceHandle MinusInnerProductMatrixInt4( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute), + reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute)}, + }; + if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) { + return nullptr; + } + return distance_table[ailego_ctz(m)][ailego_ctz(n)]; +} + +/*! Inner Product Metric + */ +class InnerProductMetric : public IndexMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + IndexMeta::MetaType mt = meta.meta_type(); + if (mt != IndexMeta::MetaType::MT_DENSE) { + return IndexError_Unsupported; + } + + IndexMeta::DataType dt = meta.data_type(); + if (dt != IndexMeta::DataType::DT_FP16 && + dt != IndexMeta::DataType::DT_FP32 && + dt != IndexMeta::DataType::DT_INT8 && + dt != IndexMeta::DataType::DT_INT4) { + return IndexError_Unsupported; + } + if (IndexMeta::UnitSizeof(dt) != meta.unit_size()) { + return IndexError_Unsupported; + } + + meta_type_ = mt; + data_type_ = dt; + params_ = index_params; + + return 0; + } + + //! Cleanup Metric + int cleanup(void) override { + return 0; + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta) const override { + return (meta.data_type() == data_type_ && + meta.unit_size() == IndexMeta::UnitSizeof(data_type_)); + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return (qmeta.data_type() == data_type_ && + qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) && + qmeta.dimension() == meta.dimension()); + } + + //! Retrieve distance function for query + MatrixDistance distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute); + + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute); + + case IndexMeta::DataType::DT_INT8: + return reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute); + + case IndexMeta::DataType::DT_INT4: + return reinterpret_cast( + ailego::MinusInnerProductMatrix::Compute); + + default: + return nullptr; + } + } + + //! Retrieve sparse distance function for query + MatrixSparseDistance sparse_distance(void) const override { + return reinterpret_cast( + ailego::MinusInnerProductSparseMatrix::Compute); + } + + //! Retrieve distance function for index features + MatrixDistance distance_matrix(size_t m, size_t n) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return MinusInnerProductMatrixFp16(m, n); + + case IndexMeta::DataType::DT_FP32: + return MinusInnerProductMatrixFp32(m, n); + + case IndexMeta::DataType::DT_INT8: + return MinusInnerProductMatrixInt8(m, n); + + case IndexMeta::DataType::DT_INT4: + return MinusInnerProductMatrixInt4(m, n); + + default: + return nullptr; + } + } + + //! Retrieve distance function for query + MatrixBatchDistance batch_distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + case IndexMeta::DataType::DT_INT8: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + case IndexMeta::DataType::DT_INT4: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + default: + return nullptr; + } + } + + //! Normalize result + void normalize(float *score) const override { + *score = -(*score); + } + + //! Denormalize threshold + void denormalize(float *score) const override { + *score = -(*score); + } + + //! Retrieve if it supports normalization + bool support_normalize(void) const override { + return true; + } + + //! Retrieve params of Metric + const ailego::Params ¶ms(void) const override { + return params_; + } + + //! Retrieve query measure object of this index measure + Pointer query_metric(void) const override { + return nullptr; + } + + private: + IndexMeta::MetaType meta_type_{IndexMeta::MetaType::MT_DENSE}; + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32}; + ailego::Params params_{}; +}; + +/*! Normalized Cosine Metric + */ +class NormalizedCosineMetric : public InnerProductMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + IndexMeta::DataType dt = meta.data_type(); + if (dt != IndexMeta::DataType::DT_FP16 && + dt != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + + InnerProductMetric::init(meta, index_params); + + return 0; + } + + //! Normalize result + void normalize(float *score) const override { + *score = 1 + (*score); + } + + //! Denormalize threshold + + void denormalize(float *score) const override { + *score -= 1; + } +}; + +/*! Inner Product Sparse Metric + */ +class InnerProductSparseMetric : public IndexMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + IndexMeta::DataType dt = meta.data_type(); + if (dt != IndexMeta::DataType::DT_FP16 && + dt != IndexMeta::DataType::DT_FP32) { + return IndexError_Unsupported; + } + + if (IndexMeta::UnitSizeof(dt) != meta.unit_size()) { + return IndexError_Unsupported; + } + + data_type_ = dt; + params_ = index_params; + + return 0; + } + + //! Cleanup Metric + int cleanup(void) override { + return 0; + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta) const override { + return (meta.data_type() == data_type_ && + meta.unit_size() == IndexMeta::UnitSizeof(data_type_)); + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return (qmeta.data_type() == data_type_ && + qmeta.data_type() == meta.data_type() && + qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) && + qmeta.unit_size() == meta.unit_size()); + } + + //! Retrieve distance function for query + MatrixDistance distance(void) const override { + return nullptr; + } + + //! Retrieve sparse distance function for query + MatrixSparseDistance sparse_distance(void) const override { + switch (data_type_) { + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::MinusInnerProductSparseMatrix::Compute); + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::MinusInnerProductSparseMatrix::Compute); + default: + return nullptr; + } + } + + //! Normalize result + void normalize(float *score) const override { + *score = -(*score); + } + + //! Denormalize threshold + void denormalize(float *score) const override { + *score = -(*score); + } + + //! Retrieve if it supports normalization + bool support_normalize(void) const override { + return true; + } + + //! Retrieve params of Metric + const ailego::Params ¶ms(void) const override { + return params_; + } + + //! Retrieve query measure object of this index measure + Pointer query_metric(void) const override { + return nullptr; + } + + private: + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32}; + ailego::Params params_{}; +}; + +INDEX_FACTORY_REGISTER_METRIC_ALIAS(InnerProduct, InnerProductMetric); +INDEX_FACTORY_REGISTER_METRIC_ALIAS(NormalizedCosine, NormalizedCosineMetric); + +INDEX_FACTORY_REGISTER_METRIC_ALIAS(InnerProductSparse, + InnerProductSparseMetric); + +} // namespace core +} // namespace zvec diff --git a/src/core/metric/metric_params.h b/src/core/metric/metric_params.h new file mode 100644 index 0000000..05f8db9 --- /dev/null +++ b/src/core/metric/metric_params.h @@ -0,0 +1,42 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include + +namespace zvec { +namespace core { + +//! MipsEuclideanMetric +static const std::string MIPS_EUCLIDEAN_METRIC_M_VALUE = + "mips_euclidean.metric.m_value"; +static const std::string MIPS_EUCLIDEAN_METRIC_U_VALUE = + "mips_euclidean.metric.u_value"; +static const std::string MIPS_EUCLIDEAN_METRIC_MAX_L2_NORM = + "mips_euclidean.metric.max_l2_norm"; +static const std::string MIPS_EUCLIDEAN_METRIC_INJECTION_TYPE = + "mips_euclidean.metric.injection_type"; + +//! QuantizedInteger Metric +static const std::string QUANTIZED_INTEGER_METRIC_ORIGIN_METRIC_NAME = + "proxima.quantized_integer.metric.origin_metric_name"; +static const std::string QUANTIZED_INTEGER_METRIC_ORIGIN_METRIC_PARAMS = + "proxima.quantized_integer.metric.origin_metric_params"; + +//! UniformInt8 Metric +static const std::string UNIFORM_INT8_METRIC_ORIGIN_METRIC_NAME = + "proxima.uniform_int8.metric.origin_metric_name"; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/metric/mips_euclidean_metric.cc b/src/core/metric/mips_euclidean_metric.cc new file mode 100644 index 0000000..280bfa7 --- /dev/null +++ b/src/core/metric/mips_euclidean_metric.cc @@ -0,0 +1,683 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include +#include +#include +#include +#include +#include "metric_params.h" + +namespace zvec { +namespace core { + +/*! Mips Squared Euclidean Metric + */ +template +class MipsSquaredEuclideanMetric : public IndexMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + data_type_ = meta.data_type(); + dimension_ = meta.dimension(); + + int injection_type = static_cast(kDefaultInjectionType); + index_params.get(MIPS_EUCLIDEAN_METRIC_INJECTION_TYPE, &injection_type); + if (injection_type >= static_cast(Injection::kNumInjections)) { + LOG_WARN("Unsupported injection_type %u, using '%s' instead", + injection_type, InjectionName(0)); + injection_type = static_cast(Injection::kLocalizedSpherical); + } + injection_ = static_cast(injection_type); + LOG_DEBUG( + "Initializing MipsSquaredEuclideanMetric with injection %s" + " type %d dimension %d", + InjectionName(injection_), data_type_, dimension_); + + float max_l2_norm = 0.0f; + float u_value = 0.0f; + index_params.get(MIPS_EUCLIDEAN_METRIC_M_VALUE, &m_value_); + index_params.get(MIPS_EUCLIDEAN_METRIC_U_VALUE, &u_value); + index_params.get(MIPS_EUCLIDEAN_METRIC_MAX_L2_NORM, &max_l2_norm); + CheckAndFixM(injection_, &m_value_); + CheckAndFixU(injection_, m_value_, &u_value); + + squared_u_value_ = u_value * u_value; + max_squared_l2_norm_ = max_l2_norm * max_l2_norm; + if (injection_ == Injection::kIdentity || + injection_ == Injection::kLocalizedSpherical) { + eta_ = 0.0f; + } else if (max_squared_l2_norm_ < std::numeric_limits::epsilon()) { + eta_ = kDefaultEta; + } else { + eta_ = squared_u_value_ / max_squared_l2_norm_; + } + + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + squared_norm2_handle_ = reinterpret_cast( + ailego::SquaredNorm2Matrix::Compute); + break; + + case IndexMeta::DataType::DT_FP16: + squared_norm2_handle_ = reinterpret_cast( + ailego::SquaredNorm2Matrix::Compute); + break; + + case IndexMeta::DataType::DT_INT8: + squared_norm2_handle_ = reinterpret_cast( + ailego::SquaredNorm2Matrix::Compute); + break; + + case IndexMeta::DataType::DT_INT4: + squared_norm2_handle_ = reinterpret_cast( + ailego::SquaredNorm2Matrix::Compute); + break; + + default: + return IndexError_Unsupported; + } + + query_metric_ = IndexFactory::CreateMetric(kQueryMetric); + if (!query_metric_) { + LOG_ERROR("Failed to create metric %s", kQueryMetric); + return IndexError_NoExist; + } + int ret = query_metric_->init(meta, ailego::Params()); + if (ret != 0) { + LOG_ERROR("Failed to initialize metric %s", kQueryMetric); + return ret; + } + params_ = index_params; + return 0; + } + + //! Cleanup Metric + int cleanup(void) override { + eta_ = 0.0f; + m_value_ = 0; + squared_u_value_ = 0.0f; + max_squared_l2_norm_ = 0.0f; + query_metric_.reset(); + return 0; + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta) const override { + return (meta.data_type() == data_type_ && + meta.unit_size() == IndexMeta::UnitSizeof(data_type_)); + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return (qmeta.data_type() == data_type_ && + qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) && + qmeta.dimension() == meta.dimension()); + } + + //! Retrieve distance function for query + MatrixBatchDistance batch_distance() const override { + MatrixDistance dist_func = distance(); + + return + [=](const void **m, const void *q, size_t num, size_t dim, float *out) { + for (size_t i = 0; i < num; ++i) { + dist_func(m[i], q, dim, out + i); + } + }; + } + + + //! Retrieve distance function for query + MatrixDistance distance(void) const override { + if (injection_ == Injection::kLocalizedSpherical) { + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, 0.0f, out); + }; + + case IndexMeta::DataType::DT_FP16: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix:: + Compute(reinterpret_cast(m), + reinterpret_cast(q), dim, 0.0f, + out); + }; + + case IndexMeta::DataType::DT_INT8: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, 0.0f, out); + }; + + case IndexMeta::DataType::DT_INT4: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, 0.0f, out); + }; + + default: + return nullptr; + } + } + + if (injection_ == Injection::kRepeatedQuadratic) { + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, m_value_, eta_, out); + }; + + case IndexMeta::DataType::DT_FP16: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix:: + Compute(reinterpret_cast(m), + reinterpret_cast(q), dim, + m_value_, eta_, out); + }; + + case IndexMeta::DataType::DT_INT8: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, m_value_, eta_, out); + }; + + case IndexMeta::DataType::DT_INT4: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, m_value_, eta_, out); + }; + + default: + return nullptr; + } + } + + if (injection_ == Injection::kSpherical) { + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, eta_, out); + }; + + case IndexMeta::DataType::DT_FP16: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix:: + Compute(reinterpret_cast(m), + reinterpret_cast(q), dim, eta_, + out); + }; + + case IndexMeta::DataType::DT_INT8: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, eta_, out); + }; + + case IndexMeta::DataType::DT_INT4: + return [&](const void *m, const void *q, size_t dim, float *out) { + ailego::MipsSquaredEuclideanDistanceMatrix::Compute( + reinterpret_cast(m), + reinterpret_cast(q), dim, eta_, out); + }; + + default: + return nullptr; + } + } + + if (injection_ == Injection::kIdentity) { + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + return reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_FP16: + return reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_INT8: + return reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute); + + case IndexMeta::DataType::DT_INT4: + return reinterpret_cast( + ailego::SquaredEuclideanDistanceMatrix::Compute); + + default: + return nullptr; + } + } + return nullptr; + } + + //! Retrieve distance function for query + MatrixSparseDistance sparse_distance(void) const override { + if (injection_ == Injection::kLocalizedSpherical) { + return [&](const void *m_sparse, const void *q_sparse, float *out) { + ailego::MipsSquaredEuclideanSparseDistanceMatrix::Compute( + m_sparse, q_sparse, out); + }; + } + + if (injection_ == Injection::kRepeatedQuadratic) { + LOG_ERROR( + "Repeated Quadratic is not supported in MipsEuclideanMetric for " + "Hybrid Vector!"); + + return nullptr; + } + + if (injection_ == Injection::kSpherical) { + LOG_ERROR( + "Spherical is not supported in MipsEuclideanMetric for Hybrid " + "Vector!"); + + return nullptr; + } + + if (injection_ == Injection::kIdentity) { + LOG_ERROR( + "Identity is not supported in MipsEuclideanMetric for Hybrid " + "Vector!"); + + return nullptr; + } + + return nullptr; + } + + //! Retrieve matrix distance function for index features + MatrixDistance distance_matrix(size_t m, size_t n) const override { + if (injection_ == Injection::kLocalizedSpherical) { + SphericalHandle compute; + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + compute = + DistanceMatrixCompute(m, n); + break; + case IndexMeta::DataType::DT_FP16: + compute = + DistanceMatrixCompute(m, n); + break; + case IndexMeta::DataType::DT_INT8: + compute = + DistanceMatrixCompute(m, n); + break; + case IndexMeta::DataType::DT_INT4: + compute = + DistanceMatrixCompute(m, n); + break; + default: + return nullptr; + } + return [=](const void *d, const void *q, size_t dim, float *out) { + compute(d, q, dim, 0.0f, out); + }; + } + + if (injection_ == Injection::kRepeatedQuadratic) { + RepeatedQuadraticHandle compute; + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + compute = + DistanceMatrixCompute(m, n); + break; + case IndexMeta::DataType::DT_FP16: + compute = + DistanceMatrixCompute( + m, n); + break; + case IndexMeta::DataType::DT_INT8: + compute = + DistanceMatrixCompute(m, n); + break; + case IndexMeta::DataType::DT_INT4: + compute = + DistanceMatrixCompute(m, n); + break; + default: + return nullptr; + } + return [=](const void *d, const void *q, size_t dim, float *out) { + compute(d, q, dim, m_value_, eta_, out); + }; + } + + if (injection_ == Injection::kSpherical) { + SphericalHandle compute; + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + compute = + DistanceMatrixCompute(m, n); + break; + case IndexMeta::DataType::DT_FP16: + compute = + DistanceMatrixCompute(m, n); + break; + case IndexMeta::DataType::DT_INT8: + compute = + DistanceMatrixCompute(m, n); + break; + case IndexMeta::DataType::DT_INT4: + compute = + DistanceMatrixCompute(m, n); + break; + default: + return nullptr; + } + return [=](const void *d, const void *q, size_t dim, float *out) { + compute(d, q, dim, eta_, out); + }; + } + + if (injection_ == Injection::kIdentity) { + switch (data_type_) { + case IndexMeta::DataType::DT_FP32: + return DistanceMatrixCompute(m, n); + case IndexMeta::DataType::DT_FP16: + return DistanceMatrixCompute(m, + n); + case IndexMeta::DataType::DT_INT8: + return DistanceMatrixCompute(m, n); + case IndexMeta::DataType::DT_INT4: + return DistanceMatrixCompute(m, n); + default: + return nullptr; + } + } + return nullptr; + } + + //! Normalize result + void normalize(float *score) const override { + query_metric_->normalize(score); + } + + //! Denormalize threshold + void denormalize(float *score) const override { + query_metric_->denormalize(score); + } + + //! Retrieve if it supports normalization + bool support_normalize(void) const override { + return query_metric_->support_normalize(); + } + + //! Retrieve params of Metric + const ailego::Params ¶ms(void) const override { + return params_; + } + + //! Train the metric + int train(const void *vec, size_t dim) override { + if (eta_ == 0.0f) { // No global norm scaling => no training. + return 0; + } + if (!squared_norm2_handle_) { + return IndexError_Unsupported; + } + + float score; + squared_norm2_handle_(vec, dim, &score); + if (score > max_squared_l2_norm_) { + max_squared_l2_norm_ = score; + const float max_l2_norm = std::sqrt(score); + params_.set(MIPS_EUCLIDEAN_METRIC_MAX_L2_NORM, max_l2_norm); + if (max_squared_l2_norm_ < 1.0 && + max_squared_l2_norm_ > squared_u_value_) { + squared_u_value_ = max_squared_l2_norm_; + params_.set(MIPS_EUCLIDEAN_METRIC_U_VALUE, max_l2_norm); + } + eta_ = squared_u_value_ / max_squared_l2_norm_; + } + return 0; + } + + //! Retrieve if it supports training + bool support_train(void) const override { + // No global norm scaling => eta_ == 0 => no training. + return eta_ != 0.0f; + } + + //! Retrieve query metric object of this index metric + Pointer query_metric(void) const override { + return query_metric_; + } + + private: + //! Type of MipsSquaredEuclideanDistanceMatrix::Compute overloaded for + // Spherical injection and LocalizedSpherical nonmetric. + template + using SphericalHandle = void (*)(const T *m, const T *q, size_t dim, + float eta, float *out); + + //! Type of MipsSquaredEuclideanDistanceMatrix::Compute overloaded for + // RepeatedQuadratic injection. + template + using RepeatedQuadraticHandle = void (*)(const T *m, const T *q, size_t dim, + size_t m_value, float eta, + float *out); + + //! Type of squared L2 norm function. + using SquaredNorm2Handle = void (*)(const void *m, size_t dim, float *out); + + enum struct Injection { // Type of injective mapping into Euclidean space. + kLocalizedSpherical = 0, // spherical with pair-only max-norm + kSpherical = 1, // require global scaling/training + kRepeatedQuadratic = 2, // require global scaling/training + kIdentity = 3, // plain Euclidean distance + kNumInjections = 4 + }; + + static const char *InjectionName(int injection) { + static const char *injection_names[] = {"LocalizedSpherical", "Spherical", + "RepeatedQuadratic", "Identity"}; + if (injection >= 0 && + injection < static_cast(Injection::kNumInjections)) { + return injection_names[injection]; + } + return "Invalid"; + } + + static const char *InjectionName(Injection injection) { + return InjectionName(static_cast(injection)); + } + + // Checks (and fixes) `*m_value`, no. additional dimensions for injection. + // `dim` is the original dimension, used ONLY by RepeatedQuadratic + // injection, where dim = 1 induces the default *m_value = 3. It's + // positioned last to allow other injections to skip it. + // Returns true if `*m_value` is modified. + static bool CheckAndFixM(Injection injection, uint32_t *m_value) { + if (injection == Injection::kRepeatedQuadratic) { + if (*m_value == 0) { + *m_value = 3u; // Recommend value in paper (3.5 Practical + // Recommendation of Parameters) + return true; + } + } else if (injection == Injection::kSpherical) { + if (*m_value != 1) { + if (*m_value != 0) { + LOG_WARN("M value (%u) set to 1 for Spherical injection", *m_value); + } + *m_value = 1; + return true; + } + } else { // kLocalizedSpherical, kIdentity, or kInvalid + if (*m_value != 0) { + LOG_WARN("M value (%u) set to 0 for %s injections", *m_value, + InjectionName(injection)); + *m_value = 0; + return true; + } + } + return false; + } + + // Checks and fixes `*u_value`, global L2 norm scalar. + // `m_value` is no. additional dimensions, used ONLY by RepeatedQuadratic + // injection. It's positioned last to allow other injections to skip it. + // Returns true if `*u_value` is set to a new value. + static bool CheckAndFixU(Injection injection, uint32_t m_value, + float *u_value) { + if (injection == Injection::kRepeatedQuadratic) { + if (*u_value <= std::numeric_limits::epsilon() || + *u_value >= 1.0) { + // Try computing a default U value + constexpr float kLogError = -5.0; // log_10(distance_error) + float new_u_value = std::pow(10, kLogError / (1 << (m_value + 1))); + if (*u_value != 0) { + LOG_WARN("U value (%f) set to %f for RepeatedQuadratic injection", + *u_value, new_u_value); + } + *u_value = new_u_value; + return true; + } else if (std::pow(*u_value, (1 << m_value)) < + std::numeric_limits::epsilon()) { + LOG_WARN( + "U value %f is too small, may cause loss of distance precision", + *u_value); + } + } else if (injection == Injection::kSpherical) { + // Spherical injection requires ||x'|| <= 1.0 for computing + // std::sqrt(1 - ||x'||^2), x' = u_value * x / max_norm. Set u_value + // to slightly < 1.0 in case of precision loss in float computation. + if (*u_value <= std::numeric_limits::epsilon() || + *u_value >= 1.0) { + static constexpr float kSphericalUValue = 1.0f - 1e-3; + if (*u_value != 0.0f) { + LOG_WARN("U value (%f) set to %f for Spherical injection", *u_value, + kSphericalUValue); + } + *u_value = kSphericalUValue; + return true; + } + } else { // kLocalizedSpherical, kIdentity, or kInvalid + if (*u_value != 1.0) { + if (*u_value != 0) { + LOG_WARN("U value (%f) set to 1.0 for %s injection", *u_value, + InjectionName(injection)); + } + *u_value = 1.0; + return true; + } + } + return false; + } + + private: + //! Type of basic DistanceMatrix::Compute function with typed parameter. + template + using TypedDistanceHandle = void (*)(const T *m, const T *q, size_t dim, + float *out); + + //! Returns m x n distance matrix compute function. + // Handle is used to resolve potential DistanceMatrix::Compute overload. + template