From 6524fabd802d1ed50e2c7889e3814ff027eef9e8 Mon Sep 17 00:00:00 2001 From: sanyi Date: Tue, 30 Dec 2025 11:02:17 +0800 Subject: [PATCH] Initial commit --- .clang-format | 31 + .github/ISSUE_TEMPLATE/bug_report.yml | 73 + .github/ISSUE_TEMPLATE/config.yml | 9 + .github/ISSUE_TEMPLATE/feature_request.yml | 55 + .github/codecov.yml | 32 + .github/workflows/build_wheel.yml | 57 + .../docker/Dockerfile.ubuntu18.10-glibc228 | 83 + .github/workflows/linux_x64_docker_ci.yml | 124 + .github/workflows/mac_arm64_ci.yml | 85 + .github/workflows/nightly_coverage.yml | 95 + .gitignore | 50 + .gitmodules | 41 + .pre-commit-config.yaml | 30 + CMakeLists.txt | 60 + CODE_OF_CONDUCT.md | 41 + CONTRIBUTING.md | 104 + LICENSE | 201 + README.md | 99 + cmake/bazel.cmake | 2338 +++++++++ cmake/option.cmake | 585 +++ cmake/utils.cmake | 28 + pyproject.toml | 242 + python/tests/detail/distance_helper.py | 302 ++ python/tests/detail/doc_helper.py | 377 ++ python/tests/detail/fixture_helper.py | 319 ++ 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 | 792 +++ python/tests/detail/test_collection_ddl.py | 1557 ++++++ python/tests/detail/test_collection_dml.py | 1180 +++++ python/tests/detail/test_collection_dql.py | 1329 +++++ .../tests/detail/test_collection_exception.py | 328 ++ python/tests/detail/test_collection_open.py | 960 ++++ python/tests/detail/test_db_config.py | 307 ++ python/tests/test_collection.py | 1045 ++++ python/tests/test_convert.py | 584 +++ python/tests/test_doc.py | 269 + python/tests/test_embedding.py | 131 + python/tests/test_params.py | 415 ++ python/tests/test_query_executor.py | 297 ++ python/tests/test_reranker.py | 184 + python/tests/test_schema.py | 246 + python/tests/test_typing.py | 141 + python/tests/test_util.py | 94 + python/zvec/__init__.py | 133 + python/zvec/__init__.pyi | 172 + python/zvec/common/__init__.py | 18 + python/zvec/common/constants.py | 22 + python/zvec/executor/__init__.py | 26 + python/zvec/executor/query_executor.py | 307 ++ python/zvec/extension/__init__.py | 26 + python/zvec/extension/embedding.py | 188 + python/zvec/extension/rerank.py | 343 ++ python/zvec/model/__init__.py | 22 + python/zvec/model/collection.py | 381 ++ python/zvec/model/convert.py | 54 + python/zvec/model/doc.py | 173 + python/zvec/model/param/__init__.py | 42 + python/zvec/model/param/__init__.pyi | 692 +++ python/zvec/model/param/vector_query.py | 80 + 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 | 300 ++ 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 | 226 + scripts/README.md | 0 scripts/gcov.sh | 41 + src/CMakeLists.txt | 13 + src/ailego/CMakeLists.txt | 28 + 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 | 1218 +++++ src/ailego/algorithm/lloyd_cluster.h | 363 ++ src/ailego/buffer/buffer_manager.cc | 579 +++ src/ailego/buffer/buffer_manager.h | 269 + src/ailego/container/bitmap.cc | 443 ++ src/ailego/container/bitmap.h | 503 ++ src/ailego/container/blob.h | 131 + src/ailego/container/bloom_filter.h | 279 ++ src/ailego/container/cube.h | 414 ++ src/ailego/container/heap.h | 234 + src/ailego/container/hypercube.h | 268 + src/ailego/container/params.cc | 139 + src/ailego/container/params.h | 776 +++ src/ailego/container/reservoir.h | 140 + src/ailego/container/vector.h | 1055 ++++ src/ailego/container/vector_array.h | 555 +++ src/ailego/encoding/json.h | 17 + src/ailego/encoding/json/mod_json.c | 3591 ++++++++++++++ src/ailego/encoding/json/mod_json.h | 1382 ++++++ src/ailego/encoding/json/mod_json_plus.h | 3446 +++++++++++++ src/ailego/hash/crc32c.cc | 636 +++ src/ailego/hash/crc32c.h | 35 + src/ailego/hash/jump_hash.h | 35 + src/ailego/internal/cpu_features.cc | 421 ++ src/ailego/internal/cpu_features.h | 208 + src/ailego/internal/platform.h | 488 ++ src/ailego/io/file.cc | 726 +++ src/ailego/io/file.h | 309 ++ src/ailego/io/file_lock.cc | 80 + src/ailego/io/file_lock.h | 84 + src/ailego/io/file_writer.h | 127 + src/ailego/io/mmap_file.h | 255 + src/ailego/io/pid_file.cc | 65 + src/ailego/io/pid_file.h | 64 + src/ailego/logger/logger.cc | 76 + src/ailego/logger/logger.h | 175 + src/ailego/math/cosine_distance_matrix.h | 94 + src/ailego/math/distance.h | 280 ++ src/ailego/math/distance_matrix.h | 21 + 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 +++++ src/ailego/math/distance_matrix_fp16.i | 861 ++++ src/ailego/math/distance_matrix_fp32.i | 1599 ++++++ src/ailego/math/distance_matrix_int32.i | 859 ++++ src/ailego/math/distance_matrix_int64.i | 472 ++ src/ailego/math/distance_matrix_popcnt.i | 2565 ++++++++++ src/ailego/math/distance_utility.h | 65 + src/ailego/math/euclidean_distance_matrix.h | 2639 ++++++++++ .../math/euclidean_distance_matrix_fp16.cc | 551 +++ .../math/euclidean_distance_matrix_fp32.cc | 920 ++++ .../math/euclidean_distance_matrix_int4.cc | 801 +++ .../math/euclidean_distance_matrix_int8.cc | 884 ++++ src/ailego/math/hamming_distance_matrix.cc | 1128 +++++ src/ailego/math/hamming_distance_matrix.h | 1158 +++++ src/ailego/math/inner_product_matrix.h | 2991 ++++++++++++ src/ailego/math/inner_product_matrix_fp16.cc | 1931 ++++++++ src/ailego/math/inner_product_matrix_fp32.cc | 1192 +++++ src/ailego/math/inner_product_matrix_int4.cc | 802 +++ src/ailego/math/inner_product_matrix_int8.cc | 841 ++++ src/ailego/math/matrix_define.i | 1850 +++++++ src/ailego/math/matrix_utility.i | 265 + .../math/mips_euclidean_distance_matrix.h | 1191 +++++ .../mips_euclidean_distance_matrix_fp16.cc | 409 ++ .../mips_euclidean_distance_matrix_fp32.cc | 684 +++ .../mips_euclidean_distance_matrix_int4.cc | 358 ++ .../mips_euclidean_distance_matrix_int8.cc | 358 ++ src/ailego/math/norm1_matrix.h | 148 + src/ailego/math/norm1_matrix_fp16.cc | 84 + src/ailego/math/norm1_matrix_fp32.cc | 77 + src/ailego/math/norm2_matrix.h | 424 ++ src/ailego/math/norm2_matrix_fp16.cc | 81 + src/ailego/math/norm2_matrix_fp32.cc | 76 + 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 | 57 + src/ailego/math_batch/distance_batch.h | 51 + .../math_batch/inner_product_distance_batch.h | 151 + .../inner_product_distance_batch_impl.h | 129 + .../inner_product_distance_batch_impl_fp16.h | 240 + .../inner_product_distance_batch_impl_int8.h | 225 + src/ailego/math_batch/utils.h | 24 + 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/parallel/thread_pool.h | 402 ++ src/ailego/parallel/thread_queue.h | 291 ++ src/ailego/pattern/closure.h | 530 ++ src/ailego/pattern/defer.h | 24 + src/ailego/pattern/expected.hpp | 2605 ++++++++++ src/ailego/pattern/factory.h | 172 + src/ailego/pattern/scope_guard.h | 177 + src/ailego/pattern/singleton.h | 51 + src/ailego/string/string_concat_helper.h | 120 + src/ailego/string/string_view.h | 64 + src/ailego/utility/bit_string_helper.h | 118 + src/ailego/utility/bitset_helper.cc | 2048 ++++++++ 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 | 386 ++ src/ailego/utility/file_helper.h | 98 + src/ailego/utility/float_helper.cc | 1158 +++++ src/ailego/utility/float_helper.h | 236 + src/ailego/utility/math_helper.h | 109 + src/ailego/utility/matrix_helper.h | 70 + src/ailego/utility/memory_helper.cc | 397 ++ src/ailego/utility/memory_helper.h | 54 + src/ailego/utility/process_helper.cc | 322 ++ src/ailego/utility/process_helper.h | 54 + src/ailego/utility/string_helper.cc | 123 + src/ailego/utility/string_helper.h | 322 ++ src/ailego/utility/string_helper_impl.h | 173 + src/ailego/utility/time_helper.cc | 215 + src/ailego/utility/time_helper.h | 206 + src/ailego/utility/type_helper.h | 128 + src/ailego/version.cc | 48 + src/ailego/version.h | 33 + src/ailego/version.i | 293 ++ src/binding/CMakeLists.txt | 8 + src/binding/python/CMakeLists.txt | 59 + src/binding/python/binding.cc | 33 + src/binding/python/exports.mac | 1 + src/binding/python/exports.map | 6 + .../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_schema.h | 35 + src/binding/python/include/python_type.h | 38 + .../python/model/common/python_config.cc | 190 + .../python/model/param/python_param.cc | 1222 +++++ src/binding/python/model/python_collection.cc | 218 + src/binding/python/model/python_doc.cc | 452 ++ .../python/model/schema/python_schema.cc | 162 + src/binding/python/typing/python_type.cc | 226 + src/core/CMakeLists.txt | 10 + src/core/algorithm/CMakeLists.txt | 9 + src/core/algorithm/cluster/CMakeLists.txt | 10 + src/core/algorithm/cluster/cluster_params.h | 207 + src/core/algorithm/cluster/kmeans_cluster.cc | 915 ++++ src/core/algorithm/cluster/linear_seeker.cc | 45 + src/core/algorithm/cluster/linear_seeker.h | 98 + .../algorithm/cluster/opt_kmeans_cluster.cc | 1346 +++++ 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/diskann_builder.h | 13 + src/core/algorithm/diskann/diskann_searcher.h | 13 + src/core/algorithm/flat/CMakeLists.txt | 11 + src/core/algorithm/flat/flat_builder.cc | 284 ++ src/core/algorithm/flat/flat_builder.h | 101 + .../algorithm/flat/flat_distance_matrix.h | 184 + src/core/algorithm/flat/flat_index_format.h | 184 + src/core/algorithm/flat/flat_searcher.cc | 252 + src/core/algorithm/flat/flat_searcher.h | 182 + .../algorithm/flat/flat_searcher_context.h | 1107 +++++ .../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 | 244 + .../algorithm/flat/flat_streamer_dumper.h | 243 + .../algorithm/flat/flat_streamer_entity.cc | 994 ++++ .../algorithm/flat/flat_streamer_entity.h | 407 ++ .../algorithm/flat/flat_streamer_provider.h | 100 + src/core/algorithm/flat/flat_utility.h | 201 + src/core/algorithm/flat_sparse/CMakeLists.txt | 11 + .../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 | 221 + .../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 | 151 + .../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 | 407 ++ .../flat_sparse/flat_sparse_streamer.h | 172 + .../flat_sparse_streamer_entity.cc | 1082 +++++ .../flat_sparse/flat_sparse_streamer_entity.h | 265 + .../flat_sparse/flat_sparse_utility.h | 58 + src/core/algorithm/hnsw/CMakeLists.txt | 11 + src/core/algorithm/hnsw/hnsw_algorithm.cc | 540 ++ src/core/algorithm/hnsw/hnsw_algorithm.h | 134 + src/core/algorithm/hnsw/hnsw_builder.cc | 404 ++ src/core/algorithm/hnsw/hnsw_builder.h | 93 + .../algorithm/hnsw/hnsw_builder_entity.cc | 198 + src/core/algorithm/hnsw/hnsw_builder_entity.h | 138 + src/core/algorithm/hnsw/hnsw_chunk.cc | 220 + src/core/algorithm/hnsw/hnsw_chunk.h | 140 + src/core/algorithm/hnsw/hnsw_context.cc | 296 ++ src/core/algorithm/hnsw/hnsw_context.h | 530 ++ .../algorithm/hnsw/hnsw_dist_calculator.h | 218 + src/core/algorithm/hnsw/hnsw_entity.cc | 362 ++ src/core/algorithm/hnsw/hnsw_entity.h | 527 ++ src/core/algorithm/hnsw/hnsw_index_hash.h | 231 + src/core/algorithm/hnsw/hnsw_index_provider.h | 133 + src/core/algorithm/hnsw/hnsw_params.h | 124 + src/core/algorithm/hnsw/hnsw_searcher.cc | 460 ++ src/core/algorithm/hnsw/hnsw_searcher.h | 139 + .../algorithm/hnsw/hnsw_searcher_entity.cc | 515 ++ .../algorithm/hnsw/hnsw_searcher_entity.h | 156 + src/core/algorithm/hnsw/hnsw_streamer.cc | 848 ++++ src/core/algorithm/hnsw/hnsw_streamer.h | 228 + .../algorithm/hnsw/hnsw_streamer_entity.cc | 701 +++ .../algorithm/hnsw/hnsw_streamer_entity.h | 525 ++ src/core/algorithm/hnsw_sparse/CMakeLists.txt | 11 + .../hnsw_sparse/hnsw_sparse_algorithm.cc | 544 +++ .../hnsw_sparse/hnsw_sparse_algorithm.h | 131 + .../hnsw_sparse/hnsw_sparse_builder.cc | 521 ++ .../hnsw_sparse/hnsw_sparse_builder.h | 102 + .../hnsw_sparse/hnsw_sparse_builder_entity.cc | 289 ++ .../hnsw_sparse/hnsw_sparse_builder_entity.h | 161 + .../hnsw_sparse/hnsw_sparse_chunk.cc | 217 + .../algorithm/hnsw_sparse/hnsw_sparse_chunk.h | 143 + .../hnsw_sparse/hnsw_sparse_context.cc | 294 ++ .../hnsw_sparse/hnsw_sparse_context.h | 522 ++ .../hnsw_sparse/hnsw_sparse_dist_calculator.h | 184 + .../hnsw_sparse/hnsw_sparse_entity.cc | 440 ++ .../hnsw_sparse/hnsw_sparse_entity.h | 705 +++ .../hnsw_sparse/hnsw_sparse_index_hash.h | 232 + .../hnsw_sparse/hnsw_sparse_index_provider.h | 168 + .../hnsw_sparse/hnsw_sparse_params.h | 116 + .../hnsw_sparse/hnsw_sparse_searcher.cc | 564 +++ .../hnsw_sparse/hnsw_sparse_searcher.h | 158 + .../hnsw_sparse_searcher_entity.cc | 630 +++ .../hnsw_sparse/hnsw_sparse_searcher_entity.h | 184 + .../hnsw_sparse/hnsw_sparse_streamer.cc | 952 ++++ .../hnsw_sparse/hnsw_sparse_streamer.h | 225 + .../hnsw_sparse_streamer_entity.cc | 1001 ++++ .../hnsw_sparse/hnsw_sparse_streamer_entity.h | 542 +++ src/core/algorithm/ivf/CMakeLists.txt | 10 + src/core/algorithm/ivf/ivf_builder.cc | 829 ++++ src/core/algorithm/ivf/ivf_builder.h | 316 ++ src/core/algorithm/ivf/ivf_centroid_index.cc | 621 +++ 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 | 973 ++++ src/core/algorithm/ivf/ivf_entity.h | 351 ++ src/core/algorithm/ivf/ivf_index_format.h | 167 + src/core/algorithm/ivf/ivf_index_provider.h | 110 + src/core/algorithm/ivf/ivf_params.h | 88 + src/core/algorithm/ivf/ivf_searcher.cc | 317 ++ src/core/algorithm/ivf/ivf_searcher.h | 114 + src/core/algorithm/ivf/ivf_searcher_context.h | 223 + src/core/algorithm/ivf/ivf_streamer.cc | 317 ++ src/core/algorithm/ivf/ivf_streamer.h | 125 + src/core/algorithm/ivf/ivf_utility.h | 177 + src/core/framework/CMakeLists.txt | 11 + src/core/framework/index_builder.h | 61 + src/core/framework/index_bundle.h | 274 ++ src/core/framework/index_cluster.cc | 158 + src/core/framework/index_cluster.h | 291 ++ src/core/framework/index_context.cc | 27 + src/core/framework/index_context.h | 265 + src/core/framework/index_converter.cc | 42 + src/core/framework/index_converter.h | 225 + src/core/framework/index_document.h | 317 ++ src/core/framework/index_dumper.h | 163 + src/core/framework/index_error.cc | 73 + src/core/framework/index_error.h | 170 + src/core/framework/index_factory.cc | 260 + src/core/framework/index_factory.h | 286 ++ src/core/framework/index_features.h | 604 +++ src/core/framework/index_filter.h | 74 + src/core/framework/index_flow.cc | 879 ++++ src/core/framework/index_flow.h | 646 +++ src/core/framework/index_format.h | 176 + src/core/framework/index_framework.h | 31 + src/core/framework/index_groupby.h | 54 + src/core/framework/index_helper.cc | 254 + src/core/framework/index_helper.h | 59 + src/core/framework/index_holder.h | 1760 +++++++ src/core/framework/index_logger.cc | 72 + src/core/framework/index_logger.h | 168 + src/core/framework/index_mapping.cc | 504 ++ src/core/framework/index_mapping.h | 206 + src/core/framework/index_memory.h | 266 + src/core/framework/index_meta.cc | 208 + src/core/framework/index_meta.h | 723 +++ src/core/framework/index_metric.h | 143 + src/core/framework/index_module.h | 67 + src/core/framework/index_packer.h | 221 + src/core/framework/index_plugin.cc | 59 + src/core/framework/index_plugin.h | 116 + src/core/framework/index_provider.h | 74 + src/core/framework/index_reducer.h | 229 + src/core/framework/index_refiner.h | 92 + src/core/framework/index_reformer.h | 134 + src/core/framework/index_runner.h | 743 +++ src/core/framework/index_searcher.h | 57 + src/core/framework/index_segment_storage.h | 236 + src/core/framework/index_stats.h | 68 + src/core/framework/index_storage.h | 271 ++ src/core/framework/index_streamer.h | 56 + src/core/framework/index_threads.h | 172 + src/core/framework/index_trainer.h | 111 + src/core/framework/index_unpacker.h | 304 ++ src/core/framework/index_version.cc | 33 + src/core/framework/index_version.h | 30 + src/core/interface/CMakeLists.txt | 10 + src/core/interface/constants.h | 13 + src/core/interface/index.cc | 773 +++ src/core/interface/index.h | 321 ++ src/core/interface/index_factory.cc | 116 + src/core/interface/index_factory.h | 225 + src/core/interface/index_param.cc | 206 + src/core/interface/index_param.h | 340 ++ src/core/interface/index_param_builders.h | 341 ++ src/core/interface/indexes/flat_index.cc | 74 + src/core/interface/indexes/hnsw_index.cc | 121 + src/core/interface/indexes/ivf_index.cc | 292 ++ src/core/interface/utils/utils.h | 103 + src/core/metric/CMakeLists.txt | 11 + src/core/metric/cosine_metric.cc | 235 + src/core/metric/euclidean_metric.cc | 1098 +++++ src/core/metric/hamming_metric.cc | 237 + src/core/metric/inner_product_metric.cc | 534 ++ src/core/metric/metric_params.h | 38 + src/core/metric/mips_euclidean_metric.cc | 683 +++ src/core/metric/quantized_integer_metric.cc | 323 ++ .../metric/quantized_integer_metric_batch.h | 354 ++ .../metric/quantized_integer_metric_matrix.h | 254 + src/core/mixed_reducer/CMakeLists.txt | 10 + src/core/mixed_reducer/mixed_reducer_params.h | 42 + .../mixed_reducer/mixed_streamer_reducer.cc | 579 +++ .../mixed_reducer/mixed_streamer_reducer.h | 115 + src/core/quantizer/CMakeLists.txt | 11 + src/core/quantizer/binary_converter.cc | 231 + src/core/quantizer/binary_reformer.cc | 108 + src/core/quantizer/cosine_converter.cc | 405 ++ src/core/quantizer/cosine_reformer.cc | 287 ++ src/core/quantizer/half_float_converter.cc | 370 ++ src/core/quantizer/half_float_reformer.cc | 238 + .../quantizer/integer_quantizer_converter.cc | 612 +++ .../quantizer/integer_quantizer_reformer.cc | 486 ++ src/core/quantizer/mips_converter.cc | 660 +++ src/core/quantizer/mips_reformer.cc | 463 ++ src/core/quantizer/quantizer_params.h | 133 + src/core/quantizer/record_quantizer.h | 131 + src/core/utility/CMakeLists.txt | 11 + src/core/utility/basic_refiner.cc | 277 ++ src/core/utility/buffer_storage.cc | 442 ++ src/core/utility/file_dumper.cc | 130 + src/core/utility/file_read_storage.cc | 503 ++ src/core/utility/memory_dumper.cc | 115 + src/core/utility/memory_read_storage.cc | 249 + src/core/utility/mmap_file_read_storage.cc | 284 ++ src/core/utility/mmap_file_storage.cc | 344 ++ src/core/utility/sparse_utility.h | 419 ++ src/core/utility/utility_params.h | 248 + src/core/utility/visit_filter.h | 427 ++ src/db/CMakeLists.txt | 22 + src/db/collection.cc | 1887 +++++++ src/db/collection.h | 111 + src/db/common/CMakeLists.txt | 14 + src/db/common/cgroup_util.cc | 467 ++ src/db/common/cgroup_util.h | 101 + src/db/common/concurrent_roaring_bitmap.cc | 211 + src/db/common/concurrent_roaring_bitmap.h | 321 ++ src/db/common/config.cc | 130 + src/db/common/config.h | 172 + src/db/common/constants.h | 105 + src/db/common/error_code.cc | 155 + src/db/common/error_code.h | 212 + src/db/common/file_helper.cc | 128 + src/db/common/file_helper.h | 260 + 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 | 75 + src/db/common/profiler.h | 201 + src/db/common/rocbsdb_context.cc | 595 +++ src/db/common/rocksdb_context.h | 121 + src/db/common/status.cc | 56 + src/db/common/status.h | 181 + src/db/common/typedef.h | 201 + src/db/common/utils.cc | 23 + src/db/common/utils.h | 19 + src/db/index/CMakeLists.txt | 16 + src/db/index/column/column_indexer.h | 25 + src/db/index/column/common/index_results.h | 68 + .../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 | 267 + .../column/inverted_column/inverted_indexer.h | 119 + .../inverted_column/inverted_rocksdb_merger.h | 115 + .../inverted_column/inverted_search_result.h | 199 + .../combined_vector_column_indexer.cc | 227 + .../combined_vector_column_indexer.h | 60 + .../column/vector_column/engine_helper.hpp | 349 ++ .../vector_column/vector_column_indexer.cc | 207 + .../vector_column/vector_column_indexer.h | 131 + .../vector_column/vector_column_params.h | 113 + .../vector_column/vector_index_results.h | 282 ++ src/db/index/common/delete_store.h | 152 + src/db/index/common/doc.cc | 1230 +++++ src/db/index/common/doc.h | 381 ++ 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.h | 332 ++ src/db/index/common/meta.h | 434 ++ src/db/index/common/options.h | 61 + src/db/index/common/proto_converter.cc | 282 ++ src/db/index/common/proto_converter.h | 62 + src/db/index/common/query_params.h | 150 + src/db/index/common/schema.cc | 461 ++ src/db/index/common/schema.h | 463 ++ src/db/index/common/stats.h | 70 + src/db/index/common/type.h | 140 + src/db/index/common/type_helper.cc | 49 + src/db/index/common/type_helper.h | 510 ++ src/db/index/common/version_manager.cc | 301 ++ 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 | 4322 +++++++++++++++++ src/db/index/segment/segment.h | 183 + src/db/index/segment/segment_helper.cc | 830 ++++ src/db/index/segment/segment_helper.h | 226 + src/db/index/segment/segment_manager.cc | 160 + src/db/index/segment/segment_manager.h | 56 + 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 | 381 ++ .../index/storage/bufferpool_forward_store.h | 131 + 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 | 185 + 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_writer.cc | 185 + src/db/index/storage/parquet_writer.h | 55 + src/db/index/storage/store_helper.h | 954 ++++ src/db/index/storage/wal/local_wal_file.cc | 245 + 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 | 68 + src/db/proto/zvec.proto | 174 + src/db/sqlengine/CMakeLists.txt | 25 + src/db/sqlengine/analyzer/query_analyzer.cc | 552 +++ 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 | 116 + src/db/sqlengine/analyzer/query_info.h | 375 ++ .../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 | 369 ++ .../sqlengine/analyzer/query_node_walker.cc | 696 +++ 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/generic_node.h | 89 + src/db/sqlengine/common/group_by.h | 46 + src/db/sqlengine/common/util.cc | 91 + 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 | 279 ++ 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 | 202 + src/db/sqlengine/parser/select_info.h | 87 + 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 | 112 + src/db/sqlengine/parser/sql_info_helper.h | 33 + .../parser/zvec_cached_sql_parser.cc | 234 + .../sqlengine/parser/zvec_cached_sql_parser.h | 57 + src/db/sqlengine/parser/zvec_parser.cc | 155 + 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 | 166 + src/db/sqlengine/planner/doc_filter.h | 71 + .../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 | 626 +++ src/db/sqlengine/planner/query_planner.h | 69 + src/db/sqlengine/planner/segment_node.cc | 85 + src/db/sqlengine/planner/segment_node.h | 55 + .../sqlengine/planner/vector_recall_node.cc | 286 ++ 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 | 574 +++ src/db/sqlengine/sqlengine_impl.h | 75 + tests/CMakeLists.txt | 6 + tests/ailego/CMakeLists.txt | 17 + .../algorithm/integer_quantizer_test.cc | 564 +++ tests/ailego/algorithm/kmeans_test.cc | 417 ++ tests/ailego/buffer/buffer_manager_test.cc | 132 + 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 | 178 + tests/ailego/io/file_lock_test.cc | 43 + tests/ailego/io/file_test.cc | 338 ++ tests/ailego/io/mmap_file_test.cc | 161 + tests/ailego/io/pid_file_test.cc | 35 + 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/hamming_distance_matrix_test.cc | 1084 +++++ .../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 | 984 ++++ ...ips_euclidean_distance_matrix_fp32_test.cc | 957 ++++ ...ips_euclidean_distance_matrix_int4_test.cc | 995 ++++ ...ips_euclidean_distance_matrix_int8_test.cc | 962 ++++ 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 | 150 + .../ailego/parallel/multi_thread_list_test.cc | 260 + tests/ailego/parallel/semaphore_test.cc | 103 + tests/ailego/parallel/thread_pool_test.cc | 162 + tests/ailego/parallel/thread_queue_test.cc | 119 + 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/process_helper_test.cc | 34 + tests/ailego/utility/string_helper_test.cc | 479 ++ tests/ailego/utility/time_helper_test.cc | 79 + tests/ailego/utility/type_helper_test.cc | 106 + tests/ailego/version_test.cc | 23 + tests/core/CMakeLists.txt | 9 + tests/core/algorithm/CMakeLists.txt | 9 + tests/core/algorithm/cluster/CMakeLists.txt | 14 + .../algorithm/cluster/kmeans_cluster_test.cc | 106 + .../cluster/opt_kmeans_cluster_test.cc | 602 +++ tests/core/algorithm/flat/CMakeLists.txt | 14 + .../core/algorithm/flat/flat_builder_test.cc | 457 ++ .../algorithm/flat/flat_searcher_test.cpp | 1492 ++++++ .../flat/flat_streamer_buffer_test.cpp | 362 ++ .../flat/flat_streamer_buffer_time_test.cpp | 140 + .../core/algorithm/flat/flat_streamer_test.cc | 1214 +++++ .../core/algorithm/flat_sparse/CMakeLists.txt | 14 + .../flat_sparse/flat_sparse_builder_test.cc | 301 ++ .../flat_sparse/flat_sparse_searcher_test.cc | 822 ++++ .../flat_sparse_streamer_buffer_test.cc | 232 + .../flat_sparse/flat_sparse_streamer_test.cc | 1729 +++++++ tests/core/algorithm/hnsw/CMakeLists.txt | 14 + .../core/algorithm/hnsw/hnsw_builder_test.cc | 543 +++ .../algorithm/hnsw/hnsw_searcher_test.cpp | 2775 +++++++++++ .../hnsw/hnsw_streamer_buffer_test.cpp | 271 ++ .../core/algorithm/hnsw/hnsw_streamer_test.cc | 3796 +++++++++++++++ .../core/algorithm/hnsw_sparse/CMakeLists.txt | 14 + .../hnsw_sparse/hnsw_sparse_builder_test.cc | 470 ++ .../hnsw_sparse/hnsw_sparse_searcher_test.cc | 1781 +++++++ .../hnsw_sparse_streamer_buffer_test.cpp | 361 ++ .../hnsw_sparse/hnsw_sparse_streamer_test.cc | 2487 ++++++++++ tests/core/algorithm/ivf/CMakeLists.txt | 14 + tests/core/algorithm/ivf/ivf_builder_test.cc | 464 ++ tests/core/algorithm/ivf/ivf_searcher_test.cc | 3423 +++++++++++++ tests/core/framework/CMakeLists.txt | 14 + tests/core/interface/CMakeLists.txt | 15 + tests/core/interface/index_interface_test.cc | 1350 +++++ tests/core/metric/CMakeLists.txt | 14 + tests/core/metric/cosine_metric_test.cc | 606 +++ tests/core/metric/euclidean_metric_test.cc | 142 + tests/core/metric/hamming_metric_test.cc | 77 + .../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 | 823 ++++ tests/core/utility/CMakeLists.txt | 18 + tests/core/utility/buffer_storage_test.cpp | 72 + 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 | 100 + tests/db/CMakeLists.txt | 45 + tests/db/collection_test.cc | 4269 ++++++++++++++++ tests/db/common/CMakeLists.txt | 32 + tests/db/common/config_test.cc | 239 + tests/db/common/status_test.cc | 160 + tests/db/index/CMakeLists.txt | 41 + ...erted_column_indexer_array_numbers_test.cc | 351 ++ .../inverted_column_indexer_bool_test.cc | 366 ++ ...rted_column_indexer_cyclic_numbers_test.cc | 592 +++ ..._column_indexer_sequential_numbers_test.cc | 777 +++ .../inverted_column_indexer_string_test.cc | 381 ++ .../inverted_indexer_util_test.cc | 259 + .../column/vector_column_indexer_test.cc | 2694 ++++++++++ .../index/common/db_proto_converter_test.cc | 473 ++ tests/db/index/common/db_type_helper_test.cc | 291 ++ tests/db/index/common/doc_test.cc | 1347 +++++ tests/db/index/common/index_params_test.cc | 189 + tests/db/index/common/meta_test.cc | 252 + tests/db/index/common/query_params_test.cc | 83 + tests/db/index/common/schema_test.cc | 826 ++++ tests/db/index/common/version_manager_test.cc | 302 ++ .../segment/column_merging_reader_test.cc | 437 ++ tests/db/index/segment/segment_helper_test.cc | 694 +++ tests/db/index/segment/segment_test.cc | 2004 ++++++++ .../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 | 1157 +++++ 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 | 735 +++ tests/db/index/utils/utils.cc | 796 +++ tests/db/index/utils/utils.h | 219 + tests/db/sqlengine/CMakeLists.txt | 42 + tests/db/sqlengine/contain_test.cc | 484 ++ tests/db/sqlengine/forward_recall_test.cc | 938 ++++ tests/db/sqlengine/invert_recall_test.cc | 789 +++ tests/db/sqlengine/like_test.cc | 372 ++ tests/db/sqlengine/mock_segment.h | 513 ++ tests/db/sqlengine/optimizer_test.cc | 322 ++ tests/db/sqlengine/query_info_test.cc | 679 +++ tests/db/sqlengine/recall_base.h | 269 + tests/db/sqlengine/simple_rewriter_test.cc | 650 +++ tests/db/sqlengine/sqlengine_test.cc | 170 + tests/db/sqlengine/test_helper.h | 112 + tests/db/sqlengine/vector_recall_test.cc | 216 + thirdparty/CMakeLists.txt | 27 + thirdparty/CRoaring/CMakeLists.txt | 17 + thirdparty/CRoaring/CRoaring-2.0.4 | 1 + thirdparty/antlr/CMakeLists.txt | 19 + thirdparty/antlr/antlr4 | 1 + thirdparty/antlr/antlr4.patch | 32 + thirdparty/arrow/CMakeLists.txt | 92 + thirdparty/arrow/apache-arrow-21.0.0 | 1 + thirdparty/arrow/arrow.patch | 13 + thirdparty/gflags/CMakeLists.txt | 51 + thirdparty/gflags/gflags-2.2.2 | 1 + thirdparty/glog/CMakeLists.txt | 28 + thirdparty/glog/glog-0.5.0 | 1 + thirdparty/glog/glog.patch | 51 + thirdparty/googletest/CMakeLists.txt | 15 + thirdparty/googletest/googletest-1.10.0 | 1 + thirdparty/lz4/CMakeLists.txt | 35 + thirdparty/lz4/lz4-1.9.4 | 1 + thirdparty/magic_enum/CMakeLists.txt | 4 + thirdparty/magic_enum/magic_enum-0.9.7 | 1 + thirdparty/protobuf/CMakeLists.txt | 57 + thirdparty/protobuf/protobuf-3.21.12 | 1 + thirdparty/rocksdb/CMakeLists.txt | 18 + thirdparty/rocksdb/rocksdb-8.1.1 | 1 + thirdparty/sparsehash/CMakeLists.txt | 11 + thirdparty/sparsehash/sparseconfig.h | 46 + thirdparty/sparsehash/sparsehash-2.0.4 | 1 + thirdparty/yaml-cpp/CMakeLists.txt | 11 + thirdparty/yaml-cpp/yaml-cpp-0.6.3 | 1 + tools/CMakeLists.txt | 8 + tools/core/CMakeLists.txt | 59 + tools/core/README.md | 140 + tools/core/bench.cc | 913 ++++ tools/core/bench_original.cc | 1066 ++++ tools/core/bench_result.h | 124 + tools/core/convert_cohere_parquet.py | 187 + tools/core/filter_result_cache.h | 77 + tools/core/flow.h | 555 +++ tools/core/helper.h | 298 ++ tools/core/index_meta_helper.h | 131 + tools/core/local_builder.cc | 1235 +++++ tools/core/local_builder_original.cc | 1183 +++++ tools/core/meta_segment_common.h | 34 + tools/core/recall.cc | 1808 +++++++ tools/core/recall_original.cc | 2093 ++++++++ tools/core/txt2vecs.cc | 723 +++ tools/core/txt_input_reader.h | 549 +++ tools/core/vecs_common.h | 50 + tools/core/vecs_index_holder.h | 388 ++ tools/core/vecs_reader.h | 378 ++ 841 files changed, 293919 insertions(+) create mode 100644 .clang-format create mode 100644 .github/ISSUE_TEMPLATE/bug_report.yml create mode 100644 .github/ISSUE_TEMPLATE/config.yml create mode 100644 .github/ISSUE_TEMPLATE/feature_request.yml create mode 100644 .github/codecov.yml create mode 100644 .github/workflows/build_wheel.yml create mode 100644 .github/workflows/docker/Dockerfile.ubuntu18.10-glibc228 create mode 100644 .github/workflows/linux_x64_docker_ci.yml create mode 100644 .github/workflows/mac_arm64_ci.yml create mode 100644 .github/workflows/nightly_coverage.yml create mode 100644 .gitignore create mode 100644 .gitmodules create mode 100644 .pre-commit-config.yaml create mode 100644 CMakeLists.txt create mode 100644 CODE_OF_CONDUCT.md create mode 100644 CONTRIBUTING.md create mode 100644 LICENSE create mode 100644 README.md create mode 100644 cmake/bazel.cmake create mode 100644 cmake/option.cmake create mode 100644 cmake/utils.cmake 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_db_config.py create mode 100644 python/tests/test_collection.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_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/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/embedding.py create mode 100644 python/zvec/extension/rerank.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/vector_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/gcov.sh 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/buffer_manager.cc create mode 100644 src/ailego/buffer/buffer_manager.h create mode 100644 src/ailego/container/bitmap.cc create mode 100644 src/ailego/container/bitmap.h create mode 100644 src/ailego/container/blob.h create mode 100644 src/ailego/container/bloom_filter.h create mode 100644 src/ailego/container/cube.h create mode 100644 src/ailego/container/heap.h create mode 100644 src/ailego/container/hypercube.h create mode 100644 src/ailego/container/params.cc create mode 100644 src/ailego/container/params.h create mode 100644 src/ailego/container/reservoir.h create mode 100644 src/ailego/container/vector.h create mode 100644 src/ailego/container/vector_array.h create mode 100644 src/ailego/encoding/json.h create mode 100644 src/ailego/encoding/json/mod_json.c create mode 100644 src/ailego/encoding/json/mod_json.h create mode 100644 src/ailego/encoding/json/mod_json_plus.h create mode 100644 src/ailego/hash/crc32c.cc create mode 100644 src/ailego/hash/crc32c.h create mode 100644 src/ailego/hash/jump_hash.h create mode 100644 src/ailego/internal/cpu_features.cc create mode 100644 src/ailego/internal/cpu_features.h create mode 100644 src/ailego/internal/platform.h create mode 100644 src/ailego/io/file.cc create mode 100644 src/ailego/io/file.h 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/io/mmap_file.h create mode 100644 src/ailego/io/pid_file.cc create mode 100644 src/ailego/io/pid_file.h create mode 100644 src/ailego/logger/logger.cc create mode 100644 src/ailego/logger/logger.h 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_fp16.i create mode 100644 src/ailego/math/distance_matrix_fp32.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_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.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_fp32.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_int4.cc create mode 100644 src/ailego/math/euclidean_distance_matrix_int8.cc create mode 100644 src/ailego/math/hamming_distance_matrix.cc create mode 100644 src/ailego/math/hamming_distance_matrix.h 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_fp32.cc create mode 100644 src/ailego/math/inner_product_matrix_int4.cc create mode 100644 src/ailego/math/inner_product_matrix_int8.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.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_fp32.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_int4.cc create mode 100644 src/ailego/math/mips_euclidean_distance_matrix_int8.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/inner_product_distance_batch.h create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl.h create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl_fp16.h create mode 100644 src/ailego/math_batch/inner_product_distance_batch_impl_int8.h create mode 100644 src/ailego/math_batch/utils.h 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/parallel/thread_pool.h create mode 100644 src/ailego/parallel/thread_queue.h create mode 100644 src/ailego/pattern/closure.h create mode 100644 src/ailego/pattern/defer.h create mode 100644 src/ailego/pattern/expected.hpp create mode 100644 src/ailego/pattern/factory.h create mode 100644 src/ailego/pattern/scope_guard.h create mode 100644 src/ailego/pattern/singleton.h create mode 100644 src/ailego/string/string_concat_helper.h create mode 100644 src/ailego/string/string_view.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/file_helper.h create mode 100644 src/ailego/utility/float_helper.cc create mode 100644 src/ailego/utility/float_helper.h 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/process_helper.cc create mode 100644 src/ailego/utility/process_helper.h create mode 100644 src/ailego/utility/string_helper.cc create mode 100644 src/ailego/utility/string_helper.h create mode 100644 src/ailego/utility/string_helper_impl.h create mode 100644 src/ailego/utility/time_helper.cc create mode 100644 src/ailego/utility/time_helper.h create mode 100644 src/ailego/utility/type_helper.h 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/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_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/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/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/diskann_builder.h create mode 100644 src/core/algorithm/diskann/diskann_searcher.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_builder.cc create mode 100644 src/core/algorithm/hnsw/hnsw_builder.h create mode 100644 src/core/algorithm/hnsw/hnsw_builder_entity.cc create mode 100644 src/core/algorithm/hnsw/hnsw_builder_entity.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_searcher.cc create mode 100644 src/core/algorithm/hnsw/hnsw_searcher.h create mode 100644 src/core/algorithm/hnsw/hnsw_searcher_entity.cc create mode 100644 src/core/algorithm/hnsw/hnsw_searcher_entity.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_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/framework/CMakeLists.txt create mode 100644 src/core/framework/index_builder.h create mode 100644 src/core/framework/index_bundle.h create mode 100644 src/core/framework/index_cluster.cc create mode 100644 src/core/framework/index_cluster.h create mode 100644 src/core/framework/index_context.cc create mode 100644 src/core/framework/index_context.h create mode 100644 src/core/framework/index_converter.cc create mode 100644 src/core/framework/index_converter.h create mode 100644 src/core/framework/index_document.h create mode 100644 src/core/framework/index_dumper.h create mode 100644 src/core/framework/index_error.cc create mode 100644 src/core/framework/index_error.h create mode 100644 src/core/framework/index_factory.cc create mode 100644 src/core/framework/index_factory.h create mode 100644 src/core/framework/index_features.h create mode 100644 src/core/framework/index_filter.h create mode 100644 src/core/framework/index_flow.cc create mode 100644 src/core/framework/index_flow.h create mode 100644 src/core/framework/index_format.h create mode 100644 src/core/framework/index_framework.h create mode 100644 src/core/framework/index_groupby.h create mode 100644 src/core/framework/index_helper.cc create mode 100644 src/core/framework/index_helper.h create mode 100644 src/core/framework/index_holder.h create mode 100644 src/core/framework/index_logger.cc create mode 100644 src/core/framework/index_logger.h create mode 100644 src/core/framework/index_mapping.cc create mode 100644 src/core/framework/index_mapping.h create mode 100644 src/core/framework/index_memory.h create mode 100644 src/core/framework/index_meta.cc create mode 100644 src/core/framework/index_meta.h create mode 100644 src/core/framework/index_metric.h create mode 100644 src/core/framework/index_module.h create mode 100644 src/core/framework/index_packer.h create mode 100644 src/core/framework/index_plugin.cc create mode 100644 src/core/framework/index_plugin.h create mode 100644 src/core/framework/index_provider.h create mode 100644 src/core/framework/index_reducer.h create mode 100644 src/core/framework/index_refiner.h create mode 100644 src/core/framework/index_reformer.h create mode 100644 src/core/framework/index_runner.h create mode 100644 src/core/framework/index_searcher.h create mode 100644 src/core/framework/index_segment_storage.h create mode 100644 src/core/framework/index_stats.h create mode 100644 src/core/framework/index_storage.h create mode 100644 src/core/framework/index_streamer.h create mode 100644 src/core/framework/index_threads.h create mode 100644 src/core/framework/index_trainer.h create mode 100644 src/core/framework/index_unpacker.h create mode 100644 src/core/framework/index_version.cc create mode 100644 src/core/framework/index_version.h create mode 100644 src/core/interface/CMakeLists.txt create mode 100644 src/core/interface/constants.h create mode 100644 src/core/interface/index.cc create mode 100644 src/core/interface/index.h create mode 100644 src/core/interface/index_factory.cc create mode 100644 src/core/interface/index_factory.h create mode 100644 src/core/interface/index_param.cc create mode 100644 src/core/interface/index_param.h create mode 100644 src/core/interface/index_param_builders.h 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/ivf_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/hamming_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/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/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/utility/CMakeLists.txt create mode 100644 src/core/utility/basic_refiner.cc 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/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/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/collection.h 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/config.h 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/rocbsdb_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/status.h 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/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/doc.h 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.h create mode 100644 src/db/index/common/meta.h create mode 100644 src/db/index/common/options.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_params.h create mode 100644 src/db/index/common/schema.cc create mode 100644 src/db/index/common/schema.h create mode 100644 src/db/index/common/stats.h create mode 100644 src/db/index/common/type.h 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_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/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/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/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 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/buffer/buffer_manager_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/pid_file_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/hamming_distance_matrix_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/process_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/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/opt_kmeans_cluster_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.cpp 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_builder_test.cc create mode 100644 tests/core/algorithm/hnsw/hnsw_searcher_test.cpp create mode 100644 tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cpp create mode 100644 tests/core/algorithm/hnsw/hnsw_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.cc 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/framework/CMakeLists.txt create mode 100644 tests/core/interface/CMakeLists.txt 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/hamming_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/utility/CMakeLists.txt create mode 100644 tests/core/utility/buffer_storage_test.cpp 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/index/CMakeLists.txt 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_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_test.cc 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/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/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 thirdparty/CMakeLists.txt create mode 100644 thirdparty/CRoaring/CMakeLists.txt create mode 160000 thirdparty/CRoaring/CRoaring-2.0.4 create mode 100644 thirdparty/antlr/CMakeLists.txt create mode 160000 thirdparty/antlr/antlr4 create mode 100644 thirdparty/antlr/antlr4.patch create mode 100644 thirdparty/arrow/CMakeLists.txt create mode 160000 thirdparty/arrow/apache-arrow-21.0.0 create mode 100644 thirdparty/arrow/arrow.patch create mode 100644 thirdparty/gflags/CMakeLists.txt create mode 160000 thirdparty/gflags/gflags-2.2.2 create mode 100644 thirdparty/glog/CMakeLists.txt create mode 160000 thirdparty/glog/glog-0.5.0 create mode 100644 thirdparty/glog/glog.patch create mode 100644 thirdparty/googletest/CMakeLists.txt create mode 160000 thirdparty/googletest/googletest-1.10.0 create mode 100644 thirdparty/lz4/CMakeLists.txt create mode 160000 thirdparty/lz4/lz4-1.9.4 create mode 100644 thirdparty/magic_enum/CMakeLists.txt create mode 160000 thirdparty/magic_enum/magic_enum-0.9.7 create mode 100644 thirdparty/protobuf/CMakeLists.txt create mode 160000 thirdparty/protobuf/protobuf-3.21.12 create mode 100644 thirdparty/rocksdb/CMakeLists.txt create mode 160000 thirdparty/rocksdb/rocksdb-8.1.1 create mode 100644 thirdparty/sparsehash/CMakeLists.txt create mode 100644 thirdparty/sparsehash/sparseconfig.h create mode 160000 thirdparty/sparsehash/sparsehash-2.0.4 create mode 100644 thirdparty/yaml-cpp/CMakeLists.txt create mode 160000 thirdparty/yaml-cpp/yaml-cpp-0.6.3 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 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/.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..98d35ad --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,9 @@ +blank_issues_enabled: false +contact_links: + - name: Documentation + url: https://github.com/your-org/zvec#readme + about: Check the README, build guide, and API docs first. + + - name: Python API Examples + url: https://github.com/your-org/zvec/tree/main/examples + about: See working usage examples. diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml new file mode 100644 index 0000000..ccdda16 --- /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: ["enhancement"] +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/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/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml new file mode 100644 index 0000000..ac8b19c --- /dev/null +++ b/.github/workflows/build_wheel.yml @@ -0,0 +1,57 @@ +name: Build Wheels + +on: + workflow_dispatch: + +jobs: + build_wheels: + name: Build wheels on self-hosted manylinux_2_28 + runs-on: linux_x64 + + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Set up Python (for cibuildwheel controller) + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install cibuildwheel + run: | + pip install --upgrade pip + pip install cibuildwheel==2.17.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 TestPyPI + if: success() && github.event_name == 'workflow_dispatch' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }} + TWINE_REPOSITORY_URL: https://test.pypi.org/legacy/ + run: | + pip install twine + twine upload --skip-existing --verbose wheelhouse/*.whl + + - name: (Optional) Install and test from TestPyPI + if: success() && github.event_name == 'workflow_dispatch' + run: | + # Create a clean venv + python -m venv test_env + source test_env/bin/activate + pip install --upgrade pip + # Install from TestPyPI (must allow pre-releases if version has dev/alpha) + pip install --index-url https://test.pypi.org/simple/ \ + --extra-index-url https://pypi.org/simple/ \ + zvec + # Run a simple smoke test + python -c "import zvec; print('Import OK:', zvec.__version__)" + shell: bash \ No newline at end of file diff --git a/.github/workflows/docker/Dockerfile.ubuntu18.10-glibc228 b/.github/workflows/docker/Dockerfile.ubuntu18.10-glibc228 new file mode 100644 index 0000000..503829e --- /dev/null +++ b/.github/workflows/docker/Dockerfile.ubuntu18.10-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/linux_x64_docker_ci.yml b/.github/workflows/linux_x64_docker_ci.yml new file mode 100644 index 0000000..7e42eeb --- /dev/null +++ b/.github/workflows/linux_x64_docker_ci.yml @@ -0,0 +1,124 @@ +name: Zvec LinuxX64 CI + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + workflow_dispatch: + +jobs: + build: + runs-on: linux_x64 + + strategy: + matrix: + python-version: ['3.10'] + fail-fast: false + + container: + image: zvec-registry.cn-hongkong.cr.aliyuncs.com/zvec/zvec:0.0.2 + options: --user root + + steps: + - name: Activate Conda environment + run: | + # Map version to env name + if [[ "${{ matrix.python-version }}" == "3.10" ]]; then + ENV_NAME="py310" + elif [[ "${{ matrix.python-version }}" == "3.11" ]]; then + ENV_NAME="py311" + elif [[ "${{ matrix.python-version }}" == "3.12" ]]; then + ENV_NAME="py312" + else + echo "Unsupported Python version" + exit 1 + fi + echo "CONDA_ENV_NAME=$ENV_NAME" >> $GITHUB_ENV + source /opt/miniforge3/bin/activate $ENV_NAME + + pip install --upgrade pip pytest pytest-cov ruff + + python --version + shell: bash + + - name: Prepare clean build directory + run: | + export CLEAN_WORKSPACE="/tmp/zvec" + mkdir -p "$CLEAN_WORKSPACE" + cd "$CLEAN_WORKSPACE" + + git config --global --add safe.directory "$CLEAN_WORKSPACE" + git clone --recursive "https://x-access-token:${GITHUB_TOKEN}@github.com/${GITHUB_REPOSITORY}.git" . + + if [ -n "${{ github.event.pull_request.head.sha }}" ]; then + git checkout "${{ github.event.pull_request.head.sha }}" + else + git checkout "${{ github.sha }}" + fi + + echo "CLEAN_WORKSPACE=$CLEAN_WORKSPACE" >> $GITHUB_ENV + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + shell: bash + + - name: Run Ruff Linter + run: | + source /opt/miniforge3/bin/activate ${{ env.CONDA_ENV_NAME }} + cd "$CLEAN_WORKSPACE" + ruff check . + shell: bash + + - name: Run Ruff Formatter Check (ensure code is formatted) + run: | + source /opt/miniforge3/bin/activate ${{ env.CONDA_ENV_NAME }} + cd "$CLEAN_WORKSPACE" + ruff format --check . + shell: bash + + - name: Run clang-format Check + run: | + source /opt/miniforge3/bin/activate ${{ env.CONDA_ENV_NAME }} + cd "$CLEAN_WORKSPACE" + + 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-18 --dry-run --Werror $CPP_FILES + shell: bash + + - name: Install Python dependencies and build package + run: | + source /opt/miniforge3/bin/activate ${{ env.CONDA_ENV_NAME }} + cd "$CLEAN_WORKSPACE" + + NPROC=$(nproc 2>/dev/null || echo $(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 2)) + echo "ParallelGroup: Using $NPROC parallel jobs for build" + + CMAKE_GENERATOR="Unix Makefiles" \ + CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \ + pip install -v . + shell: bash + + - name: Run Python Tests with Coverage + run: | + source /opt/miniforge3/bin/activate ${{ env.CONDA_ENV_NAME }} + cd "$CLEAN_WORKSPACE" + python -m pytest python/tests/ --cov=zvec --cov-report=xml --no-cov-on-fail + shell: bash + + - name: Run Cpp Tests with Coverage + run: | + cd "$CLEAN_WORKSPACE/build" + make unittest -j$(nproc) + shell: bash diff --git a/.github/workflows/mac_arm64_ci.yml b/.github/workflows/mac_arm64_ci.yml new file mode 100644 index 0000000..2b507cd --- /dev/null +++ b/.github/workflows/mac_arm64_ci.yml @@ -0,0 +1,85 @@ +name: Zvec MacArm64 CI + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + workflow_dispatch: + +jobs: + build: + runs-on: mac_m1_arm + + steps: + - name: Run Ruff Linter + run: | + cd "$CLEAN_WORKSPACE" + ruff check . + shell: bash + + - name: Run Ruff Formatter Check (ensure code is formatted) + run: | + cd "$CLEAN_WORKSPACE" + ruff format --check . + shell: bash + + - name: Run clang-format Check + run: | + cd "$CLEAN_WORKSPACE" + + 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 + + - name: Prepare clean build directory + run: | + export CLEAN_WORKSPACE="/tmp/zvec" + rm -rf "$CLEAN_WORKSPACE" + mkdir -p "$CLEAN_WORKSPACE" + cd "$CLEAN_WORKSPACE" + + git config --global --add safe.directory "$CLEAN_WORKSPACE" + git clone --recursive "https://x-access-token:${GITHUB_TOKEN}@github.com/${GITHUB_REPOSITORY}.git" . + + if [ -n "${{ github.event.pull_request.head.sha }}" ]; then + git checkout "${{ github.event.pull_request.head.sha }}" + else + git checkout "${{ github.sha }}" + fi + + echo "CLEAN_WORKSPACE=$CLEAN_WORKSPACE" >> $GITHUB_ENV + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + shell: bash + + - name: Install Python dependencies and build package + run: | + cd "$CLEAN_WORKSPACE" + pip install --upgrade pip pytest pytest-cov ruff + pip install -v . + shell: bash + + - name: Run Python Tests with Coverage + run: | + cd "$CLEAN_WORKSPACE" + python -m pytest python/tests/ --cov=zvec --cov-report=xml --no-cov-on-fail + shell: bash + + - name: Run Cpp Tests with Coverage + run: | + cd "$CLEAN_WORKSPACE/build" + make unittest -j 16 + shell: bash \ No newline at end of file diff --git a/.github/workflows/nightly_coverage.yml b/.github/workflows/nightly_coverage.yml new file mode 100644 index 0000000..86ee2c8 --- /dev/null +++ b/.github/workflows/nightly_coverage.yml @@ -0,0 +1,95 @@ +name: Nightly Coverage Report + +on: + schedule: + # Runs daily at 00:00 CST (China Standard Time) = 16:00 UTC + - cron: '0 16 * * *' + + workflow_dispatch: + +jobs: + coverage: + runs-on: linux_x64 + + strategy: + matrix: + python-version: ['3.10'] + fail-fast: false + + container: + image: zvec-registry.cn-hongkong.cr.aliyuncs.com/zvec/zvec:0.0.2 + options: --user root + + steps: + - name: Activate Conda environment + run: | + if [[ "${{ matrix.python-version }}" == "3.10" ]]; then + ENV_NAME="py310" + elif [[ "${{ matrix.python-version }}" == "3.11" ]]; then + ENV_NAME="py311" + elif [[ "${{ matrix.python-version }}" == "3.12" ]]; then + ENV_NAME="py312" + else + echo "Unsupported Python version" + exit 1 + fi + echo "CONDA_ENV_NAME=$ENV_NAME" >> $GITHUB_ENV + source /opt/miniforge3/bin/activate "$ENV_NAME" + python --version + shell: bash + + - name: Prepare clean build directory + run: | + export CLEAN_WORKSPACE="/tmp/zvec" + mkdir -p "$CLEAN_WORKSPACE" + cd "$CLEAN_WORKSPACE" + + git config --global --add safe.directory "$CLEAN_WORKSPACE" + git clone --recursive "https://x-access-token:${GITHUB_TOKEN}@github.com/${GITHUB_REPOSITORY}.git" . + + git checkout main # Always use main for nightly + + echo "CLEAN_WORKSPACE=$CLEAN_WORKSPACE" >> $GITHUB_ENV + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + shell: bash + + - name: Build with COVERAGE config + run: | + source /opt/miniforge3/bin/activate "${{ env.CONDA_ENV_NAME }}" + cd "$CLEAN_WORKSPACE" + pip install --upgrade pip pytest pytest-cov + + NPROC=$(nproc 2>/dev/null || echo 2) + echo "Using $NPROC parallel jobs" + + CMAKE_GENERATOR="Unix Makefiles" \ + CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \ + pip install -v . \ + --config-settings="cmake.build-type=COVERAGE" + shell: bash + + - name: Run Python Tests with Coverage + run: | + source /opt/miniforge3/bin/activate "${{ env.CONDA_ENV_NAME }}" + cd "$CLEAN_WORKSPACE" + python -m pytest python/tests/ --cov=zvec --cov-report=xml + shell: bash + + - name: Run C++ Tests and Generate Coverage + run: | + cd "$CLEAN_WORKSPACE/build" + make unittest -j$(nproc) # Run all (nightly can afford it) + cd "$CLEAN_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@v4 + with: + files: ${{ env.CLEAN_WORKSPACE }}/proxima-zvec-filtered.lcov.info,${{ env.CLEAN_WORKSPACE }}/coverage.xml + flags: python,cpp,nightly + name: nightly-linux-py${{ matrix.python-version }} + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..75b45dc --- /dev/null +++ b/.gitignore @@ -0,0 +1,50 @@ +.* +*~ +bazel-* +/*build* +bin/* +lib/* +var/* +thirdparty +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 +!.circleci +!.drone.yml +sdk/python/dist/ +compile_commands.json +dist +html +*.lcov.info + +# 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-* diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..5086d05 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,41 @@ +[submodule "thirdparty/googletest/googletest-1.10.0"] + path = thirdparty/googletest/googletest-1.10.0 + url = https://github.com/google/googletest.git +[submodule "thirdparty/sparsehash/sparsehash-2.0.4"] + path = thirdparty/sparsehash/sparsehash-2.0.4 + url = https://github.com/sparsehash/sparsehash.git + ignore = untracked +[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 +[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 = all +[submodule "thirdparty/CRoaring/CRoaring-2.0.4"] + path = thirdparty/CRoaring/CRoaring-2.0.4 + url = https://github.com/RoaringBitmap/CRoaring.git +[submodule "thirdparty/glog/glog-0.5.0"] + path = thirdparty/glog/glog-0.5.0 + url = https://github.com/google/glog.git + ignore = all +[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 = all +[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 = all diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..4b9beef --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,30 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.14.4 + hooks: + - id: ruff-check + name: Ruff (lint + fix) + args: [--config=pyproject.toml, --fix] + types_or: [python, pyi] + - id: ruff-format + name: Ruff (format) + args: [--config=pyproject.toml] + types_or: [python, pyi] + + + - repo: https://github.com/pre-commit/mirrors-clang-format + rev: v18.1.8 + hooks: + - id: clang-format + name: clang-format (C++) + args: [--style=file, --fallback-style=none] + types_or: [c++, c] + exclude: ^thirdparty/.*$ + + + - repo: https://github.com/gitleaks/gitleaks + rev: v8.21.4 + hooks: + - id: gitleaks + name: Detect hardcoded secrets (gitleaks) + description: Scans staged files for secrets using gitleaks rules diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..31ccec8 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,60 @@ +cmake_minimum_required(VERSION 3.13) +cmake_policy(SET CMP0077 NEW) +project(zvec) +set(CC_CXX_STANDARD 17) + +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Werror=return-type") +set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Werror=return-type") + +if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + 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_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_directories(${PROJECT_ROOT_DIR}/src) + +option(BUILD_PYTHON_BINDINGS "Build Python bindings using pybind11" OFF) +message(STATUS "BUILD_PYTHON_BINDINGS:${BUILD_PYTHON_BINDINGS}") + +option(BUILD_TOOLS "Build tools" ON) +message(STATUS "BUILD_TOOLS:${BUILD_TOOLS}") + +cc_directory(thirdparty) +cc_directories(src) +cc_directories(tests) + +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}") + install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}) +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..6a25890 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,104 @@ +# 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 + +### Prerequisites +- Python ≥ 3.9 +- CMake ≥ 3.15 (`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/your-org/zvec.git +cd zvec +``` + +> 💡 **Submodule tip**: If you forgot `--recursive`, run: +> `git submodule update --init --recursive` + +### 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/graphscope/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/README.md b/README.md new file mode 100644 index 0000000..c9beead --- /dev/null +++ b/README.md @@ -0,0 +1,99 @@ +

+ + + zvec logo + +

+ +

+ Code Coverage + Linux x64 CI + Latest Release + License +

+ +# Zvec + +**Zvec** is an open-source, in-process vector database — lightweight, lightning-fast, and designed to embed directly into applications. Built on **Proxima** (Alibaba's battle-tested vector search engine), it delivers production-grade, low-latency, scalable similarity search with minimal setup. + +[📚 Quick Start](https://zvec.org/en/docs/quickstart/) | 🌐 [Website](https://zvec.org/en/) | [📖 Documentation](https://zvec.org/en/docs/) | [📊 Benchmarks](https://zvec.org/en/docs/benchmarks/) + +## 💫 Features + +- **Blazing Fast**: Searches billions of vectors in milliseconds. +- **Simple, Just Works**: Install with `pip install zvec` and start searching in seconds. No servers, no config, no fuss. +- **Dense + Sparse Vectors**: Work with both dense and sparse embeddings, with native support for multi-vector queries in a single call. +- **Hybrid Search**: Combine semantic similarity with structured filters for precise results. +- **Runs Anywhere**: As an in-process library, Zvec runs wherever your code runs — notebooks, servers, CLI tools, or even edge devices. + +## 📦 Installation + +Install Zvec from PyPI with a single command: + +```bash +pip install zvec +``` + +**Requirements**: + +- Python 3.9 +- **Supported platforms**: Linux (x86_64) or macOS (ARM64/x86_64) + +### Building from Source + +If you prefer to build Zvec from source, you can clone the repository and install it locally. +**Note**: *The repository includes Git submodules, so cloning may take a few minutes depending on your network*. + +```bash +# Clone the repository +git clone --recurse-submodules https://github.com/GraphScope/zvec.git +cd zvec + +# Install from source +pip install . +``` + +**Build Requirements**: + +- Python 3.9 or higher +- **Supported platforms**: Linux (x86_64) or macOS (ARM64/x86_64) +- C++17 compatible compiler + +For detailed build instructions, check the [Building from Source](https://zvec.org/en/docs/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.VectorQuery("embedding", vector=[0.4, 0.3, 0.3, 0.1]), + topk=10 +) + +# Results: list of {'id': str, 'score': float, ...}, sorted by relevance +print(results) +``` + +## ❤️ 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/cmake/bazel.cmake b/cmake/bazel.cmake new file mode 100644 index 0000000..7226ce1 --- /dev/null +++ b/cmake/bazel.cmake @@ -0,0 +1,2338 @@ +## +## 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 [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.1 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) + add_custom_target( + unittest + COMMAND ${CMAKE_CTEST_COMMAND} --output-on-failure + --build-config $ + ) +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(NOT MSVC) + # Use color in diagnostics + set( + _COMPILER_FLAGS + "$<$:-fcolor-diagnostics;-stdlib=libc++>" + "$<$:-fcolor-diagnostics>" + "$<$:-fdiagnostics-color=always>" + ) + add_compile_options( + "$<$:${_COMPILER_FLAGS}>" + "$<$:${_COMPILER_FLAGS}>" + ) + unset(_COMPILER_FLAGS) +else() + # Replace the default compiling flags + 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 + ) + foreach(COMPILER_FLAG ${_COMPILER_FLAGS}) + string(REPLACE "/MT" "/MD" ${COMPILER_FLAG} "${${COMPILER_FLAG}}") + string(REGEX REPLACE "/W[0-9]" "" ${COMPILER_FLAG} "${${COMPILER_FLAG}}") + endforeach() + unset(_COMPILER_FLAGS) + add_definitions(-D_CRT_SECURE_NO_WARNINGS) + # Build shared library as default + set(BUILD_SHARED_LIBS ON) +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>>" + ) + +# C/C++ strict compile flags +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_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_ASAN_COMPILE_FLAGS} + ${BAZEL_CC_COVERAGE_COMPILE_FLAGS} + ) +endif() + + +# C/C++ strict link flags +set( + BAZEL_CC_STRICT_LINK_FLAGS + "$<$:-stdlib=libc++>" + ${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 + "$<$:-stdlib=libc++>" + ${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} $ + ) + add_library( + ${_NAME} SHARED ${_OPTION} $ + ) + 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 + "$" + ) + endif() + endforeach() + + if(LIBS_DEPS) + add_dependencies(${_NAME} ${LIBS_DEPS}) + 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}) + 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) + + 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") + list(APPEND LINK_LIBS -Wl,--whole-archive ${LIB} -Wl,--no-whole-archive) + else() + list(APPEND LINK_LIBS -Wl,-force_load ${LIB}) + endif() + else() + # Microsoft Visual C++ + list(APPEND LINK_LIBS /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(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" "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 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() + + if(CC_ARGS_SHARED AND CC_ARGS_STATIC) + _add_library(${CC_ARGS_NAME} "${EXCLUDE_OPTION}" "${CC_ARGS_SRCS}") + elseif(CC_ARGS_SHARED) + add_library(${CC_ARGS_NAME} SHARED ${EXCLUDE_OPTION} ${CC_ARGS_SRCS}) + elseif(CC_ARGS_STATIC) + add_library(${CC_ARGS_NAME} STATIC ${EXCLUDE_OPTION} ${CC_ARGS_SRCS}) + else() + add_library(${CC_ARGS_NAME} ${EXCLUDE_OPTION} ${CC_ARGS_SRCS}) + 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}) + + 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() + add_executable(${CC_ARGS_NAME} EXCLUDE_FROM_ALL ${CC_ARGS_SRCS}) + + _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}) + add_custom_target( + unittest.${CC_ARGS_NAME} + COMMAND $ "${CC_ARGS_ARGS}" + WORKING_DIRECTORY ${PROJECT_BINARY_DIR} + DEPENDS ${CC_ARGS_NAME} + ) + add_test( + NAME ${CC_ARGS_NAME} + COMMAND $ "${CC_ARGS_ARGS}" + WORKING_DIRECTORY ${PROJECT_BINARY_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") + set(FIND_GTEST_LIBS "gtest;gtest_main" CACHE STRING "GTest libraries") + 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") + set(FIND_GMOCK_LIBS "gmock;gmock_main" CACHE STRING "GMock libraries") + 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}}) + 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}) + add_custom_target( + unittest.${CUDA_ARGS_NAME} + COMMAND $ "${CUDA_ARGS_ARGS}" + WORKING_DIRECTORY ${PROJECT_BINARY_DIR} + DEPENDS ${CUDA_ARGS_NAME} + ) + add_test( + NAME ${CUDA_ARGS_NAME} + COMMAND $ "${CUDA_ARGS_ARGS}" + WORKING_DIRECTORY ${PROJECT_BINARY_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.1)\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/option.cmake b/cmake/option.cmake new file mode 100644 index 0000000..7983be5 --- /dev/null +++ b/cmake/option.cmake @@ -0,0 +1,585 @@ +## 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_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) + +## 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) + +## Intel Instructions +option(ENABLE_SSE "Enable Intel SSE instructions" OFF) +option(ENABLE_SSE2 "Enable Intel SSE2 instructions" OFF) +option(ENABLE_SSE3 "Enable Intel SSE3 instructions" OFF) +option(ENABLE_SSSE3 "Enable Intel SSSE3 instructions" OFF) +option(ENABLE_SSE4.1 "Enable Intel SSE4.1 instructions" OFF) +option(ENABLE_SSE4.2 "Enable Intel SSE4.2 instructions" OFF) +option(ENABLE_AVX "Enable Intel AVX instructions" OFF) +option(ENABLE_AVX2 "Enable Intel AVX2 instructions" OFF) +option(ENABLE_AVX512F "Enable Intel AVX512F instructions" OFF) +option(ENABLE_AVX512BW "Enable Intel AVX512BW instructions" OFF) +option(ENABLE_AVX512CD "Enable Intel AVX512CD instructions" OFF) +option(ENABLE_AVX512DQ "Enable Intel AVX512DQ instructions" OFF) +option(ENABLE_AVX512ER "Enable Intel AVX512ER instructions" OFF) +option(ENABLE_AVX512IFMA52 "Enable Intel AVX512IFMA52 instructions" OFF) +option(ENABLE_AVX512PF "Enable Intel AVX512PF instructions" OFF) +option(ENABLE_AVX512VL "Enable Intel AVX512VL instructions" OFF) +option(ENABLE_AVX512VPOPCNTDQ "Enable Intel AVX512VPOPCNTDQ instructions" OFF) +option(ENABLE_AVX5124FMAPS "Enable Intel AVX5124FMAPS instructions" OFF) +option(ENABLE_AVX5124VNNIW "Enable Intel AVX5124VNNIW instructions" OFF) +option(ENABLE_AVX512BF16 "Enable Intel AVX512BF16 instructions" OFF) +option(ENABLE_AVX512BITALG "Enable Intel AVX512BITALG instructions" OFF) +option(ENABLE_AVX512VBMI "Enable Intel AVX512VBMI instructions" OFF) +option(ENABLE_AVX512VBMI2 "Enable Intel AVX512VBMI2 instructions" OFF) +option(ENABLE_AVX512VNNI "Enable Intel AVX512VNNI instructions" OFF) +option(ENABLE_AVX512VP2INTERSECT "Enable Intel AVX512VP2INTERSECT instructions" OFF) +option(ENABLE_AVX512FP16 "Enable Intel AVX512FP16 instructions" OFF) + +option(ENABLE_AVXVNNI "Enable Intel AVXVNNI instructions" OFF) +option(ENABLE_FMA "Enable Intel FMA instructions" OFF) +option(ENABLE_F16C "Enable Intel F16C instructions" 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) + +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() + +include(CheckCCompilerFlag) +if(NOT MSVC) + if(ENABLE_ARMV8.6A) + CHECK_C_COMPILER_FLAG("-march=armv8.6-a" COMPILER_SUPPORT_ARMV8.6A) + if(COMPILER_SUPPORT_ARMV8.6A) + _AppendFlags(CMAKE_C_FLAGS -march=armv8.6-a) + _AppendFlags(CMAKE_CXX_FLAGS -march=armv8.6-a) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=armv8.6-a\'.") + endif() + endif() + + if(ENABLE_ARMV8.5A) + CHECK_C_COMPILER_FLAG("-march=armv8.5-a" COMPILER_SUPPORT_ARMV8.5A) + if(COMPILER_SUPPORT_ARMV8.5A) + _AppendFlags(CMAKE_C_FLAGS -march=armv8.5-a) + _AppendFlags(CMAKE_CXX_FLAGS -march=armv8.5-a) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=armv8.5-a\'.") + endif() + endif() + + if(ENABLE_ARMV8.4A) + CHECK_C_COMPILER_FLAG("-march=armv8.4-a" COMPILER_SUPPORT_ARMV8.4A) + if(COMPILER_SUPPORT_ARMV8.4A) + _AppendFlags(CMAKE_C_FLAGS -march=armv8.4-a) + _AppendFlags(CMAKE_CXX_FLAGS -march=armv8.4-a) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=armv8.4-a\'.") + endif() + endif() + + if(ENABLE_ARMV8.3A) + CHECK_C_COMPILER_FLAG("-march=armv8.3-a" COMPILER_SUPPORT_ARMV8.3A) + if(COMPILER_SUPPORT_ARMV8.3A) + _AppendFlags(CMAKE_C_FLAGS -march=armv8.3-a) + _AppendFlags(CMAKE_CXX_FLAGS -march=armv8.3-a) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=armv8.3-a\'.") + endif() + endif() + + if(ENABLE_ARMV8.2A) + CHECK_C_COMPILER_FLAG("-march=armv8.2-a" COMPILER_SUPPORT_ARMV8.2A) + if(COMPILER_SUPPORT_ARMV8.2A) + _AppendFlags(CMAKE_C_FLAGS -march=armv8.2-a) + _AppendFlags(CMAKE_CXX_FLAGS -march=armv8.2-a) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=armv8.2-a\'.") + endif() + endif() + + if(ENABLE_ARMV8.1A) + CHECK_C_COMPILER_FLAG("-march=armv8.1-a" COMPILER_SUPPORT_ARMV8.1A) + if(COMPILER_SUPPORT_ARMV8.1A) + _AppendFlags(CMAKE_C_FLAGS -march=armv8.1-a) + _AppendFlags(CMAKE_CXX_FLAGS -march=armv8.1-a) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=armv8.1-a\'.") + endif() + endif() + + if(ENABLE_ARMV8A) + CHECK_C_COMPILER_FLAG("-march=armv8-a" COMPILER_SUPPORT_ARMV8A) + if(COMPILER_SUPPORT_ARMV8A) + _AppendFlags(CMAKE_C_FLAGS -march=armv8-a) + _AppendFlags(CMAKE_CXX_FLAGS -march=armv8-a) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=armv8-a\'.") + endif() + endif() + + if(ENABLE_ZEN1) + CHECK_C_COMPILER_FLAG("-march=znver1" COMPILER_SUPPORT_ZEN1) + if(COMPILER_SUPPORT_ZEN1) + _AppendFlags(CMAKE_C_FLAGS -march=znver1) + _AppendFlags(CMAKE_CXX_FLAGS -march=znver1) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=znver1\'.") + endif() + endif() + + if(ENABLE_ZEN2) + CHECK_C_COMPILER_FLAG("-march=znver2" COMPILER_SUPPORT_ZEN2) + if(COMPILER_SUPPORT_ZEN2) + _AppendFlags(CMAKE_C_FLAGS -march=znver2) + _AppendFlags(CMAKE_CXX_FLAGS -march=znver2) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=znver2\'.") + endif() + endif() + + if(ENABLE_ZEN3) + CHECK_C_COMPILER_FLAG("-march=znver3" COMPILER_SUPPORT_ZEN3) + if(COMPILER_SUPPORT_ZEN3) + _AppendFlags(CMAKE_C_FLAGS -march=znver3) + _AppendFlags(CMAKE_CXX_FLAGS -march=znver3) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=znver3\'.") + endif() + endif() + + if(ENABLE_NEHALEM) + CHECK_C_COMPILER_FLAG("-march=nehalem" COMPILER_SUPPORT_NEHALEM) + if(COMPILER_SUPPORT_NEHALEM) + _AppendFlags(CMAKE_C_FLAGS -march=nehalem) + _AppendFlags(CMAKE_CXX_FLAGS -march=nehalem) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=nehalem\'.") + endif() + endif() + + if(ENABLE_SANDYBRIDGE) + CHECK_C_COMPILER_FLAG("-march=sandybridge" COMPILER_SUPPORT_SANDYBRIDGE) + if(COMPILER_SUPPORT_SANDYBRIDGE) + _AppendFlags(CMAKE_C_FLAGS -march=sandybridge) + _AppendFlags(CMAKE_CXX_FLAGS -march=sandybridge) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=sandybridge\'.") + endif() + endif() + + if(ENABLE_HASWELL) + CHECK_C_COMPILER_FLAG("-march=haswell" COMPILER_SUPPORT_HASWELL) + if(COMPILER_SUPPORT_HASWELL) + _AppendFlags(CMAKE_C_FLAGS -march=haswell) + _AppendFlags(CMAKE_CXX_FLAGS -march=haswell) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=haswell\'.") + endif() + endif() + + if(ENABLE_BROADWELL) + CHECK_C_COMPILER_FLAG("-march=broadwell" COMPILER_SUPPORT_BROADWELL) + if(COMPILER_SUPPORT_BROADWELL) + _AppendFlags(CMAKE_C_FLAGS -march=broadwell) + _AppendFlags(CMAKE_CXX_FLAGS -march=broadwell) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=broadwell\'.") + endif() + endif() + + if(ENABLE_SKYLAKE) + CHECK_C_COMPILER_FLAG("-march=skylake" COMPILER_SUPPORT_SKYLAKE) + if(COMPILER_SUPPORT_SKYLAKE) + _AppendFlags(CMAKE_C_FLAGS -march=skylake) + _AppendFlags(CMAKE_CXX_FLAGS -march=skylake) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=skylake\'.") + endif() + endif() + + if(ENABLE_SKYLAKE_AVX512) + CHECK_C_COMPILER_FLAG("-march=skylake-avx512" COMPILER_SUPPORT_SKYLAKE_AVX512) + if(COMPILER_SUPPORT_SKYLAKE_AVX512) + _AppendFlags(CMAKE_C_FLAGS -march=skylake-avx512) + _AppendFlags(CMAKE_CXX_FLAGS -march=skylake-avx512) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=skylake-avx512\'.") + endif() + endif() + + if(ENABLE_SAPPHIRERAPIDS) + CHECK_C_COMPILER_FLAG("-march=sapphirerapids" COMPILER_SUPPORT_SAPPHIRERAPIDS) + if(COMPILER_SUPPORT_SAPPHIRERAPIDS) + _AppendFlags(CMAKE_C_FLAGS -march=sapphirerapids) + _AppendFlags(CMAKE_CXX_FLAGS -march=sapphirerapids) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=sapphirerapids\'.") + endif() + endif() + + if(ENABLE_EMERALDRAPIDS) + CHECK_C_COMPILER_FLAG("-march=emeraldrapids" COMPILER_SUPPORT_EMERALDRAPID) + if(COMPILER_SUPPORT_EMERALDRAPID) + _AppendFlags(CMAKE_C_FLAGS -march=emeraldrapids) + _AppendFlags(CMAKE_CXX_FLAGS -march=emeraldrapids) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=emeraldrapids\'.") + endif() + endif() + + if(ENABLE_GRANITERAPIDS) + CHECK_C_COMPILER_FLAG("-march=graniterapids" COMPILER_SUPPORT_GRANITERAPIDS) + if(COMPILER_SUPPORT_GRANITERAPIDS) + _AppendFlags(CMAKE_C_FLAGS -march=graniterapids) + _AppendFlags(CMAKE_CXX_FLAGS -march=graniterapids) + else() + message(SEND_ERROR "Compiler cannot support flag \'-march=graniterapids\'.") + endif() + endif() + + if(ENABLE_AVX512BW) + CHECK_C_COMPILER_FLAG("-mavx512bw" COMPILER_SUPPORT_AVX512BW) + if(COMPILER_SUPPORT_AVX512BW) + _AppendFlags(CMAKE_C_FLAGS -mavx512bw) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512bw) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512bw\'.") + endif() + endif() + + if(ENABLE_AVX512CD) + CHECK_C_COMPILER_FLAG("-mavx512cd" COMPILER_SUPPORT_AVX512CD) + if(COMPILER_SUPPORT_AVX512CD) + _AppendFlags(CMAKE_C_FLAGS -mavx512cd) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512cd) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512cd\'.") + endif() + endif() + + if(ENABLE_AVX512DQ) + CHECK_C_COMPILER_FLAG("-mavx512dq" COMPILER_SUPPORT_AVX512DQ) + if(COMPILER_SUPPORT_AVX512DQ) + _AppendFlags(CMAKE_C_FLAGS -mavx512dq) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512dq) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512dq\'.") + endif() + endif() + + if(ENABLE_AVX512ER) + CHECK_C_COMPILER_FLAG("-mavx512er" COMPILER_SUPPORT_AVX512ER) + if(COMPILER_SUPPORT_AVX512ER) + _AppendFlags(CMAKE_C_FLAGS -mavx512er) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512er) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512er\'.") + endif() + endif() + + if(ENABLE_AVX512IFMA52) + CHECK_C_COMPILER_FLAG("-mavx512ifma52" COMPILER_SUPPORT_AVX512IFMA52) + if(COMPILER_SUPPORT_AVX512IFMA52) + _AppendFlags(CMAKE_C_FLAGS -mavx512ifma52) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512ifma52) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512ifma52\'.") + endif() + endif() + + if(ENABLE_AVX512PF) + CHECK_C_COMPILER_FLAG("-mavx512pf" COMPILER_SUPPORT_AVX512PF) + if(COMPILER_SUPPORT_AVX512PF) + _AppendFlags(CMAKE_C_FLAGS -mavx512pf) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512pf) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512pf\'.") + endif() + endif() + + if(ENABLE_AVX512VL) + CHECK_C_COMPILER_FLAG("-mavx512vl" COMPILER_SUPPORT_AVX512VL) + if(COMPILER_SUPPORT_AVX512VL) + _AppendFlags(CMAKE_C_FLAGS -mavx512vl) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512vl) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512vl\'.") + endif() + endif() + + if(ENABLE_AVX512VPOPCNTDQ) + CHECK_C_COMPILER_FLAG("-mavx512vpopcntdq" COMPILER_SUPPORT_AVX512VPOPCNTDQ) + if(COMPILER_SUPPORT_AVX512VPOPCNTDQ) + _AppendFlags(CMAKE_C_FLAGS -mavx512vpopcntdq) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512vpopcntdq) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512vpopcntdq\'.") + endif() + endif() + + if(ENABLE_AVX5124FMAPS) + CHECK_C_COMPILER_FLAG("-mavx5124fmaps" COMPILER_SUPPORT_AVX5124FMAPS) + if(COMPILER_SUPPORT_AVX5124FMAPS) + _AppendFlags(CMAKE_C_FLAGS -mavx5124fmaps) + _AppendFlags(CMAKE_CXX_FLAGS -mavx5124fmaps) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx5124fmaps\'.") + endif() + endif() + + if(ENABLE_AVX5124VNNIW) + CHECK_C_COMPILER_FLAG("-mavx5124vnniw" COMPILER_SUPPORT_AVX5124VNNIW) + if(COMPILER_SUPPORT_AVX5124VNNIW) + _AppendFlags(CMAKE_C_FLAGS -mavx5124vnniw) + _AppendFlags(CMAKE_CXX_FLAGS -mavx5124vnniw) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx5124vnniw\'.") + endif() + endif() + + if(ENABLE_AVX512BF16) + CHECK_C_COMPILER_FLAG("-mavx512bf16" COMPILER_SUPPORT_AVX512BF16) + if(COMPILER_SUPPORT_AVX512BF16) + _AppendFlags(CMAKE_C_FLAGS -mavx512bf16) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512bf16) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512bf16\'.") + endif() + endif() + + if(ENABLE_AVX512BITALG) + CHECK_C_COMPILER_FLAG("-mavx512bitalg" COMPILER_SUPPORT_AVX512BITALG) + if(COMPILER_SUPPORT_AVX512BITALG) + _AppendFlags(CMAKE_C_FLAGS -mavx512bitalg) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512bitalg) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512bitalg\'.") + endif() + endif() + + if(ENABLE_AVX512VBMI) + CHECK_C_COMPILER_FLAG("-mavx512vbmi" COMPILER_SUPPORT_AVX512VBMI) + if(COMPILER_SUPPORT_AVX512VBMI) + _AppendFlags(CMAKE_C_FLAGS -mavx512vbmi) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512vbmi) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512vbmi\'.") + endif() + endif() + + if(ENABLE_AVX512VBMI2) + CHECK_C_COMPILER_FLAG("-mavx512vbmi2" COMPILER_SUPPORT_AVX512VBMI2) + if(COMPILER_SUPPORT_AVX512VBMI2) + _AppendFlags(CMAKE_C_FLAGS -mavx512vbmi2) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512vbmi2) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512vbmi2\'.") + endif() + endif() + + if(ENABLE_AVX512VNNI) + CHECK_C_COMPILER_FLAG("-mavx512vnni" COMPILER_SUPPORT_AVX512VNNI) + if(COMPILER_SUPPORT_AVX512VNNI) + _AppendFlags(CMAKE_C_FLAGS -mavx512vnni) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512vnni) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512vnni\'.") + endif() + endif() + + if(ENABLE_AVX512VP2INTERSECT) + CHECK_C_COMPILER_FLAG("-mavx512vp2intersect" COMPILER_SUPPORT_AVX512VP2INTERSECT) + if(COMPILER_SUPPORT_AVX512VP2INTERSECT) + _AppendFlags(CMAKE_C_FLAGS -mavx512vp2intersect) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512vp2intersect) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512vp2intersect\'.") + endif() + endif() + + if(ENABLE_AVX512FP16) + CHECK_C_COMPILER_FLAG("-mavx512fp16" COMPILER_SUPPORT_AVX512FP16) + if(COMPILER_SUPPORT_AVX512FP16) + _AppendFlags(CMAKE_C_FLAGS -mavx512fp16) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512fp16) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512fp16\'.") + endif() + endif() + + if(ENABLE_AVXVNNI) + CHECK_C_COMPILER_FLAG("-mavxvnni" COMPILER_SUPPORT_AVXVNNI) + if(COMPILER_SUPPORT_AVXVNNI) + _AppendFlags(CMAKE_C_FLAGS -mavxvnni) + _AppendFlags(CMAKE_CXX_FLAGS -mavxvnni) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavxvnni\'.") + endif() + endif() + + if(ENABLE_AVX512F) + CHECK_C_COMPILER_FLAG("-mavx512f" COMPILER_SUPPORT_AVX512F) + if(COMPILER_SUPPORT_AVX512F) + _AppendFlags(CMAKE_C_FLAGS -mavx512f) + _AppendFlags(CMAKE_CXX_FLAGS -mavx512f) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx512f\'.") + endif() + elseif(ENABLE_AVX2) + CHECK_C_COMPILER_FLAG("-mavx2" COMPILER_SUPPORT_AVX2) + if(COMPILER_SUPPORT_AVX2) + _AppendFlags(CMAKE_C_FLAGS -mavx2) + _AppendFlags(CMAKE_CXX_FLAGS -mavx2) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx2\'.") + endif() + elseif(ENABLE_AVX) + CHECK_C_COMPILER_FLAG("-mavx" COMPILER_SUPPORT_AVX) + if(COMPILER_SUPPORT_AVX) + _AppendFlags(CMAKE_C_FLAGS -mavx) + _AppendFlags(CMAKE_CXX_FLAGS -mavx) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mavx\'.") + endif() + elseif(ENABLE_SSE4.2) + CHECK_C_COMPILER_FLAG("-msse4.2" COMPILER_SUPPORT_SSE4.2) + if(COMPILER_SUPPORT_SSE4.2) + _AppendFlags(CMAKE_C_FLAGS -msse4.2) + _AppendFlags(CMAKE_CXX_FLAGS -msse4.2) + else() + message(SEND_ERROR "Compiler cannot support flag \'-msse4.2\'.") + endif() + elseif(ENABLE_SSE4.1) + CHECK_C_COMPILER_FLAG("-msse4.1" COMPILER_SUPPORT_SSE4.1) + if(COMPILER_SUPPORT_SSE4.1) + _AppendFlags(CMAKE_C_FLAGS -msse4.1) + _AppendFlags(CMAKE_CXX_FLAGS -msse4.1) + else() + message(SEND_ERROR "Compiler cannot support flag \'-msse4.1\'.") + endif() + elseif(ENABLE_SSSE3) + CHECK_C_COMPILER_FLAG("-mssse3" COMPILER_SUPPORT_SSSE3) + if(COMPILER_SUPPORT_SSSE3) + _AppendFlags(CMAKE_C_FLAGS -mssse3) + _AppendFlags(CMAKE_CXX_FLAGS -mssse3) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mssse3\'.") + endif() + elseif(ENABLE_SSE3) + CHECK_C_COMPILER_FLAG("-msse3" COMPILER_SUPPORT_SSE3) + if(COMPILER_SUPPORT_SSE3) + _AppendFlags(CMAKE_C_FLAGS -msse3) + _AppendFlags(CMAKE_CXX_FLAGS -msse3) + else() + message(SEND_ERROR "Compiler cannot support flag \'-msse3\'.") + endif() + elseif(ENABLE_SSE2) + CHECK_C_COMPILER_FLAG("-msse2" COMPILER_SUPPORT_SSE2) + if(COMPILER_SUPPORT_SSE2) + _AppendFlags(CMAKE_C_FLAGS -msse2) + _AppendFlags(CMAKE_CXX_FLAGS -msse2) + else() + message(SEND_ERROR "Compiler cannot support flag \'-msse2\'.") + endif() + elseif(ENABLE_SSE) + CHECK_C_COMPILER_FLAG("-msse" COMPILER_SUPPORT_SSE) + if(COMPILER_SUPPORT_SSE) + _AppendFlags(CMAKE_C_FLAGS -msse) + _AppendFlags(CMAKE_CXX_FLAGS -msse) + else() + message(SEND_ERROR "Compiler cannot support flag \'-msse\'.") + endif() + endif() + + if(ENABLE_FMA) + CHECK_C_COMPILER_FLAG("-mfma" COMPILER_SUPPORT_FMA) + if(COMPILER_SUPPORT_FMA) + _AppendFlags(CMAKE_C_FLAGS -mfma) + _AppendFlags(CMAKE_CXX_FLAGS -mfma) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mfma\'.") + endif() + endif() + + if(ENABLE_F16C) + CHECK_C_COMPILER_FLAG("-mf16c" COMPILER_SUPPORT_F16C) + if(COMPILER_SUPPORT_F16C) + _AppendFlags(CMAKE_C_FLAGS -mf16c) + _AppendFlags(CMAKE_CXX_FLAGS -mf16c) + else() + message(SEND_ERROR "Compiler cannot support flag \'-mf16c\'.") + endif() + endif() +else() + if(ENABLE_AVX512F OR ENABLE_AVX512CD OR ENABLE_AVX512BW OR ENABLE_AVX512DQ OR ENABLE_AVX512VL) + CHECK_C_COMPILER_FLAG("/arch:AVX512" COMPILER_SUPPORT_AVX512F) + if(COMPILER_SUPPORT_AVX512F) + _AppendFlags(CMAKE_C_FLAGS /arch:AVX512) + _AppendFlags(CMAKE_CXX_FLAGS /arch:AVX512) + else() + message(SEND_ERROR "Compiler cannot support flag \'/arch:AVX512\'.") + endif() + elseif(ENABLE_AVX2) + CHECK_C_COMPILER_FLAG("/arch:AVX2" COMPILER_SUPPORT_AVX2) + if(COMPILER_SUPPORT_AVX2) + _AppendFlags(CMAKE_C_FLAGS /arch:AVX2) + _AppendFlags(CMAKE_CXX_FLAGS /arch:AVX2) + else() + message(SEND_ERROR "Compiler cannot support flag \'/arch:AVX2\'.") + endif() + elseif(ENABLE_AVX) + CHECK_C_COMPILER_FLAG("/arch:AVX" COMPILER_SUPPORT_AVX) + if(COMPILER_SUPPORT_AVX) + _AppendFlags(CMAKE_C_FLAGS /arch:AVX) + _AppendFlags(CMAKE_CXX_FLAGS /arch:AVX) + else() + message(SEND_ERROR "Compiler cannot support flag \'/arch:AVX\'.") + endif() + elseif(ENABLE_SSE4.2 OR ENABLE_SSE4.1 OR ENABLE_SSE3 OR ENABLE_SSSE3 OR ENABLE_SSE2) + CHECK_C_COMPILER_FLAG("/arch:SSE2" COMPILER_SUPPORT_SSE2) + if(COMPILER_SUPPORT_SSE2) + _AppendFlags(CMAKE_C_FLAGS /arch:SSE2) + _AppendFlags(CMAKE_CXX_FLAGS /arch:SSE2) + endif() + elseif(ENABLE_SSE) + CHECK_C_COMPILER_FLAG("/arch:SSE" COMPILER_SUPPORT_SSE) + if(COMPILER_SUPPORT_SSE) + _AppendFlags(CMAKE_C_FLAGS /arch:SSE) + _AppendFlags(CMAKE_CXX_FLAGS /arch:SSE) + endif() + endif() +endif() + +# OpenMP support +if(ENABLE_OPENMP) + find_package(OpenMP REQUIRED) + if(OPENMP_FOUND) + _AppendFlags(CMAKE_C_FLAGS ${OpenMP_C_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..faccff0 --- /dev/null +++ b/cmake/utils.cmake @@ -0,0 +1,28 @@ +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 patch -p1 -i "${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/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9f2ef57 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,242 @@ +###################################################################################################### +# 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" +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", + "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/graphscope/zvec" +Repository = "https://github.com/graphscope/zvec" +"Bug Tracker" = "https://github.com/graphscope/zvec/issues" +"Documentation" = "https://zvec.org" + +[project.optional-dependencies] +test = [ + "pytest >=8.0", + "pytest-cov >=4.1", + "pytest-mock >=3.12", + "cibuildwheel == 2.17.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 == 2.17.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"] +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.15,<4.0" +ninja.version = ">=1.11" +cmake.build-type = "Release" +install.strip = true # Strip symbols in release builds to reduce wheel size + +# 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" +ENABLE_SKYLAKE_AVX512 = { env = "ENABLE_SKYLAKE_AVX512", default = "OFF" } +ENABLE_SAPPHIRERAPIDS = { env = "ENABLE_SAPPHIRERAPIDS", default = "OFF" } + +###################################################################################################### +# TESTING & QUALITY +###################################################################################################### +[tool.pytest.ini_options] +minversion = "8.0" +addopts = [ + "-ra", + "--showlocals", + "--strict-markers", + "--strict-config", + "--tb=short", +] +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-*", +] +build-frontend = "build" +test-requires = ["pytest", "numpy"] +build-verbosity = 1 + +[tool.cibuildwheel.linux] +archs = ["x86_64"] +test-command = "cd {project} && pytest python/tests -v --tb=short" +manylinux-x86_64-image = "manylinux_2_28" +skip = "*musllinux*" + +[tool.cibuildwheel.macos] +archs = ["arm64"] +environment = { MACOSX_DEPLOYMENT_TARGET = "11.0" } +###################################################################################################### +# CODE QUALITY & FORMATTING (Ruff) +###################################################################################################### +[tool.ruff] +target-version = "py310" +line-length = 88 +exclude = [ + "build/", + "dist/", + ".git/", + ".venv/", + "venv/", +] + +[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 +] +fixable = ["ALL"] +unfixable = [] + +[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__) +] + +[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..263107d --- /dev/null +++ b/python/tests/detail/distance_helper.py @@ -0,0 +1,302 @@ +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: + vec1 = [np.float16(a) for a in vec1] + vec2 = [np.float16(b) for b in vec2] + + 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 0.0 + + return 1 - dot_product / (magnitude1 * magnitude2) + + +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: + return sum(np.float16(a) * np.float16(b) for a, b in zip(vec1, vec2)) + 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: + return sum((np.float16(a) - np.float16(b)) ** 2 for a, b in zip(vec1, vec2)) + return sum((a - b) ** 2 for a, b in zip(vec1, vec2)) + + +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()): + 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 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..f720b23 --- /dev/null +++ b/python/tests/detail/doc_helper.py @@ -0,0 +1,377 @@ +from zvec import CollectionSchema, Doc + +from support_helper import * + +import numpy as np +from typing import Literal, Optional, Union, Tuple + +import random +import string + + +def generate_constant_vector( + i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32" +): + if dtype == "int8": + vec = [i % 128] * dimension + vec[i % dimension] = (i + 1) % 128 + else: + vec = [i / 256.0] * dimension + vec[i % dimension] = (i + 1) / 256.0 + + 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_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_update_doc(i: int, schema: CollectionSchema) -> Doc: + 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) + 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) + 1, (i + 1) + 2] + elif field.data_type == DataType.ARRAY_UINT32: + doc_fields[field.name] = [i + 1, (i + 1) + 1, (i + 1) + 2] + elif field.data_type == DataType.ARRAY_INT64: + doc_fields[field.name] = [i + 1, (i + 1) + 1, (i + 1) + 2] + elif field.data_type == DataType.ARRAY_UINT64: + doc_fields[field.name] = [i + 1, (i + 1) + 1, (i + 1) + 2] + elif field.data_type == DataType.ARRAY_FLOAT: + doc_fields[field.name] = [ + float((i + 1) + 0.1), + float((i + 1) + 1.1), + float((i + 1) + 2.1), + ] + elif field.data_type == DataType.ARRAY_DOUBLE: + doc_fields[field.name] = [ + float((i + 1) + 0.11), + float((i + 1) + 1.11), + float((i + 1) + 2.11), + ] + elif field.data_type == DataType.ARRAY_STRING: + doc_fields[field.name] = [ + f"test_{i + 1}", + f"test_{(i + 1) + 1}", + f"test_{(i + 1) + 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 + 1, DEFAULT_VECTOR_DIMENSION, "float16" + ) + elif vector.data_type == DataType.VECTOR_FP32: + doc_vectors[vector.name] = generate_constant_vector( + i + 1, DEFAULT_VECTOR_DIMENSION, "float32" + ) + elif vector.data_type == DataType.VECTOR_INT8: + doc_vectors[vector.name] = generate_constant_vector( + i + 1, + DEFAULT_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}") + 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), 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}") + + 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..272b44e --- /dev/null +++ b/python/tests/detail/fixture_helper.py @@ -0,0 +1,319 @@ +import pytest +import logging + +from typing import Any, Generator + +import zvec +from zvec import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + IVFIndexParam, + FieldSchema, + VectorSchema, + CollectionSchema, + Collection, + Doc, + VectorQuery, +) + +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() + + 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(): + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + + 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 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 +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..dcfffd7 --- /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 + +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..af175b6 --- /dev/null +++ b/python/tests/detail/test_collection_create_and_open.py @@ -0,0 +1,792 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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", +] +invalid_path_list = [ + "invalid:path", + "", + "test_collection_with_spaces ", + "test@#$%collection", +] + + +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..6fba8cb --- /dev/null +++ b/python/tests/detail/test_collection_ddl.py @@ -0,0 +1,1557 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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( + VectorQuery(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) diff --git a/python/tests/detail/test_collection_dml.py b/python/tests/detail/test_collection_dml.py new file mode 100644 index 0000000..e4ccad6 --- /dev/null +++ b/python/tests/detail/test_collection_dml.py @@ -0,0 +1,1180 @@ +import logging +import pytest + +from zvec import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + FieldSchema, + VectorSchema, + CollectionSchema, + Collection, + Doc, + VectorQuery, + 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( + VectorQuery(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 == 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( + VectorQuery(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 == doc.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( + VectorQuery(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..8078ac6 --- /dev/null +++ b/python/tests/detail/test_collection_dql.py @@ -0,0 +1,1329 @@ +# Copyright 2025-present the zvec 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 zvec.typing import DataType, StatusCode, MetricType, QuantizeType +from zvec.model import Collection, Doc, VectorQuery +from zvec.model.param import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + FlatIndexParam, + IVFIndexParam, + HnswQueryParam, + IVFQueryParam, +) + + +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 * + + +# ==================== 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( + VectorQuery(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( + VectorQuery(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)}" + ) + + +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(VectorQuery(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( + VectorQuery(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( + filter=filter, + vectors=VectorQuery(field_name=v, vector=query_vector), + topk=topk, + ) + elif topk and not filter: + query_result = full_collection_new.query( + VectorQuery(field_name=v, vector=query_vector), topk=topk + ) + elif not topk and filter: + query_result = full_collection_new.query( + filter=filter, + vectors=VectorQuery(field_name=v, vector=query_vector), + ) + else: + query_result = full_collection_new.query( + VectorQuery(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( + VectorQuery(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( + VectorQuery(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(VectorQuery(field_name=v, vector=doc_vectors[v])) + + rrf_reranker = RrfReRanker(topn=3) + multi_query_result = full_collection.query( + vectors=multi_query_vectors, + 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-10, ( + 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) + + weighted_reranker = WeightedReRanker( + topn=3, weights=weights, metric=MetricType.IP + ) + + single_query_results = {} + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + single_query_results[v] = full_collection.query( + VectorQuery(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(VectorQuery(field_name=v, vector=doc_vectors[v])) + + multi_query_result = full_collection.query( + vectors=multi_query_vectors, + 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-10, ( + 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"]) + 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( + filter=filter, + vectors=VectorQuery( + field_name=v, vector=query_vector, param=HnswQueryParam(ef=ef) + ), + 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( + filter=filter, + vectors=VectorQuery( + field_name=v, vector=query_vector, param=HnswQueryParam(ef=ef) + ), + 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"]) + 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( + filter=filter, + vectors=VectorQuery( + field_name=v, + vector=query_vector, + param=IVFQueryParam(nprobe=nprobe), + ), + 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( + # filter=filter, + vectors=VectorQuery( + field_name=v, + vector=query_vector, + param=IVFQueryParam(nprobe=nprobe), + ), + 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( + filter=filter, + vectors=VectorQuery( + field_name=v, vector=query_vector, param=HnswIndexParam() + ), + ) + 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: VectorQuery( + 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: VectorQuery( + 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: VectorQuery( + 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: VectorQuery( + field_name="vector_fp32_field", vector=[] + ), + "Expected exception for empty dense vector data", + ), + ( + "Invalid dimension for dense vector", + lambda ref_dense_vector: VectorQuery( + 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: VectorQuery( + field_name="vector_fp32_field", id="999" + ), # Non-existent ID + "Expected exception for non-existent document ID", + ), + ( + "Both vector and id specified (invalid combination)", + lambda ref_dense_vector: VectorQuery( + field_name="vector_fp32_field", vector=ref_dense_vector, id="5" + ), + "Expected exception for specifying both vector and id", + ), + ( + "Neither vector nor id specified", + lambda ref_dense_vector: VectorQuery( + 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(vectors=[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( + filter=filter, + vectors=VectorQuery(field_name=v, vector=query_vector), + 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( + VectorQuery( + 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..cfde4f4 --- /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, + VectorSchema, + VectorQuery, +) + + +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_vectorquery_field_name(self, test_collection: Collection): + with pytest.raises(Exception) as exc_info: + result = test_collection.query(vectors=[VectorQuery()]) + assert exc_info.value is not None, ( + "Expected exception for missing VectorQuery 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..a275c92 --- /dev/null +++ b/python/tests/detail/test_collection_open.py @@ -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. + + +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( + VectorQuery(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( + VectorQuery(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(), + ), + ], + 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"}, + 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" + + # 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_db_config.py b/python/tests/detail/test_db_config.py new file mode 100644 index 0000000..87bde65 --- /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_runtime_error(self): + zvec.init() + with pytest.raises(RuntimeError): + 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..7d021d6 --- /dev/null +++ b/python/tests/test_collection.py @@ -0,0 +1,1045 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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, + InvertIndexParam, + LogLevel, + LogType, + VectorSchema, + StatusCode, + IndexOption, + IndexType, + VectorQuery, + OptimizeOption, +) + +# ==================== 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( + "sparse", 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, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}, + ) + + +@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, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}, + ) + 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) == 2 + assert stats.index_completeness["dense"] == 1 + assert stats.index_completeness["sparse"] == 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, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}, + ) + 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, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}, + ) + with pytest.raises(ValueError) as e: + # ValueError: doc validate failed: field[id] is configured not nullable, + # but doc does not contain this field + test_collection.insert(doc) + assert "field[id] is configured not nullable" in str(e.value) + + # without name + doc = Doc( + id="0", + fields={ + "id": 1, + }, + vectors={"dense": [1 + 0.1] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}, + ) + with pytest.raises(ValueError) as e: + test_collection.insert(doc) + assert "field[name] is configured not nullable" 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, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}, + ) + 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: doc validate failed: field[id] is configured not nullable, + # but doc does not contain this field + 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_vector_and_id( + self, collection_with_single_doc: Collection, single_doc: Doc + ): + with pytest.raises(ValueError): + collection_with_single_doc.query( + VectorQuery( + field_name="dense", + id=single_doc.id, + vector=single_doc.vector("dense"), + ) + ) + + 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 = VectorQuery( + field_name="dense", vector=multiple_docs[0].vector("dense"), param=[1, 2, 3] + ) + with pytest.raises(TypeError): + result = collection_with_multiple_docs.query( + filter="id in (1)", topk=100, vectors=query + ) + + def test_collection_query_by_id( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + result = collection_with_multiple_docs.query( + VectorQuery(field_name="dense", id=multiple_docs[0].id) + ) + assert len(result) == 10 + + def test_collection_query_multi_vector_with_same_field( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + with pytest.raises(ValueError): + collection_with_multiple_docs.query( + [ + VectorQuery( + field_name="dense", vector=multiple_docs[0].vector("dense") + ), + VectorQuery( + field_name="dense", vector=multiple_docs[0].vector("dense") + ), + ] + ) + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_by_dense_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_by_sparse_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_by_dense_vector_with_filter( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_by_sparse_vector_with_filter( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_with_rrf_reranker_by_multi_dense_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_with_rrf_reranker_by_multi_sparse_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_with_rrf_reranker_by_hybrid_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_with_weighted_reranker_by_multi_dense_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_with_weighted_reranker_by_multi_sparse_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass + + @pytest.mark.skip(reason="TODO: This test case is pending implementation") + def test_collection_query_with_weighted_reranker_by_hybrid_vector( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + pass diff --git a/python/tests/test_convert.py b/python/tests/test_convert.py new file mode 100644 index 0000000..f30eada --- /dev/null +++ b/python/tests/test_convert.py @@ -0,0 +1,584 @@ +from __future__ import annotations + +import math + +import pytest +from _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..d4a60ff --- /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 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..0eb5d6b --- /dev/null +++ b/python/tests/test_embedding.py @@ -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. +from __future__ import annotations + +import os +from http import HTTPStatus +from unittest.mock import MagicMock, patch + + +import pytest +from zvec.extension import QwenEmbeddingFunction + + +# ---------------------------- +# QwenEmbeddingFunction Test Case +# ---------------------------- +class TestQwenEmbeddingFunction: + def test_init_with_api_key(self): + # Test initialization with explicit API key + embedding_func = QwenEmbeddingFunction(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 = QwenEmbeddingFunction(dimension=128) + assert embedding_func._api_key == "env_key" + + def test_init_without_api_key(self): + # Test initialization without API key raises ValueError + with pytest.raises(ValueError, match="DashScope API key is required"): + QwenEmbeddingFunction(dimension=128) + + @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"): + QwenEmbeddingFunction(dimension=128) + + def test_model_property(self): + embedding_func = QwenEmbeddingFunction(dimension=128, api_key="test_key") + assert embedding_func.model == "text-embedding-v4" + + embedding_func = QwenEmbeddingFunction( + dimension=128, model="custom-model", api_key="test_key" + ) + assert embedding_func.model == "custom-model" + + @patch("zvec.extension.embedding.require_module") + def test_embed_with_empty_text(self, mock_require_module): + # Test embed method with empty text raises ValueError + embedding_func = QwenEmbeddingFunction(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.embedding.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 = QwenEmbeddingFunction(dimension=128, api_key="test_key") + 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=128, + output_type="dense", + ) + + @patch("zvec.extension.embedding.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 = QwenEmbeddingFunction(dimension=128, api_key="test_key") + + with pytest.raises(ValueError): + embedding_func.embed("test text") + + @patch("zvec.extension.embedding.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 = QwenEmbeddingFunction(dimension=128, api_key="test_key") + + with pytest.raises(ValueError): + embedding_func.embed("test text") + + @pytest.mark.skip(reason="Qwen Embedding is not available in CI") + def test_embed(self): + # Test embedding with invalid dimension + embedding_func = QwenEmbeddingFunction(dimension=128, api_key="xxx") + dense = embedding_func("test text") + assert len(dense) == 128 diff --git a/python/tests/test_params.py b/python/tests/test_params.py new file mode 100644 index 0000000..b067a25 --- /dev/null +++ b/python/tests/test_params.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. +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, + VectorQuery, + IndexType, + MetricType, + QuantizeType, + DataType, + VectorSchema, +) + +from _zvec.param import _VectorQuery + +# ---------------------------- +# 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 == 100 + 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 == 0 + 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 + + def test_custom(self): + param = HnswQueryParam(ef=10, is_using_refiner=True, radius=30, is_linear=True) + assert param.ef == 10 + assert param.is_using_refiner == True + assert param.radius == 30 + assert param.is_linear == True + + 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 TestVectorQuery: + def test_init_with_valid_id(self): + vq = VectorQuery(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 = VectorQuery(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): + VectorQuery(field_name="embedding", id="doc123", vector=[0.1])._validate() + + def test_init_without_field_name_raises_error(self): + with pytest.raises(ValueError): + VectorQuery(field_name=None)._validate() + + def test_has_id_returns_true_when_id_set(self): + vq = VectorQuery(field_name="embedding", id="doc123") + assert vq.has_id() + + def test_has_id_returns_false_when_no_id(self): + vq = VectorQuery(field_name="embedding", vector=[0.1]) + assert not vq.has_id() + + def test_has_vector_returns_true_with_non_empty_vector(self): + vq = VectorQuery(field_name="embedding", vector=[0.1]) + assert vq.has_vector() + + def test_validate_fails_on_both_id_and_vector(self): + vq = VectorQuery(field_name="test", id="doc123", vector=[0.1]) + with pytest.raises(ValueError): + vq._validate() diff --git a/python/tests/test_query_executor.py b/python/tests/test_query_executor.py new file mode 100644 index 0000000..53666c1 --- /dev/null +++ b/python/tests/test_query_executor.py @@ -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. +from __future__ import annotations + +from typing import Dict, Union +from unittest.mock import MagicMock + +import numpy as np +import math +from _zvec.param import _VectorQuery + +import pytest +from zvec.executor.query_executor import ( + MultiVectorQueryExecutor, + NoVectorQueryExecutor, + QueryContext, + QueryExecutor, + QueryExecutorFactory, + SingleVectorQueryExecutor, + VectorQuery, +) +from zvec import RrfReRanker, HnswQueryParam, CollectionSchema, VectorSchema, DataType + + +# ---------------------------- +# Mock Vector Schema +# ---------------------------- +class MockVectorSchema(VectorSchema): + def __init__(self, name="test_vector"): + self._name = name + + @property + def name(self): + return self._name + + def _get_object(self): + return MagicMock() + + +# ---------------------------- +# 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 TestVectorQuery: + def test_init(self): + query = VectorQuery(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 = VectorQuery( + 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 = VectorQuery(field_name="test_field") + assert not query.has_id() + + query = VectorQuery(field_name="test_field", id="test_id") + assert query.has_id() + + def test_has_vector(self): + query = VectorQuery(field_name="test_field") + assert not query.has_vector() + + query = VectorQuery(field_name="test_field", vector=[]) + assert not query.has_vector() + + query = VectorQuery(field_name="test_field", vector=[1, 2, 3]) + assert query.has_vector() + + def test_validate_dense_fp16_convert(self): + v = _VectorQuery() + 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 = _VectorQuery() + 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 = _VectorQuery() + 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 = _VectorQuery() + 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 = _VectorQuery() + 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 = _VectorQuery() + 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 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 + assert ctx.core_vectors == [] + + def test_properties(self): + queries = [VectorQuery(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_core_vectors_setter(self): + ctx = QueryContext(topk=10) + core_vectors = [MagicMock()] + ctx.core_vectors = core_vectors + assert ctx.core_vectors == core_vectors + + +class TestNoVectorQueryExecutor: + def test_init(self): + schema = MockCollectionSchema() + executor = NoVectorQueryExecutor(schema) + assert isinstance(executor, QueryExecutor) + + def test_do_validate_with_queries(self): + schema = MockCollectionSchema() + executor = NoVectorQueryExecutor(schema) + ctx = QueryContext(topk=10, queries=[VectorQuery(field_name="test")]) + + with pytest.raises( + ValueError, match="Collection does not support query with vector or id" + ): + executor._do_validate(ctx) + + def test_do_validate_without_queries(self): + schema = MockCollectionSchema() + executor = NoVectorQueryExecutor(schema) + ctx = QueryContext(topk=10) + + executor._do_validate(ctx) + + def test_do_build(self): + schema = MockCollectionSchema() + executor = NoVectorQueryExecutor(schema) + ctx = QueryContext(topk=5, filter="test_filter") + + result = executor._do_build(ctx, MagicMock()) + assert len(result) == 1 + assert result[0].topk == 5 + assert result[0].filter == "test_filter" + + +class TestSingleVectorQueryExecutor: + def test_init(self): + schema = MockCollectionSchema() + executor = SingleVectorQueryExecutor(schema) + assert isinstance(executor, NoVectorQueryExecutor) + + def test_do_validate_multiple_queries(self): + schema = MockCollectionSchema() + executor = SingleVectorQueryExecutor(schema) + queries = [VectorQuery(field_name="test1"), VectorQuery(field_name="test2")] + ctx = QueryContext(topk=10, queries=queries) + + with pytest.raises( + ValueError, + match="Collection has only one vector field, cannot query with multiple vectors", + ): + executor._do_validate(ctx) + + def test_do_build_without_queries(self): + schema = MockCollectionSchema() + executor = SingleVectorQueryExecutor(schema) + ctx = QueryContext(topk=5) + + result = executor._do_build(ctx, MagicMock()) + assert len(result) == 1 + assert result[0].topk == 5 + + +class TestMultiVectorQueryExecutor: + def test_init(self): + schema = MockCollectionSchema() + executor = MultiVectorQueryExecutor(schema) + assert isinstance(executor, SingleVectorQueryExecutor) + + def test_do_validate_multiple_queries_without_reranker(self): + schema = MockCollectionSchema() + executor = MultiVectorQueryExecutor(schema) + queries = [VectorQuery(field_name="test1"), VectorQuery(field_name="test2")] + ctx = QueryContext(topk=10, queries=queries) + + with pytest.raises( + ValueError, match="Reranker is required for multi-vector query" + ): + executor._do_validate(ctx) + + def test_do_validate_multiple_queries_with_reranker(self): + schema = MockCollectionSchema() + executor = MultiVectorQueryExecutor(schema) + queries = [VectorQuery(field_name="test1"), VectorQuery(field_name="test2")] + reranker = RrfReRanker() + ctx = QueryContext(topk=10, queries=queries, reranker=reranker) + + executor._do_validate(ctx) + + +class TestQueryExecutorFactory: + def test_create_no_vectors(self): + schema = MockCollectionSchema() + executor = QueryExecutorFactory.create(schema) + assert isinstance(executor, NoVectorQueryExecutor) + + def test_create_single_vector(self): + schema = MockCollectionSchema(vectors=MockVectorSchema()) + executor = QueryExecutorFactory.create(schema) + assert isinstance(executor, SingleVectorQueryExecutor) + + def test_create_multiple_vectors(self): + schema = MockCollectionSchema( + vectors={"test1": MockVectorSchema(), "test2": MockVectorSchema()} + ) + executor = QueryExecutorFactory.create(schema) + assert isinstance(executor, MultiVectorQueryExecutor) diff --git a/python/tests/test_reranker.py b/python/tests/test_reranker.py new file mode 100644 index 0000000..5b2c177 --- /dev/null +++ b/python/tests/test_reranker.py @@ -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. +from __future__ import annotations + +from unittest.mock import patch +import pytest +import math + +from zvec import RrfReRanker, WeightedReRanker, Doc, MetricType + + +# ---------------------------- +# RrfRanker Test Case +# ---------------------------- +class TestRrfReRanker: + def test_init(self): + reranker = RrfReRanker( + query="test", topn=5, rerank_field="content", rank_constant=100 + ) + assert reranker.query == "test" + assert reranker.topn == 5 + assert reranker.rerank_field == "content" + assert reranker.rank_constant == 100 + + def test_rrf_score(self): + reranker = RrfReRanker(query="test", rank_constant=60) + # 根据公式 1.0 / (k + rank + 1)，其中k=60 + assert reranker._rrf_score(0) == 1.0 / (60 + 0 + 1) + assert reranker._rrf_score(1) == 1.0 / (60 + 1 + 1) + assert reranker._rrf_score(10) == 1.0 / (60 + 10 + 1) + + def test_rerank(self): + reranker = RrfReRanker(query="test", topn=3) + + 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 = {"vector1": [doc1, doc2, doc3], "vector2": [doc3, doc1, doc4]} + + results = reranker.rerank(query_results) + + assert len(results) <= reranker.topn + + for doc in results: + assert hasattr(doc, "score") + + scores = [doc.score for doc in results] + assert scores == sorted(scores, reverse=True) + + +# ---------------------------- +# WeightedRanker Test Case +# ---------------------------- +class TestWeightedReRanker: + def test_init(self): + weights = {"vector1": 0.7, "vector2": 0.3} + reranker = WeightedReRanker( + query="test", + topn=5, + rerank_field="content", + metric=MetricType.L2, + weights=weights, + ) + assert reranker.query == "test" + assert reranker.topn == 5 + assert reranker.rerank_field == "content" + assert reranker.metric == MetricType.L2 + assert reranker.weights == weights + + def test_normalize_score(self): + reranker = WeightedReRanker(query="test") + + score = reranker._normalize_score(1.0, MetricType.L2) + expected = 1.0 - 2 * math.atan(1.0) / math.pi + assert score == expected + + score = reranker._normalize_score(1.0, MetricType.IP) + expected = 0.5 + math.atan(1.0) / math.pi + assert score == expected + + score = reranker._normalize_score(1.0, MetricType.COSINE) + expected = 1.0 - 1.0 / 2.0 + assert score == expected + + with pytest.raises(ValueError, match="Unsupported metric type"): + reranker._normalize_score(1.0, "unsupported_metric") + + def test_rerank(self): + weights = {"vector1": 0.7, "vector2": 0.3} + reranker = WeightedReRanker( + query="test", topn=3, weights=weights, metric=MetricType.L2 + ) + + doc1 = Doc(id="1", score=0.8) + doc2 = Doc(id="2", score=0.7) + doc3 = Doc(id="3", score=0.9) + + query_results = {"vector1": [doc1, doc2], "vector2": [doc2, doc3]} + + results = reranker.rerank(query_results) + + assert len(results) <= reranker.topn + + for doc in results: + assert hasattr(doc, "score") + + scores = [doc.score for doc in results] + assert scores == sorted(scores, reverse=True) + + +# # ---------------------------- +# # QwenReRanker Test Case +# # ---------------------------- +# class TestQwenReRanker: +# def test_init_without_query(self): +# with pytest.raises(ValueError): +# QwenReRanker() +# +# def test_init_without_api_key(self): +# with patch.dict(os.environ, {"DASHSCOPE_API_KEY": ""}): +# 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") +# assert reranker.query == "test" +# assert reranker._api_key == "test_key" +# +# def test_model_property(self): +# reranker = QwenReRanker(query="test", api_key="test_key") +# assert reranker.model == "gte-rerank-v2" +# +# reranker = QwenReRanker(query="test", model="custom-model", api_key="test_key") +# assert reranker.model == "custom-model" +# +# def test_rerank_empty_results(self): +# reranker = QwenReRanker(query="test", api_key="test_key") +# results = reranker.rerank({}) +# assert results == [] +# +# def test_rerank_no_documents(self): +# reranker = QwenReRanker(query="test", api_key="test_key") +# query_results = {"vector1": [Doc(id="1")]} +# with pytest.raises(ValueError, match="No documents to rerank"): +# reranker.rerank(query_results) +# +# @pytest.mark.skip(reason="Qwen ReRanker is not available in CI") +# def test_rerank_success(self): +# reranker = QwenReRanker( +# topn=3, +# query="test", +# api_key="*", +# rerank_field="content", +# ) +# query_results = { +# "vector1": [ +# Doc(id="1", fields={"content": "This is a test document."}), +# Doc(id="2", fields={"content": "Another test document."}), +# Doc(id="3", fields={"content": "Yet another test document."}), +# Doc(id="4", fields={"content": "One more test document."}), +# ], +# "vector2": [ +# Doc(id="5", fields={"content": "This is a test document2."}), +# Doc(id="6", fields={"content": "Another test document2."}), +# Doc(id="7", fields={"content": "Yet another test document2."}), +# Doc(id="8", fields={"content": "One more test document2."}), +# ], +# } +# results = reranker.rerank(query_results) +# assert len(results) == 3 diff --git a/python/tests/test_schema.py b/python/tests/test_schema.py new file mode 100644 index 0000000..f828d73 --- /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 == 100 + 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..bac8926 --- /dev/null +++ b/python/tests/test_util.py @@ -0,0 +1,94 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 + + +def test_require_module_with_openai(): + with pytest.raises(ImportError) as exc_info: + require_module("openai") diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py new file mode 100644 index 0000000..ec35829 --- /dev/null +++ b/python/zvec/__init__.py @@ -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. + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from importlib.metadata import PackageNotFoundError + + +# ============================== +# Public API — grouped by category +# ============================== + +from . import model as model + +# —— Extensions & typing —— +from .extension import DenseEmbeddingFunction, ReRanker, RrfReRanker, WeightedReRanker +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 + +# —— Query & index parameters —— +from .model.param import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + FlatIndexParam, + HnswIndexParam, + HnswQueryParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + IVFQueryParam, + OptimizeOption, +) +from .model.param.vector_query import 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", + # Core classes + "Collection", + "Doc", + # Schema + "CollectionSchema", + "FieldSchema", + "VectorSchema", + "CollectionStats", + # Parameters + "VectorQuery", + "InvertIndexParam", + "HnswIndexParam", + "FlatIndexParam", + "IVFIndexParam", + "CollectionOption", + "IndexOption", + "OptimizeOption", + "AddColumnOption", + "AlterColumnOption", + "HnswQueryParam", + "IVFQueryParam", + # Extensions + "ReRanker", + "DenseEmbeddingFunction", + "RrfReRanker", + "WeightedReRanker", + # Typing + "DataType", + "MetricType", + "QuantizeType", + "IndexType", + "LogLevel", + "LogType", + "Status", + "StatusCode", + # Tools + "require_module", +] + +# ============================== +# 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..7e536c7 --- /dev/null +++ b/python/zvec/__init__.pyi @@ -0,0 +1,172 @@ +""" +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 +from .model.param import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + FlatIndexParam, + HnswIndexParam, + HnswQueryParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + IVFQueryParam, + OptimizeOption, +) +from .model.param.vector_query import 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", + "DenseEmbeddingFunction", + "Doc", + "FieldSchema", + "FlatIndexParam", + "HnswIndexParam", + "HnswQueryParam", + "IVFIndexParam", + "IVFQueryParam", + "IndexOption", + "IndexType", + "InvertIndexParam", + "LogLevel", + "LogType", + "MetricType", + "OptimizeOption", + "QuantizeType", + "ReRanker", + "ReRanker", + "RrfReRanker", + "Status", + "StatusCode", + "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: str, + arg1: schema._FieldSchema, + arg2: str, + arg3: 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, arg0: collections.abc.Sequence[str]) -> 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._VectorQuery) -> 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 __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..56b82fd --- /dev/null +++ b/python/zvec/common/constants.py @@ -0,0 +1,22 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 + +DenseVectorType = Union[list[float], list[int], np.ndarray] +SparseVectorType = dict[int, float] +VectorType = Optional[Union[DenseVectorType, SparseVectorType]] diff --git a/python/zvec/executor/__init__.py b/python/zvec/executor/__init__.py new file mode 100644 index 0000000..2ed019e --- /dev/null +++ b/python/zvec/executor/__init__.py @@ -0,0 +1,26 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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, + QueryExecutorFactory, +) + +__all__ = [ + "QueryContext", + "QueryExecutor", + "QueryExecutorFactory", +] diff --git a/python/zvec/executor/query_executor.py b/python/zvec/executor/query_executor.py new file mode 100644 index 0000000..fa9705c --- /dev/null +++ b/python/zvec/executor/query_executor.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. +from __future__ import annotations + +import os +from abc import ABC, abstractmethod +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Optional, Union, final + +import numpy as np +from _zvec import _Collection +from _zvec.param import _VectorQuery + +from ..extension import ReRanker, RrfReRanker, WeightedReRanker +from ..model.convert import convert_to_py_doc +from ..model.doc import Doc +from ..model.param.vector_query import VectorQuery +from ..model.schema import CollectionSchema +from ..typing import DataType + +__all__ = [ + "QueryContext", + "QueryExecutor", + "QueryExecutorFactory", +] + +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[VectorQuery]] = 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 + + # core vectors + self._core_vectors = [] + + @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 + + @property + def core_vectors(self): + return self._core_vectors + + @core_vectors.setter + def core_vectors(self, core_vectors: list[_VectorQuery]): + self._core_vectors = core_vectors + + +class QueryExecutor(ABC): + def __init__(self, schema: CollectionSchema): + self._schema = schema + self._concurrency = max(1, int(os.getenv("ZVEC_QUERY_CONCURRENCY", "1"))) + + @abstractmethod + def _do_validate(self, ctx: QueryContext) -> None: + pass + + @abstractmethod + def _do_build( + self, ctx: QueryContext, collection: _Collection + ) -> list[_VectorQuery]: + pass + + def _do_build_query_wo_vector(self, ctx: QueryContext) -> _VectorQuery: + core_vector = _VectorQuery() + core_vector.topk = ctx.topk + core_vector.include_vector = ctx.include_vector + if ctx.filter: + core_vector.filter = ctx.filter + if ctx.output_fields: + core_vector.output_fields = ctx.output_fields + return core_vector + + def _do_build_query_with_vector( + self, ctx: QueryContext, query: VectorQuery, collection: _Collection + ) -> _VectorQuery: + core_vector = self._do_build_query_wo_vector(ctx) + core_vector.field_name = query.field_name + if query.param: + core_vector.query_params = query.param + + 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 output_fields + core_vector.output_fields = ctx.output_fields + + # set vector + if query.has_vector(): + vec_data = query.vector + else: + fetched = collection.Fetch([query.id]) + doc = next(iter(fetched.values())) + if not doc: + return core_vector + vec_data = doc.get_any(vector_schema.name, vector_schema.data_type) + + target_dtype = DTYPE_MAP.get(vector_schema.data_type.value) + core_vector.set_vector( + vector_schema._get_object(), + convert_to_numpy(vec_data, target_dtype) if target_dtype else vec_data, + ) + return core_vector + + def _do_execute( + self, vectors: list[_VectorQuery], collection: _Collection + ) -> dict[str, list[Doc]]: + query_cnt = len(vectors) + if query_cnt == 0: + raise ValueError("No query to execute") + + if len(vectors) == 1 or self._concurrency == 1: + results = {} + for query in vectors: + docs = collection.Query(query) + results[query.field_name] = [ + convert_to_py_doc(doc, self._schema) for doc in docs + ] + return results + + results = {} + with ThreadPoolExecutor(max_workers=self._concurrency) as executor: + future_to_query = { + executor.submit(collection.Query, query): query.field_name + for query in vectors + } + + for future in as_completed(future_to_query): + field_name = future_to_query[future] + try: + docs = future.result() + results[field_name] = [ + convert_to_py_doc(doc, self._schema) for doc in docs + ] + except Exception as e: + raise e + return results + + def _do_merge_rerank_results( + self, ctx: QueryContext, docs_map: dict[str, list[Doc]] + ) -> list[Doc]: + query_result_cnt = len(docs_map) if docs_map else 0 + if query_result_cnt == 0: + raise ValueError("Query results is none and dost not to rerank") + if query_result_cnt == 1: + if not ctx.reranker or isinstance( + ctx.reranker, (RrfReRanker, WeightedReRanker) + ): + return next(iter(docs_map.values())) + return ctx.reranker.rerank(docs_map) + return ctx.reranker.rerank(docs_map) + + @final + def execute(self, ctx: QueryContext, collection: _Collection) -> list[Doc]: + # 1. validate query + self._do_validate(ctx) + # 2. build query vector + query_vectors = self._do_build(ctx, collection) + if not query_vectors: + raise ValueError("No query to execute") + # 3. execute query + docs = self._do_execute(query_vectors, collection) + # 4. merge and rerank result + return self._do_merge_rerank_results(ctx, docs) + + +class NoVectorQueryExecutor(QueryExecutor): + def __init__(self, schema: CollectionSchema): + super().__init__(schema) + + def _do_validate(self, ctx: QueryContext) -> None: + if len(ctx.queries) > 0: + raise ValueError("Collection does not support query with vector or id") + + def _do_build( + self, ctx: QueryContext, _collection: _Collection + ) -> list[_VectorQuery]: + return [self._do_build_query_wo_vector(ctx)] + + +class SingleVectorQueryExecutor(NoVectorQueryExecutor): + def __init__(self, schema: CollectionSchema) -> None: + super().__init__(schema) + + def _do_validate(self, ctx: QueryContext) -> None: + if len(ctx.queries) > 1: + raise ValueError( + "Collection has only one vector field, cannot query with multiple vectors" + ) + for query in ctx.queries: + query._validate() + + def _do_build( + self, ctx: QueryContext, collection: _Collection + ) -> list[_VectorQuery]: + if len(ctx.queries) == 0: + return [self._do_build_query_wo_vector(ctx)] + vectors = [] + for query in ctx.queries: + vectors.append(self._do_build_query_with_vector(ctx, query, collection)) + return vectors + + +class MultiVectorQueryExecutor(SingleVectorQueryExecutor): + def __init__(self, schema: CollectionSchema) -> None: + super().__init__(schema) + + def _do_validate(self, ctx: QueryContext) -> None: + if len(ctx.queries) > 1 and ctx.reranker is None: + raise ValueError("Reranker is required for multi-vector query") + seen_fields = set() + for query in ctx.queries: + query._validate() + field = query.field_name + if field in seen_fields: + raise ValueError(f"Query field name '{field}' appears more than once") + seen_fields.add(field) + + def _do_execute( + self, vectors: list[_VectorQuery], collection: _Collection + ) -> dict[str, list[Doc]]: + return super()._do_execute(vectors, collection) + + +class QueryExecutorFactory: + @staticmethod + def create(schema: CollectionSchema) -> QueryExecutor: + vectors = schema.vectors + if len(vectors) == 0: + return NoVectorQueryExecutor(schema) + if len(vectors) == 1: + return SingleVectorQueryExecutor(schema) + return MultiVectorQueryExecutor(schema) diff --git a/python/zvec/extension/__init__.py b/python/zvec/extension/__init__.py new file mode 100644 index 0000000..83421b5 --- /dev/null +++ b/python/zvec/extension/__init__.py @@ -0,0 +1,26 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 .embedding import DenseEmbeddingFunction, QwenEmbeddingFunction +from .rerank import QwenReRanker, ReRanker, RrfReRanker, WeightedReRanker + +__all__ = [ + "DenseEmbeddingFunction", + "QwenEmbeddingFunction", + "QwenReRanker", + "ReRanker", + "RrfReRanker", + "WeightedReRanker", +] diff --git a/python/zvec/extension/embedding.py b/python/zvec/extension/embedding.py new file mode 100644 index 0000000..1bbb096 --- /dev/null +++ b/python/zvec/extension/embedding.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. +from __future__ import annotations + +import os +from abc import ABC, abstractmethod +from functools import lru_cache +from http import HTTPStatus +from typing import Optional, Union + +from ..tool import require_module +from ..typing import DataType + + +class DenseEmbeddingFunction(ABC): + """Abstract base class for dense vector embedding functions. + + Dense embedding functions map text to fixed-length real-valued vectors. + Subclasses must implement the ``embed()`` method. + + Args: + dimension (int): Dimensionality of the output embedding vector. + data_type (DataType, optional): Numeric type of the embedding. + Defaults to ``DataType.VECTOR_FP32``. + + Note: + This class is callable: ``embedding_func("text")`` is equivalent to + ``embedding_func.embed("text")``. + """ + + def __init__(self, dimension: int, data_type: DataType = DataType.VECTOR_FP32): + self._dimension = dimension + self._data_type = data_type + + @property + def dimension(self) -> int: + """int: The expected dimensionality of the embedding vector.""" + return self._dimension + + @property + def data_type(self) -> DataType: + """DataType: The numeric data type of the embedding (e.g., VECTOR_FP32).""" + return self._data_type + + @abstractmethod + def embed(self, text: str) -> list[Union[int, float]]: + """Generate a dense embedding vector for the input text. + + Args: + text (str): Input text to embed. + + Returns: + list[Union[int, float]]: A list of numbers representing the embedding. + Length must equal ``self.dimension``. + """ + raise NotImplementedError + + def __call__(self, text: str) -> list[Union[int, float]]: + return self.embed(text) + + +class SparseEmbeddingFunction(ABC): + """Abstract base class for sparse vector embedding functions. + + Sparse embedding functions map text to a dictionary of {index: weight}, + where only non-zero dimensions are stored. + + Note: + Subclasses must implement the ``embed()`` method. + """ + + @abstractmethod + def embed(self, text: str) -> dict[int, float]: + """Generate a sparse embedding for the input text. + Args: + text (str): Input text to embed. + + Returns: + dict[int, float]: Mapping from dimension index to non-zero weight. + """ + raise NotImplementedError + + +class QwenEmbeddingFunction(DenseEmbeddingFunction): + """Dense embedding function using Qwen (DashScope) Text Embedding API. + + This implementation uses the DashScope service to generate embeddings + via Qwen's text embedding models (e.g., ``text-embedding-v4``). + + Args: + dimension (int): Desired embedding dimension (e.g., 1024). + model (str, optional): DashScope embedding model name. + Defaults to ``"text-embedding-v4"``. + api_key (Optional[str], optional): DashScope API key. If not provided, + reads from ``DASHSCOPE_API_KEY`` environment variable. + + Raises: + ValueError: If API key is missing or input text is invalid. + + Note: + Requires the ``dashscope`` Python package. + Embedding results are cached using ``functools.lru_cache`` (maxsize=10). + """ + + def __init__( + self, + dimension: int, + model: str = "text-embedding-v4", + api_key: Optional[str] = None, + ): + super().__init__(dimension, DataType.VECTOR_FP32) + 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") + + @property + def model(self) -> str: + """str: The DashScope embedding model name in use.""" + return self._model + + def _connection(self): + dashscope = require_module("dashscope") + dashscope.api_key = self._api_key + return dashscope + + @lru_cache(maxsize=10) + def embed(self, text: str) -> list[Union[int, float]]: + """ + Generate embedding for a given text using Qwen (via DashScope). + + Args: + text (str): Input text to embed. Must be non-empty and valid string. + + Returns: + list[Union[int, float]]: The dense embedding vector. + + Raises: + ValueError: If input is invalid or API response is malformed. + RuntimeError: If network or internal error occurs during API call. + """ + if not isinstance(text, str): + raise TypeError(f"Expected 'text' to be str, got {type(text).__name__}") + + text = text.strip() + if not text: + raise ValueError("Input text cannot be empty or whitespace only") + + resp = self._connection().TextEmbedding.call( + model=self.model, input=text, dimension=self.dimension, output_type="dense" + ) + + if resp.status_code != HTTPStatus.OK: + error_msg = getattr(resp, "message", "Unknown error") + error_detail = f"Status={resp.status_code}, Message={error_msg}" + raise ValueError(f"QwenEmbedding failed: {error_detail}") + + output = getattr(resp, "output", None) + if not isinstance(output, dict): + raise ValueError("Invalid response: missing or malformed 'output' field") + + embeddings = output.get("embeddings") + if not isinstance(embeddings, list): + raise ValueError( + "Invalid response: 'embeddings' field is missing or not a list" + ) + + if len(embeddings) != 1: + raise ValueError( + f"Expected 1 embedding, got {len(embeddings)}. Response: {resp}" + ) + + first_emb = embeddings[0] + if not isinstance(first_emb, dict): + raise ValueError("Invalid response: embedding item is not a dictionary") + + return list(first_emb.get("embedding")) diff --git a/python/zvec/extension/rerank.py b/python/zvec/extension/rerank.py new file mode 100644 index 0000000..021f6ed --- /dev/null +++ b/python/zvec/extension/rerank.py @@ -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. +from __future__ import annotations + +import heapq +import math +import os +from abc import ABC, abstractmethod +from collections import defaultdict +from http import HTTPStatus +from typing import Optional + +from ..model.doc import Doc +from ..tool import require_module +from ..typing import MetricType + + +class ReRanker(ABC): + """Abstract base class for re-ranking search results. + + Re-rankers refine the output of one or more vector queries by applying + a secondary scoring strategy. They are used in the ``query()`` method of + ``Collection`` via the ``reranker`` parameter. + + Args: + query (Optional[str], optional): Query text used for re-ranking. + Required for LLM-based re-rankers. Defaults to None. + topn (int, optional): Number of top documents to return after re-ranking. + Defaults to 10. + rerank_field (Optional[str], optional): Field name used as input for + re-ranking (e.g., document title or body). Defaults to None. + + Note: + Subclasses must implement the ``rerank()`` method. + """ + + def __init__( + self, + query: Optional[str] = None, + topn: int = 10, + rerank_field: Optional[str] = None, + ): + self._query = query + self._topn = topn + self._rerank_field = rerank_field + + @property + def topn(self) -> int: + """int: Number of top documents to return after re-ranking.""" + return self._topn + + @property + def query(self) -> str: + """str: Query text used for re-ranking.""" + return self._query + + @property + def rerank_field(self) -> Optional[str]: + """Optional[str]: Field name used as re-ranking input.""" + return self._rerank_field + + @abstractmethod + def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]: + """Re-rank documents from one or more vector queries. + + Args: + query_results (dict[str, list[Doc]]): Mapping from vector field name + to list of retrieved documents (sorted by relevance). + + Returns: + list[Doc]: Re-ranked list of documents (length ≤ ``topn``), + with updated ``score`` fields. + """ + raise NotImplementedError + + +class RrfReRanker(ReRanker): + """Re-ranker using Reciprocal Rank Fusion (RRF). + + RRF combines results from multiple queries without requiring relevance scores. + It assigns higher weight to documents that appear early in multiple result lists. + + The RRF score for a document at rank ``r`` is: ``1 / (k + r + 1)``, + where ``k`` is the rank constant. + + Args: + query (Optional[str], optional): Ignored by RRF. Defaults to None. + topn (int, optional): Number of top documents to return. Defaults to 10. + rerank_field (Optional[str], optional): Ignored by RRF. Defaults to None. + rank_constant (int, optional): Smoothing constant ``k`` in RRF formula. + Larger values reduce the impact of early ranks. Defaults to 60. + """ + + def __init__( + self, + query: Optional[str] = None, + topn: int = 10, + rerank_field: Optional[str] = None, + rank_constant: int = 60, + ): + super().__init__(query, topn, rerank_field) + self._rank_constant = rank_constant + + @property + def rank_constant(self) -> int: + return self._rank_constant + + def _rrf_score(self, rank: int): + return 1.0 / (self._rank_constant + rank + 1) + + def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]: + """Apply Reciprocal Rank Fusion to combine multiple query results. + + Args: + query_results (dict[str, list[Doc]]): Results from one or more vector queries. + + Returns: + list[Doc]: Re-ranked documents with RRF scores in the ``score`` field. + """ + rrf_scores: dict[str, float] = defaultdict(float) + id_to_doc: dict[str, Doc] = {} + + for _, query_result in query_results.items(): + for rank, doc in enumerate(query_result): + doc_id = doc.id + rrf_score = self._rrf_score(rank) + rrf_scores[doc_id] += rrf_score + if doc_id not in id_to_doc: + id_to_doc[doc_id] = doc + + top_docs = heapq.nlargest(self.topn, rrf_scores.items(), key=lambda x: x[1]) + results = [] + for doc_id, rrf_score in top_docs: + doc = id_to_doc[doc_id] + new_doc = doc._replace(score=rrf_score) + results.append(new_doc) + return results + + +class WeightedReRanker(ReRanker): + """Re-ranker that combines scores from multiple vector fields using weights. + + Each vector field's relevance score is normalized based on its metric type, + then scaled by a user-provided weight. Final scores are summed across fields. + + Args: + query (Optional[str], optional): Ignored. Defaults to None. + topn (int, optional): Number of top documents to return. Defaults to 10. + rerank_field (Optional[str], optional): Ignored. Defaults to None. + metric (MetricType, optional): Distance metric used for score normalization. + Defaults to ``MetricType.L2``. + weights (Optional[dict[str, float]], optional): Weight per vector field. + Fields not listed use weight 1.0. Defaults to None. + + Note: + Supported metrics: L2, IP, COSINE. Scores are normalized to [0, 1]. + """ + + def __init__( + self, + query: Optional[str] = None, + topn: int = 10, + rerank_field: Optional[str] = None, + metric: MetricType = MetricType.L2, + weights: Optional[dict[str, float]] = None, + ): + super().__init__(query, topn, rerank_field) + self._weights = weights + self._metric = metric + + @property + def weights(self) -> dict[str, float]: + """dict[str, float]: Weight mapping for vector fields.""" + return self._weights + + @property + def metric(self) -> MetricType: + """MetricType: Distance metric used for score normalization.""" + return self._metric + + def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]: + """Combine scores from multiple vector fields using weighted sum. + + Args: + query_results (dict[str, list[Doc]]): Results per vector field. + + Returns: + list[Doc]: Re-ranked documents with combined scores in ``score`` field. + """ + weighted_scores: dict[str, float] = defaultdict(float) + id_to_doc: dict[str, Doc] = {} + + for vector_name, query_result in query_results.items(): + for _, doc in enumerate(query_result): + doc_id = doc.id + weighted_score = self._normalize_score( + doc.score, self.metric + ) * self.weights.get(vector_name, 1.0) + weighted_scores[doc_id] += weighted_score + if doc_id not in id_to_doc: + id_to_doc[doc_id] = doc + + top_docs = heapq.nlargest( + self.topn, weighted_scores.items(), key=lambda x: x[1] + ) + results = [] + for doc_id, weighted_score in top_docs: + doc = id_to_doc[doc_id] + new_doc = doc._replace(score=weighted_score) + results.append(new_doc) + return results + + def _normalize_score(self, 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") + + +class QwenReRanker(ReRanker): + """Re-ranker using Qwen (DashScope) LLM-based re-ranking API. + + This re-ranker sends documents to the DashScope TextReRank service for + cross-encoder style re-ranking based on semantic relevance to the query. + + Args: + query (str): Query text for semantic re-ranking. **Required**. + topn (int, optional): Number of top documents to return. Defaults to 10. + rerank_field (str): Field name containing document text for re-ranking. + **Required**. + model (str, optional): DashScope re-ranking model name. + Defaults to ``"gte-rerank-v2"``. + api_key (Optional[str], optional): DashScope API key. If not provided, + reads from ``DASHSCOPE_API_KEY`` environment variable. + + Raises: + ValueError: If ``query`` is missing, ``rerank_field`` is missing, + or API key is not provided. + + Note: + Requires the ``dashscope`` Python package. + Documents without content in ``rerank_field`` are skipped. + """ + + def __init__( + self, + query: Optional[str] = None, + topn: int = 10, + rerank_field: Optional[str] = None, + model: str = "gte-rerank-v2", + api_key: Optional[str] = None, + ): + super().__init__(query, topn, rerank_field) + if not query: + raise ValueError("Query is required for reranking") + 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") + + @property + def model(self) -> str: + """str: DashScope re-ranking model name.""" + return self._model + + def _connection(self): + dashscope = require_module("dashscope") + dashscope.api_key = self._api_key + return dashscope + + def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]: + """Re-rank documents using Qwen's TextReRank API. + + Args: + query_results (dict[str, list[Doc]]): Results from vector search. + + Returns: + list[Doc]: Re-ranked documents with relevance scores from Qwen. + + Raises: + ValueError: If API call fails or no valid documents are found. + """ + if not query_results: + return [] + + id_to_doc: dict[str, Doc] = {} + doc_ids = [] + contents = [] + + for _, query_result in query_results.items(): + for doc in query_result: + doc_id = doc.id + if doc_id in id_to_doc: + continue + + 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") + + resp = self._connection().TextReRank.call( + model=self.model, + query=self.query, + documents=list(contents), + top_n=self.topn, + return_documents=False, + ) + + if resp.status_code != HTTPStatus.OK: + raise ValueError( + f"QwenReranker failed with status {resp.status_code}: {resp.message}" + ) + + results = [] + for item in resp.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/model/__init__.py b/python/zvec/model/__init__.py new file mode 100644 index 0000000..a10c33d --- /dev/null +++ b/python/zvec/model/__init__.py @@ -0,0 +1,22 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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.vector_query import VectorQuery +from .schema.collection_schema import CollectionSchema +from .schema.field_schema import FieldSchema + +__all__ = ["Collection", "CollectionSchema", "Doc", "FieldSchema", "VectorQuery"] diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py new file mode 100644 index 0000000..d50534e --- /dev/null +++ b/python/zvec/model/collection.py @@ -0,0 +1,381 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 + +from _zvec import _Collection + +from ..executor import QueryContext, QueryExecutorFactory +from ..extension import ReRanker +from ..typing import Status +from .convert import convert_to_cpp_doc, convert_to_py_doc +from .doc import Doc +from .param import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + FlatIndexParam, + HnswIndexParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + OptimizeOption, +) +from .param.vector_query import VectorQuery +from .schema import CollectionSchema, CollectionStats, FieldSchema + +__all__ = ["Collection"] + +_VECTOR_INDEX_TYPES = (HnswIndexParam, IVFIndexParam, FlatIndexParam) + + +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 = QueryExecutorFactory.create(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, IVFIndexParam, FlatIndexParam, InvertIndexParam + ], + 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. + + Args: + field_name (str): Name of the field to index. + index_param (Union[HnswIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam]): + Index configuration. + option (Optional[IndexOption], optional): Index creation options. + Defaults to ``IndexOption()``. + + Raises: + ValueError: If a vector index is applied to a non-vector field. + """ + if index_param in _VECTOR_INDEX_TYPES and not self.schema.vector(field_name): + supported_types = ", ".join(cls.__name__ for cls in _VECTOR_INDEX_TYPES) + raise ValueError( + f"Cannot apply vector index to non-vector field '{field_name}'. " + f"The field must be of vector type to use index types like {supported_types}." + ) + self._obj.CreateIndex(field_name, index_param, option) + self._schema = CollectionSchema._from_core(self._obj.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()) + + 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.name, field_schema._get_object(), expression, option + ) + self._schema = CollectionSchema._from_core(self._obj.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()) + + 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()) + + # ========== Collection DDL Methods ========== + 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 + + 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 + + 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 + + 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]]) -> dict[str, Doc]: + """Retrieve documents by ID. + + Args: + ids (Union[str, list[str]]): Document IDs to fetch. + + 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) + 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, + vectors: Optional[Union[VectorQuery, list[VectorQuery]]] = None, + *, + topk: int = 10, + filter: Optional[str] = None, + include_vector: bool = False, + output_fields: Optional[list[str]] = None, + reranker: Optional[ReRanker] = None, + ) -> list[Doc]: + """Perform vector similarity search with optional filtering and re-ranking. + + At least one `VectorQuery` must be provided. + + Args: + vectors (Optional[Union[VectorQuery, list[VectorQuery]]], optional): + One or more vector queries. Defaults to None. + 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: + list[Doc]: Top-k matching documents, sorted by relevance score. + + Examples: + >>> from zvec import VectorQuery + >>> results = collection.query( + ... vectors=VectorQuery("embedding", vector=[0.1, 0.2]), + ... topk=5, + ... filter="category == 'tech'", + ... output_fields=["title", "url"] + ... ) + """ + ctx = QueryContext( + topk=topk, + filter=filter, + queries=[vectors] if isinstance(vectors, VectorQuery) else vectors, + 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..2eac08c --- /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 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..722293d --- /dev/null +++ b/python/zvec/model/doc.py @@ -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. +from __future__ import annotations + +import json +from typing import Any, Optional + +from ..common import VectorType + +__all__ = [ + "Doc", +] + + +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 diff --git a/python/zvec/model/param/__init__.py b/python/zvec/model/param/__init__.py new file mode 100644 index 0000000..4dbeb24 --- /dev/null +++ b/python/zvec/model/param/__init__.py @@ -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. +from __future__ import annotations + +from _zvec.param import ( + AddColumnOption, + AlterColumnOption, + CollectionOption, + FlatIndexParam, + HnswIndexParam, + HnswQueryParam, + IndexOption, + InvertIndexParam, + IVFIndexParam, + IVFQueryParam, + OptimizeOption, +) + +__all__ = [ + "AddColumnOption", + "AlterColumnOption", + "CollectionOption", + "FlatIndexParam", + "HnswIndexParam", + "HnswQueryParam", + "IVFIndexParam", + "IVFQueryParam", + "IndexOption", + "InvertIndexParam", + "OptimizeOption", +] diff --git a/python/zvec/model/param/__init__.pyi b/python/zvec/model/param/__init__.pyi new file mode 100644 index 0000000..a1a9dcc --- /dev/null +++ b/python/zvec/model/param/__init__.pyi @@ -0,0 +1,692 @@ +""" +This module contains the params of Zvec +""" + +from __future__ import annotations + +import collections +import typing + +import _zvec.typing + +__all__: list[str] = [ + "AddColumnOption", + "AlterColumnOption", + "CollectionOption", + "FlatIndexParam", + "HnswIndexParam", + "HnswQueryParam", + "IVFIndexParam", + "IVFQueryParam", + "IndexOption", + "IndexParam", + "InvertIndexParam", + "OptimizeOption", + "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``. + + 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.typing.MetricType = ..., + quantize_type: _zvec.typing.QuantizeType = ..., + ) -> 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). + """ + + 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 100. + 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 + ... ) + >>> print(params) + {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8'} + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + metric_type: _zvec.typing.MetricType = ..., + m: typing.SupportsInt = 100, + ef_construction: typing.SupportsInt = 500, + quantize_type: _zvec.typing.QuantizeType = ..., + ) -> 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. + """ + +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. + + 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, + ) -> None: + """ + 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. + """ + 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. + If set to 0, the system will auto-select a reasonable value based on data size. + Default is 0 (auto). + 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.typing.MetricType = ..., + n_list: typing.SupportsInt = 0, + n_iters: typing.SupportsInt = 10, + use_soar: bool = False, + quantize_type: _zvec.typing.QuantizeType = ..., + ) -> 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). Set to 0 for auto. + Defaults to 0. + 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. + """ + + 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 (0 = auto). + """ + + @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 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.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.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 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. + """ + + 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.typing.MetricType: + """ + MetricType: Distance metric (e.g., IP, COSINE, L2). + """ + + @property + def quantize_type(self) -> _zvec.typing.QuantizeType: + """ + QuantizeType: Vector quantization type (e.g., FP16, INT8). + """ + +class _VectorQuery: + 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, arg0: ..., arg1: typing.Any) -> None: ... + @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/vector_query.py b/python/zvec/model/param/vector_query.py new file mode 100644 index 0000000..97d105a --- /dev/null +++ b/python/zvec/model/param/vector_query.py @@ -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. +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Union + +from ...common import VectorType +from . import HnswQueryParam, IVFQueryParam + +__all__ = ["VectorQuery"] + + +@dataclass(frozen=True) +class VectorQuery: + """Represents a vector search query for a specific field in a collection. + + A `VectorQuery` can be constructed using either a document ID (to look up + its vector) or an explicit vector. It may optionally include index-specific + query parameters to control search behavior (e.g., `ef` for HNSW, `nprobe` for IVF). + + Exactly one of `id` or `vector` should be provided. If both are given, + behavior is implementation-defined (typically `id` takes precedence). + + Attributes: + field_name (str): Name of the vector 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, IVFQueryParam]], optional): + Index-specific query parameters. Default is None. + + Examples: + >>> import zvec + >>> # Query by ID + >>> q1 = zvec.VectorQuery(field_name="embedding", id="doc123") + >>> # Query by vector + >>> q2 = zvec.VectorQuery( + ... field_name="embedding", + ... vector=[0.1, 0.2, 0.3], + ... param=HnswQueryParam(ef=300) + ... ) + """ + + field_name: str + id: Optional[str] = None + vector: VectorType = None + param: Optional[Union[HnswQueryParam, IVFQueryParam]] = 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 _validate(self) -> None: + if self.field_name is None: + raise ValueError("Field name cannot be empty") + if self.id and self.vector: + raise ValueError("Cannot provide both id and vector") diff --git a/python/zvec/model/schema/__init__.py b/python/zvec/model/schema/__init__.py new file mode 100644 index 0000000..5ff532f --- /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.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..30ba20f --- /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.param +import _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.typing.DataType, + nullable: bool = False, + dimension: typing.SupportsInt = 0, + index_param: _zvec.param.IndexParam = None, + ) -> None: ... + def __ne__(self, arg0: _FieldSchema) -> bool: ... + @property + def data_type(self) -> _zvec.typing.DataType: ... + @property + def dimension(self) -> int: ... + @property + def index_param(self) -> typing.Any: ... + @property + def index_type(self) -> _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..e07095b --- /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.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..da193dd --- /dev/null +++ b/python/zvec/model/schema/field_schema.py @@ -0,0 +1,300 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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.schema import _FieldSchema + +from zvec.model.param import ( + FlatIndexParam, + HnswIndexParam, + 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[InvertIndexParam], optional): Inverted index + parameters for this field. Only applicable to fields that support + indexing (e.g., scalar fields used in filtering). Defaults to None. + + Examples: + >>> from zvec.typing import DataType + >>> from zvec.model.param import InvertIndexParam + >>> id_field = FieldSchema( + ... name="id", + ... data_type=DataType.INT64, + ... nullable=False, + ... index_param=InvertIndexParam(enable_range_optimization=True) + ... ) + """ + + def __init__( + self, + name: str, + data_type: DataType, + nullable: bool = False, + index_param: Optional[InvertIndexParam] = 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[InvertIndexParam]: + """Optional[InvertIndexParam]: Inverted 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, 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, IVFIndexParam, FlatIndexParam]: + """Union[HnswIndexParam, 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..f4deb3c --- /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(".")[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..da83c44 --- /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.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..114fb49 --- /dev/null +++ b/python/zvec/zvec.py @@ -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. +from __future__ import annotations + +from typing import Optional + +from _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] = 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, + 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, + memory_limit_mb: Optional[int] = 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``. + 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. + + 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 memory_limit_mb is not None: + config_dict["memory_limit_mb"] = memory_limit_mb + + 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 (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): 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/gcov.sh b/scripts/gcov.sh new file mode 100755 index 0000000..3625b6e --- /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 || exit 1 +eval $(echo lcov -r $project_name.lcov.info -o $project_name-filtered.lcov.info $filter_list) || 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/src/CMakeLists.txt b/src/CMakeLists.txt new file mode 100644 index 0000000..c516187 --- /dev/null +++ b/src/CMakeLists.txt @@ -0,0 +1,13 @@ +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(core) +cc_directory(db) +if(BUILD_PYTHON_BINDINGS) + cc_directory(binding) +endif() diff --git a/src/ailego/CMakeLists.txt b/src/ailego/CMakeLists.txt new file mode 100644 index 0000000..b01df97 --- /dev/null +++ b/src/ailego/CMakeLists.txt @@ -0,0 +1,28 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +find_package(Threads REQUIRED) + +if(UNIX AND NOT APPLE) + 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() + +cc_library( + NAME zvec_ailego STATIC STRICT PACKED + SRCS ${ALL_SRCS} + LIBS ${EXTRA_LIBS} + Arrow::arrow_static + Arrow::parquet_static + VERSION "${GIT_SRCS_VER}" +) \ No newline at end of file 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..89057fe --- /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..0b07293 --- /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..07a77e5 --- /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..7c003c9 --- /dev/null +++ b/src/ailego/algorithm/kmeans.h @@ -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. + +#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_{}; +}; + +/*! Binary K-Means Context + */ +template +class BinaryKmeansContext { + 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(std::is_same::value || + std::is_same::value, + "ValueType must be uint32_t or uint64_t"); + + /*! K-Means Context Cluster + */ + class Cluster { + public: + //! Constructor + Cluster(size_t dim) : accum_(dim, 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; + + const uint8_t *arr = reinterpret_cast(vec); + for (size_t i = 0; i != dim; ++i) { + if (arr[i >> 3] & (1u << (i & 7))) { + accum_[i] += 1; + } + } + mutex_.unlock(); + } + + //! Retrieve the centroid of vectors + void centroid(ValueType *out, size_t dim) const { + ailego_check_with(dim == accum_.size(), "Unmatched dimension"); + + uint8_t *arr = reinterpret_cast(out); + size_t half = count_ >> 1; + for (size_t i = 0; i != dim; ++i) { + if (accum_[i] > half) { + arr[i >> 3] |= static_cast(1 << (i & 0x7)); + } else { + arr[i >> 3] &= ~static_cast(1 << (i & 0x7)); + } + } + } + + //! 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 ValueType *m, const ValueType *q, size_t dim, + float *out) { + HammingDistanceMatrix::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) { + HammingDistanceMatrix::Compute(m, q, dim, out); + } + + //! Transpose a matrix + static void MatrixTranspose(const ValueType *src, size_t dim, T *dst) { + MatrixHelper::Transpose( + src, (dim >> 3) / sizeof(ValueType), dst); + } + + //! Reverse transpose a matrix + static void MatrixReverseTranspose(const ValueType *src, size_t dim, T *dst) { + MatrixHelper::ReverseTranspose( + src, (dim >> 3) / sizeof(ValueType), dst); + } + + //! Compute Norm2 + static void Norm2(ValueType * /*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>; + +/*! Binary K-Means cluster algorithm + */ +template > +using BinaryKmeans = LloydCluster>; + +/*! Numerical K-Means cluster algorithm + */ +template > +using NumericalInnerProductKmeans = + LloydCluster>; + +/*! Nibble K-Means cluster algorithm + */ +template > +using NibbleInnerProductKmeans = + LloydCluster>; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/algorithm/lloyd_cluster.h b/src/ailego/algorithm/lloyd_cluster.h new file mode 100644 index 0000000..4227633 --- /dev/null +++ b/src/ailego/algorithm/lloyd_cluster.h @@ -0,0 +1,363 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! 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) { + float scores[BatchCount]; + + 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); + + 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); + + 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::vector scores(BatchCount * BatchCount, 0); + 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; + }; + + float nearest_scores[BatchCount]; + size_t nearest_indexes[BatchCount]; + + 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, nearest_indexes + BatchCount, 0); + std::fill(nearest_scores, nearest_scores + 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[0]); + + 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[0]); + + 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 \ No newline at end of file diff --git a/src/ailego/buffer/buffer_manager.cc b/src/ailego/buffer/buffer_manager.cc new file mode 100644 index 0000000..15111b8 --- /dev/null +++ b/src/ailego/buffer/buffer_manager.cc @@ -0,0 +1,579 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "buffer_manager.h" +#include +#include +#include +#include +#include +#include + + +namespace zvec { + + +namespace ailego { + + +namespace { + + +struct IDHash { + size_t operator()(const BufferID &buffer_id) const { + size_t hash = std::hash{}(static_cast(buffer_id.type)); + hash = hash ^ (std::hash{}(buffer_id.file_id)); + if (buffer_id.type == BufferID::TYPE::PARQUET) { + hash = hash * 31 + std::hash{}(buffer_id.parquet().column); + hash = hash * 31 + std::hash{}(buffer_id.parquet().row_group); + } else if (buffer_id.type == BufferID::TYPE::VECTOR) { + hash = hash * 31 + std::hash{}(buffer_id.vector().offset); + } + return hash; + } +}; + + +struct IDEqual { + bool operator()(const BufferID &a, const BufferID &b) const { + if (a.type != b.type) { + return false; + } + if (a.file_name != b.file_name) { + return false; + } + if (a.file_id != b.file_id) { + return false; + } + if (a.mtime != b.mtime) { + return false; + } + if (a.type == BufferID::TYPE::PARQUET) { + return a.parquet().column == b.parquet().column && + a.parquet().row_group == b.parquet().row_group; + } else if (a.type == BufferID::TYPE::VECTOR) { + return a.vector().offset == b.vector().offset; + } else { + return false; + } + } +}; + + +} // namespace + + +struct BufferManager::BufferContext { + BufferContext(const BufferID &id, BufferPool *p) : id(id), pool(p) {}; + BufferContext(const BufferContext &) = delete; + BufferContext(BufferContext &&) = delete; + BufferContext &operator=(const BufferContext &) = delete; + BufferContext &operator=(BufferContext &&) = delete; + + + ~BufferContext() { + if (vector) { + free(vector); + } + } + + + typedef std::unique_ptr Pointer; + + + enum State : uint32_t { + IDLE = 0, // Empty and not held by any users, not in LRU + RESERVED = 1, // Pinned by a user but no data yet, not in LRU + IN_USE = 2, // Pinned by a user and data is present, not in LRU + CACHED = 3, // Data is present but not held by any users, in LRU + ERROR = 4 // Something went wrong, not in LRU + }; + + + // Identifier for the buffer + BufferID id; + + // Current state + State state{IDLE}; + + // The size of the buffer + uint32_t size{0}; + + // Handle of the file backing this buffer + File file; + + // The number of external references to this buffer (via pin/unpin) + std::atomic refs_buf{0}; + + // The number of external references to this context (via BufferHandle) + std::atomic refs_context{0}; + + BufferPool *pool{nullptr}; + + // A shared pointer to the buffers allocated for arrow parquet data + std::shared_ptr arrow{nullptr}; + + // Guard original arrow buffers to prevent premature deletion + std::vector> arrow_refs{}; + + // A pointer to the buffer allocated for vector data + void *vector{nullptr}; + + // Doubly linked LRU list + BufferContext *next{nullptr}; + BufferContext *prev{nullptr}; + + + // Return a string representation of the status + const std::string status_string() const; + + // Populate the buffer with parquet data + arrow::Status read_arrow_parquet(); + + // Populate the buffer with vector data + bool read_vector(); +}; + + +const std::string BufferManager::BufferContext::status_string() const { + std::string msg{id.to_string() + ": "}; + switch (state) { + case State::IDLE: { + msg += "Idle"; + break; + } + case State::RESERVED: { + msg += "Reserved"; + break; + } + case State::IN_USE: { + msg += "In use"; + break; + } + case State::CACHED: { + msg += "Cached"; + break; + } + case State::ERROR: { + msg += "Error"; + break; + } + } + return msg; +} + + +arrow::Status BufferManager::BufferContext::read_arrow_parquet() { + // TODO: file handler and memory pool can be optimized + arrow::MemoryPool *mem_pool = arrow::default_memory_pool(); + + // Open file + std::shared_ptr input; + const auto &file_name = id.file_name; + ARROW_ASSIGN_OR_RAISE(input, arrow::io::ReadableFile::Open(file_name)); + + // Open reader + std::unique_ptr reader; + ARROW_ASSIGN_OR_RAISE(reader, parquet::arrow::OpenFile(input, mem_pool)); + + // Perform read + int row_group = id.parquet().row_group; + int column = id.parquet().column; + auto s = reader->RowGroup(row_group)->Column(column)->Read(&arrow); + if (!s.ok()) { + LOG_ERROR("Failed to read parquet file[%s]", file_name.c_str()); + arrow = nullptr; + return s; + } + + // Compute the memory usage and hijack Arrow's buffers with our implementation + for (auto &array : arrow->chunks()) { + auto &buffers = array->data()->buffers; + for (size_t buf_idx = 0; buf_idx < buffers.size(); ++buf_idx) { + if (buffers[buf_idx] == nullptr) { + continue; + } + // Keep references to original buffers to prevent premature deletion + arrow_refs.emplace_back(buffers[buf_idx]); + size += buffers[buf_idx]->capacity(); + // Create hijacked buffer with custom deleter that notifies us when Arrow + // is finished with the buffer + std::shared_ptr hijacked_buffer( + buffers[buf_idx].get(), BufferManager::ArrowBufferDeleter(this)); + buffers[buf_idx] = hijacked_buffer; + } + } + + return arrow::Status::OK(); +} + + +bool BufferManager::BufferContext::read_vector() { + const auto &file_name = id.file_name; + if (!file.is_valid()) { + if (!File::IsExist(file_name)) { + LOG_ERROR("File[%s] does not exist", file_name.c_str()); + return false; + } + if (!File::IsRegular(file_name)) { + LOG_ERROR("[%s] is not a regular file", file_name.c_str()); + return false; + } + if (!file.open(file_name.c_str(), true, false)) { + LOG_ERROR("Failed to open file[%s]", file_name.c_str()); + return false; + } + } + AILEGO_DEFER([this] { file.close(); }); + uint32_t len = id.vector().length; + auto ret = posix_memalign((void **)&vector, 64, len); // 64-byte alignment + if (ret != 0 || vector == nullptr) { + LOG_ERROR("Failed to allocate buffer for file[%s]", file_name.c_str()); + return false; + } + uint32_t offset = id.vector().offset; + if (file.read(offset, vector, len) != len) { + LOG_ERROR("Failed to read file[%s]", file_name.c_str()); + free(vector); + vector = nullptr; + return false; + } + size = len; + return true; +} + + +// Thread-safe buffer pool implementation. +// +// BufferContext states: +// 1. Must exist in the lookup (hash) table. +// 2. LRU list presence: +// - In LRU: holds memory but not pinned by any users +// - Not in LRU: either holds memory pinned by users, or doesn't hold memory +// 3. External references: when an external user acquires a context and pins the +// memory, that context is removed from LRU list; when they unpins the +// memory, that context is moved to LRU list if it was the last reference. +// +// Any operation on the hash table is protected by mutex_table_. +// Any change to context state and LRU list is protected by mutex_context_. +// +class BufferManager::BufferPool { + public: + explicit BufferPool(uint64_t limit) : limit_(limit) { + sentinel_.next = &sentinel_; + sentinel_.prev = &sentinel_; + } + + + BufferContext *acquire_locked(BufferID &id) { + std::lock_guard lock(mutex_context_); + if (auto iter = table_.find(id); iter != table_.end()) { + return iter->second.get(); + } + auto [iter, _] = + table_.emplace(id, std::make_unique(id, this)); + return iter->second.get(); + } + + + void try_release_context_locked(BufferContext *context) { + if (context->refs_context.load() != 0) { + return; + } + std::lock_guard lock(mutex_table_); + if (context->refs_context.load() != 0) { + return; + } + if (context->state == BufferContext::State::IDLE) { + table_.erase(context->id); + } + } + + + void pin_locked(BufferContext *ctx) { + std::lock_guard lock(mutex_context_); + if (ctx->state == BufferContext::State::IDLE) { + return pin_at_IDLE(ctx); + } + if (ctx->state == BufferContext::State::IN_USE) { + return pin_at_IN_USE(ctx); + } + if (ctx->state == BufferContext::State::CACHED) { + return pin_at_CACHED(ctx); + } + if (ctx->state == BufferContext::State::ERROR) { + return; + } + } + + + bool unpin_locked(BufferContext *ctx) { + uint32_t prev_refs = ctx->refs_buf.fetch_sub(1); + if (prev_refs > 1) { + return false; + } + std::lock_guard lock(mutex_context_); + if (ctx->refs_buf.load() == 0 && + ctx->state != BufferContext::State::CACHED) { + ctx->state = BufferContext::State::CACHED; + LRU_insert(ctx); + return true; + } else { + return false; + } + } + + + void LRU_insert_locked(BufferContext *context) { + std::lock_guard lock(mutex_context_); + LRU_insert(context); + } + + + void LRU_remove_locked(BufferContext *context) { + std::lock_guard lock(mutex_context_); + LRU_remove(context); + } + + + uint64_t usage() const { + return usage_; + } + + + private: + void pin_at_IDLE(BufferContext *ctx) { + ctx->state = BufferContext::State::RESERVED; + + while (usage_ >= limit_) { + // The tail of LRU list is the least recently used context + BufferContext *victim = sentinel_.prev; + if (victim == &sentinel_) { // No victim could be found + ctx->state = BufferContext::State::ERROR; + return; + } + if (victim->state == BufferContext::State::ERROR) { + LRU_remove(victim); + try_release_context_locked(ctx); + continue; + } + if (victim->id.type == BufferID::TYPE::PARQUET) { + victim->arrow_refs.clear(); + } else { + free(victim->vector); + victim->vector = nullptr; + } + victim->state = BufferContext::State::IDLE; + LRU_remove(victim); + try_release_context_locked(ctx); + usage_ -= victim->size; + } + + if (ctx->id.type == BufferID::TYPE::PARQUET) { + if (ctx->read_arrow_parquet().ok()) { + ctx->state = BufferContext::State::IN_USE; + ctx->refs_buf.fetch_add(ctx->arrow_refs.size()); + usage_ += ctx->size; + } else { + LOG_ERROR("Failed to read to %s", ctx->id.to_string().c_str()); + ctx->state = BufferContext::State::ERROR; + } + } else { + if (ctx->read_vector()) { + ctx->state = BufferContext::State::IN_USE; + ctx->refs_buf.fetch_add(1); + usage_ += ctx->size; + } else { + LOG_ERROR("Failed to read to %s", ctx->id.to_string().c_str()); + ctx->state = BufferContext::State::ERROR; + } + } + } + + + void pin_at_IN_USE(BufferContext *ctx) { + if (ctx->id.type == BufferID::TYPE::PARQUET) { + ctx->refs_buf.fetch_add(ctx->arrow_refs.size()); + } else { + ctx->refs_buf.fetch_add(1); + } + } + + + void pin_at_CACHED(BufferContext *ctx) { + if (ctx->id.type == BufferID::TYPE::PARQUET) { + ctx->refs_buf.fetch_add(ctx->arrow_refs.size()); + } else { + ctx->refs_buf.fetch_add(1); + } + LRU_remove(ctx); + ctx->state = BufferContext::State::IN_USE; + } + + + void LRU_insert(BufferContext *context) { + if (context->refs_buf > 0) { + return; // Already pinned, should not be evicted + } + if (context->next != nullptr || context->prev != nullptr) { + return; + } + // Insert the context to the head of LRU list + context->next = sentinel_.next; + context->prev = &sentinel_; + sentinel_.next = context; + context->next->prev = context; + inactive_ += context->size; + } + + + void LRU_remove(BufferContext *context) { + if (context->next == nullptr) { + return; // Not in LRU list + } + context->next->prev = context->prev; + context->prev->next = context->next; + context->next = nullptr; + context->prev = nullptr; + inactive_ -= context->size; + } + + private: + using Table = + std::unordered_map; + + uint64_t limit_; + std::atomic usage_{0}; + std::atomic inactive_{0}; + + Table table_{}; + std::mutex mutex_table_{}; + BufferContext sentinel_{BufferID{}, this}; // LRU list sentinel + std::mutex mutex_context_{}; +}; + + +BufferManager::ArrowBufferDeleter::ArrowBufferDeleter(BufferContext *c) + : context(c) {} + + +void BufferManager::ArrowBufferDeleter::operator()(arrow::Buffer *) { + context->pool->unpin_locked(context); +} + + +BufferHandle::BufferHandle(BufferContext *context) : context_(context) { + if (context_ != nullptr) { + pool_ = context_->pool; + context_->refs_context.fetch_add(1); + } +} + + +BufferHandle::~BufferHandle() { + if (context_ != nullptr) { + uint32_t prev_refs = context_->refs_context.fetch_sub(1); + if (prev_refs > 1) { + return; + } + if (context_->state == BufferContext::State::IDLE) { + pool_->try_release_context_locked(context_); + } + } +} + + +std::shared_ptr BufferHandle::pin_parquet_data() { + pool_->pin_locked(context_); + return context_->arrow; +} + + +void *BufferHandle::pin_vector_data() { + if (!context_) { + return nullptr; + } + pool_->pin_locked(context_); + return context_->vector; +} + + +bool BufferHandle::unpin_vector_data() { + if (!context_) { + return true; + } + return pool_->unpin_locked(context_); +} + + +uint32_t BufferHandle::references() const { + return context_->refs_buf.load(); +} + + +uint32_t BufferHandle::size() const { + return context_->size; +} + + +void BufferManager::init(uint64_t limit, uint32_t num_shards) { + pools_.clear(); + uint64_t limit_per_shard = ailego_align(limit / num_shards, 4096); + for (uint32_t i = 0; i < num_shards; ++i) { + auto pool = new BufferPool(limit_per_shard); + pools_.push_back(pool); + } + LOG_INFO( + "BufferManager initialized with [%u] buffer pools, [%zu] bytes memory " + "limit per pool, total memory limit [%zu] bytes", + num_shards, (size_t)limit_per_shard, (size_t)limit); +} + + +BufferHandle BufferManager::acquire(BufferID &buffer_id) { + static IDHash id_hash{}; + auto hash_val = id_hash(buffer_id); + auto ctx = pools_[hash_val % pools_.size()]->acquire_locked(buffer_id); + return BufferHandle(ctx); +} + + +std::unique_ptr BufferManager::acquire_ptr(BufferID &buffer_id) { + static IDHash id_hash{}; + auto hash_val = id_hash(buffer_id); + auto ctx = pools_[hash_val % pools_.size()]->acquire_locked(buffer_id); + return std::make_unique(ctx); +} + + +uint64_t BufferManager::total_size_in_bytes() const { + uint64_t total_usage = 0; + for (auto pool : pools_) { + total_usage += pool->usage(); + } + return total_usage; +} + + +BufferManager::~BufferManager() { + for (auto pool : pools_) { + delete pool; + } +} + + +} // namespace ailego + + +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/buffer/buffer_manager.h b/src/ailego/buffer/buffer_manager.h new file mode 100644 index 0000000..46ee6aa --- /dev/null +++ b/src/ailego/buffer/buffer_manager.h @@ -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. + + +#pragma once + + +#include +#include +#include +#include +#include +#include + +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunused-parameter" +#pragma clang diagnostic ignored "-Wshadow" +#elif defined(__GNUC__) || defined(__GNUG__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-parameter" +#pragma GCC diagnostic ignored "-Wshadow" +#endif + +#include + +#ifdef __clang__ +#pragma clang diagnostic pop +#elif defined(__GNUC__) || defined(__GNUG__) +#pragma GCC diagnostic pop +#endif + +#include "ailego/pattern/singleton.h" + + +namespace zvec { + + +namespace ailego { + + +struct BufferID; +class BufferManager; +class BufferHandle; + + +struct BufferID { + struct ParquetPos { + int column; + int row_group; + }; + struct VectorPos { + uint32_t offset; + uint32_t length; + }; + union Position { + explicit Position() = default; + ParquetPos forward; + VectorPos vector; + }; + enum TYPE { + PARQUET = 1, + VECTOR = 2, + UNKNOWN = 0, + }; + + + static std::uint64_t getLastModifiedNs(const std::filesystem::path &p) { + auto ftime = std::filesystem::last_write_time(p); + return static_cast(ftime.time_since_epoch().count()); + } + + // Cross-platform helper to get nanosecond modification time + // static long get_st_mtime_nsec(const struct stat &file_stat) { + // #ifdef __APPLE__ + // return file_stat.st_mtim.tv_nsec; + // #else + // return file_stat.st_mtim.tv_nsec; + // #endif + // } + + static BufferID ParquetID(const std::string &file_name, int column, + int row_group) { + BufferID buffer_id{}; + buffer_id.type = TYPE::PARQUET; + buffer_id.file_name = file_name; + buffer_id.pos.forward.column = column; + buffer_id.pos.forward.row_group = row_group; + struct stat file_stat; + if (stat(file_name.c_str(), &file_stat) == 0) { + // file_stat.st_ino contains the inode number + // file_stat.st_dev contains the device ID + // Together they uniquely identify a file + buffer_id.file_id = file_stat.st_ino; + std::filesystem::path p(file_name); + buffer_id.mtime = getLastModifiedNs(p); + } + return buffer_id; + } + + static BufferID VectorID(const std::string &file_name, uint32_t offset, + uint32_t length) { + BufferID buffer_id{}; + buffer_id.type = TYPE::VECTOR; + buffer_id.file_name = file_name; + struct stat file_stat; + if (stat(file_name.c_str(), &file_stat) == 0) { + buffer_id.file_id = file_stat.st_ino; + std::filesystem::path p(file_name); + buffer_id.mtime = getLastModifiedNs(p); + } + buffer_id.pos.vector.offset = offset; + buffer_id.pos.vector.length = length; + return buffer_id; + } + + explicit BufferID() = default; + + // Type of the file backing this buffer + TYPE type{UNKNOWN}; + + // Name of the file backing this buffer + std::string file_name{}; + + // Unique file id + uint64_t file_id{}; + + long mtime{}; + + // To identify which part of the backing file should be loaded into the buffer + Position pos{}; + + + // Get the forward ID + const inline struct ParquetPos &parquet() const { + return pos.forward; + } + + + // Get the vector ID + const inline struct VectorPos &vector() const { + return pos.vector; + } + + + // Get debug string + const std::string to_string() const { + std::string msg{"Buffer["}; + if (type == TYPE::PARQUET) { + msg += "parquet: " + file_name + "[" + std::to_string(file_id) + "]" + + ", column: " + std::to_string(parquet().column) + + ", row_group: " + std::to_string(parquet().row_group); + } else if (type == TYPE::VECTOR) { + msg += "vector: " + file_name + "[" + std::to_string(file_id) + "]" + + ", offset: " + std::to_string(vector().offset); + } else { + msg += "unknown"; + } + msg += ", mtime: " + std::to_string(mtime); + msg += "]"; + return msg; + } +}; + + +// Thread-safe LRU buffer implementation. +class BufferManager : public Singleton { + friend BufferHandle; + + public: + void init(uint64_t limit, uint32_t num_shards = 1); + + BufferHandle acquire(BufferID &buffer_id); + + std::unique_ptr acquire_ptr(BufferID &buffer_id); + + uint64_t total_size_in_bytes() const; + + ~BufferManager(); + + private: + struct BufferContext; + + class BufferPool; + + // Custom deleter for Arrow buffer that automatically notifies us when the + // buffer is no longer referenced by Arrow + struct ArrowBufferDeleter { + explicit ArrowBufferDeleter(BufferContext *c); + BufferContext *context; + // Only reduces the reference count but does not actually release the + // buffer, since the buffer memory is managed by the BufferManager. + void operator()(arrow::Buffer *); + }; + + std::vector pools_; +}; + + +class BufferHandle { + public: + typedef std::unique_ptr Pointer; + + explicit BufferHandle(BufferManager::BufferContext *context = nullptr); + BufferHandle(const BufferHandle &) = delete; + BufferHandle(BufferHandle &&) = default; + BufferHandle &operator=(const BufferHandle &) = delete; + BufferHandle &operator=(BufferHandle &&) = default; + + + ~BufferHandle(); + + + // Pin parquet data in memory by allocating arrow buffers of appropriate size + // and reading data from the backing file. + // The lifecycle of the allocated memory is automatically managed through + // shared pointers. The buffers are guaranteed to be held until they are not + // referenced. + // Returns a pointer to the loaded ChunkedArray in Arrow format. + std::shared_ptr pin_parquet_data(); + + + // Pin vector data in memory by allocating a buffer of appropriate size and + // loading data from the backing file. + // The memory is guaranteed to be held until unpin() is called. The caller + // must call unpin() to release the memory when it is no longer needed. + // Returns a raw memory address. + void *pin_vector_data(); + + + // Reduce the reference count for this vector buffer. + // Returns true if this was the last reference. + // When reference count is zero, the buffer is moved to the eviction list and + // becomes eligible for removal under memory pressure. + bool unpin_vector_data(); + + + // Get the current reference count. + uint32_t references() const; + + + // Get the buffer size. + uint32_t size() const; + + + private: + using BufferContext = BufferManager::BufferContext; + using BufferPool = BufferManager::BufferPool; + + BufferContext *context_; + BufferPool *pool_; +}; + + +} // namespace ailego + + +} // namespace zvec diff --git a/src/ailego/container/bitmap.cc b/src/ailego/container/bitmap.cc new file mode 100644 index 0000000..2202554 --- /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 = NULL; + 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 = NULL; + } + ++shrink_count; + } + for (; iter != array_.rend(); ++iter) { + if ((*iter) && (*iter)->test_none()) { + delete (*iter); + *iter = NULL; + } + } + 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, NULL); + } + + 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, NULL); + } + + 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, NULL); + } + + 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 = NULL; + } + } + } + for (size_t i = overlap; i < array_.size(); ++i) { + Bucket *&dst = array_[i]; + delete dst; + dst = NULL; + } +} + +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 = NULL; + + 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 = NULL; + + 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..c95dcf7 --- /dev/null +++ b/src/ailego/container/bitmap.h @@ -0,0 +1,503 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! 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/blob.h b/src/ailego/container/blob.h new file mode 100644 index 0000000..2e45f7c --- /dev/null +++ b/src/ailego/container/blob.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 + +namespace zvec { +namespace ailego { + +/*! AiLego Blob Wrap + */ +class BlobWrap { + public: + //! Constructor + BlobWrap(void) : buffer_(nullptr), size_(0u) {} + + //! Constructor + BlobWrap(const BlobWrap &rhs) : buffer_(rhs.buffer_), size_(rhs.size_) {} + + //! Constructor + BlobWrap(BlobWrap &&rhs) : buffer_(rhs.buffer_), size_(rhs.size_) { + rhs.buffer_ = nullptr; + rhs.size_ = 0u; + } + + //! Constructor + BlobWrap(const void *buf, size_t len) + : buffer_(const_cast(buf)), size_(len) {} + + //! Constructor + BlobWrap(const std::string &buf) + : buffer_(const_cast(buf.data())), size_(buf.size()) {} + + //! Destructor + ~BlobWrap(void) {} + + //! Assignment + BlobWrap &operator=(const BlobWrap &rhs) { + buffer_ = rhs.buffer_; + size_ = rhs.size_; + return *this; + } + + //! Assignment + BlobWrap &operator=(BlobWrap &&rhs) { + buffer_ = rhs.buffer_; + size_ = rhs.size_; + rhs.buffer_ = nullptr; + rhs.size_ = 0u; + return *this; + } + + //! Test if the blob is valid + bool is_valid(void) const { + return (buffer_ && size_); + } + + //! Mount a buffer as blob + void mount(void *buf, size_t len) { + buffer_ = buf; + size_ = len; + } + + //! Mount a string as blob + void mount(std::string &buf) { + buffer_ = const_cast(buf.data()); + size_ = buf.size(); + } + + //! Umount the buffer of blob + void umount(void) { + buffer_ = nullptr; + size_ = 0u; + } + + //! Retrieve buffer of blob + void *buffer(void) { + return buffer_; + } + + //! Retrieve buffer of blob + const void *buffer(void) const { + return buffer_; + } + + //! Retrieve size of blob + size_t size(void) const { + return size_; + } + + //! Copy a buffer into blob + void copy(const void *buf, size_t len) { + memcpy(buffer_, buf, std::min(size_, len)); + } + + //! Copy a blob to blob + void copy(const BlobWrap &rhs) { + memcpy(buffer_, rhs.buffer_, std::min(size_, rhs.size_)); + } + + //! Copy a string to blob + void copy(const std::string &str) { + memcpy(buffer_, str.data(), std::min(size_, str.size())); + } + + //! Zero the buffer of blob + void zero(void) { + memset(buffer_, 0, size_); + } + + private: + void *buffer_; + size_t size_; +}; + +} // 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..714032e --- /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/cube.h b/src/ailego/container/cube.h new file mode 100644 index 0000000..d9ed202 --- /dev/null +++ b/src/ailego/container/cube.h @@ -0,0 +1,414 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { +namespace internal { + +/*! Cube Policy + */ +struct CubePolicy { + //! Destructor + virtual ~CubePolicy(void) {} + + //! Assign `src` to `dst` + virtual void assign(const void *src, void **dst) = 0; + + //! Cleanup value + virtual void cleanup(void **val) = 0; + + //! Clone value of `src` to `dst` + virtual void clone(void *const *src, void **dst) = 0; + + //! Move `src` to `dst` + virtual void move(void *src, void **dst) = 0; + + //! Retrieve size + virtual size_t size(void) const = 0; + + //! Retrieve type information + virtual const std::type_info &type(void) const = 0; + + //! Retrieve value + virtual void *value(void **src) = 0; + + //! Retrieve value + virtual const void *value(void *const *src) const = 0; +}; + +/*! Small Cube Policy + */ +template +struct SmallCubePolicy : public CubePolicy { + //! Assign `src` to `dst` + void assign(const void *src, void **dst) { + new (dst) T(*reinterpret_cast(src)); + } + + //! Cleanup value + void cleanup(void **val) { + reinterpret_cast(val)->~T(); + } + + //! Clone value of `src` to `dst` + void clone(void *const *src, void **dst) { + new (dst) T(*reinterpret_cast(src)); + } + + //! Move `src` to `dst` + void move(void *src, void **dst) { + new (dst) T(std::move(*reinterpret_cast(src))); + } + + //! Retrieve size + size_t size(void) const { + return sizeof(T); + } + + //! Retrieve type information + const std::type_info &type(void) const { + return typeid(T); + } + + //! Retrieve value + void *value(void **src) { + return reinterpret_cast(src); + } + + //! Retrieve value + const void *value(void *const *src) const { + return reinterpret_cast(src); + } +}; + +/*! Large Cube Policy + */ +template +struct LargeCubePolicy : public CubePolicy { + //! Assign `src` to `dst` + void assign(const void *src, void **dst) { + *dst = new T(*reinterpret_cast(src)); + } + + //! Cleanup value + void cleanup(void **val) { + delete (reinterpret_cast(*val)); + } + + //! Clone value of `src` to `dst` + void clone(void *const *src, void **dst) { + *dst = new T(**reinterpret_cast(src)); + } + + //! Move `src` to `dst` + void move(void *src, void **dst) { + *dst = new T(std::move(*reinterpret_cast(src))); + } + + //! Retrieve size + size_t size(void) const { + return sizeof(T); + } + + //! Retrieve type information + const std::type_info &type(void) const { + return typeid(T); + } + + //! Retrieve value + void *value(void **src) { + return *src; + } + + //! Retrieve value + const void *value(void *const *src) const { + return *src; + } +}; + +/*! Policy Selector + */ +template +struct PolicySelector { + typedef LargeCubePolicy Type; +}; + +/*! Policy Selector + */ +template +struct PolicySelector< + T, typename std::enable_if::type> { + typedef SmallCubePolicy Type; +}; + +} // namespace internal + +/*! Cube class + */ +class Cube { + public: + //! Constructor + Cube(void) : policy_(Cube::Policy()), object_(nullptr) {} + + //! Constructor + template + Cube(const T &rhs) : policy_(Cube::Policy()), object_(nullptr) { + policy_->assign(&rhs, &object_); + } + + //! Constructor + template ::value && + !std::is_same::value>::type> + Cube(T &&rhs) : policy_(Cube::Policy()), object_(nullptr) { + policy_->move(&rhs, &object_); + } + + //! Constructor + Cube(const char *str) + : policy_(Cube::Policy()), object_(nullptr) { + std::string rhs(str); + policy_->move(&rhs, &object_); + } + + //! Constructor + Cube(char str[]) : policy_(Cube::Policy()), object_(nullptr) { + std::string rhs(str); + policy_->move(&rhs, &object_); + } + + //! Constructor + Cube(const Cube &rhs) : policy_(rhs.policy_), object_(nullptr) { + policy_->clone(&rhs.object_, &object_); + } + + //! Constructor + Cube(Cube &&rhs) : policy_(rhs.policy_), object_(rhs.object_) { + rhs.policy_ = Cube::Policy(); + rhs.object_ = nullptr; + } + + //! Destructor + ~Cube(void) { + policy_->cleanup(&object_); + } + + //! Assignment + template + Cube &operator=(const T &rhs) { + this->assign(rhs); + return *this; + } + + //! Assignment + template ::value && + !std::is_same::value>::type> + Cube &operator=(T &&rhs) { + this->assign(std::forward(rhs)); + return *this; + } + + //! Assignment + Cube &operator=(const Cube &rhs) { + this->assign(rhs); + return *this; + } + + //! Assignment + Cube &operator=(Cube &&rhs) { + this->assign(std::forward(rhs)); + return *this; + } + + //! Assignment + Cube &operator=(const char *str) { + this->assign(str); + return *this; + } + + //! Assignment + Cube &operator=(char str[]) { + this->assign(str); + return *this; + } + + //! Retrieve object in original type + template + operator T &() { + return this->cast(); + } + + //! Retrieve object in original type + template + operator const T &() const { + return this->cast(); + } + + //! Assign content + template + void assign(const T &rhs) { + policy_->cleanup(&object_); + policy_ = Cube::Policy(); + policy_->assign(&rhs, &object_); + } + + //! Assign content + template ::value && + !std::is_same::value>::type> + void assign(T &&rhs) { + policy_->cleanup(&object_); + policy_ = Cube::Policy(); + policy_->move(&rhs, &object_); + } + + //! Assign content from another Cube + void assign(const Cube &rhs) { + policy_->cleanup(&object_); + policy_ = rhs.policy_; + policy_->clone(&rhs.object_, &object_); + } + + //! Assign content from another Cube + void assign(Cube &&rhs) { + if (this != &rhs) { + policy_->cleanup(&object_); + policy_ = rhs.policy_; + object_ = rhs.object_; + rhs.policy_ = Cube::Policy(); + rhs.object_ = nullptr; + } + } + + //! Assign content + void assign(const char *str) { + policy_->cleanup(&object_); + policy_ = Cube::Policy(); + std::string rhs(str); + policy_->move(&rhs, &object_); + } + + //! Assign content + void assign(char str[]) { + policy_->cleanup(&object_); + policy_ = Cube::Policy(); + std::string rhs(str); + policy_->move(&rhs, &object_); + } + + //! Swap the content with another Cube + Cube &swap(Cube &rhs) { + std::swap(policy_, rhs.policy_); + std::swap(object_, rhs.object_); + return *this; + } + + //! Cast to the original type + template + T &cast(void) { + if (policy_ != Cube::Policy()) { + throw std::bad_cast(); + } + return *reinterpret_cast(policy_->value(&object_)); + } + + //! Cast to the original type + template + const T &cast(void) const { + if (policy_ != Cube::Policy()) { + throw std::bad_cast(); + } + return *reinterpret_cast(policy_->value(&object_)); + } + + //! Cast to the original type (unsafe) + template + T &unsafe_cast(void) { + return *reinterpret_cast(policy_->value(&object_)); + } + + //! Cast to the original type (unsafe) + template + const T &unsafe_cast(void) const { + return *reinterpret_cast(policy_->value(&object_)); + } + + //! Test if the Cube is empty + bool empty(void) const { + return (policy_ == Cube::Policy()); + } + + //! Reset Cube allocated memory + void reset(void) { + policy_->cleanup(&object_); + policy_ = Cube::Policy(); + object_ = nullptr; + } + + //! Test if the Cube is compatible with another one + bool compatible(const Cube &rhs) const { + return (policy_ == rhs.policy_ || policy_->type() == rhs.policy_->type()); + } + + //! Test if the Cube is compatible with another one + template + bool compatible(void) const { + return (policy_ == Cube::Policy() || + policy_->type() == Cube::Policy()->type()); + } + + //! Retrieve size + size_t size(void) const { + return (!this->empty() ? policy_->size() : 0u); + } + + //! Retrieve type information + const std::type_info &type(void) const { + return (!this->empty() ? policy_->type() : typeid(void)); + } + + protected: + /*! Empty Policy + */ + struct EmptyPolicy {}; + + //! Make a static policy object + template + static internal::CubePolicy *MakePolicy(void) { + static typename internal::PolicySelector::Type policy; + return (&policy); + } + + //! Retrieve a static policy object + template + static internal::CubePolicy *Policy(void) { + return MakePolicy::type>(); + } + + private: + //! Members + internal::CubePolicy *policy_; + void *object_; +}; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/container/heap.h b/src/ailego/container/heap.h new file mode 100644 index 0000000..fce0367 --- /dev/null +++ b/src/ailego/container/heap.h @@ -0,0 +1,234 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Heap Adapter + */ +template , + typename TBase = std::vector> +class Heap : public TBase { + public: + //! Constructor + Heap(void) + : TBase(), limit_(std::numeric_limits::max()), compare_() {} + + //! Constructor + template + Heap(size_t max, Args &&...args) + : TBase(), + limit_(std::max(max, 1u)), + compare_(std::forward(args)...) { + TBase::reserve(limit_); + } + + //! Constructor + Heap(const Heap &rhs) + : TBase(rhs), limit_(rhs.limit_), compare_(rhs.compare_) {} + + //! Constructor + Heap(Heap &&rhs) + : TBase(std::move(rhs)), + limit_(rhs.limit_), + compare_(std::move(rhs.compare_)) {} + + //! Constructor + Heap(const TBase &rhs) + : TBase(rhs), limit_(std::numeric_limits::max()), compare_() { + std::make_heap(TBase::begin(), TBase::end(), compare_); + } + + //! Constructor + Heap(TBase &&rhs) + : TBase(std::move(rhs)), + limit_(std::numeric_limits::max()), + compare_() { + std::make_heap(TBase::begin(), TBase::end(), compare_); + } + + //! Assignment + Heap &operator=(const Heap &rhs) { + TBase::operator=(static_cast(rhs)); + limit_ = rhs.limit_; + compare_ = rhs.compare_; + return *this; + } + + //! Assignment + Heap &operator=(Heap &&rhs) { + TBase::operator=(std::move(static_cast(rhs))); + limit_ = rhs.limit_; + compare_ = std::move(rhs.compare_); + return *this; + } + + //! Exchange the content + void swap(Heap &rhs) { + TBase::swap(static_cast(rhs)); + std::swap(limit_, rhs.limit_); + std::swap(compare_, rhs.compare_); + } + + //! Pop the front element + void pop(void) { + if (TBase::size() > 1) { + auto last = TBase::end() - 1; + this->replace_heap(TBase::begin(), last, std::move(*last)); + } + TBase::pop_back(); + } + + //! Insert a new element into the heap + template + void emplace(TArgs &&...args) { + if (this->full()) { + typename std::remove_reference::type val(std::forward(args)...); + + auto first = TBase::begin(); + if (compare_(val, *first)) { + this->replace_heap(first, TBase::end(), std::move(val)); + } + } else { + TBase::emplace_back(std::forward(args)...); + std::push_heap(TBase::begin(), TBase::end(), compare_); + } + } + + //! Insert a new element into the heap + void push(const T &val) { + if (this->full()) { + auto first = TBase::begin(); + if (compare_(val, *first)) { + this->replace_heap(first, TBase::end(), val); + } + } else { + TBase::push_back(val); + std::push_heap(TBase::begin(), TBase::end(), compare_); + } + } + + //! Insert a new element into the heap + void push(T &&val) { + if (this->full()) { + auto first = TBase::begin(); + if (compare_(val, *first)) { + this->replace_heap(first, TBase::end(), std::move(val)); + } + } else { + TBase::push_back(std::move(val)); + std::push_heap(TBase::begin(), TBase::end(), compare_); + } + } + + //! Retrieve the limit of heap + size_t limit(void) const { + return limit_; + } + + //! Limit the heap with max size + void limit(size_t max) { + limit_ = std::max(max, 1u); + TBase::reserve(limit_); + } + + //! Unlimit the size of heap + void unlimit(void) { + limit_ = std::numeric_limits::max(); + } + + //! Check whether the heap is full + bool full(void) const { + return (TBase::size() == limit_); + } + + //! Update the heap + void update(void) { + std::make_heap(TBase::begin(), TBase::end(), compare_); + while (limit_ < TBase::size()) { + this->pop(); + } + } + + //! Sort the elements in the heap + void sort(void) { + std::sort(TBase::begin(), TBase::end(), compare_); + } + + protected: + //! Replace the top element of heap + template + void replace_heap(TRandomIterator first, TRandomIterator last, TValue &&val) { + using _DistanceType = + typename std::iterator_traits::difference_type; + + _DistanceType hole = 0; + _DistanceType count = _DistanceType(last - first); + + if (count > 1) { + _DistanceType child = (hole << 1) + 1; + + while (child < count) { + _DistanceType right_child = child + 1; + + if (right_child < count && + compare_(*(first + child), *(first + right_child))) { + child = right_child; + } + if (!compare_(val, *(first + child))) { + break; + } + *(first + hole) = std::move(*(first + child)); + hole = child; + child = (hole << 1) + 1; + } + } + *(first + hole) = std::forward(val); + } + + private: + size_t limit_; + TCompare compare_; +}; + +/*! Key Value Heap Comparer + */ +template > +struct KeyValueHeapComparer { + //! Function call + bool operator()(const std::pair &lhs, + const std::pair &rhs) const { + return compare_(lhs.second, rhs.second); + } + + private: + TCompare compare_; +}; + +/*! Key Value Heap + */ +template > +using KeyValueHeap = + Heap, KeyValueHeapComparer>; + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/container/hypercube.h b/src/ailego/container/hypercube.h new file mode 100644 index 0000000..490714e --- /dev/null +++ b/src/ailego/container/hypercube.h @@ -0,0 +1,268 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "cube.h" + +namespace zvec { +namespace ailego { + +/*! Hypercube + */ +class Hypercube { + public: + //! Constructor + Hypercube(void) : cubes_() {} + + //! Constructor + Hypercube(const Hypercube &rhs) : cubes_(rhs.cubes_) {} + + //! Constructor + Hypercube(Hypercube &&rhs) : cubes_() { + cubes_.swap(rhs.cubes_); + } + + //! Destructor + ~Hypercube(void) {} + + //! Assignment + Hypercube &operator=(const Hypercube &rhs) { + cubes_ = rhs.cubes_; + return *this; + } + + //! Assignment + Hypercube &operator=(Hypercube &&rhs) { + cubes_ = std::move(rhs.cubes_); + return *this; + } + + //! Overloaded operator [] + Cube &operator[](const std::string &key) { + return cubes_[key]; + } + + //! Overloaded operator [] + Cube &operator[](std::string &&key) { + return cubes_[std::forward(key)]; + } + + //! Test if the element is exist + bool has(const std::string &key) const { + return (cubes_.find(key) != cubes_.end()); + } + + //! Test if the hyper cube is empty + bool empty(void) const { + return cubes_.empty(); + } + + //! Insert a key-value pair into map + bool insert(const std::string &key, Cube &&val) { + return cubes_.emplace(key, std::forward(val)).second; + } + + //! Insert a key-value pair into map + bool insert(std::string &&key, Cube &&val) { + return cubes_ + .emplace(std::forward(key), std::forward(val)) + .second; + } + + //! Insert a key-value pair into map + template + bool insert(const std::string &key, T &&val) { + return cubes_.emplace(key, Cube(std::forward(val))).second; + } + + //! Insert a key-value pair into map + template + bool insert(std::string &&key, T &&val) { + return cubes_ + .emplace(std::forward(key), Cube(std::forward(val))) + .second; + } + + //! Insert or assign a key-value pair to map + void insert_or_assign(const std::string &key, Cube &&val) { + auto it = cubes_.lower_bound(key); + if (it != cubes_.end() && it->first == key) { + it->second = std::forward(val); + } else { + cubes_.emplace_hint(it, key, std::forward(val)); + } + } + + //! Insert or assign a key-value pair to map + void insert_or_assign(std::string &&key, Cube &&val) { + auto it = cubes_.lower_bound(key); + if (it != cubes_.end() && it->first == key) { + it->second = std::forward(val); + } else { + cubes_.emplace_hint(it, std::forward(key), + std::forward(val)); + } + } + + //! Insert or assign a key-value pair to map + template + void insert_or_assign(const std::string &key, T &&val) { + auto it = cubes_.lower_bound(key); + if (it != cubes_.end() && it->first == key) { + it->second = Cube(std::forward(val)); + } else { + cubes_.emplace_hint(it, key, Cube(std::forward(val))); + } + } + + //! Insert or assign a key-value pair to map + template + void insert_or_assign(std::string &&key, T &&val) { + auto it = cubes_.lower_bound(key); + if (it != cubes_.end() && it->first == key) { + it->second = Cube(std::forward(val)); + } else { + cubes_.emplace_hint(it, std::forward(key), + Cube(std::forward(val))); + } + } + + //! Clear the map + void clear(void) { + cubes_.clear(); + } + + //! Swap the map + void swap(Hypercube &rhs) { + cubes_.swap(rhs.cubes_); + } + + //! Erase the pair via a key + bool erase(const std::string &key) { + auto iter = cubes_.find(key); + if (iter != cubes_.end()) { + cubes_.erase(iter); + return true; + } + return false; + } + + //! Retrieve the value via a key + bool get(const std::string &key, Cube *out) const { + auto iter = cubes_.find(key); + if (iter != cubes_.end()) { + *out = iter->second; + return true; + } + return false; + } + + //! Retrieve the value via a key + Cube *get(const std::string &key) { + auto iter = cubes_.find(key); + if (iter != cubes_.end()) { + return &iter->second; + } + return nullptr; + } + + //! Retrieve the value via a key + const Cube *get(const std::string &key) const { + auto iter = cubes_.find(key); + if (iter != cubes_.end()) { + return &iter->second; + } + return nullptr; + } + + //! Retrieve the value via a key + template + bool get(const std::string &key, T *out) const { + auto iter = cubes_.find(key); + if (iter != cubes_.end()) { + if (iter->second.compatible()) { + *out = iter->second.unsafe_cast(); + return true; + } + } + return false; + } + + //! Retrieve the value via a key + template + T &get(const std::string &key, T &def) { + auto iter = cubes_.find(key); + if (iter != cubes_.end()) { + if (iter->second.compatible()) { + return iter->second.unsafe_cast(); + } + } + return def; + } + + //! Retrieve the value via a key + template + const T &get(const std::string &key, const T &def) const { + auto iter = cubes_.find(key); + if (iter != cubes_.end()) { + if (iter->second.compatible()) { + return iter->second.unsafe_cast(); + } + } + return def; + } + + //! Merge another hyper cube + void merge(const Hypercube &rhs) { + for (const auto &it : rhs.cubes_) { + auto iter = cubes_.find(it.first); + if (iter != cubes_.end()) { + iter->second = it.second; + } else { + cubes_.emplace(it.first, it.second); + } + } + } + + //! Merge another hyper cube + void merge(Hypercube &&rhs) { + for (auto &it : rhs.cubes_) { + auto iter = cubes_.find(it.first); + if (iter != cubes_.end()) { + iter->second = std::move(it.second); + } else { + cubes_.emplace(std::move(it.first), std::move(it.second)); + } + } + } + + //! Retrieve the cubes + const std::map &cubes(void) const { + return cubes_; + } + + //! Retrieve the cubes + std::map *mutable_cubes(void) { + return &cubes_; + } + + private: + std::map cubes_; +}; + +} // 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..cd62e7b --- /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 "params.h" +#include +#include +#include "ailego/logger/logger.h" + +//! 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/params.h b/src/ailego/container/params.h new file mode 100644 index 0000000..464bb6d --- /dev/null +++ b/src/ailego/container/params.h @@ -0,0 +1,776 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +//! Trying compatible with T +#define _TRYING_COMPATIBLE(cube, T, out) \ + if (cube->compatible()) \ + return ( \ + *out = static_cast::type>( \ + cube->unsafe_cast()), \ + true) + +//! Trying compatible with T (Boolean) +#define _TRYING_COMPATIBLE_BOOL(cube, T, out) \ + if (cube->compatible()) return (*out = !!cube->unsafe_cast(), true) + +//! Trying compatible with T (String) +#define _TRYING_COMPATIBLE_STRING(cube, T, out) \ + if (cube->compatible()) \ + return (out->assign(std::to_string(cube->unsafe_cast())), true) + +//! Trying convert string +#define _TRYING_CONVERT_STRING(cube, out) \ + if (cube->compatible()) \ + return (*out = Params::StringCast::type>( \ + cube->unsafe_cast()), \ + true) + +/*! Index Params + */ +class Params { + public: + //! Constructor + Params(void) : hypercube_() {} + + //! Constructor + Params(const Params &rhs) : hypercube_(rhs.hypercube_) {} + + //! Constructor + Params(Params &&rhs) : hypercube_() { + hypercube_.swap(rhs.hypercube_); + } + + //! Destructor + ~Params(void) {} + + //! Assignment + Params &operator=(const Params &rhs) { + hypercube_ = rhs.hypercube_; + return *this; + } + + //! Assignment + Params &operator=(Params &&rhs) { + hypercube_.swap(rhs.hypercube_); + return *this; + } + + //! Overloaded operator [] + ailego::Cube &operator[](const std::string &key) { + return hypercube_[key]; + } + + //! Overloaded operator [] + ailego::Cube &operator[](std::string &&key) { + return hypercube_[std::move(key)]; + } + + //! Test if the element is exist + bool has(const std::string &key) const { + return hypercube_.has(key); + } + + //! Test if the map is empty + bool empty(void) const { + return hypercube_.empty(); + } + + //! Clear the map + void clear(void) { + hypercube_.clear(); + } + + //! Erase the pair via a key + bool erase(const std::string &key) { + return hypercube_.erase(key); + } + + //! Merge another index params + void merge(const Params &rhs) { + hypercube_.merge(rhs.hypercube_); + } + + //! Merge another index params + void merge(Params &&rhs) { + hypercube_.merge(std::move(rhs.hypercube_)); + } + + //! Set the value of key in T + template + bool insert(const std::string &key, T &&val) { + return hypercube_.insert(key, std::forward(val)); + } + + //! Set the value of key in T + template + bool insert(std::string &&key, T &&val) { + return hypercube_.insert(std::forward(key), + std::forward(val)); + } + + //! Set the value of key in T + template + void set(const std::string &key, T &&val) { + hypercube_.insert_or_assign(key, std::forward(val)); + } + + //! Set the value of key in T + template + void set(std::string &&key, T &&val) { + hypercube_.insert_or_assign(std::forward(key), + std::forward(val)); + } + + //! Retrieve the value in boolean + bool get(const std::string &key, bool *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE_BOOL(cube, char, out); + _TRYING_COMPATIBLE_BOOL(cube, unsigned char, out); + _TRYING_COMPATIBLE_BOOL(cube, signed char, out); + _TRYING_COMPATIBLE_BOOL(cube, short int, out); + _TRYING_COMPATIBLE_BOOL(cube, unsigned short int, out); + _TRYING_COMPATIBLE_BOOL(cube, int, out); + _TRYING_COMPATIBLE_BOOL(cube, unsigned int, out); + _TRYING_COMPATIBLE_BOOL(cube, long int, out); + _TRYING_COMPATIBLE_BOOL(cube, unsigned long int, out); + _TRYING_COMPATIBLE_BOOL(cube, long long int, out); + _TRYING_COMPATIBLE_BOOL(cube, unsigned long long int, out); + _TRYING_COMPATIBLE_BOOL(cube, float, out); + _TRYING_COMPATIBLE_BOOL(cube, double, out); + _TRYING_COMPATIBLE_BOOL(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'char' + bool get(const std::string &key, char *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'unsigned char' + bool get(const std::string &key, unsigned char *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'signed char' + bool get(const std::string &key, signed char *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'short int' + bool get(const std::string &key, short int *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'unsigned short int' + bool get(const std::string &key, unsigned short int *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'int' + bool get(const std::string &key, int *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'unsigned int' + bool get(const std::string &key, unsigned int *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'long int' + bool get(const std::string &key, long int *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'unsigned long int' + bool get(const std::string &key, unsigned long int *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'long long int' + bool get(const std::string &key, long long int *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'unsigned long long int' + bool get(const std::string &key, unsigned long long int *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'float' + bool get(const std::string &key, float *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'double' + bool get(const std::string &key, double *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in 'long double' + bool get(const std::string &key, long double *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, long double, out); + _TRYING_COMPATIBLE(cube, double, out); + _TRYING_COMPATIBLE(cube, float, out); + _TRYING_COMPATIBLE(cube, long long int, out); + _TRYING_COMPATIBLE(cube, unsigned long long int, out); + _TRYING_COMPATIBLE(cube, long int, out); + _TRYING_COMPATIBLE(cube, unsigned long int, out); + _TRYING_COMPATIBLE(cube, int, out); + _TRYING_COMPATIBLE(cube, unsigned int, out); + _TRYING_COMPATIBLE(cube, short int, out); + _TRYING_COMPATIBLE(cube, unsigned short int, out); + _TRYING_COMPATIBLE(cube, char, out); + _TRYING_COMPATIBLE(cube, unsigned char, out); + _TRYING_COMPATIBLE(cube, signed char, out); + _TRYING_COMPATIBLE(cube, bool, out); + _TRYING_CONVERT_STRING(cube, out); + } + return false; + } + + //! Retrieve the value in string + bool get(const std::string &key, std::string *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, std::string, out); + _TRYING_COMPATIBLE_STRING(cube, bool, out); + _TRYING_COMPATIBLE_STRING(cube, char, out); + _TRYING_COMPATIBLE_STRING(cube, unsigned char, out); + _TRYING_COMPATIBLE_STRING(cube, signed char, out); + _TRYING_COMPATIBLE_STRING(cube, short int, out); + _TRYING_COMPATIBLE_STRING(cube, unsigned short int, out); + _TRYING_COMPATIBLE_STRING(cube, int, out); + _TRYING_COMPATIBLE_STRING(cube, unsigned int, out); + _TRYING_COMPATIBLE_STRING(cube, long int, out); + _TRYING_COMPATIBLE_STRING(cube, unsigned long int, out); + _TRYING_COMPATIBLE_STRING(cube, long long int, out); + _TRYING_COMPATIBLE_STRING(cube, unsigned long long int, out); + _TRYING_COMPATIBLE_STRING(cube, float, out); + _TRYING_COMPATIBLE_STRING(cube, double, out); + _TRYING_COMPATIBLE_STRING(cube, long double, out); + } + return false; + } + + //! Retrieve the value in T + template + bool get(const std::string &key, T *out) const { + const ailego::Cube *cube = hypercube_.get(key); + if (cube) { + _TRYING_COMPATIBLE(cube, T, out); + } + return false; + } + + //! Retrieve the value in boolean + bool get_as_bool(const std::string &key) const { + bool result = false; + this->get(key, &result); + return result; + } + + //! Retrieve the value in int8 + int8_t get_as_int8(const std::string &key) const { + int8_t result = 0; + this->get(key, &result); + return result; + } + + //! Retrieve the value in int16 + int16_t get_as_int16(const std::string &key) const { + int16_t result = 0; + this->get(key, &result); + return result; + } + + //! Retrieve the value in int32 + int32_t get_as_int32(const std::string &key) const { + int32_t result = 0; + this->get(key, &result); + return result; + } + + //! Retrieve the value in int64 + int64_t get_as_int64(const std::string &key) const { + int64_t result = 0; + this->get(key, &result); + return result; + } + + //! Retrieve the value in uint8 + uint8_t get_as_uint8(const std::string &key) const { + uint8_t result = 0; + this->get(key, &result); + return result; + } + + //! Retrieve the value in uint16 + uint16_t get_as_uint16(const std::string &key) const { + uint16_t result = 0; + this->get(key, &result); + return result; + } + + //! Retrieve the value in uint32 + uint32_t get_as_uint32(const std::string &key) const { + uint32_t result = 0; + this->get(key, &result); + return result; + } + + //! Retrieve the value in uint64 + uint64_t get_as_uint64(const std::string &key) const { + uint64_t result = 0; + this->get(key, &result); + return result; + } + + //! Retrieve the value in float + float get_as_float(const std::string &key) const { + float result = 0.0f; + this->get(key, &result); + return result; + } + + //! Retrieve the value in double + double get_as_double(const std::string &key) const { + double result = 0.0f; + this->get(key, &result); + return result; + } + + //! Retrieve the value in string + std::string get_as_string(const std::string &key) const { + std::string result; + this->get(key, &result); + return result; + } + + //! Retrieve the debug string + std::string debug_string(void) const { + std::string str; + SerializeToBuffer(*this, &str); + return str; + } + + //! Retrieve the map of parameters + const ailego::Hypercube &hypercube(void) const { + return hypercube_; + } + + //! Retrieve the map of parameters + ailego::Hypercube *mutable_hypercube(void) { + return &hypercube_; + } + + //! Parse parameters from buffer (Json format) + static bool ParseFromBuffer(const std::string &buf, Params *params); + + //! Parse parameters from OS environment + static void ParseFromEnvironment(Params *params); + + //! Serialize parameters into buffer + static void SerializeToBuffer(const Params ¶ms, std::string *buf); + + protected: + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return std::strtof(str.c_str(), nullptr); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return std::strtod(str.c_str(), nullptr); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return std::strtold(str.c_str(), nullptr); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return static_cast(std::strtol(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return static_cast(std::strtol(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return static_cast(std::strtoul(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return static_cast(std::strtol(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return static_cast(std::strtol(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return static_cast(std::strtol(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return static_cast(std::strtoll(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, + T>::type { + return static_cast( + std::strtoul(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + return static_cast(std::strtoul(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, + T>::type { + return static_cast( + std::strtoul(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, + T>::type { + return static_cast( + std::strtoull(str.c_str(), nullptr, 0)); + } + + //! Convert string type to another type + template + static auto StringCast(const std::string &str) -> + typename std::enable_if::value, T>::type { + if (str.empty()) { + return false; + } + char c = str[0]; + if (c == 'Y' || c == 'T' || c == 'y' || c == 't') { + return true; + } + return !!std::strtof(str.c_str(), nullptr); + } + + private: + ailego::Hypercube hypercube_; +}; + +#undef _TRYING_COMPATIBLE +#undef _TRYING_COMPATIBLE +#undef _TRYING_COMPATIBLE_BOOL +#undef _TRYING_COMPATIBLE_STRING +#undef _TRYING_CONVERT_STRING + +} // namespace ailego +} // namespace zvec \ No newline at end of file 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.h b/src/ailego/container/vector.h new file mode 100644 index 0000000..98cdef3 --- /dev/null +++ b/src/ailego/container/vector.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 + +namespace zvec { +namespace ailego { + +/*! Fixed Vector + */ +template +class FixedVector { + public: + enum { MAX_SIZE = N }; + + //! Constructor + template + FixedVector(U... vals) : data_{vals...} {} + + //! Overloaded operator [] + T &operator[](size_t i) { + return data_[i]; + } + + //! Overloaded operator [] + constexpr const T &operator[](size_t i) const { + return data_[i]; + } + + //! Retrieve data pointer + T *data(void) { + return data_; + } + + //! Retrieve data pointer + const T *data(void) const { + return data_; + } + + //! Retrieve count of elements in vector + constexpr size_t size(void) const { + return MAX_SIZE; + } + + //! Convert a array pointer to vector pointer + static FixedVector *Cast(T arr[N]) { + return reinterpret_cast *>(arr); + } + + //! Convert a array pointer to vector pointer + static const FixedVector *Cast(const T arr[N]) { + return reinterpret_cast *>(arr); + } + + private: + //! Data member + T data_[N]; +}; + +/*! Numerical Vector Adapter + */ +template ::value>::type> +class NumericalVector : public TBase { + public: + typedef typename std::remove_cv::type ValueType; + typedef ValueType *iterator; + typedef const ValueType *const_iterator; + + //! Constructor + NumericalVector(void) : TBase() {} + + //! Constructor + explicit NumericalVector(size_t dim) : TBase() { + this->resize(dim); + } + + //! Constructor + NumericalVector(size_t dim, const ValueType &val) : TBase() { + this->resize(dim, val); + } + + //! Constructor + NumericalVector(const NumericalVector &rhs) : TBase(rhs) {} + + //! Constructor + NumericalVector(NumericalVector &&rhs) : TBase(std::forward(rhs)) {} + + //! Constructor + NumericalVector(const TBase &rhs) : TBase(rhs) { + if (TBase::size() % sizeof(T) != 0) { + throw std::length_error("Unmatched length"); + } + } + + //! Constructor + NumericalVector(TBase &&rhs) : TBase(std::move(rhs)) { + if (TBase::size() % sizeof(T) != 0) { + throw std::length_error("Unmatched length"); + } + } + + //! Constructor + NumericalVector(std::initializer_list il) : TBase() { + for (const auto &it : il) { + TBase::append(reinterpret_cast(&it), sizeof(ValueType)); + } + } + + //! Assignment + NumericalVector &operator=(const NumericalVector &rhs) { + TBase::operator=(static_cast(rhs)); + return *this; + } + + //! Assignment + NumericalVector &operator=(NumericalVector &&rhs) { + TBase::operator=(std::move(static_cast(rhs))); + return *this; + } + + //! Assignment + NumericalVector &operator=(const TBase &rhs) { + TBase::operator=(rhs); + return *this; + } + + //! Assignment + NumericalVector &operator=(TBase &&rhs) { + TBase::operator=(std::move(rhs)); + return *this; + } + + //! Overloaded operator [] + ValueType &operator[](size_t i) { + return *(this->data() + i); + } + + //! Overloaded operator [] + const ValueType &operator[](size_t i) const { + return *(this->data() + i); + } + + //! Appends a copy of value + NumericalVector &append(const ValueType &val) { + TBase::append(reinterpret_cast(&val), sizeof(ValueType)); + return *this; + } + + //! Append a copy of value + void append(std::initializer_list il) { + for (const auto &it : il) { + TBase::append(reinterpret_cast(&it), sizeof(ValueType)); + } + } + + //! Assign content to vector + void assign(const ValueType *vec, size_t len) { + TBase::assign(reinterpret_cast(vec), len * sizeof(ValueType)); + } + + //! Assign content to vector + void assign(size_t n, const ValueType &val) { + this->clear(); + this->resize(n, val); + } + + //! Assign content to vector + void assign(std::initializer_list il) { + this->clear(); + for (const auto &it : il) { + TBase::append(reinterpret_cast(&it), sizeof(ValueType)); + } + } + + //! Retrieve element + ValueType &at(size_t i) { + return *(this->data() + i); + } + + //! Retrieve element + const ValueType &at(size_t i) const { + return *(this->data() + i); + } + + //! Access last element + ValueType &back(void) { + return *(this->rbegin()); + } + + //! Access last element + const ValueType &back(void) const { + return *(this->rbegin()); + } + + //! Retrieve iterator to beginning + iterator begin(void) { + return this->data(); + } + + //! Retrieve iterator to beginning + const_iterator begin(void) const { + return this->data(); + } + + //! Retrieve size of allocated storage + size_t capacity(void) const { + return (TBase::capacity() / sizeof(ValueType)); + } + + //! Clear the vector + void clear(void) { + TBase::clear(); + } + + //! Retrieve pointer of data + ValueType *data(void) { + return reinterpret_cast(&(TBase::operator[](0))); + } + + //! Retrieve pointer of data + const ValueType *data(void) const { + return reinterpret_cast(TBase::data()); + } + + //! Test if vector is empty + bool empty(void) const { + return TBase::empty(); + } + + //! An iterator to the past-the-end + iterator end(void) { + return (this->data() + this->size()); + } + + //! An iterator to the past-the-end + const_iterator end(void) const { + return (this->data() + this->size()); + } + + //! Access first element + ValueType &front(void) { + return *(this->begin()); + } + + //! Access first element + const ValueType &front(void) const { + return *(this->begin()); + } + + //! Request a change in capacity + void reserve(size_t n) { + TBase::reserve(n * sizeof(ValueType)); + } + + //! Resize the vector to a length of n elements + void resize(size_t n) { + TBase::resize(n * sizeof(ValueType)); + } + + //! Resize the vector to a length of n elements + void resize(size_t n, const ValueType &val) { + size_t count = this->size(); + + TBase::resize(n * sizeof(ValueType)); + for (size_t i = count; i < n; ++i) { + *(this->data() + i) = val; + } + } + + //! Retrieve dimension of vector + size_t size(void) const { + return (TBase::size() / sizeof(ValueType)); + } + + //! Retrieve dimension of vector + size_t dimension(void) const { + return (TBase::size() / sizeof(ValueType)); + } + + //! Retrieve size of vector in bytes + size_t bytes(void) const { + return TBase::size(); + } + + //! Swap vector values + void swap(NumericalVector &vec) { + TBase::swap(static_cast(vec)); + } +}; + +/*! Nibble Vector Adapter + */ +template ::value>::type> +class NibbleVector : public TBase { + public: + //! Type of value + using ValueType = typename std::remove_cv::type; + using StoreType = typename std::make_unsigned::type; + + //! const_iterator of Nibble Vector + class const_iterator { + public: + //! Constructor + const_iterator(void) : i_(0), owner_(nullptr) {} + + //! Constructor + const_iterator(const NibbleVector *owner, size_t i) + : i_(i), owner_(owner) {} + + //! Equality + bool operator==(const const_iterator &rhs) const { + return (i_ == rhs.i_); + } + + //! No equality + bool operator!=(const const_iterator &rhs) const { + return (i_ != rhs.i_); + } + + //! Increment (Prefix) + const_iterator &operator++() { + ++i_; + return *this; + } + + //! Increment (Suffix) + const_iterator operator++(int) { + const_iterator tmp = *this; + ++i_; + return tmp; + } + + //! Decrement (Prefix) + const_iterator &operator--() { + --i_; + return *this; + } + + //! Decrement (Suffix) + const_iterator operator--(int) { + const_iterator tmp = *this; + --i_; + return tmp; + } + + //! operator "+=" + const_iterator &operator+=(size_t offset) { + i_ += offset; + return *this; + } + + //! operator "-=" + const_iterator &operator-=(size_t offset) { + i_ -= offset; + return *this; + } + + //! Indirection (Signed integral) + ValueType operator*() const { + return owner_->element(i_); + } + + private: + size_t i_; + const NibbleVector *owner_; + }; + + //! Constructor + NibbleVector(void) : TBase() {} + + //! Constructor + explicit NibbleVector(size_t dim) : TBase() { + this->resize(dim); + } + + //! Constructor + NibbleVector(size_t dim, ValueType val) : TBase() { + this->resize(dim, val); + } + + //! Constructor + NibbleVector(const NibbleVector &rhs) : TBase(rhs) {} + + //! Constructor + NibbleVector(NibbleVector &&rhs) : TBase(std::forward(rhs)) {} + + //! Constructor + NibbleVector(const TBase &rhs) : TBase(rhs) {} + + //! Constructor + NibbleVector(TBase &&rhs) : TBase(std::move(rhs)) {} + + //! Constructor + NibbleVector(std::initializer_list il) : TBase() { + this->resize(il.size()); + + size_t index = 0; + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + + for (auto val : il) { + arr[index >> 1] |= ((uint8_t)(val & 0xf) << ((index & 1) << 2)); + ++index; + } + } + + //! Assignment + NibbleVector &operator=(const NibbleVector &rhs) { + TBase::operator=(static_cast(rhs)); + return *this; + } + + //! Assignment + NibbleVector &operator=(NibbleVector &&rhs) { + TBase::operator=(std::move(static_cast(rhs))); + return *this; + } + + //! Assignment + NibbleVector &operator=(const TBase &rhs) { + TBase::operator=(rhs); + return *this; + } + + //! Assignment + NibbleVector &operator=(TBase &&rhs) { + TBase::operator=(std::move(rhs)); + return *this; + } + + //! Overloaded operator [] (Signed integral) + ValueType operator[](size_t i) const { + return this->at(i); + } + + //! Appends a copy of value + NibbleVector &append(ValueType lo, ValueType hi) { + TBase::push_back(((uint8_t)(hi & 0xf) << 4) | (uint8_t)(lo & 0xf)); + return *this; + } + + //! Append a copy of value + void append(std::initializer_list il) { + size_t index = this->size(); + this->resize(index + il.size()); + + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + for (auto val : il) { + arr[index >> 1] |= ((uint8_t)(val & 0xf) << ((index & 1) << 2)); + ++index; + } + } + + //! Assign content to vector + void assign(const ValueType *vec, size_t len) { + this->clear(); + this->resize(len); + + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + for (size_t i = 0; i != len; ++i) { + arr[i >> 1] |= ((uint8_t)(vec[i] & 0xf) << ((i & 1) << 2)); + } + } + + //! Assign content to vector + void assign(size_t n, ValueType val) { + this->clear(); + this->resize(n, val); + } + + //! Assign content to vector + void assign(std::initializer_list il) { + this->clear(); + this->resize(il.size()); + + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + size_t index = 0; + + for (auto val : il) { + arr[index >> 1] |= ((uint8_t)(val & 0xf) << ((index & 1) << 2)); + ++index; + } + } + + //! Set a element + void set(size_t i, ValueType val) { + uint8_t *it = reinterpret_cast(&(TBase::operator[](i >> 1))); + if (i & 1) { + *it = (*it & 0x0f) | ((uint8_t)(val & 0xf) << 4); + } else { + *it = (*it & 0xf0) | (uint8_t)(val & 0xf); + } + } + + //! Retrieve element + ValueType at(size_t i) const { + return this->element(i); + } + + //! Access last element + ValueType back(void) const { + return this->at(this->size() - 1); + } + + //! Retrieve iterator to beginning + const_iterator begin(void) const { + return const_iterator(this, 0); + } + + //! Retrieve size of allocated storage + size_t capacity(void) const { + return (TBase::capacity() << 1); + } + + //! Clear the vector + void clear(void) { + TBase::clear(); + } + + //! Retrieve pointer of data + StoreType *data(void) { + return reinterpret_cast(&(TBase::operator[](0))); + } + + //! Retrieve pointer of data + const StoreType *data(void) const { + return reinterpret_cast(TBase::data()); + } + + //! Test if vector is empty + bool empty(void) const { + return TBase::empty(); + } + + //! An iterator to the past-the-end + const_iterator end(void) const { + return const_iterator(this, this->size()); + } + + //! Access first element + ValueType front(void) const { + return this->at(0); + } + + //! Request a change in capacity + void reserve(size_t n) { + TBase::reserve((n + (sizeof(ValueType) << 1) - 1) / + (sizeof(ValueType) << 1) * sizeof(ValueType)); + } + + //! Resize the vector to a length of n elements + void resize(size_t n) { + TBase::resize((n + (sizeof(ValueType) << 1) - 1) / + (sizeof(ValueType) << 1) * sizeof(ValueType)); + } + + //! Resize the vector to a length of n elements + void resize(size_t n, ValueType val) { + TBase::resize((n + (sizeof(ValueType) << 1) - 1) / + (sizeof(ValueType) << 1) * sizeof(ValueType), + ((uint8_t)(val & 0xf) << 4) | (uint8_t)(val & 0xf)); + } + + //! Retrieve dimension of vector + size_t size(void) const { + return (TBase::size() << 1); + } + + //! Retrieve dimension of vector + size_t dimension(void) const { + return (TBase::size() << 1); + } + + //! Retrieve size of vector in bytes + size_t bytes(void) const { + return TBase::size(); + } + + //! Swap vector values + void swap(NibbleVector &vec) { + TBase::swap(static_cast(vec)); + } + + protected: + //! Retrieve element (Signed integral) + template + auto element(size_t i) const -> + typename std::enable_if::value, U>::type { + const uint8_t *arr = reinterpret_cast(TBase::data()); + return (static_cast(arr[i >> 1] << (~(i << 2) & 4)) >> 4); + } + + //! Retrieve element (Unsigned integral) + template + auto element(size_t i) const -> + typename std::enable_if::value, U>::type { + const uint8_t *arr = reinterpret_cast(TBase::data()); + return ((arr[i >> 1] >> ((i & 1) << 2)) & 0xf); + } +}; + +/*! Binary Vector Adapter + */ +template ::value>::type> +class BinaryVector : public TBase { + public: + //! Type of value + using ValueType = typename std::remove_cv::type; + + //! const_iterator of Binary Vector + class const_iterator { + public: + //! Constructor + const_iterator(void) : i_(0), arr_(nullptr) {} + + //! Constructor + const_iterator(const void *buf, size_t i) + : i_(i), arr_(reinterpret_cast(buf)) {} + + //! Equality + bool operator==(const const_iterator &rhs) const { + return (i_ == rhs.i_); + } + + //! No equality + bool operator!=(const const_iterator &rhs) const { + return (i_ != rhs.i_); + } + + //! Increment (Prefix) + const_iterator &operator++() { + ++i_; + return *this; + } + + //! Increment (Suffix) + const_iterator operator++(int) { + const_iterator tmp = *this; + ++i_; + return tmp; + } + + //! Decrement (Prefix) + const_iterator &operator--() { + --i_; + return *this; + } + + //! Decrement (Suffix) + const_iterator operator--(int) { + const_iterator tmp = *this; + --i_; + return tmp; + } + + //! operator "+=" + const_iterator &operator+=(size_t offset) { + i_ += offset; + return *this; + } + + //! operator "-=" + const_iterator &operator-=(size_t offset) { + i_ -= offset; + return *this; + } + + //! Indirection (eg. *iter) + bool operator*() const { + return ((arr_[i_ >> 3] & (1u << (i_ & 7))) != 0); + } + + private: + size_t i_; + const uint8_t *arr_; + }; + + //! Constructor + BinaryVector(void) : TBase() {} + + //! Constructor + explicit BinaryVector(size_t dim) : TBase() { + this->resize(dim); + } + + //! Constructor + BinaryVector(size_t dim, bool val) : TBase() { + this->resize(dim, val); + } + + //! Constructor + BinaryVector(const BinaryVector &rhs) : TBase(rhs) {} + + //! Constructor + BinaryVector(BinaryVector &&rhs) : TBase(std::move(rhs)) {} + + //! Constructor + BinaryVector(const TBase &rhs) : TBase(rhs) { + if (TBase::size() % sizeof(T) != 0) { + throw std::length_error("Unmatched length"); + } + } + + //! Constructor + BinaryVector(TBase &&rhs) : TBase(std::move(rhs)) { + if (TBase::size() % sizeof(T) != 0) { + throw std::length_error("Unmatched length"); + } + } + + //! Constructor + BinaryVector(std::initializer_list il) : TBase() { + this->resize(il.size()); + + size_t index = 0; + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + + for (auto val : il) { + if (val) { + arr[index >> 3] |= (uint8_t)(1u << (index & 7)); + } + ++index; + } + } + + //! Assignment + BinaryVector &operator=(const BinaryVector &rhs) { + TBase::operator=(static_cast(rhs)); + return *this; + } + + //! Assignment + BinaryVector &operator=(BinaryVector &&rhs) { + TBase::operator=(std::move(static_cast(rhs))); + return *this; + } + + //! Assignment + BinaryVector &operator=(const TBase &rhs) { + TBase::operator=(rhs); + return *this; + } + + //! Assignment + BinaryVector &operator=(TBase &&rhs) { + TBase::operator=(std::move(rhs)); + return *this; + } + + //! Overloaded operator [] + bool operator[](size_t i) const { + const uint8_t *arr = reinterpret_cast(TBase::data()); + return ((arr[i >> 3] & (1u << (i & 7))) != 0); + } + + //! Assign content to vector + void assign(const bool *vec, size_t len) { + this->clear(); + this->resize(len); + + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + for (size_t i = 0; i < len; ++i) { + bool val = vec[i]; + if (val) { + arr[i >> 3] |= (1u << (i & 7)); + } + } + } + + //! Assign content to vector + void assign(size_t n, bool val) { + this->clear(); + this->resize(n, val); + } + + //! Assign content to vector + void assign(std::initializer_list il) { + this->clear(); + this->resize(il.size()); + + size_t index = 0; + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + for (auto val : il) { + if (val) { + arr[index >> 3] |= (uint8_t)(1u << (index & 7)); + } + ++index; + } + } + + //! Retrieve element + bool at(size_t i) const { + const uint8_t *arr = reinterpret_cast(TBase::data()); + return ((arr[i >> 3] & (1u << (i & 7))) != 0); + } + + //! Set a bit + void set(size_t i) { + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + arr[i >> 3] |= (uint8_t)(1u << (i & 7)); + } + + //! Reset a bit + void reset(size_t i) { + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + arr[i >> 3] &= (uint8_t)(~(1u << (i & 7))); + } + + //! Toggle a bit + void flip(size_t i) { + uint8_t *arr = reinterpret_cast(&(TBase::operator[](0))); + arr[i >> 3] ^= (uint8_t)(1u << (i & 7)); + } + + //! Access last element + bool back(void) const { + return this->at(this->size() - 1); + } + + //! Retrieve const_iterator to beginning + const_iterator begin(void) const { + return const_iterator(this->data(), 0); + } + + //! Retrieve size of allocated storage + size_t capacity(void) const { + return (TBase::capacity() << 3); + } + + //! Clear the vector + void clear(void) { + TBase::clear(); + } + + //! Retrieve pointer of data + ValueType *data(void) { + return reinterpret_cast(&(TBase::operator[](0))); + } + + //! Retrieve pointer of data + const ValueType *data(void) const { + return reinterpret_cast(TBase::data()); + } + + //! Test if vector is empty + bool empty(void) const { + return TBase::empty(); + } + + //! An const_iterator to the past-the-end + const_iterator end(void) const { + return const_iterator(this->data(), this->size()); + } + + //! Access first element + bool front(void) const { + return this->at(0); + } + + //! Request a change in capacity + void reserve(size_t n) { + TBase::reserve((n + (sizeof(ValueType) << 3) - 1) / + (sizeof(ValueType) << 3) * sizeof(ValueType)); + } + + //! Resize the vector to a length of n elements + void resize(size_t n) { + TBase::resize((n + (sizeof(ValueType) << 3) - 1) / + (sizeof(ValueType) << 3) * sizeof(ValueType)); + } + + //! Resize the vector to a length of n elements + void resize(size_t n, bool val) { + TBase::resize((n + (sizeof(ValueType) << 3) - 1) / + (sizeof(ValueType) << 3) * sizeof(ValueType), + val ? 0xffu : 0u); + } + + //! Retrieve dimension of vector + size_t size(void) const { + return (TBase::size() << 3); + } + + //! Retrieve dimension of vector + size_t dimension(void) const { + return (TBase::size() << 3); + } + + //! Retrieve size of vector in bytes + size_t bytes(void) const { + return TBase::size(); + } + + //! Swap vector values + void swap(BinaryVector &vec) { + TBase::swap(static_cast(vec)); + } +}; + +/*! Hybrid Vector Adapter + */ +template ::value>::type> +class HybridVector : public NumericalVector { + public: + typedef typename std::remove_cv::type ValueType; + typedef ValueType *iterator; + typedef const ValueType *const_iterator; + + //! Constructor + HybridVector(void) = default; + + //! Constructor + explicit HybridVector(size_t dim) : NumericalVector(dim) {} + + //! Constructor + HybridVector(const HybridVector &rhs) + : NumericalVector(rhs), + sparse_count_(rhs.sparse_count_), + sparse_indices_(rhs.sparse_indices_), + sparse_data_(rhs.sparse_data_) {} + + //! Constructor + HybridVector(HybridVector &&rhs) + : NumericalVector(std::forward>(rhs)), + sparse_count_(rhs.sparse_count_), + sparse_indices_(std::move(rhs.sparse_indices_)), + sparse_data_(std::move(rhs.sparse_data_)) {} + + //! Assignment + HybridVector &operator=(const HybridVector &rhs) { + NumericalVector::operator=( + static_cast &>(rhs)); + sparse_count_ = rhs.sparse_count_; + sparse_indices_ = rhs.sparse_indices_; + sparse_data_ = rhs.sparse_data_; + + return *this; + } + + //! Assignment + HybridVector &operator=(HybridVector &&rhs) { + NumericalVector::operator=( + std::move(static_cast &&>(rhs))); + sparse_count_ = rhs.sparse_count_; + sparse_indices_ = std::move(rhs.sparse_indices_); + sparse_data_ = std::move(rhs.sparse_data_); + + return *this; + } + + size_t sparse_count() const { + return sparse_count_; + } + + const uint32_t *sparse_indices() const { + return reinterpret_cast(sparse_indices_.data()); + } + + const ValueType *sparse_data(void) const { + return reinterpret_cast(sparse_data_.data()); + } + + //! Request a change in capacity + void resize_for_sparse(size_t n) { + sparse_indices_.resize(n * sizeof(uint32_t)); + sparse_data_.resize(n * sizeof(ValueType)); + } + + int add_sparses(const NumericalVector &indexes, + const NumericalVector &values) { + sparse_count_ = indexes.size(); + + sparse_indices_ = (const std::string &)indexes; + sparse_data_ = (const std::string &)values; + + return 0; + } + + private: + //! Data Member + size_t sparse_count_; + std::string sparse_indices_; + std::string sparse_data_; +}; + +/*! Sparse Vector Adapter + */ +template +class SparseVector { + public: + typedef typename std::remove_cv::type ValueType; + typedef ValueType *iterator; + typedef const ValueType *const_iterator; + + //! Constructor + SparseVector(void) = default; + + size_t sparse_count() const { + return sparse_count_; + } + + const uint32_t *sparse_indices() const { + return reinterpret_cast(sparse_indices_.data()); + } + + const ValueType *sparse_data(void) const { + return reinterpret_cast(sparse_data_.data()); + } + + //! Request a change in capacity + void resize_for_sparse(size_t n) { + sparse_indices_.resize(n * sizeof(uint32_t)); + sparse_data_.resize(n * sizeof(ValueType)); + } + + int add_sparses(const NumericalVector &indexes, + const NumericalVector &values) { + sparse_count_ = indexes.size(); + + sparse_indices_ = (const std::string &)indexes; + sparse_data_ = (const std::string &)values; + + return 0; + } + + private: + //! Data Member + size_t sparse_count_; + std::string sparse_indices_; + std::string sparse_data_; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/container/vector_array.h b/src/ailego/container/vector_array.h new file mode 100644 index 0000000..7819a23 --- /dev/null +++ b/src/ailego/container/vector_array.h @@ -0,0 +1,555 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "vector.h" + +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.h b/src/ailego/encoding/json.h new file mode 100644 index 0000000..6f8d34f --- /dev/null +++ b/src/ailego/encoding/json.h @@ -0,0 +1,17 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "json/mod_json_plus.h" diff --git a/src/ailego/encoding/json/mod_json.c b/src/ailego/encoding/json/mod_json.c new file mode 100644 index 0000000..aff816f --- /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 "mod_json.h" +#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/encoding/json/mod_json.h b/src/ailego/encoding/json/mod_json.h new file mode 100644 index 0000000..d812c28 --- /dev/null +++ b/src/ailego/encoding/json/mod_json.h @@ -0,0 +1,1382 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 !defined(__cplusplus) && defined(_MSC_VER) +#if !defined(inline) +#define inline __inline +#endif +#endif + +#if defined(__cplusplus) +extern "C" { +#endif + +#define MOD_JSON_FALSE (false) +#define MOD_JSON_TRUE (true) +#define MOD_JSON_INFINITY (INFINITY) + +/*! JSON Type + */ +enum mod_json_type { + mod_json_type_null = 0, + mod_json_type_boolean = 1, + mod_json_type_integer = 2, + mod_json_type_float = 3, + mod_json_type_string = 4, + mod_json_type_array = 5, + mod_json_type_object = 6 +}; + +/*! JSON Token State + */ +enum mod_json_state { + mod_json_state_null = 0, + mod_json_state_start = 1, + mod_json_state_finish = 2, + mod_json_state_array_start = 3, + mod_json_state_array_half = 4, + mod_json_state_array_finish = 5, + mod_json_state_object_start = 6, + mod_json_state_object_half1 = 7, + mod_json_state_object_half2 = 8, + mod_json_state_object_finish = 9, + mod_json_state_max = 10 +}; + +/*! JSON Token Error Code + */ +enum mod_json_error { + mod_json_error_null = 0, + mod_json_error_invalid = 1, + mod_json_error_state = 2, + mod_json_error_empty = 3, + mod_json_error_break = 4, + mod_json_error_depth = 5, + mod_json_error_trunc = 6, + mod_json_error_start = 7, + mod_json_error_array = 8, + mod_json_error_object = 9, + mod_json_error_key = 10, + mod_json_error_value = 11, + mod_json_error_quote = 12 +}; + +/*! JSON Token Event + */ +enum mod_json_event { + mod_json_event_null = 0, + mod_json_event_field = 1, + mod_json_event_object = 2, + mod_json_event_array = 3, + mod_json_event_boolean = 4, + mod_json_event_integer = 5, + mod_json_event_float = 6, + mod_json_event_string = 7 +}; + +typedef unsigned int mod_json_size_t; +typedef int mod_json_ssize_t; +typedef bool mod_json_boolean_t; +typedef char mod_json_char_t; +typedef const char mod_json_cchar_t; +typedef unsigned char mod_json_uchar_t; +typedef long long mod_json_integer_t; +typedef double mod_json_float_t; +typedef void mod_json_void_t; +typedef enum mod_json_type mod_json_type_t; +typedef union mod_json_any mod_json_any_t; +typedef struct mod_json_value mod_json_value_t; +typedef struct mod_json_string mod_json_string_t; +typedef struct mod_json_array mod_json_array_t; +typedef struct mod_json_object mod_json_object_t; +typedef struct mod_json_pair mod_json_pair_t; +typedef struct mod_json_option mod_json_option_t; +typedef enum mod_json_state mod_json_state_t; +typedef enum mod_json_error mod_json_error_t; +typedef enum mod_json_event mod_json_event_t; +typedef struct mod_json_token mod_json_token_t; + +/*! Callback function when parsing JSON + */ +typedef int (*mod_json_event_proc)(mod_json_token_t *tok, mod_json_void_t *val, + mod_json_size_t len); + +/*! JSON Any + */ +union mod_json_any { + mod_json_object_t *c_obj; + mod_json_array_t *c_arr; + mod_json_string_t *c_str; + mod_json_float_t c_float; + mod_json_boolean_t c_bool; + mod_json_integer_t c_int; +}; + +/*! JSON Value + */ +struct mod_json_value { + mod_json_ssize_t refer; + mod_json_type_t type; + mod_json_any_t data; +}; + +/*! JSON String + */ +struct mod_json_string { + mod_json_ssize_t refer; + mod_json_size_t size; + mod_json_char_t *first; + mod_json_char_t *last; +}; + +/*! JSON Array + */ +struct mod_json_array { + mod_json_ssize_t refer; + mod_json_size_t size; + mod_json_value_t **first; + mod_json_value_t **last; +}; + +/*! JSON Pair + */ +struct mod_json_pair { + mod_json_string_t *key; + mod_json_value_t *val; +}; + +/*! JSON Object + */ +struct mod_json_object { + mod_json_ssize_t refer; + mod_json_size_t size; + mod_json_pair_t *first; + mod_json_pair_t *last; +}; + +#define MOD_JSON_COMMENT 0x0001 /* Enable comments */ +#define MOD_JSON_UNSTRICT 0x0002 /* Enable loose JSON string */ +#define MOD_JSON_SIMPLE 0x0004 /* Enable simple format */ +#define MOD_JSON_SQUOTE 0x0008 /* Enable single quotes support */ + +/*! JSON Option + */ +struct mod_json_option { + mod_json_size_t options; + mod_json_size_t object_depth; + mod_json_size_t array_depth; +}; + +/** + * \brief Create and set a JSON null value + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_set_null(void); + +/** + * \brief Create and set a JSON object value + * \param obj The value to be assigned + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_set_object(mod_json_object_t *obj); + +/** + * \brief Create and set a JSON array value + * \param arr The value to be assigned + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_set_array(mod_json_array_t *arr); + +/** + * \brief Create and set a JSON string value + * \param str The value to be assigned + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_set_string(mod_json_string_t *str); + +/** + * \brief Create and set a JSON string buffer + * \param buf The pointer of string buffer + * \param len The length of string buffer + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_set_buffer(mod_json_cchar_t *buf, + mod_json_size_t len); + +/** + * \brief Create and set a JSON integer value + * \param num The value to be assigned + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_set_integer(mod_json_integer_t num); + +/** + * \brief Create and set a JSON float value + * \param dbl The value to be assigned + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_set_float(mod_json_float_t dbl); + +/** + * \brief Create and set a JSON boolean value + * \param bol The value to be assigned + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_set_boolean(mod_json_boolean_t bol); + +/** + * \brief Assign a JSON value as null + * \param val The pointer of value + */ +void mod_json_value_assign_null(mod_json_value_t *val); + +/** + * \brief Assign a JSON value as a object + * \param val The pointer of value + * \param obj The value to be assigned + */ +void mod_json_value_assign_object(mod_json_value_t *val, + mod_json_object_t *obj); + +/** + * \brief Assign a JSON value as an array + * \param val The pointer of value + * \param arr The value to be assigned + */ +void mod_json_value_assign_array(mod_json_value_t *val, mod_json_array_t *arr); + +/** + * \brief Assign a JSON value as a string + * \param val The pointer of value + * \param str The value to be assigned + */ +void mod_json_value_assign_string(mod_json_value_t *val, + mod_json_string_t *str); + +/** + * \brief Assign a JSON value as an integer + * \param val The pointer of value + * \param num The value to be assigned + */ +void mod_json_value_assign_integer(mod_json_value_t *val, + mod_json_integer_t num); + +/** + * \brief Assign a JSON value as a float + * \param val The pointer of value + * \param dbl The value to be assigned + */ +void mod_json_value_assign_float(mod_json_value_t *val, mod_json_float_t dbl); + +/** + * \brief Assign a JSON value as a boolean + * \param val The pointer of value + * \param bol The value to be assigned + */ +void mod_json_value_assign_boolean(mod_json_value_t *val, + mod_json_boolean_t bol); + +/** + * \brief Assign a new JSON value + * \param dst The pointer of destination value (can't be null) + * \param src The pointer of source value (can be null) + */ +void mod_json_value_assign(mod_json_value_t *dst, mod_json_value_t *src); + +/** + * \brief Merge a JSON value into another one + * \param dst The pointer of destination value (can't be null) + * \param src The pointer of source value (can be null) + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_value_merge(mod_json_value_t *dst, mod_json_value_t *src); + +/** + * \brief Retrieve object of a JSON value + * \param val The pointer of value + * \return Null indicates unmatched type or empty. + */ +mod_json_object_t *mod_json_value_object(mod_json_value_t *val); + +/** + * \brief Retrieve array of a JSON value + * \param val The pointer of value + * \return Null indicates unmatched type or empty. + */ +mod_json_array_t *mod_json_value_array(mod_json_value_t *val); + +/** + * \brief Retrieve string of a JSON value + * \param val The pointer of value + * \return Null indicates unmatched type or empty. + */ +mod_json_string_t *mod_json_value_string(mod_json_value_t *val); + +/** + * \brief Retrieve c-string of a JSON value + * \param val The pointer of value + * \return Null indicates unmatched type or empty. + */ +mod_json_cchar_t *mod_json_value_cstring(mod_json_value_t *val); + +/** + * \brief Retrieve float of a JSON value + * \param val The pointer of value + * \return It will try converting the unmatched + value to float. If nothing be done, + returns zero by default. + */ +mod_json_float_t mod_json_value_float(mod_json_value_t *val); + +/** + * \brief Retrieve boolean of a JSON value + * \param val The pointer of value + * \return If string, object or array is not empty, + number(integer or float) does not equal + to zero, it returns true. + */ +mod_json_boolean_t mod_json_value_boolean(mod_json_value_t *val); + +/** + * \brief Retrieve integer of a JSON value + * \param val The pointer of value + * \return It will try converting the unmatched + value to integer. If nothing be done, + returns zero by default. + */ +mod_json_integer_t mod_json_value_integer(mod_json_value_t *val); + +/** + * \brief Clone a JSON value + * \param val The pointer of value + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_value_clone(mod_json_value_t *val); + +/** + * \brief Retrieve non-zero if they are equal + * \param lhs The pointer of left value + * \param rhs The pointer of right value + * \return 1 indicates true, 0 indicates false. + */ +mod_json_boolean_t mod_json_value_is_equal(mod_json_value_t *lhs, + mod_json_value_t *rhs); + +/** + * \brief Unset or destroy a JSON value + * \param val The pointer of value + */ +void mod_json_value_unset(mod_json_value_t *val); + +/** + * \brief Increase reference count of a JSON value + * \param val The pointer of value + * \return The original pointer of value + */ +static inline mod_json_value_t *mod_json_value_get(mod_json_value_t *val) { + ++val->refer; + return val; +} + +/** + * \brief Decrease reference count of a JSON value + * \param val The pointer of value + * \return The new number of refer-counter + */ +static inline mod_json_ssize_t mod_json_value_put(mod_json_value_t *val) { + return (--val->refer); +} + +/** + * \brief Retrieve refer-counter of a JSON value + * \param val The pointer of value + * \return The number of refer-counter + */ +static inline mod_json_ssize_t mod_json_value_refer(mod_json_value_t *val) { + return (val ? val->refer : -1); +} + +/** + * \brief Set the refer-counter as leaked + * \param val The pointer of value + */ +static inline void mod_json_value_set_leaked(mod_json_value_t *val) { + val->refer = 0; +} + +/** + * \brief Retrieve non-zero if refer-counter is leaked + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_leaked( + mod_json_value_t *val) { + return (val->refer <= 0); +} + +/** + * \brief Retrieve non-zero if refer-counter is shared + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_shared( + mod_json_value_t *val) { + return (val->refer > 1); +} + +/** + * \brief Grab (get or clone) a JSON value + * \param val The pointer of value + * \return Null indicates failure + */ +static inline mod_json_value_t *mod_json_value_grab(mod_json_value_t *val) { + /* Is it leaked? */ + if (!mod_json_value_is_leaked(val)) { + return mod_json_value_get(val); + } + return mod_json_value_clone(val); +} + +/** + * \brief Retrieve type of a JSON value + * \param val The pointer of value + * \return The code of type + */ +static inline mod_json_type_t mod_json_value_type(mod_json_value_t *val) { + return (val->type); +} + +/** + * \brief Retrieve non-zero if a JSON value is null + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_null(mod_json_value_t *val) { + return (val ? val->type == mod_json_type_null : MOD_JSON_TRUE); +} + +/** + * \brief Retrieve non-zero if it is a JSON array + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_array( + mod_json_value_t *val) { + return (val ? val->type == mod_json_type_array : MOD_JSON_FALSE); +} + +/** + * \brief Retrieve non-zero if it is a JSON object + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_object( + mod_json_value_t *val) { + return (val ? val->type == mod_json_type_object : MOD_JSON_FALSE); +} + +/** + * \brief Retrieve non-zero if it is a JSON string + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_string( + mod_json_value_t *val) { + return (val ? val->type == mod_json_type_string : MOD_JSON_FALSE); +} + +/** + * \brief Retrieve non-zero if it is a JSON float + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_float( + mod_json_value_t *val) { + return (val ? val->type == mod_json_type_float : MOD_JSON_FALSE); +} + +/** + * \brief Retrieve non-zero if it is a JSON boolean + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_boolean( + mod_json_value_t *val) { + return (val ? val->type == mod_json_type_boolean : MOD_JSON_FALSE); +} + +/** + * \brief Retrieve non-zero if it is a JSON integer + * \param val The pointer of value + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_value_is_integer( + mod_json_value_t *val) { + return (val ? val->type == mod_json_type_integer : MOD_JSON_FALSE); +} + +/** + * \brief Request a change in capacity + * \param str The pointer of string + * \param n The requested size of capacity + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_string_reserve(mod_json_string_t *str, mod_json_size_t n); + +/** + * \brief Create and set a JSON string + * \param cstr The pointer of c-string + * \param len The length of c-string + * \return Null indicates failure. + */ +mod_json_string_t *mod_json_string_set(mod_json_cchar_t *cstr, + mod_json_size_t len); + +/** + * \brief Assign new content to a JSON string + * \param str The pointer of string + * \param cstr The pointer of c-string + * \param len The length of c-string + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_string_assign(mod_json_string_t *str, mod_json_cchar_t *cstr, + mod_json_size_t len); + +/** + * \brief Clone a JSON string + * \param str The pointer of string + * \return Null indicates failure. + */ +static inline mod_json_string_t *mod_json_string_clone(mod_json_string_t *str) { + return (str ? mod_json_string_set(str->first, + (mod_json_size_t)(str->last - str->first)) + : (mod_json_string_t *)0); +} + +/** + * \brief Unset or destroy a JSON string + * \param str The pointer of string + */ +void mod_json_string_unset(mod_json_string_t *str); + +/** + * \brief Reset a JSON string + * \param str The pointer of string + */ +void mod_json_string_reset(mod_json_string_t *str); + +/** + * \brief Append a c-string to a JSON string + * \param str The pointer of string + * \param cstr The pointer of c-string + * \param len The length of c-string + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_string_append(mod_json_string_t *str, mod_json_cchar_t *cstr, + mod_json_size_t len); + +/** + * \brief Add a copy of a JSON string + * \param str The main string + * \param val The appended string + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_string_add(mod_json_string_t *str, mod_json_string_t *val); + +/** + * \brief Retrieve HASH of a JSON string + * \param str The pointer of string + * \return The value of HASH + */ +mod_json_size_t mod_json_string_hash(mod_json_string_t *str); + +/** + * \brief Compare two JSON strings (case sensitive) + * \param str1 The first string + * \param str2 The second string + * \return 0 indicates equal. + */ +int mod_json_string_compare(mod_json_string_t *str1, mod_json_string_t *str2); + +/** + * \brief Convert a JSON string to an integer + * \param str The pointer of string + * \return If nothing be done, returns zero by default. + */ +mod_json_integer_t mod_json_string_integer(mod_json_string_t *str); + +/** + * \brief Convert a JSON string to a float + * \param str The pointer of string + * \return If nothing be done, returns zero by default. + */ +mod_json_float_t mod_json_string_float(mod_json_string_t *str); + +/** + * \brief Encode a JSON string + * \param src The pointer of source string + * \return Null indicates failure. + */ +mod_json_string_t *mod_json_string_encode(mod_json_string_t *src); + +/** + * \brief Decode a JSON string + * \param src The pointer of source string + * \return Null indicates failure. + */ +mod_json_string_t *mod_json_string_decode(mod_json_string_t *src); + +/** + * \brief Increase reference count of a JSON string + * \param str The pointer of string + * \return The original pointer of string + */ +static inline mod_json_string_t *mod_json_string_get(mod_json_string_t *str) { + ++str->refer; + return str; +} + +/** + * \brief Decrease reference count of a JSON string + * \param str The pointer of string + * \return The new number of refer-counter + */ +static inline mod_json_ssize_t mod_json_string_put(mod_json_string_t *str) { + return (--str->refer); +} + +/** + * \brief Retrieve refer-counter of a JSON string + * \param str The pointer of string + * \return The number of refer-counter + */ +static inline mod_json_ssize_t mod_json_string_refer(mod_json_string_t *str) { + return (str ? str->refer : -1); +} + +/** + * \brief Set the refer-counter as leaked + * \param str The pointer of string + */ +static inline void mod_json_string_set_leaked(mod_json_string_t *str) { + str->refer = 0; +} + +/** + * \brief Retrieve non-zero if refer-counter is leaked + * \param str The pointer of string + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_string_is_leaked( + mod_json_string_t *str) { + return (str->refer <= 0); +} + +/** + * \brief Retrieve non-zero if refer-counter is shared + * \param str The pointer of string + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_string_is_shared( + mod_json_string_t *str) { + return (str->refer > 1); +} + +/** + * \brief Grab (get or clone) a JSON string + * \param str The pointer of string + * \return Null indicates failure + */ +static inline mod_json_string_t *mod_json_string_grab(mod_json_string_t *str) { + /* Is it leaked? */ + if (!mod_json_string_is_leaked(str)) { + return mod_json_string_get(str); + } + return mod_json_string_clone(str); +} + +/** + * \brief Retrieve c-string of a JSON string + * \param str The pointer of string + * \return The pointer of c-string + */ +static inline mod_json_cchar_t *mod_json_string_cstr(mod_json_string_t *str) { + return (str ? str->first : (mod_json_cchar_t *)0); +} + +/** + * \brief Retrieve data pointer of a JSON string + * \param str The pointer of string + * \return The pointer of data + */ +static inline mod_json_char_t *mod_json_string_data(mod_json_string_t *str) { + return (str ? str->first : (mod_json_char_t *)0); +} + +/** + * \brief Retrieve capacity of a JSON string + * \param str The pointer of string + * \return The size of allocated storage + */ +static inline mod_json_size_t mod_json_string_capacity(mod_json_string_t *str) { + return (str ? (str->size - 1) : 0); +} + +/** + * \brief Retrieve length of a JSON string + * \param str The pointer of string + * \return The length of string + */ +static inline mod_json_size_t mod_json_string_length(mod_json_string_t *str) { + return (str ? (mod_json_size_t)(str->last - str->first) : 0); +} + +/** + * \brief Retrieve non-zero if a JSON string is empty + * \param str The pointer of string + * \return 0 indicates non-empty + */ +static inline mod_json_boolean_t mod_json_string_empty(mod_json_string_t *str) { + return (mod_json_string_length(str) == 0); +} + +/** + * \brief Create and set a JSON array + * \param size The initialized size of array + * \return Null indicates failure. + */ +mod_json_array_t *mod_json_array_set(mod_json_size_t size); + +/** + * \brief Clone a JSON array + * \param arr The pointer of array + * \return Null indicates failure. + */ +mod_json_array_t *mod_json_array_clone(mod_json_array_t *arr); + +/** + * \brief Retrieve non-zero if they are equal + * \param lhs The pointer of left array + * \param rhs The pointer of right array + * \return 1 indicates true, 0 indicates false. + */ +mod_json_boolean_t mod_json_array_is_equal(mod_json_array_t *lhs, + mod_json_array_t *rhs); + +/** + * \brief Unset or destroy a JSON array + * \param arr The pointer of array + */ +void mod_json_array_unset(mod_json_array_t *arr); + +/** + * \brief Reset a JSON array + * \param arr The pointer of array + */ +void mod_json_array_reset(mod_json_array_t *arr); + +/** + * \brief Create and set a JSON array (default parameters) + * \param size The initialized size of array + * \return Null indicates failure. + */ +static inline mod_json_array_t *mod_json_array_set_default(void) { + return mod_json_array_set(0); +} + +/** + * \brief Increase reference count of a JSON array + * \param arr The pointer of array + * \return The original pointer of array + */ +static inline mod_json_array_t *mod_json_array_get(mod_json_array_t *arr) { + ++arr->refer; + return arr; +} + +/** + * \brief Decrease reference count of a JSON array + * \param str The pointer of array + * \return The new number of refer-counter + */ +static inline mod_json_ssize_t mod_json_array_put(mod_json_array_t *arr) { + return (--arr->refer); +} + +/** + * \brief Retrieve refer-counter of a JSON array + * \param arr The pointer of array + * \return The number of refer-counter + */ +static inline mod_json_ssize_t mod_json_array_refer(mod_json_array_t *arr) { + return (arr ? arr->refer : -1); +} + +/** + * \brief Set the refer-counter as leaked + * \param arr The pointer of array + */ +static inline void mod_json_array_set_leaked(mod_json_array_t *arr) { + arr->refer = 0; +} + +/** + * \brief Retrieve non-zero if refer-counter is leaked + * \param arr The pointer of array + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_array_is_leaked( + mod_json_array_t *arr) { + return (arr->refer <= 0); +} + +/** + * \brief Retrieve non-zero if refer-counter is shared + * \param arr The pointer of array + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_array_is_shared( + mod_json_array_t *arr) { + return (arr->refer > 1); +} + +/** + * \brief Grab (get or clone) a JSON array + * \param arr The pointer of array + * \return Null indicates failure + */ +static inline mod_json_array_t *mod_json_array_grab(mod_json_array_t *arr) { + /* Is it leaked? */ + if (!mod_json_array_is_leaked(arr)) { + return mod_json_array_get(arr); + } + return mod_json_array_clone(arr); +} + +/** + * \brief Retrieve count of elements in a JSON array + * \param arr The pointer of array + * \return The count of elements + */ +static inline mod_json_size_t mod_json_array_count(mod_json_array_t *arr) { + return (arr ? (mod_json_size_t)(arr->last - arr->first) : 0); +} + +/** + * \brief Retrieve capacity of a JSON array + * \param arr The pointer of array + * \return The size of allocated storage + */ +static inline mod_json_size_t mod_json_array_capacity(mod_json_array_t *arr) { + return (arr ? arr->size : 0); +} + +/** + * \brief Retrieve non-zero if a JSON array is empty + * \param arr The pointer of array + * \return 0 indicates non-empty + */ +static inline mod_json_boolean_t mod_json_array_empty(mod_json_array_t *arr) { + return (mod_json_array_count(arr) == 0); +} + +/** + * \brief Retrieve the begin of a JSON array + * \param arr The pointer of array + * \return The pointer of begin + */ +static inline mod_json_value_t **mod_json_array_begin(mod_json_array_t *arr) { + return (arr->first); +} + +/** + * \brief Retrieve the reverse begin of a JSON array + * \param arr The pointer of array + * \return The pointer of reverse begin + */ +static inline mod_json_value_t **mod_json_array_rbegin(mod_json_array_t *arr) { + return (arr->last - 1); +} + +/** + * \brief Retrieve the end of a JSON array + * \param arr The pointer of array + * \return The pointer of end + */ +static inline mod_json_value_t **mod_json_array_end(mod_json_array_t *arr) { + return (arr->last); +} + +/** + * \brief Retrieve the reverse end of a JSON array + * \param arr The pointer of array + * \return The pointer of reverse end + */ +static inline mod_json_value_t **mod_json_array_rend(mod_json_array_t *arr) { + return (arr->first - 1); +} + +/** + * \brief Request a change in capacity + * \param arr The pointer of array + * \param n The requested size of capacity + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_array_reserve(mod_json_array_t *arr, mod_json_size_t n); + +/** + * \brief Reverse the order of the elements in an array + * \param arr The pointer of array + * \return 0 indicates success, -1 indicates failure. + */ +void mod_json_array_reverse(mod_json_array_t *arr); + +/** + * \brief Push a value into a JSON array + * \param arr The pointer of array + * \param val The pointer of value + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_array_push(mod_json_array_t *arr, mod_json_value_t *val); + +/** + * \brief Pop the last element from a JSON array + * \param arr The pointer of array + */ +void mod_json_array_pop(mod_json_array_t *arr); + +/** + * \brief Remove the first element of a JSON array + * \param arr The pointer of array + */ +void mod_json_array_shift(mod_json_array_t *arr); + +/** + * \brief Retrieve a value in JSON array + * \param arr The pointer of array + * \param id The index (start from zero) + * \return Null indicates no one be found. + */ +mod_json_value_t *mod_json_array_at(mod_json_array_t *arr, mod_json_size_t id); + +/** + * \brief Merge a JSON array into another one + * \param dst The pointer of destination array (can't be null) + * \param src The pointer of source array (can't be null) + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_array_merge(mod_json_array_t *dst, mod_json_array_t *src); + +/** + * \brief Resize a JSON array so that it contains n elements + * \param arr The pointer of array + * \param n The new size, expressed in number of elements + * \param val The pointer of value assigned (can be null) + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_array_resize(mod_json_array_t *arr, mod_json_size_t n, + mod_json_value_t *val); + +/** + * \brief Retrieve key of a JSON pair + * \param pair The pointer of pair + * \return The key of pair + */ +static inline mod_json_string_t *mod_json_pair_key(mod_json_pair_t *pair) { + return (pair->key); +} + +/** + * \brief Retrieve value of a JSON pair + * \param pair The pointer of pair + * \return The value of pair + */ +static inline mod_json_value_t *mod_json_pair_value(mod_json_pair_t *pair) { + return (pair->val); +} + +/** + * \brief Create and set a JSON object + * \param size The initialized size of object + * \return Null indicates failure. + */ +mod_json_object_t *mod_json_object_set(mod_json_size_t size); + +/** + * \brief Clone a JSON object + * \param obj The pointer of object + * \return Null indicates failure. + */ +mod_json_object_t *mod_json_object_clone(mod_json_object_t *obj); + +/** + * \brief Retrieve non-zero if they are equal + * \param lhs The pointer of left object + * \param rhs The pointer of right object + * \return 1 indicates true, 0 indicates false. + */ +mod_json_boolean_t mod_json_object_is_equal(mod_json_object_t *lhs, + mod_json_object_t *rhs); + +/** + * \brief Unset or destroy a JSON object + * \param obj The pointer of object + */ +void mod_json_object_unset(mod_json_object_t *obj); + +/** + * \brief Reset a JSON object + * \param obj The pointer of object + */ +void mod_json_object_reset(mod_json_object_t *obj); + +/** + * \brief Create and set a JSON object (default parameters) + * \return Null indicates failure. + */ +static inline mod_json_object_t *mod_json_object_set_default(void) { + return mod_json_object_set(0); +} + +/** + * \brief Increase reference count of a JSON object + * \param obj The pointer of object + * \return The original pointer of object + */ +static inline mod_json_object_t *mod_json_object_get(mod_json_object_t *obj) { + ++obj->refer; + return obj; +} + +/** + * \brief Decrease reference count of a JSON object + * \param str The pointer of object + * \return The new number of refer-counter + */ +static inline mod_json_ssize_t mod_json_object_put(mod_json_object_t *obj) { + return (--obj->refer); +} + +/** + * \brief Retrieve refer-counter of a JSON object + * \param obj The pointer of object + * \return The number of refer-counter + */ +static inline mod_json_ssize_t mod_json_object_refer(mod_json_object_t *obj) { + return (obj ? obj->refer : -1); +} + +/** + * \brief Set the refer-counter as leaked + * \param obj The pointer of object + */ +static inline void mod_json_object_set_leaked(mod_json_object_t *obj) { + obj->refer = 0; +} + +/** + * \brief Retrieve non-zero if refer-counter is leaked + * \param obj The pointer of object + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_object_is_leaked( + mod_json_object_t *obj) { + return (obj->refer <= 0); +} + +/** + * \brief Retrieve non-zero if refer-counter is shared + * \param obj The pointer of object + * \return 1 indicates TRUE, 0 indicates FALSE + */ +static inline mod_json_boolean_t mod_json_object_is_shared( + mod_json_object_t *obj) { + return (obj->refer > 1); +} + +/** + * \brief Grab (get or clone) a JSON object + * \param obj The pointer of object + * \return Null indicates failure + */ +static inline mod_json_object_t *mod_json_object_grab(mod_json_object_t *obj) { + /* Is it leaked? */ + if (!mod_json_object_is_leaked(obj)) { + return mod_json_object_get(obj); + } + return mod_json_object_clone(obj); +} + +/** + * \brief Retrieve count of elements in a JSON object + * \param obj The pointer of object + * \return The count of elements + */ +static inline mod_json_size_t mod_json_object_count(mod_json_object_t *obj) { + return (obj ? (mod_json_size_t)(obj->last - obj->first) : 0); +} + +/** + * \brief Retrieve non-zero if a JSON object is empty + * \param obj The pointer of object + * \return 0 indicates non-empty + */ +static inline mod_json_boolean_t mod_json_object_empty(mod_json_object_t *obj) { + return (mod_json_object_count(obj) == 0); +} + +/** + * \brief Retrieve the begin of a JSON object + * \param obj The pointer of object + * \return The pointer of begin + */ +static inline mod_json_pair_t *mod_json_object_begin(mod_json_object_t *obj) { + return (obj->first); +} + +/** + * \brief Retrieve the reverse begin of a JSON object + * \param obj The pointer of object + * \return The pointer of reverse begin + */ +static inline mod_json_pair_t *mod_json_object_rbegin(mod_json_object_t *obj) { + return (obj->last - 1); +} + +/** + * \brief Retrieve the end of a JSON object + * \param obj The pointer of object + * \return The pointer of end + */ +static inline mod_json_pair_t *mod_json_object_end(mod_json_object_t *obj) { + return (obj->last); +} + +/** + * \brief Retrieve the reverse end of a JSON object + * \param obj The pointer of object + * \return The pointer of reverse end + */ +static inline mod_json_pair_t *mod_json_object_rend(mod_json_object_t *obj) { + return (obj->first - 1); +} + +/** + * \brief Insert a pair into a JSON object + * \param obj The pointer of object + * \param key The string of key + * \param val The pointer of value + * \return The pair inserted, Null indicates failure. + */ +mod_json_pair_t *mod_json_object_insert(mod_json_object_t *obj, + mod_json_string_t *key, + mod_json_value_t *val); + +/** + * \brief Assign a pair into a JSON object + * \param obj The pointer of object + * \param key The string of key + * \param val The pointer of value + * \return The pair assigned, Null indicates failure. + */ +mod_json_pair_t *mod_json_object_assign(mod_json_object_t *obj, + mod_json_string_t *key, + mod_json_value_t *val); + +/** + * \brief Touch a pair in a JSON object + * \param obj The pointer of object + * \param key The c-string of key + */ +mod_json_pair_t *mod_json_object_touch(mod_json_object_t *obj, + mod_json_cchar_t *key); + +/** + * \brief Erase a pair from a JSON object + * \param obj The pointer of object + * \param key The c-string of key + */ +void mod_json_object_erase(mod_json_object_t *obj, mod_json_cchar_t *key); + +/** + * \brief Get a value in a JSON object + * \param obj The pointer of object + * \param key The c-string of key + * \return Null indicates failure. + */ +mod_json_value_t *mod_json_object_at(mod_json_object_t *obj, + mod_json_cchar_t *key); + +/** + * \brief Find a pair in a JSON object + * \param obj The pointer of object + * \param key The c-string of key + * \return Null indicates failure. + */ +mod_json_pair_t *mod_json_object_find(mod_json_object_t *obj, + mod_json_cchar_t *key); + +/** + * \brief Merge a JSON object into another one + * \param dst The pointer of destination object (can't be null) + * \param src The pointer of source object (can't be null) + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_object_merge(mod_json_object_t *dst, mod_json_object_t *src); + +/** + * \brief Create a JSON token + * \param opt The options of parser + * \return The pointer of token, Null indicates failure. + */ +mod_json_token_t *mod_json_token_create(mod_json_option_t *opt); + +/** + * \brief Destroy a JSON token + * \param tok The pointer of token + */ +void mod_json_token_destroy(mod_json_token_t *tok); + +/** + * \brief Parse a c-string with a JSON token + * \param tok The pointer of token + * \param cstr The pointer of c-string + * \return 0 indicates success, -1 indicates failure. + */ +int mod_json_token_parse(mod_json_token_t *tok, mod_json_cchar_t *cstr); + +/** + * \brief Retrieve error of a JSON token + * \param tok The pointer of token + * \return The code of error + */ +mod_json_error_t mod_json_token_error(mod_json_token_t *tok); + +/** + * \brief Retrieve error context of a JSON token + * \param tok The pointer of token + * \return The pointer of context, null indicates non-errors + */ +mod_json_cchar_t *mod_json_token_context(mod_json_token_t *tok); + +/** + * \brief Retrieve state of a JSON token + * \param tok The pointer of token + * \return The value of state + */ +mod_json_state_t mod_json_token_state(mod_json_token_t *tok); + +/** + * \brief Retrieve object depth of a JSON token + * \param tok The pointer of token + * \return The value of object depth + */ +mod_json_size_t mod_json_token_object_depth(mod_json_token_t *tok); + +/** + * \brief Retrieve array depth of a JSON token + * \param tok The pointer of token + * \return The value of array depth + */ +mod_json_size_t mod_json_token_array_depth(mod_json_token_t *tok); + +/** + * \brief Retrieve max object depth of a JSON token + * \param tok The pointer of token + * \return The value of max object depth + */ +mod_json_size_t mod_json_token_max_object_depth(mod_json_token_t *tok); + +/** + * \brief Retrieve max array depth of a JSON token + * \param tok The pointer of token + * \return The value of max array depth + */ +mod_json_size_t mod_json_token_max_array_depth(mod_json_token_t *tok); + +/** + * \brief Retrieve depth of a JSON token + * \param tok The pointer of token + * \return The value of depth + */ +mod_json_size_t mod_json_token_depth(mod_json_token_t *tok); + +/** + * \brief Retrieve max depth of a JSON token + * \param tok The pointer of token + * \return The value of max depth + */ +mod_json_size_t mod_json_token_max_depth(mod_json_token_t *tok); + +/** + * \brief Retrieve parameter of a JSON token + * \param tok The pointer of token + * \return The value of parameter + */ +mod_json_void_t *mod_json_token_param(mod_json_token_t *tok); + +/** + * \brief Set parameter of a JSON token + * \param tok The pointer of token + * \param param The value of parameter + */ +void mod_json_token_set_param(mod_json_token_t *tok, mod_json_void_t *param); + +/** + * \brief Register callback function of a JSON token + * \param tok The pointer of token + * \param proc The pointer of callback function + */ +void mod_json_token_set_event(mod_json_token_t *tok, mod_json_event_proc proc); + +/** + * \brief Retrieve event code of a JSON token + * \param tok The pointer of token + * \return The code of event + */ +mod_json_event_t mod_json_token_event(mod_json_token_t *tok); + +/** + * \brief Parse a c-string with a JSON token + * \param tok The pointer of token + * \param cstr The pointer of c-string + * \return The pointer of value, Null indicates failure. + */ +mod_json_value_t *mod_json_parse(mod_json_token_t *tok, mod_json_cchar_t *cstr); + +/** + * \brief Parse a c-string simply + * \param cstr The pointer of c-string + * \param opts The options of parser + * \return The pointer of value, Null indicates failure. + */ +mod_json_value_t *mod_json_parse_simply(mod_json_cchar_t *cstr, + mod_json_size_t opts); + +/** + * \brief Dump a JSON value in string + * \param val The pointer of value + * \return Null indicates failure. + */ +mod_json_string_t *mod_json_dump(mod_json_value_t *val); + +#if defined(__cplusplus) +} /* extern "C" */ +#endif diff --git a/src/ailego/encoding/json/mod_json_plus.h b/src/ailego/encoding/json/mod_json_plus.h new file mode 100644 index 0000000..e705a71 --- /dev/null +++ b/src/ailego/encoding/json/mod_json_plus.h @@ -0,0 +1,3446 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "mod_json.h" + +namespace zvec { +namespace ailego { + +/*! JSON String + */ +class JsonString { + public: + typedef mod_json_size_t size_type; + typedef mod_json_ssize_t ssize_type; + typedef mod_json_float_t float_type; + typedef mod_json_integer_t integer_type; + + //! Constructor + JsonString(void) : str_(0) {} + + //! Constructor + JsonString(const JsonString &rhs) : str_(0) { + if (rhs.str_) { + str_ = mod_json_string_grab(rhs.str_); + } + } + +#if __cplusplus >= 201103L + //! Constructor + JsonString(JsonString &&rhs) : str_(rhs.str_) { + rhs.str_ = 0; + } +#endif + + //! Constructor + JsonString(const char *cstr) { + str_ = cstr ? mod_json_string_set(cstr, (mod_json_size_t)std::strlen(cstr)) + : 0; + } + + //! Constructor + JsonString(const char *cstr, size_type len) { + str_ = mod_json_string_set(cstr, len); + } + + //! Constructor + JsonString(const std::string &str) { + str_ = mod_json_string_set(str.c_str(), (mod_json_size_t)str.size()); + } + + //! Destructor + ~JsonString(void) { + mod_json_string_unset(str_); + } + + //! Assign new contents to the string, replacing its current content + JsonString &operator=(const JsonString &rhs) { + this->assign(rhs); + return *this; + } + +#if __cplusplus >= 201103L + //! Assign new contents to the string, replacing its current content + JsonString &operator=(JsonString &&rhs) { + this->assign(std::move(rhs)); + return *this; + } +#endif + + //! Assign new contents to the string, replacing its current content + JsonString &operator=(const char *cstr) { + this->assign(cstr); + return *this; + } + + //! Assign new contents to the string, replacing its current content + JsonString &operator=(const std::string &rhs) { + this->assign(rhs); + return *this; + } + + //! Append a JSON string + JsonString &operator+=(const JsonString &str) { + this->append(str); + return *this; + } + + //! Append a c-style string + JsonString &operator+=(const char *cstr) { + this->append(cstr); + return *this; + } + + //! Append a character to string + JsonString &operator+=(char c) { + this->append(c); + return *this; + } + + //! Equality + bool operator==(const JsonString &rhs) const { + return (mod_json_string_compare(str_, rhs.str_) == 0); + } + + //! No equality + bool operator!=(const JsonString &rhs) const { + return !(*this == rhs); + } + + //! Retrieve the character at index n + char &operator[](size_type n) { + if (!copy_and_leak()) { + throw std::runtime_error("JsonString::operator[]"); + } + return *(str_->first + n); + } + + //! Retrieve the character at index n + const char &operator[](size_type n) const { + return *(str_->first + n); + } + + //! Retrieve non-zero if the string is valid + bool is_valid(void) const { + return (str_ != (mod_json_string_t *)0); + } + + //! Retrieve non-zero if the string is empty + bool empty(void) const { + return mod_json_string_empty(str_); + } + + //! Assign a JSON string + void assign(const JsonString &rhs) { + mod_json_string_unset(str_); + str_ = rhs.str_ ? mod_json_string_grab(rhs.str_) : 0; + } + +#if __cplusplus >= 201103L + //! Assign a JSON string + void assign(JsonString &&rhs) { + mod_json_string_unset(str_); + str_ = rhs.str_; + rhs.str_ = 0; + } +#endif + + //! Assign a c-style string + void assign(const char *cstr) { + if (cstr) { + if (!copy_on_write() || + mod_json_string_assign(str_, cstr, + (mod_json_size_t)std::strlen(cstr)) != 0) { + throw std::runtime_error("JsonString::assign"); + } + } + } + + //! Assign a c-style string + void assign(const char *cstr, size_type len) { + if (!copy_on_write() || mod_json_string_assign(str_, cstr, len) != 0) { + throw std::runtime_error("JsonString::assign"); + } + } + + //! Assign a STL-style string + void assign(const std::string &str) { + if (!copy_on_write() || + mod_json_string_assign(str_, str.c_str(), + (mod_json_size_t)str.size()) != 0) { + throw std::runtime_error("JsonString::assign"); + } + } + + //! Append a JSON string + void append(const JsonString &str) { + if (str.str_) { + if (!copy_on_write() || mod_json_string_add(str_, str.str_) != 0) { + throw std::runtime_error("JsonString::append"); + } + } + } + + //! Append a c-style string + void append(const char *cstr) { + if (cstr) { + if (!copy_on_write() || + mod_json_string_append(str_, cstr, + (mod_json_size_t)std::strlen(cstr)) != 0) { + throw std::runtime_error("JsonString::append"); + } + } + } + + //! Append a c-style string + void append(const char *cstr, size_type len) { + if (!copy_on_write() || mod_json_string_append(str_, cstr, len) != 0) { + throw std::runtime_error("JsonString::append"); + } + } + + //! Append a STL-style string + void append(const std::string &str) { + if (!copy_on_write() || + mod_json_string_append(str_, str.c_str(), + (mod_json_size_t)str.size()) != 0) { + throw std::runtime_error("JsonString::append"); + } + } + + //! Append a character to string + void append(char c) { + if (!copy_on_write() || mod_json_string_append(str_, &c, 1) != 0) { + throw std::runtime_error("JsonString::append"); + } + } + + //! Retrieve the character at index n + char &at(size_type n) { + if (this->size() <= n) { + throw std::out_of_range("JsonString::at"); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonString::at"); + } + return *(str_->first + n); + } + + //! Retrieve the character at index n + const char &at(size_type n) const { + if (this->size() <= n) { + throw std::out_of_range("JsonString::at"); + } + return *(str_->first + n); + } + + //! Request a change in capacity + void reserve(size_type n) { + if (!copy_on_write() || mod_json_string_reserve(str_, n) != 0) { + throw std::runtime_error("JsonString::reserve"); + } + } + + //! Clear the JSON string + void clear(void) { + mod_json_string_unset(str_); + str_ = 0; + } + + //! Exchange the content with another JSON string + void swap(JsonString &rhs) { + mod_json_string_t *str = str_; + str_ = rhs.str_; + rhs.str_ = str; + } + + //! Retrieve the data pointer + char *data(void) { + return mod_json_string_data(str_); + } + + //! Retrieve the data pointer + const char *data(void) const { + return mod_json_string_data(str_); + } + + //! Retrieve HASH of a JSON string + size_type hash(void) const { + return mod_json_string_hash(str_); + } + + //! Compare two JSON strings (case sensitive) + int compare(const JsonString &rhs) const { + return mod_json_string_compare(str_, rhs.str_); + } + + //! Compare two strings (case sensitive) + int compare(const char *cstr) const { + const char *self = this->c_str(); + if (self && cstr) { + return std::strcmp(self, cstr); + } + + // particular case + if (!self && cstr) { + return -1; + } else if (self && !cstr) { + return 1; + } + return 0; + } + + // Encode a JSON string + JsonString encode(void) const { + JsonString ret; + ret.str_ = mod_json_string_encode(str_); + return ret; + } + + // Decode a JSON string + JsonString decode(void) const { + JsonString ret; + ret.str_ = mod_json_string_decode(str_); + return ret; + } + + //! Retrieve the capacity of string + size_type capacity(void) const { + return mod_json_string_capacity(str_); + } + + //! Retrieve the length of string + size_type size(void) const { + return mod_json_string_length(str_); + } + + //! Retrieve the length of string + size_type length(void) const { + return mod_json_string_length(str_); + } + + //! Retrieve refer-counter of string + ssize_type refer(void) const { + return mod_json_string_refer(str_); + } + + //! Retrieve the c-style string + const char *c_str(void) const { + return mod_json_string_cstr(str_); + } + + //! Convert string to float + float_type as_float(void) const { + return mod_json_string_float(str_); + } + + //! Convert string to integer + integer_type as_integer(void) const { + return mod_json_string_integer(str_); + } + + //! Retrieve string as a STL string + std::string as_stl_string(void) const { + if (!this->empty()) { + return std::string(this->data(), this->size()); + } + return std::string(); + } + + protected: + //! Clone the string for writing + bool copy_on_write(void) { + if (str_) { + if (mod_json_string_is_shared(str_)) { + mod_json_string_put(str_); + str_ = mod_json_string_clone(str_); + } + } else { + str_ = mod_json_string_set("", 0); + } + return (str_ != 0); + } + + //! Clone the value and leak it + bool copy_and_leak(void) { + if (copy_on_write()) { + mod_json_string_set_leaked(str_); + return true; + } + return false; + } + + private: + mod_json_string_t *str_; +}; + +class JsonArray; +class JsonObject; + +/*! JSON Value + */ +class JsonValue { + public: + typedef mod_json_size_t size_type; + typedef mod_json_ssize_t ssize_type; + typedef mod_json_float_t float_type; + typedef mod_json_integer_t integer_type; + + //! Constructor + JsonValue(void) : val_(0) {} + + //! Constructor + explicit JsonValue(const bool &val) { + val_ = mod_json_value_set_boolean((mod_json_boolean_t)val); + } + + //! Constructor + explicit JsonValue(const signed char &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const char &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const short int &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const int &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const long int &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const long long int &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const float &val) { + val_ = mod_json_value_set_float((mod_json_float_t)val); + } + + //! Constructor + explicit JsonValue(const double &val) { + val_ = mod_json_value_set_float((mod_json_float_t)val); + } + + //! Constructor + explicit JsonValue(const long double &val) { + val_ = mod_json_value_set_float((mod_json_float_t)val); + } + + //! Constructor + explicit JsonValue(const unsigned char &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const unsigned short int &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const unsigned int &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const unsigned long int &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + explicit JsonValue(const unsigned long long int &val) { + val_ = mod_json_value_set_integer((mod_json_integer_t)val); + } + + //! Constructor + JsonValue(const JsonString &val) { + val_ = mod_json_value_set_string(*(mod_json_string_t **)&val); + } + + //! Constructor + JsonValue(const char *val) { + val_ = mod_json_value_set_buffer( + val, val ? (mod_json_size_t)std::strlen(val) : 0); + } + + //! Constructor + JsonValue(const char *val, size_type len) { + val_ = mod_json_value_set_buffer(val, len); + } + + //! Constructor + JsonValue(const std::string &val) { + val_ = mod_json_value_set_buffer(val.data(), (mod_json_size_t)val.size()); + } + + //! Constructor + JsonValue(const JsonArray &val) { + val_ = mod_json_value_set_array(*(mod_json_array_t **)&val); + } + + //! Constructor + JsonValue(const JsonObject &val) { + val_ = mod_json_value_set_object(*(mod_json_object_t **)&val); + } + + //! Constructor + JsonValue(const JsonValue &rhs) : val_(0) { + if (rhs.val_) { + val_ = mod_json_value_grab(rhs.val_); + } + } + +#if __cplusplus >= 201103L + //! Constructor + JsonValue(JsonValue &&rhs) : val_(rhs.val_) { + rhs.val_ = 0; + } +#endif + + //! Destructor + ~JsonValue(void) { + mod_json_value_unset(val_); + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const JsonValue &rhs) { + this->assign(rhs); + return *this; + } + +#if __cplusplus >= 201103L + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(JsonValue &&rhs) { + this->assign(std::move(rhs)); + return *this; + } +#endif + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const bool &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const signed char &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const char &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const short int &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const int &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const long int &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const long long int &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const float &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const double &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const long double &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const unsigned char &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const unsigned short int &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const unsigned int &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const unsigned long int &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const unsigned long long int &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const JsonString &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const char *val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const std::string &val) { + this->assign(val); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const JsonArray &arr) { + this->assign(arr); + return *this; + } + + //! Assign new contents to the value, replacing its current content + JsonValue &operator=(const JsonObject &obj) { + this->assign(obj); + return *this; + } + + //! Equality + bool operator==(const JsonValue &rhs) const { + return mod_json_value_is_equal(val_, rhs.val_); + } + + //! No equality + bool operator!=(const JsonValue &rhs) const { + return !(*this == rhs); + } + + //! Treat self value as object by force, retrieving value of a key + JsonValue &operator[](const char *key) { + return this->get_value(key); + } + + //! Retrieve a reference of value by a key + JsonValue operator[](const char *key) const { + return this->get_value(key); + } + + //! Treat self value as object by force, retrieving value of a key + JsonValue &operator[](const JsonString &key) { + return this->get_value(key.c_str()); + } + + //! Retrieve a reference of value by a key + JsonValue operator[](const JsonString &key) const { + return this->get_value(key.c_str()); + } + + //! Treat self value as object by force, retrieving value of a key + JsonValue &operator[](const std::string &key) { + return this->get_value(key.c_str()); + } + + //! Retrieve a reference of value by a key + JsonValue operator[](const std::string &key) const { + return this->get_value(key.c_str()); + } + + //! Treat self value as array by force, retrieving value at index n + JsonValue &operator[](size_type n) { + return this->get_value(n); + } + + //! Retrieve a reference of value at index n + JsonValue operator[](size_type n) const { + return this->get_value(n); + } + + //! Retrieve non-zero if the value is valid + bool is_valid(void) const { + return (val_ != (mod_json_value_t *)0); + } + + //! Retrieve non-zero if the value is a object + bool is_object(void) const { + return mod_json_value_is_object(val_); + } + + //! Retrieve non-zero if the value is an array + bool is_array(void) const { + return mod_json_value_is_array(val_); + } + + //! Retrieve non-zero if the value is a string + bool is_string(void) const { + return mod_json_value_is_string(val_); + } + + //! Retrieve non-zero if the value is null + bool is_null(void) const { + return mod_json_value_is_null(val_); + } + + //! Retrieve non-zero if the value is a float + bool is_float(void) const { + return mod_json_value_is_float(val_); + } + + //! Retrieve non-zero if the value is an integer + bool is_integer(void) const { + return mod_json_value_is_integer(val_); + } + + //! Retrieve non-zero if the value is a boolean + bool is_boolean(void) const { + return mod_json_value_is_boolean(val_); + } + + //! Assign new contents to the value, replacing its current content + void assign(const JsonValue &rhs) { + mod_json_value_unset(val_); + val_ = rhs.val_ ? mod_json_value_grab(rhs.val_) : 0; + } + +#if __cplusplus >= 201103L + //! Assign new contents to the value, replacing its current content + void assign(JsonValue &&rhs) { + mod_json_value_unset(val_); + val_ = rhs.val_; + rhs.val_ = 0; + } +#endif + + //! Assign new contents to the value, replacing its current content + void assign(const bool &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_boolean(val_, (mod_json_boolean_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const signed char &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const char &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const short int &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const int &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const long int &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const long long int &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const float &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_float(val_, (mod_json_float_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const double &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_float(val_, (mod_json_float_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const long double &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_float(val_, (mod_json_float_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const unsigned char &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const unsigned short int &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const unsigned int &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const unsigned long int &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const unsigned long long int &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_integer(val_, (mod_json_integer_t)val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const JsonString &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_string(val_, *(mod_json_string_t **)&val); + } + + //! Assign new contents to the value, replacing its current content + void assign(const char *val) { + JsonString str(val); + if (!str.is_valid() || !copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_string(val_, *(mod_json_string_t **)&str); + } + + //! Assign new contents to the value, replacing its current content + void assign(const char *val, size_type len) { + JsonString str(val, len); + if (!str.is_valid() || !copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_string(val_, *(mod_json_string_t **)&str); + } + + //! Assign new contents to the value, replacing its current content + void assign(const std::string &val) { + JsonString str(val); + if (!str.is_valid() || !copy_on_write()) { + throw std::runtime_error("JsonValue::assign"); + } + mod_json_value_assign_string(val_, *(mod_json_string_t **)&str); + } + + //! Assign new contents to the value, replacing its current content + void assign(const JsonArray &arr); + + //! Assign new contents to the value, replacing its current content + void assign(const JsonObject &obj); + + //! Retrieve refer-counter of JSON value + ssize_type refer(void) const { + return mod_json_value_refer(val_); + } + + //! Retrieve value as JSON format string + JsonString as_json_string(void) const { + mod_json_string_t *tmp = mod_json_dump(val_); + JsonString ret = *reinterpret_cast(&tmp); + if (tmp) { + mod_json_string_unset(tmp); + } + return ret; + } + + //! Retrieve value as a STL string + std::string as_stl_string(void) const { + if (is_string()) { + return to_string().as_stl_string(); + } + return std::string(); + } + + //! Retrieve value as JSON string + const JsonString &as_string(void) const { + if (!is_string()) { + throw std::logic_error("JsonValue::as_string"); + } + return to_string(); + } + + //! Retrieve value as c-style string + const char *as_c_string(void) const { + return mod_json_value_cstring(val_); + } + + //! Retrieve value as JSON string + JsonString &as_string(void) { + if (!is_string()) { + throw std::logic_error("JsonValue::as_string"); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonValue::as_string"); + } + return to_string(); + } + + //! Retrieve value as JSON array + const JsonArray &as_array(void) const { + if (!is_array()) { + throw std::logic_error("JsonValue::as_array"); + } + return to_array(); + } + + //! Retrieve value as JSON array + JsonArray &as_array(void) { + if (!is_array()) { + throw std::logic_error("JsonValue::as_array"); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonValue::as_array"); + } + return to_array(); + } + + //! Retrieve value as JSON object + const JsonObject &as_object(void) const { + if (!is_object()) { + throw std::logic_error("JsonValue::as_object"); + } + return to_object(); + } + + //! Retrieve value as JSON object + JsonObject &as_object(void) { + if (!is_object()) { + throw std::logic_error("JsonValue::as_object"); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonValue::as_object"); + } + return to_object(); + } + + //! Retrieve value as float + float_type as_float(void) const { + return mod_json_value_float(val_); + } + + //! Retrieve value as integer + integer_type as_integer(void) const { + return mod_json_value_integer(val_); + } + + //! Retrieve value as boolean + bool as_bool(void) const { + return mod_json_value_boolean(val_); + } + + //! Exchange the content with another JSON value + void swap(JsonValue &rhs) { + mod_json_value_t *val = val_; + val_ = rhs.val_; + rhs.val_ = val; + } + + //! Merge another JSON value + void merge(const JsonValue &rhs) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::merge"); + } + mod_json_value_merge(val_, rhs.val_); + } + + //! Parse a sting as a JSON value + bool parse(const char *str) { + mod_json_token_t *tok = mod_json_token_create(NULL); + + if (tok) { + mod_json_value_t *jval = mod_json_parse(tok, str); + + mod_json_token_destroy(tok); + if (jval) { + *this = *reinterpret_cast(&jval); + mod_json_value_unset(jval); + return is_valid(); + } + } + return false; + } + + //! Parse a sting as a JSON value + bool parse(const JsonString &str) { + return this->parse(str.c_str()); + } + + //! Parse a sting as a JSON value + bool parse(const std::string &str) { + return this->parse(str.c_str()); + } + + protected: + //! Clone the value for writing + bool copy_on_write(void) { + if (val_) { + if (mod_json_value_is_shared(val_)) { + mod_json_value_put(val_); + val_ = mod_json_value_clone(val_); + } + } else { + val_ = mod_json_value_set_null(); + } + return (val_ != 0); + } + + //! Clone the value and leak it + bool copy_and_leak(void) { + if (copy_on_write()) { + mod_json_value_set_leaked(val_); + return true; + } + return false; + } + + //! Convert value to JSON object + JsonObject &to_object(void); + + //! Convert value to JSON object + const JsonObject &to_object(void) const; + + //! Convert value to JSON array + JsonArray &to_array(void); + + //! Convert value to JSON array + const JsonArray &to_array(void) const; + + //! Convert value to JSON string + JsonString &to_string(void); + + //! Convert value to JSON string + const JsonString &to_string(void) const; + + //! Treat self value as object by force, retrieving value of a key + JsonValue &get_value(const char *key); + + //! Retrieve a reference of value by a key + JsonValue get_value(const char *key) const; + + //! Treat self value as array by force, retrieving value at index n + JsonValue &get_value(size_type n); + + //! Retrieve a reference of value at index n + JsonValue get_value(size_type n) const; + + //! Set the new array to the value, replacing its current content + void set_value(const JsonArray &val); + + //! Set the new object to the value, replacing its current content + void set_value(const JsonObject &val); + + private: + mod_json_value_t *val_; +}; + +/*! JSON Array + */ +class JsonArray { + public: + typedef mod_json_size_t size_type; + typedef mod_json_ssize_t ssize_type; + + class iterator; + class const_iterator; + class reverse_iterator; + class const_reverse_iterator; + + /*! Const iterator of JSON Array + */ + class const_iterator { + public: + //! Constructor + const_iterator(void) : iter_(0) {} + + //! Equality + bool operator==(const const_iterator &rhs) const { + return (iter_ == rhs.iter_); + } + + //! No equality + bool operator!=(const const_iterator &rhs) const { + return (iter_ != rhs.iter_); + } + + //! Increment (Prefix) + const_iterator &operator++() { + ++iter_; + return *this; + } + + //! Increment (Suffix) + const_iterator operator++(int) { + const_iterator tmp = *this; + ++iter_; + return tmp; + } + + //! Decrement (Prefix) + const_iterator &operator--() { + --iter_; + return *this; + } + + //! Decrement (Suffix) + const_iterator operator--(int) { + const_iterator tmp = *this; + --iter_; + return tmp; + } + + //! Indirection (eg. *iter) + const JsonValue &operator*() const { + return *reinterpret_cast(iter_); + } + + //! Structure dereference (eg. iter->) + const JsonValue *operator->() const { + return reinterpret_cast(iter_); + } + + //! Retrieve as const reverse iterator + operator const_reverse_iterator() const { + return const_reverse_iterator(iter_); + } + + protected: + friend class JsonArray; + friend class JsonArray::iterator; + friend class JsonArray::reverse_iterator; + friend class JsonArray::const_reverse_iterator; + + //! Constructor for friends + const_iterator(mod_json_value_t *const *iter) : iter_(iter) {} + + private: + mod_json_value_t *const *iter_; + }; + + /*! iterator of JSON Array + */ + class iterator { + public: + //! Constructor + iterator(void) : iter_(0) {} + + //! Equality + bool operator==(const iterator &rhs) const { + return (iter_ == rhs.iter_); + } + + //! No equality + bool operator!=(const iterator &rhs) const { + return (iter_ != rhs.iter_); + } + + //! Increment (Prefix) + iterator &operator++() { + ++iter_; + return *this; + } + + //! Increment (Suffix) + iterator operator++(int) { + iterator tmp = *this; + ++iter_; + return tmp; + } + + //! Decrement (Prefix) + iterator &operator--() { + --iter_; + return *this; + } + + //! Decrement (Suffix) + iterator operator--(int) { + iterator tmp = *this; + --iter_; + return tmp; + } + + //! Indirection (eg. *iter) + JsonValue &operator*() const { + return *reinterpret_cast(iter_); + } + + //! Structure dereference (eg. iter->) + JsonValue *operator->() const { + return reinterpret_cast(iter_); + } + + //! Retrieve as const iterator + operator const_iterator() const { + return const_iterator(iter_); + } + + //! Retrieve as reverse iterator + operator reverse_iterator() const { + return reverse_iterator(iter_); + } + + //! Retrieve as const reverse iterator + operator const_reverse_iterator() const { + return const_reverse_iterator(iter_); + } + + protected: + friend class JsonArray; + friend class JsonArray::reverse_iterator; + + //! Constructor for friends + iterator(mod_json_value_t **iter) : iter_(iter) {} + + private: + mod_json_value_t **iter_; + }; + + /*! Const Reverse iterator of JSON Array + */ + class const_reverse_iterator { + public: + //! Constructor + const_reverse_iterator(void) : iter_(0) {} + + //! Equality + bool operator==(const const_reverse_iterator &rhs) const { + return (iter_ == rhs.iter_); + } + + //! No equality + bool operator!=(const const_reverse_iterator &rhs) const { + return (iter_ != rhs.iter_); + } + + //! Increment (Prefix) + const_reverse_iterator &operator++() { + --iter_; + return *this; + } + + //! Increment (Suffix) + const_reverse_iterator operator++(int) { + const_reverse_iterator tmp = *this; + --iter_; + return tmp; + } + + //! Decrement (Prefix) + const_reverse_iterator &operator--() { + ++iter_; + return *this; + } + + //! Decrement (Suffix) + const_reverse_iterator operator--(int) { + const_reverse_iterator tmp = *this; + ++iter_; + return tmp; + } + + //! Indirection (eg. *iter) + const JsonValue &operator*() const { + return *reinterpret_cast(iter_); + } + + //! Structure dereference (eg. iter->) + const JsonValue *operator->() const { + return reinterpret_cast(iter_); + } + + //! Retrieve as const iterator + operator const_iterator() const { + return const_iterator(iter_); + } + + protected: + friend class JsonArray; + friend class JsonArray::iterator; + friend class JsonArray::const_iterator; + friend class JsonArray::reverse_iterator; + + //! Constructor for friends + const_reverse_iterator(mod_json_value_t *const *iter) : iter_(iter) {} + + private: + mod_json_value_t *const *iter_; + }; + + /*! Reverse iterator of JSON Array + */ + class reverse_iterator { + public: + //! Constructor + reverse_iterator(void) : iter_(0) {} + + //! Equality + bool operator==(const reverse_iterator &rhs) const { + return (iter_ == rhs.iter_); + } + + //! No equality + bool operator!=(const reverse_iterator &rhs) const { + return (iter_ != rhs.iter_); + } + + //! Increment (Prefix) + reverse_iterator &operator++() { + --iter_; + return *this; + } + + //! Increment (Suffix) + reverse_iterator operator++(int) { + reverse_iterator tmp = *this; + --iter_; + return tmp; + } + + //! Decrement (Prefix) + reverse_iterator &operator--() { + ++iter_; + return *this; + } + + //! Decrement (Suffix) + reverse_iterator operator--(int) { + reverse_iterator tmp = *this; + ++iter_; + return tmp; + } + + //! Indirection (eg. *iter) + JsonValue &operator*() const { + return *reinterpret_cast(iter_); + } + + //! Structure dereference (eg. iter->) + JsonValue *operator->() const { + return reinterpret_cast(iter_); + } + + //! Retrieve as iterator + operator iterator() const { + return iterator(iter_); + } + + //! Retrieve as const iterator + operator const_iterator() const { + return const_iterator(iter_); + } + + //! Retrieve as const reverse iterator + operator const_reverse_iterator() const { + return const_reverse_iterator(iter_); + } + + protected: + friend class JsonArray; + friend class JsonArray::iterator; + + //! Constructor for friends + reverse_iterator(mod_json_value_t **iter) : iter_(iter) {} + + private: + mod_json_value_t **iter_; + }; + + //! Constructor + JsonArray(void) : arr_(0) {} + + //! Constructor + JsonArray(const JsonArray &rhs) : arr_(0) { + if (rhs.arr_) { + arr_ = mod_json_array_grab(rhs.arr_); + } + } + +#if __cplusplus >= 201103L + //! Constructor + JsonArray(JsonArray &&rhs) : arr_(rhs.arr_) { + rhs.arr_ = 0; + } +#endif + + //! Destructor + ~JsonArray(void) { + mod_json_array_unset(arr_); + } + + //! Assign new contents to the array, replacing its current content + JsonArray &operator=(const JsonArray &rhs) { + this->assign(rhs); + return *this; + } + +#if __cplusplus >= 201103L + //! Assign new contents to the array, replacing its current content + JsonArray &operator=(JsonArray &&rhs) { + this->assign(std::move(rhs)); + return *this; + } +#endif + + //! Equality + bool operator==(const JsonArray &rhs) const { + return mod_json_array_is_equal(arr_, rhs.arr_); + } + + //! No equality + bool operator!=(const JsonArray &rhs) const { + return !(*this == rhs); + } + + //! Retrieve the value at index n, if no one exists, throw an exception. + JsonValue &operator[](size_type n) { + return this->at(n); + } + + //! Retrieve the value at index n, if no one exists, return a null value. + JsonValue operator[](size_type n) const { + return ((n < this->size()) ? this->get_value(n) : JsonValue()); + } + + //! Retrieve non-zero if the array is valid + bool is_valid(void) const { + return (arr_ != (mod_json_array_t *)0); + } + + //! Retrieve non-zero if the array is empty + bool empty(void) const { + return mod_json_array_empty(arr_); + } + + //! Retrieve the size of JSON array + size_type size(void) const { + return mod_json_array_count(arr_); + } + + //! Retrieve the capacity of JSON array + size_type capacity(void) const { + return mod_json_array_capacity(arr_); + } + + //! Retrieve refer-counter of JSON array + ssize_type refer(void) const { + return mod_json_array_refer(arr_); + } + + //! Assign new contents to the array, replacing its current content + void assign(const JsonArray &rhs) { + mod_json_array_unset(arr_); + arr_ = rhs.arr_ ? mod_json_array_grab(rhs.arr_) : 0; + } + +#if __cplusplus >= 201103L + //! Assign new contents to the array, replacing its current content + void assign(JsonArray &&rhs) { + mod_json_array_unset(arr_); + arr_ = rhs.arr_; + rhs.arr_ = 0; + } +#endif + + //! Request a change in capacity + void reserve(size_type n) { + if (!copy_on_write() || mod_json_array_reserve(arr_, n) != 0) { + throw std::runtime_error("JsonArray::reserve"); + } + } + + //! Reverse the order of the elements + void reverse(void) { + if (arr_ && copy_on_write()) { + mod_json_array_reverse(arr_); + } + } + + //! Push a value to array + void push(const JsonValue &val) { + JsonValue tmp(val); + + if (!copy_on_write() || + mod_json_array_push(arr_, *((mod_json_value_t **)&tmp)) != 0) { + throw std::runtime_error("JsonArray::push"); + } + } + + //! Pop the last element from array + void pop(void) { + if (arr_) { + if (!copy_on_write()) { + throw std::runtime_error("JsonArray::pop"); + } + mod_json_array_pop(arr_); + } + } + + //! Remove the first element of array + void shift(void) { + if (arr_) { + if (!copy_on_write()) { + throw std::runtime_error("JsonArray::shift"); + } + mod_json_array_shift(arr_); + } + } + + //! Retrieve the value at index n + JsonValue &at(size_type n) { + if (this->size() <= n) { + throw std::out_of_range("JsonArray::at"); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonArray::at"); + } + return this->get_value(n); + } + + //! Retrieve the value at index n + const JsonValue &at(size_type n) const { + if (this->size() <= n) { + throw std::out_of_range("JsonArray::at"); + } + return this->get_value(n); + } + + //! Retrieve a reference to the first element + JsonValue &front(void) { + if (this->size() <= 0) { + throw std::out_of_range("JsonArray::front"); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonArray::front"); + } + return this->get_value(0); + } + + //! Retrieve a reference to the first element + const JsonValue &front(void) const { + if (this->size() <= 0) { + throw std::out_of_range("JsonArray::front"); + } + return this->get_value(0); + } + + //! Retrieve a reference to the last element + JsonValue &back(void) { + if (this->size() <= 0) { + throw std::out_of_range("JsonArray::back"); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonArray::back"); + } + return this->get_value(this->size() - 1); + } + + //! Retrieve a reference to the last element + const JsonValue &back(void) const { + if (this->size() <= 0) { + throw std::out_of_range("JsonArray::back"); + } + return this->get_value(this->size() - 1); + } + + //! Clear the JSON array + void clear(void) { + mod_json_array_unset(arr_); + arr_ = 0; + } + + //! Exchange the content with another JSON array + void swap(JsonArray &rhs) { + mod_json_array_t *arr = arr_; + arr_ = rhs.arr_; + rhs.arr_ = arr; + } + + //! Merge another JSON array + void merge(const JsonArray &rhs) { + if (!copy_on_write()) { + throw std::runtime_error("JsonArray::merge"); + } + mod_json_array_merge(arr_, rhs.arr_); + } + + //! Resize a JSON array so that it contains n elements + void resize(size_type n, const JsonValue &val = JsonValue()) { + if (!copy_on_write() || + mod_json_array_resize(arr_, n, *((mod_json_value_t **)&val)) != 0) { + throw std::runtime_error("JsonArray::resize"); + } + } + + //! Retrieve an iterator pointing to the first element + iterator begin(void) { + if (copy_and_leak()) { + return iterator(mod_json_array_begin(arr_)); + } + return iterator(); + } + + //! Retrieve a const iterator pointing to the first element + const_iterator begin(void) const { + if (arr_) { + return const_iterator(mod_json_array_begin(arr_)); + } + return const_iterator(); + } + + //! Retrieve a const iterator pointing to the first element + const_iterator cbegin(void) const { + if (arr_) { + return const_iterator(mod_json_array_begin(arr_)); + } + return const_iterator(); + } + + //! Retrieve a reverse iterator pointing to the last element + reverse_iterator rbegin(void) { + if (copy_and_leak()) { + return reverse_iterator(mod_json_array_rbegin(arr_)); + } + return reverse_iterator(); + } + + //! Retrieve a const reverse iterator pointing to the last element + const_reverse_iterator rbegin(void) const { + if (arr_) { + return const_reverse_iterator(mod_json_array_rbegin(arr_)); + } + return const_reverse_iterator(); + } + + //! Retrieve a const reverse iterator pointing to the last element + const_reverse_iterator crbegin(void) const { + if (arr_) { + return const_reverse_iterator(mod_json_array_rbegin(arr_)); + } + return const_reverse_iterator(); + } + + //! Retrieve an iterator pointing to the past-the-end element + iterator end(void) { + if (copy_and_leak()) { + return iterator(mod_json_array_end(arr_)); + } + return iterator(); + } + + //! Retrieve a const iterator pointing to the past-the-end element + const_iterator end(void) const { + if (arr_) { + return const_iterator(mod_json_array_end(arr_)); + } + return const_iterator(); + } + + //! Retrieve a const iterator pointing to the past-the-end element + const_iterator cend(void) const { + if (arr_) { + return const_iterator(mod_json_array_end(arr_)); + } + return const_iterator(); + } + + //! Retrieve a reverse pointing to the past-the-end element + reverse_iterator rend(void) { + if (copy_and_leak()) { + return reverse_iterator(mod_json_array_rend(arr_)); + } + return reverse_iterator(); + } + + //! Retrieve a const reverse pointing to the past-the-end element + const_reverse_iterator rend(void) const { + if (arr_) { + return const_reverse_iterator(mod_json_array_rend(arr_)); + } + return const_reverse_iterator(); + } + + //! Retrieve a const reverse pointing to the past-the-end element + const_reverse_iterator crend(void) const { + if (arr_) { + return const_reverse_iterator(mod_json_array_rend(arr_)); + } + return const_reverse_iterator(); + } + + protected: + //! Clone the array for writing + bool copy_on_write(void) { + if (arr_) { + if (mod_json_array_is_shared(arr_)) { + mod_json_array_put(arr_); + arr_ = mod_json_array_clone(arr_); + } + } else { + arr_ = mod_json_array_set_default(); + } + return (arr_ != 0); + } + + //! Clone the array and leak it + bool copy_and_leak(void) { + if (copy_on_write()) { + mod_json_array_set_leaked(arr_); + return true; + } + return false; + } + + //! Retrieve the value at index n + JsonValue &get_value(size_type n) { + return *reinterpret_cast(arr_->first + n); + } + + //! Retrieve the value at index n + const JsonValue &get_value(size_type n) const { + return *reinterpret_cast(arr_->first + n); + } + + private: + mod_json_array_t *arr_; +}; + +/*! JSON Pair + */ +class JsonPair { + public: + //! Constructor + JsonPair(void) : pair_(0) {} + + //! Retrieve non-zero if the pair is valid + bool is_valid(void) const { + return (pair_ != (mod_json_pair_t *)0); + } + + //! Retrieve the key of pair + const JsonString &key(void) const { + return *reinterpret_cast(&pair_->key); + } + + //! Retrieve the value of pair + JsonValue &value(void) { + return *reinterpret_cast(&pair_->val); + } + + //! Retrieve the value of pair + const JsonValue &value(void) const { + return *reinterpret_cast(&pair_->val); + } + + protected: + friend class JsonObject; + + //! Constructor for friends + JsonPair(mod_json_pair_t *pair) : pair_(pair) {} + + //! Constructor for friends + JsonPair(const JsonPair &rhs) : pair_(rhs.pair_) {} + + private: + mod_json_pair_t *pair_; +}; + +/*! JSON Object + */ +class JsonObject { + public: + typedef mod_json_size_t size_type; + typedef mod_json_ssize_t ssize_type; + + class iterator; + class const_iterator; + class reverse_iterator; + class const_reverse_iterator; + + /*! Const iterator of JSON Object + */ + class const_iterator { + public: + //! Constructor + const_iterator(void) : iter_(0) {} + + //! Equality + bool operator==(const const_iterator &rhs) const { + return (iter_ == rhs.iter_); + } + + //! No equality + bool operator!=(const const_iterator &rhs) const { + return (iter_ != rhs.iter_); + } + + //! Increment (Prefix) + const_iterator &operator++() { + ++iter_; + return *this; + } + + //! Increment (Suffix) + const_iterator operator++(int) { + const_iterator tmp = *this; + ++iter_; + return tmp; + } + + //! Decrement (Prefix) + const_iterator &operator--() { + --iter_; + return *this; + } + + //! Decrement (Suffix) + const_iterator operator--(int) { + const_iterator tmp = *this; + --iter_; + return tmp; + } + + //! Indirection (eg. *iter) + const JsonPair &operator*() const { + return *reinterpret_cast(&iter_); + } + + //! Structure dereference (eg. iter->) + const JsonPair *operator->() const { + return reinterpret_cast(&iter_); + } + + //! Retrieve as const reverse iterator + operator const_reverse_iterator() const { + return const_reverse_iterator(iter_); + } + + protected: + friend class JsonObject; + friend class JsonObject::iterator; + friend class JsonObject::reverse_iterator; + friend class JsonObject::const_reverse_iterator; + + //! Constructor for friends + const_iterator(const mod_json_pair_t *iter) : iter_(iter) {} + + private: + const mod_json_pair_t *iter_; + }; + + /*! iterator of JSON Object + */ + class iterator { + public: + //! Constructor + iterator(void) : iter_(0) {} + + //! Equality + bool operator==(const iterator &rhs) const { + return (iter_ == rhs.iter_); + } + + //! No equality + bool operator!=(const iterator &rhs) const { + return (iter_ != rhs.iter_); + } + + //! Increment (Prefix) + iterator &operator++() { + ++iter_; + return *this; + } + + //! Increment (Suffix) + iterator operator++(int) { + iterator tmp = *this; + ++iter_; + return tmp; + } + + //! Decrement (Prefix) + iterator &operator--() { + --iter_; + return *this; + } + + //! Decrement (Suffix) + iterator operator--(int) { + iterator tmp = *this; + --iter_; + return tmp; + } + + //! Indirection (eg. *iter) + JsonPair &operator*() const { + return *reinterpret_cast((mod_json_pair_t **)&iter_); + } + + //! Structure dereference (eg. iter->) + JsonPair *operator->() const { + return reinterpret_cast((mod_json_pair_t **)&iter_); + } + + //! Retrieve as const iterator + operator const_iterator() const { + return const_iterator(iter_); + } + + //! Retrieve as reverse iterator + operator reverse_iterator() const { + return reverse_iterator(iter_); + } + + //! Retrieve as const reverse iterator + operator const_reverse_iterator() const { + return const_reverse_iterator(iter_); + } + + protected: + friend class JsonObject; + friend class JsonObject::reverse_iterator; + + //! Constructor for friends + iterator(mod_json_pair_t *iter) : iter_(iter) {} + + private: + mod_json_pair_t *iter_; + }; + + /*! Const Reverse iterator of JSON Object + */ + class const_reverse_iterator { + public: + //! Constructor + const_reverse_iterator(void) : iter_(0) {} + + //! Equality + bool operator==(const const_reverse_iterator &rhs) const { + return (iter_ == rhs.iter_); + } + + //! No equality + bool operator!=(const const_reverse_iterator &rhs) const { + return (iter_ != rhs.iter_); + } + + //! Increment (Prefix) + const_reverse_iterator &operator++() { + --iter_; + return *this; + } + + //! Increment (Suffix) + const_reverse_iterator operator++(int) { + const_reverse_iterator tmp = *this; + --iter_; + return tmp; + } + + //! Decrement (Prefix) + const_reverse_iterator &operator--() { + ++iter_; + return *this; + } + + //! Decrement (Suffix) + const_reverse_iterator operator--(int) { + const_reverse_iterator tmp = *this; + ++iter_; + return tmp; + } + + //! Indirection (eg. *iter) + const JsonPair &operator*() const { + return *reinterpret_cast(&iter_); + } + + //! Structure dereference (eg. iter->) + const JsonPair *operator->() const { + return reinterpret_cast(&iter_); + } + + //! Retrieve as const iterator + operator const_iterator() const { + return const_iterator(iter_); + } + + protected: + friend class JsonObject; + friend class JsonObject::iterator; + friend class JsonObject::const_iterator; + friend class JsonObject::reverse_iterator; + + //! Constructor for friends + const_reverse_iterator(const mod_json_pair_t *iter) : iter_(iter) {} + + private: + const mod_json_pair_t *iter_; + }; + + /*! iterator of JSON Object + */ + class reverse_iterator { + public: + //! Constructor + reverse_iterator(void) : iter_(0) {} + + //! Equality + bool operator==(const reverse_iterator &rhs) const { + return (iter_ == rhs.iter_); + } + + //! No equality + bool operator!=(const reverse_iterator &rhs) const { + return (iter_ != rhs.iter_); + } + + //! Increment (Prefix) + reverse_iterator &operator++() { + --iter_; + return *this; + } + + //! Increment (Suffix) + reverse_iterator operator++(int) { + reverse_iterator tmp = *this; + --iter_; + return tmp; + } + + //! Decrement (Prefix) + reverse_iterator &operator--() { + ++iter_; + return *this; + } + + //! Decrement (Suffix) + reverse_iterator operator--(int) { + reverse_iterator tmp = *this; + ++iter_; + return tmp; + } + + //! Indirection (eg. *iter) + JsonPair &operator*() const { + return *reinterpret_cast((mod_json_pair_t **)&iter_); + } + + //! Structure dereference (eg. iter->) + JsonPair *operator->() const { + return reinterpret_cast((mod_json_pair_t **)&iter_); + } + + //! Retrieve as iterator + operator iterator() const { + return iterator(iter_); + } + + //! Retrieve as const iterator + operator const_iterator() const { + return const_iterator(iter_); + } + + //! Retrieve as const reverse iterator + operator const_reverse_iterator() const { + return const_reverse_iterator(iter_); + } + + protected: + friend class JsonObject; + friend class JsonArray::iterator; + + //! Constructor for friends + reverse_iterator(mod_json_pair_t *iter) : iter_(iter) {} + + private: + mod_json_pair_t *iter_; + }; + + //! Constructor + JsonObject(void) : obj_(0) {} + + //! Constructor + JsonObject(const JsonObject &rhs) : obj_(0) { + if (rhs.obj_) { + obj_ = mod_json_object_grab(rhs.obj_); + } + } + +#if __cplusplus >= 201103L + //! Constructor + JsonObject(JsonObject &&rhs) : obj_(rhs.obj_) { + rhs.obj_ = 0; + } +#endif + + //! Destructor + ~JsonObject(void) { + mod_json_object_unset(obj_); + } + + //! Assign new contents to the object, replacing its current content + JsonObject &operator=(const JsonObject &rhs) { + this->assign(rhs); + return *this; + } + +#if __cplusplus >= 201103L + //! Assign new contents to the object, replacing its current content + JsonObject &operator=(JsonObject &&rhs) { + this->assign(std::move(rhs)); + return *this; + } +#endif + + //! Equality + bool operator==(const JsonObject &rhs) const { + return mod_json_object_is_equal(obj_, rhs.obj_); + } + + //! No equality + bool operator!=(const JsonObject &rhs) const { + return !(*this == rhs); + } + + //! Retrieve the value of a key, if no one exists, create a new one. + JsonValue &operator[](const char *key) { + if (!key) { + throw std::invalid_argument("JsonObject::operator[]"); + } + + if (!copy_and_leak()) { + throw std::runtime_error("JsonObject::operator[]"); + } + + JsonPair pair(mod_json_object_touch(obj_, key)); + if (!pair.is_valid()) { + throw std::runtime_error("JsonObject::operator[]"); + } + return pair.value(); + } + + //! Retrieve the value of a key, if no one exists, return a null value. + JsonValue operator[](const char *key) const { + if (!key) { + throw std::invalid_argument("JsonObject::operator[]"); + } + + JsonPair pair(mod_json_object_find(obj_, key)); + return (pair.is_valid() ? pair.value() : JsonValue()); + } + + //! Retrieve the value of a key, if no one exists, create a new one. + JsonValue &operator[](const JsonString &key) { + return (*this)[key.c_str()]; + } + + //! Retrieve the value of a key, if no one exists, return a null value. + JsonValue operator[](const JsonString &key) const { + return (*this)[key.c_str()]; + } + + //! Retrieve non-zero if the object is valid + bool is_valid(void) const { + return (obj_ != (mod_json_object_t *)0); + } + + //! Retrieve non-zero if the object is empty + bool empty(void) const { + return mod_json_object_empty(obj_); + } + + //! Retrieve the size of JSON object + size_type size(void) const { + return mod_json_object_count(obj_); + } + + //! Retrieve refer-counter of JSON object + ssize_type refer(void) const { + return mod_json_object_refer(obj_); + } + + //! Assign new contents to the object, replacing its current content + void assign(const JsonObject &rhs) { + mod_json_object_unset(obj_); + obj_ = rhs.obj_ ? mod_json_object_grab(rhs.obj_) : 0; + } + +#if __cplusplus >= 201103L + //! Assign new contents to the object, replacing its current content + void assign(JsonObject &&rhs) { + mod_json_object_unset(obj_); + obj_ = rhs.obj_; + rhs.obj_ = 0; + } +#endif + + //! Clear the JSON object + void clear(void) { + mod_json_object_unset(obj_); + obj_ = 0; + } + + //! Set the value of a key + bool set(const JsonString &key, const JsonValue &val) { + JsonValue tmp(val); + if (!copy_on_write()) { + throw std::runtime_error("JsonObject::set"); + } + return (mod_json_object_insert(obj_, *(mod_json_string_t **)&key, + *(mod_json_value_t **)&tmp) != + (mod_json_pair_t *)0); + } + + //! Retrieve the value of a key + bool get(const char *key, JsonValue *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = pair.value(); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, JsonString *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid() || !pair.value().is_string()) { + return false; + } + *val = pair.value().as_string(); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, std::string *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid() || !pair.value().is_string()) { + return false; + } + *val = pair.value().as_stl_string(); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, JsonArray *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid() || !pair.value().is_array()) { + return false; + } + *val = pair.value().as_array(); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, JsonObject *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid() || !pair.value().is_object()) { + return false; + } + *val = pair.value().as_object(); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, bool *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = pair.value().as_bool(); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, signed char *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, char *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, short int *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, int *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, long int *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, long long int *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, unsigned char *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, unsigned short int *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, unsigned int *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, unsigned long int *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, unsigned long long int *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_integer()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, float *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_float()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, double *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_float()); + return true; + } + + //! Retrieve the value of a key + bool get(const char *key, long double *val) const { + const JsonPair pair(mod_json_object_find(obj_, key)); + if (!pair.is_valid()) { + return false; + } + *val = static_cast(pair.value().as_float()); + return true; + } + + //! Retrieve the value of a key + template + bool get(const JsonString &key, T *val) const { + return this->get(key.c_str(), val); + } + + //! Retrieve the value of a key + template + bool get(const std::string &key, T *val) const { + return this->get(key.c_str(), val); + } + + //! Delete a key-value pair from JSON object + void unset(const char *key) { + if (obj_ && key) { + if (!copy_on_write()) { + throw std::runtime_error("JsonObject::unset"); + } + mod_json_object_erase(obj_, key); + } + } + + //! Retrieve non-zero if the key exists in JSON object + bool has(const char *key) const { + return (mod_json_object_find(obj_, key) != (mod_json_pair_t *)0); + } + + //! Exchange the content with another JSON object + void swap(JsonObject &rhs) { + mod_json_object_t *obj = obj_; + obj_ = rhs.obj_; + rhs.obj_ = obj; + } + + //! Merge another JSON object + void merge(const JsonObject &rhs) { + if (!copy_on_write()) { + throw std::runtime_error("JsonObject::merge"); + } + mod_json_object_merge(obj_, rhs.obj_); + } + + //! Retrieve an iterator pointing to the first element + iterator begin(void) { + if (copy_and_leak()) { + return iterator(mod_json_object_begin(obj_)); + } + return iterator(); + } + + //! Retrieve a const iterator pointing to the first element + const_iterator begin(void) const { + if (obj_) { + return const_iterator(mod_json_object_begin(obj_)); + } + return const_iterator(); + } + + //! Retrieve a const iterator pointing to the first element + const_iterator cbegin(void) const { + if (obj_) { + return const_iterator(mod_json_object_begin(obj_)); + } + return const_iterator(); + } + + //! Retrieve a reverse iterator pointing to the last element + reverse_iterator rbegin(void) { + if (copy_and_leak()) { + return reverse_iterator(mod_json_object_rbegin(obj_)); + } + return reverse_iterator(); + } + + //! Retrieve a const reverse iterator pointing to the last element + const_reverse_iterator rbegin(void) const { + if (obj_) { + return const_reverse_iterator(mod_json_object_rbegin(obj_)); + } + return const_reverse_iterator(); + } + + //! Retrieve a const reverse iterator pointing to the last element + const_reverse_iterator crbegin(void) const { + if (obj_) { + return const_reverse_iterator(mod_json_object_rbegin(obj_)); + } + return const_reverse_iterator(); + } + + //! Retrieve an iterator pointing to the past-the-end element + iterator end(void) { + if (copy_and_leak()) { + return iterator(mod_json_object_end(obj_)); + } + return iterator(); + } + + //! Retrieve a const iterator pointing to the past-the-end element + const_iterator end(void) const { + if (obj_) { + return const_iterator(mod_json_object_end(obj_)); + } + return const_iterator(); + } + + //! Retrieve a const iterator pointing to the past-the-end element + const_iterator cend(void) const { + if (obj_) { + return const_iterator(mod_json_object_end(obj_)); + } + return const_iterator(); + } + + //! Retrieve a reverse pointing to the past-the-end element + reverse_iterator rend(void) { + if (copy_and_leak()) { + return reverse_iterator(mod_json_object_rend(obj_)); + } + return reverse_iterator(); + } + + //! Retrieve a const reverse pointing to the past-the-end element + const_reverse_iterator rend(void) const { + if (obj_) { + return const_reverse_iterator(mod_json_object_rend(obj_)); + } + return const_reverse_iterator(); + } + + //! Retrieve a const reverse pointing to the past-the-end element + const_reverse_iterator crend(void) const { + if (obj_) { + return const_reverse_iterator(mod_json_object_rend(obj_)); + } + return const_reverse_iterator(); + } + + protected: + //! Clone the object for writing + bool copy_on_write(void) { + if (obj_) { + if (mod_json_object_is_shared(obj_)) { + mod_json_object_put(obj_); + obj_ = mod_json_object_clone(obj_); + } + } else { + obj_ = mod_json_object_set_default(); + } + return (obj_ != 0); + } + + //! Clone the object and leak it + bool copy_and_leak(void) { + if (copy_on_write()) { + mod_json_object_set_leaked(obj_); + return true; + } + return false; + } + + private: + mod_json_object_t *obj_; +}; + +//! Assign new contents to the value, replacing its current content +inline void JsonValue::assign(const JsonArray &arr) { + this->set_value(arr); +} + +//! Assign new contents to the value, replacing its current content +inline void JsonValue::assign(const JsonObject &obj) { + this->set_value(obj); +} + +//! Convert value to JSON object +inline JsonObject &JsonValue::to_object(void) { + return *reinterpret_cast(&val_->data.c_obj); +} + +//! Convert value to JSON object +inline const JsonObject &JsonValue::to_object(void) const { + return *reinterpret_cast(&val_->data.c_obj); +} + +//! Convert value to JSON array +inline JsonArray &JsonValue::to_array(void) { + return *reinterpret_cast(&val_->data.c_arr); +} + +//! Convert value to JSON array +inline const JsonArray &JsonValue::to_array(void) const { + return *reinterpret_cast(&val_->data.c_arr); +} + +//! Convert value to JSON string +inline JsonString &JsonValue::to_string(void) { + return *reinterpret_cast(&val_->data.c_str); +} + +//! Convert value to JSON string +inline const JsonString &JsonValue::to_string(void) const { + return *reinterpret_cast(&val_->data.c_str); +} + +//! Treat self value as object by force, retrieving value of a key +inline JsonValue &JsonValue::get_value(const char *key) { + if (!is_object()) { + *this = JsonObject(); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonValue::get_value"); + } + return (to_object())[key]; +} + +//! Retrieve a reference of value by a key +inline JsonValue JsonValue::get_value(const char *key) const { + return (is_object() ? (to_object())[key] : JsonValue()); +} + +//! Treat self value as array by force, retrieving value at index n +inline JsonValue &JsonValue::get_value(size_type n) { + if (!is_array()) { + throw std::logic_error("JsonValue::get_value"); + } + if (!copy_and_leak()) { + throw std::runtime_error("JsonValue::get_value"); + } + return (to_array())[n]; +} + +//! Retrieve a reference of value at index n +inline JsonValue JsonValue::get_value(size_type n) const { + return (is_array() ? (to_array())[n] : JsonValue()); +} + +//! Set the new array to the value, replacing its current content +inline void JsonValue::set_value(const JsonArray &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::set_value"); + } + mod_json_value_assign_array(val_, *(mod_json_array_t **)&val); +} + +//! Set the new object to the value, replacing its current content +inline void JsonValue::set_value(const JsonObject &val) { + if (!copy_on_write()) { + throw std::runtime_error("JsonValue::set_value"); + } + mod_json_value_assign_object(val_, *(mod_json_object_t **)&val); +} + +/*! JSON Parser + */ +class JsonParser { + public: + typedef mod_json_size_t size_type; + + //! Constructor + JsonParser(void) + : state_(mod_json_state_null), error_(mod_json_error_null), context_(0) { + option_.options = 0; + option_.object_depth = 0; + option_.array_depth = 0; + } + + //! Destructor + ~JsonParser(void) {} + + //! Set the max object depth + void set_object_depth(size_type depth) { + option_.object_depth = depth; + } + + //! Set the max array depth + void set_array_depth(size_type depth) { + option_.array_depth = depth; + } + + //! Enable/Disable comments + void set_comment(bool enable = true) { + if (enable) { + option_.options |= MOD_JSON_COMMENT; + } else { + option_.options &= ~MOD_JSON_COMMENT; + } + } + + //! Enable/Disable loose strings + void set_unstrict(bool enable = true) { + if (enable) { + option_.options |= MOD_JSON_UNSTRICT; + } else { + option_.options &= ~MOD_JSON_UNSTRICT; + } + } + + //! Enable/Disable simple format + void set_simple(bool enable = true) { + if (enable) { + option_.options |= MOD_JSON_SIMPLE; + } else { + option_.options &= ~MOD_JSON_SIMPLE; + } + } + + //! Enable/Disable single quotes support + void set_squote(bool enable = true) { + if (enable) { + option_.options |= MOD_JSON_SQUOTE; + } else { + option_.options &= ~MOD_JSON_SQUOTE; + } + } + + //! Convert a sting to a JSON value + bool parse(const char *str, JsonValue *out) { + mod_json_token_t *tok; + + state_ = mod_json_state_null; + error_ = mod_json_error_null; + context_ = str; + + tok = mod_json_token_create(&option_); + if (tok) { + mod_json_value_t *jval; + + jval = mod_json_parse(tok, str); + + /* save information of token */ + state_ = mod_json_token_state(tok); + error_ = mod_json_token_error(tok); + context_ = mod_json_token_context(tok); + mod_json_token_destroy(tok); + + if (jval) { + *out = *reinterpret_cast(&jval); + mod_json_value_unset(jval); + + return out->is_valid(); + } + } + return false; + } + + //! Retrieve the error code of parser + int error(void) const { + return (int)error_; + } + + //! Retrieve the state code of parser + int state(void) const { + return (int)state_; + } + + //! Retrieve the context of parser + const char *context(void) const { + return context_; + } + + private: + mod_json_option_t option_; + mod_json_state_t state_; + mod_json_error_t error_; + mod_json_cchar_t *context_; +}; + +/*! JSON Dumper + */ +class JsonDumper { + public: + //! Constructor + JsonDumper(void) : str_() {} + + //! Destructor + ~JsonDumper(void) {} + + //! Dump a JSON value to string + bool dump(const JsonValue &val) { + mod_json_string_t *str; + + str = mod_json_dump(*((mod_json_value_t **)&val)); + str_ = *reinterpret_cast(&str); + if (str) { + mod_json_string_unset(str); + return true; + } + return false; + } + + //! Retrieve result of dumper + JsonString &result(void) { + return str_; + } + + //! Retrieve result of dumper + const JsonString &result(void) const { + return str_; + } + + private: + JsonString str_; +}; + +//! Equality +static inline bool operator==(const ailego::JsonString &lhs, const char *rhs) { + const char *self = lhs.c_str(); + if (self == rhs) { + return true; + } + + if (self && rhs) { + return (std::strcmp(self, rhs) == 0); + } + return false; +} + +//! Equality +static inline bool operator==(const char *lhs, const ailego::JsonString &rhs) { + return (rhs == lhs); +} + +//! Equality +static inline bool operator==(const ailego::JsonString &lhs, + const std::string &rhs) { + std::size_t ls = lhs.size(); + std::size_t rs = rhs.size(); + if (ls == 0 && rs == 0) { + return true; + } + + if (ls == rs) { + const char *ld = lhs.data(); + const char *rd = rhs.data(); + + if (ld && rd) { + return (std::memcmp(ld, rd, ls) == 0); + } + } + return false; +} + +//! Equality +static inline bool operator==(const std::string &lhs, const JsonString &rhs) { + return (rhs == lhs); +} + +//! Equality +static inline bool operator==(const JsonString &lhs, const JsonValue &rhs) { + return (rhs.is_string() ? lhs == rhs.as_string() : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const JsonString &rhs) { + return (lhs.is_string() ? lhs.as_string() == rhs : false); +} + +//! Equality +static inline bool operator==(const JsonArray &lhs, const JsonValue &rhs) { + return (rhs.is_array() ? lhs == rhs.as_array() : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const JsonArray &rhs) { + return (lhs.is_array() ? lhs.as_array() == rhs : false); +} + +//! Equality +static inline bool operator==(const JsonObject &lhs, const JsonValue &rhs) { + return (rhs.is_object() ? lhs == rhs.as_object() : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const JsonObject &rhs) { + return (lhs.is_object() ? lhs.as_object() == rhs : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const bool &rhs) { + return (lhs.is_boolean() ? lhs.as_bool() == rhs : false); +} + +//! Equality +static inline bool operator==(const bool &lhs, const JsonValue &rhs) { + return (rhs.is_boolean() ? lhs == rhs.as_bool() : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const signed char &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const signed char &lhs, const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const char &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const char &lhs, const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const short int &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const short int &lhs, const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const int &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const int &lhs, const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const long int &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const long int &lhs, const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const long long int &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const long long int &lhs, const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const float &rhs) { + if (lhs.is_float()) { + double diff = static_cast(lhs.as_float() - rhs); + return ((diff < DBL_EPSILON) && (diff > -DBL_EPSILON)); + } + return false; +} + +//! Equality +static inline bool operator==(const float &lhs, const JsonValue &rhs) { + if (rhs.is_float()) { + double diff = static_cast(rhs.as_float() - lhs); + return ((diff < DBL_EPSILON) && (diff > -DBL_EPSILON)); + } + return false; +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const double &rhs) { + if (lhs.is_float()) { + double diff = static_cast(lhs.as_float() - rhs); + return ((diff < DBL_EPSILON) && (diff > -DBL_EPSILON)); + } + return false; +} + +//! Equality +static inline bool operator==(const double &lhs, const JsonValue &rhs) { + if (rhs.is_float()) { + double diff = static_cast(rhs.as_float() - lhs); + return ((diff < DBL_EPSILON) && (diff > -DBL_EPSILON)); + } + return false; +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const long double &rhs) { + if (lhs.is_float()) { + double diff = static_cast(lhs.as_float() - rhs); + return ((diff < DBL_EPSILON) && (diff > -DBL_EPSILON)); + } + return false; +} + +//! Equality +static inline bool operator==(const long double &lhs, const JsonValue &rhs) { + if (rhs.is_float()) { + double diff = static_cast(rhs.as_float() - lhs); + return ((diff < DBL_EPSILON) && (diff > -DBL_EPSILON)); + } + return false; +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const unsigned char &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const unsigned char &lhs, const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, + const unsigned short int &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const unsigned short int &lhs, + const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const unsigned int &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const unsigned int &lhs, const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, + const unsigned long int &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const unsigned long int &lhs, + const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, + const unsigned long long int &rhs) { + return (lhs.is_integer() + ? lhs.as_integer() == static_cast(rhs) + : false); +} + +//! Equality +static inline bool operator==(const unsigned long long int &lhs, + const JsonValue &rhs) { + return (rhs.is_integer() + ? static_cast(lhs) == rhs.as_integer() + : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const char *rhs) { + return (lhs.is_string() ? lhs.as_string() == rhs : false); +} + +//! Equality +static inline bool operator==(const char *lhs, const JsonValue &rhs) { + return (rhs.is_string() ? lhs == rhs.as_string() : false); +} + +//! Equality +static inline bool operator==(const JsonValue &lhs, const std::string &rhs) { + return (lhs.is_string() ? lhs.as_string() == rhs : false); +} + +//! Equality +static inline bool operator==(const std::string &lhs, const JsonValue &rhs) { + return (rhs.is_string() ? lhs == rhs.as_string() : false); +} + +//! No equality +static inline bool operator!=(const JsonString &lhs, const char *rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const char *lhs, const JsonString &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonString &lhs, const std::string &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const std::string &lhs, const JsonString &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonString &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const JsonString &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonArray &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const JsonArray &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonObject &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const JsonObject &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const bool &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const bool &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const signed char &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const signed char &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const char &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const char &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const short int &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const short int &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const int &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const int &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const long int &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const long int &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const long long int &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const long long int &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const float &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const float &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const double &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const double &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const long double &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const long double &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const unsigned char &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const unsigned char &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, + const unsigned short int &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const unsigned short int &lhs, + const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const unsigned int &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const unsigned int &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, + const unsigned long int &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const unsigned long int &lhs, + const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, + const unsigned long long int &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const unsigned long long int &lhs, + const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const char *rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const char *lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const JsonValue &lhs, const std::string &rhs) { + return !(lhs == rhs); +} + +//! No equality +static inline bool operator!=(const std::string &lhs, const JsonValue &rhs) { + return !(lhs == rhs); +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/hash/crc32c.cc b/src/ailego/hash/crc32c.cc new file mode 100644 index 0000000..20157df --- /dev/null +++ b/src/ailego/hash/crc32c.cc @@ -0,0 +1,636 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "crc32c.h" +#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/hash/crc32c.h b/src/ailego/hash/crc32c.h new file mode 100644 index 0000000..c365af3 --- /dev/null +++ b/src/ailego/hash/crc32c.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 + +namespace zvec { +namespace ailego { + +/*! Crc32c Hash + */ +struct Crc32c { + //! Compute the CRC32C checksum for the source data buffer + static uint32_t Hash(const void *data, size_t len, uint32_t crc); + + //! Compute the CRC32C checksum for the source data buffer + static inline uint32_t Hash(const void *data, size_t len) { + return Hash(data, len, 0u); + } +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/hash/jump_hash.h b/src/ailego/hash/jump_hash.h new file mode 100644 index 0000000..c958938 --- /dev/null +++ b/src/ailego/hash/jump_hash.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 + +namespace zvec { +namespace ailego { + +//! Jump consistent hash algorithm (https://arxiv.org/pdf/1406.2294.pdf) +static inline int32_t JumpHash(uint64_t key, int32_t num_buckets) { + int64_t b = 1, j = 0; + while (j < num_buckets) { + b = j; + key = key * 2862933555777941757ULL + 1; + j = (int64_t)(double(b + 1) * + (double(1LL << 31) / double((key >> 33) + 1))); + } + return (int32_t)b; +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/internal/cpu_features.cc b/src/ailego/internal/cpu_features.cc new file mode 100644 index 0000000..7b16e3a --- /dev/null +++ b/src/ailego/internal/cpu_features.cc @@ -0,0 +1,421 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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" + +#if !defined(_MSC_VER) && !defined(__ARM_ARCH) +#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) +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(__ARM_ARCH) +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, NULL) >= 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 + ; +} + +} // 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..eaa39a6 --- /dev/null +++ b/src/ailego/internal/cpu_features.h @@ -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. + +#pragma once + +#include "platform.h" + +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_; +}; + +} // namespace internal +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/internal/platform.h b/src/ailego/internal/platform.h new file mode 100644 index 0000000..ecfe673 --- /dev/null +++ b/src/ailego/internal/platform.h @@ -0,0 +1,488 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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(_WIN32) || defined(_WIN64) +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_MSC_VER) +#include +#else +#include +#include +#if defined(__x86_64__) || defined(__i386) +#include +#endif +#if defined(__ARM_NEON) +#include +#endif +#if defined(__ARM_FEATURE_CRC32) +#include +#endif +#endif + +#if defined(__cplusplus) +extern "C" { +#endif + +#ifndef NDEBUG +#define AILEGO_DEBUG +#endif + +//! Fixed Intel intrinsics macro in MSVC +#if defined(_MSC_VER) +#if (_M_IX86_FP == 2 || defined(_M_AMD64) || defined(_M_X64)) +#define __SSE__ 1 +#define __SSE2__ 1 +#if _MSC_VER >= 1500 +#define __SSE3__ 1 +#define __SSSE3__ 1 +#define __SSE4_1__ 1 +#define __SSE4_2__ 1 +#endif +#elif _M_IX86_FP == 1 +#define __SSE__ 1 +#endif +#endif // _MSC_VER + +#if defined(_WIN32) || defined(_WIN64) +#if defined(_WIN64) +#define AILEGO_M64 +#else +#define AILEGO_M32 +#endif +#endif + +#if defined(__GNUC__) +#if defined(__x86_64__) || defined(__aarch64__) || defined(__ppc64__) +#define AILEGO_M64 +#else +#define AILEGO_M32 +#endif +#endif + +#ifndef AILEGO_ALIGNED +#if defined(_MSC_VER) +#define AILEGO_ALIGNED(x) __declspec(align(x)) +#define AILEGO_DEPRECATED __declspec(deprecated) +#elif defined(__GNUC__) +#define AILEGO_ALIGNED(x) __attribute__((aligned(x))) +#define AILEGO_DEPRECATED __attribute__((deprecated)) +#else +#define AILEGO_ALIGNED(x) +#define AILEGO_DEPRECATED +#endif +#endif + +//! Add 'inline' for MSVC +#if defined(_MSC_VER) && !defined(__cplusplus) +#if !defined(inline) +#define inline __inline +#endif +#endif + +//! Add 'ssize_t' for MSVC +#if defined(_MSC_VER) +typedef intptr_t ssize_t; +#endif + +#if defined(_MSC_VER) +//! Returns the number of trailing 0-bits in x +static inline int ailego_ctz32(uint32_t x) { + unsigned long r = 0; + _BitScanForward(&r, x); + return (int)r; +} + +//! Returns the number of leading 0-bits in x +static inline int ailego_clz32(uint32_t x) { + unsigned long r = 0; + _BitScanReverse(&r, x); + return (31 - (int)r); +} + +#if defined(AILEGO_M64) +//! Returns the number of trailing 0-bits in x +static inline int ailego_ctz64(uint64_t x) { + unsigned long r = 0; + _BitScanForward64(&r, x); + return (int)r; +} + +//! Returns the number of leading 0-bits in x +static inline int ailego_clz64(uint64_t x) { + unsigned long r = 0; + _BitScanReverse64(&r, x); + return (63 - (int)r); +} +#else +//! Returns the number of trailing 0-bits in x +static inline int ailego_ctz64(uint64_t x) { + unsigned long r = 0; + unsigned long m = (unsigned long)x; + _BitScanForward(&r, m); + if (r == 0) { + m = (unsigned long)(x >> 32); + _BitScanForward(&r, m); + if (r != 0) { + r += 32; + } + } + return (int)r; +} + +//! Returns the number of leading 0-bits in x +static inline int ailego_clz64(uint64_t x) { + unsigned long r = 0; + unsigned long m = (unsigned long)(x >> 32); + _BitScanReverse(&r, m); + if (r != 0) { + return (31 - (int)r); + } + m = (unsigned long)x; + _BitScanReverse(&r, m); + return (63 - (int)r); +} +#endif // AILEGO_M64 + +//! Counts the number of one bits +#define ailego_popcount32(x) (__popcnt(x)) +#define ailego_popcount64(x) (__popcnt64(x)) +#define ailego_likely(x) (x) +#define ailego_unlikely(x) (x) +#ifdef __SSE__ +#define ailego_prefetch(p) _mm_prefetch((p), 0) +#else +#define ailego_prefetch(p) ((void)(p)) +#endif +#else // !_MSC_VER +#define ailego_ctz32(x) (__builtin_ctz(x)) +#define ailego_ctz64(x) (__builtin_ctzll(x)) +#define ailego_clz32(x) (__builtin_clz(x)) +#define ailego_clz64(x) (__builtin_clzll(x)) +#define ailego_popcount32(x) (__builtin_popcount(x)) +#define ailego_popcount64(x) (__builtin_popcountl(x)) +#define ailego_likely(x) (__builtin_expect(!!(x), 1)) +#define ailego_unlikely(x) (__builtin_expect(!!(x), 0)) +#define ailego_prefetch(p) (__builtin_prefetch((p))) +#endif // _MSC_VER + +#if defined(AILEGO_M64) +#define ailego_ctz ailego_ctz64 +#define ailego_clz ailego_clz64 +#define ailego_popcount ailego_popcount64 +#else +#define ailego_ctz ailego_ctz32 +#define ailego_clz ailego_clz32 +#define ailego_popcount ailego_popcount32 +#endif // AILEGO_M64 + +#if defined(__arm__) || defined(__aarch64__) +// ARMv7 Architecture Reference Manual (for YIELD) +// ARM Compiler toolchain Compiler Reference (for __yield() instrinsic) +#if defined(__CC_ARM) +#define ailego_yield() __yield() +#else +#define ailego_yield() __asm__ __volatile__("yield") +#endif // __CC_ARM +#elif defined(__SSE2__) +#define ailego_yield() _mm_pause() +#else +#define ailego_yield() ((void)0) +#endif // __arm__ || __aarch64__ + +#if defined(_MSC_VER) +#define ailego_aligned_malloc(SIZE, ALIGN) \ + _aligned_malloc((size_t)(SIZE), (ALIGN)) +#define ailego_aligned_free _aligned_free +#else // !_MSC_VER +#if defined(_ISOC11_SOURCE) +#define ailego_aligned_malloc(SIZE, ALIGN) \ + aligned_alloc((ALIGN), (size_t)(SIZE)) +#else // !_ISOC11_SOURCE +#define ailego_aligned_malloc(SIZE, ALIGN) \ + ailego_posix_malloc((size_t)(SIZE), (ALIGN)) +#endif // _ISOC11_SOURCE +#define ailego_aligned_free free +#endif // _MSC_VER + +#if !defined(__SANITIZE_ADDRESS__) +#if defined(__has_feature) +#if __has_feature(address_sanitizer) +#define __SANITIZE_ADDRESS__ 1 +#endif // address_sanitizer +#endif // __has_feature +#endif // !__SANITIZE_ADDRESS__ + +#if !defined(__SANITIZE_ADDRESS__) +#if !defined(ailego_malloc) +#if defined(__AVX512F__) +#define ailego_malloc(SIZE) ailego_aligned_malloc((SIZE), 64) +#elif defined(__AVX__) +#define ailego_malloc(SIZE) ailego_aligned_malloc((SIZE), 32) +#elif defined(__SSE__) +#define ailego_malloc(SIZE) ailego_aligned_malloc((SIZE), 16) +#elif defined(__ARM_NEON) +#define ailego_malloc(SIZE) ailego_aligned_malloc((SIZE), 16) +#endif +#endif // !ailego_malloc +#if (defined(__SSE__) || defined(__ARM_NEON)) && !defined(ailego_free) +#define ailego_free ailego_aligned_free +#endif +#endif // !__SANITIZE_ADDRESS__ + +#ifndef ailego_malloc +#define ailego_malloc(SIZE) malloc((size_t)(SIZE)) +#endif +#ifndef ailego_free +#define ailego_free free +#endif + +#ifndef ailego_offsetof +#define ailego_offsetof(TYPE, MEMBER) ((size_t) & ((TYPE *)0)->MEMBER) +#endif + +#ifndef ailego_align +#define ailego_align(SIZE, BOUND) (((SIZE) + ((BOUND) - 1)) & ~((BOUND) - 1)) +#endif + +#ifndef ailego_align8 +#define ailego_align8(SIZE) ailego_align(SIZE, 8) +#endif + +#ifndef ailego_min +#define ailego_min(A, B) (((A) < (B)) ? (A) : (B)) +#endif + +#ifndef ailego_max +#define ailego_max(A, B) (((A) > (B)) ? (A) : (B)) +#endif + +#ifndef ailego_malloc_object +#define ailego_malloc_object(TYPE) ((TYPE *)ailego_malloc(sizeof(TYPE))) +#endif +#ifndef ailego_malloc_array +#define ailego_malloc_array(TYPE, SIZE) \ + ((TYPE *)ailego_malloc(SIZE * sizeof(TYPE))) +#endif + +#ifndef ailego_minus_if_ne_zero +#define ailego_minus_if_ne_zero(COND) \ + if (ailego_unlikely((COND) != 0)) return (-1) +#endif + +#ifndef ailego_zero_if_ne_zero +#define ailego_zero_if_ne_zero(COND) \ + if (ailego_unlikely((COND) != 0)) return (0) +#endif + +#ifndef ailego_null_if_ne_zero +#define ailego_null_if_ne_zero(COND) \ + if (ailego_unlikely((COND) != 0)) return (NULL) +#endif + +#ifndef ailego_false_if_ne_zero +#define ailego_false_if_ne_zero(COND) \ + if (ailego_unlikely((COND) != 0)) return (false) +#endif + +#ifndef ailego_return_if_ne_zero +#define ailego_return_if_ne_zero(COND) \ + if (ailego_unlikely((COND) != 0)) return +#endif + +#ifndef ailego_break_if_ne_zero +#define ailego_break_if_ne_zero(COND) \ + if (ailego_unlikely((COND) != 0)) break +#endif + +#ifndef ailego_continue_if_ne_zero +#define ailego_continue_if_ne_zero(COND) \ + if (ailego_unlikely((COND) != 0)) continue +#endif + +#ifndef ailego_do_if_ne_zero +#define ailego_do_if_ne_zero(COND) if (ailego_unlikely((COND) != 0)) +#endif + +#ifndef ailego_minus_if_lt_zero +#define ailego_minus_if_lt_zero(COND) \ + if (ailego_unlikely((COND) < 0)) return (-1) +#endif + +#ifndef ailego_zero_if_lt_zero +#define ailego_zero_if_lt_zero(COND) \ + if (ailego_unlikely((COND) < 0)) return (0) +#endif + +#ifndef ailego_null_if_lt_zero +#define ailego_null_if_lt_zero(COND) \ + if (ailego_unlikely((COND) < 0)) return (NULL) +#endif + +#ifndef ailego_false_if_lt_zero +#define ailego_false_if_lt_zero(COND) \ + if (ailego_unlikely((COND) < 0)) return (false) +#endif + +#ifndef ailego_return_if_lt_zero +#define ailego_return_if_lt_zero(COND) \ + if (ailego_unlikely((COND) < 0)) return +#endif + +#ifndef ailego_break_if_lt_zero +#define ailego_break_if_lt_zero(COND) \ + if (ailego_unlikely((COND) < 0)) break +#endif + +#ifndef ailego_continue_if_lt_zero +#define ailego_continue_if_lt_zero(COND) \ + if (ailego_unlikely((COND) < 0)) continue +#endif + +#ifndef ailego_do_if_lt_zero +#define ailego_do_if_lt_zero(COND) if (ailego_unlikely((COND) < 0)) +#endif + +#ifndef ailego_minus_if_false +#define ailego_minus_if_false(COND) \ + if (ailego_unlikely(!(COND))) return (-1) +#endif + +#ifndef ailego_zero_if_false +#define ailego_zero_if_false(COND) \ + if (ailego_unlikely(!(COND))) return (0) +#endif + +#ifndef ailego_null_if_false +#define ailego_null_if_false(COND) \ + if (ailego_unlikely(!(COND))) return (NULL) +#endif + +#ifndef ailego_false_if_false +#define ailego_false_if_false(COND) \ + if (ailego_unlikely(!(COND))) return (false) +#endif + +#ifndef ailego_return_if_false +#define ailego_return_if_false(COND) \ + if (ailego_unlikely(!(COND))) return +#endif + +#ifndef ailego_break_if_false +#define ailego_break_if_false(COND) \ + if (ailego_unlikely(!(COND))) break +#endif + +#ifndef ailego_continue_if_false +#define ailego_continue_if_false(COND) \ + if (ailego_unlikely(!(COND))) continue +#endif + +#ifndef ailego_do_if_false +#define ailego_do_if_false(COND) if (ailego_unlikely(!(COND))) +#endif + +#ifndef ailego_compile_assert +#define ailego_compile_assert(COND, MSG) \ + typedef char Static_Assertion_##MSG[(!!(COND)) * 2 - 1] +#endif + +#ifndef ailego_static_assert3 +#define ailego_static_assert3(COND, LINE) \ + ailego_compile_assert(COND, At_Line_##LINE) +#endif + +#ifndef ailego_static_assert2 +#define ailego_static_assert2(COND, LINE) ailego_static_assert3(COND, LINE) +#endif + +#ifndef ailego_static_assert +#define ailego_static_assert(COND) ailego_static_assert2(COND, __LINE__) +#endif + +//! Abort and report if an assertion is failed +#ifndef ailego_assert_abort +#define ailego_assert_abort(COND, MSG) \ + (void)(ailego_likely(COND) || (ailego_assert_report(__FILE__, __FUNCTION__, \ + __LINE__, #COND, (MSG)), \ + abort(), 0)) +#endif + +#ifdef AILEGO_DEBUG +#ifndef ailego_assert +#define ailego_assert(COND) ailego_assert_abort(COND, "") +#endif +#ifndef ailego_assert_with +#define ailego_assert_with(COND, MSG) ailego_assert_abort(COND, MSG) +#endif +#else // !AILEGO_DEBUG +#ifndef ailego_assert +#define ailego_assert(COND) ((void)0) +#endif +#ifndef ailego_assert_with +#define ailego_assert_with(COND, MSG) ((void)0) +#endif +#endif // AILEGO_DEBUG + +#ifndef ailego_check +#define ailego_check(COND) ailego_assert_abort(COND, "") +#endif +#ifndef ailego_check_with +#define ailego_check_with(COND, MSG) ailego_assert_abort(COND, MSG) +#endif + +#ifndef _MSC_VER +//! Allocates memory on a specified alignment boundary +static inline void *ailego_posix_malloc(size_t size, size_t align) { + void *ptr; + ailego_null_if_ne_zero(posix_memalign(&ptr, align, size)); + return ptr; +} +#endif + +//! Report an assertion is failed +static inline void ailego_assert_report(const char *file, const char *func, + int line, const char *cond, + const char *msg) { + fprintf(stderr, "Assertion failed: (%s) in %s(), %s line %d. %s\n", cond, + func, file, line, msg); +} + +// because the behavior of conversion from negative float to unsigned integer +// is undefined (on arm, result will be zero), it's necessary to convert it +// to signed integer firstly +static inline uint8_t static_cast_from_float_to_uint8(float data) { + return static_cast(static_cast(data)); +} + +static inline uint16_t static_cast_from_float_to_uint16(float data) { + return static_cast(static_cast(data)); +} + +#ifdef __cplusplus +} /* extern "C" */ +#endif diff --git a/src/ailego/io/file.cc b/src/ailego/io/file.cc new file mode 100644 index 0000000..bfdefc8 --- /dev/null +++ b/src/ailego/io/file.cc @@ -0,0 +1,726 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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.h" +#if !defined(_WIN64) && !defined(_WIN32) +#include +#include +#include +#include +#include +#include +#include +#else +#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); + +#if defined(MAP_POPULATE) + if (opts & File::MMAP_POPULATE) { + prot |= MAP_POPULATE; + } +#endif + int flags = (opts & File::MMAP_SHARED) ? MAP_SHARED : MAP_PRIVATE; +#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); + +#if defined(MAP_POPULATE) + if (opts & File::MMAP_POPULATE) { + prot |= MAP_POPULATE; + } +#endif + int flags = (opts & File::MMAP_SHARED) ? MAP_SHARED | MAP_ANONYMOUS + : MAP_PRIVATE | MAP_ANONYMOUS; +#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 + +//! Create a local file +bool File::create(const char *path, size_t len, bool direct) { + ailego_false_if_false(native_handle_ == File::InvalidHandle && path); + + // Try opening or creating the file + HANDLE file_handle = + CreateFileA(path, GENERIC_WRITE | GENERIC_READ, FILE_SHARE_READ, 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 = CreateFileA( + path, GENERIC_WRITE | GENERIC_READ, FILE_SHARE_READ, 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); + + // Try opening the file + DWORD flags = FILE_ATTRIBUTE_NORMAL; + if (direct) { + flags |= FILE_FLAG_NO_BUFFERING; + } + HANDLE file_handle = + CreateFileA(path, (rdonly ? GENERIC_READ : GENERIC_READ | GENERIC_WRITE), + FILE_SHARE_READ, 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) { + LARGE_INTEGER file_size; + file_size.QuadPart = len; + + // Create map object + HANDLE file_mapping = CreateFileMapping( + handle, nullptr, + ((opts & File::MMAP_READONLY) ? PAGE_READONLY : PAGE_READWRITE), + file_size.HighPart, file_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; + } + file_size.QuadPart = off; + + // Map the whole file to memory and close handle + void *addr = MapViewOfFile(file_mapping, desired_access, file_size.HighPart, + file_size.LowPart, 0); + CloseHandle(file_mapping); + + 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::MemoryMap(size_t, int) { + return nullptr; +} + +void *File::MemoryRemap(void *, size_t, void *, size_t) { + return nullptr; +} + +void File::MemoryUnmap(void *addr, size_t /*len*/) { + ailego_return_if_false(addr); + UnmapViewOfFile(addr); +} + +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 \ No newline at end of file diff --git a/src/ailego/io/file.h b/src/ailego/io/file.h new file mode 100644 index 0000000..54f6b6d --- /dev/null +++ b/src/ailego/io/file.h @@ -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. + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +/*! File Utility + */ +class File { + public: + //! Native Handle in OS + typedef FileHelper::NativeHandle NativeHandle; + + //! Invalid Handle + static constexpr NativeHandle InvalidHandle = (NativeHandle)(-1); + + //! Specifies the position in a file to use for seeking. + enum struct Origin { Begin = 0, Current = 1, End = 2 }; + + //! Options of memory mapping + enum { + MMAP_READONLY = 1, + MMAP_SHARED = 2, + MMAP_LOCKED = 4, + MMAP_WARMUP = 8, + MMAP_POPULATE = 16, + MMAP_HUGE_PAGE = 32, + }; + + //! Constructor + File(void) : native_handle_(File::InvalidHandle), read_only_(false) {} + + //! Constructor + File(File &&rhs) { + read_only_ = rhs.read_only_; + native_handle_ = rhs.native_handle_; + rhs.read_only_ = false; + rhs.native_handle_ = File::InvalidHandle; + } + + //! Destructor + ~File(void) { + this->close(); + } + + //! Assignment + File &operator=(File &&rhs) { + read_only_ = rhs.read_only_; + native_handle_ = rhs.native_handle_; + rhs.read_only_ = false; + rhs.native_handle_ = File::InvalidHandle; + return *this; + } + + //! Test if the file is valid + bool is_valid(void) const { + return (native_handle_ != File::InvalidHandle); + } + + //! Retrieve non-zero if memory region is read only + bool read_only(void) const { + return read_only_; + } + + //! Retrieve native handle + NativeHandle native_handle(void) const { + return native_handle_; + } + + //! Create a local file + bool create(const char *path, size_t size, bool direct); + + //! Open a local file + bool open(const char *path, bool rdonly, bool direct); + + //! Close the local file + void close(void); + + //! Reset the file + void reset(void); + + //! Write data into the file + size_t write(const void *data, size_t len); + + //! Write data into the file + size_t write(ssize_t off, const void *data, size_t len); + + //! Read data from the file + size_t read(void *buf, size_t len); + + //! Read data from the file + size_t read(ssize_t off, void *buf, size_t len); + + //! Synchronize memory with physical storage + bool flush(void); + + //! Sets the current position of the file to the given value + bool seek(ssize_t off, Origin origin); + + //! Truncate the file to a specified length + bool truncate(size_t len); + + //! Retrieve size of file + size_t size(void) const; + + //! Retrieve offset of file + ssize_t offset(void) const; + + //! Create a local file + bool create(const char *path, size_t len) { + return this->create(path, len, false); + } + + //! Create a local file + bool create(const std::string &path, size_t len, bool direct) { + return this->create(path.c_str(), len, direct); + } + + //! Create a local file + bool create(const std::string &path, size_t len) { + return this->create(path.c_str(), len); + } + + //! Open a local file + bool open(const char *path, bool rdonly) { + return this->open(path, rdonly, false); + } + + //! Open a local file + bool open(const std::string &path, bool rdonly, bool direct) { + return this->open(path.c_str(), rdonly, direct); + } + + //! Open a local file + bool open(const std::string &path, bool rdonly) { + return this->open(path.c_str(), rdonly); + } + + //! Map a region of file into memory + void *map(ssize_t off, size_t len, int opts) { + if (read_only_) { + opts |= File::MMAP_READONLY; + } + return File::MemoryMap(native_handle_, off, len, opts); + } + + //! Map a region of file into memory + static void *MemoryMap(NativeHandle handle, ssize_t off, size_t len, + int opts); + + //! Map an anonymous region into memory + static void *MemoryMap(size_t len, int opts); + + //! Remap the region into memory + static void *MemoryRemap(void *oldptr, size_t oldsize, void *newptr, + size_t newsize); + + //! Unmap a mapping region + static void MemoryUnmap(void *addr, size_t len); + + //! Synchronize a memory map + static bool MemoryFlush(void *addr, size_t len); + + //! Lock the memory region into RAM + static bool MemoryLock(void *addr, size_t len); + + //! Unlock the memory region in RAM + static bool MemoryUnlock(void *addr, size_t len); + + //! Warm up a memory region + static void MemoryWarmup(void *addr, size_t len); + + //! Delete a name and possibly the file it refers to + static bool Delete(const char *path) { + return FileHelper::DeleteFile(path); + } + + //! Delete a name and possibly the file it refers to + static bool Delete(const std::string &path) { + return FileHelper::DeleteFile(path.c_str()); + } + + //! Change the name or location of a file + static bool Rename(const char *oldpath, const char *newpath) { + return FileHelper::RenameFile(oldpath, newpath); + } + + //! Change the name or location of a file + static bool Rename(const std::string &oldpath, const std::string &newpath) { + return FileHelper::RenameFile(oldpath.c_str(), newpath.c_str()); + } + + //! Retrieve the base name from a path + static const char *BaseName(const char *path) { + return FileHelper::BaseName(path); + } + + //! Retrieve the base name from a path + static const char *BaseName(const std::string &path) { + return BaseName(path.c_str()); + } + + //! Make directories' path + static bool MakePath(const char *path) { + return FileHelper::MakePath(path); + } + + //! Make directories' path + static bool MakePath(const std::string &path) { + return FileHelper::MakePath(path.c_str()); + } + + //! Remove a file or a directory (includes files & subdirectories) + static bool RemovePath(const char *path) { + return FileHelper::RemovePath(path); + } + + //! Remove a file or a directory (includes files & subdirectories) + static bool RemovePath(const std::string &path) { + return FileHelper::RemovePath(path.c_str()); + } + + //! Remove a directory (includes files & subdirectories) + static bool RemoveDirectory(const char *path) { + return FileHelper::RemoveDirectory(path); + } + + //! Remove a directory (includes files & subdirectories) + static bool RemoveDirectory(const std::string &path) { + return FileHelper::RemoveDirectory(path.c_str()); + } + + //! Retrieve non-zero if the path exists + static bool IsExist(const char *path) { + return FileHelper::IsExist(path); + } + + //! Retrieve non-zero if the path exists + static bool IsExist(const std::string &path) { + return FileHelper::IsExist(path.c_str()); + } + + //! Retrieve non-zero if the path is a regular file + static bool IsRegular(const char *path) { + return FileHelper::IsRegular(path); + } + + //! Retrieve non-zero if the path is a regular file + static bool IsRegular(const std::string &path) { + return FileHelper::IsRegular(path.c_str()); + } + + //! Retrieve non-zero if the path is a directory + static bool IsDirectory(const char *path) { + return FileHelper::IsDirectory(path); + } + + //! Retrieve non-zero if the path is a directory + static bool IsDirectory(const std::string &path) { + return FileHelper::IsDirectory(path.c_str()); + } + + //! Retrieve non-zero if the path is a symbolic link + static bool IsSymbolicLink(const char *path) { + return FileHelper::IsSymbolicLink(path); + } + + //! Retrieve non-zero if the path is a symbolic link + static bool IsSymbolicLink(const std::string &path) { + return FileHelper::IsSymbolicLink(path.c_str()); + } + + //! Retrieve non-zero if two paths are pointing to the same file + static bool IsSame(const char *path1, const char *path2) { + return FileHelper::IsSame(path1, path2); + } + + //! Retrieve non-zero if two paths are pointing to the same file + static bool IsSame(const std::string &path1, const std::string &path2) { + return FileHelper::IsSame(path1.c_str(), path2.c_str()); + } + + private: + //! Disable them + File(const File &) = delete; + File &operator=(const File &) = delete; + + //! Members + NativeHandle native_handle_; + bool read_only_; +}; + +} // 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..5e79f51 --- /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 "file.h" + +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/io/mmap_file.h b/src/ailego/io/mmap_file.h new file mode 100644 index 0000000..286f50a --- /dev/null +++ b/src/ailego/io/mmap_file.h @@ -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. + +#pragma once + +#include "file.h" + +namespace zvec { +namespace ailego { + +/*! Memory Mapping File + */ +class MMapFile { + public: + //! Constructor + MMapFile(void) + : read_only_(false), region_(nullptr), region_size_(0), offset_(0) {} + + //! Constructor + MMapFile(MMapFile &&rhs) { + read_only_ = rhs.read_only_; + region_ = rhs.region_; + region_size_ = rhs.region_size_; + offset_ = rhs.offset_; + rhs.read_only_ = false; + rhs.region_ = nullptr; + rhs.region_size_ = 0; + rhs.offset_ = 0; + } + + //! Destructor + ~MMapFile(void) { + this->close(); + } + + //! Assignment + MMapFile &operator=(MMapFile &&rhs) { + read_only_ = rhs.read_only_; + region_ = rhs.region_; + region_size_ = rhs.region_size_; + offset_ = rhs.offset_; + rhs.read_only_ = false; + rhs.region_ = nullptr; + rhs.region_size_ = 0; + rhs.offset_ = 0; + return *this; + } + + //! Test if the file is valid + bool is_valid(void) const { + return (region_ != nullptr); + } + + //! Retrieve non-zero if memory region is read only + bool read_only(void) const { + return read_only_; + } + + //! Create a memory mapping file + bool create(const char *path, size_t len) { + ailego_false_if_false(!region_ && path); + + File file; + ailego_false_if_false(file.create(path, len)); + + region_ = File::MemoryMap(file.native_handle(), 0, len, File::MMAP_SHARED); + ailego_false_if_false(region_); + + read_only_ = false; + region_size_ = len; + return true; + } + + //! Create a memory mapping file + bool create(const std::string &path, size_t len) { + return this->create(path.c_str(), len); + } + + //! Open a memory mapping file + bool open(const char *path, bool rdonly, bool shared) { + ailego_false_if_false(!region_ && path); + + File file; + ailego_false_if_false(file.open(path, rdonly, false)); + + size_t len = file.size(); + int opts = 0; + if (rdonly) { + opts |= File::MMAP_READONLY; + } + if (shared) { + opts |= File::MMAP_SHARED; + } + region_ = File::MemoryMap(file.native_handle(), 0, len, opts); + ailego_false_if_false(region_); + + read_only_ = rdonly; + region_size_ = len; + return true; + } + + //! Open a memory mapping file + bool open(const std::string &path, bool rdonly, bool shared) { + return this->open(path.c_str(), rdonly, shared); + } + + //! Open a memory mapping file + bool open(const char *path, bool rdonly) { + return this->open(path, rdonly, false); + } + + //! Open a memory mapping file + bool open(const std::string &path, bool rdonly) { + return this->open(path, rdonly, false); + } + + //! Close the memory mapping file + void close(void) { + File::MemoryUnmap(region_, region_size_); + region_ = nullptr; + region_size_ = 0; + offset_ = 0; + } + + //! Synchronize memory with physical storage + bool flush(void) { + return File::MemoryFlush(region_, region_size_); + } + + //! Lock the memory region into RAM + bool lock(void) { + return File::MemoryLock(region_, region_size_); + } + + //! Unlock the memory region in RAM + bool unlock(void) { + return File::MemoryUnlock(region_, region_size_); + } + + //! Warm up the memory region + void warmup(void) { + File::MemoryWarmup(region_, region_size_); + } + + //! Reset the file + void reset(void) { + offset_ = 0; + } + + //! Write data into the storage + size_t write(const void *data, size_t len) { + if (offset_ + len > region_size_) { + len = region_size_ - offset_; + } + memcpy((uint8_t *)region_ + offset_, data, len); + offset_ += len; + return len; + } + + //! Write data into the storage + size_t write(size_t off, const void *data, size_t len) { + if (off + len > region_size_) { + if (off > region_size_) { + off = region_size_; + } + len = region_size_ - off; + } + memcpy((uint8_t *)region_ + off, data, len); + return len; + } + + //! Read data from the storage (Zero-copy) + size_t read(const void **data, size_t len) { + if (offset_ + len > region_size_) { + len = region_size_ - offset_; + } + *data = (uint8_t *)region_ + offset_; + offset_ += len; + return len; + } + + //! Read data from the storage (Zero-copy) + size_t read(size_t off, const void **data, size_t len) { + if (off + len > region_size_) { + if (off > region_size_) { + off = region_size_; + } + len = region_size_ - off; + } + *data = (uint8_t *)region_ + off; + return len; + } + + //! Read data from the storage + size_t read(void *data, size_t len) { + if (offset_ + len > region_size_) { + len = region_size_ - offset_; + } + memcpy(data, (uint8_t *)region_ + offset_, len); + offset_ += len; + return len; + } + + //! Read data from the storage + size_t read(size_t off, void *data, size_t len) { + if (off + len > region_size_) { + if (off > region_size_) { + off = region_size_; + } + len = region_size_ - off; + } + memcpy(data, (uint8_t *)region_ + off, len); + return len; + } + + //! Retrieve memory region of file + void *region(void) const { + return region_; + } + + //! Retrieve region size of file + size_t size(void) const { + return region_size_; + } + + //! Retrieve offset of file + size_t offset(void) const { + return offset_; + } + + private: + //! Disable them + MMapFile(const MMapFile &) = delete; + MMapFile &operator=(const MMapFile &) = delete; + + //! Members + bool read_only_; + void *region_; + size_t region_size_; + size_t offset_; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/io/pid_file.cc b/src/ailego/io/pid_file.cc new file mode 100644 index 0000000..143716d --- /dev/null +++ b/src/ailego/io/pid_file.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 "pid_file.h" +#include + +namespace zvec { +namespace ailego { + +bool PidFile::open(const char *path) { + if (file_.is_valid() || !path || !(*path)) { + return false; + } + + if (!File::IsExist(path)) { + if (!file_.create(path, 0, false)) { + return false; + } + if (!FileLock::TryLock(file_.native_handle())) { + file_.close(); + return false; + } + } else { + if (!file_.open(path, false, false)) { + return false; + } + if (!FileLock::TryLock(file_.native_handle())) { + file_.close(); + return false; + } + file_.truncate(0); + } + + std::string str(std::to_string(ProcessHelper::SelfPid())); + file_.write(str.data(), str.size()); + return true; +} + +void PidFile::close(void) { + if (file_.is_valid()) { + FileLock::Unlock(file_.native_handle()); + + std::string path; + FileHelper::GetFilePath(file_.native_handle(), &path); + file_.close(); + + if (!path.empty()) { + File::Delete(path); + } + } +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/io/pid_file.h b/src/ailego/io/pid_file.h new file mode 100644 index 0000000..fb9f3ba --- /dev/null +++ b/src/ailego/io/pid_file.h @@ -0,0 +1,64 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "file_lock.h" + +namespace zvec { +namespace ailego { + +/*! Process ID File + */ +class PidFile { + public: + //! Constructor + PidFile(void) {} + + //! Constructor + PidFile(PidFile &&rhs) : file_(std::move(rhs.file_)) {} + + //! Assignment + PidFile &operator=(PidFile &&rhs) { + file_ = std::move(rhs.file_); + return *this; + } + + //! Test if the pid file is valid + bool is_valid(void) const { + return file_.is_valid(); + } + + //! Open a local file + bool open(const char *path); + + //! Open a local file + bool open(const std::string &path) { + return this->open(path.c_str()); + } + + //! Close the pid file + void close(void); + + private: + //! Disable them + PidFile(const PidFile &) = delete; + PidFile &operator=(const PidFile &) = 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..69bf2ba --- /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 "logger.h" +#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/logger/logger.h b/src/ailego/logger/logger.h new file mode 100644 index 0000000..bf921f4 --- /dev/null +++ b/src/ailego/logger/logger.h @@ -0,0 +1,175 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "ailego/container/params.h" +#include "ailego/pattern/factory.h" + +//! Register Index Logger +#define FACTORY_REGISTER_LOGGER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, zvec::ailego::Logger, __IMPL__, \ + ##__VA_ARGS__) + +//! Register Index Logger +#define FACTORY_REGISTER_LOGGER(__IMPL__, ...) \ + FACTORY_REGISTER_LOGGER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +#define PROXIMA_LOG_IMPL(level, format, ...) \ + do { \ + if (zvec::ailego::LoggerBroker::IsLevelEnabled(level)) { \ + zvec::ailego::LoggerBroker::Log(level, __FILE__, __LINE__, format, \ + ##__VA_ARGS__); \ + } \ + } while (0) + +//! Log Debug Message +#ifndef LOG_DEBUG +#define LOG_DEBUG(format, ...) \ + PROXIMA_LOG_IMPL(zvec::ailego::Logger::LEVEL_DEBUG, format, ##__VA_ARGS__) +#endif + +//! Log Information Message +#ifndef LOG_INFO +#define LOG_INFO(format, ...) \ + PROXIMA_LOG_IMPL(zvec::ailego::Logger::LEVEL_INFO, format, ##__VA_ARGS__) +#endif + +//! Log Warn Message +#ifndef LOG_WARN +#define LOG_WARN(format, ...) \ + PROXIMA_LOG_IMPL(zvec::ailego::Logger::LEVEL_WARN, format, ##__VA_ARGS__) +#endif + +//! Log Error Message +#ifndef LOG_ERROR +#define LOG_ERROR(format, ...) \ + PROXIMA_LOG_IMPL(zvec::ailego::Logger::LEVEL_ERROR, format, ##__VA_ARGS__) +#endif + +//! Log Fatal Message +#ifndef LOG_FATAL +#define LOG_FATAL(format, ...) \ + PROXIMA_LOG_IMPL(zvec::ailego::Logger::LEVEL_FATAL, format, ##__VA_ARGS__) +#endif + +namespace zvec { +namespace ailego { + +/*! Index Logger + */ +struct Logger { + //! Index Logger Pointer + typedef std::shared_ptr Pointer; + + static const int LEVEL_DEBUG; + static const int LEVEL_INFO; + static const int LEVEL_WARN; + static const int LEVEL_ERROR; + static const int LEVEL_FATAL; + + //! Retrieve string of level + static const char *LevelString(int level) { + static const char *info[] = {"DEBUG", " INFO", " WARN", "ERROR", "FATAL"}; + if (level < (int)(sizeof(info) / sizeof(info[0]))) { + return info[level]; + } + return ""; + } + + //! Retrieve symbol of level + static char LevelSymbol(int level) { + static const char info[5] = {'D', 'I', 'W', 'E', 'F'}; + if (level < (int)(sizeof(info) / sizeof(info[0]))) { + return info[level]; + } + return ' '; + } + + //! Destructor + virtual ~Logger(void) {} + + //! Initialize Logger + virtual int init(const Params ¶ms) = 0; + + //! Cleanup Logger + virtual int cleanup(void) = 0; + + //! Log Message + virtual void log(int level, const char *file, int line, const char *format, + va_list args) = 0; +}; + +/*! Index Logger Broker + */ +class LoggerBroker { + public: + //! Register Logger + static Logger::Pointer Register(Logger::Pointer logger) { + Logger::Pointer ret = std::move(logger_); + logger_ = std::move(logger); + return ret; + } + + //! Register Logger with init params + static int Register(Logger::Pointer logger, const ailego::Params ¶ms) { + //! Cleanup the previous, before initizlizing the new one + if (logger_) { + logger_->cleanup(); + } + logger_ = std::move(logger); + return logger_->init(params); + } + + //! Unregister Logger + static void Unregister(void) { + logger_ = nullptr; + } + + //! Set Level of Logger + static void SetLevel(int level) { + logger_level_ = level; + } + + //! Check if log level is enabled + static bool IsLevelEnabled(int level) { + return logger_level_ <= level && logger_; + } + + //! Log Message + __attribute__((format(printf, 4, 5))) static void Log( + int level, const char *file, int line, const char *format, ...) { + if (IsLevelEnabled(level)) { + va_list args; + va_start(args, format); + logger_->log(level, file, line, format, args); + va_end(args); + } + } + + private: + //! Disable them + LoggerBroker(void) = delete; + LoggerBroker(const LoggerBroker &) = delete; + LoggerBroker(LoggerBroker &&) = delete; + + //! Members + static int logger_level_; + static Logger::Pointer logger_; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/cosine_distance_matrix.h b/src/ailego/math/cosine_distance_matrix.h new file mode 100644 index 0000000..8f4aabd --- /dev/null +++ b/src/ailego/math/cosine_distance_matrix.h @@ -0,0 +1,94 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "distance_utility.h" +#include "inner_product_matrix.h" +#include "norm2_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..c5934dc --- /dev/null +++ b/src/ailego/math/distance.h @@ -0,0 +1,280 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 hamming distance between two vectors (BINARY) + static float Hamming(const uint32_t *lhs, const uint32_t *rhs, size_t dim) { + float result; + HammingDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + +#if defined(AILEGO_M64) + //! Compute the hamming distance between two vectors (BINARY) + static float Hamming(const uint64_t *lhs, const uint64_t *rhs, size_t dim) { + float result; + HammingDistanceMatrix::Compute(lhs, rhs, dim, &result); + return result; + } + +#else + //! Compute the hamming distance between two vectors (BINARY) + static float Hamming(const uint64_t *lhs, const uint64_t *rhs, size_t dim) { + float result; + HammingDistanceMatrix::Compute( + reinterpret_cast(lhs), + reinterpret_cast(rhs), dim, &result); + return result; + } +#endif + + //! 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..760f6f0 --- /dev/null +++ b/src/ailego/math/distance_matrix.h @@ -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. + +#pragma once + +#include "cosine_distance_matrix.h" +#include "euclidean_distance_matrix.h" +#include "hamming_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..43cc68d --- /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_epi64(*(const __m64 *)(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_epi64(*(const __m64 *)(m))); \ + __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(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_epi64(*(const __m64 *)(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_epi64(*(const __m64 *)(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_epi64(*(const __m64 *)(m))); \ + __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(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_fp16.i b/src/ailego/math/distance_matrix_fp16.i new file mode 100644 index 0000000..8cd558e --- /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" + +#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_epi64((__m64)(_MASK), *((const __m64 *)(lhs)))); \ + __m256 ymm_rhs = _mm256_cvtph_ps( \ + _mm_set_epi64((__m64)(_MASK), *((const __m64 *)(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((short)(_MASK), (short)(_MASK), (short)(_MASK), \ + (short)(_MASK), (short)(_MASK), (short)(_MASK), \ + (short)(_MASK), *((const short *)(lhs)))); \ + __m256 ymm_rhs = _mm256_cvtph_ps( \ + _mm_set_epi16((short)(_MASK), (short)(_MASK), (short)(_MASK), \ + (short)(_MASK), (short)(_MASK), (short)(_MASK), \ + (short)(_MASK), *((const short *)(rhs)))); \ + _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_epi64(*(const __m64 *)(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_epi64(*(const __m64 *)(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..4f610d6 --- /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_int32.i b/src/ailego/math/distance_matrix_int32.i new file mode 100644 index 0000000..867d92d --- /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..6360309 --- /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_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..070745a --- /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..cb0cf2c --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix.h @@ -0,0 +1,2639 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! 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; + } +}; + +/*! 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) { + 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))); + } +}; + +/*! 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) { + 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))); + } +}; + +/*! 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)]); + } +}; + +#if !defined(__SSE4_1__) +/*! Squared Euclidean Distance Matrix (INT4, M=1, N=1) + */ +template <> +struct SquaredEuclideanDistanceMatrix { + //! 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 & 1) && out); + + 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)]; + } + *out = sum; + } +}; +#endif // !__SSE4_1__ + +/*! 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); + } +}; + +#if !defined(__SSE4_1__) +/*! Euclidean Distance Matrix (INT4, M=1, N=1) + */ +template <> +struct EuclideanDistanceMatrix { + //! 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 & 1) && out); + + 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)]; + } + *out = std::sqrt(sum); + } +}; +#endif // !__SSE4_1__ + +#if defined(__SSE__) || defined(__ARM_NEON) +/*! Squared Euclidean Distance Matrix (FP32, M=1, N=1) + */ +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 (FP32, M=2, N=1) + */ +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 (FP32, M=2, N=2) + */ +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 (FP32, M=4, N=1) + */ +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 (FP32, M=4, N=2) + */ +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 (FP32, M=4, N=4) + */ +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 (FP32, M=8, N=1) + */ +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 (FP32, M=8, N=2) + */ +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 (FP32, M=8, N=4) + */ +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 (FP32, M=8, N=8) + */ +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 (FP32, M=16, N=1) + */ +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 (FP32, M=16, N=2) + */ +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 (FP32, M=16, N=4) + */ +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 (FP32, M=16, N=8) + */ +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 (FP32, M=16, N=16) + */ +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 (FP32, M=32, N=1) + */ +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 (FP32, M=32, N=2) + */ +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 (FP32, M=32, N=4) + */ +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 (FP32, M=32, N=8) + */ +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 (FP32, M=32, N=16) + */ +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 (FP32, M=32, N=32) + */ +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); +}; +#endif // __SSE__ || __ARM_NEON + +#if defined(__SSE__) || (defined(__ARM_NEON) && (defined(__aarch64__))) +/*! Euclidean Distance Matrix (FP32, M=1, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=2, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=2, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=4, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=4, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=4, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=8, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=8, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=8, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=8, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=16, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=16, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=16, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=16, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=16, N=16) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=32, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=32, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=32, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=32, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=32, N=16) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP32, M=32, N=32) + */ +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); +}; +#endif // __SSE__ || __ARM_NEON && __aarch64__ + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +/*! Squared Euclidean Distance Matrix (FP16, M=1, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=1, N=1) + */ +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); +}; + +#if !defined(__ARM_NEON) +/*! Squared Euclidean Distance Matrix (FP16, M=2, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=2, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=4, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=4, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=4, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=8, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=8, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=8, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=8, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=16, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=16, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=16, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=16, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=16, N=16) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=32, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=32, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=32, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=32, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=32, N=16) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (FP16, M=32, N=32) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=2, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=2, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=4, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=4, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=4, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=8, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=8, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=8, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=8, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=16, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=16, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=16, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=16, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=16, N=16) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=32, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=32, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=32, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=32, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=32, N=16) + */ +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); +}; + +/*! Euclidean Distance Matrix (FP16, M=32, N=32) + */ +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); +}; +#endif // !__ARM_NEON +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +#if defined(__SSE4_1__) +/*! Squared Euclidean Distance Matrix (INT8, M=1, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=2, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=2, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=4, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=4, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=4, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=8, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=8, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=8, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=8, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=16, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=16, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=16, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=16, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=16, N=16) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=32, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=32, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=32, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=32, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=32, N=16) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT8, M=32, N=32) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=1, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=2, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=2, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=4, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=4, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=4, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=8, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=8, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=8, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=8, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=16, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=16, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=16, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=16, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=16, N=16) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=32, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=32, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=32, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=32, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=32, N=16) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT8, M=32, N=32) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=1, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=2, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=2, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=4, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=4, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=4, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=8, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=8, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=8, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=8, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=16, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=16, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=16, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=16, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=16, N=16) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=32, N=1) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=32, N=2) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=32, N=4) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=32, N=8) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=32, N=16) + */ +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); +}; + +/*! Squared Euclidean Distance Matrix (INT4, M=32, N=32) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=1, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=2, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=2, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=4, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=4, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=4, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=8, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=8, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=8, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=8, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=16, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=16, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=16, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=16, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=16, N=16) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=32, N=1) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=32, N=2) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=32, N=4) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=32, N=8) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=32, N=16) + */ +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); +}; + +/*! Euclidean Distance Matrix (INT4, M=32, N=32) + */ +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); +}; +#endif // __SSE4_1__ + +/*! 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.cc b/src/ailego/math/euclidean_distance_matrix_fp16.cc new file mode 100644 index 0000000..366642a --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp16.cc @@ -0,0 +1,551 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#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_FP32_STEP_NEON SSD_FP32_NEON +#define ACCUM_FP16_STEP_GENERAL SSD_FP16_GENERAL +#define ACCUM_FP16_STEP_NEON SSD_FP16_NEON + +//! 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 (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_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 (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); \ + } + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) + +#if defined(__AVX512FP16__) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceAVX512FP16(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 + +//! 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) + ACCUM_FP16_1X1_NEON(m, q, dim, out, 0ull, ) +#elif defined(__AVX512FP16__) + *out = SquaredEuclideanDistanceAVX512FP16(m, q, dim); +#elif defined(__AVX512F__) + ACCUM_FP16_1X1_AVX512(m, q, dim, out, 0ull, ) +#else + ACCUM_FP16_1X1_AVX(m, q, dim, out, 0ull, ) +#endif // __AVX512F__ +} + +//! 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) { +#if defined(__ARM_NEON) + ACCUM_FP16_1X1_NEON(m, q, dim, out, 0ull, std::sqrt) +#elif defined(__AVX512FP16__) + *out = std::sqrt(SquaredEuclideanDistanceAVX512FP16(m, q, dim)); +#elif defined(__AVX512F__) + ACCUM_FP16_1X1_AVX512(m, q, dim, out, 0ull, std::sqrt) +#else + ACCUM_FP16_1X1_AVX(m, q, dim, out, 0ull, std::sqrt) +#endif // __AVX512F__ +} + +#if !defined(__ARM_NEON) +//! Compute the distance between matrix and query (FP16, M=2, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_2X1_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=2, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_2X2_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_4X1_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_4X2_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_4X4_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_8X1_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_8X2_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_8X4_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { + ACCUM_FP16_8X8_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=16, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X1_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X1_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X2_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X2_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X4_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X4_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X8_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X8_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=16) +void SquaredEuclideanDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X16_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X16_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X1_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X1_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X2_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X2_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X4_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X4_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X8_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X8_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=16) +void SquaredEuclideanDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X16_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X16_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=32) +void SquaredEuclideanDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X32_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X32_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=2, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_2X1_AVX(m, q, dim, out, _mm_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=2, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_2X2_AVX(m, q, dim, out, _mm_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_4X1_AVX(m, q, dim, out, _mm_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_4X2_AVX(m, q, dim, out, _mm_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_4X4_AVX(m, q, dim, out, _mm_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_8X1_AVX(m, q, dim, out, _mm256_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_8X2_AVX(m, q, dim, out, _mm256_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_8X4_AVX(m, q, dim, out, _mm256_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_8X8_AVX(m, q, dim, out, _mm256_sqrt_ps) +} + +//! Compute the distance between matrix and query (FP16, M=16, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X1_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_16X1_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X2_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_16X2_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X4_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_16X4_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X8_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_16X8_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=16) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X16_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_16X16_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X1_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_32X1_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X2_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_32X2_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X4_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_32X4_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X8_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_32X8_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=16) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X16_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_32X16_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=32) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X32_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#else + ACCUM_FP16_32X32_AVX(m, q, dim, out, _mm256_sqrt_ps) +#endif // __AVX512F__ +} +#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/euclidean_distance_matrix_fp32.cc b/src/ailego/math/euclidean_distance_matrix_fp32.cc new file mode 100644 index 0000000..9b49be3 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_fp32.cc @@ -0,0 +1,920 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#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_FP32_STEP_NEON SSD_FP32_NEON + +//! 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 (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 (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); \ + } + +#if defined(__ARM_NEON) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceNEON(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) { + 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) + } + return result; +} +#endif // __ARM_NEON + +#if defined(__SSE__) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceSSE(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; +} +#endif // __SSE__ + +#if defined(__AVX__) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceAVX(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; +} +#endif // __AVX__ + +#if defined(__AVX512F__) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceAVX512(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); +} +#endif + +#if defined(__SSE__) || defined(__ARM_NEON) +//! 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) + *out = SquaredEuclideanDistanceNEON(m, q, dim); +#else +#if defined(__AVX512F__) + if (dim > 15) { + *out = SquaredEuclideanDistanceAVX512(m, q, dim); + return; + } +#endif // __AVX512F__ +#if defined(__AVX__) + if (dim > 7) { + *out = SquaredEuclideanDistanceAVX(m, q, dim); + return; + } +#endif // __AVX__ + *out = SquaredEuclideanDistanceSSE(m, q, dim); +#endif // __ARM_NEON +} + +//! Compute the distance between matrix and query (FP32, M=2, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_2X1_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_2X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_2X1_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=2, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_2X2_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_2X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_2X2_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X1_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_4X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_4X1_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X2_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_4X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_4X2_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X4_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_4X4_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_4X4_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X1_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_8X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_8X1_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X2_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_8X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_8X2_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X4_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_8X4_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_8X4_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X8_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_8X8_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_8X8_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=16, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X1_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X1_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X1_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X2_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X2_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X2_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X4_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X4_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X4_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X4_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X8_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X8_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X8_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X8_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=16) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X16_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X16_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X16_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X16_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X1_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X1_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X1_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X2_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X2_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X2_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X4_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X4_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X4_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X4_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X8_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X8_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X8_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X8_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=16) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X16_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X16_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X16_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X16_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=32) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X32_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X32_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X32_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X32_SSE(m, q, dim, out, ) +#endif +} +#endif // __SSE__ || __ARM_NEON + +#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__)) +//! 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) { +#if defined(__ARM_NEON) + *out = std::sqrt(SquaredEuclideanDistanceNEON(m, q, dim)); +#else +#if defined(__AVX512F__) + if (dim > 15) { + *out = std::sqrt(SquaredEuclideanDistanceAVX512(m, q, dim)); + return; + } +#endif // __AVX512F__ +#if defined(__AVX__) + if (dim > 7) { + *out = std::sqrt(SquaredEuclideanDistanceAVX(m, q, dim)); + return; + } +#endif // __AVX__ + *out = std::sqrt(SquaredEuclideanDistanceSSE(m, q, dim)); +#endif // __ARM_NEON +} + +//! Compute the distance between matrix and query (FP32, M=2, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_2X1_NEON(m, q, dim, out, vsqrt_f32) +#elif defined(__AVX__) + ACCUM_FP32_2X1_AVX(m, q, dim, out, _mm_sqrt_ps) +#else + ACCUM_FP32_2X1_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=2, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_2X2_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX__) + ACCUM_FP32_2X2_AVX(m, q, dim, out, _mm_sqrt_ps) +#else + ACCUM_FP32_2X2_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X1_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX__) + ACCUM_FP32_4X1_AVX(m, q, dim, out, _mm_sqrt_ps) +#else + ACCUM_FP32_4X1_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X2_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX__) + ACCUM_FP32_4X2_AVX(m, q, dim, out, _mm_sqrt_ps) +#else + ACCUM_FP32_4X2_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X4_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX__) + ACCUM_FP32_4X4_AVX(m, q, dim, out, _mm_sqrt_ps) +#else + ACCUM_FP32_4X4_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X1_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX__) + ACCUM_FP32_8X1_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_8X1_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X2_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX__) + ACCUM_FP32_8X2_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_8X2_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X4_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX__) + ACCUM_FP32_8X4_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_8X4_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X8_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX__) + ACCUM_FP32_8X8_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_8X8_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=16, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X1_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X1_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_16X1_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_16X1_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X2_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X2_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_16X2_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_16X2_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X4_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X4_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_16X4_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_16X4_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X8_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X8_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_16X8_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_16X8_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=16) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X16_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X16_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_16X16_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_16X16_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X1_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X1_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_32X1_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_32X1_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X2_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X2_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_32X2_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_32X2_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X4_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X4_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_32X4_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_32X4_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X8_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X8_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_32X8_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_32X8_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=16) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X16_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X16_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_32X16_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_32X16_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=32) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X32_NEON(m, q, dim, out, vsqrtq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X32_AVX512(m, q, dim, out, _mm512_sqrt_ps) +#elif defined(__AVX__) + ACCUM_FP32_32X32_AVX(m, q, dim, out, _mm256_sqrt_ps) +#else + ACCUM_FP32_32X32_SSE(m, q, dim, out, _mm_sqrt_ps) +#endif +} +#endif // __SSE__ || __ARM_NEON && __aarch64__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_int4.cc b/src/ailego/math/euclidean_distance_matrix_int4.cc new file mode 100644 index 0000000..bef4321 --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_int4.cc @@ -0,0 +1,801 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#define ACCUM_INT4_STEP_SSE SSD_INT4_SSE +#define ACCUM_INT4_STEP_AVX SSD_INT4_AVX + +#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__ + +#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 (GENERAL) +#define SSD_INT4_GENERAL(m, q, sum) \ + sum += Int4SquaredDiffTable[(((m) << 4) & 0xf0) | (((q) >> 0) & 0xf)] + \ + Int4SquaredDiffTable[(((m) >> 0) & 0xf0) | (((q) >> 4) & 0xf)]; + +//! 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); \ + } + +//! 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_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); \ + } + +//! 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); \ + } + +//! Compute the square root of value (SSE) +#define SQRT_FP32_SSE(v, ...) _mm_sqrt_ps(_mm_cvtepi32_ps(v)) + +//! 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)) + +#if defined(__SSE4_1__) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceSSE(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; +} +#endif // __SSE4_1__ + +#if defined(__AVX2__) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceAVX(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; +} +#endif // __AVX2__ + +#if defined(__SSE4_1__) +//! 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 (dim > 63) { + *out = SquaredEuclideanDistanceAVX(m, q, dim >> 1); + return; + } +#endif // __AVX2__ + *out = SquaredEuclideanDistanceSSE(m, q, dim >> 1); +} + +//! Compute the distance between matrix and query (INT4, M=2, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_2X1_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_2X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=2, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_2X2_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_2X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X1_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_4X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X2_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_4X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X4_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_4X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_8X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_8X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_8X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_8X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=16) +void SquaredEuclideanDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X16_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X16_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=16) +void SquaredEuclideanDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X16_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X16_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=32) +void SquaredEuclideanDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X32_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X32_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! 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) { +#if defined(__AVX2__) + if (dim > 63) { + *out = std::sqrt(SquaredEuclideanDistanceAVX(m, q, dim >> 1)); + return; + } +#endif // __AVX2__ + *out = std::sqrt(SquaredEuclideanDistanceSSE(m, q, dim >> 1)); +} + +//! Compute the distance between matrix and query (INT4, M=2, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_2X1_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT4_2X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=2, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_2X2_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT4_2X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X1_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT4_4X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X2_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT4_4X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X4_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT4_4X4_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X1_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_8X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X2_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_8X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X4_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_8X4_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X8_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_8X8_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X1_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_16X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X2_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_16X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X4_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_16X4_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X8_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_16X8_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=16) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X16_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_16X16_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X1_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_32X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X2_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_32X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X4_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_32X4_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X8_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_32X8_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=16) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X16_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_32X16_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=32) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X32_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT4_32X32_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/euclidean_distance_matrix_int8.cc b/src/ailego/math/euclidean_distance_matrix_int8.cc new file mode 100644 index 0000000..cd14d2d --- /dev/null +++ b/src/ailego/math/euclidean_distance_matrix_int8.cc @@ -0,0 +1,884 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#define ACCUM_INT8_STEP_SSE SSD_INT8_SSE +#define ACCUM_INT8_STEP_AVX SSD_INT8_AVX + +#if defined(__SSE4_1__) +static const __m128i ONES_INT16_SSE = _mm_set1_epi32(0x00010001); +#endif // __SSE4_1__ + +#if defined(__AVX2__) +static const __m256i ONES_INT16_AVX = _mm256_set1_epi32(0x00010001); +#endif // __AVX2__ + +//! 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 (SSE) +#define SQRT_FP32_SSE(v, ...) _mm_sqrt_ps(_mm_cvtepi32_ps(v)) + +//! 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)) + +#if defined(__SSE4_1__) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceSSE(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; +} +#endif // __SSE4_1__ + +#if defined(__AVX2__) +//! Squared Euclidean Distance +static inline float SquaredEuclideanDistanceAVX(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; +} +#endif // __AVX2__ + +#if defined(__SSE4_1__) +//! 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 (dim > 31) { + *out = SquaredEuclideanDistanceAVX(m, q, dim); + return; + } +#endif // __AVX2__ + *out = SquaredEuclideanDistanceSSE(m, q, dim); +} + +//! Compute the distance between matrix and query (INT8, M=2, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_2X1_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_2X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=2, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_2X2_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_2X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X1_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_4X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X2_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_4X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X4_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_4X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_8X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_8X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_8X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_8X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=16) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X16_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X16_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=1) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=2) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=4) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=8) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=16) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X16_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X16_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=32) +void SquaredEuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X32_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X32_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! 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) { +#if defined(__AVX2__) + if (dim > 31) { + *out = std::sqrt(SquaredEuclideanDistanceAVX(m, q, dim)); + return; + } +#endif // __AVX2__ + *out = std::sqrt(SquaredEuclideanDistanceSSE(m, q, dim)); +} + +//! Compute the distance between matrix and query (INT8, M=2, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_2X1_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT8_2X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=2, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_2X2_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT8_2X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X1_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT8_4X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X2_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT8_4X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X4_AVX(m, q, dim, out, SQRT_FP32_SSE) +#else + ACCUM_INT8_4X4_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X1_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_8X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X2_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_8X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X4_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_8X4_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X8_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_8X8_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X1_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_16X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X2_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_16X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X4_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_16X4_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X8_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_16X8_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=16) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X16_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_16X16_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=1) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X1_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_32X1_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=2) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X2_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_32X2_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=4) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X4_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_32X4_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=8) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X8_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_32X8_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=16) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X16_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_32X16_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=32) +void EuclideanDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X32_AVX(m, q, dim, out, SQRT_FP32_AVX) +#else + ACCUM_INT8_32X32_SSE(m, q, dim, out, SQRT_FP32_SSE) +#endif // __AVX2__ +} +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/hamming_distance_matrix.cc b/src/ailego/math/hamming_distance_matrix.cc new file mode 100644 index 0000000..7e7d058 --- /dev/null +++ b/src/ailego/math/hamming_distance_matrix.cc @@ -0,0 +1,1128 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "hamming_distance_matrix.h" +#include "distance_matrix_popcnt.i" + +namespace zvec { +namespace ailego { + +#define POPCNT_UINT32_STEP1_SSE HAMMING_UINT32_STEP1_SSE +#define POPCNT_UINT32_STEP2_SSE HAMMING_UINT32_STEP2_SSE +#define POPCNT_UINT32_STEP3_SSE HAMMING_UINT32_STEP3_SSE +#define POPCNT_UINT32_STEP1_AVX HAMMING_UINT32_STEP1_AVX +#define POPCNT_UINT32_STEP2_AVX HAMMING_UINT32_STEP2_AVX +#define POPCNT_UINT32_STEP3_AVX HAMMING_UINT32_STEP3_AVX +#define POPCNT_UINT64_STEP1_AVX HAMMING_UINT64_STEP1_AVX +#define POPCNT_UINT64_STEP2_AVX HAMMING_UINT64_STEP2_AVX + +//! Calculate population count (Step 1 SSE) +#define HAMMING_UINT32_STEP1_SSE(xmm_m, xmm_q, xmm_sum) \ + xmm_sum = _mm_add_epi8( \ + VerticalPopCount_INT8_V128(_mm_xor_si128(xmm_m, xmm_q)), xmm_sum); + +//! Calculate population count (Step 2 SSE) +#define HAMMING_UINT32_STEP2_SSE(xmm_m, xmm_q, xmm_sum) \ + xmm_sum = _mm_add_epi16( \ + VerticalPopCount_INT16_V128(_mm_xor_si128(xmm_m, xmm_q)), xmm_sum); + +//! Calculate population count (Step 3 SSE) +#define HAMMING_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum) \ + xmm_sum = _mm_add_epi32( \ + VerticalPopCount_INT32_V128(_mm_xor_si128(xmm_m, xmm_q)), xmm_sum); + +//! Calculate population count (Step 1 AVX) +#define HAMMING_UINT32_STEP1_AVX(ymm_m, ymm_q, ymm_sum) \ + ymm_sum = _mm256_add_epi8( \ + VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum); + +//! Calculate population count (Step 2 AVX) +#define HAMMING_UINT32_STEP2_AVX(ymm_m, ymm_q, ymm_sum) \ + ymm_sum = _mm256_add_epi16( \ + VerticalPopCount_INT16_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum); + +//! Calculate population count (Step 3 AVX) +#define HAMMING_UINT32_STEP3_AVX(ymm_m, ymm_q, ymm_sum) \ + ymm_sum = _mm256_add_epi32( \ + VerticalPopCount_INT32_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum); + +//! Calculate population count (Step 1 AVX) +#define HAMMING_UINT64_STEP1_AVX(ymm_m, ymm_q, ymm_sum) \ + ymm_sum = _mm256_add_epi8( \ + VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum); + +//! Calculate population count (Step 2 AVX) +#define HAMMING_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum) \ + ymm_sum = _mm256_add_epi64( \ + VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum); + +#if defined(__AVX512VL__) && defined(__AVX512DQ__) +#define CONVERT_UINT64_TO_FP32(v, ...) _mm256_cvtepu64_ps(v) +#elif defined(__AVX2__) +static const __m256i CONVERT_UINT32_MASK_AVX = + _mm256_set_epi32(0, 0, 0, 0, 6, 4, 2, 0); + +#define CONVERT_UINT64_TO_FP32(v, ...) \ + _mm_cvtepi32_ps(_mm256_castsi256_si128( \ + _mm256_permutevar8x32_epi32(v, CONVERT_UINT32_MASK_AVX))) +#endif // __AVX512VL__ && __AVX512DQ__ + +#define SQRT_UINT64_TO_FP32(v, ...) _mm_sqrt_ps(CONVERT_UINT64_TO_FP32(v)) +#define SQRT_UINT32_TO_FP32_SSE(v, ...) _mm_sqrt_ps(_mm_cvtepi32_ps(v)) +#define SQRT_UINT32_TO_FP32_AVX(v, ...) _mm256_sqrt_ps(_mm256_cvtepi32_ps(v)) + +#if defined(__AVX2__) +static inline size_t HammingDistanceAVX(const uint32_t *lhs, + const uint32_t *rhs, size_t size) { + __m256i ymm_sum_0 = _mm256_setzero_si256(); + __m256i ymm_sum_1 = _mm256_setzero_si256(); + + const uint32_t *lhs_0 = lhs + ((size >> 4) << 4); + const uint32_t *lhs_1 = (size > 496 ? lhs + 496 : lhs_0); + const uint32_t *lhs_2 = lhs + size; + + if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) { + for (; lhs != lhs_1; lhs += 16, rhs += 16) { + __m256i ymm_lhs_0 = _mm256_load_si256((__m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_load_si256((__m256i *)(lhs + 8)); + __m256i ymm_rhs_0 = _mm256_load_si256((__m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_load_si256((__m256i *)(rhs + 8)); + + ymm_sum_0 = _mm256_add_epi8( + VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_lhs_0, ymm_rhs_0)), + ymm_sum_0); + ymm_sum_1 = _mm256_add_epi8( + VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_lhs_1, ymm_rhs_1)), + ymm_sum_1); + } + ymm_sum_0 = _mm256_sad_epu8(ymm_sum_0, POPCNT_ZERO_AVX); + ymm_sum_1 = _mm256_sad_epu8(ymm_sum_1, POPCNT_ZERO_AVX); + + for (; lhs != lhs_0; lhs += 16, rhs += 16) { + __m256i ymm_lhs_0 = _mm256_load_si256((__m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_load_si256((__m256i *)(lhs + 8)); + __m256i ymm_rhs_0 = _mm256_load_si256((__m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_load_si256((__m256i *)(rhs + 8)); + + ymm_sum_0 = _mm256_add_epi64( + VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs_0, ymm_rhs_0)), + ymm_sum_0); + ymm_sum_1 = _mm256_add_epi64( + VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs_1, ymm_rhs_1)), + ymm_sum_1); + } + + if (lhs_2 >= lhs + 8) { + __m256i ymm_lhs = _mm256_load_si256((__m256i *)(lhs)); + __m256i ymm_rhs = _mm256_load_si256((__m256i *)(rhs)); + ymm_sum_0 = _mm256_add_epi64( + VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs, ymm_rhs)), + ymm_sum_0); + lhs += 8; + rhs += 8; + } + } else { + for (; lhs != lhs_1; lhs += 16, rhs += 16) { + __m256i ymm_lhs_0 = _mm256_loadu_si256((__m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_loadu_si256((__m256i *)(lhs + 8)); + __m256i ymm_rhs_0 = _mm256_loadu_si256((__m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_loadu_si256((__m256i *)(rhs + 8)); + + ymm_sum_0 = _mm256_add_epi8( + VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_lhs_0, ymm_rhs_0)), + ymm_sum_0); + ymm_sum_1 = _mm256_add_epi8( + VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_lhs_1, ymm_rhs_1)), + ymm_sum_1); + } + ymm_sum_0 = _mm256_sad_epu8(ymm_sum_0, POPCNT_ZERO_AVX); + ymm_sum_1 = _mm256_sad_epu8(ymm_sum_1, POPCNT_ZERO_AVX); + + for (; lhs != lhs_0; lhs += 16, rhs += 16) { + __m256i ymm_lhs_0 = _mm256_loadu_si256((__m256i *)(lhs + 0)); + __m256i ymm_lhs_1 = _mm256_loadu_si256((__m256i *)(lhs + 8)); + __m256i ymm_rhs_0 = _mm256_loadu_si256((__m256i *)(rhs + 0)); + __m256i ymm_rhs_1 = _mm256_loadu_si256((__m256i *)(rhs + 8)); + + ymm_sum_0 = _mm256_add_epi64( + VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs_0, ymm_rhs_0)), + ymm_sum_0); + ymm_sum_1 = _mm256_add_epi64( + VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs_1, ymm_rhs_1)), + ymm_sum_1); + } + + if (lhs_2 >= lhs + 8) { + __m256i ymm_lhs = _mm256_loadu_si256((__m256i *)(lhs)); + __m256i ymm_rhs = _mm256_loadu_si256((__m256i *)(rhs)); + ymm_sum_0 = _mm256_add_epi64( + VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs, ymm_rhs)), + ymm_sum_0); + lhs += 8; + rhs += 8; + } + } + + size_t count = + (size_t)HorizontalAdd_INT64_V256(_mm256_add_epi64(ymm_sum_0, ymm_sum_1)); + switch (lhs_2 - lhs) { + case 7: + count += ailego_popcount32(lhs[6] ^ rhs[6]); + /* FALLTHRU */ + case 6: + count += ailego_popcount32(lhs[5] ^ rhs[5]); + /* FALLTHRU */ + case 5: + count += ailego_popcount32(lhs[4] ^ rhs[4]); + /* FALLTHRU */ + case 4: + count += ailego_popcount32(lhs[3] ^ rhs[3]); + /* FALLTHRU */ + case 3: + count += ailego_popcount32(lhs[2] ^ rhs[2]); + /* FALLTHRU */ + case 2: + count += ailego_popcount32(lhs[1] ^ rhs[1]); + /* FALLTHRU */ + case 1: + count += ailego_popcount32(lhs[0] ^ rhs[0]); + } + return count; +} + +static inline size_t HammingDistanceAVX(const uint64_t *lhs, + const uint64_t *rhs, size_t size) { + return HammingDistanceAVX(reinterpret_cast(lhs), + reinterpret_cast(rhs), + (size << 1)); +} +#endif // __AVX2__ + +#if defined(AILEGO_M64) +static inline size_t HammingDistance(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 += ailego_popcount64(*(uint64_t *)(&lhs[6]) ^ *(uint64_t *)(&rhs[6])); + count += ailego_popcount64(*(uint64_t *)(&lhs[4]) ^ *(uint64_t *)(&rhs[4])); + count += ailego_popcount64(*(uint64_t *)(&lhs[2]) ^ *(uint64_t *)(&rhs[2])); + count += ailego_popcount64(*(uint64_t *)(&lhs[0]) ^ *(uint64_t *)(&rhs[0])); + } + switch (last - last_aligned) { + case 7: + count += ailego_popcount32(lhs[6] ^ rhs[6]); + /* FALLTHRU */ + case 6: + count += ailego_popcount32(lhs[5] ^ rhs[5]); + /* FALLTHRU */ + case 5: + count += ailego_popcount32(lhs[4] ^ rhs[4]); + /* FALLTHRU */ + case 4: + count += ailego_popcount32(lhs[3] ^ rhs[3]); + /* FALLTHRU */ + case 3: + count += ailego_popcount32(lhs[2] ^ rhs[2]); + /* FALLTHRU */ + case 2: + count += ailego_popcount32(lhs[1] ^ rhs[1]); + /* FALLTHRU */ + case 1: + count += ailego_popcount32(lhs[0] ^ rhs[0]); + } + return count; +} + +static inline size_t HammingDistance(const uint64_t *lhs, const uint64_t *rhs, + size_t size) { + const uint64_t *last = lhs + size; + const uint64_t *last_aligned = lhs + ((size >> 2) << 2); + size_t count = 0; + + for (; lhs != last_aligned; lhs += 4, rhs += 4) { + count += ailego_popcount64(lhs[3] ^ rhs[3]); + count += ailego_popcount64(lhs[2] ^ rhs[2]); + count += ailego_popcount64(lhs[1] ^ rhs[1]); + count += ailego_popcount64(lhs[0] ^ rhs[0]); + } + switch (last - last_aligned) { + case 3: + count += ailego_popcount64(lhs[2] ^ rhs[2]); + /* FALLTHRU */ + case 2: + count += ailego_popcount64(lhs[1] ^ rhs[1]); + /* FALLTHRU */ + case 1: + count += ailego_popcount64(lhs[0] ^ rhs[0]); + } + return count; +} +#else +static inline size_t HammingDistance(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 += ailego_popcount32(lhs[3] ^ rhs[3]); + count += ailego_popcount32(lhs[2] ^ rhs[2]); + count += ailego_popcount32(lhs[1] ^ rhs[1]); + count += ailego_popcount32(lhs[0] ^ rhs[0]); + } + switch (last - last_aligned) { + case 3: + count += ailego_popcount32(lhs[2] ^ rhs[2]); + /* FALLTHRU */ + case 2: + count += ailego_popcount32(lhs[1] ^ rhs[1]); + /* FALLTHRU */ + case 1: + count += ailego_popcount32(lhs[0] ^ rhs[0]); + } + return count; +} +#endif // AILEGO_M64 + +//! Compute the distance between matrix and query (UINT32, M=1, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + if (cnt > 63) { + *out = static_cast(HammingDistanceAVX(m, q, cnt)); + return; + } +#endif + *out = static_cast(HammingDistance(m, q, cnt)); +} + +#if defined(__SSSE3__) +//! Compute the distance between matrix and query (UINT32, M=2, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_2X1_AVX(m, q, cnt, out, _mm_cvtepi32_ps) +#else + POPCNT_UINT32_2X1_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=2, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_2X2_AVX(m, q, cnt, out, _mm_cvtepi32_ps) +#else + POPCNT_UINT32_2X2_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=4, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_4X1_AVX(m, q, cnt, out, _mm_cvtepi32_ps) +#else + POPCNT_UINT32_4X1_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=4, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_4X2_AVX(m, q, cnt, out, _mm_cvtepi32_ps) +#else + POPCNT_UINT32_4X2_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=4, N=4) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_4X4_AVX(m, q, cnt, out, _mm_cvtepi32_ps) +#else + POPCNT_UINT32_4X4_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=8, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_8X1_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_8X1_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=8, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_8X2_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_8X2_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=8, N=4) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_8X4_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_8X4_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=8, N=8) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_8X8_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_8X8_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X1_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_16X1_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X2_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_16X2_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=4) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X4_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_16X4_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=8) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X8_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_16X8_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=16) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X16_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_16X16_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X1_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_32X1_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X2_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_32X2_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=4) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X4_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_32X4_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=8) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X8_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_32X8_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=16) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X16_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_32X16_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=32) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X32_AVX(m, q, cnt, out, _mm256_cvtepi32_ps) +#else + POPCNT_UINT32_32X32_SSE(m, q, cnt, out, _mm_cvtepi32_ps) +#endif +} +#endif // __SSSE3__ + +#if defined(AILEGO_M64) +//! Compute the distance between matrix and query (UINT64, M=1, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); +#if defined(__AVX2__) + if (cnt > 31) { + *out = static_cast(HammingDistanceAVX(m, q, cnt)); + return; + } +#endif + *out = static_cast(HammingDistance(m, q, cnt)); +} + +#if defined(__AVX2__) +//! Compute the distance between matrix and query (UINT64, M=2, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_2X1_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=2, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_2X2_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=4, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_4X1_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=4, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_4X2_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=4, N=4) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_4X4_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=8, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_8X1_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=8, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_8X2_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=8, N=4) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_8X4_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=8, N=8) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_8X8_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X1_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X2_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=4) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X4_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=8) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X8_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=16) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X16_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=1) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X1_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=2) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X2_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=4) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X4_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=8) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X8_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=16) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X16_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=32) +void HammingDistanceMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X32_AVX(m, q, cnt, out, CONVERT_UINT64_TO_FP32) +} +#endif // __AVX2__ +#endif // AILEGO_M64 + +//! Compute the distance between matrix and query (UINT32, M=1, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + if (cnt > 63) { + *out = std::sqrt(static_cast(HammingDistanceAVX(m, q, cnt))); + return; + } +#endif + *out = std::sqrt(static_cast(HammingDistance(m, q, cnt))); +} + +#if defined(__SSSE3__) +//! Compute the distance between matrix and query (UINT32, M=2, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_2X1_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#else + POPCNT_UINT32_2X1_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=2, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_2X2_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#else + POPCNT_UINT32_2X2_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=4, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_4X1_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#else + POPCNT_UINT32_4X1_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=4, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_4X2_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#else + POPCNT_UINT32_4X2_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=4, N=4) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_4X4_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#else + POPCNT_UINT32_4X4_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=8, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_8X1_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_8X1_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=8, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_8X2_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_8X2_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=8, N=4) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_8X4_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_8X4_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=8, N=8) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_8X8_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_8X8_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X1_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_16X1_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X2_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_16X2_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=4) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X4_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_16X4_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=8) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X8_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_16X8_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=16, N=16) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_16X16_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_16X16_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X1_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_32X1_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X2_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_32X2_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=4) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X4_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_32X4_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=8) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X8_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_32X8_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=16) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X16_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_32X16_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (UINT32, M=32, N=32) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 5); +#if defined(__AVX2__) + POPCNT_UINT32_32X32_AVX(m, q, cnt, out, SQRT_UINT32_TO_FP32_AVX) +#else + POPCNT_UINT32_32X32_SSE(m, q, cnt, out, SQRT_UINT32_TO_FP32_SSE) +#endif +} +#endif // __SSSE3__ + +#if defined(AILEGO_M64) +//! Compute the distance between matrix and query (UINT64, M=1, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); +#if defined(__AVX2__) + if (cnt > 31) { + *out = std::sqrt(static_cast(HammingDistanceAVX(m, q, cnt))); + return; + } +#endif + *out = std::sqrt(static_cast(HammingDistance(m, q, cnt))); +} + +#if defined(__AVX2__) +//! Compute the distance between matrix and query (UINT64, M=2, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_2X1_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=2, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_2X2_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=4, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_4X1_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=4, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_4X2_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=4, N=4) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_4X4_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=8, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_8X1_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=8, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_8X2_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=8, N=4) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_8X4_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=8, N=8) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_8X8_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X1_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X2_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=4) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X4_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=8) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X8_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=16, N=16) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_16X16_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=1) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X1_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=2) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X2_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=4) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X4_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=8) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X8_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=16) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X16_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} + +//! Compute the distance between matrix and query (UINT64, M=32, N=32) +void HammingSquareRootDistanceMatrix::Compute( + const ValueType *m, const ValueType *q, size_t dim, float *out) { + size_t cnt = (dim >> 6); + POPCNT_UINT64_32X32_AVX(m, q, cnt, out, SQRT_UINT64_TO_FP32) +} +#endif // __AVX2__ +#endif // AILEGO_M64 + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/hamming_distance_matrix.h b/src/ailego/math/hamming_distance_matrix.h new file mode 100644 index 0000000..ace8397 --- /dev/null +++ b/src/ailego/math/hamming_distance_matrix.h @@ -0,0 +1,1158 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Hamming Distance Matrix + */ +template // NOTE: useless 'typename=void' to avoid clang + // compile error +struct HammingDistanceMatrix; + +/*! Hamming Distance Matrix (UINT32) + */ +template +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_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 & 31) && out); + + size_t cnt = (dim >> 5); + if (cnt > 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(ailego_popcount32(m_val ^ q[j])); + r += M; + } + } + m += M; + q += N; + } + + for (size_t k = 1; k < cnt; ++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 += static_cast(ailego_popcount32(m_val ^ q[j])); + r += M; + } + } + m += M; + q += N; + } + } +}; + +/*! Hamming Distance Matrix (UINT32, M=1, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +#if defined(__SSSE3__) +/*! Hamming Distance Matrix (UINT32, M=2, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=2, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=4, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=4, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=4, N=4) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=8, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=8, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=8, N=4) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=8, N=8) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=16, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=16, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=16, N=4) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=16, N=8) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=16, N=16) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=32, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=32, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=32, N=4) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=32, N=8) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=32, N=16) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT32, M=32, N=32) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; +#endif // __SSSE3__ + +#if defined(AILEGO_M64) +/*! Hamming Distance Matrix (UINT64) + */ +template +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_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 & 63) && out); + + size_t cnt = (dim >> 6); + if (cnt > 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(ailego_popcount64(m_val ^ q[j])); + r += M; + } + } + m += M; + q += N; + } + + for (size_t k = 1; k < cnt; ++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 += static_cast(ailego_popcount64(m_val ^ q[j])); + r += M; + } + } + m += M; + q += N; + } + } +}; + +/*! Hamming Distance Matrix (UINT64, M=1, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +#if defined(__AVX2__) +/*! Hamming Distance Matrix (UINT64, M=2, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=2, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=4, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=4, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=4, N=4) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=8, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=8, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=8, N=4) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=8, N=8) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=16, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=16, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=16, N=4) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=16, N=8) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=16, N=16) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=32, N=1) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=32, N=2) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=32, N=4) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=32, N=8) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=32, N=16) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Distance Matrix (UINT64, M=32, N=32) + */ +template <> +struct HammingDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; +#endif // __AVX2__ +#endif // AILEGO_M64 + +/*! Hamming Square Root Distance Matrix + */ +template +struct HammingSquareRootDistanceMatrix { + //! 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); + + HammingDistanceMatrix::Compute(m, q, dim, out); + for (size_t i = 0; i < N * M; ++i) { + float val = *out; + *out++ = std::sqrt(val); + } + } +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=1, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +#if defined(__SSSE3__) +/*! Hamming Square Root Distance Matrix (UINT32, M=2, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=2, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=4, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=4, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=4, N=4) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=8, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=8, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=8, N=4) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=8, N=8) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=16, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=16, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=16, N=4) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=16, N=8) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=16, N=16) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=32, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=32, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=32, N=4) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=32, N=8) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=32, N=16) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT32, M=32, N=32) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint32_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; +#endif // __SSSE3__ + +#if defined(AILEGO_M64) +/*! Hamming Square Root Distance Matrix (UINT64, M=1, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +#if defined(__AVX2__) +/*! Hamming Square Root Distance Matrix (UINT64, M=2, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=2, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=4, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=4, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=4, N=4) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=8, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=8, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=8, N=4) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=8, N=8) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=16, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=16, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=16, N=4) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=16, N=8) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=16, N=16) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=32, N=1) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=32, N=2) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=32, N=4) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=32, N=8) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=32, N=16) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; + +/*! Hamming Square Root Distance Matrix (UINT64, M=32, N=32) + */ +template <> +struct HammingSquareRootDistanceMatrix { + //! Type of value + using ValueType = uint64_t; + + //! Compute the distance between matrix and query + static void Compute(const ValueType *m, const ValueType *q, size_t dim, + float *out); +}; +#endif // __AVX2__ +#endif // AILEGO_M64 + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/inner_product_matrix.h b/src/ailego/math/inner_product_matrix.h new file mode 100644 index 0000000..5bd7e0a --- /dev/null +++ b/src/ailego/math/inner_product_matrix.h @@ -0,0 +1,2991 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Inner Product Matrix + */ +template +struct InnerProductMatrix; + +/*! 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; + } +}; + +/*! 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) { + 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)); + } +}; + +/*! 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) { + 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)); + } +}; + +/*! 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)]); + } +}; + +#if !defined(__SSE4_1__) +/*! Inner Product Matrix (INT4, M=1, N=1) + */ +template <> +struct InnerProductMatrix { + //! 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 & 1) && out); + + 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)]; + } + *out = sum; + } +}; +#endif // !__SSE4_1__ + +template +struct MinusInnerProductMatrix; + +/*! 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; + } +}; + +/*! 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) { + 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)); + } +}; + +/*! 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) { + 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)); + } +}; + +/*! 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)]); + } +}; + +#if !defined(__SSE4_1__) +/*! Minus Inner Product Matrix (INT4, M=1, N=1) + */ +template <> +struct MinusInnerProductMatrix { + //! 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 & 1) && out); + + 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)]; + } + *out = sum; + } +}; +#endif // !__SSE4_1__ + +#if defined(__SSE__) || defined(__ARM_NEON) +/*! Inner Product Matrix (FP32, M=1, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=2, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=2, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=4, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=4, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=4, N=4) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=8, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=8, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=8, N=4) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=8, N=8) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=16, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=16, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=16, N=4) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=16, N=8) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=16, N=16) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=32, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=32, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=32, N=4) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=32, N=8) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=32, N=16) + */ +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); +}; + +/*! Inner Product Matrix (FP32, M=32, N=32) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=1, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=2, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=2, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=4, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=4, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=4, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=8, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=8, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=8, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=8, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=16, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=16, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=16, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=16, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=16, N=16) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=32, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=32, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=32, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=32, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=32, N=16) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP32, M=32, N=32) + */ +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); +}; +#endif // __SSE__ || __ARM_NEON + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +/*! Inner Product Matrix (FP16, M=1, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=1, N=1) + */ +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); +}; + +#if !defined(__ARM_NEON) +/*! Inner Product Matrix (FP16, M=2, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=2, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=4, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=4, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=4, N=4) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=8, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=8, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=8, N=4) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=8, N=8) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=16, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=16, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=16, N=4) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=16, N=8) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=16, N=16) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=32, N=1) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=32, N=2) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=32, N=4) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=32, N=8) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=32, N=16) + */ +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); +}; + +/*! Inner Product Matrix (FP16, M=32, N=32) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=2, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=2, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=4, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=4, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=4, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=8, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=8, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=8, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=8, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=16, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=16, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=16, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=16, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=16, N=16) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=32, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=32, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=32, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=32, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=32, N=16) + */ +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); +}; + +/*! Minus Inner Product Matrix (FP16, M=32, N=32) + */ +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); +}; +#endif // !__ARM_NEON +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +#if defined(__SSE4_1__) +/*! Inner Product Matrix (INT8, M=1, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=2, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=2, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=4, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=4, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=4, N=4) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=8, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=8, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=8, N=4) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=8, N=8) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=16, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=16, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=16, N=4) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=16, N=8) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=16, N=16) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=32, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=32, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=32, N=4) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=32, N=8) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=32, N=16) + */ +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); +}; + +/*! Inner Product Matrix (INT8, M=32, N=32) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=1, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=2, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=2, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=4, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=4, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=4, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=8, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=8, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=8, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=8, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=16, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=16, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=16, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=16, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=16, N=16) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=32, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=32, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=32, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=32, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=32, N=16) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT8, M=32, N=32) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=1, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=2, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=2, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=4, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=4, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=4, N=4) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=8, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=8, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=8, N=4) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=8, N=8) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=16, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=16, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=16, N=4) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=16, N=8) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=16, N=16) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=32, N=1) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=32, N=2) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=32, N=4) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=32, N=8) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=32, N=16) + */ +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); +}; + +/*! Inner Product Matrix (INT4, M=32, N=32) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=1, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=2, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=2, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=4, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=4, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=4, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=8, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=8, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=8, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=8, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=16, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=16, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=16, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=16, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=16, N=16) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=32, N=1) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=32, N=2) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=32, N=4) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=32, N=8) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=32, N=16) + */ +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); +}; + +/*! Minus Inner Product Matrix (INT4, M=32, N=32) + */ +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); +}; +#endif // __SSE4_1__ + +template +struct MinusInnerProductSparseMatrix { + //! Type of value + using ValueType = typename std::remove_cv::type; + + static constexpr uint32_t SEGMENT_ID_BITS = 16; + static constexpr uint32_t SEGMENT_ID_MASK = 0xFFFF; + + 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} {} + }; + + static inline void transform_sparse_format(uint32_t sparse_count, + const uint32_t *sparse_index, + const void *sparse_value, + std::string &buffer); + + static inline 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) { + 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; + } + + 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 += 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; + } + } + + *out = -sum; + } +}; + +template +float MinusInnerProductSparseMatrix::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 +void MinusInnerProductSparseMatrix::transform_sparse_format( + uint32_t sparse_count, const uint32_t *sparse_index, + const void *sparse_value, std::string &buffer) { + uint32_t unit_size = sizeof(T); + + 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(T)); + + 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; + } +} + +#if defined(__SSE4_1__) +template <> +float MinusInnerProductSparseMatrix::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); + +template <> +float MinusInnerProductSparseMatrix:: + 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); +#endif + +#if defined(__AVX512FP16__) +template <> +float MinusInnerProductSparseMatrix:: + 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); +#endif + +} // 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..6e2e8a6 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp16.cc @@ -0,0 +1,1931 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +#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_FP16_STEP_GENERAL FMA_FP16_GENERAL +#define ACCUM_FP16_STEP_NEON FMA_FP16_NEON + +#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__ + +#if defined(__SSE__) +static const __m128 NEGZEROS_FP32_SSE = _mm_set1_ps(-0.0f); +#endif // __SSE__ + +#if defined(__AVX__) +static const __m256 NEGZEROS_FP32_AVX = _mm256_set1_ps(-0.0f); +#endif // __AVX__ + +#if defined(__AVX512F__) +static const __m512 NEGZEROS_FP32_AVX512 = _mm512_set1_ps(-0.0f); +#endif // __AVX512F__ + +//! Reverse sign of value (GENERAL) +#define NEGATE_FP32_GENERAL(v) -(v) + +#define NEGATE_FP32_SSE(v, ...) _mm_xor_ps(v, NEGZEROS_FP32_SSE) + +//! Reverse sign of value (AVX) +#define NEGATE_FP32_AVX(v, ...) _mm256_xor_ps(v, NEGZEROS_FP32_AVX) + +//! Reverse sign of value (AVX512) +#define NEGATE_FP32_AVX512(v, ...) _mm512_xor_ps(v, NEGZEROS_FP32_AVX512) + +//! 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 (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); + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) + +#if defined(__AVX512FP16__) +//! Inner Product +static inline float InnerProductAVX512FP16(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); +} + +#endif + +//! 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) + ACCUM_FP16_1X1_NEON(m, q, dim, out, 0ull, ) +#elif defined(__AVX512FP16__) + *out = InnerProductAVX512FP16(m, q, dim); +#elif defined(__AVX512F__) + ACCUM_FP16_1X1_AVX512(m, q, dim, out, 0ull, ) +#else + ACCUM_FP16_1X1_AVX(m, q, dim, out, 0ull, ) +#endif +} + +//! 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) + ACCUM_FP16_1X1_NEON(m, q, dim, out, 0ull, NEGATE_FP32_GENERAL) +#elif defined(__AVX512FP16__) + *out = -InnerProductAVX512FP16(m, q, dim); +#elif defined(__AVX512F__) + ACCUM_FP16_1X1_AVX512(m, q, dim, out, 0ull, NEGATE_FP32_GENERAL) +#else + ACCUM_FP16_1X1_AVX(m, q, dim, out, 0ull, NEGATE_FP32_GENERAL) +#endif +} + +#if !defined(__ARM_NEON) +//! Compute the distance between matrix and query (FP16, M=2, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_2X1_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=2, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_2X2_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_4X1_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_4X2_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_4X4_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_8X1_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_8X2_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_8X4_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { + ACCUM_FP16_8X8_AVX(m, q, dim, out, ) +} + +//! Compute the distance between matrix and query (FP16, M=16, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X1_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X1_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X2_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X2_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X4_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X4_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X8_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X8_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=16) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X16_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_16X16_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X1_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X1_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X2_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X2_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X4_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X4_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X8_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X8_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=16) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X16_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X16_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=32) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X32_AVX512(m, q, dim, out, ) +#else + ACCUM_FP16_32X32_AVX(m, q, dim, out, ) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=2, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_2X1_AVX(m, q, dim, out, NEGATE_FP32_SSE) +} + +//! Compute the distance between matrix and query (FP16, M=2, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_2X2_AVX(m, q, dim, out, NEGATE_FP32_SSE) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_4X1_AVX(m, q, dim, out, NEGATE_FP32_SSE) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_4X2_AVX(m, q, dim, out, NEGATE_FP32_SSE) +} + +//! Compute the distance between matrix and query (FP16, M=4, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_4X4_AVX(m, q, dim, out, NEGATE_FP32_SSE) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_8X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_8X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_8X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +} + +//! Compute the distance between matrix and query (FP16, M=8, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { + ACCUM_FP16_8X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +} + +//! Compute the distance between matrix and query (FP16, M=16, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X1_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_16X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X2_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_16X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X4_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_16X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X8_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_16X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=16, N=16) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_16X16_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_16X16_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X1_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_32X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X2_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_32X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X4_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_32X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X8_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_32X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=16) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X16_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_32X16_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} + +//! Compute the distance between matrix and query (FP16, M=32, N=32) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX512F__) + ACCUM_FP16_32X32_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#else + ACCUM_FP16_32X32_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#endif // __AVX512F__ +} +#endif // !__ARM_NEON +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +// 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 InnerProductSparseInSegmentAVX(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; +} + +#elif defined(__AVX512FP16__) +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 InnerProductSparseInSegmentAVX512FP16(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])); + __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; +} + +#else +float InnerProductSparseInSegment(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; +} +#endif // __AVX512FP16__ + +template <> +float MinusInnerProductSparseMatrix:: + 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(__AVX__) + return InnerProductSparseInSegmentAVX(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +#elif defined(__AVX512FP16__) + return InnerProductSparseInSegmentAVX512FP16(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +#else + return InnerProductSparseInSegment(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +#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..a93ba9b --- /dev/null +++ b/src/ailego/math/inner_product_matrix_fp32.cc @@ -0,0 +1,1192 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "inner_product_matrix.h" + +#if DEBUG_PRINT +#include + +static void inline print_data16(__m128i *data) { + uint16_t buffer[16]; + + memcpy(buffer, data, sizeof(buffer)); + + // std::cout << "result equals: " << std::endl; + + for (int i = 0; i < 8; i++) { + uint64_t value = buffer[i]; + if (i < 8) { + std::cout << value << ", "; + } else { + std::cout << value << std::endl; + } + } +} +#endif + +namespace zvec { +namespace ailego { + +#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 + +#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__ + +#if defined(__SSE__) +static const __m128 NEGZEROS_FP32_SSE = _mm_set1_ps(-0.0f); +#endif // __SSE__ + +#if defined(__AVX__) +static const __m256 NEGZEROS_FP32_AVX = _mm256_set1_ps(-0.0f); +#endif // __AVX__ + +#if defined(__AVX512F__) +static const __m512 NEGZEROS_FP32_AVX512 = _mm512_set1_ps(-0.0f); +#endif // __AVX512F__ + +//! Reverse sign of value (SSE) +#define NEGATE_FP32_SSE(v, ...) _mm_xor_ps(v, NEGZEROS_FP32_SSE) + +//! Reverse sign of value (AVX) +#define NEGATE_FP32_AVX(v, ...) _mm256_xor_ps(v, NEGZEROS_FP32_AVX) + +//! Reverse sign of value (AVX512) +#define NEGATE_FP32_AVX512(v, ...) _mm512_xor_ps(v, NEGZEROS_FP32_AVX512) + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_FP32_GENERAL(m, q, sum) sum += (m * q); + +//! 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 (NEON) +#define FMA_FP32_NEON(v_m, v_q, v_sum) v_sum = vfmaq_f32(v_sum, v_m, v_q); + +#if defined(__ARM_NEON) +//! Inner Product +static inline float InnerProductNEON(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; +} +#endif // __ARM_NEON + +#if defined(__SSE__) +//! Inner Product +static inline float InnerProductSSE(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; +} + +#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 InnerProductSparseInSegmentSSE(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))); + } + + float __attribute__((aligned(16))) 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 InnerProductSparseInSegment(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__ + +template <> +float MinusInnerProductSparseMatrix::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 InnerProductSparseInSegmentSSE(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +#else + return InnerProductSparseInSegment(m_sparse_count, m_sparse_index, + m_sparse_value, q_sparse_count, + q_sparse_index, q_sparse_value); +#endif +} + +#if defined(__AVX__) +//! Inner Product +static inline float InnerProductAVX(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; +} +#endif // __AVX__ + +#if defined(__AVX512F__) +//! Inner Product +static inline float InnerProductAVX512(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); +} +#endif + +#if defined(__SSE__) || defined(__ARM_NEON) +//! Compute the distance between matrix and query (FP32, M=1, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + *out = InnerProductNEON(m, q, dim); +#else +#if defined(__AVX512F__) + if (dim > 15) { + *out = InnerProductAVX512(m, q, dim); + return; + } +#endif // __AVX512F__ +#if defined(__AVX__) + if (dim > 7) { + *out = InnerProductAVX(m, q, dim); + return; + } +#endif // __AVX__ + *out = InnerProductSSE(m, q, dim); +#endif // __ARM_NEON +} + +//! Compute the distance between matrix and query (FP32, M=2, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_2X1_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_2X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_2X1_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=2, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_2X2_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_2X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_2X2_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X1_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_4X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_4X1_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X2_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_4X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_4X2_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X4_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_4X4_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_4X4_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X1_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_8X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_8X1_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X2_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_8X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_8X2_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X4_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_8X4_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_8X4_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X8_NEON(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_8X8_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_8X8_SSE(m, q, dim, out, ) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=16, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X1_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X1_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X1_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X2_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X2_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X2_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X4_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X4_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X4_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X4_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X8_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X8_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X8_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X8_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=16) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X16_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_16X16_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_16X16_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_16X16_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X1_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X1_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X1_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X1_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X2_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X2_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X2_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X2_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X4_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X4_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X4_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X4_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X8_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X8_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X8_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X8_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=16) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X16_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X16_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X16_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X16_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=32) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X32_NEON(m, q, dim, out, ) +#elif defined(__AVX512F__) + ACCUM_FP32_32X32_AVX512(m, q, dim, out, ) +#elif defined(__AVX__) + ACCUM_FP32_32X32_AVX(m, q, dim, out, ) +#else + ACCUM_FP32_32X32_SSE(m, q, dim, out, ) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=1, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + *out = -InnerProductNEON(m, q, dim); +#else +#if defined(__AVX512F__) + if (dim > 15) { + *out = -InnerProductAVX512(m, q, dim); + return; + } +#endif // __AVX512F__ +#if defined(__AVX__) + if (dim > 7) { + *out = -InnerProductAVX(m, q, dim); + return; + } +#endif // __AVX__ + *out = -InnerProductSSE(m, q, dim); +#endif // __ARM_NEON +} + +//! Compute the distance between matrix and query (FP32, M=2, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_2X1_NEON(m, q, dim, out, vneg_f32) +#elif defined(__AVX__) + ACCUM_FP32_2X1_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_FP32_2X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=2, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_2X2_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX__) + ACCUM_FP32_2X2_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_FP32_2X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X1_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX__) + ACCUM_FP32_4X1_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_FP32_4X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X2_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX__) + ACCUM_FP32_4X2_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_FP32_4X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=4, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_4X4_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX__) + ACCUM_FP32_4X4_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_FP32_4X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X1_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX__) + ACCUM_FP32_8X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_8X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X2_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX__) + ACCUM_FP32_8X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_8X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X4_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX__) + ACCUM_FP32_8X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_8X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=8, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_8X8_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX__) + ACCUM_FP32_8X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_8X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX__ +} + +//! Compute the distance between matrix and query (FP32, M=16, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X1_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X1_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_16X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_16X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X2_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X2_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_16X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_16X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X4_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X4_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_16X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_16X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X8_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X8_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_16X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_16X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=16, N=16) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_16X16_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_16X16_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_16X16_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_16X16_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X1_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X1_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_32X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_32X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X2_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X2_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_32X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_32X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X4_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X4_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_32X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_32X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X8_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X8_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_32X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_32X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=16) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X16_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X16_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_32X16_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_32X16_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} + +//! Compute the distance between matrix and query (FP32, M=32, N=32) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__ARM_NEON) + ACCUM_FP32_32X32_NEON(m, q, dim, out, vnegq_f32) +#elif defined(__AVX512F__) + ACCUM_FP32_32X32_AVX512(m, q, dim, out, NEGATE_FP32_AVX512) +#elif defined(__AVX__) + ACCUM_FP32_32X32_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_FP32_32X32_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif +} +#endif // __SSE__ || __ARM_NEON + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_int4.cc b/src/ailego/math/inner_product_matrix_int4.cc new file mode 100644 index 0000000..2ae5eb4 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_int4.cc @@ -0,0 +1,802 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +#define ACCUM_INT4_STEP_SSE FMA_INT4_SSE +#define ACCUM_INT4_STEP_AVX FMA_INT4_AVX + +#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 __m128 NEGZEROS_FP32_SSE = _mm_set1_ps(-0.0f); +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 __m256 NEGZEROS_FP32_AVX = _mm256_set1_ps(-0.0f); +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(m, q, sum) \ + sum += Int4MulTable[(((m) << 4) & 0xf0) | (((q) >> 0) & 0xf)] + \ + Int4MulTable[(((m) >> 0) & 0xf0) | (((q) >> 4) & 0xf)]; + +//! 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); \ + } + +//! Reverse sign of value (SSE) +#define NEGATE_FP32_SSE(v, ...) \ + _mm_xor_ps(_mm_cvtepi32_ps(v), NEGZEROS_FP32_SSE) + +//! Reverse sign of value (AVX) +#define NEGATE_FP32_AVX(v, ...) \ + _mm256_xor_ps(_mm256_cvtepi32_ps(v), NEGZEROS_FP32_AVX) + +//! Reverse sign of value (AVX512) +#define NEGATE_FP32_AVX512(v, ...) \ + _mm512_xor_ps(_mm512_cvtepi32_ps(v), NEGZEROS_FP32_AVX512) + +#if defined(__SSE4_1__) +//! Inner Product +static inline float InnerProductSSE(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; +} +#endif // __SSE4_1__ + +#if defined(__AVX2__) +//! Inner Product +static inline float InnerProductAVX(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; +} +#endif // __AVX2__ + +#if defined(__SSE4_1__) +//! Compute the distance between matrix and query (INT4, M=1, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + if (dim > 63) { + *out = InnerProductAVX(m, q, dim >> 1); + return; + } +#endif // __AVX2__ + *out = InnerProductSSE(m, q, dim >> 1); +} + +//! Compute the distance between matrix and query (INT4, M=2, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_2X1_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_2X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=2, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_2X2_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_2X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X1_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_4X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X2_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_4X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X4_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT4_4X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_8X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_8X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_8X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_8X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=16) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X16_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_16X16_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=16) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X16_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X16_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=32) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X32_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT4_32X32_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=1, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + if (dim > 63) { + *out = -InnerProductAVX(m, q, dim >> 1); + return; + } +#endif // __AVX2__ + *out = -InnerProductSSE(m, q, dim >> 1); +} + +//! Compute the distance between matrix and query (INT4, M=2, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_2X1_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT4_2X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=2, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_2X2_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT4_2X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X1_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT4_4X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X2_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT4_4X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=4, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_4X4_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT4_4X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_8X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_8X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_8X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=8, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_8X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_8X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_16X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_16X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_16X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_16X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=16, N=16) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_16X16_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_16X16_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_32X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_32X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_32X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_32X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=16) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X16_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_32X16_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT4, M=32, N=32) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT4_32X32_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT4_32X32_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/inner_product_matrix_int8.cc b/src/ailego/math/inner_product_matrix_int8.cc new file mode 100644 index 0000000..a307f69 --- /dev/null +++ b/src/ailego/math/inner_product_matrix_int8.cc @@ -0,0 +1,841 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "inner_product_matrix.h" + +namespace zvec { +namespace ailego { + +#define ACCUM_INT8_STEP_SSE FMA_INT8_SSE +#define ACCUM_INT8_STEP_AVX FMA_INT8_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__ + +#if defined(__SSE__) +static const __m128 NEGZEROS_FP32_SSE = _mm_set1_ps(-0.0f); +#endif // __SSE__ + +#if defined(__AVX__) +static const __m256 NEGZEROS_FP32_AVX = _mm256_set1_ps(-0.0f); +#endif // __AVX__ + +#if defined(__AVX512F__) +static const __m512 NEGZEROS_FP32_AVX512 = _mm512_set1_ps(-0.0f); +#endif // __AVX512F__ + +#if defined(__SSE4_1__) +static const __m128i ONES_INT16_SSE = _mm_set1_epi32(0x00010001); +#endif // __SSE4_1__ + +#if defined(__AVX2__) +static const __m256i ONES_INT16_AVX = _mm256_set1_epi32(0x00010001); +#endif // __AVX2__ + +//! Reverse sign of value (SSE) +#define NEGATE_FP32_SSE(v, ...) \ + _mm_xor_ps(_mm_cvtepi32_ps(v), NEGZEROS_FP32_SSE) + +//! Reverse sign of value (AVX) +#define NEGATE_FP32_AVX(v, ...) \ + _mm256_xor_ps(_mm256_cvtepi32_ps(v), NEGZEROS_FP32_AVX) + +//! Reverse sign of value (AVX512) +#define NEGATE_FP32_AVX512(v, ...) \ + _mm512_xor_ps(_mm512_cvtepi32_ps(v), NEGZEROS_FP32_AVX512) + +//! 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); + +#if defined(__SSE4_1__) +//! Inner Product +static inline float InnerProductSSE(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; +} +#endif // __SSE4_1__ + +#if defined(__AVX2__) +//! Inner Product +static inline float InnerProductAVX(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; +} +#endif // __AVX2__ + +#if defined(__SSE4_1__) +//! Compute the distance between matrix and query (INT8, M=1, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + if (dim > 31) { + *out = InnerProductAVX(m, q, dim); + return; + } +#endif // __AVX2__ + *out = InnerProductSSE(m, q, dim); +} + +//! Compute the distance between matrix and query (INT8, M=2, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_2X1_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_2X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=2, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_2X2_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_2X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X1_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_4X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X2_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_4X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X4_AVX(m, q, dim, out, _mm_cvtepi32_ps) +#else + ACCUM_INT8_4X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_8X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_8X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_8X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_8X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=16) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X16_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_16X16_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=1) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X1_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X1_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=2) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X2_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X2_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=4) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X4_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X4_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=8) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X8_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X8_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=16) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X16_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X16_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=32) +void InnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, size_t dim, + float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X32_AVX(m, q, dim, out, _mm256_cvtepi32_ps) +#else + ACCUM_INT8_32X32_SSE(m, q, dim, out, _mm_cvtepi32_ps) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=1, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + if (dim > 31) { + *out = -InnerProductAVX(m, q, dim); + return; + } +#endif // __AVX2__ + *out = -InnerProductSSE(m, q, dim); +} + +//! Compute the distance between matrix and query (INT8, M=2, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_2X1_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT8_2X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=2, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_2X2_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT8_2X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X1_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT8_4X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X2_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT8_4X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=4, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_4X4_AVX(m, q, dim, out, NEGATE_FP32_SSE) +#else + ACCUM_INT8_4X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_8X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_8X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_8X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=8, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_8X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_8X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_16X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_16X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_16X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_16X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=16, N=16) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_16X16_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_16X16_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=1) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X1_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_32X1_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=2) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X2_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_32X2_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=4) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X4_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_32X4_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=8) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X8_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_32X8_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=16) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X16_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_32X16_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} + +//! Compute the distance between matrix and query (INT8, M=32, N=32) +void MinusInnerProductMatrix::Compute(const ValueType *m, + const ValueType *q, + size_t dim, float *out) { +#if defined(__AVX2__) + ACCUM_INT8_32X32_AVX(m, q, dim, out, NEGATE_FP32_AVX) +#else + ACCUM_INT8_32X32_SSE(m, q, dim, out, NEGATE_FP32_SSE) +#endif // __AVX2__ +} +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file 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..c8c4297 --- /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. + +#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__) +static const __m256i POPCNT_MASK1_INT8_AVX = _mm256_set1_epi8(0x0f); +static const __m256i POPCNT_MASK1_INT16_AVX = _mm256_set1_epi16(1); +static const __m256i POPCNT_MASK2_INT16_AVX = _mm256_set1_epi16(0xff); +static const __m256i POPCNT_MASK1_INT32_AVX = _mm256_set1_epi32(0xff); +static const __m256i POPCNT_ZERO_AVX = _mm256_setzero_si256(); +static const __m256i 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 \ No newline at end of file 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..162228c --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix.h @@ -0,0 +1,1191 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "distance_utility.h" + +namespace zvec { +namespace ailego { + +/*! 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; + } +}; + +/*! 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; + } + + float u2[M]; + float v2[N]; + 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; + } + + float u2[M]; + float v2[N]; + 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; + } + + float u2[M]; + 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; + } + + float u2[M]; + 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; + } + + float u2[M]; + float v2[N]; + 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; + } + + float u2[M]; + float v2[N]; + 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; + } + + float u2[M]; + 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; + } + + float u2[M]; + 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; + } + + float u2[M]; + float v2[N]; + 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; + } + + float u2[M]; + float v2[N]; + 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; + } +}; + +#if !defined(__SSE4_1__) +/*! Mips Squared Euclidean Distance Matrix (INT4, M=1, N=1) + */ +template <> +struct MipsSquaredEuclideanDistanceMatrix { + //! 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 & 1) && out); + + 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)]; + } + *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 && !(dim & 1) && out); + + 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; + } + *out = sum; + } + + 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)); + } +}; +#endif // !__SSE4_1__ + +/*! 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; + } + + float u2[M]; + 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; + } + + float u2[M]; + 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; + } +}; + +#if defined(__SSE__) || defined(__ARM_NEON) +/*! Mips Squared Euclidean Distance Matrix (FP32, M=1, N=1) + */ +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); +}; +#endif // __SSE__ || __ARM_NEON + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +/*! Mips Squared Euclidean Distance Matrix (FP16, M=1, N=1) + */ +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); +}; +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +#if defined(__SSE4_1__) +/*! Mips Squared Euclidean Distance Matrix (INT8, M=1, N=1) + */ +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); +}; + +/*! Mips Squared Euclidean Distance Matrix (INT4, M=1, N=1) + */ +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); +}; +#endif + +/*! 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; +} + +#if defined(__SSE4_1__) +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); + +#endif + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16.cc new file mode 100644 index 0000000..f4ede59 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16.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 "distance_matrix_accum_fp16.i" +#include "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +//! 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; \ + } + +#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 +static inline float InnerProductAndSquaredNormNEON(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 +static inline float InnerProductAndSquaredNormNEON(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 +#endif // __ARM_NEON && __aarch64__ + +#if defined(__AVX__) && defined(__F16C__) +#if defined(__AVX512F__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormAVX512(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; +} +#else +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormAVX(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; +} +#endif // __AVX512F__ +#endif // __AVX__ && __F16C__ + +#if (defined(__F16C__) && defined(__AVX__)) || \ + (defined(__ARM_NEON) && defined(__aarch64__)) +//! 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) { + float u2; + float v2; + float sum; + +#if defined(__ARM_NEON) + sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2); +#elif defined(__AVX512F__) + sum = InnerProductAndSquaredNormAVX512(p, q, dim, &u2, &v2); +#else + sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2); +#endif + + *out = ComputeSphericalInjection(sum, u2, v2, 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) { + float u2; + float v2; + float sum; + +#if defined(__ARM_NEON) + sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2); +#elif defined(__AVX512F__) + sum = InnerProductAndSquaredNormAVX512(p, q, dim, &u2, &v2); +#else + sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2); +#endif + + 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; +} +#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec \ No newline at end of file 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..f14117a --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32.cc @@ -0,0 +1,684 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +//! Calculate Fused-Multiply-Add (GENERAL) +#define FMA_FP32_GENERAL(lhs, rhs, sum, norm1, norm2) \ + { \ + sum += (lhs) * (rhs); \ + norm1 += (lhs) * (lhs); \ + norm2 += (rhs) * (rhs); \ + } + +//! 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); + +#if defined(__ARM_NEON) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormNEON(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 + +#if defined(__SSE__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormSSE(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; +} +#endif // __SSE__ + +#if defined(__AVX__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormAVX(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; +} +#endif // __AVX__ + +#if defined(__AVX512F__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormAVX512(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); +} +#endif // __AVX512F__ + +#if defined(__SSE__) +//! 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) { + float u2; + float v2; + float sum; + +#if defined(__AVX512F__) + if (dim > 15) { + sum = InnerProductAndSquaredNormAVX512(p, q, dim, &u2, &v2); + } else +#endif // __AVX512F__ +#if defined(__AVX__) + if (dim > 7) { + sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2); + } else +#endif // __AVX__ + { + sum = InnerProductAndSquaredNormSSE(p, q, dim, &u2, &v2); + } + + *out = ComputeSphericalInjection(sum, u2, v2, 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) { + float u2; + float v2; + float sum; + +#if defined(__AVX512F__) + if (dim > 15) { + sum = InnerProductAndSquaredNormAVX512(p, q, dim, &u2, &v2); + } else +#endif // __AVX512F__ +#if defined(__AVX__) + if (dim > 7) { + sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2); + } else +#endif // __AVX__ + { + sum = InnerProductAndSquaredNormSSE(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; +} +#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))); + } + + float __attribute__((aligned(16))) 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__ + +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 +} + +#if defined(__ARM_NEON) +//! 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) { + float u2; + float v2; + float sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2); + + *out = ComputeSphericalInjection(sum, u2, v2, 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) { + float u2; + float v2; + float sum = InnerProductAndSquaredNormNEON(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; +} +#endif //__ARM_NEON + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int4.cc b/src/ailego/math/mips_euclidean_distance_matrix_int4.cc new file mode 100644 index 0000000..dc74b2d --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_int4.cc @@ -0,0 +1,358 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "inner_product_matrix.h" +#include "mips_euclidean_distance_matrix.h" +#include "norm_matrix.h" + +namespace zvec { +namespace ailego { + +#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)); \ + } + +//! 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) + +//! 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); \ + } + +//! 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); \ + } + +#if defined(__SSE4_1__) +#if defined(__AVX2__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormAVX(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; +} +#else +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormSSE(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; +} +#endif // __AVX2__ + +//! 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) { + float u2; + float v2; + float sum; + +#if defined(__AVX2__) + sum = InnerProductAndSquaredNormAVX(p, q, dim >> 1, &u2, &v2); +#else + sum = InnerProductAndSquaredNormSSE(p, q, dim >> 1, &u2, &v2); +#endif + + *out = ComputeSphericalInjection(sum, u2, v2, 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) { + float u2; + float v2; + float sum; + +#if defined(__AVX2__) + sum = InnerProductAndSquaredNormAVX(p, q, dim >> 1, &u2, &v2); +#else + sum = InnerProductAndSquaredNormSSE(p, q, dim >> 1, &u2, &v2); +#endif + + 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; +} +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int8.cc b/src/ailego/math/mips_euclidean_distance_matrix_int8.cc new file mode 100644 index 0000000..accf3d1 --- /dev/null +++ b/src/ailego/math/mips_euclidean_distance_matrix_int8.cc @@ -0,0 +1,358 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "mips_euclidean_distance_matrix.h" + +namespace zvec { +namespace ailego { + +#if defined(__SSE4_1__) +static const __m128i ONES_INT16_SSE = _mm_set1_epi32(0x00010001); +#endif // __SSE4_1__ + +#if defined(__AVX2__) +static const __m256i ONES_INT16_AVX = _mm256_set1_epi32(0x00010001); +#endif // __AVX2__ + +//! 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 (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) + +//! 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) + +#if defined(__SSE4_1__) +#if defined(__AVX2__) +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormAVX(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; +} +#else +//! Compute the Inner Product between p and q, and each Squared L2-Norm value +static inline float InnerProductAndSquaredNormSSE(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; +} +#endif // __AVX2__ + +//! 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) { + float u2; + float v2; + float sum; + +#if defined(__AVX2__) + sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2); +#else + sum = InnerProductAndSquaredNormSSE(p, q, dim, &u2, &v2); +#endif + + *out = ComputeSphericalInjection(sum, u2, v2, 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) { + float u2; + float v2; + float sum; + +#if defined(__AVX2__) + sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2); +#else + sum = InnerProductAndSquaredNormSSE(p, q, dim, &u2, &v2); +#endif + + 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; +} +#endif // __SSE4_1__ + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/norm1_matrix.h b/src/ailego/math/norm1_matrix.h new file mode 100644 index 0000000..3e23fa4 --- /dev/null +++ b/src/ailego/math/norm1_matrix.h @@ -0,0 +1,148 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! 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..2238605 --- /dev/null +++ b/src/ailego/math/norm1_matrix_fp16.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 "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, ) +#elif defined(__AVX512F__) + NORM_FP16_1_AVX512(m, dim, out, ) +#else + 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/norm1_matrix_fp32.cc b/src/ailego/math/norm1_matrix_fp32.cc new file mode 100644 index 0000000..de4205b --- /dev/null +++ b/src/ailego/math/norm1_matrix_fp32.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 "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__) +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); + +#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, ) +#elif defined(__AVX512F__) + NORM_FP32_1_AVX512(m, dim, out, ) +#elif defined(__AVX__) + NORM_FP32_1_AVX(m, dim, out, ) +#else + NORM_FP32_1_SSE(m, dim, out, ) +#endif +} +#endif // __SSE__ || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/math/norm2_matrix.h b/src/ailego/math/norm2_matrix.h new file mode 100644 index 0000000..9d976cf --- /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..5c0579e --- /dev/null +++ b/src/ailego/math/norm2_matrix_fp16.cc @@ -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. + +#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) +#elif defined(__AVX512F__) + NORM_FP16_1_AVX512(m, dim, out, std::sqrt) +#else + 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, ) +#elif defined(__AVX512F__) + NORM_FP16_1_AVX512(m, dim, out, ) +#else + 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..b5dcf2f --- /dev/null +++ b/src/ailego/math/norm2_matrix_fp32.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 "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) +#elif defined(__AVX512F__) + NORM_FP32_1_AVX512(m, dim, out, std::sqrt) +#elif defined(__AVX__) + NORM_FP32_1_AVX(m, dim, out, std::sqrt) +#else + 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, ) +#elif defined(__AVX512F__) + NORM_FP32_1_AVX512(m, dim, out, ) +#elif defined(__AVX__) + NORM_FP32_1_AVX(m, dim, out, ) +#else + NORM_FP32_1_SSE(m, dim, out, ) +#endif +} +#endif // __SSE__ || (__ARM_NEON && __aarch64__) + +} // namespace ailego +} // namespace zvec \ No newline at end of file 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..fec0742 --- /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_epi64((__m64)(0ull), *((const __m64 *)(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..345a879 --- /dev/null +++ b/src/ailego/math_batch/cosine_distance_batch.h @@ -0,0 +1,57 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "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..8301ac6 --- /dev/null +++ b/src/ailego/math_batch/distance_batch.h @@ -0,0 +1,51 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "cosine_distance_batch.h" +#include "inner_product_distance_batch.h" +#include "utils.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); + } + + _ComputeBatch(m, q, num, dim, out); + } +}; + +} // namespace zvec::ailego \ No newline at end of file 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..c8c5217 --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch.h @@ -0,0 +1,151 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "inner_product_distance_batch_impl.h" +#include "inner_product_distance_batch_impl_fp16.h" +#include "inner_product_distance_batch_impl_int8.h" +#include "utils.h" + +namespace zvec::ailego::DistanceBatch { + +template +struct InnerProductDistanceBatch; + +// 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_fallback(query, ptrs, prefetch_ptrs, dim, sums); + } + static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc() { + return nullptr; + } +}; + +template +struct InnerProductDistanceBatchImpl { + using ValueType = float; + static void compute_one_to_many( + const ValueType *query, const ValueType **ptrs, + std::array &prefetch_ptrs, size_t dim, + float *sums) { +#if defined(__AVX2__) + return compute_one_to_many_avx2_fp32( + query, ptrs, prefetch_ptrs, dim, sums); +#else + return compute_one_to_many_fallback(query, ptrs, prefetch_ptrs, dim, sums); +#endif + } + + static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc() { + return nullptr; + } +}; + +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) { +// #if defined(__AVX512BW__) // TODO: this version is problematic +// return compute_one_to_many_avx512_int8( +// query, ptrs, prefetch_ptrs, dim, sums); +#if defined(__AVX512VNNI__) + return compute_one_to_many_avx512_vnni_int8( + query, ptrs, prefetch_ptrs, dim, sums); +#elif defined(__AVX2__) + return compute_one_to_many_avx2_int8( + query, ptrs, prefetch_ptrs, dim, sums); +#else + return compute_one_to_many_fallback(query, ptrs, prefetch_ptrs, dim, sums); +#endif + } + + static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc() { +#if defined(__AVX512VNNI__) + return compute_one_to_many_avx512_vnni_int8_query_preprocess; +#endif + return nullptr; + } +}; + +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) { +#if defined(__AVX512FP16__) + return compute_one_to_many_avx512fp16_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +#elif defined(__AVX512F__) + return compute_one_to_many_avx512f_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +#elif defined(__AVX2__) + return compute_one_to_many_avx_fp16( + query, ptrs, prefetch_ptrs, dim, sums); +#else + return compute_one_to_many_fallback(query, ptrs, prefetch_ptrs, dim, sums); +#endif + } +}; + +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(); + } +}; + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl.h b/src/ailego/math_batch/inner_product_distance_batch_impl.h new file mode 100644 index 0000000..bfe8fc6 --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl.h @@ -0,0 +1,129 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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::ailego::DistanceBatch { + +template +static void compute_one_to_many_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]); + } +} + +#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 v[0] + 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_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; + } + std::array res; + 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]; + } +} +#endif + + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_fp16.h b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16.h new file mode 100644 index 0000000..64bd1ae --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16.h @@ -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. + +#pragma once + +#include +#include +#include +#include +#include + +namespace zvec::ailego::DistanceBatch { + +#if defined(__AVX512FP16__) + +template +static std::enable_if_t, void> +compute_one_to_many_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]); + } +} + +#elif defined(__AVX512F__) + +template +static std::enable_if_t, void> +compute_one_to_many_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)); + } + } +} + +#elif defined(__AVX2__) + +template +static std::enable_if_t, void> +compute_one_to_many_av2_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)); + } + } +} + +#endif + + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_int8.h b/src/ailego/math_batch/inner_product_distance_batch_impl_int8.h new file mode 100644 index 0000000..e37388f --- /dev/null +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_int8.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 + +namespace zvec::ailego::DistanceBatch { + + +#if defined(__AVX512VNNI__) + +static void compute_one_to_many_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_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]); + } + } + 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) { + 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) { +// __m512i accs[dp_batch]; +// 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); +// } +// } +// __m512i q_lo = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(q, 0)); +// __m512i q_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(q, 1)); +// __m512i data_lo[dp_batch]; +// __m512i data_hi[dp_batch]; +// 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)); +// } +// __m512i prod_lo[dp_batch]; +// __m512i prod_hi[dp_batch]; +// 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]); +// } +// } + +#elif defined(__AVX2__) + +template +static std::enable_if_t, void> +compute_one_to_many_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])); + } + } + std::array temp_results; + 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]); + } +} + +#endif + + +} // namespace zvec::ailego::DistanceBatch \ No newline at end of file diff --git a/src/ailego/math_batch/utils.h b/src/ailego/math_batch/utils.h new file mode 100644 index 0000000..15f37a4 --- /dev/null +++ b/src/ailego/math_batch/utils.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 +#include + +namespace zvec::ailego::DistanceBatch { + +typedef void (*DistanceBatchQueryPreprocessFunc)(void *query, size_t dim); + +} // 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..7bbd457 --- /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..263a1f4 --- /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..fff3788 --- /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 "thread_pool.h" + +#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/parallel/thread_pool.h b/src/ailego/parallel/thread_pool.h new file mode 100644 index 0000000..8bf56c0 --- /dev/null +++ b/src/ailego/parallel/thread_pool.h @@ -0,0 +1,402 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 ailego { + +/*! Thread Pool + */ +class ThreadPool { + public: + /*! Thread Pool Task Group + */ + class TaskGroup : public std::enable_shared_from_this { + public: + using Pointer = std::shared_ptr; + + //! Constructor + TaskGroup(ThreadPool *pool) : pool_(pool) {} + + //! Push a task to the queue + void enqueue(const ClosureHandler &handle) { + pool_->enqueue(handle, this->shared_from_this(), nullptr); + } + + //! Push a task to the queue + void enqueue(ClosureHandler &&handle) { + pool_->enqueue(std::move(handle), this->shared_from_this(), nullptr); + } + + //! Submit a task to the queue + void submit(ClosureHandler &&handle) { + return enqueue_and_wake(std::move(handle)); + } + + //! Push a task to the queue + void enqueue_and_wake(const ClosureHandler &handle) { + pool_->enqueue_and_wake(handle, this->shared_from_this(), nullptr); + } + + //! Push a task to the queue + void enqueue_and_wake(ClosureHandler &&handle) { + pool_->enqueue_and_wake(std::move(handle), this->shared_from_this(), + nullptr); + } + + //! Execute a function as a task in pool + template + void execute_and_wait(TArgs &&...args) { + ThreadPool::TaskControl ctrl; + pool_->enqueue_and_wake(Closure::New(std::forward(args)...), + this->shared_from_this(), &ctrl); + ctrl.wait(); + } + + //! Execute a function as a task in pool + template + void execute(TArgs &&...args) { + this->enqueue_and_wake(Closure::New(std::forward(args)...)); + } + + //! Wait until all tasks in group finished + void wait_finish(void) { + std::unique_lock lock(mutex_); + cond_.wait(lock, [this]() { return this->is_finished(); }); + } + + //! Check if the group is finished + bool is_finished(void) const { + return (active_count_ == 0 && pending_count_ == 0); + } + + //! Retrieve count of pending tasks in group + size_t pending_count(void) const { + return pending_count_.load(std::memory_order_relaxed); + } + + //! Retrieve count of active tasks in group + size_t active_count(void) const { + return active_count_.load(std::memory_order_relaxed); + } + + protected: + friend class ThreadPool; + + //! Mark a task enqueued + void mark_task_enqueued(void) { + ++pending_count_; + } + + //! Mark a task actived + void mark_task_actived(void) { + std::lock_guard lock(mutex_); + ++active_count_; + --pending_count_; + } + + //! Notify a task finished + void notify(void) { + std::lock_guard lock(mutex_); + if (--active_count_ == 0 && pending_count_ == 0) { + cond_.notify_all(); + } + } + + private: + //! Members + ThreadPool *pool_{nullptr}; + std::atomic_uint active_count_{0}; + std::atomic_uint pending_count_{0}; + std::mutex mutex_{}; + std::condition_variable cond_{}; + }; + + //! Constructor + explicit ThreadPool(uint32_t size, bool binding); + + //! Constructor + explicit ThreadPool(bool binding) + : ThreadPool{std::max(std::thread::hardware_concurrency(), 1u), binding} { + } + + //! Constructor + ThreadPool(void) : ThreadPool{false} {} + + //! Destructor + ~ThreadPool(void) { + this->stop(); + + // Join all threads + for (auto it = pool_.begin(); it != pool_.end(); ++it) { + if (it->joinable()) { + it->join(); + } + } + } + + //! Retrieve thread count in pool + size_t count(void) const { + return pool_.size(); + } + + //! Stop all threads + void stop(void) { + // Set stop flag as ture, then wake all threads + stopping_ = true; + std::lock_guard lock(queue_mutex_); + work_cond_.notify_all(); + } + + //! Push a task to the queue + void enqueue(const ClosureHandler &handle) { + this->enqueue(handle, nullptr); + } + + //! Push a task to the queue + void enqueue(ClosureHandler &&handle) { + this->enqueue(std::move(handle), nullptr); + } + + //! Push a task to the queue + void enqueue_and_wake(const ClosureHandler &handle) { + this->enqueue_and_wake(handle, nullptr); + } + + //! Push a task to the queue + void enqueue_and_wake(ClosureHandler &&handle) { + this->enqueue_and_wake(std::move(handle), nullptr); + } + + //! Execute a function as a task in pool + template + void execute_and_wait(TArgs &&...args) { + ThreadPool::TaskControl ctrl; + this->enqueue_and_wake(Closure::New(std::forward(args)...), &ctrl); + ctrl.wait(); + } + + //! Execute a function as a task in pool + template + void execute(TArgs &&...args) { + this->enqueue_and_wake(Closure::New(std::forward(args)...)); + } + + //! Wake any one thread + void wake_any(void) { + std::lock_guard lock(queue_mutex_); + work_cond_.notify_one(); + } + + //! Wake all threads + void wake_all(void) { + std::lock_guard lock(queue_mutex_); + work_cond_.notify_all(); + } + + //! Wait until all threads finished processing + void wait_finish(void) { + std::unique_lock lock(wait_mutex_); + finished_cond_.wait(lock, [this]() { return this->is_finished(); }); + } + + //! Wait until all threads stopped processing + void wait_stop(void) { + std::unique_lock lock(wait_mutex_); + stopped_cond_.wait(lock, [this]() { return this->is_stopped(); }); + } + + //! Make a task group + TaskGroup::Pointer make_group(void) { + return std::make_shared(this); + } + + //! Check if the pool is finished + bool is_finished(void) const { + return (active_count_ == 0 && pending_count_ == 0); + } + + //! Check if the pool is stopped + bool is_stopped(void) const { + return (worker_count_ == 0); + } + + //! Retrieve count of worker in pool + size_t worker_count(void) const { + return worker_count_.load(std::memory_order_relaxed); + } + + //! Retrieve count of pending tasks in pool + size_t pending_count(void) const { + return pending_count_.load(std::memory_order_relaxed); + } + + //! Retrieve count of active tasks in pool + size_t active_count(void) const { + return active_count_.load(std::memory_order_relaxed); + } + + //! Get the thread index via thread id + int indexof(const std::thread::id &thread_id) const { + for (size_t i = 0; i < pool_.size(); ++i) { + if (pool_[i].get_id() == thread_id) { + return static_cast(i); + } + } + return -1; + } + + //! Get the current work thread index + int indexof_this(void) const { + return this->indexof(std::this_thread::get_id()); + } + + //! Bind threads to processors + void bind(void); + + //! Unbind threads of processors + void unbind(void); + + protected: + //! Thread task control + class TaskControl { + public: + //! Notify task finished + void notify(void) { + finished_ = true; + std::lock_guard lock(mutex_); + cond_.notify_one(); + } + + //! Wait until task finished + void wait(void) { + std::unique_lock lock(mutex_); + cond_.wait(lock, [this]() { return finished_.load(); }); + } + + private: + std::atomic_bool finished_{false}; + std::mutex mutex_{}; + std::condition_variable cond_{}; + }; + + //! Thread task + struct Task { + // Constructor + Task(const ClosureHandler &h, TaskControl *c) : handle(h), control(c) {} + + // Constructor + Task(ClosureHandler &&h, TaskControl *c) + : handle(std::move(h)), control(c) {} + + // Constructor + Task(const ClosureHandler &h, TaskGroup::Pointer &&g, TaskControl *c) + : handle(h), group(std::move(g)), control(c) {} + + // Constructor + Task(ClosureHandler &&h, TaskGroup::Pointer &&g, TaskControl *c) + : handle(std::move(h)), group(std::move(g)), control(c) {} + + // Constructor + Task(void) {} + + //! Members + ClosureHandler handle{}; + TaskGroup::Pointer group{nullptr}; + TaskControl *control{nullptr}; + }; + + //! Thread worker callback + void worker(void); + + //! Pick a task from queue + bool picking(Task *task); + + //! Push a task to the queue + template + void enqueue(T &&handle, TaskControl *ctrl) { + if (handle) { + std::lock_guard lock(queue_mutex_); + ++pending_count_; + queue_.emplace(std::forward(handle), ctrl); + } + } + + //! Push a task to the queue with group + template + void enqueue(T &&handle, TaskGroup::Pointer &&group, TaskControl *ctrl) { + if (handle) { + std::lock_guard lock(queue_mutex_); + ++pending_count_; + group->mark_task_enqueued(); + queue_.emplace(std::forward(handle), std::move(group), ctrl); + } + } + + //! Push a task to the queue + template + void enqueue_and_wake(T &&handle, TaskControl *ctrl) { + if (handle) { + std::lock_guard lock(queue_mutex_); + ++pending_count_; + queue_.emplace(std::forward(handle), ctrl); + work_cond_.notify_one(); + } + } + + //! Push a task to the queue with group + template + void enqueue_and_wake(T &&handle, TaskGroup::Pointer &&group, + TaskControl *ctrl) { + if (handle) { + std::lock_guard lock(queue_mutex_); + ++pending_count_; + group->mark_task_enqueued(); + queue_.emplace(std::forward(handle), std::move(group), ctrl); + work_cond_.notify_one(); + } + } + + private: + //! Disable them + ThreadPool(const ThreadPool &) = delete; + ThreadPool(ThreadPool &&) = delete; + ThreadPool &operator=(const ThreadPool &) = delete; + + //! Members + std::queue queue_{}; + std::atomic_bool stopping_{false}; + std::atomic_uint worker_count_{0}; + std::atomic_uint active_count_{0}; + std::atomic_uint pending_count_{0}; + std::mutex queue_mutex_{}; + std::mutex wait_mutex_{}; + std::condition_variable work_cond_{}; + std::condition_variable finished_cond_{}; + std::condition_variable stopped_cond_{}; + std::vector pool_{}; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/parallel/thread_queue.h b/src/ailego/parallel/thread_queue.h new file mode 100644 index 0000000..3940e72 --- /dev/null +++ b/src/ailego/parallel/thread_queue.h @@ -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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! Thread Queue (One Thread One Queue) + */ +class ThreadQueue { + public: + /*! Thread Worker (One Thread One Worker) + */ + class ThreadWorker { + public: + //! Constructor + ThreadWorker(ThreadQueue *owner) : owner_(owner) {} + + //! Destructor + ~ThreadWorker(void) { + // Join the current thread + if (thread_.joinable()) { + thread_.join(); + } + } + + //! Push a task to the queue + template + void enqueue(T &&handle) { + std::lock_guard lock(mutex_); + queue_.emplace(std::forward(handle)); + } + + //! Push a task to the queue + template + void enqueue_and_wake(T &&handle) { + std::lock_guard lock(mutex_); + queue_.emplace(std::forward(handle)); + cond_.notify_one(); + } + + //! Execute a function as a task + template + void execute(TArgs &&...args) { + this->enqueue_and_wake(Closure::New(std::forward(args)...)); + } + + //! Push a task to the queue with high priority + template + void enqueue_high_priority(T &&handle) { + std::lock_guard lock(mutex_); + queue_high_priority_.emplace(std::forward(handle)); + } + + //! Push a task to the queue with high priority and wake + template + void enqueue_high_priority_and_wake(T &&handle) { + std::lock_guard lock(mutex_); + queue_high_priority_.emplace(std::forward(handle)); + cond_.notify_one(); + } + + //! Execute a task in high priority + template + void execute_high_priority(TArgs &&...args) { + this->enqueue_and_wake(Closure::New(std::forward(args)...)); + } + + //! Wake the thread + void wake(void) { + std::lock_guard lock(mutex_); + cond_.notify_one(); + } + + //! Notify thread stopped + void stop(void) { + // Set stop flag as ture, then wake the thread + stopping_ = true; + std::lock_guard lock(mutex_); + cond_.notify_one(); + } + + protected: + //! Thread worker callback + void worker(void) { + owner_->mark_worker_started(); + + ClosureHandler task; + while (this->picking(&task)) { + // Run the task + if (task) { + task->run(); + task = nullptr; + } + } + owner_->mark_worker_stopped(); + } + + //! Pick a task from queue + bool picking(ClosureHandler *task) { + std::unique_lock latch(mutex_); + cond_.wait(latch, [this]() { + return (queue_.size() > 0 || queue_high_priority_.size() > 0 || + stopping_); + }); + if (stopping_) { + return false; + } + + if (!queue_high_priority_.empty()) { + *task = std::move(queue_high_priority_.front()); + queue_high_priority_.pop(); + } else { + *task = std::move(queue_.front()); + queue_.pop(); + } + + return true; + } + + private: + //! Disable them + ThreadWorker(void) = delete; + ThreadWorker(ThreadWorker &&) = delete; + ThreadWorker(const ThreadWorker &) = delete; + ThreadWorker &operator=(const ThreadWorker &) = delete; + + //! Members + ThreadQueue *owner_{nullptr}; + std::queue queue_{}; + std::queue queue_high_priority_{}; + std::atomic_bool stopping_{false}; + std::mutex mutex_{}; + std::condition_variable cond_{}; + std::thread thread_{&ThreadWorker::worker, this}; + }; + + //! Constructor + ThreadQueue(void) + : ThreadQueue{std::max(std::thread::hardware_concurrency(), 1u)} {} + + //! Constructor + explicit ThreadQueue(uint32_t size) { + for (uint32_t i = 0u; i < size; ++i) { + threads_.emplace_back(new ThreadWorker(this)); + } + } + + //! Destructor + ~ThreadQueue(void) { + this->stop(); + // Cleanup threads + for (auto it = threads_.begin(); it != threads_.end(); ++it) { + delete *it; + } + } + + //! operator [] + ThreadWorker &operator[](size_t i) { + return *(threads_[i]); + } + + //! Stop the thread + void stop(void) { + // Stop all workers + for (auto it = threads_.begin(); it != threads_.end(); ++it) { + (*it)->stop(); + } + } + + //! Wake all worker threads + void wake(void) { + for (auto it = threads_.begin(); it != threads_.end(); ++it) { + (*it)->wake(); + } + } + + //! Wait until all threads stopped processing + void wait_stop(void) { + std::unique_lock lock(wait_mutex_); + stopped_cond_.wait(lock, [this]() { return this->is_stopped(); }); + } + + //! Check if the pool is stopped + bool is_stopped(void) const { + return (worker_count_ == 0); + } + + //! Retrieve count of worker in queue + size_t worker_count(void) const { + return worker_count_.load(std::memory_order_relaxed); + } + + //! Retrieve thread count in queue + size_t count(void) const { + return threads_.size(); + } + + //! Push a task to the queue + template + void enqueue(uint64_t key, T &&handle) { + threads_[JumpHash(key, static_cast(threads_.size()))]->enqueue( + std::forward(handle)); + } + + //! Push a task to the queue + template + void enqueue_and_wake(uint64_t key, T &&handle) { + threads_[JumpHash(key, static_cast(threads_.size()))] + ->enqueue_and_wake(std::forward(handle)); + } + + //! Execute a function as a task in pool + template + void execute(uint64_t key, TArgs &&...args) { + this->enqueue_and_wake(key, Closure::New(std::forward(args)...)); + } + + //! Push a task to the queue with high priority + template + void enqueue_high_priority(uint64_t key, T &&handle) { + threads_[JumpHash(key, static_cast(threads_.size()))] + ->enqueue_high_priority(std::forward(handle)); + } + + //! Push a task to the queue with high priority and wake + template + void enqueue_high_priority_and_wake(uint64_t key, T &&handle) { + threads_[JumpHash(key, static_cast(threads_.size()))] + ->enqueue_high_priority_and_wake(std::forward(handle)); + } + + //! Execute a function as a task in pool with high priority + template + void execute_high_priority(uint64_t key, TArgs &&...args) { + this->enqueue_high_priority_and_wake( + key, Closure::New(std::forward(args)...)); + } + + protected: + //! Mark a worker started + void mark_worker_started(void) { + ++worker_count_; + } + + //! Mark a worker stopped + void mark_worker_stopped(void) { + // Decrease count of workers + std::lock_guard lock(wait_mutex_); + if (--worker_count_ == 0) { + stopped_cond_.notify_all(); + } + } + + private: + //! Disable them + ThreadQueue(const ThreadQueue &) = delete; + ThreadQueue(ThreadQueue &&) = delete; + ThreadQueue &operator=(const ThreadQueue &) = delete; + + //! Members + std::atomic_uint worker_count_{0}; + std::mutex wait_mutex_{}; + std::condition_variable stopped_cond_{}; + std::vector threads_{}; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/pattern/closure.h b/src/ailego/pattern/closure.h new file mode 100644 index 0000000..df624ac --- /dev/null +++ b/src/ailego/pattern/closure.h @@ -0,0 +1,530 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Callback Validator (declaration) + */ +template +struct CallbackValidator; + +/*! Callback Validator (function pointer) + */ +template +struct CallbackValidator { + enum { Value = true }; +}; + +/*! Callback Validator (function) + */ +template +struct CallbackValidator : CallbackValidator { +}; + +/*! Callback Validator (member function pointer) + */ +template +struct CallbackValidator + : CallbackValidator {}; + +/*! Callback Validator (constable member function pointer) + */ +template +struct CallbackValidator + : CallbackValidator {}; + +/*! Callback Validator (volatile member function pointer) + */ +template +struct CallbackValidator + : CallbackValidator {}; + +/*! Callback Validator (constable volatile member function pointer) + */ +template +struct CallbackValidator + : CallbackValidator {}; + +/*! Callback Validator + */ +template +struct CallbackValidator { + protected: + using FalseType = long; + using TrueType = char; + + //! Check if the class contains operator() + template + static TrueType &Validate(decltype(&T::operator())); + + //! Check if the class contains operator() + template + static FalseType &Validate(...); + + public: + enum { Value = (sizeof(Validate(nullptr)) == sizeof(TrueType)) }; +}; + +/*! Callback Validator (left reference) + */ +template +struct CallbackValidator : CallbackValidator {}; + +/*! Callback Validator (right reference) + */ +template +struct CallbackValidator : CallbackValidator {}; + +/*! Callback Traits (declaration) + */ +template +struct CallbackTraits; + +/*! Callback Traits (function pointer) + */ +template +struct CallbackTraits { + using Type = R (*)(TParams...); + using ResultType = R; + using TupleType = std::tuple::type...>; + + //! Callback Traits Parameter + template + struct Parameter { + using Type = typename std::tuple_element>::type; + }; + + //! Number of parameters + enum { Arity = sizeof...(TParams) }; +}; + +/*! Callback Traits (function) + */ +template +struct CallbackTraits : CallbackTraits { + using Type = R (*)(TParams...); +}; + +/*! Callback Traits (member function pointer) + */ +template +struct CallbackTraits + : CallbackTraits { + using Type = R (T::*)(TParams...); +}; + +/*! Callback Traits (constable member function pointer) + */ +template +struct CallbackTraits + : CallbackTraits { + using Type = R (T::*)(TParams...) const; +}; + +/*! Callback Traits (volatile member function pointer) + */ +template +struct CallbackTraits + : CallbackTraits { + using Type = R (T::*)(TParams...) volatile; +}; + +/*! Callback Traits (constable volatile member function pointer) + */ +template +struct CallbackTraits + : CallbackTraits { + using Type = R (T::*)(TParams...) const volatile; +}; + +/*! Callback Traits + */ +template +struct CallbackTraits : CallbackTraits { + using Type = TFunc; +}; + +/*! Callback Traits (left reference) + */ +template +struct CallbackTraits : CallbackTraits {}; + +/*! Callback Traits (right reference) + */ +template +struct CallbackTraits : CallbackTraits {}; + +/*! Callback Functor + */ +template +struct CallbackFunctor { + using Traits = CallbackTraits; + using Type = typename Traits::Type; + using ResultType = typename Traits::ResultType; + using TupleType = typename Traits::TupleType; + + //! Tuple Index Maker + template + struct TupleIndexMaker : TupleIndexMaker {}; + + //! Tuple Index + template + struct TupleIndex {}; + + //! Tuple Index Maker (special) + template + struct TupleIndexMaker<0, I...> { + using Type = TupleIndex; + }; + + //! Run the callback function + template + static ResultType Run(Type &impl, TupleType &tuple, TupleIndex) { + return (impl)(std::forward::Type>( + std::get(tuple))...); + } + + //! Run the callback member function + template + static ResultType Run(T *obj, Type &impl, TupleType &tuple, + TupleIndex) { + return (obj->*impl)( + std::forward::Type>( + std::get(tuple))...); + } + + //! Run the callback function + static ResultType Run(Type &impl, TupleType &tuple) { + return Run(impl, tuple, typename TupleIndexMaker::Type()); + } + + //! Run the callback member function + template + static ResultType Run(T *obj, Type &impl, TupleType &tuple) { + return Run(obj, impl, tuple, + typename TupleIndexMaker::Type()); + } +}; + +/*! Callback Object + */ +template +struct CallbackObject { + using Type = typename std::remove_reference::type; +}; + +/*! Callback (declaration) + */ +template +class Callback; + +/*! Callback (void) + */ +template <> +class Callback { + public: + using Pointer = std::shared_ptr>; + + //! Destructor + virtual ~Callback(void) {} + + //! Function call + void operator()(void) { + this->run(); + } + + //! Run the callback function + virtual void run(void) = 0; + + //! Create callback closure (member function pointer) + template + static typename Callback::Pointer New(T *obj, R (T::*impl)(TParams...), + TArgs &&...args); + + //! Create callback closure (constable member function pointer) + template + static typename Callback::Pointer New(const T *obj, + R (T::*impl)(TParams...) const, + TArgs &&...args); + + //! Create callback closure (volatile member function pointer) + template + static typename Callback::Pointer New(volatile T *obj, + R (T::*impl)(TParams...) volatile, + TArgs &&...args); + + //! Create callback closure (constable volatile member function pointer) + template + static typename Callback::Pointer New(const volatile T *obj, + R (T::*impl)(TParams...) + const volatile, + TArgs &&...args); + + //! Create callback closure (function) + template < + typename TFunc, typename... TArgs, + typename = typename std::enable_if::Value>::type> + static typename Callback::ResultType>::Pointer + New(TFunc &&impl, TArgs &&...args); +}; + +/*! Callback + */ +template +class Callback : public Callback { + public: + using Pointer = std::shared_ptr>; + using Callback::run; + + //! Function call + void operator()(void) { + this->run(); + } + + //! Function call with return + void operator()(R *r) { + this->run(r); + } + + //! Run the callback function + virtual void run(R *) = 0; + + protected: + //! Constructor + Callback(void) {}; +}; + +/*! Callback Implementation + */ +template +class CallbackImpl : public Callback { + public: + using Object = CallbackObject; + using Functor = CallbackFunctor; + + //! Constructor + template + CallbackImpl(typename Object::Type *obj, const typename Functor::Type &impl, + TArgs &&...args) + : obj_(obj), impl_(impl), tuple_(std::forward(args)...) {} + + //! Constructor + template + CallbackImpl(typename Object::Type *obj, typename Functor::Type &&impl, + TArgs &&...args) + : obj_(obj), + impl_(std::move(impl)), + tuple_(std::forward(args)...) {} + + //! Run the callback function + void run(void) override { + Functor::Run(obj_, impl_, tuple_); + } + + //! Run the callback function + void run(R *r) override { + *r = Functor::Run(obj_, impl_, tuple_); + } + + protected: + //! Disable them + CallbackImpl(void) = delete; + CallbackImpl(const CallbackImpl &) = delete; + CallbackImpl(CallbackImpl &&) = delete; + CallbackImpl &operator=(const CallbackImpl &) = delete; + + private: + typename Object::Type *obj_; + typename Functor::Type impl_; + typename Functor::TupleType tuple_; +}; + +/*! Callback Implementation + */ +template +class CallbackImpl : public Callback { + public: + using Object = CallbackObject; + using Functor = CallbackFunctor; + + //! Constructor + template + CallbackImpl(typename Object::Type *obj, const typename Functor::Type &impl, + TArgs &&...args) + : obj_(obj), impl_(impl), tuple_(std::forward(args)...) {} + + //! Constructor + template + CallbackImpl(typename Object::Type *obj, typename Functor::Type &&impl, + TArgs &&...args) + : obj_(obj), + impl_(std::move(impl)), + tuple_(std::forward(args)...) {} + + //! Run the callback function + void run(void) override { + Functor::Run(obj_, impl_, tuple_); + } + + protected: + //! Disable them + CallbackImpl(void) = delete; + CallbackImpl(const CallbackImpl &) = delete; + CallbackImpl(CallbackImpl &&) = delete; + CallbackImpl &operator=(const CallbackImpl &) = delete; + + private: + typename Object::Type *obj_; + typename Functor::Type impl_; + typename Functor::TupleType tuple_; +}; + +/*! Callback Implementation + */ +template +class CallbackImpl : public Callback { + public: + using Functor = CallbackFunctor; + + //! Constructor + template + CallbackImpl(const typename Functor::Type &impl, TArgs &&...args) + : impl_(impl), tuple_(std::forward(args)...) {} + + //! Constructor + template + CallbackImpl(typename Functor::Type &&impl, TArgs &&...args) + : impl_(std::move(impl)), tuple_(std::forward(args)...) {} + + //! Run the callback function + void run(void) override { + Functor::Run(impl_, tuple_); + } + + //! Run the callback function + void run(R *r) override { + *r = Functor::Run(impl_, tuple_); + } + + protected: + //! Disable them + CallbackImpl(void) = delete; + CallbackImpl(const CallbackImpl &) = delete; + CallbackImpl(CallbackImpl &&) = delete; + CallbackImpl &operator=(const CallbackImpl &) = delete; + + private: + typename Functor::Type impl_; + typename Functor::TupleType tuple_; +}; + +/*! Callback Implementation + */ +template +class CallbackImpl : public Callback { + public: + using Functor = CallbackFunctor; + + //! Constructor + template + CallbackImpl(const typename Functor::Type &impl, TArgs &&...args) + : impl_(impl), tuple_(std::forward(args)...) {} + + //! Constructor + template + CallbackImpl(typename Functor::Type &&impl, TArgs &&...args) + : impl_(std::move(impl)), tuple_(std::forward(args)...) {} + + //! Run the callback function + void run(void) override { + Functor::Run(impl_, tuple_); + } + + protected: + //! Disable them + CallbackImpl(void) = delete; + CallbackImpl(const CallbackImpl &) = delete; + CallbackImpl(CallbackImpl &&) = delete; + CallbackImpl &operator=(const CallbackImpl &) = delete; + + private: + typename Functor::Type impl_; + typename Functor::TupleType tuple_; +}; + +//! Create callback closure (member function pointer) +template +typename Callback::Pointer Callback::New(T *obj, + R (T::*impl)(TParams...), + TArgs &&...args) { + return std::make_shared>( + obj, impl, std::forward(args)...); +} + +//! Create callback closure (constable member function pointer) +template +typename Callback::Pointer Callback::New(const T *obj, + R (T::*impl)(TParams...) + const, + TArgs &&...args) { + return std::make_shared>( + obj, impl, std::forward(args)...); +} + +//! Create callback closure (volatile member function pointer) +template +typename Callback::Pointer Callback::New( + volatile T *obj, R (T::*impl)(TParams...) volatile, TArgs &&...args) { + return std::make_shared>( + obj, impl, std::forward(args)...); +} + +//! Create callback closure (constable volatile member function pointer) +template +typename Callback::Pointer Callback::New(const volatile T *obj, + R (T::*impl)(TParams...) + const volatile, + TArgs &&...args) { + return std::make_shared>( + obj, impl, std::forward(args)...); +} + +//! Create callback closure (function) +template +typename Callback::ResultType>::Pointer +Callback::New(TFunc &&impl, TArgs &&...args) { + return std::make_shared::ResultType, decltype(impl)>>( + std::forward(impl), std::forward(args)...); +} + +//! Callback Handler +template +using CallbackHandler = typename Callback::Pointer; + +//! Closure +using Closure = Callback; + +//! Closure Handler +using ClosureHandler = Closure::Pointer; + +} // namespace ailego +} // namespace zvec 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/expected.hpp b/src/ailego/pattern/expected.hpp new file mode 100644 index 0000000..8f660e0 --- /dev/null +++ b/src/ailego/pattern/expected.hpp @@ -0,0 +1,2605 @@ +/// +// expected - An implementation of std::expected with extensions +// Written in 2017 by Sy Brand (tartanllama@gmail.com, @TartanLlama) +// +// Documentation available at http://tl.tartanllama.xyz/ +// +// To the extent possible under law, the author(s) have dedicated all +// copyright and related and neighboring rights to this software to the +// public domain worldwide. This software is distributed without any warranty. +// +// You should have received a copy of the CC0 Public Domain Dedication +// along with this software. If not, see +// . +/// + +#ifndef TL_EXPECTED_HPP +#define TL_EXPECTED_HPP + +#define TL_EXPECTED_VERSION_MAJOR 1 +#define TL_EXPECTED_VERSION_MINOR 2 +#define TL_EXPECTED_VERSION_PATCH 0 + +#include +#include +#include +#include + +#if defined(__EXCEPTIONS) || defined(_CPPUNWIND) +#define TL_EXPECTED_EXCEPTIONS_ENABLED +#endif + +#if (defined(_MSC_VER) && _MSC_VER == 1900) +#define TL_EXPECTED_MSVC2015 +#define TL_EXPECTED_MSVC2015_CONSTEXPR +#else +#define TL_EXPECTED_MSVC2015_CONSTEXPR constexpr +#endif + +#if (defined(__GNUC__) && __GNUC__ == 4 && __GNUC_MINOR__ <= 9 && \ + !defined(__clang__)) +#define TL_EXPECTED_GCC49 +#endif + +#if (defined(__GNUC__) && __GNUC__ == 5 && __GNUC_MINOR__ <= 4 && \ + !defined(__clang__)) +#define TL_EXPECTED_GCC54 +#endif + +#if (defined(__GNUC__) && __GNUC__ == 5 && __GNUC_MINOR__ <= 5 && \ + !defined(__clang__)) +#define TL_EXPECTED_GCC55 +#endif + +#if !defined(TL_ASSERT) +// can't have assert in constexpr in C++11 and GCC 4.9 has a compiler bug +#if (TL_CPLUSPLUS > 201103L) && !defined(TL_EXPECTED_GCC49) +#include +#define TL_ASSERT(x) assert(x) +#else +#define TL_ASSERT(x) +#endif +#endif + +#if (defined(__GNUC__) && __GNUC__ == 4 && __GNUC_MINOR__ <= 9 && \ + !defined(__clang__)) +// GCC < 5 doesn't support overloading on const&& for member functions + +#define TL_EXPECTED_NO_CONSTRR +// GCC < 5 doesn't support some standard C++11 type traits +#define TL_EXPECTED_IS_TRIVIALLY_COPY_CONSTRUCTIBLE(T) \ + std::has_trivial_copy_constructor +#define TL_EXPECTED_IS_TRIVIALLY_COPY_ASSIGNABLE(T) \ + std::has_trivial_copy_assign + +// This one will be different for GCC 5.7 if it's ever supported +#define TL_EXPECTED_IS_TRIVIALLY_DESTRUCTIBLE(T) \ + std::is_trivially_destructible + +// GCC 5 < v < 8 has a bug in is_trivially_copy_constructible which breaks +// std::vector for non-copyable types +#elif (defined(__GNUC__) && __GNUC__ < 8 && !defined(__clang__)) +#ifndef TL_GCC_LESS_8_TRIVIALLY_COPY_CONSTRUCTIBLE_MUTEX +#define TL_GCC_LESS_8_TRIVIALLY_COPY_CONSTRUCTIBLE_MUTEX +namespace tl { +namespace detail { +template +struct is_trivially_copy_constructible + : std::is_trivially_copy_constructible {}; +#ifdef _GLIBCXX_VECTOR +template +struct is_trivially_copy_constructible> : std::false_type {}; +#endif +} // namespace detail +} // namespace tl +#endif + +#define TL_EXPECTED_IS_TRIVIALLY_COPY_CONSTRUCTIBLE(T) \ + tl::detail::is_trivially_copy_constructible +#define TL_EXPECTED_IS_TRIVIALLY_COPY_ASSIGNABLE(T) \ + std::is_trivially_copy_assignable +#define TL_EXPECTED_IS_TRIVIALLY_DESTRUCTIBLE(T) \ + std::is_trivially_destructible +#else +#define TL_EXPECTED_IS_TRIVIALLY_COPY_CONSTRUCTIBLE(T) \ + std::is_trivially_copy_constructible +#define TL_EXPECTED_IS_TRIVIALLY_COPY_ASSIGNABLE(T) \ + std::is_trivially_copy_assignable +#define TL_EXPECTED_IS_TRIVIALLY_DESTRUCTIBLE(T) \ + std::is_trivially_destructible +#endif + +#ifdef _MSVC_LANG +#define TL_CPLUSPLUS _MSVC_LANG +#else +#define TL_CPLUSPLUS __cplusplus +#endif + +#if TL_CPLUSPLUS > 201103L +#define TL_EXPECTED_CXX14 +#endif + +#ifdef TL_EXPECTED_GCC49 +#define TL_EXPECTED_GCC49_CONSTEXPR +#else +#define TL_EXPECTED_GCC49_CONSTEXPR constexpr +#endif + +#if (TL_CPLUSPLUS == 201103L || defined(TL_EXPECTED_MSVC2015) || \ + defined(TL_EXPECTED_GCC49)) +#define TL_EXPECTED_11_CONSTEXPR +#else +#define TL_EXPECTED_11_CONSTEXPR constexpr +#endif + +#if TL_CPLUSPLUS >= 201703L +#define TL_EXPECTED_NODISCARD [[nodiscard]] +#else +#define TL_EXPECTED_NODISCARD +#endif + +namespace tl { +template +class TL_EXPECTED_NODISCARD expected; + +#ifndef TL_MONOSTATE_INPLACE_MUTEX +#define TL_MONOSTATE_INPLACE_MUTEX +class monostate {}; + +struct in_place_t { + explicit in_place_t() = default; +}; +static constexpr in_place_t in_place{}; +#endif + +template +class unexpected { + public: + static_assert(!std::is_same::value, "E must not be void"); + + unexpected() = delete; + constexpr explicit unexpected(const E &e) : m_val(e) {} + + constexpr explicit unexpected(E &&e) : m_val(std::move(e)) {} + + template ::value>::type * = nullptr> + constexpr explicit unexpected(Args &&...args) + : m_val(std::forward(args)...) {} + template < + class U, class... Args, + typename std::enable_if &, Args &&...>::value>::type * = nullptr> + constexpr explicit unexpected(std::initializer_list l, Args &&...args) + : m_val(l, std::forward(args)...) {} + + constexpr const E &value() const & { + return m_val; + } + TL_EXPECTED_11_CONSTEXPR E &value() & { + return m_val; + } + TL_EXPECTED_11_CONSTEXPR E &&value() && { + return std::move(m_val); + } + constexpr const E &&value() const && { + return std::move(m_val); + } + + private: + E m_val; +}; + +#ifdef __cpp_deduction_guides +template +unexpected(E) -> unexpected; +#endif + +template +constexpr bool operator==(const unexpected &lhs, const unexpected &rhs) { + return lhs.value() == rhs.value(); +} +template +constexpr bool operator!=(const unexpected &lhs, const unexpected &rhs) { + return lhs.value() != rhs.value(); +} +template +constexpr bool operator<(const unexpected &lhs, const unexpected &rhs) { + return lhs.value() < rhs.value(); +} +template +constexpr bool operator<=(const unexpected &lhs, const unexpected &rhs) { + return lhs.value() <= rhs.value(); +} +template +constexpr bool operator>(const unexpected &lhs, const unexpected &rhs) { + return lhs.value() > rhs.value(); +} +template +constexpr bool operator>=(const unexpected &lhs, const unexpected &rhs) { + return lhs.value() >= rhs.value(); +} + +template +unexpected::type> make_unexpected(E &&e) { + return unexpected::type>(std::forward(e)); +} + +struct unexpect_t { + unexpect_t() = default; +}; +static constexpr unexpect_t unexpect{}; + +namespace detail { +template +[[noreturn]] TL_EXPECTED_11_CONSTEXPR void throw_exception(E &&e) { +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + throw std::forward(e); +#else + (void)e; +#ifdef _MSC_VER + __assume(0); +#else + __builtin_unreachable(); +#endif +#endif +} + +#ifndef TL_TRAITS_MUTEX +#define TL_TRAITS_MUTEX +// C++14-style aliases for brevity +template +using remove_const_t = typename std::remove_const::type; +template +using remove_reference_t = typename std::remove_reference::type; +template +using decay_t = typename std::decay::type; +template +using enable_if_t = typename std::enable_if::type; +template +using conditional_t = typename std::conditional::type; + +// std::conjunction from C++17 +template +struct conjunction : std::true_type {}; +template +struct conjunction : B {}; +template +struct conjunction + : std::conditional, B>::type {}; + +#if defined(_LIBCPP_VERSION) && __cplusplus == 201103L +#define TL_TRAITS_LIBCXX_MEM_FN_WORKAROUND +#endif + +// In C++11 mode, there's an issue in libc++'s std::mem_fn +// which results in a hard-error when using it in a noexcept expression +// in some cases. This is a check to workaround the common failing case. +#ifdef TL_TRAITS_LIBCXX_MEM_FN_WORKAROUND +template +struct is_pointer_to_non_const_member_func : std::false_type {}; +template +struct is_pointer_to_non_const_member_func + : std::true_type {}; +template +struct is_pointer_to_non_const_member_func + : std::true_type {}; +template +struct is_pointer_to_non_const_member_func + : std::true_type {}; +template +struct is_pointer_to_non_const_member_func + : std::true_type {}; +template +struct is_pointer_to_non_const_member_func + : std::true_type {}; +template +struct is_pointer_to_non_const_member_func + : std::true_type {}; + +template +struct is_const_or_const_ref : std::false_type {}; +template +struct is_const_or_const_ref : std::true_type {}; +template +struct is_const_or_const_ref : std::true_type {}; +#endif + +// std::invoke from C++17 +// https://stackoverflow.com/questions/38288042/c11-14-invoke-workaround +template < + typename Fn, typename... Args, +#ifdef TL_TRAITS_LIBCXX_MEM_FN_WORKAROUND + typename = enable_if_t::value && + is_const_or_const_ref::value)>, +#endif + typename = enable_if_t>::value>, int = 0> +constexpr auto invoke(Fn &&f, Args &&...args) noexcept( + noexcept(std::mem_fn(f)(std::forward(args)...))) + -> decltype(std::mem_fn(f)(std::forward(args)...)) { + return std::mem_fn(f)(std::forward(args)...); +} + +template >::value>> +constexpr auto invoke(Fn &&f, Args &&...args) noexcept( + noexcept(std::forward(f)(std::forward(args)...))) + -> decltype(std::forward(f)(std::forward(args)...)) { + return std::forward(f)(std::forward(args)...); +} + +// std::invoke_result from C++17 +template +struct invoke_result_impl; + +template +struct invoke_result_impl< + F, + decltype(detail::invoke(std::declval(), std::declval()...), void()), + Us...> { + using type = + decltype(detail::invoke(std::declval(), std::declval()...)); +}; + +template +using invoke_result = invoke_result_impl; + +template +using invoke_result_t = typename invoke_result::type; + +#if defined(_MSC_VER) && _MSC_VER <= 1900 +// TODO make a version which works with MSVC 2015 +template +struct is_swappable : std::true_type {}; + +template +struct is_nothrow_swappable : std::true_type {}; +#else +// https://stackoverflow.com/questions/26744589/what-is-a-proper-way-to-implement-is-swappable-to-test-for-the-swappable-concept +namespace swap_adl_tests { +// if swap ADL finds this then it would call std::swap otherwise (same +// signature) +struct tag {}; + +template +tag swap(T &, T &); +template +tag swap(T (&a)[N], T (&b)[N]); + +// helper functions to test if an unqualified swap is possible, and if it +// becomes std::swap +template +std::false_type can_swap(...) noexcept(false); +template (), std::declval()))> +std::true_type can_swap(int) noexcept(noexcept(swap(std::declval(), + std::declval()))); + +template +std::false_type uses_std(...); +template +std::is_same(), std::declval())), tag> +uses_std(int); + +template +struct is_std_swap_noexcept + : std::integral_constant::value && + std::is_nothrow_move_assignable::value> {}; + +template +struct is_std_swap_noexcept : is_std_swap_noexcept {}; + +template +struct is_adl_swap_noexcept + : std::integral_constant(0))> {}; +} // namespace swap_adl_tests + +template +struct is_swappable + : std::integral_constant< + bool, + decltype(detail::swap_adl_tests::can_swap(0))::value && + (!decltype(detail::swap_adl_tests::uses_std(0))::value || + (std::is_move_assignable::value && + std::is_move_constructible::value))> {}; + +template +struct is_swappable + : std::integral_constant< + bool, + decltype(detail::swap_adl_tests::can_swap(0))::value && + (!decltype(detail::swap_adl_tests::uses_std( + 0))::value || + is_swappable::value)> {}; + +template +struct is_nothrow_swappable + : std::integral_constant< + bool, + is_swappable::value && + ((decltype(detail::swap_adl_tests::uses_std(0))::value && + detail::swap_adl_tests::is_std_swap_noexcept::value) || + (!decltype(detail::swap_adl_tests::uses_std(0))::value && + detail::swap_adl_tests::is_adl_swap_noexcept::value))> {}; +#endif +#endif + +// Trait for checking if a type is a tl::expected +template +struct is_expected_impl : std::false_type {}; +template +struct is_expected_impl> : std::true_type {}; +template +using is_expected = is_expected_impl>; + +template +using expected_enable_forward_value = detail::enable_if_t< + std::is_constructible::value && + !std::is_same, in_place_t>::value && + !std::is_same, detail::decay_t>::value && + !std::is_same, detail::decay_t>::value>; + +template +using expected_enable_from_other = detail::enable_if_t< + std::is_constructible::value && + std::is_constructible::value && + !std::is_constructible &>::value && + !std::is_constructible &&>::value && + !std::is_constructible &>::value && + !std::is_constructible &&>::value && + !std::is_convertible &, T>::value && + !std::is_convertible &&, T>::value && + !std::is_convertible &, T>::value && + !std::is_convertible &&, T>::value>; + +template +using is_void_or = conditional_t::value, std::true_type, U>; + +template +using is_copy_constructible_or_void = + is_void_or>; + +template +using is_move_constructible_or_void = + is_void_or>; + +template +using is_copy_assignable_or_void = is_void_or>; + +template +using is_move_assignable_or_void = is_void_or>; + +} // namespace detail + +namespace detail { +struct no_init_t {}; +static constexpr no_init_t no_init{}; + +// Implements the storage of the values, and ensures that the destructor is +// trivial if it can be. +// +// This specialization is for where neither `T` or `E` is trivially +// destructible, so the destructors must be called on destruction of the +// `expected` +template ::value, + bool = std::is_trivially_destructible::value> +struct expected_storage_base { + constexpr expected_storage_base() : m_val(T{}), m_has_val(true) {} + constexpr expected_storage_base(no_init_t) : m_no_init(), m_has_val(false) {} + + template ::value> * = + nullptr> + constexpr expected_storage_base(in_place_t, Args &&...args) + : m_val(std::forward(args)...), m_has_val(true) {} + + template &, Args &&...>::value> * = nullptr> + constexpr expected_storage_base(in_place_t, std::initializer_list il, + Args &&...args) + : m_val(il, std::forward(args)...), m_has_val(true) {} + template ::value> * = + nullptr> + constexpr explicit expected_storage_base(unexpect_t, Args &&...args) + : m_unexpect(std::forward(args)...), m_has_val(false) {} + + template &, Args &&...>::value> * = nullptr> + constexpr explicit expected_storage_base(unexpect_t, + std::initializer_list il, + Args &&...args) + : m_unexpect(il, std::forward(args)...), m_has_val(false) {} + + ~expected_storage_base() { + if (m_has_val) { + m_val.~T(); + } else { + m_unexpect.~unexpected(); + } + } + union { + T m_val; + unexpected m_unexpect; + char m_no_init; + }; + bool m_has_val; +}; + +// This specialization is for when both `T` and `E` are trivially-destructible, +// so the destructor of the `expected` can be trivial. +template +struct expected_storage_base { + constexpr expected_storage_base() : m_val(T{}), m_has_val(true) {} + constexpr expected_storage_base(no_init_t) : m_no_init(), m_has_val(false) {} + + template ::value> * = + nullptr> + constexpr expected_storage_base(in_place_t, Args &&...args) + : m_val(std::forward(args)...), m_has_val(true) {} + + template &, Args &&...>::value> * = nullptr> + constexpr expected_storage_base(in_place_t, std::initializer_list il, + Args &&...args) + : m_val(il, std::forward(args)...), m_has_val(true) {} + template ::value> * = + nullptr> + constexpr explicit expected_storage_base(unexpect_t, Args &&...args) + : m_unexpect(std::forward(args)...), m_has_val(false) {} + + template &, Args &&...>::value> * = nullptr> + constexpr explicit expected_storage_base(unexpect_t, + std::initializer_list il, + Args &&...args) + : m_unexpect(il, std::forward(args)...), m_has_val(false) {} + + expected_storage_base(const expected_storage_base &) = default; + expected_storage_base(expected_storage_base &&) = default; + expected_storage_base &operator=(const expected_storage_base &) = default; + expected_storage_base &operator=(expected_storage_base &&) = default; + ~expected_storage_base() = default; + union { + T m_val; + unexpected m_unexpect; + char m_no_init; + }; + bool m_has_val; +}; + +// T is trivial, E is not. +template +struct expected_storage_base { + constexpr expected_storage_base() : m_val(T{}), m_has_val(true) {} + TL_EXPECTED_MSVC2015_CONSTEXPR expected_storage_base(no_init_t) + : m_no_init(), m_has_val(false) {} + + template ::value> * = + nullptr> + constexpr expected_storage_base(in_place_t, Args &&...args) + : m_val(std::forward(args)...), m_has_val(true) {} + + template &, Args &&...>::value> * = nullptr> + constexpr expected_storage_base(in_place_t, std::initializer_list il, + Args &&...args) + : m_val(il, std::forward(args)...), m_has_val(true) {} + template ::value> * = + nullptr> + constexpr explicit expected_storage_base(unexpect_t, Args &&...args) + : m_unexpect(std::forward(args)...), m_has_val(false) {} + + template &, Args &&...>::value> * = nullptr> + constexpr explicit expected_storage_base(unexpect_t, + std::initializer_list il, + Args &&...args) + : m_unexpect(il, std::forward(args)...), m_has_val(false) {} + + expected_storage_base(const expected_storage_base &) = default; + expected_storage_base(expected_storage_base &&) = default; + expected_storage_base &operator=(const expected_storage_base &) = default; + expected_storage_base &operator=(expected_storage_base &&) = default; + ~expected_storage_base() { + if (!m_has_val) { + m_unexpect.~unexpected(); + } + } + + union { + T m_val; + unexpected m_unexpect; + char m_no_init; + }; + bool m_has_val; +}; + +// E is trivial, T is not. +template +struct expected_storage_base { + constexpr expected_storage_base() : m_val(T{}), m_has_val(true) {} + constexpr expected_storage_base(no_init_t) : m_no_init(), m_has_val(false) {} + + template ::value> * = + nullptr> + constexpr expected_storage_base(in_place_t, Args &&...args) + : m_val(std::forward(args)...), m_has_val(true) {} + + template &, Args &&...>::value> * = nullptr> + constexpr expected_storage_base(in_place_t, std::initializer_list il, + Args &&...args) + : m_val(il, std::forward(args)...), m_has_val(true) {} + template ::value> * = + nullptr> + constexpr explicit expected_storage_base(unexpect_t, Args &&...args) + : m_unexpect(std::forward(args)...), m_has_val(false) {} + + template &, Args &&...>::value> * = nullptr> + constexpr explicit expected_storage_base(unexpect_t, + std::initializer_list il, + Args &&...args) + : m_unexpect(il, std::forward(args)...), m_has_val(false) {} + + expected_storage_base(const expected_storage_base &) = default; + expected_storage_base(expected_storage_base &&) = default; + expected_storage_base &operator=(const expected_storage_base &) = default; + expected_storage_base &operator=(expected_storage_base &&) = default; + ~expected_storage_base() { + if (m_has_val) { + m_val.~T(); + } + } + union { + T m_val; + unexpected m_unexpect; + char m_no_init; + }; + bool m_has_val; +}; + +// `T` is `void`, `E` is trivially-destructible +template +struct expected_storage_base { +#if __GNUC__ <= 5 +// no constexpr for GCC 4/5 bug +#else + TL_EXPECTED_MSVC2015_CONSTEXPR +#endif + expected_storage_base() : m_has_val(true) {} + + constexpr expected_storage_base(no_init_t) : m_val(), m_has_val(false) {} + + constexpr expected_storage_base(in_place_t) : m_has_val(true) {} + + template ::value> * = + nullptr> + constexpr explicit expected_storage_base(unexpect_t, Args &&...args) + : m_unexpect(std::forward(args)...), m_has_val(false) {} + + template &, Args &&...>::value> * = nullptr> + constexpr explicit expected_storage_base(unexpect_t, + std::initializer_list il, + Args &&...args) + : m_unexpect(il, std::forward(args)...), m_has_val(false) {} + + expected_storage_base(const expected_storage_base &) = default; + expected_storage_base(expected_storage_base &&) = default; + expected_storage_base &operator=(const expected_storage_base &) = default; + expected_storage_base &operator=(expected_storage_base &&) = default; + ~expected_storage_base() = default; + struct dummy {}; + union { + unexpected m_unexpect; + dummy m_val; + }; + bool m_has_val; +}; + +// `T` is `void`, `E` is not trivially-destructible +template +struct expected_storage_base { + constexpr expected_storage_base() : m_dummy(), m_has_val(true) {} + constexpr expected_storage_base(no_init_t) : m_dummy(), m_has_val(false) {} + + constexpr expected_storage_base(in_place_t) : m_dummy(), m_has_val(true) {} + + template ::value> * = + nullptr> + constexpr explicit expected_storage_base(unexpect_t, Args &&...args) + : m_unexpect(std::forward(args)...), m_has_val(false) {} + + template &, Args &&...>::value> * = nullptr> + constexpr explicit expected_storage_base(unexpect_t, + std::initializer_list il, + Args &&...args) + : m_unexpect(il, std::forward(args)...), m_has_val(false) {} + + expected_storage_base(const expected_storage_base &) = default; + expected_storage_base(expected_storage_base &&) = default; + expected_storage_base &operator=(const expected_storage_base &) = default; + expected_storage_base &operator=(expected_storage_base &&) = default; + ~expected_storage_base() { + if (!m_has_val) { + m_unexpect.~unexpected(); + } + } + + union { + unexpected m_unexpect; + char m_dummy; + }; + bool m_has_val; +}; + +// This base class provides some handy member functions which can be used in +// further derived classes +template +struct expected_operations_base : expected_storage_base { + using expected_storage_base::expected_storage_base; + + template + void construct(Args &&...args) noexcept { + new (std::addressof(this->m_val)) T(std::forward(args)...); + this->m_has_val = true; + } + + template + void construct_with(Rhs &&rhs) noexcept { + new (std::addressof(this->m_val)) T(std::forward(rhs).get()); + this->m_has_val = true; + } + + template + void construct_error(Args &&...args) noexcept { + new (std::addressof(this->m_unexpect)) + unexpected(std::forward(args)...); + this->m_has_val = false; + } + +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + + // These assign overloads ensure that the most efficient assignment + // implementation is used while maintaining the strong exception guarantee. + // The problematic case is where rhs has a value, but *this does not. + // + // This overload handles the case where we can just copy-construct `T` + // directly into place without throwing. + template ::value> + * = nullptr> + void assign(const expected_operations_base &rhs) noexcept { + if (!this->m_has_val && rhs.m_has_val) { + geterr().~unexpected(); + construct(rhs.get()); + } else { + assign_common(rhs); + } + } + + // This overload handles the case where we can attempt to create a copy of + // `T`, then no-throw move it into place if the copy was successful. + template ::value && + std::is_nothrow_move_constructible::value> + * = nullptr> + void assign(const expected_operations_base &rhs) noexcept { + if (!this->m_has_val && rhs.m_has_val) { + T tmp = rhs.get(); + geterr().~unexpected(); + construct(std::move(tmp)); + } else { + assign_common(rhs); + } + } + + // This overload is the worst-case, where we have to move-construct the + // unexpected value into temporary storage, then try to copy the T into place. + // If the construction succeeds, then everything is fine, but if it throws, + // then we move the old unexpected value back into place before rethrowing the + // exception. + template ::value && + !std::is_nothrow_move_constructible::value> + * = nullptr> + void assign(const expected_operations_base &rhs) { + if (!this->m_has_val && rhs.m_has_val) { + auto tmp = std::move(geterr()); + geterr().~unexpected(); + +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + try { + construct(rhs.get()); + } catch (...) { + geterr() = std::move(tmp); + throw; + } +#else + construct(rhs.get()); +#endif + } else { + assign_common(rhs); + } + } + + // These overloads do the same as above, but for rvalues + template ::value> + * = nullptr> + void assign(expected_operations_base &&rhs) noexcept { + if (!this->m_has_val && rhs.m_has_val) { + geterr().~unexpected(); + construct(std::move(rhs).get()); + } else { + assign_common(std::move(rhs)); + } + } + + template ::value> + * = nullptr> + void assign(expected_operations_base &&rhs) { + if (!this->m_has_val && rhs.m_has_val) { + auto tmp = std::move(geterr()); + geterr().~unexpected(); +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + try { + construct(std::move(rhs).get()); + } catch (...) { + geterr() = std::move(tmp); + throw; + } +#else + construct(std::move(rhs).get()); +#endif + } else { + assign_common(std::move(rhs)); + } + } + +#else + + // If exceptions are disabled then we can just copy-construct + void assign(const expected_operations_base &rhs) noexcept { + if (!this->m_has_val && rhs.m_has_val) { + geterr().~unexpected(); + construct(rhs.get()); + } else { + assign_common(rhs); + } + } + + void assign(expected_operations_base &&rhs) noexcept { + if (!this->m_has_val && rhs.m_has_val) { + geterr().~unexpected(); + construct(std::move(rhs).get()); + } else { + assign_common(std::move(rhs)); + } + } + +#endif + + // The common part of move/copy assigning + template + void assign_common(Rhs &&rhs) { + if (this->m_has_val) { + if (rhs.m_has_val) { + get() = std::forward(rhs).get(); + } else { + destroy_val(); + construct_error(std::forward(rhs).geterr()); + } + } else { + if (!rhs.m_has_val) { + geterr() = std::forward(rhs).geterr(); + } + } + } + + bool has_value() const { + return this->m_has_val; + } + + TL_EXPECTED_11_CONSTEXPR T &get() & { + return this->m_val; + } + constexpr const T &get() const & { + return this->m_val; + } + TL_EXPECTED_11_CONSTEXPR T &&get() && { + return std::move(this->m_val); + } +#ifndef TL_EXPECTED_NO_CONSTRR + constexpr const T &&get() const && { + return std::move(this->m_val); + } +#endif + + TL_EXPECTED_11_CONSTEXPR unexpected &geterr() & { + return this->m_unexpect; + } + constexpr const unexpected &geterr() const & { + return this->m_unexpect; + } + TL_EXPECTED_11_CONSTEXPR unexpected &&geterr() && { + return std::move(this->m_unexpect); + } +#ifndef TL_EXPECTED_NO_CONSTRR + constexpr const unexpected &&geterr() const && { + return std::move(this->m_unexpect); + } +#endif + + TL_EXPECTED_11_CONSTEXPR void destroy_val() { + get().~T(); + } +}; + +// This base class provides some handy member functions which can be used in +// further derived classes +template +struct expected_operations_base : expected_storage_base { + using expected_storage_base::expected_storage_base; + + template + void construct() noexcept { + this->m_has_val = true; + } + + // This function doesn't use its argument, but needs it so that code in + // levels above this can work independently of whether T is void + template + void construct_with(Rhs &&) noexcept { + this->m_has_val = true; + } + + template + void construct_error(Args &&...args) noexcept { + new (std::addressof(this->m_unexpect)) + unexpected(std::forward(args)...); + this->m_has_val = false; + } + + template + void assign(Rhs &&rhs) noexcept { + if (!this->m_has_val) { + if (rhs.m_has_val) { + geterr().~unexpected(); + construct(); + } else { + geterr() = std::forward(rhs).geterr(); + } + } else { + if (!rhs.m_has_val) { + construct_error(std::forward(rhs).geterr()); + } + } + } + + bool has_value() const { + return this->m_has_val; + } + + TL_EXPECTED_11_CONSTEXPR unexpected &geterr() & { + return this->m_unexpect; + } + constexpr const unexpected &geterr() const & { + return this->m_unexpect; + } + TL_EXPECTED_11_CONSTEXPR unexpected &&geterr() && { + return std::move(this->m_unexpect); + } +#ifndef TL_EXPECTED_NO_CONSTRR + constexpr const unexpected &&geterr() const && { + return std::move(this->m_unexpect); + } +#endif + + TL_EXPECTED_11_CONSTEXPR void destroy_val() { + // no-op + } +}; + +// This class manages conditionally having a trivial copy constructor +// This specialization is for when T and E are trivially copy constructible +template ::value && + TL_EXPECTED_IS_TRIVIALLY_COPY_CONSTRUCTIBLE(E)::value, + bool = (is_copy_constructible_or_void::value && + std::is_copy_constructible::value)> +struct expected_copy_base : expected_operations_base { + using expected_operations_base::expected_operations_base; +}; + +// This specialization is for when T or E are non-trivially copy constructible +template +struct expected_copy_base : expected_operations_base { + using expected_operations_base::expected_operations_base; + + expected_copy_base() = default; + expected_copy_base(const expected_copy_base &rhs) + : expected_operations_base(no_init) { + if (rhs.has_value()) { + this->construct_with(rhs); + } else { + this->construct_error(rhs.geterr()); + } + } + + expected_copy_base(expected_copy_base &&rhs) = default; + expected_copy_base &operator=(const expected_copy_base &rhs) = default; + expected_copy_base &operator=(expected_copy_base &&rhs) = default; +}; + +// This class manages conditionally having a trivial move constructor +// Unfortunately there's no way to achieve this in GCC < 5 AFAIK, since it +// doesn't implement an analogue to std::is_trivially_move_constructible. We +// have to make do with a non-trivial move constructor even if T is trivially +// move constructible +#ifndef TL_EXPECTED_GCC49 +template >::value && + std::is_trivially_move_constructible::value> +struct expected_move_base : expected_copy_base { + using expected_copy_base::expected_copy_base; +}; +#else +template +struct expected_move_base; +#endif +template +struct expected_move_base : expected_copy_base { + using expected_copy_base::expected_copy_base; + + expected_move_base() = default; + expected_move_base(const expected_move_base &rhs) = default; + + expected_move_base(expected_move_base &&rhs) noexcept( + std::is_nothrow_move_constructible::value) + : expected_copy_base(no_init) { + if (rhs.has_value()) { + this->construct_with(std::move(rhs)); + } else { + this->construct_error(std::move(rhs.geterr())); + } + } + expected_move_base &operator=(const expected_move_base &rhs) = default; + expected_move_base &operator=(expected_move_base &&rhs) = default; +}; + +// This class manages conditionally having a trivial copy assignment operator +template < + class T, class E, + bool = + is_void_or< + T, conjunction>::value && + TL_EXPECTED_IS_TRIVIALLY_COPY_ASSIGNABLE(E)::value && + TL_EXPECTED_IS_TRIVIALLY_COPY_CONSTRUCTIBLE(E)::value && + TL_EXPECTED_IS_TRIVIALLY_DESTRUCTIBLE(E)::value, + bool = (is_copy_constructible_or_void::value && + std::is_copy_constructible::value && + is_copy_assignable_or_void::value && + std::is_copy_assignable::value)> +struct expected_copy_assign_base : expected_move_base { + using expected_move_base::expected_move_base; +}; + +template +struct expected_copy_assign_base : expected_move_base { + using expected_move_base::expected_move_base; + + expected_copy_assign_base() = default; + expected_copy_assign_base(const expected_copy_assign_base &rhs) = default; + + expected_copy_assign_base(expected_copy_assign_base &&rhs) = default; + expected_copy_assign_base &operator=(const expected_copy_assign_base &rhs) { + this->assign(rhs); + return *this; + } + expected_copy_assign_base &operator=(expected_copy_assign_base &&rhs) = + default; +}; + +// This class manages conditionally having a trivial move assignment operator +// Unfortunately there's no way to achieve this in GCC < 5 AFAIK, since it +// doesn't implement an analogue to std::is_trivially_move_assignable. We have +// to make do with a non-trivial move assignment operator even if T is trivially +// move assignable +#ifndef TL_EXPECTED_GCC49 +template < + class T, class E, + bool = is_void_or< + T, conjunction, + std::is_trivially_move_constructible, + std::is_trivially_move_assignable>>::value && + std::is_trivially_destructible::value && + std::is_trivially_move_constructible::value && + std::is_trivially_move_assignable::value> +struct expected_move_assign_base : expected_copy_assign_base { + using expected_copy_assign_base::expected_copy_assign_base; +}; +#else +template +struct expected_move_assign_base; +#endif + +template +struct expected_move_assign_base + : expected_copy_assign_base { + using expected_copy_assign_base::expected_copy_assign_base; + + expected_move_assign_base() = default; + expected_move_assign_base(const expected_move_assign_base &rhs) = default; + + expected_move_assign_base(expected_move_assign_base &&rhs) = default; + + expected_move_assign_base &operator=(const expected_move_assign_base &rhs) = + default; + + expected_move_assign_base &operator=( + expected_move_assign_base + &&rhs) noexcept(std::is_nothrow_move_constructible::value && + std::is_nothrow_move_assignable::value) { + this->assign(std::move(rhs)); + return *this; + } +}; + +// expected_delete_ctor_base will conditionally delete copy and move +// constructors depending on whether T is copy/move constructible +template ::value && + std::is_copy_constructible::value), + bool EnableMove = (is_move_constructible_or_void::value && + std::is_move_constructible::value)> +struct expected_delete_ctor_base { + expected_delete_ctor_base() = default; + expected_delete_ctor_base(const expected_delete_ctor_base &) = default; + expected_delete_ctor_base(expected_delete_ctor_base &&) noexcept = default; + expected_delete_ctor_base &operator=(const expected_delete_ctor_base &) = + default; + expected_delete_ctor_base &operator=(expected_delete_ctor_base &&) noexcept = + default; +}; + +template +struct expected_delete_ctor_base { + expected_delete_ctor_base() = default; + expected_delete_ctor_base(const expected_delete_ctor_base &) = default; + expected_delete_ctor_base(expected_delete_ctor_base &&) noexcept = delete; + expected_delete_ctor_base &operator=(const expected_delete_ctor_base &) = + default; + expected_delete_ctor_base &operator=(expected_delete_ctor_base &&) noexcept = + default; +}; + +template +struct expected_delete_ctor_base { + expected_delete_ctor_base() = default; + expected_delete_ctor_base(const expected_delete_ctor_base &) = delete; + expected_delete_ctor_base(expected_delete_ctor_base &&) noexcept = default; + expected_delete_ctor_base &operator=(const expected_delete_ctor_base &) = + default; + expected_delete_ctor_base &operator=(expected_delete_ctor_base &&) noexcept = + default; +}; + +template +struct expected_delete_ctor_base { + expected_delete_ctor_base() = default; + expected_delete_ctor_base(const expected_delete_ctor_base &) = delete; + expected_delete_ctor_base(expected_delete_ctor_base &&) noexcept = delete; + expected_delete_ctor_base &operator=(const expected_delete_ctor_base &) = + default; + expected_delete_ctor_base &operator=(expected_delete_ctor_base &&) noexcept = + default; +}; + +// expected_delete_assign_base will conditionally delete copy and move +// constructors depending on whether T and E are copy/move constructible + +// assignable +template ::value && + std::is_copy_constructible::value && + is_copy_assignable_or_void::value && + std::is_copy_assignable::value), + bool EnableMove = (is_move_constructible_or_void::value && + std::is_move_constructible::value && + is_move_assignable_or_void::value && + std::is_move_assignable::value)> +struct expected_delete_assign_base { + expected_delete_assign_base() = default; + expected_delete_assign_base(const expected_delete_assign_base &) = default; + expected_delete_assign_base(expected_delete_assign_base &&) noexcept = + default; + expected_delete_assign_base &operator=(const expected_delete_assign_base &) = + default; + expected_delete_assign_base &operator=( + expected_delete_assign_base &&) noexcept = default; +}; + +template +struct expected_delete_assign_base { + expected_delete_assign_base() = default; + expected_delete_assign_base(const expected_delete_assign_base &) = default; + expected_delete_assign_base(expected_delete_assign_base &&) noexcept = + default; + expected_delete_assign_base &operator=(const expected_delete_assign_base &) = + default; + expected_delete_assign_base &operator=( + expected_delete_assign_base &&) noexcept = delete; +}; + +template +struct expected_delete_assign_base { + expected_delete_assign_base() = default; + expected_delete_assign_base(const expected_delete_assign_base &) = default; + expected_delete_assign_base(expected_delete_assign_base &&) noexcept = + default; + expected_delete_assign_base &operator=(const expected_delete_assign_base &) = + delete; + expected_delete_assign_base &operator=( + expected_delete_assign_base &&) noexcept = default; +}; + +template +struct expected_delete_assign_base { + expected_delete_assign_base() = default; + expected_delete_assign_base(const expected_delete_assign_base &) = default; + expected_delete_assign_base(expected_delete_assign_base &&) noexcept = + default; + expected_delete_assign_base &operator=(const expected_delete_assign_base &) = + delete; + expected_delete_assign_base &operator=( + expected_delete_assign_base &&) noexcept = delete; +}; + +// This is needed to be able to construct the expected_default_ctor_base which +// follows, while still conditionally deleting the default constructor. +struct default_constructor_tag { + explicit constexpr default_constructor_tag() = default; +}; + +// expected_default_ctor_base will ensure that expected has a deleted default +// constructor if T is not default constructible. +// This specialization is for when T is default constructible +template ::value || std::is_void::value> +struct expected_default_ctor_base { + constexpr expected_default_ctor_base() noexcept = default; + constexpr expected_default_ctor_base( + expected_default_ctor_base const &) noexcept = default; + constexpr expected_default_ctor_base(expected_default_ctor_base &&) noexcept = + default; + expected_default_ctor_base &operator=( + expected_default_ctor_base const &) noexcept = default; + expected_default_ctor_base &operator=( + expected_default_ctor_base &&) noexcept = default; + + constexpr explicit expected_default_ctor_base(default_constructor_tag) {} +}; + +// This specialization is for when T is not default constructible +template +struct expected_default_ctor_base { + constexpr expected_default_ctor_base() noexcept = delete; + constexpr expected_default_ctor_base( + expected_default_ctor_base const &) noexcept = default; + constexpr expected_default_ctor_base(expected_default_ctor_base &&) noexcept = + default; + expected_default_ctor_base &operator=( + expected_default_ctor_base const &) noexcept = default; + expected_default_ctor_base &operator=( + expected_default_ctor_base &&) noexcept = default; + + constexpr explicit expected_default_ctor_base(default_constructor_tag) {} +}; +} // namespace detail + +template +class bad_expected_access : public std::exception { + public: + explicit bad_expected_access(E e) : m_val(std::move(e)) {} + + virtual const char *what() const noexcept override { + return "Bad expected access"; + } + + const E &error() const & { + return m_val; + } + E &error() & { + return m_val; + } + const E &&error() const && { + return std::move(m_val); + } + E &&error() && { + return std::move(m_val); + } + + private: + E m_val; +}; + +/// An `expected` object is an object that contains the storage for +/// another object and manages the lifetime of this contained object `T`. +/// Alternatively it could contain the storage for another unexpected object +/// `E`. The contained object may not be initialized after the expected object +/// has been initialized, and may not be destroyed before the expected object +/// has been destroyed. The initialization state of the contained object is +/// tracked by the expected object. +template +class TL_EXPECTED_NODISCARD expected + : private detail::expected_move_assign_base, + private detail::expected_delete_ctor_base, + private detail::expected_delete_assign_base, + private detail::expected_default_ctor_base { + static_assert(!std::is_reference::value, "T must not be a reference"); + static_assert(!std::is_same::type>::value, + "T must not be in_place_t"); + static_assert(!std::is_same::type>::value, + "T must not be unexpect_t"); + static_assert( + !std::is_same>::type>::value, + "T must not be unexpected"); + static_assert(!std::is_reference::value, "E must not be a reference"); + + T *valptr() { + return std::addressof(this->m_val); + } + const T *valptr() const { + return std::addressof(this->m_val); + } + unexpected *errptr() { + return std::addressof(this->m_unexpect); + } + const unexpected *errptr() const { + return std::addressof(this->m_unexpect); + } + + template ::value> * = nullptr> + TL_EXPECTED_11_CONSTEXPR U &val() { + return this->m_val; + } + TL_EXPECTED_11_CONSTEXPR unexpected &err() { + return this->m_unexpect; + } + + template ::value> * = nullptr> + constexpr const U &val() const { + return this->m_val; + } + constexpr const unexpected &err() const { + return this->m_unexpect; + } + + using impl_base = detail::expected_move_assign_base; + using ctor_base = detail::expected_default_ctor_base; + + public: + typedef T value_type; + typedef E error_type; + typedef unexpected unexpected_type; + +#if defined(TL_EXPECTED_CXX14) && !defined(TL_EXPECTED_GCC49) && \ + !defined(TL_EXPECTED_GCC54) && !defined(TL_EXPECTED_GCC55) + template + TL_EXPECTED_11_CONSTEXPR auto and_then(F &&f) & { + return and_then_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR auto and_then(F &&f) && { + return and_then_impl(std::move(*this), std::forward(f)); + } + template + constexpr auto and_then(F &&f) const & { + return and_then_impl(*this, std::forward(f)); + } + +#ifndef TL_EXPECTED_NO_CONSTRR + template + constexpr auto and_then(F &&f) const && { + return and_then_impl(std::move(*this), std::forward(f)); + } +#endif + +#else + template + TL_EXPECTED_11_CONSTEXPR auto and_then( + F &&f) & -> decltype(and_then_impl(std::declval(), + std::forward(f))) { + return and_then_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR auto and_then( + F &&f) && -> decltype(and_then_impl(std::declval(), + std::forward(f))) { + return and_then_impl(std::move(*this), std::forward(f)); + } + template + constexpr auto and_then( + F &&f) const & -> decltype(and_then_impl(std::declval(), + std::forward(f))) { + return and_then_impl(*this, std::forward(f)); + } + +#ifndef TL_EXPECTED_NO_CONSTRR + template + constexpr auto and_then(F &&f) + const && -> decltype(and_then_impl(std::declval(), + std::forward(f))) { + return and_then_impl(std::move(*this), std::forward(f)); + } +#endif +#endif + +#if defined(TL_EXPECTED_CXX14) && !defined(TL_EXPECTED_GCC49) && \ + !defined(TL_EXPECTED_GCC54) && !defined(TL_EXPECTED_GCC55) + template + TL_EXPECTED_11_CONSTEXPR auto map(F &&f) & { + return expected_map_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR auto map(F &&f) && { + return expected_map_impl(std::move(*this), std::forward(f)); + } + template + constexpr auto map(F &&f) const & { + return expected_map_impl(*this, std::forward(f)); + } + template + constexpr auto map(F &&f) const && { + return expected_map_impl(std::move(*this), std::forward(f)); + } +#else + template + TL_EXPECTED_11_CONSTEXPR decltype(expected_map_impl( + std::declval(), std::declval())) + map(F &&f) & { + return expected_map_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR decltype(expected_map_impl(std::declval(), + std::declval())) + map(F &&f) && { + return expected_map_impl(std::move(*this), std::forward(f)); + } + template + constexpr decltype(expected_map_impl(std::declval(), + std::declval())) + map(F &&f) const & { + return expected_map_impl(*this, std::forward(f)); + } + +#ifndef TL_EXPECTED_NO_CONSTRR + template + constexpr decltype(expected_map_impl(std::declval(), + std::declval())) + map(F &&f) const && { + return expected_map_impl(std::move(*this), std::forward(f)); + } +#endif +#endif + +#if defined(TL_EXPECTED_CXX14) && !defined(TL_EXPECTED_GCC49) && \ + !defined(TL_EXPECTED_GCC54) && !defined(TL_EXPECTED_GCC55) + template + TL_EXPECTED_11_CONSTEXPR auto transform(F &&f) & { + return expected_map_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR auto transform(F &&f) && { + return expected_map_impl(std::move(*this), std::forward(f)); + } + template + constexpr auto transform(F &&f) const & { + return expected_map_impl(*this, std::forward(f)); + } + template + constexpr auto transform(F &&f) const && { + return expected_map_impl(std::move(*this), std::forward(f)); + } +#else + template + TL_EXPECTED_11_CONSTEXPR decltype(expected_map_impl( + std::declval(), std::declval())) + transform(F &&f) & { + return expected_map_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR decltype(expected_map_impl(std::declval(), + std::declval())) + transform(F &&f) && { + return expected_map_impl(std::move(*this), std::forward(f)); + } + template + constexpr decltype(expected_map_impl(std::declval(), + std::declval())) + transform(F &&f) const & { + return expected_map_impl(*this, std::forward(f)); + } + +#ifndef TL_EXPECTED_NO_CONSTRR + template + constexpr decltype(expected_map_impl(std::declval(), + std::declval())) + transform(F &&f) const && { + return expected_map_impl(std::move(*this), std::forward(f)); + } +#endif +#endif + +#if defined(TL_EXPECTED_CXX14) && !defined(TL_EXPECTED_GCC49) && \ + !defined(TL_EXPECTED_GCC54) && !defined(TL_EXPECTED_GCC55) + template + TL_EXPECTED_11_CONSTEXPR auto map_error(F &&f) & { + return map_error_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR auto map_error(F &&f) && { + return map_error_impl(std::move(*this), std::forward(f)); + } + template + constexpr auto map_error(F &&f) const & { + return map_error_impl(*this, std::forward(f)); + } + template + constexpr auto map_error(F &&f) const && { + return map_error_impl(std::move(*this), std::forward(f)); + } +#else + template + TL_EXPECTED_11_CONSTEXPR decltype(map_error_impl(std::declval(), + std::declval())) + map_error(F &&f) & { + return map_error_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR decltype(map_error_impl(std::declval(), + std::declval())) + map_error(F &&f) && { + return map_error_impl(std::move(*this), std::forward(f)); + } + template + constexpr decltype(map_error_impl(std::declval(), + std::declval())) + map_error(F &&f) const & { + return map_error_impl(*this, std::forward(f)); + } + +#ifndef TL_EXPECTED_NO_CONSTRR + template + constexpr decltype(map_error_impl(std::declval(), + std::declval())) + map_error(F &&f) const && { + return map_error_impl(std::move(*this), std::forward(f)); + } +#endif +#endif +#if defined(TL_EXPECTED_CXX14) && !defined(TL_EXPECTED_GCC49) && \ + !defined(TL_EXPECTED_GCC54) && !defined(TL_EXPECTED_GCC55) + template + TL_EXPECTED_11_CONSTEXPR auto transform_error(F &&f) & { + return map_error_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR auto transform_error(F &&f) && { + return map_error_impl(std::move(*this), std::forward(f)); + } + template + constexpr auto transform_error(F &&f) const & { + return map_error_impl(*this, std::forward(f)); + } + template + constexpr auto transform_error(F &&f) const && { + return map_error_impl(std::move(*this), std::forward(f)); + } +#else + template + TL_EXPECTED_11_CONSTEXPR decltype(map_error_impl(std::declval(), + std::declval())) + transform_error(F &&f) & { + return map_error_impl(*this, std::forward(f)); + } + template + TL_EXPECTED_11_CONSTEXPR decltype(map_error_impl(std::declval(), + std::declval())) + transform_error(F &&f) && { + return map_error_impl(std::move(*this), std::forward(f)); + } + template + constexpr decltype(map_error_impl(std::declval(), + std::declval())) + transform_error(F &&f) const & { + return map_error_impl(*this, std::forward(f)); + } + +#ifndef TL_EXPECTED_NO_CONSTRR + template + constexpr decltype(map_error_impl(std::declval(), + std::declval())) + transform_error(F &&f) const && { + return map_error_impl(std::move(*this), std::forward(f)); + } +#endif +#endif + template + expected TL_EXPECTED_11_CONSTEXPR or_else(F &&f) & { + return or_else_impl(*this, std::forward(f)); + } + + template + expected TL_EXPECTED_11_CONSTEXPR or_else(F &&f) && { + return or_else_impl(std::move(*this), std::forward(f)); + } + + template + expected constexpr or_else(F &&f) const & { + return or_else_impl(*this, std::forward(f)); + } + +#ifndef TL_EXPECTED_NO_CONSTRR + template + expected constexpr or_else(F &&f) const && { + return or_else_impl(std::move(*this), std::forward(f)); + } +#endif + constexpr expected() = default; + constexpr expected(const expected &rhs) = default; + constexpr expected(expected &&rhs) = default; + expected &operator=(const expected &rhs) = default; + expected &operator=(expected &&rhs) = default; + + template ::value> * = + nullptr> + constexpr expected(in_place_t, Args &&...args) + : impl_base(in_place, std::forward(args)...), + ctor_base(detail::default_constructor_tag{}) {} + + template &, Args &&...>::value> * = nullptr> + constexpr expected(in_place_t, std::initializer_list il, Args &&...args) + : impl_base(in_place, il, std::forward(args)...), + ctor_base(detail::default_constructor_tag{}) {} + + template ::value> * = + nullptr, + detail::enable_if_t::value> * = + nullptr> + explicit constexpr expected(const unexpected &e) + : impl_base(unexpect, e.value()), + ctor_base(detail::default_constructor_tag{}) {} + + template < + class G = E, + detail::enable_if_t::value> * = + nullptr, + detail::enable_if_t::value> * = nullptr> + constexpr expected(unexpected const &e) + : impl_base(unexpect, e.value()), + ctor_base(detail::default_constructor_tag{}) {} + + template < + class G = E, + detail::enable_if_t::value> * = nullptr, + detail::enable_if_t::value> * = nullptr> + explicit constexpr expected(unexpected &&e) noexcept( + std::is_nothrow_constructible::value) + : impl_base(unexpect, std::move(e.value())), + ctor_base(detail::default_constructor_tag{}) {} + + template < + class G = E, + detail::enable_if_t::value> * = nullptr, + detail::enable_if_t::value> * = nullptr> + constexpr expected(unexpected &&e) noexcept( + std::is_nothrow_constructible::value) + : impl_base(unexpect, std::move(e.value())), + ctor_base(detail::default_constructor_tag{}) {} + + template ::value> * = + nullptr> + constexpr explicit expected(unexpect_t, Args &&...args) + : impl_base(unexpect, std::forward(args)...), + ctor_base(detail::default_constructor_tag{}) {} + + template &, Args &&...>::value> * = nullptr> + constexpr explicit expected(unexpect_t, std::initializer_list il, + Args &&...args) + : impl_base(unexpect, il, std::forward(args)...), + ctor_base(detail::default_constructor_tag{}) {} + + template ::value && + std::is_convertible::value)> * = + nullptr, + detail::expected_enable_from_other + * = nullptr> + explicit TL_EXPECTED_11_CONSTEXPR expected(const expected &rhs) + : ctor_base(detail::default_constructor_tag{}) { + if (rhs.has_value()) { + this->construct(*rhs); + } else { + this->construct_error(rhs.error()); + } + } + + template ::value && + std::is_convertible::value)> * = + nullptr, + detail::expected_enable_from_other + * = nullptr> + TL_EXPECTED_11_CONSTEXPR expected(const expected &rhs) + : ctor_base(detail::default_constructor_tag{}) { + if (rhs.has_value()) { + this->construct(*rhs); + } else { + this->construct_error(rhs.error()); + } + } + + template < + class U, class G, + detail::enable_if_t::value && + std::is_convertible::value)> * = nullptr, + detail::expected_enable_from_other * = nullptr> + explicit TL_EXPECTED_11_CONSTEXPR expected(expected &&rhs) + : ctor_base(detail::default_constructor_tag{}) { + if (rhs.has_value()) { + this->construct(std::move(*rhs)); + } else { + this->construct_error(std::move(rhs.error())); + } + } + + template < + class U, class G, + detail::enable_if_t<(std::is_convertible::value && + std::is_convertible::value)> * = nullptr, + detail::expected_enable_from_other * = nullptr> + TL_EXPECTED_11_CONSTEXPR expected(expected &&rhs) + : ctor_base(detail::default_constructor_tag{}) { + if (rhs.has_value()) { + this->construct(std::move(*rhs)); + } else { + this->construct_error(std::move(rhs.error())); + } + } + + template < + class U = T, + detail::enable_if_t::value> * = nullptr, + detail::expected_enable_forward_value * = nullptr> + explicit TL_EXPECTED_MSVC2015_CONSTEXPR expected(U &&v) + : expected(in_place, std::forward(v)) {} + + template < + class U = T, + detail::enable_if_t::value> * = nullptr, + detail::expected_enable_forward_value * = nullptr> + TL_EXPECTED_MSVC2015_CONSTEXPR expected(U &&v) + : expected(in_place, std::forward(v)) {} + + template < + class U = T, class G = T, + detail::enable_if_t::value> * = + nullptr, + detail::enable_if_t::value> * = nullptr, + detail::enable_if_t< + (!std::is_same, detail::decay_t>::value && + !detail::conjunction, + std::is_same>>::value && + std::is_constructible::value && + std::is_assignable::value && + std::is_nothrow_move_constructible::value)> * = nullptr> + expected &operator=(U &&v) { + if (has_value()) { + val() = std::forward(v); + } else { + err().~unexpected(); + ::new (valptr()) T(std::forward(v)); + this->m_has_val = true; + } + + return *this; + } + + template < + class U = T, class G = T, + detail::enable_if_t::value> * = + nullptr, + detail::enable_if_t::value> * = nullptr, + detail::enable_if_t< + (!std::is_same, detail::decay_t>::value && + !detail::conjunction, + std::is_same>>::value && + std::is_constructible::value && + std::is_assignable::value && + std::is_nothrow_move_constructible::value)> * = nullptr> + expected &operator=(U &&v) { + if (has_value()) { + val() = std::forward(v); + } else { + auto tmp = std::move(err()); + err().~unexpected(); + +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + try { + ::new (valptr()) T(std::forward(v)); + this->m_has_val = true; + } catch (...) { + err() = std::move(tmp); + throw; + } +#else + ::new (valptr()) T(std::forward(v)); + this->m_has_val = true; +#endif + } + + return *this; + } + + template ::value && + std::is_assignable::value> * = nullptr> + expected &operator=(const unexpected &rhs) { + if (!has_value()) { + err() = rhs; + } else { + this->destroy_val(); + ::new (errptr()) unexpected(rhs); + this->m_has_val = false; + } + + return *this; + } + + template ::value && + std::is_move_assignable::value> * = nullptr> + expected &operator=(unexpected &&rhs) noexcept { + if (!has_value()) { + err() = std::move(rhs); + } else { + this->destroy_val(); + ::new (errptr()) unexpected(std::move(rhs)); + this->m_has_val = false; + } + + return *this; + } + + template ::value> * = nullptr> + void emplace(Args &&...args) { + if (has_value()) { + val().~T(); + } else { + err().~unexpected(); + this->m_has_val = true; + } + ::new (valptr()) T(std::forward(args)...); + } + + template ::value> * = nullptr> + void emplace(Args &&...args) { + if (has_value()) { + val().~T(); + ::new (valptr()) T(std::forward(args)...); + } else { + auto tmp = std::move(err()); + err().~unexpected(); + +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + try { + ::new (valptr()) T(std::forward(args)...); + this->m_has_val = true; + } catch (...) { + err() = std::move(tmp); + throw; + } +#else + ::new (valptr()) T(std::forward(args)...); + this->m_has_val = true; +#endif + } + } + + template &, Args &&...>::value> * = nullptr> + void emplace(std::initializer_list il, Args &&...args) { + if (has_value()) { + T t(il, std::forward(args)...); + val() = std::move(t); + } else { + err().~unexpected(); + ::new (valptr()) T(il, std::forward(args)...); + this->m_has_val = true; + } + } + + template &, Args &&...>::value> * = nullptr> + void emplace(std::initializer_list il, Args &&...args) { + if (has_value()) { + T t(il, std::forward(args)...); + val() = std::move(t); + } else { + auto tmp = std::move(err()); + err().~unexpected(); + +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + try { + ::new (valptr()) T(il, std::forward(args)...); + this->m_has_val = true; + } catch (...) { + err() = std::move(tmp); + throw; + } +#else + ::new (valptr()) T(il, std::forward(args)...); + this->m_has_val = true; +#endif + } + } + + private: + using t_is_void = std::true_type; + using t_is_not_void = std::false_type; + using t_is_nothrow_move_constructible = std::true_type; + using move_constructing_t_can_throw = std::false_type; + using e_is_nothrow_move_constructible = std::true_type; + using move_constructing_e_can_throw = std::false_type; + + void swap_where_both_have_value(expected & /*rhs*/, t_is_void) noexcept { + // swapping void is a no-op + } + + void swap_where_both_have_value(expected &rhs, t_is_not_void) { + using std::swap; + swap(val(), rhs.val()); + } + + void swap_where_only_one_has_value(expected &rhs, t_is_void) noexcept( + std::is_nothrow_move_constructible::value) { + ::new (errptr()) unexpected_type(std::move(rhs.err())); + rhs.err().~unexpected_type(); + std::swap(this->m_has_val, rhs.m_has_val); + } + + void swap_where_only_one_has_value(expected &rhs, t_is_not_void) { + swap_where_only_one_has_value_and_t_is_not_void( + rhs, typename std::is_nothrow_move_constructible::type{}, + typename std::is_nothrow_move_constructible::type{}); + } + + void swap_where_only_one_has_value_and_t_is_not_void( + expected &rhs, t_is_nothrow_move_constructible, + e_is_nothrow_move_constructible) noexcept { + auto temp = std::move(val()); + val().~T(); + ::new (errptr()) unexpected_type(std::move(rhs.err())); + rhs.err().~unexpected_type(); + ::new (rhs.valptr()) T(std::move(temp)); + std::swap(this->m_has_val, rhs.m_has_val); + } + + void swap_where_only_one_has_value_and_t_is_not_void( + expected &rhs, t_is_nothrow_move_constructible, + move_constructing_e_can_throw) { + auto temp = std::move(val()); + val().~T(); +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + try { + ::new (errptr()) unexpected_type(std::move(rhs.err())); + rhs.err().~unexpected_type(); + ::new (rhs.valptr()) T(std::move(temp)); + std::swap(this->m_has_val, rhs.m_has_val); + } catch (...) { + val() = std::move(temp); + throw; + } +#else + ::new (errptr()) unexpected_type(std::move(rhs.err())); + rhs.err().~unexpected_type(); + ::new (rhs.valptr()) T(std::move(temp)); + std::swap(this->m_has_val, rhs.m_has_val); +#endif + } + + void swap_where_only_one_has_value_and_t_is_not_void( + expected &rhs, move_constructing_t_can_throw, + e_is_nothrow_move_constructible) { + auto temp = std::move(rhs.err()); + rhs.err().~unexpected_type(); +#ifdef TL_EXPECTED_EXCEPTIONS_ENABLED + try { + ::new (rhs.valptr()) T(std::move(val())); + val().~T(); + ::new (errptr()) unexpected_type(std::move(temp)); + std::swap(this->m_has_val, rhs.m_has_val); + } catch (...) { + rhs.err() = std::move(temp); + throw; + } +#else + ::new (rhs.valptr()) T(std::move(val())); + val().~T(); + ::new (errptr()) unexpected_type(std::move(temp)); + std::swap(this->m_has_val, rhs.m_has_val); +#endif + } + + public: + template + detail::enable_if_t::value && + detail::is_swappable::value && + (std::is_nothrow_move_constructible::value || + std::is_nothrow_move_constructible::value)> + swap(expected &rhs) noexcept(std::is_nothrow_move_constructible::value && + detail::is_nothrow_swappable::value && + std::is_nothrow_move_constructible::value && + detail::is_nothrow_swappable::value) { + if (has_value() && rhs.has_value()) { + swap_where_both_have_value(rhs, typename std::is_void::type{}); + } else if (!has_value() && rhs.has_value()) { + rhs.swap(*this); + } else if (has_value()) { + swap_where_only_one_has_value(rhs, typename std::is_void::type{}); + } else { + using std::swap; + swap(err(), rhs.err()); + } + } + + constexpr const T *operator->() const { + TL_ASSERT(has_value()); + return valptr(); + } + TL_EXPECTED_11_CONSTEXPR T *operator->() { + TL_ASSERT(has_value()); + return valptr(); + } + + template ::value> * = nullptr> + constexpr const U &operator*() const & { + TL_ASSERT(has_value()); + return val(); + } + template ::value> * = nullptr> + TL_EXPECTED_11_CONSTEXPR U &operator*() & { + TL_ASSERT(has_value()); + return val(); + } + template ::value> * = nullptr> + constexpr const U &&operator*() const && { + TL_ASSERT(has_value()); + return std::move(val()); + } + template ::value> * = nullptr> + TL_EXPECTED_11_CONSTEXPR U &&operator*() && { + TL_ASSERT(has_value()); + return std::move(val()); + } + + constexpr bool has_value() const noexcept { + return this->m_has_val; + } + constexpr explicit operator bool() const noexcept { + return this->m_has_val; + } + + template ::value> * = nullptr> + TL_EXPECTED_11_CONSTEXPR const U &value() const & { + if (!has_value()) + detail::throw_exception(bad_expected_access(err().value())); + return val(); + } + template ::value> * = nullptr> + TL_EXPECTED_11_CONSTEXPR U &value() & { + if (!has_value()) + detail::throw_exception(bad_expected_access(err().value())); + return val(); + } + template ::value> * = nullptr> + TL_EXPECTED_11_CONSTEXPR const U &&value() const && { + if (!has_value()) + detail::throw_exception(bad_expected_access(std::move(err()).value())); + return std::move(val()); + } + template ::value> * = nullptr> + TL_EXPECTED_11_CONSTEXPR U &&value() && { + if (!has_value()) + detail::throw_exception(bad_expected_access(std::move(err()).value())); + return std::move(val()); + } + + constexpr const E &error() const & { + TL_ASSERT(!has_value()); + return err().value(); + } + TL_EXPECTED_11_CONSTEXPR E &error() & { + TL_ASSERT(!has_value()); + return err().value(); + } + constexpr const E &&error() const && { + TL_ASSERT(!has_value()); + return std::move(err().value()); + } + TL_EXPECTED_11_CONSTEXPR E &&error() && { + TL_ASSERT(!has_value()); + return std::move(err().value()); + } + + template + constexpr T value_or(U &&v) const & { + static_assert(std::is_copy_constructible::value && + std::is_convertible::value, + "T must be copy-constructible and convertible to from U&&"); + return bool(*this) ? **this : static_cast(std::forward(v)); + } + template + TL_EXPECTED_11_CONSTEXPR T value_or(U &&v) && { + static_assert(std::is_move_constructible::value && + std::is_convertible::value, + "T must be move-constructible and convertible to from U&&"); + return bool(*this) ? std::move(**this) : static_cast(std::forward(v)); + } +}; + +namespace detail { +template +using exp_t = typename detail::decay_t::value_type; +template +using err_t = typename detail::decay_t::error_type; +template +using ret_t = expected>; + +#ifdef TL_EXPECTED_CXX14 +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + *std::declval()))> +constexpr auto and_then_impl(Exp &&exp, F &&f) { + static_assert(detail::is_expected::value, "F must return an expected"); + + return exp.has_value() + ? detail::invoke(std::forward(f), *std::forward(exp)) + : Ret(unexpect, std::forward(exp).error()); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval()))> +constexpr auto and_then_impl(Exp &&exp, F &&f) { + static_assert(detail::is_expected::value, "F must return an expected"); + + return exp.has_value() ? detail::invoke(std::forward(f)) + : Ret(unexpect, std::forward(exp).error()); +} +#else +template +struct TC; +template (), + *std::declval())), + detail::enable_if_t>::value> * = nullptr> +auto and_then_impl(Exp &&exp, F &&f) -> Ret { + static_assert(detail::is_expected::value, "F must return an expected"); + + return exp.has_value() + ? detail::invoke(std::forward(f), *std::forward(exp)) + : Ret(unexpect, std::forward(exp).error()); +} + +template ())), + detail::enable_if_t>::value> * = nullptr> +constexpr auto and_then_impl(Exp &&exp, F &&f) -> Ret { + static_assert(detail::is_expected::value, "F must return an expected"); + + return exp.has_value() ? detail::invoke(std::forward(f)) + : Ret(unexpect, std::forward(exp).error()); +} +#endif + +#ifdef TL_EXPECTED_CXX14 +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + *std::declval())), + detail::enable_if_t::value> * = nullptr> +constexpr auto expected_map_impl(Exp &&exp, F &&f) { + using result = ret_t>; + return exp.has_value() ? result(detail::invoke(std::forward(f), + *std::forward(exp))) + : result(unexpect, std::forward(exp).error()); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + *std::declval())), + detail::enable_if_t::value> * = nullptr> +auto expected_map_impl(Exp &&exp, F &&f) { + using result = expected>; + if (exp.has_value()) { + detail::invoke(std::forward(f), *std::forward(exp)); + return result(); + } + + return result(unexpect, std::forward(exp).error()); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval())), + detail::enable_if_t::value> * = nullptr> +constexpr auto expected_map_impl(Exp &&exp, F &&f) { + using result = ret_t>; + return exp.has_value() ? result(detail::invoke(std::forward(f))) + : result(unexpect, std::forward(exp).error()); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval())), + detail::enable_if_t::value> * = nullptr> +auto expected_map_impl(Exp &&exp, F &&f) { + using result = expected>; + if (exp.has_value()) { + detail::invoke(std::forward(f)); + return result(); + } + + return result(unexpect, std::forward(exp).error()); +} +#else +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + *std::declval())), + detail::enable_if_t::value> * = nullptr> + +constexpr auto expected_map_impl(Exp &&exp, + F &&f) -> ret_t> { + using result = ret_t>; + + return exp.has_value() ? result(detail::invoke(std::forward(f), + *std::forward(exp))) + : result(unexpect, std::forward(exp).error()); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + *std::declval())), + detail::enable_if_t::value> * = nullptr> + +auto expected_map_impl(Exp &&exp, F &&f) -> expected> { + if (exp.has_value()) { + detail::invoke(std::forward(f), *std::forward(exp)); + return {}; + } + + return unexpected>(std::forward(exp).error()); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval())), + detail::enable_if_t::value> * = nullptr> + +constexpr auto expected_map_impl(Exp &&exp, + F &&f) -> ret_t> { + using result = ret_t>; + + return exp.has_value() ? result(detail::invoke(std::forward(f))) + : result(unexpect, std::forward(exp).error()); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval())), + detail::enable_if_t::value> * = nullptr> + +auto expected_map_impl(Exp &&exp, F &&f) -> expected> { + if (exp.has_value()) { + detail::invoke(std::forward(f)); + return {}; + } + + return unexpected>(std::forward(exp).error()); +} +#endif + +#if defined(TL_EXPECTED_CXX14) && !defined(TL_EXPECTED_GCC49) && \ + !defined(TL_EXPECTED_GCC54) && !defined(TL_EXPECTED_GCC55) +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +constexpr auto map_error_impl(Exp &&exp, F &&f) { + using result = expected, detail::decay_t>; + return exp.has_value() + ? result(*std::forward(exp)) + : result(unexpect, detail::invoke(std::forward(f), + std::forward(exp).error())); +} +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +auto map_error_impl(Exp &&exp, F &&f) { + using result = expected, monostate>; + if (exp.has_value()) { + return result(*std::forward(exp)); + } + + detail::invoke(std::forward(f), std::forward(exp).error()); + return result(unexpect, monostate{}); +} +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +constexpr auto map_error_impl(Exp &&exp, F &&f) { + using result = expected, detail::decay_t>; + return exp.has_value() + ? result() + : result(unexpect, detail::invoke(std::forward(f), + std::forward(exp).error())); +} +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +auto map_error_impl(Exp &&exp, F &&f) { + using result = expected, monostate>; + if (exp.has_value()) { + return result(); + } + + detail::invoke(std::forward(f), std::forward(exp).error()); + return result(unexpect, monostate{}); +} +#else +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +constexpr auto map_error_impl(Exp &&exp, F &&f) + -> expected, detail::decay_t> { + using result = expected, detail::decay_t>; + + return exp.has_value() + ? result(*std::forward(exp)) + : result(unexpect, detail::invoke(std::forward(f), + std::forward(exp).error())); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +auto map_error_impl(Exp &&exp, F &&f) -> expected, monostate> { + using result = expected, monostate>; + if (exp.has_value()) { + return result(*std::forward(exp)); + } + + detail::invoke(std::forward(f), std::forward(exp).error()); + return result(unexpect, monostate{}); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +constexpr auto map_error_impl(Exp &&exp, F &&f) + -> expected, detail::decay_t> { + using result = expected, detail::decay_t>; + + return exp.has_value() + ? result() + : result(unexpect, detail::invoke(std::forward(f), + std::forward(exp).error())); +} + +template >::value> * = nullptr, + class Ret = decltype(detail::invoke(std::declval(), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +auto map_error_impl(Exp &&exp, F &&f) -> expected, monostate> { + using result = expected, monostate>; + if (exp.has_value()) { + return result(); + } + + detail::invoke(std::forward(f), std::forward(exp).error()); + return result(unexpect, monostate{}); +} +#endif + +#ifdef TL_EXPECTED_CXX14 +template (), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +constexpr auto or_else_impl(Exp &&exp, F &&f) { + static_assert(detail::is_expected::value, "F must return an expected"); + return exp.has_value() ? std::forward(exp) + : detail::invoke(std::forward(f), + std::forward(exp).error()); +} + +template (), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +detail::decay_t or_else_impl(Exp &&exp, F &&f) { + return exp.has_value() ? std::forward(exp) + : (detail::invoke(std::forward(f), + std::forward(exp).error()), + std::forward(exp)); +} +#else +template (), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +auto or_else_impl(Exp &&exp, F &&f) -> Ret { + static_assert(detail::is_expected::value, "F must return an expected"); + return exp.has_value() ? std::forward(exp) + : detail::invoke(std::forward(f), + std::forward(exp).error()); +} + +template (), + std::declval().error())), + detail::enable_if_t::value> * = nullptr> +detail::decay_t or_else_impl(Exp &&exp, F &&f) { + return exp.has_value() ? std::forward(exp) + : (detail::invoke(std::forward(f), + std::forward(exp).error()), + std::forward(exp)); +} +#endif +} // namespace detail + +template +constexpr bool operator==(const expected &lhs, + const expected &rhs) { + return (lhs.has_value() != rhs.has_value()) + ? false + : (!lhs.has_value() ? lhs.error() == rhs.error() : *lhs == *rhs); +} +template +constexpr bool operator!=(const expected &lhs, + const expected &rhs) { + return (lhs.has_value() != rhs.has_value()) + ? true + : (!lhs.has_value() ? lhs.error() != rhs.error() : *lhs != *rhs); +} +template +constexpr bool operator==(const expected &lhs, + const expected &rhs) { + return (lhs.has_value() != rhs.has_value()) + ? false + : (!lhs.has_value() ? lhs.error() == rhs.error() : true); +} +template +constexpr bool operator!=(const expected &lhs, + const expected &rhs) { + return (lhs.has_value() != rhs.has_value()) + ? true + : (!lhs.has_value() ? lhs.error() != rhs.error() : false); +} + +template +constexpr bool operator==(const expected &x, const U &v) { + return x.has_value() ? *x == v : false; +} +template +constexpr bool operator==(const U &v, const expected &x) { + return x.has_value() ? *x == v : false; +} +template +constexpr bool operator!=(const expected &x, const U &v) { + return x.has_value() ? *x != v : true; +} +template +constexpr bool operator!=(const U &v, const expected &x) { + return x.has_value() ? *x != v : true; +} + +template +constexpr bool operator==(const expected &x, const unexpected &e) { + return x.has_value() ? false : x.error() == e.value(); +} +template +constexpr bool operator==(const unexpected &e, const expected &x) { + return x.has_value() ? false : x.error() == e.value(); +} +template +constexpr bool operator!=(const expected &x, const unexpected &e) { + return x.has_value() ? true : x.error() != e.value(); +} +template +constexpr bool operator!=(const unexpected &e, const expected &x) { + return x.has_value() ? true : x.error() != e.value(); +} + +template ::value || + std::is_move_constructible::value) && + detail::is_swappable::value && + std::is_move_constructible::value && + detail::is_swappable::value> * = nullptr> +void swap(expected &lhs, + expected &rhs) noexcept(noexcept(lhs.swap(rhs))) { + lhs.swap(rhs); +} +} // namespace tl + +#endif diff --git a/src/ailego/pattern/factory.h b/src/ailego/pattern/factory.h new file mode 100644 index 0000000..b6a7d9e --- /dev/null +++ b/src/ailego/pattern/factory.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 +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +/*! Factory + */ +template +class Factory { + public: + /*! Factory Register + */ + template ::value>::type> + class Register { + public: + //! Constructor + Register(const char *key) { + Factory::Instance()->set(key, [] { return Register::Construct(); }); + } + + //! Constructor + template + Register(const char *key, TArgs &&...args) { + std::tuple tuple(std::forward(args)...); + + Factory::Instance()->set(key, [tuple] { + return Register::Construct( + tuple, typename TupleIndexMaker::Type()); + }); + } + + protected: + //! Tuple Index Maker + template + struct TupleIndexMaker : TupleIndexMaker {}; + + //! Tuple Index + template + struct TupleIndex {}; + + //! Tuple Index Maker (special) + template + struct TupleIndexMaker<0, I...> { + typedef TupleIndex Type; + }; + + //! Construct a register object + template + static TImpl *Construct(const std::tuple &tuple, + TupleIndex) { + return new (std::nothrow) TImpl(std::get(tuple)...); + } + + //! Construct a register object + static TImpl *Construct(void) { + return new (std::nothrow) TImpl(); + } + }; + + //! Produce an instance (c_ptr) + static TBase *Make(const char *key) { + return Factory::Instance()->produce(key); + } + + //! Produce an instance (shared_ptr) + static std::shared_ptr MakeShared(const char *key) { + return std::shared_ptr(Factory::Make(key)); + } + + //! Produce an instance (unique_ptr) + static std::unique_ptr MakeUnique(const char *key) { + return std::unique_ptr(Factory::Make(key)); + } + + //! Test if the class is exist + static bool Has(const char *key) { + return Factory::Instance()->has(key); + } + + //! Retrieve classes in factory + static std::vector Classes(void) { + return Factory::Instance()->classes(); + } + + protected: + //! Constructor + Factory(void) : map_() {} + + //! Retrieve the singleton factory + static Factory *Instance(void) { + static Factory factory; + return (&factory); + } + + //! Inserts a new class into map + template + void set(const char *key, TFunc &&func) { + map_[key] = std::forward(func); + } + + //! Produce an instance + TBase *produce(const char *key) { + auto iter = map_.find(key); + if (iter != map_.end()) { + return iter->second(); + } + return nullptr; + } + + //! Test if the class is exist + bool has(const char *key) { + return (map_.find(key) != map_.end()); + } + + //! Retrieve classes in factory + std::vector classes(void) const { + std::vector vec; + for (const auto &it : map_) { + vec.push_back(std::string(it.first)); + } + return vec; + } + + private: + //! Disable them + Factory(const Factory &); + Factory(Factory &&); + Factory &operator=(const Factory &); + + /*! Key Comparer + */ + struct KeyComparer { + bool operator()(const char *lhs, const char *rhs) const { + return (std::strcmp(lhs, rhs) < 0); + } + }; + + //! Don't use variable buffer as key store. + //! The key must be use a static buffer to store. + std::map, KeyComparer> map_; +}; + +//! Factory Register +#define AILEGO_FACTORY_REGISTER(__NAME__, __BASE__, __IMPL__, ...) \ + static ailego::Factory<__BASE__>::Register<__IMPL__> \ + __ailegoFactoryRegister_##__NAME__(#__NAME__, ##__VA_ARGS__) + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/pattern/scope_guard.h b/src/ailego/pattern/scope_guard.h new file mode 100644 index 0000000..f37feaa --- /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 "closure.h" + +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/pattern/singleton.h b/src/ailego/pattern/singleton.h new file mode 100644 index 0000000..f3880cd --- /dev/null +++ b/src/ailego/pattern/singleton.h @@ -0,0 +1,51 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Singleton (C++11) + */ +template +class Singleton { + public: + using ObjectType = typename std::remove_reference::type; + + //! Retrieve instance of object + static ObjectType &Instance(void) noexcept( + std::is_nothrow_constructible::value) { + // Since it's a static variable, if the class has already been created, + // it won't be created again. And it is thread-safe in C++11. + static ObjectType obj; + return obj; + } + + protected: + //! Constructor (Allow inheritance) + Singleton(void) {} + + private: + //! Disable them + Singleton(Singleton const &) = delete; + Singleton(Singleton &&) = delete; + Singleton &operator=(Singleton const &) = delete; + Singleton &operator=(Singleton &&) = delete; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/string/string_concat_helper.h b/src/ailego/string/string_concat_helper.h new file mode 100644 index 0000000..9348e0a --- /dev/null +++ b/src/ailego/string/string_concat_helper.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 "string_view.h" + +namespace zvec { +namespace ailego { +namespace internal { + +//! Helper class to convert integer and float types to string, facilitating +//! string concatenation memory allocation. +class Alphameric { + public: + //! Deals with int, int8_t, int16_t, int32_t, bool, short, signed char, non + //! class enum + Alphameric(int n) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%d", n)) {} + + //! Deals with unsigned int, uint8_t, uint16_t, uint32_t, unsigned short, + //! unsigned char + Alphameric(unsigned int n) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%u", n)) {} + + //! Deals with long, int32_t, int64_t + Alphameric(long n) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%ld", n)) {} + + //! Deals with unsigned long, uint32_t, uint64_t + Alphameric(unsigned long n) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%lu", n)) {} + + //! Deals with long long, int64_t + Alphameric(long long n) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%lld", n)) {} + + //! Deals with unsigned long long, uint64_t + Alphameric(unsigned long long n) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%llu", n)) {} + + //! Deals with float, with 6 precision digit the same as std::to_string + Alphameric(float f) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%g", f)) {} + + //! Deals with double, with 6 precision digit the same as std::to_string + Alphameric(double f) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%g", f)) {} + + //! Deals with long double, with 6 precision digit the same as std::to_string + Alphameric(long double f) + : view_(buffer_, snprintf(buffer_, sizeof(buffer_), "%Lg", f)) {} + + //! Deals with const char* + Alphameric(const char *s) : view_(s) {} + + //! Deals with std::string + Alphameric(const std::string &s) : view_(s) {} + + //! Deals with StringView + Alphameric(StringView s) : view_(s) {} + + // Use string literals ":" instead of character literals ':'. + Alphameric(char c) = delete; + Alphameric(const Alphameric &) = delete; + Alphameric &operator=(const Alphameric &) = delete; + + //! Deals with enum class with non int underlying type + template {} && !std::is_convertible{}>::type> + Alphameric(T e) + : Alphameric(static_cast::type>(e)) {} + + //! Deals with std::vector subscript reference + template ::value && + (std::is_same::reference>::value || + std::is_same::const_reference>::value)>:: + type * = nullptr> + Alphameric(T e) : Alphameric(static_cast(e)) {} + + //! string size + size_t size() const { + return view_.size(); + } + + //! string data + const char *data() const { + return view_.data(); + } + + //! string view + StringView view() const { + return view_; + } + + private: + static constexpr int kBufferSize = 32; + char buffer_[kBufferSize]; + StringView view_; +}; + +} // namespace internal +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/string/string_view.h b/src/ailego/string/string_view.h new file mode 100644 index 0000000..42d7457 --- /dev/null +++ b/src/ailego/string/string_view.h @@ -0,0 +1,64 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +//! StringView provides a lightweight view into the string data provided by +//! a `std::string`, double-quoted string literal, character array, or even +//! another `StringView`. +//! +//! A `StringView` does *not* own the string to which it +//! points, and that data cannot be modified through the view. +class StringView { + public: + //! Default constructor + StringView() = default; + + //! Construct from c-string + StringView(const char *str) + : data_(str), size_(str != nullptr ? strlen(str) : 0) {} + + //! Construct from [str, str+s) + StringView(const char *str, size_t len) : data_(str), size_(len) {} + + //! Construct from std::string + StringView(const std::string &str) : data_(str.data()), size_(str.size()) {} + + //! Retrieve data of string + const char *data() const { + return data_; + } + + //! Retrieve size of string + size_t size() const { + return size_; + } + + //! Retrieve non-zero if it is empty + bool empty() const { + return size_ == 0; + } + + private: + const char *data_{nullptr}; + size_t size_{0}; +}; + +} // 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..43fef95 --- /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..a4e61ff --- /dev/null +++ b/src/ailego/utility/bitset_helper.cc @@ -0,0 +1,2048 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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" + +#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..a759b68 --- /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..5e15241 --- /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 "file_helper.h" +#include "string_helper.h" + +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..4ae1c8e --- /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..095fb82 --- /dev/null +++ b/src/ailego/utility/file_helper.cc @@ -0,0 +1,386 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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_helper.h" + +#if defined(_WIN32) || defined(_WIN64) +#include +#else +#if defined(__APPLE__) || defined(__MACH__) +#include +#endif +#include +#include +#include +#include +#include +#include +#endif + +namespace zvec { +namespace ailego { + +bool FileHelper::GetSelfPath(std::string *path) { +#if defined(_WIN32) || defined(_WIN64) + char buf[MAX_PATH]; + DWORD len = GetModuleFileNameA(NULL, buf, MAX_PATH); +#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 (len <= 0) { + return false; + } + path->assign(buf, len); + return true; +} + +bool FileHelper::GetFilePath(NativeHandle handle, std::string *path) { +#if defined(_WIN32) || defined(_WIN64) + char buf[MAX_PATH]; + DWORD len = + GetFinalPathNameByHandleA(handle, buf, MAX_PATH, FILE_NAME_OPENED); +#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 (len <= 0) { + return false; + } + path->assign(buf, len); + return true; +} + +#if !defined(_WIN32) && !defined(_WIN64) + +static inline char *JoinFilePath(const char *prefix, const char *suffix) { + size_t prefix_len = strlen(prefix); + size_t suffix_len = strlen(suffix); + + char *path = (char *)malloc(prefix_len + suffix_len + 2); + if (path) { + memcpy(path, prefix, prefix_len); + memcpy(path + prefix_len + 1, suffix, suffix_len); + path[prefix_len] = '/'; + path[prefix_len + suffix_len + 1] = '\0'; + } + return path; +} + +bool FileHelper::GetWorkingDirectory(std::string *path) { + char buf[PATH_MAX]; + + if (!getcwd(buf, PATH_MAX)) { + return false; + } + path->assign(buf); + return !path->empty(); +} + +bool FileHelper::GetFileSize(const char *path, size_t *psz) { + struct stat buf; + if (stat(path, &buf) != 0) { + return false; + } + *psz = buf.st_size; + return true; +} + +bool FileHelper::DeleteFile(const char *path) { + // Delete a file by the path + return (unlink(path) == 0); +} + +bool FileHelper::RenameFile(const char *oldpath, const char *newpath) { + return (rename(oldpath, newpath) == 0); +} + +bool FileHelper::MakePath(const char *path) { + char pathbuf[PATH_MAX]; + char *sp, *pp; + + strncpy(pathbuf, path, sizeof(pathbuf) - 1); + pathbuf[PATH_MAX - 1] = '\0'; + + pp = pathbuf; + while ((sp = strchr(pp, '/')) != nullptr) { + // Neither root nor double slash in path + if (sp != pp) { + *sp = '\0'; + if (mkdir(pathbuf, 0755) == -1 && errno != EEXIST) { + return false; + } + *sp = '/'; + } + pp = sp + 1; + } + return !(*pp != '\0' && mkdir(pathbuf, 0755) == -1 && errno != EEXIST); +} + +bool FileHelper::RemoveDirectory(const char *path) { + DIR *dir = opendir(path); + if (!dir) { + return false; + } + + struct dirent *dent; + while ((dent = readdir(dir)) != nullptr) { + if (!strcmp(dent->d_name, ".") || !strcmp(dent->d_name, "..")) { + continue; + } + char *fullpath = JoinFilePath(path, dent->d_name); + if (!fullpath) { + continue; + } + + if (FileHelper::IsDirectory(fullpath)) { + FileHelper::RemoveDirectory(fullpath); + } else { + FileHelper::DeleteFile(fullpath); + } + free(fullpath); + } + closedir(dir); + return (rmdir(path) == 0); +} + +bool FileHelper::IsExist(const char *path) { + return (access(path, F_OK) == 0); +} + +bool FileHelper::IsRegular(const char *path) { + struct stat buf; + if (stat(path, &buf) != 0) { + return false; + } + return ((buf.st_mode & S_IFREG) != 0); +} + +bool FileHelper::IsDirectory(const char *path) { + struct stat buf; + if (stat(path, &buf) != 0) { + return false; + } + return ((buf.st_mode & S_IFDIR) != 0); +} + +bool FileHelper::IsSymbolicLink(const char *path) { + struct stat buf; + if (stat(path, &buf) != 0) { + return false; + } + return ((buf.st_mode & S_IFLNK) != 0); +} + +bool FileHelper::IsSame(const char *path1, const char *path2) { + char real_path1[PATH_MAX]; + char real_path2[PATH_MAX]; + if (!realpath(path1, real_path1)) { + return false; + } + if (!realpath(path2, real_path2)) { + return false; + } + return (!strcmp(real_path1, real_path2)); +} + +#else +#undef RemoveDirectory +#undef DeleteFile +#undef GetFileSize + +static inline char *JoinFilePath(const char *prefix, const char *suffix) { + size_t prefix_len = strlen(prefix); + size_t suffix_len = strlen(suffix); + + char *path = (char *)malloc(prefix_len + suffix_len + 2); + if (path) { + memcpy(path, prefix, prefix_len); + memcpy(path + prefix_len + 1, suffix, suffix_len); + path[prefix_len] = '\\'; + path[prefix_len + suffix_len + 1] = '\0'; + } + return path; +} + +bool FileHelper::GetWorkingDirectory(std::string *path) { + char buf[MAX_PATH]; + DWORD len = GetCurrentDirectoryA(MAX_PATH, buf); + + if (len <= 0) { + return false; + } + path->assign(buf, len); + return true; +} + +bool FileHelper::GetFileSize(const char *path, size_t *psz) { + HANDLE handle = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, nullptr, + OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + + LARGE_INTEGER file_size; + if (!GetFileSizeEx(handle, &file_size)) { + return false; + } + *psz = (size_t)file_size.QuadPart; + return true; +} + +bool FileHelper::DeleteFile(const char *path) { + // Delete a file by the path + return (DeleteFileA(path)); +} + +bool FileHelper::RenameFile(const char *oldpath, const char *newpath) { + return (MoveFileA(oldpath, newpath)); +} + +bool FileHelper::MakePath(const char *path) { + char pathbuf[MAX_PATH]; + char *sp, *pp; + + strncpy(pathbuf, path, sizeof(pathbuf) - 1); + pathbuf[MAX_PATH - 1] = '\0'; + + pp = pathbuf; + while ((sp = strpbrk(pp, "/\\")) != nullptr) { + // Neither root nor double slash in path + if (sp != pp) { + *sp = '\0'; + if (!CreateDirectoryA(pathbuf, nullptr) && + GetLastError() != ERROR_ALREADY_EXISTS) { + return false; + } + *sp = '\\'; + } + pp = sp + 1; + } + return !(*pp != '\0' && !CreateDirectoryA(pathbuf, nullptr) && + GetLastError() != ERROR_ALREADY_EXISTS); +} + +bool FileHelper::RemoveDirectory(const char *path) { + char *pathbuf = JoinFilePath(path, "*.*"); + ailego_false_if_false(pathbuf); + + WIN32_FIND_DATAA file_info; + HANDLE file = FindFirstFileA(pathbuf, &file_info); + + ailego_do_if_false(file != INVALID_HANDLE_VALUE) { + free(pathbuf); + FindClose(file); + return false; + } + + do { + if (!strcmp(file_info.cFileName, ".") || + !strcmp(file_info.cFileName, "..")) { + continue; + } + + char *fullpath = JoinFilePath(path, file_info.cFileName); + if (!fullpath) { + continue; + } + + if (file_info.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) { + FileHelper::RemoveDirectory(fullpath); + } else { + FileHelper::DeleteFile(fullpath); + } + free(fullpath); + } while (FindNextFileA(file, &file_info)); + + free(pathbuf); + FindClose(file); + return (!!RemoveDirectoryA(path)); +} + +bool FileHelper::IsExist(const char *path) { + DWORD attr = GetFileAttributesA(path); + return (attr != INVALID_FILE_ATTRIBUTES); +} + +bool FileHelper::IsRegular(const char *path) { + DWORD attr = GetFileAttributesA(path); + return (attr != INVALID_FILE_ATTRIBUTES && + !(attr & FILE_ATTRIBUTE_DIRECTORY)); +} + +bool FileHelper::IsDirectory(const char *path) { + DWORD attr = GetFileAttributesA(path); + return (attr != INVALID_FILE_ATTRIBUTES && (attr & FILE_ATTRIBUTE_DIRECTORY)); +} + +bool FileHelper::IsSymbolicLink(const char *path) { + DWORD attr = GetFileAttributesA(path); + return (attr != INVALID_FILE_ATTRIBUTES && + (attr & FILE_ATTRIBUTE_REPARSE_POINT)); +} + +bool FileHelper::IsSame(const char *path1, const char *path2) { + char real_path1[MAX_PATH]; + char real_path2[MAX_PATH]; + char **part_path1 = nullptr; + char **part_path2 = nullptr; + DWORD path1_size = + GetFullPathNameA(path1, sizeof(real_path1), real_path1, part_path1); + DWORD path2_size = + GetFullPathNameA(path2, sizeof(real_path2), real_path2, part_path2); + + if ((part_path1 && *part_path1 != 0) || (part_path2 && *part_path2 != 0) || + (path1_size != path2_size)) { + return false; + } + return (!strcmp(real_path1, real_path2)); +} + +#endif // !_WIN32 && !_WIN64 + +bool FileHelper::RemovePath(const char *path) { + if (FileHelper::IsDirectory(path)) { + return FileHelper::RemoveDirectory(path); + } + return FileHelper::DeleteFile(path); +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/utility/file_helper.h b/src/ailego/utility/file_helper.h new file mode 100644 index 0000000..8bce938 --- /dev/null +++ b/src/ailego/utility/file_helper.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 +#include +#include + +namespace zvec { +namespace ailego { + +/*! File Helper Module + */ +struct FileHelper { +#if defined(_WIN32) || defined(_WIN64) + //! Native Handle in Windows + typedef void *NativeHandle; +#else + //! Native Handle in POSIX + typedef int NativeHandle; +#endif + + //! Invalid Handle + static constexpr NativeHandle InvalidHandle = (NativeHandle)(-1); + + //! Retrieve the path of self process + static bool GetSelfPath(std::string *path); + + //! Retrieve the final path for the specified file + static bool GetFilePath(NativeHandle handle, std::string *path); + + //! Retrieve current working directory + static bool GetWorkingDirectory(std::string *path); + + //! Get the size of a file + static bool GetFileSize(const char *path, size_t *psz); + + //! Delete a name and possibly the file it refers to + static bool DeleteFile(const char *path); + + //! Change the name or location of a file + static bool RenameFile(const char *oldpath, const char *newpath); + + //! Make directories' path + static bool MakePath(const char *path); + + //! Remove a file or a directory (includes files & subdirectories) + static bool RemovePath(const char *path); + + //! Remove a directory (includes files & subdirectories) + static bool RemoveDirectory(const char *path); + + //! Retrieve non-zero if the path exists + static bool IsExist(const char *path); + + //! Retrieve non-zero if the path is a regular file + static bool IsRegular(const char *path); + + //! Retrieve non-zero if the path is a directory + static bool IsDirectory(const char *path); + + //! Retrieve non-zero if the path is a symbolic link + static bool IsSymbolicLink(const char *path); + + //! Retrieve non-zero if two paths are pointing to the same file + static bool IsSame(const char *path1, const char *path2); + + //! Retrieve the size of a file + static size_t FileSize(const char *path) { + size_t file_size = 0; + GetFileSize(path, &file_size); + return file_size; + } + + //! Retrieve the base name from a path + static const char *BaseName(const char *path) { + const char *output = std::strrchr(path, '/'); + if (!output) { + output = std::strrchr(path, '\\'); + } + return (output ? output + 1 : 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..6c00793 --- /dev/null +++ b/src/ailego/utility/float_helper.cc @@ -0,0 +1,1158 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "float_helper.h" + +#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(__F16C__) && defined(__AVX512F__) +static inline void convert_fp16_to_fp32(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] = _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(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] = _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(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] = _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(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] = _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); + } +} +#elif defined(__F16C__) && defined(__AVX__) +static inline void convert_fp16_to_fp32(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(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(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(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); + } +} +#elif defined(__aarch64__) +static inline float float32(uint16_t val) { + __fp16 *p = reinterpret_cast<__fp16 *>(&val); + return *p; +} + +static inline uint16_t float16(float val) { + __fp16 f = static_cast<__fp16>(val); + uint16_t *fp = reinterpret_cast(&f); + return *fp; +} + +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 *p = reinterpret_cast(&bits); + return (*p); +} + +// 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); +} + +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); + } +} +#endif // __F16C__ && __AVX512F__ + +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 \ No newline at end of file diff --git a/src/ailego/utility/float_helper.h b/src/ailego/utility/float_helper.h new file mode 100644 index 0000000..7a6d140 --- /dev/null +++ b/src/ailego/utility/float_helper.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 + +namespace zvec { +namespace ailego { + +/*! Float Helper + */ +struct FloatHelper { + //! Convert FP16 to FP32 + static float ToFP32(uint16_t val); + + //! Convert FP16 to FP32 (array) + static void ToFP32(const uint16_t *arr, size_t size, float *out); + + //! Convert FP16 to FP32 with normalization (array) + static void ToFP32(const uint16_t *arr, size_t size, float norm, float *out); + + //! Convert FP32 to FP16 + static uint16_t ToFP16(float val); + + //! Convert FP32 to FP16 (array) + static void ToFP16(const float *arr, size_t size, uint16_t *out); + + //! Convert FP32 to FP16 with normalization (array) + static void ToFP16(const float *arr, size_t size, float norm, uint16_t *out); + + //! Convert FP16 to FP32 with normalization + static inline float ToFP32(uint16_t val, float norm) { + return (FloatHelper::ToFP32(val) / norm); + } + + //! Convert FP32 to FP16 with normalization + static inline uint16_t ToFP16(float val, float norm) { + return FloatHelper::ToFP16(val / norm); + } +}; + +#if !defined(__aarch64__) +/*! Half-Precision Floating Point + */ +class Float16 { + public: + //! Constructor + Float16(void) : value_(0) {} + + //! Constructor + Float16(float val) : value_(FloatHelper::ToFP16(val)) {} + + //! Constructor + Float16(double val) : value_(FloatHelper::ToFP16(static_cast(val))) {} + + //! Assigment + Float16 &operator=(float val) { + this->value_ = FloatHelper::ToFP16(val); + return *this; + } + + //! Assigment + Float16 &operator+=(float val) { + this->value_ = FloatHelper::ToFP16(FloatHelper::ToFP32(this->value_) + val); + return *this; + } + + //! Assigment + Float16 &operator-=(float val) { + this->value_ = FloatHelper::ToFP16(FloatHelper::ToFP32(this->value_) - val); + return *this; + } + + //! Assigment + Float16 &operator*=(float val) { + this->value_ = FloatHelper::ToFP16(FloatHelper::ToFP32(this->value_) * val); + return *this; + } + + //! Assigment + Float16 &operator/=(float val) { + this->value_ = FloatHelper::ToFP16(FloatHelper::ToFP32(this->value_) / val); + return *this; + } + + //! Retrieve value in FP32 + operator float() const { + return FloatHelper::ToFP32(this->value_); + } + + //! Equal operator + bool operator==(const Float16 &rhs) const { + return this->value_ == rhs.value_; + } + + //! No equal operator + bool operator!=(const Float16 &rhs) const { + return this->value_ != rhs.value_; + } + + //! Less than operator + bool operator<(const Float16 &rhs) const { + return FloatHelper::ToFP32(this->value_) < FloatHelper::ToFP32(rhs.value_); + } + + //! Less than or equal operator + bool operator<=(const Float16 &rhs) const { + return FloatHelper::ToFP32(this->value_) <= FloatHelper::ToFP32(rhs.value_); + } + + //! Greater than operator + bool operator>(const Float16 &rhs) const { + return FloatHelper::ToFP32(this->value_) > FloatHelper::ToFP32(rhs.value_); + } + + //! Greater than or equal operator + bool operator>=(const Float16 &rhs) const { + return FloatHelper::ToFP32(this->value_) >= FloatHelper::ToFP32(rhs.value_); + } + + //! Calculate the absolute value + static inline Float16 Absolute(const Float16 &x) { + Float16 abs; + abs.value_ = static_cast(x.value_ & 0x7fff); + return abs; + } + + private: + uint16_t value_; +}; +#else +/*! Half-Precision Floating Point + */ +class Float16 { + public: + //! Constructor + Float16(void) : value_(0) {} + + //! Constructor + Float16(__fp16 val) : value_(val) {} + + //! Assigment + Float16 &operator=(__fp16 val) { + this->value_ = val; + return *this; + } + + //! Assigment + Float16 &operator+=(__fp16 val) { + this->value_ = this->value_ + val; + return *this; + } + + //! Assigment + Float16 &operator-=(__fp16 val) { + this->value_ = this->value_ - val; + return *this; + } + + //! Assigment + Float16 &operator*=(__fp16 val) { + this->value_ = this->value_ * val; + return *this; + } + + //! Assigment + Float16 &operator/=(__fp16 val) { + this->value_ = this->value_ / val; + return *this; + } + + //! Retrieve value in FP16 + operator __fp16() const { + return this->value_; + } + + //! Equal operator + bool operator==(const Float16 &rhs) const { + return this->value_ == rhs.value_; + } + + //! No equal operator + bool operator!=(const Float16 &rhs) const { + return this->value_ != rhs.value_; + } + + //! Less than operator + bool operator<(const Float16 &rhs) const { + return this->value_ < rhs.value_; + } + + //! Less than or equal operator + bool operator<=(const Float16 &rhs) const { + return this->value_ <= rhs.value_; + } + + //! Greater than operator + bool operator>(const Float16 &rhs) const { + return this->value_ > rhs.value_; + } + + //! Greater than or equal operator + bool operator>=(const Float16 &rhs) const { + return this->value_ >= rhs.value_; + } + + //! Calculate the absolute value + static inline Float16 Absolute(const Float16 &x) { + Float16 abs(x.value_); + uint16_t *p = reinterpret_cast(&abs.value_); + *p &= 0x7fff; + return abs; + } + + private: + __fp16 value_; +}; +#endif + +// Check size of Float16 +static_assert(sizeof(Float16) == 2, "Float16 must be aligned with 2 bytes"); + +} // 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..211346e --- /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 "float_helper.h" +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..2fea357 --- /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..b2f0094 --- /dev/null +++ b/src/ailego/utility/memory_helper.cc @@ -0,0 +1,397 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "file_helper.h" +#include "string_helper.h" + +#if defined(_WIN64) || defined(_WIN32) +#include +#include +#else +#if defined(__linux__) || defined(__linux) +#include +#elif defined(__APPLE__) && defined(__MACH__) +#include +#include +#endif +#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, NULL, 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, NULL, 10); + ++count; + } + break; + + case 'B': + if (strncmp(buf, "Buffers:", 8) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10); + ++count; + } + break; + + case 'C': + if (strncmp(buf, "Cached:", 7) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 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, NULL, 10); + ++count; + } else if (strncmp(buf, "MemFree:", 8) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10); + ++count; + } + break; + + case 'B': + if (strncmp(buf, "Buffers:", 8) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10); + ++count; + } + break; + + case 'C': + if (strncmp(buf, "Cached:", 7) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10); + ++count; + } + break; + + case 'S': + if (strncmp(buf, "Slab:", 5) == 0) { + avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 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, 0, 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; +} + +} // namespace ailego +} // namespace zvec \ No newline at end of file diff --git a/src/ailego/utility/memory_helper.h b/src/ailego/utility/memory_helper.h new file mode 100644 index 0000000..b145f03 --- /dev/null +++ b/src/ailego/utility/memory_helper.h @@ -0,0 +1,54 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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); + + //! 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/process_helper.cc b/src/ailego/utility/process_helper.cc new file mode 100644 index 0000000..7c3b218 --- /dev/null +++ b/src/ailego/utility/process_helper.cc @@ -0,0 +1,322 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "process_helper.h" + +#if defined(_WIN32) || defined(_WIN64) +#include +#include +#else +#include +#include +#include +#include +#include +#endif + +#include + +#if defined(__GNUC__) || defined(__GNUG__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-result" +#endif + +namespace zvec { +namespace ailego { + +#if defined(_WIN32) || defined(_WIN64) +uint32_t ProcessHelper::SelfPid(void) { + return GetCurrentProcessId(); +} + +uint32_t ProcessHelper::SelfTid(void) { + return GetCurrentThreadId(); +} + +uint32_t ProcessHelper::ParentPid(void) { + HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0); + if (snapshot == INVALID_HANDLE_VALUE) { + return (uint32_t)-1; + } + + DWORD pid = GetCurrentProcessId(); + PROCESSENTRY32 pe = {sizeof(pe)}; + for (BOOL ret = Process32First(snapshot, &pe); ret; + ret = Process32Next(snapshot, &pe)) { + if (pe.th32ProcessID == pid) { + CloseHandle(snapshot); + return pe.th32ParentProcessID; + } + } + CloseHandle(snapshot); + return (uint32_t)-1; +} + +size_t ProcessHelper::BackTrace(void **buf, size_t size) { + return RtlCaptureStackBackTrace(1, (DWORD)size, buf, nullptr); +} + +bool ProcessHelper::IsExist(uint32_t pid) { + HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0); + if (snapshot == INVALID_HANDLE_VALUE) { + return false; + } + + PROCESSENTRY32 pe = {sizeof(pe)}; + for (BOOL ret = Process32First(snapshot, &pe); ret; + ret = Process32Next(snapshot, &pe)) { + if (pe.th32ProcessID == pid) { + CloseHandle(snapshot); + return true; + } + } + CloseHandle(snapshot); + return false; +} + +void ProcessHelper::Daemon(const char *, const char *) { + // ::TODO:: +} + +void ProcessHelper::RegisterSignal(int sig, void f(int)) { + signal(sig, f); +} + +#else // !_WIN32 && !_WIN64 +uint32_t ProcessHelper::SelfPid(void) { + return getpid(); +} + +uint32_t ProcessHelper::SelfTid(void) { +#if defined(__linux) || defined(__linux__) + return syscall(SYS_gettid); +#else + uint64_t tid = (uint64_t)-1; + pthread_threadid_np(nullptr, &tid); + return (uint32_t)tid; +#endif +} + +uint32_t ProcessHelper::ParentPid(void) { + return getppid(); +} + +size_t ProcessHelper::BackTrace(void **buf, size_t size) { + size = backtrace(buf, size); + + if (size != 0) { + --size; + // Skip the current function + for (size_t i = 0; i < size; ++i) { + buf[i] = buf[i + 1]; + } + } + return size; +} + +bool ProcessHelper::IsExist(uint32_t pid) { + return (kill(pid, 0) == 0); +} + +void ProcessHelper::Daemon(const char *out, const char *err) { + // Fork off the parent process + pid_t pid = fork(); + if (pid < 0) { + exit(EXIT_FAILURE); + } + + if (pid > 0) { + // Let the parent terminate + exit(EXIT_SUCCESS); + } + + // The child process becomes session leader + if (setsid() < 0) { + exit(EXIT_FAILURE); + } + + // Ignore some signals + signal(SIGCHLD, SIG_IGN); + signal(SIGHUP, SIG_IGN); + + // Fork off for the second time + pid = fork(); + + if (pid < 0) { + exit(EXIT_FAILURE); + } + + if (pid > 0) { + // Let the parent terminate + exit(EXIT_SUCCESS); + } + + // Set new file permissions + umask(0); + + // Change the working directory to the root directory + // or another appropriated directory + chdir("/"); + + // Close all open file descriptors + for (long x = sysconf(_SC_OPEN_MAX); x >= 0; --x) { + close(x); + } + + stdin = nullptr; + stdout = nullptr; + stderr = nullptr; + + // Redirect standard output + if (out) { + stdout = fopen(out, "w+"); + } + + // Redirect standard error + if (err) { + stderr = fopen(err, "w+"); + } else if (out) { + stderr = fopen(out, "w+"); + } +} + +void ProcessHelper::RegisterSignal(int sig, void f(int)) { + struct sigaction sa; + + sigemptyset(&sa.sa_mask); + sa.sa_handler = f; + sa.sa_flags = SA_RESTART; + sigaction(sig, &sa, nullptr); +} +#endif // _WIN32 || _WIN64 + +#if defined(__linux) || defined(__linux__) +static const char *signal_names[32] = { + "NIL", // 00 NIL. + "SIGHUP", // 01 Hangup (POSIX). + "SIGINT", // 02 Interrupt (ANSI). + "SIGQUIT", // 03 Quit (POSIX). + "SIGILL", // 04 Illegal instruction (ANSI). + "SIGTRAP", // 05 Trace trap (POSIX). + "SIGABRT", // 06 Abort (ANSI). + "SIGBUS", // 07 BUS error (4.2 BSD). + "SIGFPE", // 08 Floating-point exception (ANSI). + "SIGKILL", // 09 Kill, unblockable (POSIX). + "SIGUSR1", // 10 User-defined signal 1 (POSIX). + "SIGSEGV", // 11 Segmentation violation (ANSI). + "SIGUSR2", // 12 User-defined signal 2 (POSIX). + "SIGPIPE", // 13 Broken pipe (POSIX). + "SIGALRM", // 14 Alarm clock (POSIX). + "SIGTERM", // 15 Termination (ANSI). + "SIGSTKFLT", // 16 Stack fault. + "SIGCHLD", // 17 Child status has changed (POSIX). + "SIGCONT", // 18 Continue (POSIX). + "SIGSTOP", // 19 Stop, unblockable (POSIX). + "SIGTSTP", // 20 Keyboard stop (POSIX). + "SIGTTIN", // 21 Background read from tty (POSIX). + "SIGTTOU", // 22 Background write to tty (POSIX). + "SIGURG", // 23 Urgent condition on socket (4.2 BSD). + "SIGXCPU", // 24 CPU limit exceeded (4.2 BSD). + "SIGXFSZ", // 25 File size limit exceeded (4.2 BSD). + "SIGVTALRM", // 26 Virtual alarm clock (4.2 BSD). + "SIGPROF", // 27 Profiling alarm clock (4.2 BSD). + "SIGWINCH", // 28 Window size change (4.3 BSD, Sun). + "SIGIO", // 29 I/O now possible (4.2 BSD). + "SIGPWR", // 30 Power failure restart (System V). + "SIGSYS" // 31 Bad system call. +}; + +#elif defined(__APPLE__) || defined(__MACH__) || defined(__FreeBSD__) || \ + defined(__NetBSD__) || defined(__OpenBSD__) +static const char *signal_names[33] = { + "NIL", // 00 NIL. + "SIGHUP", // 01 hangup. + "SIGINT", // 02 interrupt. + "SIGQUIT", // 03 quit. + "SIGILL", // 04 illegal instr. (not reset when caught). + "SIGTRAP", // 05 trace trap (not reset when caught). + "SIGABRT", // 06 abort(). + "SIGEMT", // 07 EMT instruction. + "SIGFPE", // 08 floating point exception. + "SIGKILL", // 09 kill (cannot be caught or ignored). + "SIGBUS", // 10 bus error. + "SIGSEGV", // 11 segmentation violation. + "SIGSYS", // 12 non-existent system call invoked. + "SIGPIPE", // 13 write on a pipe with no one to read it. + "SIGALRM", // 14 alarm clock. + "SIGTERM", // 15 software termination signal from kill. + "SIGURG", // 16 urgent condition on IO channel. + "SIGSTOP", // 17 sendable stop signal not from tty. + "SIGTSTP", // 18 stop signal from tty. + "SIGCONT", // 19 continue a stopped process. + "SIGCHLD", // 20 to parent on child stop or exit. + "SIGTTIN", // 21 to readers pgrp upon background tty read. + "SIGTTOU", // 22 like TTIN if (tp->t_local<OSTOP). + "SIGIO", // 23 input/output possible signal. + "SIGXCPU", // 24 exceeded CPU time limit. + "SIGXFSZ", // 25 exceeded file size limit. + "SIGVTALRM", // 26 virtual time alarm. + "SIGPROF", // 27 profiling time alarm. + "SIGWINCH", // 28 window size changes. + "SIGINFO", // 29 information request. + "SIGUSR1", // 30 user defined signal 1. + "SIGUSR2", // 31 user defined signal 2. + "SIGTHR" // 32 reserved by thread library. +}; + +#elif defined(_WIN32) || defined(_WIN64) +static const char *signal_names[] = { + "NIL", // 00 + "NIL", // 01 + "SIGINT", // 02 interrupt + "NIL", // 03 + "SIGILLL", // 04 illegal instruction - invalid function image + "NIL", // 05 + "NIL", // 06 + "NIL", // 07 + "SIGFPE", // 08 floating point exception + "NIL", // 09 + "NIL", // 10 + "SIGSEGV", // 11 segment violation + "NIL", // 12 + "NIL", // 13 + "NIL", // 14 + "SIGTERM", // 15 Software termination signal from kill + "NIL", // 16 + "NIL", // 17 + "NIL", // 18 + "NIL", // 19 + "NIL", // 20 + "SIGBREAK", // 21 Ctrl-Break sequence + "SIGABRT", // 22 abnormal termination triggered by abort call +}; +#endif + +void ProcessHelper::IgnoreSignal(int sig) { + signal(sig, SIG_IGN); +} + +const char *ProcessHelper::SignalName(int sig) { + if (sig >= 0 && sig < (int)(sizeof(signal_names) / sizeof(signal_names[0]))) { + return signal_names[sig]; + } + return signal_names[0]; +} + +} // namespace ailego +} // namespace zvec + +#if defined(__GNUC__) || defined(__GNUG__) +#pragma GCC diagnostic pop +#endif \ No newline at end of file diff --git a/src/ailego/utility/process_helper.h b/src/ailego/utility/process_helper.h new file mode 100644 index 0000000..f28594b --- /dev/null +++ b/src/ailego/utility/process_helper.h @@ -0,0 +1,54 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Process Helper + */ +struct ProcessHelper { + //! Retrieve current process id + static uint32_t SelfPid(void); + + //! Retrieve current thread id + static uint32_t SelfTid(void); + + //! Retrieve current parent process id + static uint32_t ParentPid(void); + + //! Retrieve a backtrace for calling program + static size_t BackTrace(void **buf, size_t size); + + //! Retrieve non-zero if the process id exists + static bool IsExist(uint32_t pid); + + //! Run the process in the background + static void Daemon(const char *out, const char *err); + + //! Ignore signal + static void IgnoreSignal(int sig); + + //! Register signal + static void RegisterSignal(int sig, void f(int)); + + //! Retrieve the name of a signal id + static const char *SignalName(int sig); +}; + +} // 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..cdd81db --- /dev/null +++ b/src/ailego/utility/string_helper.cc @@ -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. + +#include "string_helper.h" +#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; +} + +#if defined(_MSC_VER) +#define strncasecmp _strnicmp +#endif + +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/string_helper.h b/src/ailego/utility/string_helper.h new file mode 100644 index 0000000..1f54055 --- /dev/null +++ b/src/ailego/utility/string_helper.h @@ -0,0 +1,322 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "string_helper_impl.h" + +namespace zvec { +namespace ailego { + +/*! String Helper + */ +struct StringHelper { + //! Return true if the `ref` starts with the given prefix + static bool StartsWith(const std::string &ref, const std::string &prefix); + + //! Return true if the `ref` ends with the given suffix + static bool EndsWith(const std::string &ref, const std::string &suffix); + + //! Split a string into a vector of T + //! NOTE: delim better NOT contain valid symbol for T, + //! i.e. digits + - for integers, + //! digits + - E e . for floating numbers + //! otherwise there will be performance overhead. + template + static void Split(const std::string &str, char delim, std::vector *out) { + return details::SplitImpl(str, delim, out); + } + template + static void Split(const std::string &str, const char *delim, + std::vector *out) { + return details::SplitImpl(str, delim, out); + } + template + static void Split(const std::string &str, const std::string &delim, + std::vector *out) { + return details::SplitImpl(str, delim, out); + } + + template + static void Split(const std::string &str, char delim, std::vector *out, + bool skip_empty) { + return details::SplitImpl(str, delim, out, skip_empty); + } + template + static void Split(const std::string &str, const char *delim, + std::vector *out, bool skip_empty) { + return details::SplitImpl(str, delim, out, skip_empty); + } + template + static void Split(const std::string &str, const std::string &delim, + std::vector *out, bool skip_empty) { + return details::SplitImpl(str, delim, out, + skip_empty); + } + + // Trim from start (in place) + static void LeftTrim(std::string &str); + + // Trim from end (in place) + static void RightTrim(std::string &str); + + // Trim from both ends (in place) + static void Trim(std::string &str); + + // Trim from start (copying) + static std::string CopyLeftTrim(std::string str); + + // Trim from end (copying) + static std::string CopyRightTrim(std::string str); + + // Trim from both ends (copying) + static std::string CopyTrim(std::string str); + + //! Compare ignore case + static bool CompareIgnoreCase(const std::string &a, const std::string &b); + + //! Convert string to floating-point number (double) + static bool ToDouble(const std::string &str, double *val) { + char *endptr = nullptr; + *val = std::strtod(str.c_str(), &endptr); + return (endptr && *endptr == '\0'); + } + + //! Convert string to floating-point number (float) + static bool ToFloat(const std::string &str, float *val) { + char *endptr = nullptr; + *val = std::strtof(str.c_str(), &endptr); + return (endptr && *endptr == '\0'); + } + + //! Convert string to integer number (int8_t) + static bool ToInt8(const std::string &str, int8_t *val) { + char *endptr = nullptr; + *val = static_cast(std::strtol(str.c_str(), &endptr, 0)); + return (endptr && *endptr == '\0'); + } + + //! Convert string to integer number (int16_t) + static bool ToInt16(const std::string &str, int16_t *val) { + char *endptr = nullptr; + *val = static_cast(std::strtol(str.c_str(), &endptr, 0)); + return (endptr && *endptr == '\0'); + } + + //! Convert string to integer number (int32_t) + static bool ToInt32(const std::string &str, int32_t *val) { + char *endptr = nullptr; + *val = static_cast(std::strtol(str.c_str(), &endptr, 0)); + return (endptr && *endptr == '\0'); + } + + //! Convert string to integer number (int64_t) + static bool ToInt64(const std::string &str, int64_t *val) { + char *endptr = nullptr; + *val = static_cast(std::strtoll(str.c_str(), &endptr, 0)); + return (endptr && *endptr == '\0'); + } + + //! Convert string to unsigned integer number (uint8_t) + static bool ToUint8(const std::string &str, uint8_t *val) { + char *endptr = nullptr; + *val = static_cast(std::strtoul(str.c_str(), &endptr, 0)); + return (endptr && *endptr == '\0'); + } + + //! Convert string to unsigned integer number (uint16_t) + static bool ToUint16(const std::string &str, uint16_t *val) { + char *endptr = nullptr; + *val = static_cast(std::strtoul(str.c_str(), &endptr, 0)); + return (endptr && *endptr == '\0'); + } + + //! Convert string to unsigned integer number (uint32_t) + static bool ToUint32(const std::string &str, uint32_t *val) { + char *endptr = nullptr; + *val = static_cast(std::strtoul(str.c_str(), &endptr, 0)); + return (endptr && *endptr == '\0'); + } + + //! Convert string to unsigned integer number (uint64_t) + static bool ToUint64(const std::string &str, uint64_t *val) { + char *endptr = nullptr; + *val = static_cast(std::strtoull(str.c_str(), &endptr, 0)); + return (endptr && *endptr == '\0'); + } + + //! Convert floating-point number string (double) + static std::string ToString(double val) { + return std::to_string(val); + } + + //! Convert floating-point number string (float) + static std::string ToString(float val) { + return std::to_string(val); + } + + //! Convert integer number to string (int8_t) + static std::string ToString(int8_t val) { + return std::to_string(val); + } + + //! Convert integer number to string (int16_t) + static std::string ToString(int16_t val) { + return std::to_string(val); + } + + //! Convert integer number to string (int32_t) + static std::string ToString(int32_t val) { + return std::to_string(val); + } + + //! Convert integer number to string (int64_t) + static std::string ToString(int64_t val) { + return std::to_string(val); + } + + //! Convert unsigned integer number to string (uint8_t) + static std::string ToString(uint8_t val) { + return std::to_string(val); + } + + //! Convert unsigned integer number to string (uint16_t) + static std::string ToString(uint16_t val) { + return std::to_string(val); + } + + //! Convert unsigned integer number to string (uint32_t) + static std::string ToString(uint32_t val) { + return std::to_string(val); + } + + //! Convert unsigned integer number to string (uint64_t) + static std::string ToString(uint64_t val) { + return std::to_string(val); + } + + //! Concatenation of arbitrary number of std::string, c-string, integers, + //! floating point numbers with one memory allocation. + //! E.g. auto s = Concat("foo", 123, std::string("bar"), 3.14159); + //! + //! Do not do the following, use Append instead + //! str = Concat(str, ...); + //! str.append(Concat(str, ...)); + //! str += Concat(str, ...); + //! + //! NOTE: char literal(e.g. ':') is not allowed, + //! use string literal(e.g. ":") instead. + static std::string Concat() { + return {}; + } + static std::string Concat(const internal::Alphameric &a); + static std::string Concat(const internal::Alphameric &a, + const internal::Alphameric &b); + static std::string Concat(const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c); + static std::string Concat(const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c, + const internal::Alphameric &d); + // Support 5 or more arguments + template + static std::string Concat(const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c, + const internal::Alphameric &d, + const internal::Alphameric &e, const T &...args) { + std::string result; + Append(&result, a, b, c, d, e, args...); + return result; + } + + //! Append arbitrary number of std::string, c-string, integers, + //! floating point numbers to existing string with one memory allocation. + //! E.g. Append(&str, "foo", 123, std::string("bar"), 3.14159); + //! + //! WARNING: Append requires that none of the arguments be a reference to + //! destination str. + //! + //! Do not do the following + //! std::string s = "foo"; + //! Append(&s, s); + //! + //! NOTE: char literal(e.g. ':') is not allowed, + //! use string literal(e.g. ":") instead. + static void Append(std::string *) {} + static void Append(std::string *str, const internal::Alphameric &a); + static void Append(std::string *str, const internal::Alphameric &a, + const internal::Alphameric &b); + static void Append(std::string *str, const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c); + static void Append(std::string *str, const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c, + const internal::Alphameric &d); + // Support 5 or more arguments + template + static void Append(std::string *str, const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c, + const internal::Alphameric &d, + const internal::Alphameric &e, const T &...args) { + AppendViews(str, + {a.view(), b.view(), c.view(), d.view(), e.view(), + static_cast(args).view()...}); + } + + //! Append list of StringView to str. + static void AppendViews(std::string *str, + std::initializer_list views); +}; + +inline std::string StringHelper::Concat(const internal::Alphameric &a) { + std::string result; + Append(&result, a); + return result; +} + +inline std::string StringHelper::Concat(const internal::Alphameric &a, + const internal::Alphameric &b) { + std::string result; + Append(&result, a, b); + return result; +} + +inline std::string StringHelper::Concat(const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c) { + std::string result; + Append(&result, a, b, c); + return result; +} + +inline std::string StringHelper::Concat(const internal::Alphameric &a, + const internal::Alphameric &b, + const internal::Alphameric &c, + const internal::Alphameric &d) { + std::string result; + Append(&result, a, b, c, d); + return result; +} + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/string_helper_impl.h b/src/ailego/utility/string_helper_impl.h new file mode 100644 index 0000000..9817214 --- /dev/null +++ b/src/ailego/utility/string_helper_impl.h @@ -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. + +#pragma once + +#include +#include +#include + +namespace zvec { +namespace ailego { +namespace details { + +//! Convert string to integers or floating point numbers +template +static T CStringToType(const char *begin, char **endptr) { + static_assert( + std::is_same::value || std::is_same::value || + std::is_same::value || std::is_same::value || + std::is_same::value || + std::is_same::value || + std::is_same::value || std::is_same::value || + std::is_same::value || std::is_same::value, + "type not supported"); + if (std::is_same::value || std::is_same::value || + std::is_same::value) { + return static_cast(strtol(begin, endptr, 0)); + } else if (std::is_same::value) { + return static_cast(strtoll(begin, endptr, 0)); + } else if (std::is_same::value || + std::is_same::value || + std::is_same::value) { + return static_cast(strtoul(begin, endptr, 0)); + } else if (std::is_same::value) { + return static_cast(strtoull(begin, endptr, 0)); + } else if (std::is_same::value) { + return static_cast(strtof(begin, endptr)); + } else { + return static_cast(strtod(begin, endptr)); + } +} + +//! Convert [begin, end) to T +//! If [end, ) contains valid T symbol, extra overhead will be incurred by +//! constructing std::string +template +struct StringToType { + T operator()(const char *begin, const char *end) { + char *eptr = nullptr; + auto v = CStringToType(begin, &eptr); + if (eptr > end) { + // NOTE: [begin, end) is not 0 terminated + // If delimiter contains valid T symbol, eptr might point to location + // after end. + // We create string here, which is guaranteed to be 0 terminated. + std::string s{begin, end}; + return CStringToType(s.c_str(), &eptr); + } + return v; + } +}; + +//! Specialization for std::string +template <> +struct StringToType { + std::string operator()(const char *begin, const char *end) { + return {begin, end}; + } +}; + +//! Return delimiter length. +template +struct DelimiterLen { + size_t operator()(T delimiter); +}; + +//! Return delimiter length for char. +template <> +struct DelimiterLen { + size_t operator()(char) { + return 1; + } +}; + +//! Return delimiter length for const char*. +template <> +struct DelimiterLen { + size_t operator()(const char *delimiter) { + return delimiter == nullptr ? 0 : std::strlen(delimiter); + } +}; + +//! Return delimiter length for std::string. +template <> +struct DelimiterLen { + size_t operator()(const std::string &delimiter) { + return delimiter.size(); + } +}; + +//! Split implementation. +template ::value || + std::is_same::value || + std::is_same::value, + D>::type> +static void SplitImpl(const std::string &str, D delim, std::vector *out) { + StringToType func; + out->clear(); + + auto s = str.data(); + size_t delimiter_len = DelimiterLen()(delim); + if (delimiter_len != 0) { + size_t a = 0, b = str.find(delim); + while (b != std::string::npos) { + out->push_back(func(s + a, s + b)); + a = b + delimiter_len; + b = str.find(delim, a); + } + out->push_back(func(s + a, s + str.length())); + } else { + out->push_back(func(s + 0, s + str.length())); + } +} + +//! Split implementation. +template ::value || + std::is_same::value || + std::is_same::value, + D>::type> +static void SplitImpl(const std::string &str, D delim, std::vector *out, + bool skip_empty) { + StringToType func; + out->clear(); + + auto s = str.data(); + size_t delimiter_len = DelimiterLen()(delim); + if (delimiter_len != 0) { + size_t a = 0, b = str.find(delim); + while (b != std::string::npos) { + if (!skip_empty || b - a > 0) { + out->push_back(func(s + a, s + b)); + } + + a = b + delimiter_len; + b = str.find(delim, a); + } + if (!skip_empty || str.length() - a > 0) { + out->push_back(func(s + a, s + str.length())); + } + } else { + if (!skip_empty || str.length() > 0) { + out->push_back(func(s + 0, s + str.length())); + } + } +} +} // namespace details +} // 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..36b63c1 --- /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 "time_helper.h" + +#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(0); + 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(0); + 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/utility/time_helper.h b/src/ailego/utility/time_helper.h new file mode 100644 index 0000000..b253538 --- /dev/null +++ b/src/ailego/utility/time_helper.h @@ -0,0 +1,206 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Monotime + */ +struct Monotime { + //! Retrieve monotonic time in nanoseconds + static uint64_t NanoSeconds(void); + + //! Retrieve monotonic time in microseconds + static uint64_t MicroSeconds(void); + + //! Retrieve monotonic time in milliseconds + static uint64_t MilliSeconds(void); + + //! Retrieve monotonic time in seconds + static uint64_t Seconds(void); +}; + +/*! Realtime + */ +struct Realtime { + //! Retrieve system time in nanoseconds + static uint64_t NanoSeconds(void); + + //! Retrieve system time in microseconds + static uint64_t MicroSeconds(void); + + //! Retrieve system time in milliseconds + static uint64_t MilliSeconds(void); + + //! Retrieve system time in seconds + static uint64_t Seconds(void); + + //! Retrieve a timestamp as a specific local time format + static size_t Localtime(uint64_t stamp, const char *format, char *buf, + size_t len); + + //! Retrieve a timestamp as a specific GMT time format + static size_t Gmtime(uint64_t stamp, const char *format, char *buf, + size_t len); + + //! Retrieve local time in string + static size_t Localtime(const char *format, char *buf, size_t len); + + //! Retrieve GMT time in string + static size_t Gmtime(const char *format, char *buf, size_t len); + + //! Retrieve local time in string + static size_t Localtime(char *buf, size_t len) { + return Localtime("%Y-%m-%d %H:%M:%S", buf, len); + } + + //! Retrieve GMT time in string + static size_t Gmtime(char *buf, size_t len) { + return Gmtime("%Y-%m-%d %H:%M:%S", buf, len); + } + + //! Retrieve local time in string + static std::string Localtime(void) { + char str[32]; + Localtime(str, sizeof(str)); + return std::string(str); + } + + //! Retrieve GMT time in string + static std::string Gmtime(void) { + char str[32]; + Gmtime(str, sizeof(str)); + return std::string(str); + } + + //! Retrieve a timestamp as a specific local time format + static size_t Localtime(uint64_t stamp, char *buf, size_t len) { + return Localtime(stamp, "%Y-%m-%d %H:%M:%S", buf, len); + } + + //! Retrieve a timestamp as a specific GMT time format + static size_t Gmtime(uint64_t stamp, char *buf, size_t len) { + return Gmtime(stamp, "%Y-%m-%d %H:%M:%S", buf, len); + } + + //! Retrieve a timestamp as a specific local time format + static std::string Localtime(uint64_t stamp) { + char str[32]; + Localtime(stamp, str, sizeof(str)); + return std::string(str); + } + + //! Retrieve a timestamp as a specific GMT time format + static std::string Gmtime(uint64_t stamp) { + char str[32]; + Gmtime(stamp, str, sizeof(str)); + return std::string(str); + } +}; + +/*! Thread-specific CPU time + */ +struct CPUtime { + //! Retrieve CPU time in nanoseconds + static uint64_t NanoSeconds(void); + + //! Retrieve CPU time in microseconds + static uint64_t MicroSeconds(void); + + //! Retrieve CPU time in milliseconds + static uint64_t MilliSeconds(void); + + //! Retrieve CPU time in seconds + static uint64_t Seconds(void); +}; + +/*! Elapsed Time + */ +class ElapsedTime { + public: + //! Constructor + ElapsedTime(void) : stamp_(Monotime::NanoSeconds()) {} + + //! Retrieve the elapsed time in nanoseconds + uint64_t nano_seconds(void) const { + return (Monotime::NanoSeconds() - stamp_); + } + + //! Retrieve the elapsed time in milliseconds + uint64_t micro_seconds(void) const { + return (this->nano_seconds() / 1000u); + } + + //! Retrieve the elapsed time in milliseconds + uint64_t milli_seconds(void) const { + return (this->nano_seconds() / 1000000u); + } + + //! Retrieve the elapsed time in seconds + uint64_t seconds(void) const { + return (this->nano_seconds() / 1000000000u); + } + + //! Update time stamp + void reset(void) { + stamp_ = Monotime::NanoSeconds(); + } + + private: + uint64_t stamp_; +}; + +/*! Elapsed CPU Time + */ +class ElapsedCPUTime { + public: + //! Constructor + ElapsedCPUTime(void) : stamp_(CPUtime::NanoSeconds()) {} + + //! Retrieve the elapsed time in nanoseconds + uint64_t nano_seconds(void) const { + return (CPUtime::NanoSeconds() - stamp_); + } + + //! Retrieve the elapsed time in milliseconds + uint64_t micro_seconds(void) const { + return (this->nano_seconds() / 1000u); + } + + //! Retrieve the elapsed time in milliseconds + uint64_t milli_seconds(void) const { + return (this->nano_seconds() / 1000000u); + } + + //! Retrieve the elapsed time in seconds + uint64_t seconds(void) const { + return (this->nano_seconds() / 1000000000u); + } + + //! Update time stamp + void reset(void) { + stamp_ = CPUtime::NanoSeconds(); + } + + private: + uint64_t stamp_; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/utility/type_helper.h b/src/ailego/utility/type_helper.h new file mode 100644 index 0000000..576a9e3 --- /dev/null +++ b/src/ailego/utility/type_helper.h @@ -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. + +#pragma once + +#include +#include +#include "float_helper.h" + +namespace zvec { +namespace ailego { + +//! Determines if a type is an arithmetic type (includes Float16) +template +struct IsArithmetic + : std::integral_constant::value || + std::is_same::value> {}; + +//! Determines if a type is a signed arithmetic type (includes Float16) +template +struct IsSignedArithmetic + : std::integral_constant::value || + std::is_same::value> {}; + +//! Determines if a type is a unsigned arithmetic type (includes Float16) +template +struct IsUnsignedArithmetic + : std::integral_constant::value> {}; + +//! Determines if a type is a floating-point type (includes Float16) +template +struct IsFloatingPoint + : std::integral_constant::value || + std::is_same::value> {}; + +#if __GNUC__ >= 5 || defined(_MSC_VER) || defined(__clang__) +template +using IsTriviallyCopyable = std::is_trivially_copyable; +#else +template +using IsTriviallyCopyable = std::has_trivial_copy_constructor; +#endif + +#if __cplusplus >= 201703L // C++17 + +//! Determines if a type can be invoked with the specified argument types +template +using IsInvocable = std::is_invocable; + +//! Determines if a type can be invoked with the specified argument types +template +using IsInvocableWithResult = std::is_invocable_r; + +#else +//! Determines if a type can be invoked with the specified argument types +template +struct IsInvocable + : std::is_constructible, + std::reference_wrapper< + typename std::remove_reference::type> > { +}; + +//! Determines if a type can be invoked with the specified argument types +template +struct IsInvocableWithResult + : std::is_constructible, + std::reference_wrapper< + typename std::remove_reference::type> > { +}; +#endif + +//! Fixed underlying_type used with conditional +template ::value> +struct UnderlyingType { + typedef typename std::remove_cv::type type; +}; + +//! Fixed underlying_type used with conditional +template +struct UnderlyingType { + typedef typename std::underlying_type::type type; +}; + +#if __cplusplus >= 201703L // C++17 + +//! Variadic logical AND metafunction +template +using Conjunction = std::conjunction; + +//! Variadic logical OR metafunction +template +using Disjunction = std::disjunction; + +#else +//! Variadic logical AND metafunction +template +struct Conjunction : std::true_type {}; + +//! Variadic logical AND metafunction +template +struct Conjunction + : std::conditional, + std::false_type>::type {}; + +//! Variadic logical OR metafunction +template +struct Disjunction : std::false_type {}; + +//! Variadic logical OR metafunction +template +struct Disjunction + : std::conditional >::type {}; +#endif + +} // namespace ailego +} // namespace zvec 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..e46c049 --- /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 "internal/platform.h" + +#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..7dab04a --- /dev/null +++ b/src/binding/CMakeLists.txt @@ -0,0 +1,8 @@ +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(python) \ No newline at end of file diff --git a/src/binding/python/CMakeLists.txt b/src/binding/python/CMakeLists.txt new file mode 100644 index 0000000..07074c7 --- /dev/null +++ b/src/binding/python/CMakeLists.txt @@ -0,0 +1,59 @@ +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/param/python_param.cc + model/schema/python_schema.cc + model/common/python_config.cc + typing/python_type.cc +) + +pybind11_add_module(_zvec ${SRC_LISTS}) + +if (CMAKE_SYSTEM_NAME STREQUAL "Linux") + target_link_libraries(_zvec PRIVATE + -Wl,--whole-archive + $ + $ + $ + $ + $ + $ + $ + $ + $ + $ + -Wl,--no-whole-archive + zvec + ) + target_link_options(_zvec PRIVATE + "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/exports.map" + ) +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,$ + zvec + ) + target_link_libraries(_zvec PRIVATE + -Wl,-exported_symbols_list,${CMAKE_CURRENT_SOURCE_DIR}/exports.mac + ) +endif () + +target_include_directories(_zvec PRIVATE ${PYBIND11_INCLUDE_DIR} ${CMAKE_SOURCE_DIR}/src ${CMAKE_SOURCE_DIR}/src/binding/python/include) diff --git a/src/binding/python/binding.cc b/src/binding/python/binding.cc new file mode 100644 index 0000000..ed8d691 --- /dev/null +++ b/src/binding/python/binding.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 "python_collection.h" +#include "python_config.h" +#include "python_doc.h" +#include "python_param.h" +#include "python_schema.h" +#include "python_type.h" + +namespace zvec { +PYBIND11_MODULE(_zvec, m) { + m.doc() = "Zvec core module"; + + ZVecPyTyping::Initialize(m); + ZVecPyParams::Initialize(m); + ZVecPySchemas::Initialize(m); + ZVecPyConfig::Initialize(m); + ZVecPyDoc::Initialize(m); + ZVecPyCollection::Initialize(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..66d92e0 --- /dev/null +++ b/src/binding/python/exports.map @@ -0,0 +1,6 @@ +{ + global: + PyInit_*; # export only python related functions + 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..4a365bc --- /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 "db/collection.h" + +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..65f5986 --- /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 "db/common/config.h" + +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..9998db0 --- /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 "db/index/common/doc.h" + +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..ac5ab73 --- /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 "db/index/common/options.h" +#include "db/index/common/type.h" + +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_schema.h b/src/binding/python/include/python_schema.h new file mode 100644 index 0000000..eba8031 --- /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 "db/index/common/type.h" + +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..221af5f --- /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 "db/common/status.h" +#include "db/index/common/type.h" + +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..bbcbb5b --- /dev/null +++ b/src/binding/python/model/common/python_config.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 "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::DEBUG; + if (iequals(s, "info")) return GlobalConfig::LogLevel::INFO; + if (iequals(s, "warn") || iequals(s, "warning")) + return GlobalConfig::LogLevel::WARN; + if (iequals(s, "error")) return GlobalConfig::LogLevel::ERROR; + if (iequals(s, "fatal")) return GlobalConfig::LogLevel::FATAL; + 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); + } + + // initialize (contains validate) + Status status = GlobalConfig::Instance().Initialize(data); + if (!status.ok()) { + throw std::runtime_error("Initialization failed: " + status.message()); + } + return py::none(); + }); +} + + +} // 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..0a528d1 --- /dev/null +++ b/src/binding/python/model/param/python_param.cc @@ -0,0 +1,1222 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "core/interface/constants.h" +#include "db/index/common/index_params.h" +#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"; + 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"; + 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::string serialize_vector(const T *data, size_t n) { + std::string buf; + buf.resize(n * sizeof(T)); + std::memcpy(buf.data(), data, n * sizeof(T)); + return buf; +} + +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 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. +)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( + "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()); + 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"); + // 基类，不能直接实例化，用于子类 + 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 100. + 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 + ... ) + >>> print(params) + {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8'} +)pbdoc"); + hnsw_params + .def(py::init(), + 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) + .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( + "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()); + 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()) + "}"; + }) + .def(py::pickle( + [](const HnswIndexParams &self) { + return py::make_tuple(self.metric_type(), self.m(), + self.ef_construction(), self.quantize_type()); + }, + [](py::tuple t) { + if (t.size() != 4) + throw std::runtime_error("Invalid state for HnswIndexParams"); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast()); + })); + + // 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(), + py::arg("metric_type") = MetricType::IP, + py::arg("quantize_type") = QuantizeType::UNDEFINED, + 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). +)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()); + 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()) + "}"; + }) + .def(py::pickle( + [](const FlatIndexParams &self) { + return py::make_tuple(self.metric_type(), self.quantize_type()); + }, + [](py::tuple t) { + if (t.size() != 2) + throw std::runtime_error("Invalid state for FlatIndexParams"); + return std::make_shared(t[0].cast(), + t[1].cast()); + })); + + // 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. + If set to 0, the system will auto-select a reasonable value based on data size. + Default is 0 (auto). + 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(), + py::arg("metric_type") = MetricType::IP, py::arg("n_list") = 0, + py::arg("n_iters") = 10, py::arg("use_soar") = false, + py::arg("quantize_type") = QuantizeType::UNDEFINED, + 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). Set to 0 for auto. + Defaults to 0. + 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. +)pbdoc") + .def_property_readonly("n_list", &IVFIndexParams::n_list, + "int: Number of inverted lists (0 = auto).") + .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()); + 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()) + "}"; + }) + .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()); + }, + [](py::tuple t) { + if (t.size() != 5) + throw std::runtime_error("Invalid state for IVFIndexParams"); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast()); + })); +} + +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(), + 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 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. +)pbdoc") + .def_property_readonly( + "ef", [](const HnswQueryParams &self) -> int { return self.ef(); }, + "int: Size of the dynamic candidate list during HNSW search.") + .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()) + "}"; + }) + .def(py::pickle( + [](const HnswQueryParams &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 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()); + 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; + })); +} + +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) { + py::class_(m, "_VectorQuery") + .def(py::init<>()) + // properties + .def_readwrite("topk", &VectorQuery::topk_) + .def_readwrite("field_name", &VectorQuery::field_name_) + .def_readwrite("filter", &VectorQuery::filter_) + .def_readwrite("include_vector", &VectorQuery::include_vector_) + .def_readwrite("query_params", &VectorQuery::query_params_) + .def_readwrite("output_fields", &VectorQuery::output_fields_) + // vector + .def("set_vector", + [](VectorQuery &self, const FieldSchema &field_schema, + const py::object &obj) { + const DataType data_type = field_schema.data_type(); + + // dense vector + 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(); + switch (data_type) { + case DataType::VECTOR_FP32: { + self.query_vector_ = serialize_vector( + static_cast(buf.ptr), buf.size); + return; + } + case DataType::VECTOR_FP64: { + self.query_vector_ = serialize_vector( + static_cast(buf.ptr), buf.size); + return; + } + case DataType::VECTOR_INT8: { + self.query_vector_ = serialize_vector( + static_cast(buf.ptr), buf.size); + return; + } + case DataType::VECTOR_FP16: { + self.query_vector_ = serialize_vector( + static_cast(buf.ptr), buf.size); + return; + } + default: + throw py::type_error( + "Unsupported dense vector type for ndarray input: " + + std::to_string(static_cast(data_type))); + } + } + // 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.query_sparse_indices_ = std::move(indices); + self.query_sparse_values_ = 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.query_sparse_indices_ = std::move(indices); + self.query_sparse_values_ = 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()); + }) + .def( + "get_vector", + [](const VectorQuery &self, + const FieldSchema &field_schema) -> py::object { + DataType data_type = field_schema.data_type(); + if (FieldSchema::is_dense_vector_field(data_type)) { + if (self.query_vector_.empty()) { + throw std::runtime_error("No dense vector has been set"); + } + + size_t byte_size = self.query_vector_.size(); + const void *data = self.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 (self.query_sparse_indices_.empty()) { + return py::dict(); + } + + // Deserialize indices: stored as uint32_t[] + size_t indices_byte_size = self.query_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( + self.query_sparse_indices_.data()); + + // Deserialize values + switch (data_type) { + case DataType::SPARSE_VECTOR_FP32: { + if (self.query_sparse_values_.size() != n * sizeof(float)) { + throw std::runtime_error( + "Sparse FP32 values buffer size mismatch"); + } + const float *values = reinterpret_cast( + self.query_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 (self.query_sparse_values_.size() != + n * sizeof(uint16_t)) { + throw std::runtime_error( + "Sparse FP16 values buffer size mismatch"); + } + const uint16_t *raw_bits = reinterpret_cast( + self.query_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 VectorQuery &self) { + return py::make_tuple( + self.topk_, self.field_name_, self.query_vector_, + self.query_sparse_indices_, self.query_sparse_values_, + self.filter_, self.include_vector_, self.output_fields_, + self.query_params_ ? py::cast(self.query_params_) : py::none()); + }, + [](py::tuple t) { + if (t.size() != 9) + throw std::runtime_error("Invalid pickle data for VectorQuery"); + + VectorQuery obj{}; + obj.topk_ = t[0].cast(); + obj.field_name_ = t[1].cast(); + obj.query_vector_ = t[2].cast(); + obj.query_sparse_indices_ = t[3].cast(); + obj.query_sparse_values_ = t[4].cast(); + obj.filter_ = t[5].cast(); + obj.include_vector_ = t[6].cast(); + obj.output_fields_ = t[7].cast>(); + + if (!t[8].is_none()) { + obj.query_params_ = t[8].cast(); + } + return obj; + })); +} +} // namespace zvec \ No newline at end of file diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc new file mode 100644 index 0000000..211c671 --- /dev/null +++ b/src/binding/python/model/python_collection.cc @@ -0,0 +1,218 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "db/collection.h" + +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) { + auto result = + Collection::CreateAndOpen(path, schema, options); + return unwrap_expected(result); + }) + .def_static("Open", [](const std::string &path, + const CollectionOptions &options) { + auto 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) { + const auto status = self.Destroy(); + throw_if_error(status); + }) + .def("Flush", [](Collection &self) { + auto 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) { + const auto status = + self.CreateIndex(column_name, index_options, options); + throw_if_error(status); + }) + .def("DropIndex", + [](Collection &self, const std::string &column_name) { + const auto status = self.DropIndex(column_name); + throw_if_error(status); + }) + .def("Optimize", [](Collection &self, const OptimizeOptions &options) { + const auto status = self.Optimize(options); + throw_if_error(status); + }); + + // binding column ddl methods + col.def("AddColumn", + [](Collection &self, const std::string &column_name, + const FieldSchema::Ptr &column_schema, + const std::string &expression, const AddColumnOptions &options) { + const auto status = + self.AddColumn(column_name, column_schema, expression, options); + throw_if_error(status); + }) + .def("DropColumn", + [](Collection &self, std::string &column_name) { + auto 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) { + const auto 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) { + const auto result = self.Insert(docs); + return unwrap_expected(result); + }) + .def("Update", + [](Collection &self, std::vector &docs) { + const auto result = self.Update(docs); + return unwrap_expected(result); + }) + .def("Upsert", + [](Collection &self, std::vector &docs) { + const auto result = self.Upsert(docs); + return unwrap_expected(result); + }) + .def("Delete", + [](Collection &self, const std::vector &pks) { + const auto result = self.Delete(pks); + return unwrap_expected(result); + }) + .def("DeleteByFilter", [](Collection &self, const std::string &filter) { + const auto status = self.DeleteByFilter(filter); + throw_if_error(status); + }); +} + +void ZVecPyCollection::bind_dql_methods( + py::class_ &col) { + col.def("Query", + [](const Collection &self, const VectorQuery &query) { + const auto result = self.Query(query); + // return DocPtrList + return unwrap_expected(result); + }) + .def("GroupByQuery", + [](const Collection &self, const GroupByVectorQuery &query) { + const auto result = self.GroupByQuery(query); + // return GroupResults + return unwrap_expected(result); + }) + .def("Fetch", + [](const Collection &self, const std::vector &pks) { + const auto result = self.Fetch(pks); + // return DocPtrMap + return unwrap_expected(result); + }); +} + +} // 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..f54d310 --- /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 std::move(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 std::move(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 std::move(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 std::move(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/schema/python_schema.cc b/src/binding/python/model/schema/python_schema.cc new file mode 100644 index 0000000..c833a77 --- /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 "db/index/common/schema.h" +#include "db/index/common/stats.h" + +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..ee057cf --- /dev/null +++ b/src/binding/python/typing/python_type.cc @@ -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. + +#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("IVF", IndexType::IVF) + .value("FLAT", IndexType::FLAT) + .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); +} + +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..c001139 --- /dev/null +++ b/src/core/CMakeLists.txt @@ -0,0 +1,10 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +cc_directory(framework) +cc_directory(algorithm) +cc_directory(metric) +cc_directory(quantizer) +cc_directory(utility) +cc_directory(interface) +cc_directory(mixed_reducer) \ No newline at end of file diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt new file mode 100644 index 0000000..648dbef --- /dev/null +++ b/src/core/algorithm/CMakeLists.txt @@ -0,0 +1,9 @@ +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) \ No newline at end of file diff --git a/src/core/algorithm/cluster/CMakeLists.txt b/src/core/algorithm/cluster/CMakeLists.txt new file mode 100644 index 0000000..b0ccc79 --- /dev/null +++ b/src/core/algorithm/cluster/CMakeLists.txt @@ -0,0 +1,10 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +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 + 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..06cfdad --- /dev/null +++ b/src/core/algorithm/cluster/cluster_params.h @@ -0,0 +1,207 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 = + "proxima.general.cluster.count"; +static const std::string GENERAL_THREAD_COUNT = + "proxima.general.cluster.thread_count"; + +//! Optimize K-means +static const std::string OPTKMEANS_CLUSTER_COUNT = + "proxima.optkmeans.cluster.count"; +static const std::string OPTKMEANS_CLUSTER_MAX_ITERATIONS = + "proxima.optkmeans.cluster.max_iterations"; +static const std::string OPTKMEANS_CLUSTER_EPSILON = + "proxima.optkmeans.cluster.epsilon"; +static const std::string OPTKMEANS_CLUSTER_SHARD_FACTOR = + "proxima.optkmeans.cluster.shard_factor"; +static const std::string OPTKMEANS_CLUSTER_PURGE_EMPTY = + "proxima.optkmeans.cluster.purge_empty"; +static const std::string OPTKMEANS_CLUSTER_MARKOV_CHAIN_LENGTH = + "proxima.optkmeans.cluster.markov_chain_length"; +static const std::string OPTKMEANS_CLUSTER_ASSUMPTION_FREE = + "proxima.optkmeans.cluster.assumption_free"; + +//! K-means +static const std::string KMEANS_CLUSTER_COUNT = "proxima.kmeans.cluster.count"; +static const std::string KMEANS_CLUSTER_SHARD_FACTOR = + "proxima.kmeans.cluster.shard_factor"; +static const std::string KMEANS_CLUSTER_EPSILON = + "proxima.kmeans.cluster.epsilon"; +static const std::string KMEANS_CLUSTER_MAX_ITERATIONS = + "proxima.kmeans.cluster.max_iterations"; +static const std::string KMEANS_CLUSTER_PURGE_EMPTY = + "proxima.kmeans.cluster.purge_empty"; +static const std::string KMEANS_CLUSTER_BATCH = "proxima.kmeans.cluster.batch"; +static const std::string KMEANS_CLUSTER_SEEKER_CLASS = + "proxima.kmeans.cluster.seeker_class"; +static const std::string KMEANS_CLUSTER_SEEKER_PARAMS = + "proxima.kmeans.cluster.seeker_params"; + +//! Mini Batch K-means +static const std::string MINIBATCHKMEANS_CLUSTER_COUNT = + "proxima.minibatchkmeans.cluster.count"; +static const std::string MINIBATCHKMEANS_CLUSTER_SHARD_FACTOR = + "proxima.minibatchkmeans.cluster.shard_factor"; +static const std::string MINIBATCHKMEANS_CLUSTER_EPSILON = + "proxima.minibatchkmeans.cluster.epsilon"; +static const std::string MINIBATCHKMEANS_CLUSTER_MAX_ITERATIONS = + "proxima.minibatchkmeans.cluster.max_iterations"; +static const std::string MINIBATCHKMEANS_CLUSTER_PURGE_EMPTY = + "proxima.minibatchkmeans.cluster.purge_empty"; +static const std::string MINIBATCHKMEANS_CLUSTER_TRY_COUNT = + "proxima.minibatchkmeans.cluster.try_count"; +static const std::string MINIBATCHKMEANS_CLUSTER_BATCH_COUNT = + "proxima.minibatchkmeans.cluster.batch_count"; +static const std::string MINIBATCHKMEANS_CLUSTER_SEEKER_CLASS = + "proxima.minibatchkmeans.cluster.seeker_class"; +static const std::string MINIBATCHKMEANS_CLUSTER_SEEKER_PARAMS = + "proxima.minibatchkmeans.cluster.seeker_params"; + +//! K-means++ +static const std::string KMEANSPP_CLUSTER_COUNT = + "proxima.kmeanspp.cluster.count"; +static const std::string KMEANSPP_CLUSTER_SHARD_FACTOR = + "proxima.kmeanspp.cluster.shard_factor"; +static const std::string KMEANSPP_CLUSTER_CLASS = + "proxima.kmeanspp.cluster.class"; +static const std::string KMEANSPP_CLUSTER_PARAMS = + "proxima.kmeanspp.cluster.params"; + +//! K-MC2 +static const std::string KMC2_CLUSTER_COUNT = "proxima.kmc2.cluster.count"; +static const std::string KMC2_CLUSTER_SHARD_FACTOR = + "proxima.kmc2.cluster.shard_factor"; +static const std::string KMC2_CLUSTER_MARKOV_CHAIN_LENGTH = + "proxima.kmc2.cluster.markov_chain_length"; +static const std::string KMC2_CLUSTER_ASSUMPTION_FREE = + "proxima.kmc2.cluster.assumption_free"; +static const std::string KMC2_CLUSTER_CLASS = "proxima.kmc2.cluster.class"; +static const std::string KMC2_CLUSTER_PARAMS = "proxima.kmc2.cluster.params"; + +//! Bisecting K-means +static const std::string BIKMEANS_CLUSTER_COUNT = + "proxima.bikmeans.cluster.count"; +static const std::string BIKMEANS_CLUSTER_INIT_COUNT = + "proxima.bikmeans.cluster.init_count"; +static const std::string BIKMEANS_CLUSTER_PURGE_EMPTY = + "proxima.bikmeans.cluster.purge_empty"; +static const std::string BIKMEANS_CLUSTER_FIRST_CLASS = + "proxima.bikmeans.cluster.first_class"; +static const std::string BIKMEANS_CLUSTER_SECOND_CLASS = + "proxima.bikmeans.cluster.second_class"; +static const std::string BIKMEANS_CLUSTER_FIRST_PARAMS = + "proxima.bikmeans.cluster.first_params"; +static const std::string BIKMEANS_CLUSTER_SECOND_PARAMS = + "proxima.bikmeans.cluster.second_params"; + +//! K-medoids +static const std::string KMEDOIDS_CLUSTER_COUNT = + "proxima.kmedoids.cluster.count"; +static const std::string KMEDOIDS_CLUSTER_SHARD_FACTOR = + "proxima.kmedoids.cluster.shard_factor"; +static const std::string KMEDOIDS_CLUSTER_EPSILON = + "proxima.kmedoids.cluster.epsilon"; +static const std::string KMEDOIDS_CLUSTER_MAX_ITERATIONS = + "proxima.kmedoids.cluster.max_iterations"; +static const std::string KMEDOIDS_CLUSTER_PURGE_EMPTY = + "proxima.kmedoids.cluster.purge_empty"; +static const std::string KMEDOIDS_CLUSTER_BENCH_RATIO = + "proxima.kmedoids.cluster.bench_ratio"; +static const std::string KMEDOIDS_CLUSTER_ONLY_MEANS = + "proxima.kmedoids.cluster.only_means"; +static const std::string KMEDOIDS_CLUSTER_WITHOUT_MEANS = + "proxima.kmedoids.cluster.without_means"; +static const std::string KMEDOIDS_CLUSTER_SEEKER_CLASS = + "proxima.kmedoids.cluster.seeker_class"; +static const std::string KMEDOIDS_CLUSTER_SEEKER_PARAMS = + "proxima.kmedoids.cluster.seeker_params"; + +//! Stratified +static const std::string STRATIFIED_CLUSTER_COUNT = + "proxima.stratified.cluster.count"; +static const std::string STRATIFIED_CLUSTER_FIRST_CLASS = + "proxima.stratified.cluster.first_class"; +static const std::string STRATIFIED_CLUSTER_SECOND_CLASS = + "proxima.stratified.cluster.second_class"; +static const std::string STRATIFIED_CLUSTER_FIRST_COUNT = + "proxima.stratified.cluster.first_count"; +static const std::string STRATIFIED_CLUSTER_SECOND_COUNT = + "proxima.stratified.cluster.second_count"; +static const std::string STRATIFIED_CLUSTER_FIRST_PARAMS = + "proxima.stratified.cluster.first_params"; +static const std::string STRATIFIED_CLUSTER_SECOND_PARAMS = + "proxima.stratified.cluster.second_params"; +static const std::string STRATIFIED_CLUSTER_AUTO_TUNING = + "proxima.stratified.cluster.auto_tuning"; +static const std::string STRATIFIED_CLUSTER_SECOND_POOL_COUNT = + "proxima.stratified.cluster.second_pool_count"; + +//! Gap Statistics +static const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MIN = + "proxima.gapstats.cluster_estimater.k_min"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MAX = + "proxima.gapstats.cluster_estimater.k_max"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MIN_STEP = + "proxima.gapstats.cluster_estimater.k_min_step"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MAX_STEP = + "proxima.gapstats.cluster_estimater.k_max_step"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_TRY_COUNT = + "proxima.gapstats.cluster_estimater.try_count"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_SHARD_FACTOR = + "proxima.gapstats.cluster_estimater.shard_factor"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_ENABLE_MC2 = + "proxima.gapstats.cluster_estimater.enable_mc2"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_MARKOV_CHAIN_LENGTH = + "proxima.gapstats.cluster_estimater.markov_chain_length"; +static const std::string GAPSTATS_CLUSTER_ESTIMATER_CLUSTER_CLASS = + "proxima.gapstats.cluster_estimater.cluster_class"; + +static const std::string CLUSTER_TRAINER_SAMPLE_COUNT = + "proxima.cluster.trainer.sample_count"; +static const std::string CLUSTER_TRAINER_SAMPLE_RATIO = + "proxima.cluster.trainer.sample_ratio"; +static const std::string CLUSTER_TRAINER_THREAD_COUNT = + "proxima.cluster.trainer.thread_count"; +static const std::string CLUSTER_TRAINER_FILE_NAME = + "proxima.cluster.trainer.file_name"; +static const std::string CLUSTER_TRAINER_CLASS_NAME = + "proxima.cluster.trainer.class_name"; + +static const std::string STRATIFIED_TRAINER_SAMPLE_COUNT = + "proxima.stratified.trainer.sample_count"; +static const std::string STRATIFIED_TRAINER_SAMPLE_RATIO = + "proxima.stratified.trainer.sample_ratio"; +static const std::string STRATIFIED_TRAINER_THREAD_COUNT = + "proxima.stratified.trainer.thread_count"; +static const std::string STRATIFIED_TRAINER_FILE_NAME = + "proxima.stratified.trainer.file_name"; +static const std::string STRATIFIED_TRAINER_CLASS_NAME = + "proxima.stratified.trainer.class_name"; +static const std::string STRATIFIED_TRAINER_CLUSTER_COUNT = + "proxima.stratified.trainer.cluster_count"; +static const std::string STRATIFIED_TRAINER_AUTOAUNE = + "proxima.stratified.trainer.autotune"; +static const std::string STRATIFIED_TRAINER_PARAMS_IN_LEVEL_PREFIX = + "proxima.stratified.trainer.cluster_params_in_level_"; + +} // 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..d1cdfed --- /dev/null +++ b/src/core/algorithm/cluster/kmeans_cluster.cc @@ -0,0 +1,915 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_cluster.h" +#include "framework/index_error.h" +#include "framework/index_factory.h" +#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 + virtual ~KmeansCluster(void) {} + + //! Initialize Cluster + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms); + + //! Cleanup Cluster + virtual int cleanup(void); + + //! Reset Cluster + virtual int reset(void); + + //! Update Cluster + virtual int update(const ailego::Params ¶ms); + + //! Suggest dividing to K clusters + virtual void suggest(uint32_t k); + + //! Mount features + virtual int mount(IndexFeatures::Pointer feats); + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Classify + virtual int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Label + virtual int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out); + + 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()) {} + + virtual size_t count(void) const { + return centroids_.size(); + } + + virtual size_t dimension(void) const { + return feature_dimension_; + } + + virtual const void *element(size_t i) const { + return centroids_[i].feature(); + } + + virtual IndexMeta::DataType data_type(void) const { + return data_type_; + } + + virtual size_t element_size(void) const { + 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]; + 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/opt_kmeans_cluster.cc b/src/core/algorithm/cluster/opt_kmeans_cluster.cc new file mode 100644 index 0000000..b1b7c76 --- /dev/null +++ b/src/core/algorithm/cluster/opt_kmeans_cluster.cc @@ -0,0 +1,1346 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_cluster.h" +#include "framework/index_error.h" +#include "framework/index_factory.h" +#include "cluster_params.h" + +namespace zvec { +namespace core { + +/*! Optimize K-Means cluster algorithm + */ +class OptKmeansAlgorithm : public IndexCluster { + public: + //! Constructor + OptKmeansAlgorithm(void) {} + + //! Destructor + virtual ~OptKmeansAlgorithm(void) {} + + //! Initialize Cluster + int init(const IndexMeta &meta, const ailego::Params ¶ms); + + //! Mount features + virtual int mount(IndexFeatures::Pointer feats); + + //! Suggest dividing to K clusters + virtual void suggest(uint32_t k); + + //! Classify + virtual int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Label + virtual int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out); + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s) = 0; + + //! Cleanup Cluster + virtual int cleanup(void); + + //! Reset Cluster + virtual int reset(void); + + //! Update Cluster + virtual int update(const ailego::Params ¶ms); + + 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]; + 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; + case IndexMeta::DataType::DT_BINARY32: + case IndexMeta::DataType::DT_BINARY64: + if (feats->dimension() % 32 != 0) { + LOG_ERROR( + "Unsupported feature dimension %zu (dimension of binary " + "must be an integer multiple of 32).", + 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 + virtual ~NumericalKmeansAlgorithm(void) {} + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + 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 + virtual ~NibbleKmeansAlgorithm(void) {} + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + 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; +} + +/*! Binary K-Means cluster algorithm + */ +template +class BinaryKmeansAlgorithm : 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 + BinaryKmeansAlgorithm(void) {} + + //! Destructor + virtual ~BinaryKmeansAlgorithm(void) {} + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + protected: + //! update centroids + void update_centroids(IndexCluster::CentroidList ¢s, + const ailego::BinaryKmeans &algorithm); +}; + +template +void BinaryKmeansAlgorithm::update_centroids( + IndexCluster::CentroidList ¢s, + const ailego::BinaryKmeans &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() >> 3)); + } +} + +template +int BinaryKmeansAlgorithm::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::BinaryKmeans 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(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::BinaryKmeans, 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 + virtual ~NumericalInnerProductKmeansAlgorithm(void) {} + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + 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 + virtual ~NibbleInnerProductKmeansAlgorithm(void) {} + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + 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 + virtual ~OptKmeansCluster(void) {} + + //! Initialize Cluster + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms); + + //! Cleanup Cluster + virtual int cleanup(void); + + //! Reset Cluster + virtual int reset(void); + + //! Update Cluster + virtual int update(const ailego::Params ¶ms); + + //! Suggest dividing to K clusters + virtual void suggest(uint32_t k); + + //! Mount features + virtual int mount(IndexFeatures::Pointer feats); + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Classify + virtual int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Label + virtual int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out); + + 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; + } + // TODO + case IndexMeta::DataType::DT_BINARY32: { + algorithm_.reset(new (std::nothrow) BinaryKmeansAlgorithm); + break; + } + +#if defined(AILEGO_M64) + case IndexMeta::DataType::DT_BINARY64: { + algorithm_.reset(new (std::nothrow) BinaryKmeansAlgorithm); + break; + } +#endif // AILEGO_M64 + + 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..46af115 --- /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 "framework/index_framework.h" + +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..09a031f --- /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 "framework/index_cluster.h" +#include "framework/index_error.h" +#include "framework/index_factory.h" +#include "cluster_params.h" + +namespace zvec { +namespace core { + +/*! Stratified Cluster + */ +class StratifiedCluster : public IndexCluster { + public: + //! Constructor + StratifiedCluster(void) {} + + //! Destructor + virtual ~StratifiedCluster(void) {} + + //! Initialize Cluster + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) { + meta_ = meta; + this->update_params(params); + return 0; + } + + //! Cleanup Cluster + virtual int cleanup(void) { + features_.reset(); + return 0; + } + + //! Reset Cluster + virtual int reset(void) { + features_.reset(); + return 0; + } + + //! Update Cluster + virtual int update(const ailego::Params ¶ms) { + this->update_params(params); + return 0; + } + + //! Suggest dividing to K clusters + virtual void suggest(uint32_t k) { + cluster_count_ = k; + } + + //! Mount features + virtual int mount(IndexFeatures::Pointer feats) { + if (!feats) { + return IndexError_InvalidArgument; + } + if (!feats->is_matched(meta_)) { + return IndexError_Mismatch; + } + features_ = std::move(feats); + return 0; + } + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Classify + virtual int classify(IndexThreads::Pointer threads, + IndexCluster::CentroidList ¢s); + + //! Label + virtual int label(IndexThreads::Pointer threads, + const IndexCluster::CentroidList ¢s, + std::vector *out); + + 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..27a2f91 --- /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 "framework/index_error.h" +#include "framework/index_factory.h" +#include "framework/index_helper.h" +#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..08f5875 --- /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 "framework/index_cluster.h" +#include "framework/index_trainer.h" + +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..f4d021c --- /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/diskann_builder.h b/src/core/algorithm/diskann/diskann_builder.h new file mode 100644 index 0000000..2244750 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_builder.h @@ -0,0 +1,13 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. diff --git a/src/core/algorithm/diskann/diskann_searcher.h b/src/core/algorithm/diskann/diskann_searcher.h new file mode 100644 index 0000000..2244750 --- /dev/null +++ b/src/core/algorithm/diskann/diskann_searcher.h @@ -0,0 +1,13 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. diff --git a/src/core/algorithm/flat/CMakeLists.txt b/src/core/algorithm/flat/CMakeLists.txt new file mode 100644 index 0000000..4564d8e --- /dev/null +++ b/src/core/algorithm/flat/CMakeLists.txt @@ -0,0 +1,11 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +#message(STATUS "PROJECT_ROOT_DIR = ${PROJECT_ROOT_DIR}") +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 + 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..f4636d8 --- /dev/null +++ b/src/core/algorithm/flat/flat_builder.cc @@ -0,0 +1,284 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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(FLAT_BUILDER_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(FLAT_SEARCHER_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; + } + + holder_ = nullptr; + 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(PARAM_FLAT_SEGMENT_KEYS, 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(PARAM_FLAT_SEGMENT_MAPPING, 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(PARAM_FLAT_SEGMENT_FEATURES, 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(PARAM_FLAT_SEGMENT_FEATURES, 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..78af5aa --- /dev/null +++ b/src/core/algorithm/flat/flat_builder.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 "flat_utility.h" + +namespace zvec { +namespace core { + +/*! Flat Builder + */ +template +class FlatBuilder : public IndexBuilder { + public: + //! Destructor + virtual ~FlatBuilder(void) {} + + //! 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..8c3e690 --- /dev/null +++ b/src/core/algorithm/flat/flat_index_format.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 + +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_)); + memset(index_meta, 0, sizeof(index_meta)); + } + 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}; + char index_meta[0]; +}; + +/*! 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 { + uint16_t segment_id{0}; + uint16_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(uint16_t id, bool col, uint32_t off) + : segment_id(id), column_major(col), reserved(0u), offset(off) {} + + uint16_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..f9cdcfa --- /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(PARAM_FLAT_SEGMENT_MAPPING); + if (!mapping_segment) { + LOG_ERROR("Failed to fetch segment %s", + PARAM_FLAT_SEGMENT_MAPPING.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(PARAM_FLAT_SEGMENT_KEYS); + if (!keys_segment) { + LOG_ERROR("Failed to fetch segment %s", PARAM_FLAT_SEGMENT_KEYS.c_str()); + return IndexError_NoExist; + } + features_segment_ = cntr->get(PARAM_FLAT_SEGMENT_FEATURES); + if (!features_segment_) { + LOG_ERROR("Failed to fetch segment %s", PARAM_FLAT_SEGMENT_KEYS.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..e9a6151 --- /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 "framework/index_searcher.h" +#include "flat_distance_matrix.h" +#include "flat_index_format.h" + +namespace zvec { +namespace core { + +/*! Flat Searcher + */ +template +class FlatSearcher : public IndexSearcher { + public: + //! Destructor + virtual ~FlatSearcher(void) = default; + + //! Initialize Searcher + int init(const ailego::Params &index_params) override { + params_ = index_params; + read_block_size_ = FLAT_READ_BLOCK_SIZE; + index_params.get(FLAT_SEARCHER_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::map key_id_mapping_; + uint32_t magic_{IndexContext::GenerateMagic()}; + uint32_t read_block_size_{FLAT_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..f92b679 --- /dev/null +++ b/src/core/algorithm/flat/flat_searcher_context.h @@ -0,0 +1,1107 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 + virtual ~FlatSearcherContext(void) {} + + //! 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 + 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]; + } + + //! 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..4eb6383 --- /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) { + 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(FLAT_BUILDER_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_READ_BLOCK_SIZE; + params.get(FLAT_SEARCHER_READ_BLOCK_SIZE, &read_block_size_); + params.get(FLAT_SEARCHER_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 (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; + entity_.get_vector_by_key(key, block); + 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; + entity_.get_vector_by_key(key, block); + 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; + entity_.get_vector_by_key(key, block); + 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..dde1f40 --- /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); + virtual ~FlatStreamer(void); + + 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_; + } + + virtual const void *get_vector(uint64_t key) const override { + return this->get_vector_by_key(key); + } + + virtual 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_{}; + FlatStreamerEntity entity_; + bool column_major_order_{false}; + bool use_key_info_map_{true}; + uint32_t read_block_size_{0}; +}; + +} // 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..24cfd9e --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer_context.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 "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 + virtual ~FlatStreamerContext(void) = 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 + 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]; + } + + //! 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); + } + } + } + + 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 {} + + //! Reset the context + void reset(const FlatStreamer *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; + 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..7dd5c0d --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer_dumper.h @@ -0,0 +1,243 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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" +#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(PARAM_FLAT_SEGMENT_FEATURES, 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(PARAM_FLAT_SEGMENT_FEATURES, 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(PARAM_FLAT_SEGMENT_KEYS, 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(PARAM_FLAT_SEGMENT_MAPPING, 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..00ac983 --- /dev/null +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -0,0 +1,994 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "flat_utility.h" +#include "index_error.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); + 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 (!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(PARAM_FLAT_SEGMENT_FEATURES, 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(PARAM_FLAT_SEGMENT_FEATURES, 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()); + 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..8a2afd8 --- /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(PARAM_FLAT_SEGMENT_FEATURES, 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_{}; + StreamerLinearMeta meta_{}; + 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_{}; +}; + +} // 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..c84e741 --- /dev/null +++ b/src/core/algorithm/flat/flat_utility.h @@ -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. +#pragma once + +#include +#include +#include +#include "framework/index_error.h" +#include "framework/index_factory.h" +#include "framework/index_meta.h" +#include "framework/index_metric.h" + +namespace zvec { +namespace core { + +//! The default size of reading a block +static constexpr uint32_t FLAT_READ_BLOCK_SIZE = 4 * 1024 * 1024; +static const std::string FLAT_LINEAR_META_SEG_ID = "linear.linear_meta"; +static const std::string FLAT_LINEAR_LIST_HEAD_SEG_ID = + "linear.linear_list_head"; + +// reducer params +static const std::string PARAM_FLAT_REDUCER_WORKING_PATH( + "proxima.linear.reducer.working_path"); +static const std::string PARAM_FLAT_REDUCER_NUM_OF_ADD_THREADS( + "proxima.linear.reducer.num_of_add_threads"); +static const std::string PARAM_FLAT_REDUCER_INDEX_NAME( + "proxima.linear.reducer.index_name"); + +// index params +static const std::string PARAM_FLAT_SEGMENT_KEYS("linear.keys"); +static const std::string PARAM_FLAT_SEGMENT_FEATURES("linear.features"); +static const std::string PARAM_FLAT_SEGMENT_MAPPING("linear.mapping"); + +static const std::string FLAT_BUILDER_COLUMN_MAJOR_ORDER( + "proxima.linear.builder.column_major_order"); +static const std::string FLAT_SEARCHER_BATCH_SIZE( + "proxima.linear.searcher.batch_size"); +static const std::string FLAT_SEARCHER_READ_BLOCK_SIZE( + "proxima.linear.searcher.read_block_size"); +static const std::string FLAT_SEARCHER_USE_ID_MAP( + "proxima.linear.searcher.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 (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..e27d2d3 --- /dev/null +++ b/src/core/algorithm/flat_sparse/CMakeLists.txt @@ -0,0 +1,11 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +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 + 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..1ed689f --- /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..4e9f2b1 --- /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..48f5a6a --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_context.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 +#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 + virtual ~FlatSparseContext(void) = 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(size_t idx) { + 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..7b1e0f8 --- /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..e141480 --- /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..0f50b4c --- /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 "framework/index_logger.h" +#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 + virtual uint32_t sparse_count() const override { + return sparse_count_; + } + + //! Retrieve sparse indices + virtual const uint32_t *sparse_indices() const override { + return reinterpret_cast(sparse_indices_buffer_.data()); + } + + //! Retrieve sparse data + virtual const void *sparse_data() const override { + return reinterpret_cast(sparse_data_buffer_.data()); + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt() && valid_; + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + // std::cout << "iter key=" << cur_id_ << std::endl; + return entity_->get_key(cur_id_); + } + + //! Next iterator + virtual 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..f77c28d --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_search.h @@ -0,0 +1,151 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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()); + 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..4094825 --- /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..23002ec --- /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); + virtual ~FlatSparseSearcher(void); + + 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..f0d331a --- /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..32f8d8d --- /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..7d2815d --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_streamer.cc @@ -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. + +#include "flat_sparse_streamer.h" +#include +#include +#include +#include "framework/index_error.h" +#include "framework/index_meta.h" +#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("Add vector failed, dim size too larg, dim_size=%u, key=%zu", + sparse_count, (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("Add vector failed, dim size too larg, dim_size=%u, key=%zu", + sparse_count, (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..fa443dc --- /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); + virtual ~FlatSparseStreamer(void); + + 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_{}; + FlatSparseStreamerEntity entity_; + + uint32_t magic_{0U}; + Stats stats_{}; + 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..3e294a6 --- /dev/null +++ b/src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.cc @@ -0,0 +1,1082 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "ailego/utility/time_helper.h" +#include "framework/index_error.h" +#include "framework/index_logger.h" +#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..a864555 --- /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..4566b5d --- /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 = 4096; +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..f4a1054 --- /dev/null +++ b/src/core/algorithm/hnsw/CMakeLists.txt @@ -0,0 +1,11 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +cc_library( + NAME core_knn_hnsw + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework sparsehash + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + 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..9a41258 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_algorithm.cc @@ -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. +#include "hnsw_algorithm.h" +#include +#include +#include + +namespace zvec { +namespace core { + +HnswAlgorithm::HnswAlgorithm(HnswEntity &entity) + : entity_(entity), + mt_(std::chrono::system_clock::now().time_since_epoch().count()), + lock_pool_(kLockCnt) {} + +int HnswAlgorithm::cleanup() { + return 0; +} + +int HnswAlgorithm::add_node(node_id_t id, level_t level, HnswContext *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 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); + + if (ctx->group_by_search()) { + expand_neighbors_by_group(topk_heap, ctx); + } + + return 0; +} + +int HnswAlgorithm::fast_search(HnswContext *ctx) const { + auto max_level = entity_.cur_max_level(); + auto entry_point = entity_.entry_point(); + + dist_t dist = ctx->dist_calculator().dist(entry_point); + for (level_t cur_level = max_level; 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 HnswAlgorithm::select_entry_point(level_t level, node_id_t *entry_point, + dist_t *dist, HnswContext *ctx) const { + auto &entity = ctx->get_entity(); + HnswDistCalculator &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 = entity.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; + + float dists[size]; + const void *neighbor_vecs[size]; + for (uint32_t i = 0; i < size; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + + dc.batch_dist(neighbor_vecs, size, dists); + + 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; +} + +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; +} + +void HnswAlgorithm::search_neighbors(level_t level, node_id_t *entry_point, + dist_t *dist, TopkHeap &topk, + HnswContext *ctx) const { + const auto &entity = ctx->get_entity(); + HnswDistCalculator &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())++; + } + + node_id_t 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(neighbor_ids, 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 + + float dists[size]; + const void *neighbor_vecs[size]; + + for (uint32_t i = 0; i < size; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + + dc.batch_dist(neighbor_vecs, size, dists); + + 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); + } + } // end if + } // end for + } // while + + return; +} + +void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, + HnswContext *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(); + 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(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())++; + } + + node_id_t 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(neighbor_ids, size, neighbor_vec_blocks); + 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_vec_blocks[prefetch_id].data()); + } + dist_t cur_dist = dc.dist(neighbor_vec_blocks[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 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; +} + +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; +} + +} // 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..b345484 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_algorithm.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 +#include +#include "hnsw_context.h" +#include "hnsw_dist_calculator.h" +#include "hnsw_entity.h" + +namespace zvec { +namespace core { + +//! hnsw graph algorithm implement +class HnswAlgorithm { + public: + typedef std::unique_ptr UPointer; + + public: + //! Constructor + explicit HnswAlgorithm(HnswEntity &entity); + + //! Destructor + ~HnswAlgorithm() = default; + + //! Cleanup HnswAlgorithm + 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, HnswContext *ctx); + + //! do knn search in graph with lock + //! return 0 on success, or errCode in failure. results saved in ctx + int search(HnswContext *ctx) const; + + //! do knn search in graph without lock + //! return 0 on success, or errCode in failure. results saved in ctx + int fast_search(HnswContext *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, + 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 + //! 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, HnswContext *ctx) 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}; + + HnswEntity &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_builder.cc b/src/core/algorithm/hnsw/hnsw_builder.cc new file mode 100644 index 0000000..1ff0cbd --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_builder.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 "hnsw_builder.h" +#include +#include +#include +#include "framework/index_error.h" +#include "framework/index_factory.h" +#include "framework/index_logger.h" +#include "hnsw_algorithm.h" +#include "hnsw_params.h" + +namespace zvec { +namespace core { + +HnswBuilder::HnswBuilder() {} + +int HnswBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { + LOG_INFO("Begin HnswBuilder::init"); + + meta_ = meta; + auto params_copy = params; + meta_.set_builder("HnswBuilder", HnswEntity::kRevision, + std::move(params_copy)); + + size_t memory_quota = 0UL; + params.get(PARAM_HNSW_BUILDER_MEMORY_QUOTA, &memory_quota); + params.get(PARAM_HNSW_BUILDER_THREAD_COUNT, &thread_cnt_); + params.get(PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT, &neighbor_cnt_); + params.get(PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_); + float ratio = HnswEntity::kDefaultNeighborPruneRatio; + params.get(PARAM_HNSW_BUILDER_UPPER_NEIGHBOR_RATIO, &ratio); + upper_neighbor_cnt_ = ratio * neighbor_cnt_; + params.get(PARAM_HNSW_BUILDER_EFCONSTRUCTION, &ef_construction_); + params.get(PARAM_HNSW_BUILDER_SCALING_FACTOR, &scaling_factor_); + params.get(PARAM_HNSW_BUILDER_CHECK_INTERVAL_SECS, &check_interval_secs_); + ratio = HnswEntity::kDefaultNeighborPruneRatio; + params.get(PARAM_HNSW_BUILDER_NEIGHBOR_PRUNE_RATIO, &ratio); + + size_t prune_cnt = neighbor_cnt_ * ratio; + + if (ef_construction_ == 0) { + ef_construction_ = HnswEntity::kDefaultEfConstruction; + } + if (neighbor_cnt_ == 0) { + neighbor_cnt_ = HnswEntity::kDefaultNeighborCnt; + } + if (neighbor_cnt_ > kMaxNeighborCnt) { + LOG_ERROR("[%s] must be in range (0,%d]", + PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (min_neighbor_cnt_ > neighbor_cnt_) { + LOG_ERROR("[%s]-[%d] must be <= [%s]-[%d]", + PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT.c_str(), min_neighbor_cnt_, + PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), neighbor_cnt_); + return IndexError_InvalidArgument; + } + if (upper_neighbor_cnt_ == 0) { + neighbor_cnt_ = HnswEntity::kDefaultUpperNeighborCnt; + } + if (upper_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) { + LOG_ERROR("UpperNeighborCnt must be in range (0,%d)", + HnswEntity::kMaxNeighborCnt); + 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_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_BUILDER_THREAD_COUNT.c_str(), + std::thread::hardware_concurrency()); + } + if (prune_cnt == 0UL) { + prune_cnt = upper_neighbor_cnt_; + } + + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + int ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMetric init failed, ret=%d", ret); + return ret; + } + + entity_.set_vector_size(meta_.element_size()); + + entity_.set_ef_construction(ef_construction_); + entity_.set_neighbor_cnt(neighbor_cnt_); + entity_.set_min_neighbor_cnt(min_neighbor_cnt_); + entity_.set_upper_neighbor_cnt(upper_neighbor_cnt_); + entity_.set_scaling_factor(scaling_factor_); + entity_.set_memory_quota(memory_quota); + entity_.set_prune_cnt(prune_cnt); + + ret = entity_.init(); + if (ret != 0) { + return ret; + } + + alg_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_)); + + ret = alg_->init(); + if (ret != 0) { + return ret; + } + + state_ = BUILD_STATE_INITED; + LOG_INFO( + "End HnswBuilder::init, params: vectorSize=%u efConstruction=%u " + "neighborCnt=%u upperNeighborCnt=%u scalingFactor=%u " + "memoryQuota=%zu neighborPruneCnt=%zu metricName=%s ", + meta_.element_size(), ef_construction_, neighbor_cnt_, + upper_neighbor_cnt_, scaling_factor_, memory_quota, prune_cnt, + meta_.metric_name().c_str()); + + return 0; +} + +int HnswBuilder::cleanup(void) { + LOG_INFO("Begin HnswBuilder::cleanup"); + + neighbor_cnt_ = HnswEntity::kDefaultNeighborCnt; + min_neighbor_cnt_ = 0; + upper_neighbor_cnt_ = HnswEntity::kDefaultUpperNeighborCnt; + ef_construction_ = HnswEntity::kDefaultEfConstruction; + scaling_factor_ = HnswEntity::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 HnswBuilder::cleanup"); + + return 0; +} + +int HnswBuilder::train(IndexThreads::Pointer, IndexHolder::Pointer holder) { + if (state_ != BUILD_STATE_INITED) { + LOG_ERROR("Init the builder before HnswBuilder::train"); + return IndexError_NoReady; + } + + if (!holder) { + LOG_ERROR("Input holder is nullptr while training index"); + return IndexError_InvalidArgument; + } + if (!holder->is_matched(meta_)) { + LOG_ERROR("Input holder doesn't match index meta while training index"); + return IndexError_Mismatch; + } + LOG_INFO("Begin HnswBuilder::train"); + size_t trained_cost_time = 0; + size_t trained_count = 0; + + if (metric_->support_train()) { + auto start_time = ailego::Monotime::MilliSeconds(); + auto iter = holder->create_iterator(); + if (!iter) { + LOG_ERROR("Create iterator for holder failed"); + return IndexError_Runtime; + } + while (iter->is_valid()) { + int ret = metric_->train(iter->data(), meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Hnsw build measure train failed, ret=%d", ret); + return ret; + } + iter->next(); + ++trained_count; + } + trained_cost_time = ailego::Monotime::MilliSeconds() - start_time; + } + stats_.set_trained_count(trained_count); + stats_.set_trained_costtime(trained_cost_time); + state_ = BUILD_STATE_TRAINED; + + LOG_INFO("End HnswBuilder::train"); + + return 0; +} + +int HnswBuilder::train(const IndexTrainer::Pointer & /*trainer*/) { + if (state_ != BUILD_STATE_INITED) { + LOG_ERROR("Init the builder before HnswBuilder::train"); + return IndexError_NoReady; + } + + LOG_INFO("Begin HnswBuilder::train by trainer"); + + stats_.set_trained_count(0UL); + stats_.set_trained_costtime(0UL); + state_ = BUILD_STATE_TRAINED; + + LOG_INFO("End HnswBuilder::train by trainer"); + + return 0; +} + +int HnswBuilder::build(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) { + if (state_ != BUILD_STATE_TRAINED) { + LOG_ERROR("Train the index before HnswBuilder::build"); + return IndexError_NoReady; + } + + 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 HnswBuilder::build"); + + if (holder->count() != static_cast(-1)) { + LOG_DEBUG("HnswBuilder holder documents count %lu", holder->count()); + int ret = entity_.reserve_space(holder->count()); + if (ret != 0) { + LOG_ERROR("HnswBuilde reserver space failed"); + return ret; + } + } + auto iter = 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; + + const void *vec = iter->data(); + ret = entity_.add_vector(level, iter->key(), vec, &id); + if (ailego_unlikely(ret != 0)) { + return ret; + } + iter->next(); + } + // Holder is not needed, cleanup it. + holder.reset(); + + LOG_INFO("Finished save vector, start build graph..."); + + 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, &HnswBuilder::do_build, i, + threads->count(), &finished)); + } + + while (!task_group->is_finished()) { + std::unique_lock lk(mutex_); + 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(); + + stats_.set_built_count(finished.load()); + stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time); + state_ = BUILD_STATE_BUILT; + + LOG_INFO("End HnswBuilder::build"); + return 0; +} + +void HnswBuilder::do_build(node_id_t idx, size_t step_size, + std::atomic *finished) { + AILEGO_DEFER([&]() { + std::lock_guard latch(mutex_); + cond_.notify_one(); + }); + HnswContext *ctx = new (std::nothrow) + HnswContext(meta_.dimension(), metric_, + std::shared_ptr(&entity_, [](HnswEntity *) {})); + if (ailego_unlikely(ctx == nullptr)) { + if (!error_.exchange(true)) { + LOG_ERROR("Failed to create context"); + errcode_ = IndexError_NoMemory; + } + return; + } + HnswContext::Pointer auto_ptr(ctx); + ctx->set_max_scan_num(entity_.doc_cnt()); + int ret = ctx->init(HnswContext::kBuilderContext); + if (ret != 0) { + if (!error_.exchange(true)) { + LOG_ERROR("Failed to init context"); + errcode_ = IndexError_Runtime; + } + return; + } + + IndexQueryMeta qmeta(meta_.data_type(), meta_.dimension()); + for (node_id_t id = idx; id < entity_.doc_cnt(); id += step_size) { + ctx->reset_query(entity_.get_vector(id)); + 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 HnswBuilder::dump(const IndexDumper::Pointer &dumper) { + if (state_ != BUILD_STATE_BUILT) { + LOG_INFO("Build the index before HnswBuilder::dump"); + return IndexError_NoReady; + } + + LOG_INFO("Begin HnswBuilder::dump"); + + meta_.set_searcher("HnswSearcher", HnswEntity::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("HnswBuilder dump index failed"); + return ret; + } + + stats_.set_dumped_count(entity_.doc_cnt()); + stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time); + + LOG_INFO("EndHnswBuilder::dump"); + return 0; +} + +INDEX_FACTORY_REGISTER_BUILDER(HnswBuilder); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_builder.h b/src/core/algorithm/hnsw/hnsw_builder.h new file mode 100644 index 0000000..7536fe4 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_builder.h @@ -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. +#pragma once + +#include +#include "framework/index_builder.h" +#include "hnsw_algorithm.h" +#include "hnsw_builder_entity.h" + +namespace zvec { +namespace core { + +class HnswBuilder : public IndexBuilder { + public: + //! Constructor + HnswBuilder(); + + //! 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, + 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_; + } + + private: + 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 + }; + + HnswBuilderEntity entity_{}; + HnswAlgorithm::UPointer alg_; // impl graph algorithm + uint32_t thread_cnt_{0}; + uint32_t neighbor_cnt_{HnswEntity::kDefaultNeighborCnt}; + uint32_t min_neighbor_cnt_{0}; + uint32_t upper_neighbor_cnt_{HnswEntity::kDefaultUpperNeighborCnt}; + uint32_t ef_construction_{HnswEntity::kDefaultEfConstruction}; + uint32_t scaling_factor_{HnswEntity::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/hnsw_builder_entity.cc b/src/core/algorithm/hnsw/hnsw_builder_entity.cc new file mode 100644 index 0000000..380eade --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_builder_entity.cc @@ -0,0 +1,198 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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_builder_entity.h" +#include +#include +#include "utility/sparse_utility.h" + +namespace zvec { +namespace core { + +HnswBuilderEntity::HnswBuilderEntity() { + update_ep_and_level(kInvalidNodeId, 0U); +} + +int HnswBuilderEntity::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->HnswEntity::cleanup(); + + return 0; +} + +int HnswBuilderEntity::init() { + size_t size = vector_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 HnswBuilderEntity::reserve_space(size_t docs) { + if (memory_quota_ > 0 && (node_size() * docs + neighbors_size_ * docs + + sizeof(NeighborIndex) * 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); + + return 0; +} + +int HnswBuilderEntity::add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) { + if (memory_quota_ > 0 && + (vectors_buffer_.capacity() + keys_buffer_.capacity() + + neighbors_buffer_.capacity() + upper_neighbors_buffer_.capacity() + + neighbors_index_.capacity() * sizeof(NeighborIndex)) > memory_quota_) { + LOG_ERROR("Add vector failed, used memory exceed quota, cur_doc=%u", + doc_cnt()); + return IndexError_NoMemory; + } + + vectors_buffer_.append(reinterpret_cast(vec), vector_size()); + vectors_buffer_.append(padding_size_, '\0'); + keys_buffer_.append(reinterpret_cast(&key), sizeof(key)); + + // 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 HnswBuilderEntity::get_key(node_id_t id) const { + return *(reinterpret_cast(keys_buffer_.data() + + id * sizeof(key_t))); +} + +const void *HnswBuilderEntity::get_vector(node_id_t id) const { + return vectors_buffer_.data() + id * node_size(); +} + +int HnswBuilderEntity::get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const { + const void *vec = get_vector(id); + block.reset((void *)vec); + return 0; +} + +int HnswBuilderEntity::get_vector(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 HnswBuilderEntity::get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + const void *vecs[count]; + get_vector(ids, count, vecs); + for (uint32_t i = 0; i < count; ++i) { + vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i])); + } + return 0; +} + +const Neighbors HnswBuilderEntity::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 HnswBuilderEntity::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 HnswBuilderEntity::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 HnswBuilderEntity::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/hnsw_builder_entity.h b/src/core/algorithm/hnsw/hnsw_builder_entity.h new file mode 100644 index 0000000..fb303df --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_builder_entity.h @@ -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. +#pragma once + +#include +#include "hnsw_entity.h" + +namespace zvec { +namespace core { + +class HnswBuilderEntity : public HnswEntity { + public: + //! Add vector and key to hnsw entity, and local id will be saved to id + virtual int add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) override; + + //! Get primary key of the node id + virtual key_t get_key(node_id_t id) const override; + + //! Get vector feature data by key + virtual const void *get_vector(node_id_t id) const override; + + //! Batch get vectors feature data by keys + virtual int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + + virtual int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + virtual 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 + 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 + virtual const Neighbors get_neighbors(level_t level, + node_id_t id) const override; + + //! Replace node id in level's neighbors + virtual int update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) override; + + //! add a neighbor to id in graph level + virtual void add_neighbor(level_t level, node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + + //! Dump the hnsw graph to dumper + virtual int dump(const IndexDumper::Pointer &dumper) override; + + //! Cleanup the entity + virtual int cleanup(void) override; + + public: + //! Constructor + HnswBuilderEntity(); + + //! 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 + int reserve_space(size_t docs); + + //! 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) + 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: + HnswBuilderEntity(const HnswBuilderEntity &) = delete; + HnswBuilderEntity &operator=(const HnswBuilderEntity &) = delete; + + private: + friend class HnswSearcherEntity; + //! class internal used only + struct NeighborIndex { + NeighborIndex(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/hnsw_chunk.cc b/src/core/algorithm/hnsw/hnsw_chunk.cc new file mode 100644 index 0000000..166aa04 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_chunk.cc @@ -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. +#include "hnsw_chunk.h" +#include +#include +#include +#include +#include "framework/index_error.h" +#include "framework/index_helper.h" +#include "framework/index_logger.h" +#include "framework/index_streamer.h" + +namespace zvec { +namespace core { + +int ChunkBroker::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 ChunkBroker::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(&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 ChunkBroker::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 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..f042672 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_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 "framework/index_error.h" +#include "framework/index_logger.h" +#include "framework/index_storage.h" +#include "framework/index_streamer.h" + +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_MAX = 8 + }; + static constexpr size_t kDefaultChunkSeqId = 0UL; + + ChunkBroker(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: + ChunkBroker(const ChunkBroker &) = delete; + ChunkBroker &operator=(const ChunkBroker &) = delete; + + struct HnswChunkMeta { + HnswChunkMeta(void) { + memset(this, 0, sizeof(HnswChunkMeta)); + } + void clear() { + memset(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/hnsw_context.cc b/src/core/algorithm/hnsw/hnsw_context.cc new file mode 100644 index 0000000..c832bc5 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_context.cc @@ -0,0 +1,296 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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_probility_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + update_heap_.limit(entity_->neighbor_cnt() + 1); + break; + + case kSearcherContext: + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_, + negative_probility_); + 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_probility_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + + update_heap_.limit(entity_->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; + 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_probility_; + 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_probility_) > 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_probility_); + 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_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_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(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->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 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..88c6f4f --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_context.h @@ -0,0 +1,530 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_context.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 + virtual ~HnswContext(); + + public: + //! Set topk of search result + virtual void set_topk(uint32_t val) override { + topk_ = val; + topk_heap_.limit(std::max(val, ef_)); + } + + //! 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 uint32_t magic(void) const override { + return 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 this->debugging(); + } + + //! Retrieve string of debug + virtual 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 + virtual 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_; + } + + 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_) { + 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); + } + } + + 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_) { + group_results_[idx][i].mutable_docs()->emplace_back( + entity_->get_key(id), score, id, entity_->get_vector(id)); + } 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 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_probility(float v) { + negative_probility_ = 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::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_; + + 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_probility_{HnswEntity::kDefaultBFNegativeProbility}; + uint32_t ef_{HnswEntity::kDefaultEf}; + 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_; +}; + +} // 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..cba6918 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_dist_calculator.h @@ -0,0 +1,218 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_meta.h" +#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_++; + + 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_); + } + + //! 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(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); + } + + 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 = 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; + } + + 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..b521db4 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_entity.cc @@ -0,0 +1,362 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_stats.h" +#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; + + char padding[padding_size]; + memset(padding, 0, sizeof(padding)); + 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, 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, 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; + node_id_t mapping[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() <= 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; + } + 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; + node_id_t 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"); + memset(buffer, 0, sizeof(buffer)); + 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, sizeof(buffer)) != sizeof(buffer)) { + LOG_ERROR("Dump graph neighbor id=%u failed, size %lu", id, + sizeof(buffer)); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(buffer, sizeof(buffer), crc); + offset += sizeof(buffer); + } + } + 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..abbb805 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_entity.h @@ -0,0 +1,527 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "ailego/logger/logger.h" +#include "framework/index_dumper.h" +#include "framework/index_error.h" +#include "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; +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 max_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 thumb_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(this, &header, sizeof(header)); + } + + HNSWHeader &operator=(const HNSWHeader &header) { + memcpy(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(this, 0, sizeof(HNSWHeader)); + hnsw.entry_point = kInvalidNodeId; + graph.size = sizeof(GraphHeader); + hnsw.size = sizeof(HnswHeader); + } + + size_t neighbor_cnt() const { + return graph.max_neighbor_count; + } + + size_t upper_neighbor_cnt() const { + return hnsw.thumb_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; + 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 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.max_neighbor_count + : header_.hnsw.thumb_neighbor_count; + } + + //! get max neighbor size of graph level 0 + inline size_t neighbor_cnt() const { + return header_.graph.max_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.thumb_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_neighbor_cnt(size_t cnt) { + header_.graph.max_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.thumb_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; + } + + 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 kDefaultNeighborCnt = 100; + constexpr static uint32_t kDefaultUpperNeighborCnt = 50; + 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 kDefaultBFNegativeProbility = 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 kDefaultNeighborPruneRatio = + 0.5f; // prune count / neighbor count + constexpr static float kDefaultUpperNeighborRatio = + 0.5f; // upper neighbor count / neighbor count + + 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..1557dcd --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_index_hash.h @@ -0,0 +1,231 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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, 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: + Chunk::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(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) { + 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(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; + } + //! 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(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; + } + //! 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: + 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..b5a6f31 --- /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 "framework/index_provider.h" +#include "framework/index_searcher.h" +#include "framework/index_streamer.h" +#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 + virtual const void *data(void) const override { + return entity_->get_vector(cur_id_); + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt(); + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + return entity_->get_key(cur_id_); + } + + //! Next iterator + virtual 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..7a4aaad --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_params.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 + +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_RATIO( + "proxima.hnsw.builder.neighbor_prune_ratio"); +static const std::string PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT( + "proxima.hnsw.builder.max_neighbor_count"); +static const std::string PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT( + "proxima.hnsw.builder.min_neighbor_count"); +static const std::string PARAM_HNSW_BUILDER_UPPER_NEIGHBOR_RATIO( + "proxima.hnsw.builder.upper_neighbor_ratio"); +static const std::string PARAM_HNSW_BUILDER_HYBRID_VECTOR_ENABLE( + "proxima.hnsw.builder.hybrid_vector_enable"); +static const std::string PARAM_HNSW_BUILDER_HYBRID_VECTOR_SEPARATE_NEIGHBOR( + "proxima.hnsw.builder.hybrid_vector_separate_neighbor"); +static const std::string PARAM_HNSW_BUILDER_SPARSE_NEIGHBOR_RATIO( + "proxima.hnsw.builder.sparse_neighbor_ratio"); + +static const std::string PARAM_HNSW_SEARCHER_EF("proxima.hnsw.searcher.ef"); +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_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_UPPER_NEIGHBOR_RATIO( + "proxima.hnsw.streamer.upper_neighbor_ratio"); +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_RATIO( + "proxima.hnsw.streamer.neighbor_prune_ratio"); +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_HYBRID_VECTOR_ENABLE( + "proxima.hnsw.streamer.hybrid_vector_enable"); +static const std::string PARAM_HNSW_STREAMER_HYBRID_VECTOR_SEPARATE_NEIGHBOR( + "proxima.hnsw.streamer.hybrid_vector_separate_neighbor"); +static const std::string PARAM_HNSW_STREAMER_SPARSE_NEIGHBOR_RATIO( + "proxima.hnsw.streamer.sparse_neighbor_ratio"); +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"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_searcher.cc b/src/core/algorithm/hnsw/hnsw_searcher.cc new file mode 100644 index 0000000..ac3cdef --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_searcher.cc @@ -0,0 +1,460 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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_searcher.h" +#include "hnsw_algorithm.h" +#include "hnsw_index_provider.h" +#include "hnsw_params.h" + +namespace zvec { +namespace core { + +HnswSearcher::HnswSearcher() {} + +HnswSearcher::~HnswSearcher() {} + +int HnswSearcher::init(const ailego::Params &search_params) { + params_ = search_params; + params_.get(PARAM_HNSW_SEARCHER_EF, &ef_); + params_.get(PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_); + params_.get(PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_); + params_.get(PARAM_HNSW_SEARCHER_CHECK_CRC_ENABLE, &check_crc_enabled_); + params_.get(PARAM_HNSW_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE, + &neighbors_in_memory_enabled_); + params_.get(PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB, + &bf_negative_probility_); + params_.get(PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + params_.get(PARAM_HNSW_SEARCHER_FORCE_PADDING_RESULT_ENABLE, + &force_padding_topk_enabled_); + + if (ef_ == 0) { + ef_ = HnswEntity::kDefaultEf; + } + if (bf_negative_probility_ <= 0.0f || bf_negative_probility_ >= 1.0f) { + LOG_ERROR("[%s] must be in range (0,1)", + PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB.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", + ef_, max_scan_ratio_, bf_enabled_, check_crc_enabled_, + neighbors_in_memory_enabled_, bf_negative_probility_, + bruteforce_threshold_, force_padding_topk_enabled_); + + return 0; +} + +void HnswSearcher::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 HnswSearcher::cleanup() { + LOG_INFO("Begin HnswSearcher:cleanup"); + + metric_.reset(); + meta_.clear(); + stats_.clear_attributes(); + stats_.set_loaded_count(0UL); + stats_.set_loaded_costtime(0UL); + max_scan_ratio_ = HnswEntity::kDefaultScanRatio; + max_scan_num_ = 0U; + ef_ = HnswEntity::kDefaultEf; + bf_enabled_ = false; + bf_negative_probility_ = HnswEntity::kDefaultBFNegativeProbility; + bruteforce_threshold_ = HnswEntity::kDefaultBruteForceThreshold; + check_crc_enabled_ = false; + neighbors_in_memory_enabled_ = false; + entity_.cleanup(); + state_ = STATE_INIT; + + LOG_INFO("End HnswSearcher:cleanup"); + + return 0; +} + +int HnswSearcher::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 HnswSearcher: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("HnswSearcher load index failed"); + return ret; + } + + alg_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_)); + + if (metric) { + metric_ = metric; + } else { + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("CreateMetric 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("IndexMetric 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 HnswSearcher::load"); + + return 0; +} + +int HnswSearcher::unload() { + LOG_INFO("HnswSearcher 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 HnswSearcher::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_num(max_scan_num_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + + return ctx->update_context(HnswContext::kSearcherContext, meta_, metric_, + entity, magic_); +} + +int HnswSearcher::search_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; + } + 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 + int ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + ctx->resize_results(count); + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + int ret = alg_->fast_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; +} + +int HnswSearcher::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; + } + 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 + int ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + 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().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 { + for (size_t q = 0; q < count; ++q) { + 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))) { + dist_t dist = ctx->dist_calculator().dist(id); + 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 HnswSearcher::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; + } + + 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 + int ret = update_context(ctx); + if (ret != 0) { + return ret; + } + } + + ctx->clear(); + 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().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 { + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + 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); + query = static_cast(query) + qmeta.element_size(); + } + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } + + return 0; +} + +IndexSearcher::Context::Pointer HnswSearcher::create_context() const { + if (ailego_unlikely(state_ != STATE_LOADED)) { + LOG_ERROR("Load the index first before create context"); + return Context::Pointer(); + } + const HnswEntity::Pointer search_ctx_entity = entity_.clone(); + if (!search_ctx_entity) { + LOG_ERROR("Failed to create search context entity"); + return Context::Pointer(); + } + HnswContext *ctx = new (std::nothrow) + HnswContext(meta_.dimension(), metric_, search_ctx_entity); + if (ailego_unlikely(ctx == nullptr)) { + LOG_ERROR("Failed to new HnswContext"); + 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_probility(bf_negative_probility_); + 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::kSearcherContext)) != 0) { + LOG_ERROR("Init HnswContext failed"); + delete ctx; + return Context::Pointer(); + } + + return Context::Pointer(ctx); +} + +IndexProvider::Pointer HnswSearcher::create_provider(void) const { + LOG_DEBUG("HnswSearcher create provider"); + + auto entity = entity_.clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone HnswEntity failed"); + return Provider::Pointer(); + } + return Provider::Pointer( + new (std::nothrow) HnswIndexProvider(meta_, entity, "HnswSearcher")); +} + +const void *HnswSearcher::get_vector(uint64_t key) const { + return entity_.get_vector_by_key(key); +} + +INDEX_FACTORY_REGISTER_SEARCHER(HnswSearcher); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_searcher.h b/src/core/algorithm/hnsw/hnsw_searcher.h new file mode 100644 index 0000000..306e2c9 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_searcher.h @@ -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. +#pragma once + +#include "framework/index_framework.h" +#include "hnsw_searcher_entity.h" +#include "hnsw_streamer.h" + +namespace zvec { +namespace core { + +class HnswSearcher : public IndexSearcher { + public: + using ContextPointer = IndexSearcher::Context::Pointer; + + public: + HnswSearcher(void); + ~HnswSearcher(void); + + HnswSearcher(const HnswSearcher &) = delete; + HnswSearcher &operator=(const HnswSearcher &) = delete; + + protected: + //! Initialize Searcher + virtual int init(const ailego::Params ¶ms) override; + + //! Cleanup Searcher + virtual int cleanup(void) override; + + //! Load Index from storage + virtual int load(IndexStorage::Pointer container, + IndexMetric::Pointer metric) override; + + //! Unload index from storage + virtual int unload(void) override; + + //! KNN Search + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_impl(query, qmeta, 1, context); + } + + //! KNN Search + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + ContextPointer &context) const override; + + //! Linear Search + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + ContextPointer &context) const override { + return search_bf_impl(query, qmeta, 1, context); + } + + //! Linear Search + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + ContextPointer &context) const override; + + //! Linear search by primary keys + virtual 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 + virtual 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 + virtual const void *get_vector(uint64_t key) const override; + + //! Create a searcher context + virtual ContextPointer create_context() const override; + + //! Create a new iterator + virtual IndexProvider::Pointer create_provider(void) const override; + + //! Retrieve statistics + virtual const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + virtual const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve params of index + virtual const ailego::Params ¶ms(void) const override { + return params_; + } + + virtual 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(HnswContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; + + HnswSearcherEntity entity_{}; + HnswAlgorithm::UPointer alg_; // impl graph algorithm + + IndexMetric::Pointer metric_{}; + IndexMeta meta_{}; + ailego::Params params_{}; + Stats stats_; + uint32_t ef_{HnswEntity::kDefaultEf}; + uint32_t max_scan_num_{0U}; + uint32_t bruteforce_threshold_{HnswEntity::kDefaultBruteForceThreshold}; + float max_scan_ratio_{HnswEntity::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_probility_{HnswEntity::kDefaultBFNegativeProbility}; + uint32_t magic_{0U}; + + State state_{STATE_INIT}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw/hnsw_searcher_entity.cc b/src/core/algorithm/hnsw/hnsw_searcher_entity.cc new file mode 100644 index 0000000..ca092f2 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_searcher_entity.cc @@ -0,0 +1,515 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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_searcher_entity.h" +#include +#include "utility/sparse_utility.h" + +namespace zvec { +namespace core { + +HnswSearcherEntity::HnswSearcherEntity() {} + +int HnswSearcherEntity::cleanup(void) { + storage_.reset(); + vectors_.reset(); + keys_.reset(); + neighbors_.reset(); + neighbors_meta_.reset(); + neighbors_in_memory_enabled_ = false; + loaded_ = false; + + this->HnswEntity::cleanup(); + + return 0; +} + +key_t HnswSearcherEntity::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 HnswSearcherEntity::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; +} + +const void *HnswSearcherEntity::get_vector_by_key(key_t key) const { + node_id_t local_id = get_id(key); + if (ailego_unlikely(local_id == kInvalidNodeId)) { + return nullptr; + } + + return get_vector(local_id); +} + +const void *HnswSearcherEntity::get_vector(node_id_t id) const { + size_t read_size = vector_size(); + size_t offset = node_size() * id; + + const void *vec; + if (ailego_unlikely(vectors_->read(offset, &vec, read_size) != read_size)) { + LOG_ERROR("Read vector from segment failed"); + return nullptr; + } + return vec; +} + +int HnswSearcherEntity::get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const { + const void *vec = get_vector(id); + block.reset((void *)vec); + return 0; +} + +const void *HnswSearcherEntity::get_vectors() const { + const void *vec; + size_t len = node_size() * doc_cnt(); + if (vectors_->read(0, &vec, len) != len) { + LOG_ERROR("Read vectors from segment failed"); + return nullptr; + } + return vec; +} + +int HnswSearcherEntity::get_vector(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 = vector_size(); + + for (uint32_t i = 0; i < count; ++i) { + segment_datas_[i].offset = node_size() * ids[i]; + segment_datas_[i].length = read_size; + + ailego_assert_with(segment_datas_[i].offset < vectors_->data_size(), + "invalid offset"); + } + if (ailego_unlikely(!vectors_->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 HnswSearcherEntity::get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + const void *vecs[count]; + get_vector(ids, count, vecs); + for (uint32_t i = 0; i < count; ++i) { + vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i])); + } + return 0; +} + +const Neighbors HnswSearcherEntity::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 GraphNeighborMeta *m; + if (ailego_unlikely(neighbors_meta_->read(id * sizeof(GraphNeighborMeta), + (const void **)(&m), + sizeof(GraphNeighborMeta)) != + sizeof(GraphNeighborMeta))) { + 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 HnswNeighborMeta *m; + if (ailego_unlikely(upper_neighbors_meta_->read(id * sizeof(HnswNeighborMeta), + (const void **)(&m), + sizeof(HnswNeighborMeta)) != + sizeof(HnswNeighborMeta))) { + 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 HnswSearcherEntity::load(const IndexStorage::Pointer &container, + bool check_crc) { + storage_ = 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 " + "neighborCnt=%zu upperNeighborCnt=%zu scalingFactor=%zu " + "vectorSize=%zu nodeSize=%zu vectorSegmentSize=%zu keySegmentSize=%zu " + "neighborsSegmentSize=%zu neighborsMetaSegmentSize=%zu ", + doc_cnt(), entry_point(), cur_max_level(), ef_construction(), + neighbor_cnt(), upper_neighbor_cnt(), scaling_factor(), vector_size(), + node_size(), vectors_->data_size(), keys_->data_size(), + neighbors_ == nullptr ? 0 : neighbors_->data_size(), + neighbors_meta_ == nullptr ? 0 : neighbors_meta_->data_size()); + + return 0; +} + +int HnswSearcherEntity::load_segments(bool check_crc) { + //! load header + const void *data = nullptr; + HNSWHeader hd; + auto graph_hd_segment = storage_->get(kGraphHeaderSegmentId); + if (!graph_hd_segment || graph_hd_segment->data_size() < sizeof(hd.graph)) { + LOG_ERROR("Miss or invalid segment %s", kGraphHeaderSegmentId.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", kGraphHeaderSegmentId.c_str()); + return IndexError_ReadData; + } + memcpy(&hd.graph, data, sizeof(hd.graph)); + + auto hnsw_hd_segment = storage_->get(kHnswHeaderSegmentId); + if (!hnsw_hd_segment || hnsw_hd_segment->data_size() < sizeof(hd.hnsw)) { + LOG_ERROR("Miss or invalid segment %s", kHnswHeaderSegmentId.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", kHnswHeaderSegmentId.c_str()); + return IndexError_ReadData; + } + memcpy(&hd.hnsw, data, sizeof(hd.hnsw)); + *mutable_header() = hd; + segment_datas_.resize(std::max(neighbor_cnt(), upper_neighbor_cnt())); + + vectors_ = storage_->get(kGraphFeaturesSegmentId); + if (!vectors_) { + LOG_ERROR("IndexStorage get segment %s failed", + kGraphFeaturesSegmentId.c_str()); + return IndexError_InvalidFormat; + } + keys_ = storage_->get(kGraphKeysSegmentId); + if (!keys_) { + LOG_ERROR("IndexStorage get segment %s failed", + kGraphKeysSegmentId.c_str()); + return IndexError_InvalidFormat; + } + + neighbors_ = storage_->get(kGraphNeighborsSegmentId); + if (!neighbors_ || (neighbors_->data_size() == 0 && doc_cnt() > 1)) { + LOG_ERROR("IndexStorage get segment %s failed or empty", + kGraphNeighborsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + neighbors_meta_ = storage_->get(kGraphOffsetsSegmentId); + if (!neighbors_meta_ || + neighbors_meta_->data_size() < sizeof(GraphNeighborMeta) * doc_cnt()) { + LOG_ERROR("IndexStorage get segment %s failed or invalid size", + kGraphOffsetsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + upper_neighbors_ = storage_->get(kHnswNeighborsSegmentId); + if (!upper_neighbors_ || + (upper_neighbors_->data_size() == 0 && cur_max_level() > 0)) { + LOG_ERROR("IndexStorage get segment %s failed or empty", + kHnswNeighborsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + upper_neighbors_meta_ = storage_->get(kHnswOffsetsSegmentId); + if (!upper_neighbors_meta_ || upper_neighbors_meta_->data_size() < + sizeof(HnswNeighborMeta) * doc_cnt()) { + LOG_ERROR("IndexStorage get segment %s failed or invalid size", + kHnswOffsetsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + mapping_ = storage_->get(kGraphMappingSegmentId); + if (!mapping_ || mapping_->data_size() < sizeof(node_id_t) * doc_cnt()) { + LOG_ERROR("IndexStorage get segment %s failed or invalid size", + kGraphMappingSegmentId.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(vectors_); + segments.emplace_back(keys_); + + 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 HnswSearcherEntity::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 = storage_->get(kGraphOffsetsSegmentId); + if (!neighbors_meta) { + LOG_ERROR("IndexStorage get segment graph.offsets failed"); + return IndexError_InvalidArgument; + } + + const GraphNeighborMeta *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", kGraphOffsetsSegmentId.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 HnswSearcherEntity::get_fixed_neighbors( + std::vector *fixed_neighbors) const { + //! Get a new segemnt to release the buffer after loading neighbors + auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId); + if (!neighbors_meta) { + LOG_ERROR("IndexStorage get segment graph.offsets failed"); + return IndexError_InvalidArgument; + } + + const GraphNeighborMeta *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", kGraphOffsetsSegmentId.c_str()); + return IndexError_InvalidArgument; + } + + size_t fixed_neighbor_cnt = 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 HnswSearcherEntity::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 HnswEntity::Pointer HnswSearcherEntity::clone() const { + auto vectors = vectors_->clone(); + if (ailego_unlikely(!vectors)) { + LOG_ERROR("clone segment %s failed", kGraphFeaturesSegmentId.c_str()); + return HnswEntity::Pointer(); + } + auto keys = keys_->clone(); + if (ailego_unlikely(!keys)) { + LOG_ERROR("clone segment %s failed", kGraphKeysSegmentId.c_str()); + return HnswEntity::Pointer(); + } + + auto mapping = mapping_->clone(); + if (ailego_unlikely(!mapping)) { + LOG_ERROR("clone segment %s failed", kGraphMappingSegmentId.c_str()); + return HnswEntity::Pointer(); + } + + auto neighbors = neighbors_->clone(); + if (ailego_unlikely(!neighbors)) { + LOG_ERROR("clone segment %s failed", kGraphNeighborsSegmentId.c_str()); + return HnswEntity::Pointer(); + } + auto upper_neighbors = upper_neighbors_->clone(); + if (ailego_unlikely(!neighbors)) { + LOG_ERROR("clone segment %s failed", kHnswNeighborsSegmentId.c_str()); + return HnswEntity::Pointer(); + } + auto neighbors_meta = neighbors_meta_->clone(); + if (ailego_unlikely(!neighbors_meta)) { + LOG_ERROR("clone segment %s failed", kGraphOffsetsSegmentId.c_str()); + return HnswEntity::Pointer(); + } + auto upper_neighbors_meta = upper_neighbors_meta_->clone(); + if (ailego_unlikely(!upper_neighbors_meta)) { + LOG_ERROR("clone segment %s failed", kHnswOffsetsSegmentId.c_str()); + return HnswEntity::Pointer(); + } + + SegmentGroupParam neighbor_group{neighbors, neighbors_meta, upper_neighbors, + upper_neighbors_meta}; + + HnswSearcherEntity *entity = new (std::nothrow) + HnswSearcherEntity(header(), vectors, keys, mapping, neighbor_group, + fixed_neighbors_, neighbors_in_memory_enabled_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswSearcherEntity new failed"); + } + + return HnswEntity::Pointer(entity); +} + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw/hnsw_searcher_entity.h b/src/core/algorithm/hnsw/hnsw_searcher_entity.h new file mode 100644 index 0000000..dc84ccc --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_searcher_entity.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 "hnsw_builder_entity.h" +#include "hnsw_entity.h" + +namespace zvec { +namespace core { + +class HnswSearcherEntity : public HnswEntity { + 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 + HnswSearcherEntity(); + + //! 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 override; + + //! Get primary key of the node id + virtual 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 vector feature data by key + virtual const void *get_vector_by_key(key_t key) const override; + + //! Get vector feature data by id + virtual const void *get_vector(node_id_t id) const override; + + //! Get vector feature data by id + virtual int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + + virtual int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + virtual int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const override; + + //! Get all vectors + const void *get_vectors() const; + + //! Get the node id's neighbors on graph level + virtual const Neighbors get_neighbors(level_t level, + node_id_t id) const override; + + virtual int load(const IndexStorage::Pointer &container, + bool check_crc) override; + + int load_segments(bool check_crc); + + virtual 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 + HnswSearcherEntity(const HNSWHeader &hd, const SegmentPointer &vectors, + const SegmentPointer &keys, const SegmentPointer &mapping, + const SegmentGroupParam &neighbor_group, + const std::shared_ptr &fixed_neighbors, + bool neighbors_in_memory_enabled) + : HnswEntity(hd), + vectors_(vectors), + 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), + neighbors_in_memory_enabled_(neighbors_in_memory_enabled) { + segment_datas_.resize(std::max(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) + 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: + HnswSearcherEntity(const HnswSearcherEntity &) = delete; + HnswSearcherEntity &operator=(const HnswSearcherEntity &) = delete; + + private: + IndexStorage::Pointer storage_{}; + + SegmentPointer vectors_{}; + SegmentPointer keys_{}; + SegmentPointer mapping_{}; + + SegmentPointer neighbors_{}; + SegmentPointer neighbors_meta_{}; + SegmentPointer upper_neighbors_{}; + SegmentPointer upper_neighbors_meta_{}; + + 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/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc new file mode 100644 index 0000000..b14ad05 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -0,0 +1,848 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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() : entity_(stats_) {} + +HnswStreamer::~HnswStreamer() { + if (state_ == STATE_INITED) { + 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, &neighbor_cnt_); + float ratio = HnswEntity::kDefaultUpperNeighborRatio; + params.get(PARAM_HNSW_STREAMER_UPPER_NEIGHBOR_RATIO, &ratio); + upper_neighbor_cnt_ = ratio * neighbor_cnt_; + + params.get(PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_); + params.get(PARAM_HNSW_STREAMER_EF, &ef_); + 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_SCALING_FACTOR, &scaling_factor_); + 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_); + ratio = HnswEntity::kDefaultNeighborPruneRatio; + params.get(PARAM_HNSW_STREAMER_NEIGHBOR_PRUNE_RATIO, &ratio); + size_t prune_cnt = ratio * neighbor_cnt_; + 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_); + entity_.set_use_key_info_map(use_id_map_); + + 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 (neighbor_cnt_ == 0U) { + neighbor_cnt_ = HnswEntity::kDefaultNeighborCnt; + } + if (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 (upper_neighbor_cnt_ == 0U) { + upper_neighbor_cnt_ = HnswEntity::kDefaultUpperNeighborCnt; + } + if (upper_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) { + LOG_ERROR("UpperNeighborCnt must be in range (0,%d)", + HnswEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (min_neighbor_cnt_ > 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(), 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_neighbor_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; + } + + entity_.set_ef_construction(ef_construction_); + entity_.set_neighbor_cnt(neighbor_cnt_); + entity_.set_upper_neighbor_cnt(upper_neighbor_cnt_); + entity_.set_scaling_factor(scaling_factor_); + entity_.set_prune_cnt(prune_cnt); + + entity_.set_vector_size(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; + } + + LOG_DEBUG( + "Init params: maxIndexSize=%zu docsHardLimit=%zu docsSoftLimit=%zu " + "efConstruction=%u ef=%u neighborCnt=%u upperNeighborCnt=%u " + "scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu " + "bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbility=%.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_, neighbor_cnt_, upper_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_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_)); + + ret = alg_->init(); + if (ret != 0) { + return ret; + } + + 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(); + entity_.cleanup(); + + if (alg_) { + alg_->cleanup(); + } + + max_index_size_ = 0UL; + docs_hard_limit_ = HnswEntity::kDefaultDocsHardLimit; + docs_soft_limit_ = 0UL; + neighbor_cnt_ = HnswEntity::kDefaultNeighborCnt; + upper_neighbor_cnt_ = HnswEntity::kDefaultUpperNeighborCnt; + 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::kDefaultBFNegativeProbility; + max_scan_ratio_ = HnswEntity::kDefaultScanRatio; + state_ = STATE_INIT; + check_crc_enabled_ = false; + filter_same_key_ = false; + get_vector_enabled_ = false; + + return 0; +} + +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; + } + int 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); + } + + 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_->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(); + + 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(); }); + + meta_.set_searcher("HnswSearcher", HnswEntity::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 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_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_probility(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 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 +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 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; +} + + +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().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().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().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..e81c6f2 --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_streamer.h @@ -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. +#pragma once + +#include +#include "framework/index_framework.h" +#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); + virtual ~HnswStreamer(void); + + HnswStreamer(const HnswStreamer &streamer) = delete; + HnswStreamer &operator=(const HnswStreamer &streamer) = delete; + + protected: + //! Initialize Streamer + virtual int init(const IndexMeta &imeta, + const ailego::Params ¶ms) override; + + //! Cleanup Streamer + virtual int cleanup(void) override; + + //! Create a context + virtual Context::Pointer create_context(void) const override; + + //! Create a new iterator + virtual IndexProvider::Pointer create_provider(void) const override; + + //! Add a vector into index + virtual int add_impl(uint64_t pkey, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Add a vector with id into index + virtual int add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + //! Similarity search + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const override; + + //! Similarity brute force search + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const override; + + //! Linear search by primary keys + virtual 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 + virtual 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 + virtual const void *get_vector(uint64_t key) const override { + return entity_.get_vector_by_key(key); + } + + virtual int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_.get_vector_by_key(key, block); + } + + //! Fetch vector by id + virtual const void *get_vector_by_id(uint32_t id) const override { + return entity_.get_vector(id); + } + + virtual 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 + virtual int open(IndexStorage::Pointer stg) override; + + //! Close file + virtual int close(void) override; + + //! flush file + virtual int flush(uint64_t checkpoint) override; + + //! Dump index into storage + virtual int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + virtual const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta of index + virtual const IndexMeta &meta(void) const override { + return meta_; + } + + virtual 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(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(); + } + }; + + HnswStreamerEntity entity_; + HnswAlgorithm::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 neighbor_cnt_{HnswEntity::kDefaultNeighborCnt}; + uint32_t min_neighbor_cnt_{0u}; + uint32_t upper_neighbor_cnt_{HnswEntity::kDefaultUpperNeighborCnt}; + uint32_t ef_{HnswEntity::kDefaultEf}; + 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::kDefaultBFNegativeProbility}; + 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}; + + //! 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..510087f --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc @@ -0,0 +1,701 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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" +#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(); + 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(); + upper_neighbor_chunks_.clear(); + 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) { + char buffer[neighbor_size_]; + NeighborsHeader *hd = reinterpret_cast(buffer); + 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 HnswStreamerEntity::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(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; + 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 *HnswStreamerEntity::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=%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"); + 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"); + 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"); + 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); + 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 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)); + 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; + } + } + + upper_neighbor_chunks_.resize( + broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR)); + 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; + } + } + + 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 = 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(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, neighborCnt=%zu upperneighborCnt=%zu " + "efConstruction=%zu curDocCnt=%u totalVecs=%u maxLevel=%u", + 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(); + upper_neighbor_chunks_.clear(); + + 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->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 (neighbor_cnt() != hd->neighbor_cnt() || + upper_neighbor_cnt() != hd->upper_neighbor_cnt()) { + LOG_ERROR("Param neighbors:%zu:%zu mismatch index previous %zu:%zu", + neighbor_cnt(), upper_neighbor_cnt(), hd->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; + 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(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 = -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(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_, 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("HnswStreamerEntity new failed"); + } + return HnswEntity::Pointer(entity); +} + +} // 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..cd515ce --- /dev/null +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.h @@ -0,0 +1,525 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_framework.h" +#include "hnsw_chunk.h" +#include "hnsw_entity.h" +#include "hnsw_index_hash.h" +#include "hnsw_params.h" + +namespace zvec { +namespace core { + +//! HnswStreamerEntity manage vector data, pkey, and node's neighbors +class HnswStreamerEntity : public HnswEntity { + public: + //! Cleanup + //! return 0 on success, or errCode in failure + virtual int cleanup() override; + + //! Make a copy of streamer entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + virtual const HnswEntity::Pointer clone() const override; + + //! Get primary key of the node id + virtual key_t get_key(node_id_t id) const override; + + //! Get vector feature data by key + virtual const void *get_vector(node_id_t id) const override; + + //! Get vectors feature data by local ids + virtual int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + + virtual int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + virtual 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 + virtual 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 + virtual int add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) override; + + //! Add vector and id to hnsw entity + virtual int add_vector_with_id(level_t level, node_id_t id, + const void *vec) override; + + virtual 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 + virtual void add_neighbor(level_t level, node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + + //! Dump index by dumper + virtual int dump(const IndexDumper::Pointer &dumper) override; + + virtual 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 + HnswStreamerEntity(IndexStreamer::Stats &stats); + + //! Destructor + ~HnswStreamerEntity(); + + //! Get vector feature data by key + virtual const void *get_vector_by_key(key_t key) const override { + auto id = get_id(key); + return id == kInvalidNodeId ? nullptr : get_vector(id); + } + + virtual 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 neighbors size + inline size_t neighbors_size() const { + return sizeof(NeighborsHeader) + neighbor_cnt() * sizeof(node_id_t); + } + + //! Get neighbors size in hybrid mode + inline size_t neighbors_size_hybrid() const { + return sizeof(NeighborsHeader) * 2 + 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); + } + + //! Get upper neighbors size in hybrid mode + inline size_t upper_neighbors_size_hybrid() const { + return 2 * + (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 + }; + 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 + 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, + 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) + : 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_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(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 { + 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->index) >> upper_neighbor_mask_bits_; + uint32_t offset = (((meta->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; + } + 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(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.level = level; + meta.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("Chunk 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, 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; + } + + private: + HnswStreamerEntity(const HnswStreamerEntity &) = delete; + HnswStreamerEntity &operator=(const HnswStreamerEntity &) = 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}; + + 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_{}; + + ChunkBroker::Pointer broker_{}; // chunk broker +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_sparse/CMakeLists.txt b/src/core/algorithm/hnsw_sparse/CMakeLists.txt new file mode 100644 index 0000000..fe26d10 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/CMakeLists.txt @@ -0,0 +1,11 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +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 + 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..6dd47a1 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.cc @@ -0,0 +1,544 @@ + +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 + +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; +} + +int HnswSparseAlgorithm::fast_search(HnswSparseContext *ctx) const { + auto max_level = entity_.cur_max_level(); + auto entry_point = entity_.entry_point(); + + dist_t dist = ctx->dist_calculator().dist(entry_point); + for (level_t cur_level = max_level; 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())++; + } + + node_id_t 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, 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())++; + } + + node_id_t 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, 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..fa5fab5 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.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 "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 with lock + //! return 0 on success, or errCode in failure. results saved in ctx + int search(HnswSparseContext *ctx) const; + + //! do knn search in graph without lock + //! return 0 on success, or errCode in failure. results saved in ctx + int fast_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..4516e41 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc @@ -0,0 +1,521 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_error.h" +#include "framework/index_factory.h" +#include "framework/index_logger.h" +#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_MAX_NEIGHBOR_COUNT, &neighbor_cnt_); + params.get(PARAM_HNSW_SPARSE_BUILDER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_); + float ratio = HnswSparseEntity::kDefaultNeighborPruneRatio; + params.get(PARAM_HNSW_SPARSE_BUILDER_UPPER_NEIGHBOR_RATIO, &ratio); + upper_neighbor_cnt_ = ratio * neighbor_cnt_; + params.get(PARAM_HNSW_SPARSE_BUILDER_EFCONSTRUCTION, &ef_construction_); + params.get(PARAM_HNSW_SPARSE_BUILDER_SCALING_FACTOR, &scaling_factor_); + params.get(PARAM_HNSW_SPARSE_BUILDER_CHECK_INTERVAL_SECS, + &check_interval_secs_); + ratio = HnswSparseEntity::kDefaultNeighborPruneRatio; + params.get(PARAM_HNSW_SPARSE_BUILDER_NEIGHBOR_PRUNE_RATIO, &ratio); + + size_t prune_cnt = neighbor_cnt_ * ratio; + + if (ef_construction_ == 0) { + ef_construction_ = HnswSparseEntity::kDefaultEfConstruction; + } + if (neighbor_cnt_ == 0) { + neighbor_cnt_ = HnswSparseEntity::kDefaultNeighborCnt; + } + if (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_ > 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(), neighbor_cnt_); + return IndexError_InvalidArgument; + } + if (upper_neighbor_cnt_ == 0) { + neighbor_cnt_ = HnswSparseEntity::kDefaultUpperNeighborCnt; + } + if (upper_neighbor_cnt_ > HnswSparseEntity::kMaxNeighborCnt) { + LOG_ERROR("UpperNeighborCnt 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_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_neighbor_cnt(neighbor_cnt_); + entity_.set_min_neighbor_cnt(min_neighbor_cnt_); + entity_.set_upper_neighbor_cnt(upper_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 " + "neighborCnt=%u upperNeighborCnt=%u scalingFactor=%u " + "memoryQuota=%zu neighborPruneCnt=%zu measureName=%s ", + ef_construction_, neighbor_cnt_, upper_neighbor_cnt_, scaling_factor_, + memory_quota, prune_cnt, meta_.metric_name().c_str()); + + return 0; +} + +int HnswSparseBuilder::cleanup(void) { + LOG_INFO("Begin HnswSparseBuilder::cleanup"); + + neighbor_cnt_ = HnswSparseEntity::kDefaultNeighborCnt; + min_neighbor_cnt_ = 0; + upper_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperNeighborCnt; + 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)); + } + + while (!task_group->is_finished()) { + std::unique_lock lk(mutex_); + 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..7e3ba4a --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_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 "framework/index_builder.h" +#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 neighbor_cnt_{HnswSparseEntity::kDefaultNeighborCnt}; + uint32_t min_neighbor_cnt_{0}; + uint32_t upper_neighbor_cnt_{HnswSparseEntity::kDefaultUpperNeighborCnt}; + 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..a27accf --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.cc @@ -0,0 +1,289 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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("Add vector failed, dim size too larg, dim_size=%u, key=%zu", + sparse_count, (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 { + const void *vecs[count]; + get_vector_metas(ids, count, vecs); + 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..aaf97aa --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.h @@ -0,0 +1,161 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 + 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) override; + + //! Get primary key of the node id + virtual key_t get_key(node_id_t id) const override; + + //! Get vector feature data by key + virtual const void *get_vector_meta(node_id_t id) const override; + + virtual int get_vector_meta(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Batch get vectors feature data by keys + virtual int get_vector_metas(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + virtual 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 + virtual const Neighbors get_neighbors(level_t level, + node_id_t id) const override; + + //! Replace node id in level's neighbors + virtual int update_neighbors( + level_t level, node_id_t id, + const std::vector> &neighbors) override; + + //! add a neighbor to id in graph level + virtual 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 + virtual const void *get_sparse_data(uint64_t offset, + uint32_t len) const override; + //! Get sparse data from id + virtual const void *get_sparse_data(node_id_t id) const override; + + virtual int get_sparse_data(uint64_t offset, uint32_t len, + IndexStorage::MemoryBlock &block) const override; + + virtual int get_sparse_data(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get sparse data from vector + virtual std::pair get_sparse_data_from_vector( + const void *vec) const override; + + virtual int get_sparse_data_from_vector(const void *vec, + IndexStorage::MemoryBlock &block, + int &sparse_length) const override; + + //! Dump the hnsw graph to dumper + virtual int dump(const IndexDumper::Pointer &dumper) override; + + //! Cleanup the entity + virtual 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) + 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..eb0e74b --- /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 "framework/index_error.h" +#include "framework/index_helper.h" +#include "framework/index_logger.h" +#include "framework/index_streamer.h" + +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(&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..b5a0086 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.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 +#include +#include +#include "framework/index_error.h" +#include "framework/index_logger.h" +#include "framework/index_storage.h" +#include "framework/index_streamer.h" + +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(this, 0, sizeof(HnswSparseChunkMeta)); + } + void clear() { + memset(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..cad3709 --- /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_probility_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + update_heap_.limit(entity_->neighbor_cnt() + 1); + break; + + case kSparseSearcherContext: + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_, + negative_probility_); + 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_probility_); + if (ret != 0) { + LOG_ERROR("Create filter failed, mode %d", filter_mode_); + return ret; + } + + update_heap_.limit(entity_->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; + 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_probility_; + 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_probility_) > 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_probility_); + 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(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->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..4f31a6d --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h @@ -0,0 +1,522 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_context.h" +#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 + virtual ~HnswSparseContext(); + + public: + //! Set topk of search result + virtual void set_topk(uint32_t val) override { + topk_ = val; + topk_heap_.limit(std::max(val, ef_)); + } + + //! 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 uint32_t magic(void) const override { + return 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 this->debugging(); + } + + //! Retrieve string of debug + virtual 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 + virtual 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_) { + group_results_[idx][i].mutable_docs()->emplace_back( + entity_->get_key(id), score, id, entity_->get_vector_meta(id)); + } 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_probility(float v) { + negative_probility_ = 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; + } + + virtual void set_bruteforce_threshold(uint32_t v) override { + bruteforce_threshold_ = v; + } + + inline uint32_t get_bruteforce_threshold() const { + return bruteforce_threshold_; + } + + virtual void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + + virtual 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_probility_{HnswSparseEntity::kDefaultBFNegativeProbility}; + 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..4a6782b --- /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 "framework/index_meta.h" +#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()(id_t i) { + return dist(i); + } + + dist_t operator()(id_t lhs, 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..798c55c --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.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 "hnsw_sparse_entity.h" + +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; + node_id_t mapping[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() <= 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; + } + 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; + node_id_t 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"); + memset(buffer, 0, sizeof(buffer)); + 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, sizeof(buffer)) != sizeof(buffer)) { + LOG_ERROR("Dump graph neighbor id=%u failed, size %lu", id, + sizeof(buffer)); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(buffer, sizeof(buffer), crc); + offset += sizeof(buffer); + } + } + 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..8885fa4 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h @@ -0,0 +1,705 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "ailego/logger/logger.h" +#include "framework/index_dumper.h" +#include "framework/index_error.h" +#include "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; +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 max_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; + uint16_t sparse_max_neighbor_count; + uint16_t sparse_min_neighbor_count; + uint32_t total_sparse_count; + // uint8_t reserved_[4]; +}; + +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 thumb_neighbor_count; + uint32_t ef_construction; + uint32_t scaling_factor; + uint32_t max_level; + uint32_t entry_point; + uint32_t options; + uint16_t thumb_sparse_neighbor_count; + 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(this, &header, sizeof(header)); + } + + HNSWSparseHeader &operator=(const HNSWSparseHeader &header) { + memcpy(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(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.max_neighbor_count; + } + + size_t upper_neighbor_cnt() const { + return hnsw.thumb_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; + } + + size_t sparse_neighbor_cnt() const { + return graph.sparse_max_neighbor_count; + } + + size_t sparse_min_neighbor_cnt() const { + return graph.sparse_min_neighbor_count; + } + + size_t upper_sparse_neighbor_cnt() const { + return hnsw.thumb_sparse_neighbor_count; + } + + SparseGraphHeader graph; + HnswSparseHeader 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 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.max_neighbor_count + : header_.hnsw.thumb_neighbor_count; + } + + //! get max neighbor size of graph level 0 + inline size_t neighbor_cnt() const { + return header_.graph.max_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.thumb_neighbor_count; + } + + size_t sparse_neighbor_cnt() const { + return header_.graph.sparse_max_neighbor_count; + } + + size_t sparse_min_neighbor_cnt() const { + return header_.graph.sparse_min_neighbor_count; + } + + size_t upper_sparse_neighbor_cnt() const { + return header_.hnsw.thumb_sparse_neighbor_count; + } + + size_t dense_neighbor_cnt() const { + return header_.graph.max_neighbor_count - + header_.graph.sparse_max_neighbor_count; + } + + size_t dense_min_neighbor_cnt() const { + return header_.graph.min_neighbor_count - + header_.graph.sparse_min_neighbor_count; + } + + size_t upper_dense_neighbor_cnt() const { + return header_.hnsw.thumb_neighbor_count - + header_.hnsw.thumb_sparse_neighbor_count; + } + + size_t sparse_neighbor_offset() const { + return sizeof(NeighborsHeader) + dense_neighbor_cnt() * sizeof(node_id_t); + } + + size_t upper_sparse_neighbor_offset() const { + return sizeof(NeighborsHeader) + + upper_dense_neighbor_cnt() * sizeof(node_id_t); + } + + size_t dense_neighbor_size() const { + return sizeof(NeighborsHeader) + dense_neighbor_cnt() * sizeof(node_id_t); + } + + size_t sparse_neighbor_size() const { + return sizeof(NeighborsHeader) + sparse_neighbor_cnt() * sizeof(node_id_t); + } + + size_t upper_dense_neighbor_size() const { + return sizeof(NeighborsHeader) + + upper_dense_neighbor_cnt() * sizeof(node_id_t); + } + + size_t upper_sparse_neighbor_size() const { + return sizeof(NeighborsHeader) + + upper_sparse_neighbor_cnt() * sizeof(node_id_t); + } + + //! 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_neighbor_cnt(size_t cnt) { + header_.graph.max_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.thumb_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; + } + + void set_sparse_neighbor_cnt(size_t cnt) { + header_.graph.sparse_max_neighbor_count = cnt; + } + + void set_sparse_min_neighbor_cnt(size_t cnt) { + header_.graph.sparse_min_neighbor_count = cnt; + } + + + void set_upper_sparse_neighbor_cnt(size_t cnt) { + header_.hnsw.thumb_sparse_neighbor_count = cnt; + } + + 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 kDefaultNeighborCnt = 100; + constexpr static uint32_t kDefaultUpperNeighborCnt = 50; + 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 kDefaultBFNegativeProbility = 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 kDefaultNeighborPruneRatio = + 0.5f; // prune count / neighbor count + constexpr static float kDefaultUpperNeighborRatio = + 0.5f; // upper neighbor count / neighbor count + + constexpr static float kHybridNeighborRatio = 0.5f; + + constexpr static uint32_t kSparseMetaSize = 2u * sizeof(uint64_t); + constexpr static float kDefaultSparseNeighborRatio = 0.5f; + constexpr static uint32_t kSparseMaxDimSize = 4096; + 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..56731b7 --- /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 + virtual uint32_t sparse_count() const override { + return sparse_count_; + } + + //! Retrieve sparse indices + virtual const uint32_t *sparse_indices() const override { + return reinterpret_cast(sparse_indices_buffer_.data()); + } + + //! Retrieve sparse data + virtual const void *sparse_data() const override { + return reinterpret_cast(sparse_data_buffer_.data()); + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt() && valid_; + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + return entity_->get_key(cur_id_); + } + + //! Next iterator + virtual 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..b880ce9 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_params.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 + +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_RATIO( + "proxima.hnsw.sparse_builder.neighbor_prune_ratio"); +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_MIN_NEIGHBOR_COUNT( + "proxima.hnsw.sparse_builder.min_neighbor_count"); +static const std::string PARAM_HNSW_SPARSE_BUILDER_UPPER_NEIGHBOR_RATIO( + "proxima.hnsw.sparse_builder.upper_neighbor_ratio"); + +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_UPPER_NEIGHBOR_RATIO( + "proxima.hnsw.sparse_streamer.upper_neighbor_ratio"); +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_RATIO( + "proxima.hnsw.sparse_streamer.neighbor_prune_ratio"); +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..9f1ee66 --- /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_probility_); + 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_probility_ <= 0.0f || bf_negative_probility_ >= 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_probility_, + 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_probility_ = HnswSparseEntity::kDefaultBFNegativeProbility; + 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_->fast_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_probility(bf_negative_probility_); + 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..dec4090 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.h @@ -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. +#pragma once + +#include "framework/index_framework.h" +#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); + virtual ~HnswSparseSearcher(void); + + 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_probility_{HnswSparseEntity::kDefaultBFNegativeProbility}; + + 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..ddb5155 --- /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 { + const void *vecs[count]; + get_vector_metas(ids, count, vecs); + 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 " + "neighborCnt=%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(), + 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(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 = 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..0ddeca5 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.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 "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 + virtual const HnswSparseEntity::Pointer clone() const override; + + //! Get primary key of the node id + virtual 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 + virtual 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 + virtual const void *get_vector_meta(node_id_t id) const override; + + virtual int get_vector_meta(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get vector feature data by id + virtual int get_vector_metas(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + + virtual 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 + virtual const void *get_sparse_data(uint64_t offset, + uint32_t len) const override; + + //! Get sparse data from id + virtual const void *get_sparse_data(node_id_t id) const override; + + virtual int get_sparse_data(uint64_t offset, uint32_t len, + IndexStorage::MemoryBlock &block) const override; + + virtual int get_sparse_data(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get sparse data from vector + virtual std::pair get_sparse_data_from_vector( + const void *vec) const override; + + virtual 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 + virtual const Neighbors get_neighbors(level_t level, + node_id_t id) const override; + + virtual int load(const IndexStorage::Pointer &container, + bool check_crc) override; + + int load_segments(bool check_crc); + + virtual 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(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) + 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..adf2c11 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.cc @@ -0,0 +1,952 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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) { + 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, &neighbor_cnt_); + float ratio = HnswSparseEntity::kDefaultUpperNeighborRatio; + params.get(PARAM_HNSW_SPARSE_STREAMER_UPPER_NEIGHBOR_RATIO, &ratio); + upper_neighbor_cnt_ = ratio * neighbor_cnt_; + + 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_SCALING_FACTOR, &scaling_factor_); + 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_); + ratio = HnswSparseEntity::kDefaultNeighborPruneRatio; + params.get(PARAM_HNSW_SPARSE_STREAMER_NEIGHBOR_PRUNE_RATIO, &ratio); + size_t prune_cnt = ratio * neighbor_cnt_; + 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 (neighbor_cnt_ == 0U) { + neighbor_cnt_ = HnswSparseEntity::kDefaultNeighborCnt; + } + if (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 (upper_neighbor_cnt_ == 0U) { + upper_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperNeighborCnt; + } + if (upper_neighbor_cnt_ > HnswSparseEntity::kMaxNeighborCnt) { + LOG_ERROR("UpperNeighborCnt must be in range (0,%d)", + HnswSparseEntity::kMaxNeighborCnt); + return IndexError_InvalidArgument; + } + if (min_neighbor_cnt_ > 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(), + 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_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_neighbor_cnt(neighbor_cnt_); + entity_.set_upper_neighbor_cnt(upper_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 neighborCnt=%u upperNeighborCnt=%u " + "scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu " + "bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbility=%.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_, neighbor_cnt_, upper_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; + neighbor_cnt_ = HnswSparseEntity::kDefaultNeighborCnt; + upper_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperNeighborCnt; + 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::kDefaultBFNegativeProbility; + 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_probility(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("Add vector failed, dim size too larg, dim_size=%u, id=%u", + sparse_count, 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("Add vector failed, dim size too larg, dim_size=%u, key=%zu", + sparse_count, (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..20c2788 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.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 "framework/index_framework.h" +#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); + virtual ~HnswSparseStreamer(void); + + 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(); + } + }; + + HnswSparseStreamerEntity entity_; + HnswSparseAlgorithm::UPointer alg_; + IndexMeta meta_{}; + IndexMetric::Pointer metric_{}; + + IndexMetric::MatrixSparseDistance add_distance_{}; + IndexMetric::MatrixSparseDistance search_distance_{}; + Stats stats_{}; + 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 neighbor_cnt_{HnswSparseEntity::kDefaultNeighborCnt}; + uint32_t min_neighbor_cnt_{0u}; + uint32_t upper_neighbor_cnt_{HnswSparseEntity::kDefaultUpperNeighborCnt}; + 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::kDefaultBFNegativeProbility}; + 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..8299388 --- /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) { + char buffer[neighbor_size_]; + NeighborsHeader *hd = reinterpret_cast(buffer); + 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 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, neighborCnt=%zu upperneighborCnt=%zu " + "efConstruction=%zu curDocCnt=%u totalVecs=%u maxLevel=%u", + 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->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 (neighbor_cnt() != hd->neighbor_cnt() || + upper_neighbor_cnt() != hd->upper_neighbor_cnt()) { + LOG_ERROR("Param neighbors:%zu:%zu mismatch index previous %zu:%zu", + 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 = -1UL; + size_t sparse_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(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 = -1UL; + size_t sparse_chunk_offset = -1UL; + + // 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..8351d08 --- /dev/null +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.h @@ -0,0 +1,542 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_framework.h" +#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 + virtual int cleanup() override; + + //! Make a copy of streamer entity, to support thread-safe operation. + //! The segment in container cannot be read concurrenly + virtual const HnswSparseEntity::Pointer clone() const override; + + //! Get primary key of the node id + virtual key_t get_key(node_id_t id) const override; + + //! Get vector feature data by key + virtual const void *get_vector_meta(node_id_t id) const override; + + virtual int get_vector_meta(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get vectors feature data by local ids + virtual int get_vector_metas(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + virtual 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 + virtual const void *get_sparse_data(uint64_t offset, + uint32_t len) const override; + + //! Get sparse data from id + virtual const void *get_sparse_data(node_id_t id) const override; + + virtual int get_sparse_data(uint64_t offset, uint32_t len, + IndexStorage::MemoryBlock &block) const override; + virtual int get_sparse_data(node_id_t id, + IndexStorage::MemoryBlock &block) const override; + + //! Get sparse data from vector + virtual std::pair get_sparse_data_from_vector( + const void *vec) const override; + virtual int get_sparse_data_from_vector(const void *vec, + IndexStorage::MemoryBlock &block, + int &sparse_length) const override; + + //! Get sparse vector feature data by key + virtual 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 + 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 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 + virtual 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 + virtual 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 + virtual int add_vector_with_id(level_t level, node_id_t id, + const std::string &sparse_vec, + uint32_t sparse_count) override; + + virtual 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 + virtual void add_neighbor(level_t level, node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + + //! Dump index by dumper + virtual int dump(const IndexDumper::Pointer &dumper) override; + + virtual 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) + 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 + }; + 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->index) >> upper_neighbor_mask_bits_; + uint32_t offset = (((meta->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 = -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(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.level = level; + meta.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..ffcf309 --- /dev/null +++ b/src/core/algorithm/ivf/CMakeLists.txt @@ -0,0 +1,10 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +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 + 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..42bce13 --- /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 "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 + virtual ~Iterator(void) {} + + //! Retrieve pointer of data + virtual const void *data(void) const override { + return holder_->element((*elems_)[index_]); + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return index_ < elems_->size(); + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + return (*elems_)[index_]; + } + + //! Next iterator + virtual 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) + virtual size_t count(void) const override { + return elems_->size(); + } + + //! Retrieve dimension + virtual size_t dimension(void) const override { + return holder_->dimension(); + } + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const override { + return holder_->data_type(); + } + + //! Retrieve element size in bytes + virtual size_t element_size(void) const override { + return holder_->element_size(); + } + + //! Retrieve if it can multi-pass + virtual bool multipass(void) const override { + return true; + } + + //! Create a new iterator + virtual 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..72ec766 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_builder.h @@ -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. +#pragma once + +#include "framework/index_builder.h" +#include "framework/index_meta.h" +#include "ivf_centroid_index.h" + +namespace zvec { +namespace core { + +/*! IVF Builder + */ +class IVFBuilder : public IndexBuilder { + public: + //! Constructor + IVFBuilder(); + + //! Destructor + ~IVFBuilder(); + + //! Disable them + IVFBuilder(const IVFBuilder &) = delete; + IVFBuilder &operator=(const IVFBuilder &) = delete; + + public: + //! 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 file system + virtual int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + virtual 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 + virtual ~Iterator(void) {} + + //! Retrieve pointer of data + virtual const void *data(void) const override { + return holder_->element(id_); + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return id_ < holder_->count(); + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + return holder_->key(id_); + } + + //! Next iterator + virtual 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) + virtual size_t count(void) const override { + return features_->count(); + } + + //! Retrieve dimension + virtual size_t dimension(void) const override { + return features_->dimension(); + } + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const override { + return features_->data_type(); + } + + //! Retrieve element size in bytes + virtual size_t element_size(void) const override { + return features_->element_size(); + } + + //! Retrieve if it can multi-pass + virtual bool multipass(void) const override { + return true; + } + + //! Create a new iterator + virtual 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..c662f48 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_centroid_index.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 "ivf_centroid_index.h" +#include "framework/index_framework.h" +#include "metric/metric_params.h" +#include "quantizer/quantizer_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) {} + + protected: + //! Initialize Trainer + virtual int init(const IndexMeta &, const ailego::Params &) override { + return 0; + } + + //! Cleanup Trainer + virtual int cleanup(void) override { + return 0; + } + + //! Train the data + virtual int train(IndexHolder::Pointer) override { + return 0; + } + + //! Train the data + virtual int train(IndexThreads::Pointer, IndexHolder::Pointer) override { + return 0; + } + + //! Load index from file path or dir + virtual int load(IndexStorage::Pointer) override { + return 0; + } + + //! Dump index into file path or dir + virtual int dump(const IndexDumper::Pointer &) override { + return 0; + } + + //! Retrieve Index Meta + virtual const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve statistics + virtual const IndexTrainer::Stats &stats(void) const override { + return stats_; + } + + //! Retrieve the output indexes + virtual 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 + virtual int init(const ailego::Params &) override { + return 0; + } + + //! Cleanup Reformer + virtual int cleanup(void) override { + return 0; + } + + //! Load index from container + virtual int load(IndexStorage::Pointer) override { + return 0; + } + + //! Unload index + virtual int unload(void) override { + return 0; + } + + //! Transform query + virtual 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 + virtual 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 + virtual 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 + virtual int init(const ailego::Params &) override { + return 0; + } + + //! Cleanup Reformer + virtual int cleanup(void) override { + return 0; + } + + //! Load index from container + virtual int load(IndexStorage::Pointer) override { + return 0; + } + + //! Unload index + virtual int unload(void) override { + return 0; + } + + //! Transform query + virtual int transform(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const override { + return IndexError_NotImplemented; + } + + //! Transform queries + virtual 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 + virtual 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..bdbab8d --- /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 "framework/index_framework.h" +#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 + virtual ~Iterator(void) {} + + //! Retrieve pointer of data + virtual const void *data(void) const override { + return (*features_)[id_]; + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return id_ < features_->size(); + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + return id_; + } + + //! Next iterator + virtual 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) + virtual size_t count(void) const override { + return features_.size(); + } + + //! Retrieve dimension + virtual size_t dimension(void) const override { + return dimension_; + } + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const override { + return data_type_; + } + + //! Retrieve element size in bytes + virtual size_t element_size(void) const override { + return element_size_; + } + + //! Retrieve if it can multi-pass + virtual bool multipass(void) const override { + return true; + } + + //! Create a new iterator + virtual 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..7ae0fe6 --- /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 "framework/index_framework.h" + +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..781271b --- /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 "framework/index_framework.h" +#include "metric/metric_params.h" +#include "quantizer/quantizer_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_{}; + InvertedIndexHeader header_{}; + uint32_t cur_list_id_{0}; + uint32_t dumped_feature_count_{0}; + size_t dumped_features_size_{0}; + mutable size_t dumped_size_{0}; +}; + +} // 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..b096c92 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_entity.cc @@ -0,0 +1,973 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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; +} + +//! 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; + 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; + float 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); + + *(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; + float 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); + 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..184548c --- /dev/null +++ b/src/core/algorithm/ivf/ivf_entity.h @@ -0,0 +1,351 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_framework.h" +#include "metric/metric_params.h" +#include "quantizer/quantizer_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 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); + + //! 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_{}; + InvertedIndexHeader header_{}; + 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 +}; + +} // 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..61be68b --- /dev/null +++ b/src/core/algorithm/ivf/ivf_index_format.h @@ -0,0 +1,167 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_framework.h" + +namespace zvec { +namespace core { + +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]; + char index_meta[0]; +}; + +/*! 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..13de55f --- /dev/null +++ b/src/core/algorithm/ivf/ivf_index_provider.h @@ -0,0 +1,110 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_searcher.h" +#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 + virtual Iterator::Pointer create_iterator(void) override { + return Iterator::Pointer(new (std::nothrow) Iterator(entity_)); + } + + //! Retrieve count of vectors + virtual size_t count(void) const override { + return entity_->vector_count(); + } + + //! Retrieve dimension of vector + virtual size_t dimension(void) const override { + return meta_.dimension(); + } + + //! Retrieve type of vector + virtual IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + //! Retrieve vector size in bytes + virtual size_t element_size(void) const override { + return meta_.element_size(); + } + + //! Retrieve a vector using a primary key + virtual const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + //! Retrieve the owner class + virtual const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class Iterator : public IndexProvider::Iterator { + public: + Iterator(const IVFEntity::Pointer &entity) : entity_(entity) {} + + //! 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 + virtual const void *data(void) const override { + return entity_->get_vector(index_); + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return index_ < entity_->vector_count(); + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + return entity_->get_key(index_); + } + + //! Next iterator + virtual 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..a33a7aa --- /dev/null +++ b/src/core/algorithm/ivf/ivf_params.h @@ -0,0 +1,88 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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"); + +// 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..ddcb869 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_searcher.cc @@ -0,0 +1,317 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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(); + 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(); + uint32_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..37d44e9 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_searcher.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 "framework/index_searcher.h" +#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 + virtual int init(const ailego::Params ¶meters) override; + + //! Cleanup Searcher + virtual int cleanup(void) override; + + //! Load index from container + virtual int load(IndexStorage::Pointer container, + IndexMetric::Pointer metric) override; + + //! Unload index + virtual int unload(void) override; + + //! Similarity brute force search + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const override; + + //! Similarity search + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const override; + + //! Retrieve statistics + virtual const Stats &stats(void) const override; + + //! Create a searcher context + virtual Context::Pointer create_context(void) const override; + + //! Create a new iterator + virtual IndexProvider::Pointer create_provider(void) const override; + + //! Retrieve meta of index + virtual const IndexMeta &meta(void) const override { + return meta_; + } + + //! Retrieve params of index + virtual 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..959ac93 --- /dev/null +++ b/src/core/algorithm/ivf/ivf_searcher_context.h @@ -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. +#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 + virtual void set_topk(uint32_t k) override { + topk_ = k; + result_heap_.limit(topk_); + result_heap_.set_threshold(this->threshold()); + } + + //! Retrieve search result + virtual const IndexDocumentList &result(void) const override { + return results_[0]; + } + + //! Retrieve search result with index + virtual const IndexDocumentList &result(size_t idx) const override { + ailego_assert_with(results_.size() > idx, "invalid index"); + return results_[idx]; + } + + //! Retrieve mutable result with index + virtual 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 + virtual 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); + centroid_searcher_ctx_->set_topk(topk_val); + 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 + virtual 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..d9fdb8e --- /dev/null +++ b/src/core/algorithm/ivf/ivf_streamer.cc @@ -0,0 +1,317 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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(); + 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(); + uint32_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..f395f6d --- /dev/null +++ b/src/core/algorithm/ivf/ivf_streamer.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. +#ifndef __IVF_STREAMER_H__ +#define __IVF_STREAMER_H__ + +#include "framework/index_streamer.h" +#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 + virtual int init(const IndexMeta & /*meta*/, + const ailego::Params & /*params*/) override; + + //! Cleanup Searcher + virtual int cleanup(void) override; + + //! Load index from container + virtual int open(IndexStorage::Pointer storage) override; + + virtual int flush(uint64_t /*check_point*/) override { + return 0; + } + virtual int close(void) override { + return this->unload(); + } + + //! Unload index + virtual int unload(void) override; + + //! Similarity brute force search + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity brute force search + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const override; + + //! Similarity search + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Similarity search + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const override; + + //! Retrieve statistics + virtual const Stats &stats(void) const override; + + //! Create a searcher context + virtual Context::Pointer create_context(void) const override; + + //! Create a new iterator + virtual IndexProvider::Pointer create_provider(void) const override; + + //! Retrieve meta of index + virtual const IndexMeta &meta(void) const override { + return meta_; + } + + virtual 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..01cd39b --- /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, args...) \ + do { \ + if (ailego_unlikely((code) != 0)) { \ + LOG_ERROR(fmt, ##args); \ + return code; \ + } \ + } while (0) +#endif + +#ifndef ivf_assert_with_msg +#define ivf_assert_with_msg(cond, err, fmt, args...) \ + do { \ + if (ailego_unlikely(!(cond))) { \ + LOG_ERROR(fmt, ##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/framework/CMakeLists.txt b/src/core/framework/CMakeLists.txt new file mode 100644 index 0000000..edfa1ca --- /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 SHARED 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_builder.h b/src/core/framework/index_builder.h new file mode 100644 index 0000000..8863e31 --- /dev/null +++ b/src/core/framework/index_builder.h @@ -0,0 +1,61 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_helper.h" +#include "index_holder.h" +#include "index_meta.h" +#include "index_runner.h" + +namespace zvec { +namespace core { + +class IndexBuilder : public IndexRunner { + public: + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexBuilder(void) {} + + //! Initialize the builder + virtual int init(const IndexMeta & /*meta*/, + const ailego::Params & /*params*/) { + return IndexError_NotImplemented; + } + + //! Train and build the index + static int TrainAndBuild(const IndexBuilder::Pointer &builder, + IndexHolder::Pointer holder) { + auto two_pass_holder = IndexHelper::MakeTwoPassHolder(std::move(holder)); + int ret = builder->train(two_pass_holder); + if (ret == 0) { + ret = builder->build(std::move(two_pass_holder)); + } + return ret; + } + + //! Train, build and dump the index + static int TrainBuildAndDump(const IndexBuilder::Pointer &builder, + IndexHolder::Pointer holder, + const IndexDumper::Pointer &dumper) { + int ret = IndexBuilder::TrainAndBuild(builder, std::move(holder)); + if (ret == 0) { + ret = builder->dump(dumper); + } + return ret; + } +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_bundle.h b/src/core/framework/index_bundle.h new file mode 100644 index 0000000..2f34566 --- /dev/null +++ b/src/core/framework/index_bundle.h @@ -0,0 +1,274 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Index Bundle + */ +struct IndexBundle { + //! Index Bundle Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexBundle(void) {} + + //! Retrieve index buffer via key + virtual ailego::BlobWrap get(const std::string &key) const = 0; + + //! Test if the key is exist + virtual bool has(const std::string &key) const = 0; + + //! Retrieve all + virtual std::map all(void) const = 0; + + //! Retrieve the count of indexes + virtual size_t count(void) const = 0; +}; + +/*! Trivial Index Bundle + */ +class TrivialIndexBundle : public IndexBundle { + public: + //! Trivial Index Bundle Pointer + typedef std::shared_ptr Pointer; + + //! Retrieve index buffer via key + virtual ailego::BlobWrap get(const std::string &key) const { + auto iter = map_.find(key); + if (iter != map_.end()) { + return iter->second; + } + return ailego::BlobWrap(); + } + + //! Test if the key is exist + virtual bool has(const std::string &key) const { + return (map_.find(key) != map_.end()); + } + + //! Retrieve all + virtual std::map all(void) const { + return map_; + } + + //! Retrieve the count of indexes + virtual size_t count(void) const { + return map_.size(); + } + + //! Set an index buffer in bundle + void set(const std::string &key, const ailego::BlobWrap &blob) { + map_[key] = blob; + } + + //! Set an index buffer in bundle + void set(std::string &&key, const ailego::BlobWrap &blob) { + map_[std::move(key)] = blob; + } + + //! Set an index buffer in bundle + void set(const std::string &key, const void *buf, size_t len) { + map_[key] = ailego::BlobWrap(buf, len); + } + + //! Set an index buffer in bundle + void set(std::string &&key, const void *buf, size_t len) { + map_[std::move(key)] = ailego::BlobWrap(buf, len); + } + + private: + std::map map_; +}; + +/*! Memory Index Bundle + */ +class MemoryIndexBundle : public IndexBundle { + public: + //! Memory Index Bundle Pointer + typedef std::shared_ptr Pointer; + + //! Retrieve index buffer via key + virtual ailego::BlobWrap get(const std::string &key) const { + auto iter = map_.find(key); + if (iter != map_.end()) { + return ailego::BlobWrap(iter->second.data(), iter->second.size()); + } + return ailego::BlobWrap(); + } + + //! Test if the key is exist + virtual bool has(const std::string &key) const { + return (map_.find(key) != map_.end()); + } + + //! Retrieve all + virtual std::map all(void) const { + std::map result; + for (const auto &it : map_) { + result.emplace(it.first, + ailego::BlobWrap(it.second.data(), it.second.size())); + } + return result; + } + + //! Retrieve the count of indexes + virtual size_t count(void) const { + return map_.size(); + } + + //! Set an index buffer in bundle + void set(const std::string &key, const std::string &buf) { + map_[key] = buf; + } + + //! Set an index buffer in bundle + void set(std::string &&key, const std::string &buf) { + map_[std::move(key)] = buf; + } + + //! Set an index buffer in bundle + void set(const std::string &key, std::string &&buf) { + map_[key] = std::move(buf); + } + + //! Set an index buffer in bundle + void set(std::string &&key, std::string &&buf) { + map_[std::move(key)] = std::move(buf); + } + + //! Set an index buffer in bundle + void set(const std::string &key, const void *buf, size_t len) { + map_[key].assign(reinterpret_cast(buf), len); + } + + //! Set an index buffer in bundle + void set(std::string &&key, const void *buf, size_t len) { + map_[std::move(key)].assign(reinterpret_cast(buf), len); + } + + private: + std::map map_; +}; + +/*! MMap File Index Bundle + */ +class MMapFileIndexBundle : public IndexBundle { + public: + //! Memory Index Bundle Pointer + typedef std::shared_ptr Pointer; + + //! Retrieve index buffer via key + virtual ailego::BlobWrap get(const std::string &key) const { + auto iter = map_.find(key); + if (iter != map_.end()) { + return ailego::BlobWrap(iter->second.region(), iter->second.size()); + } + return ailego::BlobWrap(); + } + + //! Test if the key is exist + virtual bool has(const std::string &key) const { + return (map_.find(key) != map_.end()); + } + + //! Retrieve all + virtual std::map all(void) const { + std::map result; + for (const auto &it : map_) { + result.emplace(it.first, + ailego::BlobWrap(it.second.region(), it.second.size())); + } + return result; + } + + //! Retrieve the count of indexes + virtual size_t count(void) const { + return map_.size(); + } + + //! Create a memory mapping file in bundle + bool create(const std::string &prefix, const std::string &key, size_t len) { + ailego::MMapFile file; + if (!file.create(prefix + '/' + key, len)) { + return false; + } + map_[key] = std::move(file); + return true; + } + + //! Create a memory mapping file in bundle + bool create(const std::string &prefix, std::string &&key, size_t len) { + ailego::MMapFile file; + if (!file.create(prefix + '/' + key, len)) { + return false; + } + map_[std::move(key)] = std::move(file); + return true; + } + + //! Create a memory mapping file in bundle + bool create(const std::string &path, size_t len) { + ailego::MMapFile file; + if (!file.create(path, len)) { + return false; + } + map_[ailego::File::BaseName(path)] = std::move(file); + return true; + } + + //! Open a memory mapping file in bundle + bool open(const std::string &prefix, const std::string &key, bool rdonly) { + ailego::MMapFile file; + if (!file.open(prefix + '/' + key, rdonly)) { + return false; + } + map_[key] = std::move(file); + return true; + } + + //! Open a memory mapping file in bundle + bool open(const std::string &prefix, std::string &&key, bool rdonly) { + ailego::MMapFile file; + if (!file.open(prefix + '/' + key, rdonly)) { + return false; + } + map_[std::move(key)] = std::move(file); + return true; + } + + //! Open a memory mapping file in bundle + bool open(const std::string &path, bool rdonly) { + ailego::MMapFile file; + if (!file.open(path, rdonly)) { + return false; + } + map_[ailego::File::BaseName(path)] = std::move(file); + return true; + } + + private: + std::map map_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_cluster.cc b/src/core/framework/index_cluster.cc new file mode 100644 index 0000000..118c946 --- /dev/null +++ b/src/core/framework/index_cluster.cc @@ -0,0 +1,158 @@ +// 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 "index_cluster.h" +#include "index_bundle.h" +#include "index_error.h" + +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_cluster.h b/src/core/framework/index_cluster.h new file mode 100644 index 0000000..acb26b4 --- /dev/null +++ b/src/core/framework/index_cluster.h @@ -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. +#pragma once + +#include +#include +#include +#include "index_bundle.h" +#include "index_features.h" +#include "index_meta.h" +#include "index_module.h" +#include "index_threads.h" + +namespace zvec { +namespace core { + +/*! Index Cluster + */ +struct IndexCluster : public IndexModule { + public: + //! Index Cluster Pointer + typedef std::shared_ptr Pointer; + + /*! Index Cluster Centroid + */ + class Centroid { + public: + //! Constructor + Centroid(void) + : buffer_(), score_(0.0), follows_(0), similars_(), subitems_() {} + + //! Constructor + Centroid(const void *feat, size_t bytes) + : buffer_(std::string(reinterpret_cast(feat), bytes)), + score_(0.0), + follows_(0), + similars_(), + subitems_() {} + + //! Constructor + Centroid(const Centroid &rhs) + : buffer_(rhs.buffer_), + score_(rhs.score_), + follows_(rhs.follows_), + similars_(rhs.similars_), + subitems_(rhs.subitems_) {} + + //! Constructor + Centroid(Centroid &&rhs) + : buffer_(std::move(rhs.buffer_)), + score_(rhs.score_), + follows_(rhs.follows_), + similars_(std::move(rhs.similars_)), + subitems_(std::move(rhs.subitems_)) {} + + //! Assignment + Centroid &operator=(const Centroid &rhs) { + buffer_ = rhs.buffer_; + score_ = rhs.score_; + follows_ = rhs.follows_; + similars_ = rhs.similars_; + subitems_ = rhs.subitems_; + return *this; + } + + //! Assignment + Centroid &operator=(Centroid &&rhs) { + buffer_ = std::move(rhs.buffer_); + score_ = rhs.score_; + follows_ = rhs.follows_; + similars_ = std::move(rhs.similars_); + subitems_ = std::move(rhs.subitems_); + return *this; + } + + //! Less than + bool operator<(const Centroid &rhs) const { + return (this->score_ < rhs.score_); + } + + //! Test if matchs the meta + bool is_matched(const IndexMeta &meta) const { + if (buffer_.size() != meta.element_size()) { + return false; + } + for (const auto &it : subitems_) { + if (!it.is_matched(meta)) { + return false; + } + } + return true; + } + + //! Set feature of centroid + void set_feature(const void *feat, size_t bytes) { + buffer_.assign(std::string(reinterpret_cast(feat), bytes)); + } + + //! Set feature of centroid + template + void set_feature(const ailego::NumericalVector &feat) { + buffer_.assign(feat); + } + + //! Set feature of centroid + template + void set_feature(ailego::NumericalVector &&feat) { + buffer_.assign(std::forward>(feat)); + } + + //! Set score of centroid + void set_score(double val) { + score_ = val; + } + + //! Set follows of centroid + void set_follows(size_t count) { + follows_ = count; + } + + //! Set similars of centroid + void set_similars(const std::vector &feats) { + similars_ = feats; + } + + //! Set similars of centroid + void set_similars(std::vector &&feats) { + similars_ = std::move(feats); + } + + //! Set subitems of centroid + void set_subitems(const std::vector ¢s) { + subitems_ = cents; + } + + //! Set subitems of centroid + void set_subitems(std::vector &¢s) { + subitems_ = std::move(cents); + } + + //! Retrieve feature buffer + std::string *mutable_buffer(void) { + return &buffer_; + } + + //! Retrieve feature buffer + const std::string &buffer(void) const { + return buffer_; + } + + //! Retrieve feature vector + template + ailego::NumericalVector *mutable_vector(void) { + return static_cast *>(&buffer_); + } + + //! Retrieve feature vector + template + const ailego::NumericalVector &vector(void) const { + return static_cast &>(buffer_); + } + + //! Retrieve feature pointer + const void *feature(void) const { + return buffer_.data(); + } + + //! Retrieve size of centroid in bytes + size_t size(void) const { + return buffer_.size(); + } + + //! Retrieve score of centroid + double score(void) const { + return score_; + } + + //! Retrieve follows' count of centroid + size_t follows(void) const { + return follows_; + } + + //! Retrieve similars of centroid + const std::vector &similars(void) const { + return similars_; + } + + //! Retrieve similars of centroid + std::vector *mutable_similars(void) { + return &similars_; + } + + //! Retrieve the sub centroids + const std::vector &subitems(void) const { + return subitems_; + } + + //! Retrieve the sub centroids + std::vector *mutable_subitems(void) { + return &subitems_; + } + + //! Retrieve the count of subitems (includes children's children) + size_t subcount(void) const { + size_t total = subitems_.size(); + for (const auto &it : subitems_) { + total += it.subcount(); + } + return total; + } + + private: + //! Members + std::string buffer_; + double score_; + size_t follows_; + std::vector similars_; + std::vector subitems_; + }; + + //! Index Cluster Centroid List + typedef std::vector CentroidList; + + //! Destructor + virtual ~IndexCluster(void) {} + + //! Deserialize centroids from bundle + static int Deserialize(const IndexMeta &meta, IndexBundle::Pointer bundle, + CentroidList *cents); + + //! Serialize centroids into bundle + static int Serialize(const IndexMeta &meta, const CentroidList ¢s, + IndexBundle::Pointer *out); + + //! Initialize Cluster + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) = 0; + + //! Cleanup Cluster + virtual int cleanup(void) = 0; + + //! Reset Cluster + virtual int reset(void) = 0; + + //! Update Cluster + virtual int update(const ailego::Params ¶ms) = 0; + + //! Suggest dividing to K clusters + virtual void suggest(uint32_t k) = 0; + + //! Mount features + virtual int mount(IndexFeatures::Pointer feats) = 0; + + //! Cluster + virtual int cluster(CentroidList ¢s) { + return this->cluster(nullptr, cents); + } + + //! Cluster + virtual int cluster(IndexThreads::Pointer threads, CentroidList ¢s) = 0; + + //! Classify + virtual int classify(CentroidList ¢s) { + return this->classify(nullptr, cents); + } + + //! Classify + virtual int classify(IndexThreads::Pointer threads, CentroidList ¢s) = 0; + + //! Label + virtual int label(const CentroidList ¢s, std::vector *out) { + return this->label(nullptr, cents, out); + } + + //! Label + virtual int label(IndexThreads::Pointer threads, const CentroidList ¢s, + std::vector *out) = 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..e7ffa36 --- /dev/null +++ b/src/core/framework/index_context.cc @@ -0,0 +1,27 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_context.h" +#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_context.h b/src/core/framework/index_context.h new file mode 100644 index 0000000..20f62ba --- /dev/null +++ b/src/core/framework/index_context.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 "index_document.h" +#include "index_error.h" +#include "index_filter.h" +#include "index_groupby.h" +#include "index_metric.h" +#include "index_stats.h" + +namespace zvec { +namespace core { + +/*! Profiler + */ +struct Profiler { + Profiler() = default; + ~Profiler() = default; + + void add(const std::string &name, double time) { + timings[name] += time; + } + + std::string display() const { + std::string info = "================================================\n"; + + for (auto itr = timings.begin(); itr != timings.end(); ++itr) { + info += + itr->first + std::string(": ") + std::to_string(itr->second) + " s\n"; + } + + info += "================================================\n"; + + return info; + } + + std::map timings; +}; + +/*! Index Context + */ +class IndexContext { + public: + //! Index Context Pointer + typedef std::unique_ptr Pointer; + + //! Index Context UPointer + typedef std::unique_ptr UPointer; + + /*! Index Context Stats + */ + class Stats : public IndexStats { + public: + //! Set count of documents filtered + void set_filtered_count(size_t count) { + filtered_count_ = count; + } + + //! Set count of documents dist calced + void set_dist_calced_count(size_t count) { + dist_calced_count_ = count; + } + + //! Retrieve count of documents filtered + size_t filtered_count(void) const { + return filtered_count_; + } + + //! Retrieve count of documents dist-calced + size_t dist_calced_count(void) const { + return dist_calced_count_; + } + + //! Retrieve count of documents filtered (mutable) + size_t *mutable_filtered_count(void) { + return &filtered_count_; + } + + //! Retrieve count of documents dist-calced (mutable) + size_t *mutable_dist_calced_count(void) { + return &dist_calced_count_; + } + + void clear() { + this->clear_attributes(); + + filtered_count_ = 0u; + dist_calced_count_ = 0u; + } + + private: + //! Members + size_t filtered_count_{0u}; + size_t dist_calced_count_{0u}; + }; + + //! Constructor + IndexContext() {} + + //! Constructor + IndexContext(IndexMetric::Pointer index_metric) + : index_metric_(std::move(index_metric)) {} + + //! Destructor + virtual ~IndexContext(void) {} + + //! Set topk of search result + virtual void set_topk(uint32_t topk) = 0; + + virtual uint32_t topk() const { + return 0; + } + + virtual void set_group_params(uint32_t /*group_mum*/, + uint32_t /*group_topk*/){}; + + //! Set brute force threshold + virtual void set_bruteforce_threshold(uint32_t /*bruteforce_threshold*/) {} + + //! Set mode of debug + virtual void set_debug_mode(bool /*enable*/) {} + + //! Set fetch vector + virtual void set_fetch_vector(bool /*enable*/) {} + + //! Retrieve search result + virtual const IndexDocumentList &result(void) const = 0; + + //! Retrieve search result with index + virtual const IndexDocumentList &result(size_t /*index*/) const { + return this->result(); + } + + //! Retrieve mutable result with index + virtual IndexDocumentList *mutable_result(size_t idx) = 0; + + //! Retrieve search group result with index + virtual const IndexGroupDocumentList &group_result(void) const { + // to make it compile + static const IndexGroupDocumentList empty_list{}; + return empty_list; + }; + + //! Retrieve search group result with index + virtual const IndexGroupDocumentList &group_result(size_t /*idx*/) const { + return this->group_result(); + } + + //! Update the parameters of context + virtual int update(const ailego::Params & /*params*/) { + return IndexError_NotImplemented; + } + + //! Retrieve mode of debug + virtual bool debug_mode(void) const { + return false; + } + + //! Retrieve debug information + virtual std::string debug_string(void) const { + return std::string(); + } + + //! Retrieve magic number + virtual uint32_t magic(void) const { + return 0; + } + + //! Retrieve search filter + const IndexFilter &filter(void) const { + return filter_; + } + + //! Retrieve fetch vector + virtual bool fetch_vector(void) const { + return false; + } + + //! Reset context + virtual void reset(void) {} + + //! Set the filter of context + template + void set_filter(T &&func) { + filter_.set(std::forward(func)); + } + + //! Reset the filter of context + void reset_filter(void) { + filter_.reset(); + } + + //! Retrieve search groupby + const IndexGroupBy &group_by(void) const { + return group_by_; + } + + //! Set the groupby of context + template + void set_group_by(T &&func) { + group_by_.set(std::forward(func)); + } + + //! Reset the groupby of context + void reset_group_by(void) { + group_by_.reset(); + } + + //! Set threshold for RNN + void set_threshold(float val) { + if (index_metric_ && index_metric_->support_normalize()) { + index_metric_->denormalize(&val); + } + + threshold_ = val; + } + + //! Retrieve value of threshold for RNN + float threshold(void) const { + return threshold_; + } + + //! Reset value of threshold for RNN + void reset_threshold(void) { + threshold_ = std::numeric_limits::max(); + } + + //! Generate a global magic number + static uint32_t GenerateMagic(void); + + //! Profiler + Profiler &profiler() { + return profiler_; + } + + private: + //! Members + IndexFilter filter_{}; + IndexGroupBy group_by_{}; + float threshold_{std::numeric_limits::max()}; + + + Profiler profiler_{}; + + protected: + IndexMetric::Pointer index_metric_{nullptr}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/framework/index_converter.cc b/src/core/framework/index_converter.cc new file mode 100644 index 0000000..251ee7b --- /dev/null +++ b/src/core/framework/index_converter.cc @@ -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. +#include "index_converter.h" +#include "index_error.h" +#include "index_helper.h" + +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_converter.h b/src/core/framework/index_converter.h new file mode 100644 index 0000000..db6f7e8 --- /dev/null +++ b/src/core/framework/index_converter.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 "index_dumper.h" +#include "index_holder.h" +#include "index_meta.h" +#include "index_stats.h" + +namespace zvec { +namespace core { + +/*! Index Converter + */ +class IndexConverter : public IndexModule { + public: + //! Index Converter Pointer + typedef std::shared_ptr Pointer; + + /*! Index Converter Stats + */ + class Stats : public IndexStats { + public: + Stats() {} + Stats(const Stats &stats) { + *this = stats; + } + Stats &operator=(const Stats &stats) { + this->trained_count_.store(stats.trained_count_.load()); + this->transformed_count_.store(stats.transformed_count_.load()); + this->dumped_size_.store(stats.dumped_size_.load()); + this->discarded_count_.store(stats.discarded_count_.load()); + this->trained_costtime_.store(stats.trained_costtime_.load()); + this->transformed_costtime_.store(stats.transformed_costtime_.load()); + this->dumped_costtime_.store(stats.dumped_costtime_.load()); + return *this; + } + //! Set count of documents trained + void set_trained_count(size_t count) { + trained_count_ = count; + } + + //! Set count of documents transformed + void set_transformed_count(size_t count) { + transformed_count_ = count; + } + + //! Set size of documents dumped + void set_dumped_size(size_t size) { + dumped_size_ = size; + } + + //! Set count of documents discarded + void set_discarded_count(size_t count) { + discarded_count_ = count; + } + + //! Set time cost of documents trained + void set_trained_costtime(uint64_t cost) { + trained_costtime_ = cost; + } + + //! Set time cost of documents transformed + void set_transformed_costtime(uint64_t cost) { + transformed_costtime_ = cost; + } + + //! Set time cost of documents dumped + void set_dumped_costtime(uint64_t cost) { + dumped_costtime_ = cost; + } + + //! Retrieve count of documents trained + size_t trained_count(void) const { + return trained_count_; + } + + //! Retrieve count of documents transformed + size_t transformed_count(void) const { + return transformed_count_; + } + + //! Retrieve size of documents dumped + size_t dumped_size(void) const { + return dumped_size_; + } + + //! Retrieve count of documents discarded + size_t discarded_count(void) const { + return discarded_count_; + } + + //! Retrieve time cost of documents trained + uint64_t trained_costtime(void) const { + return trained_costtime_; + } + + //! Retrieve time cost of documents transformed + uint64_t transformed_costtime(void) const { + return transformed_costtime_; + } + + //! Retrieve time cost of documents dumped + uint64_t dumped_costtime(void) const { + return dumped_costtime_; + } + + //! Retrieve count of documents trained (mutable) + std::atomic *mutable_trained_count(void) { + return &trained_count_; + } + + //! Retrieve count of documents transformed (mutable) + std::atomic *mutable_transformed_count(void) { + return &transformed_count_; + } + + //! Retrieve size of documents dumped (mutable) + std::atomic *mutable_dumped_size(void) { + return &dumped_size_; + } + + //! Retrieve count of documents discarded (mutable) + std::atomic *mutable_discarded_count(void) { + return &discarded_count_; + } + + //! Retrieve time cost of documents trained (mutable) + std::atomic *mutable_trained_costtime(void) { + return &trained_costtime_; + } + + //! Retrieve time cost of documents transformed (mutable) + std::atomic *mutable_transformed_costtime(void) { + return &transformed_costtime_; + } + + //! Retrieve time cost of documents dumped (mutable) + std::atomic *mutable_dumped_costtime(void) { + return &dumped_costtime_; + } + + private: + //! Members + std::atomic trained_count_{0u}; + std::atomic transformed_count_{0u}; + std::atomic dumped_size_{0u}; + std::atomic discarded_count_{0u}; + std::atomic trained_costtime_{0u}; + std::atomic transformed_costtime_{0u}; + std::atomic dumped_costtime_{0u}; + }; + + //! Destructor + virtual ~IndexConverter(void) {} + + //! Initialize Converter + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) = 0; + + //! Cleanup Converter + virtual int cleanup(void) = 0; + + //! Train the data + virtual int train(IndexHolder::Pointer) { + return IndexError_NotImplemented; + } + + //! Train the data + virtual int train(IndexSparseHolder::Pointer) { + return IndexError_NotImplemented; + } + + //! Transform the data + virtual int transform(IndexHolder::Pointer) { + return IndexError_NotImplemented; + }; + + //! Transform the data + virtual int transform(IndexSparseHolder::Pointer) { + return IndexError_NotImplemented; + } + + //! Dump index into storage + virtual int dump(const IndexDumper::Pointer &dumper) = 0; + + //! Retrieve statistics + virtual const Stats &stats(void) const = 0; + + //! Retrieve a holder as result + virtual IndexHolder::Pointer result(void) const { + return nullptr; + } + + //! Retrieve a holder as result + virtual IndexSparseHolder::Pointer sparse_result(void) const { + return nullptr; + } + + //! Retrieve Index Meta + virtual const IndexMeta &meta(void) const = 0; + + //! Train and transform the index + static int TrainAndTransform(const IndexConverter::Pointer &converter, + IndexHolder::Pointer holder); + + //! Train, transform and dump the index + static int TrainTransformAndDump(const IndexConverter::Pointer &converter, + IndexHolder::Pointer holder, + const IndexDumper::Pointer &dumper); +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_document.h b/src/core/framework/index_document.h new file mode 100644 index 0000000..37ac4f8 --- /dev/null +++ b/src/core/framework/index_document.h @@ -0,0 +1,317 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_storage.h" + +namespace zvec { +namespace core { + +/*! Index Sparse Document + */ +class IndexSparseDocument { + public: + //! Constructor + IndexSparseDocument() = default; + + //! Retrieve mutable sparse count + uint32_t *mutable_sparse_count(void) { + return &sparse_count_; + } + + //! Retrieve mutable sparse indices + std::string *mutable_sparse_indices(void) { + return &sparse_indices_; + } + + //! Retrieve mutable sparse values + std::string *mutable_sparse_values(void) { + return &sparse_values_; + } + + //! Retrieve sparse count + uint32_t sparse_count(void) const { + return sparse_count_; + } + + //! Retrieve sparse indices + const std::string &sparse_indices(void) const { + return sparse_indices_; + } + + //! Retrieve sparse values + const std::string &sparse_values(void) const { + return sparse_values_; + } + + private: + uint32_t sparse_count_{0}; + std::string sparse_indices_{}; + std::string sparse_values_{}; +}; + +/*! Index Document + */ +class IndexDocument { + public: + //! Constructor + IndexDocument() = default; + + //! Constructor + IndexDocument(uint64_t k, float v) : key_(k), score_(v) {} + + //! Constructor + IndexDocument(uint64_t k, float v, uint32_t i) + : key_(k), score_(v), index_(i) {} + + //! Constructor + IndexDocument(uint64_t k, float v, uint32_t i, const void *vector) + : key_(k), score_(v), index_(i), vector_(vector) {} + + IndexDocument(uint64_t k, float v, uint32_t i, + IndexStorage::MemoryBlock vec_block) + : key_(k), score_(v), index_(i), vec_mem_block_(vec_block) { + vector_ = vec_mem_block_.data(); + has_vec_mem_block_ = true; + } + + //! Constructor + IndexDocument(uint64_t k, float v, uint32_t i, const void *vector, + IndexSparseDocument sparse_doc) + : key_(k), + score_(v), + index_(i), + vector_(vector), + sparse_doc_(std::move(sparse_doc)) {} + + IndexDocument(uint64_t k, float v, uint32_t i, + IndexStorage::MemoryBlock vec_block, + IndexSparseDocument sparse_doc) + : key_(k), + score_(v), + index_(i), + vec_mem_block_(vec_block), + sparse_doc_(std::move(sparse_doc)) { + has_vec_mem_block_ = true; + vector_ = vec_mem_block_.data(); + } + + //! Constructor + IndexDocument(const IndexDocument &rhs) + : key_(rhs.key_), + score_(rhs.score_), + index_(rhs.index_), + vector_(rhs.vector_), + sparse_doc_{rhs.sparse_doc_} { + if (rhs.has_vec_mem_block_) { + vec_mem_block_ = rhs.vec_mem_block_; + has_vec_mem_block_ = true; + } + } + + //! Assignment + IndexDocument &operator=(const IndexDocument &rhs) { + if (this != &rhs) { + key_ = rhs.key_; + score_ = rhs.score_; + index_ = rhs.index_; + vector_ = rhs.vector_; + if (rhs.has_vec_mem_block_) { + vec_mem_block_ = rhs.vec_mem_block_; + has_vec_mem_block_ = true; + } + sparse_doc_ = rhs.sparse_doc_; + } + return *this; + } + + //! Less than + bool operator<(const IndexDocument &rhs) const { + return (this->score_ < rhs.score_); + } + + //! Greater than + bool operator>(const IndexDocument &rhs) const { + return (this->score_ > rhs.score_); + } + + //! Retrieve primary key + uint64_t key(void) const { + return key_; + } + + //! Retrieve score value + float score(void) const { + return score_; + } + + //! Retrieve index id + uint32_t index(void) const { + return index_; + } + + //! Retrieve vec + const void *vector() const { + return vector_; + } + + //! Retrieve vec + const IndexSparseDocument &sparse_doc() const { + return sparse_doc_; + } + + //! Retrieve mutable primary key + uint64_t *mutable_key(void) { + return &key_; + } + + //! Retrieve mutable score value + float *mutable_score(void) { + return &score_; + } + + //! Retrieve mutable index id + uint32_t *mutable_index(void) { + return &index_; + } + + //! Retrieve primary key + void set_key(uint64_t val) { + key_ = val; + } + + //! Retrieve score value + void set_score(float val) { + score_ = val; + } + + //! Retrieve index id + void set_index(uint32_t val) { + index_ = val; + } + + private: + //! Data members + uint64_t key_{0u}; + float score_{0.0f}; + uint32_t index_{0u}; + const void *vector_{nullptr}; + bool has_vec_mem_block_{false}; + mutable IndexStorage::MemoryBlock vec_mem_block_{}; + IndexSparseDocument sparse_doc_{}; +}; + +/*! Index Document Heap + */ +class IndexDocumentHeap : public ailego::Heap { + public: + //! Constructor + IndexDocumentHeap(void) : ailego::Heap() {} + + //! Constructor + IndexDocumentHeap(size_t max) : ailego::Heap(max) {} + + //! Constructor + IndexDocumentHeap(size_t max, float val) + : ailego::Heap(max), threshold_(val) {} + + //! Constructor + IndexDocumentHeap(const IndexDocumentHeap &rhs) + : ailego::Heap(rhs), threshold_(rhs.threshold_) {} + + //! Constructor + IndexDocumentHeap(IndexDocumentHeap &&rhs) + : ailego::Heap(std::move(rhs)), + threshold_(rhs.threshold_) {} + + //! Constructor + IndexDocumentHeap(const std::vector &rhs) + : ailego::Heap(rhs) {} + + //! Constructor + IndexDocumentHeap(std::vector &&rhs) + : ailego::Heap(std::move(rhs)) {} + + //! Insert a document into the heap + void emplace(uint64_t key, float score) { + if (score <= threshold_) { + ailego::Heap::emplace(key, score); + } + } + + //! Insert a document into the heap + void emplace(uint64_t key, float score, uint32_t index) { + if (score <= threshold_) { + ailego::Heap::emplace(key, score, index); + } + } + + //! Set threshold for RNN + void set_threshold(float val) { + threshold_ = val; + } + + //! Retrieve value of threshold for RNN + float threshold(void) const { + return threshold_; + } + + private: + //! members + float threshold_{std::numeric_limits::max()}; +}; + +class GroupIndexDocument { + public: + GroupIndexDocument() = default; + + GroupIndexDocument(const std::string &group_id, + const std::vector &docs) + : group_id_(group_id), docs_(docs) {} + + const std::string &group_id() const { + return group_id_; + } + + const std::vector &docs() const { + return docs_; + } + + std::string *mutable_group_id() { + return &group_id_; + } + + std::vector *mutable_docs() { + return &docs_; + } + + void set_group_id(const std::string &group_id) { + group_id_ = group_id; + } + + private: + std::string group_id_; + std::vector docs_; +}; + +/*! Index Document List + */ +using IndexDocumentList = std::vector; +using IndexGroupDocumentList = std::vector; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_dumper.h b/src/core/framework/index_dumper.h new file mode 100644 index 0000000..40133ab --- /dev/null +++ b/src/core/framework/index_dumper.h @@ -0,0 +1,163 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_module.h" +#include "index_packer.h" + +namespace zvec { +namespace core { + +/*! Index Dumper + */ +class IndexDumper : public IndexModule { + public: + //! Index Dumper Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexDumper(void) {} + + //! Initialize dumper + virtual int init(const ailego::Params ¶ms) = 0; + + //! Cleanup dumper + virtual int cleanup(void) = 0; + + //! Create a file for dumping + virtual int create(const std::string &path) = 0; + + //! Close file + virtual int close(void) = 0; + + //! Append a segment meta into table + virtual int append(const std::string &id, size_t data_size, + size_t padding_size, uint32_t crc) = 0; + + //! Write data to the storage + virtual size_t write(const void *data, size_t len) = 0; + + //! Retrieve magic number of index + virtual uint32_t magic(void) const = 0; +}; + +/*! Index Segment Dumper + */ +class IndexSegmentDumper : public IndexDumper { + public: + //! Index Segment Dumper Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + IndexSegmentDumper(IndexDumper::Pointer dumper, std::string segid) + : segment_id_(std::move(segid)), dumper_(std::move(dumper)) {} + + //! Destructor + virtual ~IndexSegmentDumper(void) { + this->close_index(); + } + + //! Initialize dumper + int init(const ailego::Params &) override { + return 0; + } + + //! Cleanup dumper + int cleanup(void) override { + return 0; + } + + //! Create a file for dumping + int create(const std::string &segid) override { + if (dumped_size_ != 0) { + return IndexError_NoReady; + } + + auto write_data = [&](const void *buf, size_t size) { + return this->write_to_dumper(buf, size); + }; + if (!packer_.setup(write_data)) { + return IndexError_WriteData; + } + segment_id_ = segid; + return 0; + } + + //! Close file + int close(void) override { + return this->close_index(); + } + + //! Append a segment meta into table + int append(const std::string &id, size_t data_size, size_t padding_size, + uint32_t crc) override { + stab_.emplace_back(id, data_size, padding_size, crc); + return 0; + } + + //! Write data to the storage + size_t write(const void *data, size_t len) override { + auto write_data = [&](const void *buf, size_t size) { + return this->write_to_dumper(buf, size); + }; + + if (dumped_size_ == 0 && !packer_.setup(write_data)) { + return 0; + } + return packer_.pack(write_data, data, len); + } + + //! Retrieve magic number of index + uint32_t magic(void) const override { + return packer_.magic(); + } + + protected: + //! Write data to dumper + size_t write_to_dumper(const void *data, size_t len) { + size_t wrlen = dumper_->write(data, len); + dumped_size_ += wrlen; + return wrlen; + } + + //! Close index file + int close_index(void) { + if (dumped_size_ == 0) { + return 0; + } + + auto write_data = [&](const void *buf, size_t size) { + return this->write_to_dumper(buf, size); + }; + if (!packer_.finish(write_data, stab_)) { + return IndexError_WriteData; + } + stab_.clear(); + + int ret = dumper_->append(segment_id_, dumped_size_, 0, 0); + dumped_size_ = 0u; + return ret; + } + + private: + size_t dumped_size_{0}; + std::string segment_id_{}; + IndexDumper::Pointer dumper_{}; + IndexPacker packer_{}; + std::vector stab_{}; +}; + +} // 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..e7184bf --- /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 "index_error.h" + +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_error.h b/src/core/framework/index_error.h new file mode 100644 index 0000000..e95acc1 --- /dev/null +++ b/src/core/framework/index_error.h @@ -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. +#pragma once + +#include +#include "ailego/pattern/expected.hpp" + +namespace zvec { +namespace core { + +/*! Error + */ + +class ErrorCode; + +template +using Result = tl::expected; + +template +using result_unexpect = tl::unexpected; + +/*! Index Error + */ +class IndexError { + public: + /*! Index Error Code + */ + class Code { + public: + //! Constructor + Code(int val, const char *str) : value_(-val), desc_(str) { + IndexError::Instance()->emplace(this); + } + + //! Retrieve the value of code + operator int() const { + return (this->value_); + } + + //! Retrieve the value of code + int value() const { + return (this->value_); + } + + //! Retrieve the description of code + const char *desc() const { + return (this->desc_); + } + + private: + int value_; + const char *desc_; + }; + + //! Retrieve the description of code + static const char *What(int val) { + return IndexError::Instance()->what(val); + } + + protected: + //! Constructor + IndexError(void) : map_() {} + + //! Inserts a new code into map + void emplace(const IndexError::Code *code) { + map_.emplace(code->value(), code); + } + + //! Retrieve the description of code + const char *what(int val) const { + auto iter = map_.find(val); + if (iter != map_.end()) { + return iter->second->desc(); + } + return ""; + } + + //! Retrieve the singleton + static IndexError *Instance(void) { + static IndexError error; + return (&error); + } + + private: + //! Disable them + IndexError(const IndexError &) = delete; + IndexError(IndexError &&) = delete; + IndexError &operator=(const IndexError &) = delete; + + //! Error code map + std::map map_; +}; + +//! Index Error Code Define +#define INDEX_ERROR_CODE_DEFINE(__NAME__, __VAL__, __DESC__) \ + const IndexError::Code IndexError_##__NAME__((__VAL__), (__DESC__)); \ + const IndexError::Code &_IndexErrorCode_##__VAL__##_Register( \ + IndexError_##__NAME__) + +//! Index Error Code Declare +#define INDEX_ERROR_CODE_DECLARE(__NAME__) \ + extern const IndexError::Code IndexError_##__NAME__ + +//! Build-in error code +INDEX_ERROR_CODE_DECLARE(Success); // Success +INDEX_ERROR_CODE_DECLARE(Runtime); // Runtime error +INDEX_ERROR_CODE_DECLARE(Logic); // Logic error +INDEX_ERROR_CODE_DECLARE(Type); // Type error +INDEX_ERROR_CODE_DECLARE(System); // System call error +INDEX_ERROR_CODE_DECLARE(Cast); // Cast error +INDEX_ERROR_CODE_DECLARE(IO); // IO error +INDEX_ERROR_CODE_DECLARE(AuthExpired); // Auth expired error + +INDEX_ERROR_CODE_DECLARE(NotImplemented); // Not implemented +INDEX_ERROR_CODE_DECLARE(Unsupported); // Unsupported +INDEX_ERROR_CODE_DECLARE(Denied); // Permission denied +INDEX_ERROR_CODE_DECLARE(Canceled); // Operation canceled +INDEX_ERROR_CODE_DECLARE(Overflow); // Overflow +INDEX_ERROR_CODE_DECLARE(Underflow); // Underflow +INDEX_ERROR_CODE_DECLARE(OutOfRange); // Out of range +INDEX_ERROR_CODE_DECLARE(NoBuffer); // No buffer space available +INDEX_ERROR_CODE_DECLARE(NoMemory); // Not enough space +INDEX_ERROR_CODE_DECLARE(NoParamFound); // No parameter found +INDEX_ERROR_CODE_DECLARE(NoReady); // No ready +INDEX_ERROR_CODE_DECLARE(NoExist); // No exist +INDEX_ERROR_CODE_DECLARE(Exist); // Already exist +INDEX_ERROR_CODE_DECLARE(Mismatch); // Mismatch +INDEX_ERROR_CODE_DECLARE(Duplicate); // Duplicate +INDEX_ERROR_CODE_DECLARE(Uninitialized); // Uninitialized + +INDEX_ERROR_CODE_DECLARE(InvalidArgument); // Invalid argument +INDEX_ERROR_CODE_DECLARE(InvalidFormat); // Invalid format +INDEX_ERROR_CODE_DECLARE(InvalidLength); // Invalid length +INDEX_ERROR_CODE_DECLARE(InvalidChecksum); // Invalid checksum +INDEX_ERROR_CODE_DECLARE(InvalidValue); // Invalid value + +INDEX_ERROR_CODE_DECLARE(CreateDirectory); // Create directory error +INDEX_ERROR_CODE_DECLARE(OpenDirectory); // Open directory error +INDEX_ERROR_CODE_DECLARE(Serialize); // Serialize error +INDEX_ERROR_CODE_DECLARE(Deserialize); // Deserialize error +INDEX_ERROR_CODE_DECLARE(CreateFile); // Create file error +INDEX_ERROR_CODE_DECLARE(OpenFile); // Open file error +INDEX_ERROR_CODE_DECLARE(SeekFile); // Seek file error +INDEX_ERROR_CODE_DECLARE(CloseFile); // Close file error +INDEX_ERROR_CODE_DECLARE(TruncateFile); // TruncateFile file error +INDEX_ERROR_CODE_DECLARE(MMapFile); // MMap file error +INDEX_ERROR_CODE_DECLARE(FlushFile); // Flush file error +INDEX_ERROR_CODE_DECLARE(WriteData); // Write data error +INDEX_ERROR_CODE_DECLARE(ReadData); // Read data error + +INDEX_ERROR_CODE_DECLARE(PackIndex); // Read data error +INDEX_ERROR_CODE_DECLARE(UnpackIndex); // Read data error +INDEX_ERROR_CODE_DECLARE(IndexLoaded); // Index loaded +INDEX_ERROR_CODE_DECLARE(NoIndexLoaded); // No index loaded +INDEX_ERROR_CODE_DECLARE(NoTrained); // No trained +INDEX_ERROR_CODE_DECLARE(IndexFull); // Index full + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/framework/index_factory.cc b/src/core/framework/index_factory.cc new file mode 100644 index 0000000..687100f --- /dev/null +++ b/src/core/framework/index_factory.cc @@ -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. +#include "index_factory.h" + +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_factory.h b/src/core/framework/index_factory.h new file mode 100644 index 0000000..9b9332f --- /dev/null +++ b/src/core/framework/index_factory.h @@ -0,0 +1,286 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_builder.h" +#include "index_cluster.h" +#include "index_converter.h" +#include "index_dumper.h" +#include "index_logger.h" +#include "index_metric.h" +#include "index_reducer.h" +#include "index_refiner.h" +#include "index_reformer.h" +#include "index_searcher.h" +#include "index_storage.h" +#include "index_streamer.h" +#include "index_trainer.h" + +namespace zvec { +namespace core { + +/*! Index Factory + */ +struct IndexFactory { + //! Create a index Metric by name + static IndexMetric::Pointer CreateMetric(const std::string &name); + + //! Test if the Metric is exist + static bool HasMetric(const std::string &name); + + //! Retrieve all Metric classes + static std::vector AllMetrics(void); + + //! Create a index logger by name + static IndexLogger::Pointer CreateLogger(const std::string &name); + + //! Test if the logger is exist + static bool HasLogger(const std::string &name); + + //! Retrieve all logger classes + static std::vector AllLoggers(void); + + //! Create a index dumper by name + static IndexDumper::Pointer CreateDumper(const std::string &name); + + //! Test if the dumper is exist + static bool HasDumper(const std::string &name); + + //! Retrieve all dumper classes + static std::vector AllDumpers(void); + + //! Test if the container is exist + static bool HasContainer(const std::string &name); + + //! Retrieve all container classes + static std::vector AllContainers(void); + + //! Create a index storage by name + static IndexStorage::Pointer CreateStorage(const std::string &name); + + //! Test if the storage is exist + static bool HasStorage(const std::string &name); + + //! Retrieve all storage classes + static std::vector AllStorages(void); + + //! Create a index converter by name + static IndexConverter::Pointer CreateConverter(const std::string &name); + + //! Test if the converter is exist + static bool HasConverter(const std::string &name); + + //! Retrieve all converter classes + static std::vector AllConverters(void); + + //! Create a index reformer by name + static IndexReformer::Pointer CreateReformer(const std::string &name); + + //! Test if the reformer is exist + static bool HasReformer(const std::string &name); + + //! Retrieve all reformer classes + static std::vector AllReformers(void); + + //! Create a index trainer by name + static IndexTrainer::Pointer CreateTrainer(const std::string &name); + + //! Test if the trainer is exist + static bool HasTrainer(const std::string &name); + + //! Retrieve all trainer classes + static std::vector AllTrainers(void); + + //! Create a index builder by name + static IndexBuilder::Pointer CreateBuilder(const std::string &name); + + //! Test if the builder is exist + static bool HasBuilder(const std::string &name); + + //! Retrieve all builder classes + static std::vector AllBuilders(void); + + //! Create a index searcher by name + static IndexSearcher::Pointer CreateSearcher(const std::string &name); + + //! Test if the searcher is exist + static bool HasSearcher(const std::string &name); + + //! Retrieve all searcher classes + static std::vector AllSearchers(void); + + //! Create a index streamer by name + static IndexStreamer::Pointer CreateStreamer(const std::string &name); + + //! Test if the streamer is exist + static bool HasStreamer(const std::string &name); + + //! Retrieve all streamer classes + static std::vector AllStreamers(void); + + //! Create a index reducer by name + static IndexReducer::Pointer CreateReducer(const std::string &name); + + //! Test if the reducer is exist + static bool HasReducer(const std::string &name); + + //! Retrieve all reducer classes + static std::vector AllReducers(void); + + //! Create a index cluster by name + static IndexCluster::Pointer CreateCluster(const std::string &name); + + //! Test if the cluster is exist + static bool HasCluster(const std::string &name); + + //! Retrieve all cluster classes + static std::vector AllClusters(void); + + //! Create a index streamer reducer by name + static IndexStreamerReducer::Pointer CreateStreamerReducer( + const std::string &name); + + //! Test if the streamer reducer is exist + static bool HasStreamerReducer(const std::string &name); + + //! Retrieve all streamer reducer classes + static std::vector AllStreamerReducers(void); + + //! Create a refiner by name + static IndexRefiner::Pointer CreateRefiner(const std::string &name); + + //! Test if the refiner is exist + static bool HasRefiner(const std::string &name); + + //! Retrieve all refiner classes + static std::vector AllRefiners(void); +}; + +//! Register Index Metric +#define INDEX_FACTORY_REGISTER_METRIC_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexMetric, __IMPL__, ##__VA_ARGS__) + +//! Register Index Metric +#define INDEX_FACTORY_REGISTER_METRIC(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_METRIC_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Logger +#define INDEX_FACTORY_REGISTER_LOGGER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexLogger, __IMPL__, ##__VA_ARGS__) + +//! Register Index Logger +#define INDEX_FACTORY_REGISTER_LOGGER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_LOGGER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Dumper +#define INDEX_FACTORY_REGISTER_DUMPER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexDumper, __IMPL__, ##__VA_ARGS__) + +//! Register Index Dumper +#define INDEX_FACTORY_REGISTER_DUMPER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_DUMPER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Storage +#define INDEX_FACTORY_REGISTER_STORAGE_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexStorage, __IMPL__, ##__VA_ARGS__) + +//! Register Index Storage +#define INDEX_FACTORY_REGISTER_STORAGE(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_STORAGE_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Converter +#define INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexConverter, __IMPL__, ##__VA_ARGS__) + +//! Register Index Converter +#define INDEX_FACTORY_REGISTER_CONVERTER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Reformer +#define INDEX_FACTORY_REGISTER_REFORMER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexReformer, __IMPL__, ##__VA_ARGS__) + +//! Register Index Reformer +#define INDEX_FACTORY_REGISTER_REFORMER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_REFORMER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Trainer +#define INDEX_FACTORY_REGISTER_TRAINER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexTrainer, __IMPL__, ##__VA_ARGS__) + +//! Register Index Trainer +#define INDEX_FACTORY_REGISTER_TRAINER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_TRAINER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Builder +#define INDEX_FACTORY_REGISTER_BUILDER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexBuilder, __IMPL__, ##__VA_ARGS__) + +//! Register Index Builder +#define INDEX_FACTORY_REGISTER_BUILDER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_BUILDER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Searcher +#define INDEX_FACTORY_REGISTER_SEARCHER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexSearcher, __IMPL__, ##__VA_ARGS__) + +//! Register Index Searcher +#define INDEX_FACTORY_REGISTER_SEARCHER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_SEARCHER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Streamer +#define INDEX_FACTORY_REGISTER_STREAMER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexStreamer, __IMPL__, ##__VA_ARGS__) + +//! Register Index Streamer +#define INDEX_FACTORY_REGISTER_STREAMER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_STREAMER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Reducer +#define INDEX_FACTORY_REGISTER_REDUCER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexReducer, __IMPL__, ##__VA_ARGS__) + +//! Register Index Reducer +#define INDEX_FACTORY_REGISTER_REDUCER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_REDUCER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Streamer Reducer +#define INDEX_FACTORY_REGISTER_STREAMER_REDUCER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexStreamerReducer, __IMPL__, \ + ##__VA_ARGS__) + +//! Register Index Streamer Reducer +#define INDEX_FACTORY_REGISTER_STREAMER_REDUCER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_STREAMER_REDUCER_ALIAS(__IMPL__, __IMPL__, \ + ##__VA_ARGS__) + +//! Register Index Cluster +#define INDEX_FACTORY_REGISTER_CLUSTER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexCluster, __IMPL__, ##__VA_ARGS__) + +//! Register Index Cluster +#define INDEX_FACTORY_REGISTER_CLUSTER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_CLUSTER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +//! Register Index Refiner +#define INDEX_FACTORY_REGISTER_REFINER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, IndexRefiner, __IMPL__, ##__VA_ARGS__) + +//! Register Index Refiner +#define INDEX_FACTORY_REGISTER_REFINER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_REFINER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_features.h b/src/core/framework/index_features.h new file mode 100644 index 0000000..7efce6b --- /dev/null +++ b/src/core/framework/index_features.h @@ -0,0 +1,604 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_meta.h" + +namespace zvec { +namespace core { + +/*! Index Features + */ +struct IndexFeatures { + //! Index Features Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexFeatures(void) {} + + //! Retrieve feature via index + virtual const void *element(size_t i) const = 0; + + //! Retrieve count of elements + virtual size_t count(void) const = 0; + + //! Retrieve dimension + virtual size_t dimension(void) const = 0; + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const = 0; + + //! Test if it is a compacted buffer + virtual bool is_compacted(void) const { + return false; + } + + //! Retrieve pointer of compacted buffer + virtual const void *data(void) const { + return nullptr; + } + + //! Retrieve size of feature + virtual size_t element_size(void) const { + return IndexMeta::ElementSizeof(this->data_type(), this->dimension()); + } + + //! Operator [] + const void *operator[](size_t i) const { + return this->element(i); + } + + //! Test if matchs the meta + bool is_matched(const IndexMeta &meta) const { + return (meta.data_type() == this->data_type() && + meta.dimension() == this->dimension() && + meta.element_size() == this->element_size()); + } +}; + +/*! Coherent Index Features + */ +class CoherentIndexFeatures : public IndexFeatures { + public: + //! Coherent Index Features Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + CoherentIndexFeatures(void) + : features_buffer_(nullptr), + features_count_(0), + feature_size_(0), + feature_dimension_(0), + data_type_(IndexMeta::DataType::DT_UNDEFINED) {} + + //! Constructor + CoherentIndexFeatures(const IndexMeta &meta) + : features_buffer_(nullptr), + features_count_(0), + feature_size_(meta.element_size()), + feature_dimension_(meta.dimension()), + data_type_(meta.data_type()) {} + + //! Constructor + CoherentIndexFeatures(const IndexMeta &meta, const void *buf, size_t len) + : features_buffer_(buf), + features_count_(len / meta.element_size()), + feature_size_(meta.element_size()), + feature_dimension_(meta.dimension()), + data_type_(meta.data_type()) {} + + //! Mount features + void mount(const IndexMeta &meta, const void *buf, size_t len) { + features_buffer_ = buf; + data_type_ = meta.data_type(); + feature_size_ = meta.element_size(); + feature_dimension_ = meta.dimension(); + features_count_ = len / feature_size_; + } + + //! Mount features + void mount(const void *buf, size_t len) { + features_buffer_ = buf; + features_count_ = len / feature_size_; + } + + //! Retrieve count of elements + virtual size_t count(void) const { + return features_count_; + } + + //! Retrieve dimension + virtual size_t dimension(void) const { + return feature_dimension_; + } + + //! Retrieve feature via index + virtual const void *element(size_t i) const { + return (reinterpret_cast(features_buffer_) + + feature_size_ * i); + } + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const { + return data_type_; + } + + //! Test if it is a compacted buffer + virtual bool is_compacted(void) const { + return true; + } + + //! Retrieve pointer of compacted buffer + virtual const void *data(void) const { + return features_buffer_; + } + + //! Retrieve size of feature + virtual size_t element_size(void) const { + return feature_size_; + } + + private: + const void *features_buffer_; + size_t features_count_; + size_t feature_size_; + size_t feature_dimension_; + IndexMeta::DataType data_type_; +}; + +/*! Flexible Index Features + */ +class FlexibleIndexFeatures : public IndexFeatures { + public: + //! Flexible Index Features Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + FlexibleIndexFeatures(void) + : features_(nullptr), + features_count_(0), + feature_size_(0), + feature_dimension_(0), + data_type_(IndexMeta::DataType::DT_UNDEFINED) {} + + //! Constructor + FlexibleIndexFeatures(const IndexMeta &meta) + : features_(nullptr), + features_count_(0), + feature_size_(meta.element_size()), + feature_dimension_(meta.dimension()), + data_type_(meta.data_type()) {} + + //! Constructor + FlexibleIndexFeatures(const IndexMeta &meta, const void *const *feats, + size_t feats_count) + : features_(feats), + features_count_(feats_count), + feature_size_(meta.element_size()), + feature_dimension_(meta.dimension()), + data_type_(meta.data_type()) {} + + //! Mount features + void mount(const IndexMeta &meta, const void *const *feats, + size_t feats_count) { + features_ = feats; + features_count_ = feats_count; + data_type_ = meta.data_type(); + feature_size_ = meta.element_size(); + feature_dimension_ = meta.dimension(); + } + + //! Mount features + void mount(const void *const *feats, size_t feats_count) { + features_ = feats; + features_count_ = feats_count; + } + + //! Retrieve count of elements + virtual size_t count(void) const { + return features_count_; + } + + //! Retrieve dimension + virtual size_t dimension(void) const { + return feature_dimension_; + } + + //! Retrieve feature via index + virtual const void *element(size_t i) const { + return *(features_ + i); + } + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const { + return data_type_; + } + + //! Retrieve size of feature + virtual size_t element_size(void) const { + return feature_size_; + } + + private: + const void *const *features_; + size_t features_count_; + size_t feature_size_; + size_t feature_dimension_; + IndexMeta::DataType data_type_; +}; + +/*! Gap Index Features + */ +class GapIndexFeatures : public IndexFeatures { + public: + //! Gap Index Features Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + GapIndexFeatures(const IndexMeta &meta) + : features_(), + bucket_limit_(0), + features_count_(0), + feature_size_(meta.element_size()), + feature_dimension_(meta.dimension()), + data_type_(meta.data_type()) { + if (feature_size_ >= 1024 * 1024) { + bucket_limit_ = 64u; + } else { + bucket_limit_ = (1024 * 1024 * 64) / feature_size_; + } + } + + //! Constructor + GapIndexFeatures(const GapIndexFeatures &rhs) + : features_(rhs.features_), + bucket_limit_(rhs.bucket_limit_), + features_count_(rhs.features_count_), + feature_size_(rhs.feature_size_), + feature_dimension_(rhs.feature_dimension_), + data_type_(rhs.data_type_) {} + + //! Constructor + GapIndexFeatures(GapIndexFeatures &&rhs) + : features_(std::move(rhs.features_)), + bucket_limit_(rhs.bucket_limit_), + features_count_(rhs.features_count_), + feature_size_(rhs.feature_size_), + feature_dimension_(rhs.feature_dimension_), + data_type_(rhs.data_type_) {} + + //! Assignment + GapIndexFeatures &operator=(const GapIndexFeatures &rhs) { + features_ = rhs.features_; + bucket_limit_ = rhs.bucket_limit_; + features_count_ = rhs.features_count_; + feature_size_ = rhs.feature_size_; + feature_dimension_ = rhs.feature_dimension_; + data_type_ = rhs.data_type_; + return *this; + } + + //! Assignment + GapIndexFeatures &operator=(GapIndexFeatures &&rhs) { + features_ = std::move(rhs.features_); + bucket_limit_ = rhs.bucket_limit_; + features_count_ = rhs.features_count_; + feature_size_ = rhs.feature_size_; + feature_dimension_ = rhs.feature_dimension_; + data_type_ = rhs.data_type_; + return *this; + } + + //! Append a feature + void emplace(const void *feat) { + if (features_count_ % bucket_limit_ == 0) { + std::string bucket; + bucket.reserve(bucket_limit_ * feature_size_); + bucket.assign(reinterpret_cast(feat), feature_size_); + features_.push_back(std::move(bucket)); + } else { + features_[features_count_ / bucket_limit_].append( + reinterpret_cast(feat), feature_size_); + } + ++features_count_; + } + + //! Replace a feature + void replace(size_t i, const void *feat) { + std::memcpy(const_cast(features_[i / bucket_limit_].data()) + + feature_size_ * (i % bucket_limit_), + feat, feature_size_); + } + + //! Clear the features + void clear(void) { + features_.clear(); + features_count_ = 0; + } + + //! Retrieve feature via index + void *at(size_t i) { + return (const_cast(features_[i / bucket_limit_].data()) + + feature_size_ * (i % bucket_limit_)); + } + + //! Retrieve feature via index + const void *at(size_t i) const { + return (features_[i / bucket_limit_].data() + + feature_size_ * (i % bucket_limit_)); + } + + //! Retrieve count of elements + virtual size_t count(void) const { + return features_count_; + } + + //! Retrieve dimension + virtual size_t dimension(void) const { + return feature_dimension_; + } + + //! Retrieve feature via index + virtual const void *element(size_t i) const { + return this->at(i); + } + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const { + return data_type_; + } + + //! Test if it is a compacted buffer + virtual bool is_compacted(void) const { + return (features_.size() == 1u); + } + + //! Retrieve pointer of compacted buffer + virtual const void *data(void) const { + return (features_.size() == 1u ? features_.front().data() : nullptr); + } + + //! Retrieve size of feature + virtual size_t element_size(void) const { + return feature_size_; + } + + private: + //! Disable them + GapIndexFeatures(void) = delete; + + //! Members + std::vector features_; + size_t bucket_limit_; + size_t features_count_; + size_t feature_size_; + size_t feature_dimension_; + IndexMeta::DataType data_type_; +}; + +/*! Compact Index Features + */ +class CompactIndexFeatures : public IndexFeatures { + public: + //! Compact Index Features Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + CompactIndexFeatures(const IndexMeta &meta) + : features_(), + feature_size_(meta.element_size()), + feature_dimension_(meta.dimension()), + data_type_(meta.data_type()) {} + + //! Constructor + CompactIndexFeatures(const CompactIndexFeatures &rhs) + : features_(rhs.features_), + feature_size_(rhs.feature_size_), + feature_dimension_(rhs.feature_dimension_), + data_type_(rhs.data_type_) {} + + //! Constructor + CompactIndexFeatures(CompactIndexFeatures &&rhs) + : features_(std::move(rhs.features_)), + feature_size_(rhs.feature_size_), + feature_dimension_(rhs.feature_dimension_), + data_type_(rhs.data_type_) {} + + //! Assignment + CompactIndexFeatures &operator=(const CompactIndexFeatures &rhs) { + features_ = rhs.features_; + feature_size_ = rhs.feature_size_; + feature_dimension_ = rhs.feature_dimension_; + data_type_ = rhs.data_type_; + return *this; + } + + //! Assignment + CompactIndexFeatures &operator=(CompactIndexFeatures &&rhs) { + features_ = std::move(rhs.features_); + feature_size_ = rhs.feature_size_; + feature_dimension_ = rhs.feature_dimension_; + data_type_ = rhs.data_type_; + return *this; + } + + //! Append a feature + void emplace(const void *feat) { + features_.append(reinterpret_cast(feat), feature_size_); + } + + //! Replace a feature + void replace(size_t i, const void *feat) { + std::memcpy(const_cast(features_.data()) + feature_size_ * i, feat, + feature_size_); + } + + //! Resize the container + void resize(size_t n) { + features_.resize(feature_size_ * n); + } + + //! Reserve the container + void reserve(size_t n) { + features_.reserve(feature_size_ * n); + } + + //! Clear the features + void clear(void) { + features_.clear(); + } + + //! Retrieve feature via index + void *at(size_t i) { + return (const_cast(features_.data()) + feature_size_ * i); + } + + //! Retrieve feature via index + const void *at(size_t i) const { + return (features_.data() + feature_size_ * i); + } + + //! Retrieve count of elements + virtual size_t count(void) const { + return (features_.size() / feature_size_); + } + + //! Retrieve dimension + virtual size_t dimension(void) const { + return feature_dimension_; + } + + //! Retrieve feature via index + virtual const void *element(size_t i) const { + return this->at(i); + } + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const { + return data_type_; + } + + //! Test if it is a compacted buffer + virtual bool is_compacted(void) const { + return true; + } + + //! Retrieve pointer of compacted buffer + virtual const void *data(void) const { + return features_.data(); + } + + //! Retrieve size of feature + virtual size_t element_size(void) const { + return feature_size_; + } + + private: + //! Disable them + CompactIndexFeatures(void) = delete; + + //! Members + std::string features_; + size_t feature_size_; + size_t feature_dimension_; + IndexMeta::DataType data_type_; +}; + +/*! Sample Index Features + */ +template +class SampleIndexFeatures : public TBase { + public: + //! Sample Index Features Pointer + typedef std::shared_ptr> Pointer; + + //! Constructor + SampleIndexFeatures(const IndexMeta &meta, size_t cnt) + : TBase(meta), samples_(std::max(cnt, 1u)), total_(0), mt_() {} + + //! Constructor + SampleIndexFeatures(const SampleIndexFeatures &rhs) + : TBase(rhs), samples_(rhs.samples_), total_(rhs.total_), mt_() {} + + //! Constructor + SampleIndexFeatures(SampleIndexFeatures &&rhs) + : TBase(std::move(rhs)), + samples_(rhs.samples_), + total_(rhs.total_), + mt_() {} + + //! Assignment + SampleIndexFeatures &operator=(const SampleIndexFeatures &rhs) { + TBase::operator=(static_cast(rhs)); + samples_ = rhs.samples_; + total_ = rhs.total_; + return *this; + } + + //! Assignment + SampleIndexFeatures &operator=(SampleIndexFeatures &&rhs) { + TBase::operator=(std::move(static_cast(rhs))); + samples_ = rhs.samples_; + total_ = rhs.total_; + return *this; + } + + //! Retrieve count of samples + size_t samples(void) const { + return samples_; + } + + //! Retrieve count of total + size_t total(void) const { + return total_; + } + + //! Append a feature + void emplace(const void *feat) { + if (TBase::count() >= samples_) { + std::uniform_int_distribution dt(0, total_); + size_t i = dt(mt_); + + if (i < samples_) { + TBase::replace(i, feat); + } + } else { + TBase::emplace(feat); + } + ++total_; + } + + //! Clear the features + void clear(void) { + TBase::clear(); + total_ = 0; + } + + private: + //! Disable them + SampleIndexFeatures(void) = delete; + + //! Members + size_t samples_; + size_t total_; + std::mt19937 mt_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_filter.h b/src/core/framework/index_filter.h new file mode 100644 index 0000000..c345a84 --- /dev/null +++ b/src/core/framework/index_filter.h @@ -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. +#pragma once + +#include + +namespace zvec { +namespace core { + +/*! Index Filter + */ +class IndexFilter { + public: + //! Constructor + IndexFilter(void) {} + + //! Constructor + IndexFilter(const IndexFilter &rhs) : filter_(rhs.filter_) {} + + //! Constructor + IndexFilter(IndexFilter &&rhs) + : filter_(std::forward(rhs.filter_)) {} + + //! Copy assignment operator + IndexFilter &operator=(const IndexFilter &rhs) { + filter_ = rhs.filter_; + return *this; + } + + //! Copy assignment operator + IndexFilter &operator=(IndexFilter &&rhs) { + filter_ = std::forward(rhs.filter_); + return *this; + } + + //! Function call + bool operator()(uint64_t key) const { + return (filter_ ? filter_(key) : false); + } + + //! Set the filter function + template + void set(T &&func) { + filter_ = std::forward(func); + } + + //! Reset the filter function + void reset(void) { + filter_ = nullptr; + } + + //! Test if the function is valid + bool is_valid(void) const { + return (!!filter_); + } + + private: + //! Members + std::function filter_{}; +}; + +} // 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..c68a9a1 --- /dev/null +++ b/src/core/framework/index_flow.cc @@ -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 "index_flow.h" +#include +#include "index_helper.h" + +//! 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_flow.h b/src/core/framework/index_flow.h new file mode 100644 index 0000000..19d132b --- /dev/null +++ b/src/core/framework/index_flow.h @@ -0,0 +1,646 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_factory.h" + +namespace zvec { +namespace core { + +/*! Index Flow + */ +class IndexFlow { + public: + /*! Index Flow Context + */ + class Context { + public: + //! Index Flow Pointer + typedef std::unique_ptr Pointer; + + //! Index Flow UPointer + typedef std::unique_ptr UPointer; + + //! Retrieve searcher context + IndexSearcher::Context::Pointer &searcher_context(void) { + return searcher_context_; + } + + //! Set topk of search result + void set_topk(uint32_t topk) { + return searcher_context_->set_topk(topk); + } + + //! Retrieve search results + const IndexDocumentList &result(void) const { + return searcher_context_->result(); + } + + //! Retrieve search result with index + const IndexDocumentList &result(size_t index) const { + return searcher_context_->result(index); + } + + //! Set the filter of context + template + void set_filter(T &&func) { + searcher_context_->set_filter(std::forward(func)); + } + + //! Reset the filter of context + void reset_filter(void) { + searcher_context_->reset_filter(); + } + + //! Set mode of debug + void set_debug_mode(bool enable) { + searcher_context_->set_debug_mode(enable); + } + + //! Update the parameters of context + int update(const ailego::Params ¶ms) { + return searcher_context_->update(params); + } + + //! Retrieve debug information + std::string debug_string(void) const { + return searcher_context_->debug_string(); + } + + //! Retrieve magic number + uint32_t magic(void) const { + return searcher_context_->magic(); + } + + //! Retrieve mode of debug + bool debug_mode(void) const { + return searcher_context_->debug_mode(); + } + + //! Retrieve mutable features buffer + std::string *mutable_features(void) { + return &features_; + } + + //! Retrieve features buffer + const std::string &features(void) const { + return features_; + } + + protected: + friend class IndexFlow; + + //! Constructor + Context(IndexSearcher::Context::Pointer &&ctx) + : searcher_context_(std::move(ctx)) {} + + private: + IndexSearcher::Context::Pointer searcher_context_{}; + std::string features_{}; + }; + + //! Constructor + IndexFlow(void) {} + + //! Constructor + IndexFlow(IndexFlow &&rhs) + : storage_(std::move(rhs.storage_)), + reformer_(std::move(rhs.reformer_)), + searcher_(std::move(rhs.searcher_)), + metric_(std::move(rhs.metric_)), + user_reformer_(std::move(rhs.user_reformer_)), + user_searcher_(std::move(rhs.user_searcher_)), + user_metric_name_(std::move(rhs.user_metric_name_)), + user_metric_params_(std::move(rhs.user_metric_params_)) {} + + //! Assignment + IndexFlow &operator=(IndexFlow &&rhs) { + storage_ = std::move(rhs.storage_); + reformer_ = std::move(rhs.reformer_); + searcher_ = std::move(rhs.searcher_); + metric_ = std::move(rhs.metric_); + user_reformer_ = std::move(rhs.user_reformer_); + user_searcher_ = std::move(rhs.user_searcher_); + user_metric_name_ = std::move(rhs.user_metric_name_); + user_metric_params_ = std::move(rhs.user_metric_params_); + return *this; + } + + //! Retrieve index meta + const IndexMeta &meta(void) const { + return meta_; + } + + //! Retrieve index reformer + const IndexReformer::Pointer &reformer(void) const { + return reformer_; + } + + //! Retrieve index searcher + const IndexSearcher::Pointer &searcher(void) const { + return searcher_; + } + + //! Retrieve index metric + const IndexMetric::Pointer &metric(void) const { + return metric_; + } + + //! Set the index storage (user) + int set_storage(const std::string &name, const ailego::Params ¶ms); + + //! Set the index reformer (user) + int set_reformer(const std::string &name, const ailego::Params ¶ms); + + //! Set the index searcher (user) + int set_searcher(const std::string &name, const ailego::Params ¶ms); + + //! Set the index searcher (user) + int set_searcher(IndexSearcher::Pointer searcher); + + //! Set the index metric (user) + int set_metric(const std::string &name, const ailego::Params ¶ms); + + //! Load index + int load(const std::string &path); + + //! Unload index + int unload(void); + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const; + + //! Similarity brute force search + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const; + + //! Similarity search (FP16) + template ::type> + int search_bf(const ailego::Float16 *vec, size_t dim, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search (FP32) + template ::type> + int search_bf(const float *vec, size_t dim, Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search (INT8) + template ::type> + int search_bf(const int8_t *vec, size_t dim, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search (INT4) + template ::type> + int search_bf(const uint8_t *vec, size_t dim, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search (BINARY) + template ::type> + int search_bf(const uint32_t *vec, size_t dim, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search in batch (FP16) + template ::type> + int search_bf(const ailego::Float16 *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity search in batch (FP32) + template ::type> + int search_bf(const float *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity search in batch (INT8) + template ::type> + int search_bf(const int8_t *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity search in batch (INT4) + template ::type> + int search_bf(const uint8_t *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity Search in batch (BINARY) + template ::type> + int search_bf(const uint32_t *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_bf_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity search (FP16) + template ::type> + int search(const ailego::Float16 *vec, size_t dim, + Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search (FP32) + template ::type> + int search(const float *vec, size_t dim, Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search (INT8) + template ::type> + int search(const int8_t *vec, size_t dim, Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search (INT4) + template ::type> + int search(const uint8_t *vec, size_t dim, Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search (BINARY32) + template ::type> + int search(const uint32_t *vec, size_t dim, Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), context); + } + + //! Similarity search in batch (FP16) + template ::type> + int search(const ailego::Float16 *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity search in batch (FP32) + template ::type> + int search(const float *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity search in batch (INT8) + template ::type> + int search(const int8_t *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity search in batch (INT4) + template ::type> + int search(const uint8_t *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Similarity Search in batch (BINARY) + template ::type> + int search(const uint32_t *vec, size_t dim, size_t rows, + Context::Pointer &context) const { + return this->search_impl(vec, IndexQueryMeta(DT, dim), rows, context); + } + + //! Create a flow context + Context::Pointer create_context(void) const { + return Context::Pointer(new Context(searcher_->create_context())); + } + + private: + //! Disable them + IndexFlow(const IndexFlow &) = delete; + IndexFlow &operator=(const IndexFlow &) = delete; + + int load_internal(); + + //! Members + IndexMeta meta_{}; + IndexStorage::Pointer storage_{}; + IndexReformer::Pointer reformer_{}; + IndexSearcher::Pointer searcher_{}; + IndexMetric::Pointer metric_{}; + IndexReformer::Pointer user_reformer_{}; + IndexSearcher::Pointer user_searcher_{}; + std::string user_metric_name_{}; + ailego::Params user_metric_params_{}; +}; + + +/*! Index Sparse Flow + */ +class IndexSparseFlow { + public: + /*! Index Sparse Flow Context + */ + class Context { + public: + //! Index Flow Pointer + typedef std::unique_ptr Pointer; + + //! Index Flow UPointer + typedef std::unique_ptr UPointer; + + //! Retrieve searcher context + IndexSearcher::Context::Pointer &searcher_context(void) { + return searcher_context_; + } + + //! Set topk of search result + void set_topk(uint32_t topk) { + return searcher_context_->set_topk(topk); + } + + //! Retrieve search results + const IndexDocumentList &result(void) const { + return searcher_context_->result(); + } + + //! Retrieve search result with index + const IndexDocumentList &result(size_t index) const { + return searcher_context_->result(index); + } + + //! Set the filter of context + template + void set_filter(T &&func) { + searcher_context_->set_filter(std::forward(func)); + } + + //! Reset the filter of context + void reset_filter(void) { + searcher_context_->reset_filter(); + } + + //! Set mode of debug + void set_debug_mode(bool enable) { + searcher_context_->set_debug_mode(enable); + } + + //! Update the parameters of context + int update(const ailego::Params ¶ms) { + return searcher_context_->update(params); + } + + //! Retrieve debug information + std::string debug_string(void) const { + return searcher_context_->debug_string(); + } + + //! Retrieve magic number + uint32_t magic(void) const { + return searcher_context_->magic(); + } + + //! Retrieve mode of debug + bool debug_mode(void) const { + return searcher_context_->debug_mode(); + } + + //! Retrieve mutable features buffer + std::string *mutable_features(void) { + return &features_; + } + + //! Retrieve features buffer + const std::string &features(void) const { + return features_; + } + + protected: + friend class IndexSparseFlow; + + //! Constructor + Context(IndexSearcher::Context::Pointer &&ctx) + : searcher_context_(std::move(ctx)) {} + + private: + IndexSearcher::Context::Pointer searcher_context_{}; + std::string features_{}; + }; + + //! Constructor + IndexSparseFlow(void) {} + + //! Constructor + IndexSparseFlow(IndexSparseFlow &&rhs) + : storage_(std::move(rhs.storage_)), + reformer_(std::move(rhs.reformer_)), + searcher_(std::move(rhs.searcher_)), + metric_(std::move(rhs.metric_)), + user_reformer_(std::move(rhs.user_reformer_)), + user_searcher_(std::move(rhs.user_searcher_)), + user_metric_name_(std::move(rhs.user_metric_name_)), + user_metric_params_(std::move(rhs.user_metric_params_)) {} + + //! Assignment + IndexSparseFlow &operator=(IndexSparseFlow &&rhs) { + storage_ = std::move(rhs.storage_); + reformer_ = std::move(rhs.reformer_); + searcher_ = std::move(rhs.searcher_); + metric_ = std::move(rhs.metric_); + user_reformer_ = std::move(rhs.user_reformer_); + user_searcher_ = std::move(rhs.user_searcher_); + user_metric_name_ = std::move(rhs.user_metric_name_); + user_metric_params_ = std::move(rhs.user_metric_params_); + return *this; + } + + //! Retrieve index sparse meta + const IndexMeta &meta(void) const { + return meta_; + } + + //! Retrieve index reformer + const IndexReformer::Pointer &reformer(void) const { + return reformer_; + } + + //! Retrieve index searcher + const IndexSearcher::Pointer &searcher(void) const { + return searcher_; + } + + //! Retrieve index metric + const IndexMetric::Pointer &metric(void) const { + return metric_; + } + + //! Set the index storage (user) + int set_storage(const std::string &name, const ailego::Params ¶ms); + + //! Set the index reformer (user) + int set_reformer(const std::string &name, const ailego::Params ¶ms); + + //! Set the index searcher (user) + int set_searcher(const std::string &name, const ailego::Params ¶ms); + + //! Set the index searcher (user) + int set_searcher(IndexSearcher::Pointer searcher); + + //! Set the index metric (user) + int set_metric(const std::string &name, const ailego::Params ¶ms); + + //! Load index + int load(const std::string &path); + + //! Unload index + int unload(void); + + //! 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; + + //! 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; + + //! Similarity brute force search and 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; + + //! 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; + + //! Similarity search (FP16) + template ::type> + int search_bf(const uint32_t sparse_count, const uint32_t *sparse_indices, + const ailego::Float16 *sparse_query, + Context::Pointer &context) const { + return this->search_bf_impl(sparse_count, sparse_indices, sparse_query, + IndexQueryMeta(DT), context); + } + + //! Similarity search (FP32) + template ::type> + int search_bf(const uint32_t sparse_count, const uint32_t *sparse_indices, + const float *sparse_query, Context::Pointer &context) const { + return this->search_bf_impl(sparse_count, sparse_indices, sparse_query, + IndexQueryMeta(DT), context); + } + + //! Similarity search (FP16) + template ::type> + int search(const uint32_t sparse_count, const uint32_t *sparse_indices, + const ailego::Float16 *sparse_query, + Context::Pointer &context) const { + return this->search_impl(sparse_count, sparse_indices, sparse_query, + IndexQueryMeta(DT), context); + } + + //! Similarity search (FP32) + template ::type> + int search(const uint32_t sparse_count, const uint32_t *sparse_indices, + const float *sparse_query, Context::Pointer &context) const { + return this->search_impl(sparse_count, sparse_indices, sparse_query, + IndexQueryMeta(DT), context); + } + + //! Create a flow context + Context::Pointer create_context(void) const { + return Context::Pointer(new Context(searcher_->create_context())); + } + + private: + //! Disable them + IndexSparseFlow(const IndexSparseFlow &) = delete; + IndexSparseFlow &operator=(const IndexSparseFlow &) = delete; + + int load_internal(); + + //! Members + IndexMeta meta_{}; + IndexStorage::Pointer storage_{}; + IndexReformer::Pointer reformer_{}; + IndexSearcher::Pointer searcher_{}; + IndexMetric::Pointer metric_{}; + IndexReformer::Pointer user_reformer_{}; + IndexSearcher::Pointer user_searcher_{}; + std::string user_metric_name_{}; + ailego::Params user_metric_params_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_format.h b/src/core/framework/index_format.h new file mode 100644 index 0000000..0f5e996 --- /dev/null +++ b/src/core/framework/index_format.h @@ -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. +#pragma once + +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +/*! Index Format + */ +struct IndexFormat { + /*! Version Number + */ + enum { FORMAT_VERSION = 0x0002 }; + + /*! Index Format Meta Header + */ + struct MetaHeader { + uint32_t header_crc; + uint16_t reserved1_; + uint16_t version; + uint32_t revision; + uint32_t magic; + uint16_t meta_header_size; + uint16_t meta_footer_size; + uint32_t meta_footer_offset; + uint32_t content_offset; + uint32_t reserved2_; + uint64_t setup_time; + uint64_t reserved3_[3]; + }; + + static_assert(sizeof(MetaHeader) % 32 == 0, + "MetaHeader must be aligned with 32 bytes"); + + /*! Index Format Meta Footer + */ + struct MetaFooter { + uint32_t footer_crc; + uint32_t segments_meta_crc; + uint32_t content_crc; + uint32_t segment_count; + uint32_t segments_meta_size; + uint32_t reserved1_; + uint64_t content_size; + uint64_t content_padding_size; + uint64_t check_point; + uint64_t update_time; + uint64_t reserved2_[8]; + uint64_t total_size; + }; + + static_assert(sizeof(MetaFooter) % 32 == 0, + "MetaFooter must be aligned with 32 bytes"); + + /*! Index Format Segment Meta + */ + struct SegmentMeta { + uint32_t segment_id_offset; + uint32_t data_crc; + uint64_t data_index; + uint64_t data_size; + uint64_t padding_size; + }; + + static_assert(sizeof(SegmentMeta) % 32 == 0, + "SegmentMeta must be aligned with 32 bytes"); + + /*! Index Format Segment Meta Buffer + */ + class SegmentMetaBuffer { + public: + //! Constructor + SegmentMetaBuffer(uint32_t count) : capacity_(count) { + buffer_.clear(); + buffer_.resize(sizeof(SegmentMeta) * capacity_); + } + + //! Append a segment meta into buffer + bool append(const std::string &id, size_t data_size, size_t padding_size, + uint32_t data_crc) { + if (count_ >= capacity_) { + return false; + } + SegmentMeta *meta = (SegmentMeta *)buffer_.data() + count_; + meta->segment_id_offset = static_cast(buffer_.size()); + meta->data_index = offset_; + meta->data_size = data_size; + meta->data_crc = data_crc; + meta->padding_size = padding_size; + buffer_.append(id.c_str(), std::strlen(id.c_str()) + 1); + count_ += 1; + offset_ += data_size + padding_size; + return true; + } + + //! Resize the buffer + void resize(size_t val) { + buffer_.resize(val); + } + + //! Retrieve pointer of data + const void *data(void) const { + return buffer_.data(); + } + + //! Retrieve size of data + size_t size(void) const { + return buffer_.size(); + } + + //! Retrieve crc of buffer + uint32_t crc(void) const { + return ailego::Crc32c::Hash(buffer_.data(), buffer_.size(), 0); + } + + private: + //! Disable them + SegmentMetaBuffer(void) = delete; + + //! Members + std::string buffer_{}; + size_t offset_{0u}; + uint32_t capacity_{0u}; + uint32_t count_{0u}; + }; + + //! Setup meta header structure + static void SetupMetaHeader(MetaHeader *header, uint32_t footer_offset, + uint32_t content_offset) { + memset(header, 0, sizeof(MetaHeader)); + header->version = IndexFormat::FORMAT_VERSION; + header->revision = 0; + header->magic = std::random_device()(); + header->meta_header_size = sizeof(MetaHeader); + header->meta_footer_size = sizeof(MetaFooter); + header->meta_footer_offset = footer_offset; + header->content_offset = content_offset; + header->setup_time = ailego::Realtime::Seconds(); + header->header_crc = ailego::Crc32c::Hash(header, sizeof(MetaHeader), 0); + } + + //! Setup meta footer structure + static void SetupMetaFooter(MetaFooter *footer) { + memset(footer, 0, sizeof(MetaFooter)); + } + + //! Update meta footer structure + static void UpdateMetaFooter(MetaFooter *footer, uint64_t check_point) { + if (check_point != 0) { + footer->check_point = check_point; + } + footer->update_time = ailego::Realtime::Seconds(); + footer->footer_crc = 0; + footer->footer_crc = ailego::Crc32c::Hash(footer, sizeof(MetaFooter), 0); + } +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_framework.h b/src/core/framework/index_framework.h new file mode 100644 index 0000000..d4a12bc --- /dev/null +++ b/src/core/framework/index_framework.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 "index_builder.h" +#include "index_bundle.h" +#include "index_cluster.h" +#include "index_converter.h" +#include "index_factory.h" +#include "index_filter.h" +#include "index_flow.h" +#include "index_helper.h" +#include "index_holder.h" +#include "index_mapping.h" +#include "index_memory.h" +#include "index_meta.h" +#include "index_plugin.h" +#include "index_runner.h" +#include "index_searcher.h" +#include "index_trainer.h" diff --git a/src/core/framework/index_groupby.h b/src/core/framework/index_groupby.h new file mode 100644 index 0000000..51da06f --- /dev/null +++ b/src/core/framework/index_groupby.h @@ -0,0 +1,54 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Index GroupBy + */ +class IndexGroupBy { + public: + //! Function call + std::string operator()(uint64_t key) const { + return (group_by_ ? group_by_(key) : ""); + } + + //! Set the group by function + template + void set(T &&func) { + group_by_ = std::forward(func); + } + + //! Reset the group by function + void reset(void) { + group_by_ = nullptr; + } + + //! Test if the function is valid + bool is_valid(void) const { + return (!!group_by_); + } + + private: + //! Members + std::function group_by_{}; +}; + +} // namespace core +} // namespace zvec \ No newline at end of file diff --git a/src/core/framework/index_helper.cc b/src/core/framework/index_helper.cc new file mode 100644 index 0000000..a19d749 --- /dev/null +++ b/src/core/framework/index_helper.cc @@ -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. +#include "index_helper.h" +#include +#include "index_error.h" + +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(); + const void *data = nullptr; + + if (segment->read(0, &data, len) != len) { + return IndexError_ReadData; + } + 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 + virtual ~FirstPassIterator(void) {} + + //! 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 + virtual ~SecondPassIterator(void) {} + + //! 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_helper.h b/src/core/framework/index_helper.h new file mode 100644 index 0000000..41ae1f0 --- /dev/null +++ b/src/core/framework/index_helper.h @@ -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. +#pragma once + +#include "index_dumper.h" +#include "index_holder.h" +#include "index_meta.h" +#include "index_storage.h" + +namespace zvec { +namespace core { + +/*! Index Helper + */ +struct IndexHelper { + //! Serialize meta information to dumper + static int SerializeToDumper(const IndexMeta &mt, IndexDumper *dumper, + const std::string &key); + + //! Serialize meta information to storage + static int SerializeToStorage(const IndexMeta &mt, IndexStorage *storage, + const std::string &key); + + //! Derialize meta information from storage + static int DeserializeFromStorage(IndexStorage *storage, + const std::string &key, IndexMeta *out); + + //! Serialize meta information to dumper + static int SerializeToDumper(const IndexMeta &mt, IndexDumper *dumper) { + return SerializeToDumper(mt, dumper, "IndexMeta"); + } + + //! Serialize meta information to storage + static int SerializeToStorage(const IndexMeta &mt, IndexStorage *storage) { + return SerializeToStorage(mt, storage, "IndexMeta"); + } + + //! Derialize meta information from storage + static int DeserializeFromStorage(IndexStorage *storage, IndexMeta *out) { + return DeserializeFromStorage(storage, "IndexMeta", out); + } + + //! Create a proxy holder that can be traversed twice. + static IndexHolder::Pointer MakeTwoPassHolder(IndexHolder::Pointer holder); +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_holder.h b/src/core/framework/index_holder.h new file mode 100644 index 0000000..4dff9e3 --- /dev/null +++ b/src/core/framework/index_holder.h @@ -0,0 +1,1760 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "ailego/container/vector.h" +#include "index_features.h" +#include "index_meta.h" + +namespace zvec { +namespace core { + +/*! Index Holder + */ +struct IndexHolder { + //! Index Holder Pointer + typedef std::shared_ptr Pointer; + + /*! Index Holder Iterator + */ + struct Iterator { + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Destructor + virtual ~Iterator(void) {} + + //! Retrieve pointer of data + virtual const void *data(void) const = 0; + + //! Test if the iterator is valid + virtual bool is_valid(void) const = 0; + + //! Retrieve primary key + virtual uint64_t key(void) const = 0; + + //! Next iterator + virtual void next(void) = 0; + }; + + //! Destructor + virtual ~IndexHolder(void) {} + + //! Retrieve count of elements in holder (-1 indicates unknown) + virtual size_t count(void) const = 0; + + //! Retrieve dimension + virtual size_t dimension(void) const = 0; + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const = 0; + + //! Retrieve element size in bytes + virtual size_t element_size(void) const = 0; + + //! Retrieve if it can multi-pass + virtual bool multipass(void) const = 0; + + //! Create a new iterator + virtual Iterator::Pointer create_iterator(void) = 0; + + //! Test if matchs the meta + bool is_matched(const IndexMeta &meta) const { + return (this->data_type() == meta.data_type() && + this->dimension() == meta.dimension() && + this->element_size() == meta.element_size()); + } +}; + +/*! Index Hybrid Holder + */ +struct IndexHybridHolder : public IndexHolder { + //! Index Holder Pointer + typedef std::shared_ptr Pointer; + + /*! Index Holder Iterator + */ + struct Iterator : public IndexHolder::Iterator { + //! Index Holder Pointer + typedef std::unique_ptr Pointer; + + //! Destructor + virtual ~Iterator(void) {} + + //! Retrieve pointer of data + virtual const void *data(void) const = 0; + + //! Test if the iterator is valid + virtual bool is_valid(void) const = 0; + + //! Retrieve primary key + virtual uint64_t key(void) const = 0; + + //! Retrieve sparse count + virtual uint32_t sparse_count() const = 0; + + //! Retrieve sparse indicies + virtual const uint32_t *sparse_indices() const = 0; + + //! Retrieve sparse data + virtual const void *sparse_data() const = 0; + + //! Next iterator + virtual void next(void) = 0; + }; + + //! Destructor + virtual ~IndexHybridHolder(void) {} + + //! Retrieve sparse count summing up over all the docs + virtual size_t total_sparse_count(void) const = 0; + + //! Create a new hybrid iterator + virtual Iterator::Pointer create_hybrid_iterator(void) = 0; +}; + +/*! Index Sparse Holder + */ +struct IndexSparseHolder { + //! Index Sparse Holder Pointer + typedef std::shared_ptr Pointer; + + /*! Index Holder Iterator + */ + struct Iterator { + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Destructor + virtual ~Iterator(void) {} + + //! Test if the iterator is valid + virtual bool is_valid(void) const = 0; + + //! Retrieve primary key + virtual uint64_t key(void) const = 0; + + //! Retrieve sparse count + virtual uint32_t sparse_count() const = 0; + + //! Retrieve sparse indicies + virtual const uint32_t *sparse_indices() const = 0; + + //! Retrieve sparse data + virtual const void *sparse_data() const = 0; + + //! Next iterator + virtual void next(void) = 0; + }; + + //! Destructor + virtual ~IndexSparseHolder(void) {} + + //! Retrieve count of elements in holder (-1 indicates unknown) + virtual size_t count(void) const = 0; + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const = 0; + + //! Retrieve if it can multi-pass + virtual bool multipass(void) const = 0; + + //! Create a new iterator + virtual Iterator::Pointer create_iterator(void) = 0; + + //! Test if matchs the meta + bool is_matched(const IndexMeta &meta) const { + return (this->data_type() == meta.data_type()); + } + + //! Retrieve sparse count summing up over all the docs for reserving space + virtual size_t total_sparse_count(void) const = 0; +}; + +/*! One-Pass Numerical Index Holder + */ +template +class OnePassNumericalIndexHolder : public IndexHolder { + public: + /*! One-Pass Index Holder Iterator + */ + class Iterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(OnePassNumericalIndexHolder *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + virtual ~Iterator(void) {} + + //! 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: + OnePassNumericalIndexHolder *holder_{nullptr}; + typename std::list>>:: + iterator features_iter_{}; + }; + + //! Constructor + OnePassNumericalIndexHolder(size_t dim) : dimension_(dim) {} + + //! 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 IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return dimension_ * sizeof(T); + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return false; + } + + //! Create a new iterator + IndexHolder::Iterator::Pointer create_iterator(void) override { + return IndexHolder::Iterator::Pointer( + new OnePassNumericalIndexHolder::Iterator(this)); + } + + //! Append an element into holder + bool emplace(uint64_t key, const ailego::NumericalVector &vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, vec); + return true; + } + + //! Append an element into holder + bool emplace(uint64_t key, ailego::NumericalVector &&vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, std::move(vec)); + return true; + } + + private: + //! Disable them + OnePassNumericalIndexHolder(void) = delete; + + //! Members + size_t dimension_{0}; + std::list>> features_; +}; + +/*! Multi-Pass Numerical Index Holder + */ +template +class MultiPassNumericalIndexHolder : public IndexHolder { + public: + /*! Multi-Pass Index Holder Iterator + */ + class Iterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(MultiPassNumericalIndexHolder *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + virtual ~Iterator(void) {} + + //! 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 { + ++features_iter_; + } + + private: + MultiPassNumericalIndexHolder *holder_{nullptr}; + typename std::vector>>:: + iterator features_iter_{}; + }; + + //! Constructor + MultiPassNumericalIndexHolder(size_t dim) : dimension_(dim) {} + + //! 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 IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return dimension_ * sizeof(T); + } + + //! 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 MultiPassNumericalIndexHolder::Iterator(this)); + } + + //! Append an element into holder + bool emplace(uint64_t key, const ailego::NumericalVector &vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, vec); + return true; + } + + //! Append an element into holder + bool emplace(uint64_t key, ailego::NumericalVector &&vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, std::move(vec)); + return true; + } + + //! Request a change in capacity + void reserve(size_t size) { + features_.reserve(size); + } + + private: + //! Disable them + MultiPassNumericalIndexHolder(void) = delete; + + //! Members + size_t dimension_{0}; + std::vector>> features_; +}; + +/*! One-Pass Binary Index Holder + */ +template +class OnePassBinaryIndexHolder : public IndexHolder { + public: + /*! One-Pass Index Holder Iterator + */ + class Iterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(OnePassBinaryIndexHolder *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + virtual ~Iterator(void) {} + + //! 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: + OnePassBinaryIndexHolder *holder_{nullptr}; + typename std::list>>::iterator + features_iter_{}; + }; + + //! Constructor + OnePassBinaryIndexHolder(size_t dim) : dimension_(dim) {} + + //! 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 IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return (dimension_ + (sizeof(T) << 3) - 1) / (sizeof(T) << 3) * sizeof(T); + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return false; + } + + //! Create a new iterator + IndexHolder::Iterator::Pointer create_iterator(void) override { + return IndexHolder::Iterator::Pointer( + new OnePassBinaryIndexHolder::Iterator(this)); + } + + //! Append an element into holder + bool emplace(uint64_t key, const ailego::BinaryVector &vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, vec); + return true; + } + + //! Append an element into holder + bool emplace(uint64_t key, ailego::BinaryVector &&vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, std::move(vec)); + return true; + } + + private: + //! Disable them + OnePassBinaryIndexHolder(void) = delete; + + //! Members + size_t dimension_{0}; + std::list>> features_; +}; + +/*! Multi-Pass Binary Index Holder + */ +template +class MultiPassBinaryIndexHolder : public IndexHolder { + public: + /*! Multi-Pass Index Holder Iterator + */ + class Iterator : public IndexHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(MultiPassBinaryIndexHolder *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + virtual ~Iterator(void) {} + + //! 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 { + ++features_iter_; + } + + private: + MultiPassBinaryIndexHolder *holder_{nullptr}; + typename std::vector>>::iterator + features_iter_{}; + }; + + //! Constructor + MultiPassBinaryIndexHolder(size_t dim) : dimension_(dim) {} + + //! 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 IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return (dimension_ + (sizeof(T) << 3) - 1) / (sizeof(T) << 3) * sizeof(T); + } + + //! 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 MultiPassBinaryIndexHolder::Iterator(this)); + } + + //! Append an element into holder + bool emplace(uint64_t key, const ailego::BinaryVector &vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, vec); + return true; + } + + //! Append an element into holder + bool emplace(uint64_t key, ailego::BinaryVector &&vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, std::move(vec)); + return true; + } + + //! Request a change in capacity + void reserve(size_t size) { + features_.reserve(size); + } + + private: + //! Disable them + MultiPassBinaryIndexHolder(void) = delete; + + //! Members + size_t dimension_{0}; + std::vector>> features_; +}; + +/*! One-Pass Index Hybrid Holder + */ +template +class OnePassIndexHybridHolderBase : public IndexHybridHolder { + public: + /*! One-Pass Index Holder Iterator + */ + class Iterator : public IndexHybridHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(OnePassIndexHybridHolderBase *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + virtual ~Iterator(void) {} + + //! 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_++); + } + + //! Retrieve primary key + uint32_t sparse_count() const override { + return features_iter_->second.sparse_count(); + } + + //! Retrieve primary key + const uint32_t *sparse_indices() const override { + return features_iter_->second.sparse_indices(); + } + + //! Retrieve primary key + const void *sparse_data() const override { + return features_iter_->second.sparse_data(); + } + + private: + OnePassIndexHybridHolderBase *holder_{nullptr}; + typename std::list>>::iterator + features_iter_{}; + }; + + //! Constructor + OnePassIndexHybridHolderBase(size_t dim) : dimension_(dim) {} + + //! 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 IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return dimension_ * sizeof(T); + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return false; + } + + //! Create a new iterator + IndexHolder::Iterator::Pointer create_iterator(void) override { + return IndexHolder::Iterator::Pointer( + new OnePassIndexHybridHolderBase::Iterator(this)); + } + + //! Create a new hybrid iterator + IndexHybridHolder::Iterator::Pointer create_hybrid_iterator(void) override { + return IndexHybridHolder::Iterator::Pointer( + new OnePassIndexHybridHolderBase::Iterator(this)); + } + + //! Retrieve sparse count summing up over all the docs + size_t total_sparse_count(void) const override { + return total_sparse_count_; + ; + } + + //! Append an element into holder + bool emplace(uint64_t key, const ailego::HybridVector &vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, vec); + + total_sparse_count_ += vec.sparse_count(); + + return true; + } + + //! Append an element into holder + bool emplace(uint64_t key, ailego::HybridVector &&vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, std::move(vec)); + + total_sparse_count_ += vec.sparse_count(); + + return true; + } + + private: + //! Disable them + OnePassIndexHybridHolderBase(void) = delete; + + //! Members + size_t dimension_{0}; + std::list>> features_; + size_t total_sparse_count_{0}; +}; + +/*! Multi-Pass Index Hybrid Holder Base + */ +template +class MultiPassIndexHybridHolderBase : public IndexHybridHolder { + public: + /*! Multi-Pass Index Holder Iterator + */ + class Iterator : public IndexHybridHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(MultiPassIndexHybridHolderBase *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + virtual ~Iterator(void) {} + + //! 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 { + ++features_iter_; + } + + //! Retrieve primary key + uint32_t sparse_count() const override { + return features_iter_->second.sparse_count(); + } + + //! Retrieve primary key + const uint32_t *sparse_indices() const override { + return features_iter_->second.sparse_indices(); + } + + //! Retrieve primary key + const void *sparse_data() const override { + return features_iter_->second.sparse_data(); + } + + private: + MultiPassIndexHybridHolderBase *holder_{nullptr}; + typename std::vector>>::iterator + features_iter_{}; + }; + + //! Constructor + MultiPassIndexHybridHolderBase(size_t dim) : dimension_(dim) {} + + //! 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 IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve element size in bytes + size_t element_size(void) const override { + return dimension_ * sizeof(T); + } + + //! 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 MultiPassIndexHybridHolderBase::Iterator(this)); + } + + //! Create a new hybrid iterator + IndexHybridHolder::Iterator::Pointer create_hybrid_iterator(void) override { + return IndexHybridHolder::Iterator::Pointer( + new MultiPassIndexHybridHolderBase::Iterator(this)); + } + + //! Retrieve sparse count summing up over all the docs + size_t total_sparse_count(void) const override { + return 0; + } + + //! Append an element into holder + bool emplace(uint64_t key, const ailego::HybridVector &vec) { + if (vec.size() != dimension_) { + return false; + } + + features_.emplace_back(key, vec); + + total_sparse_count_ += vec.sparse_count(); + + return true; + } + + //! Append an element into holder + bool emplace(uint64_t key, ailego::HybridVector &&vec) { + if (vec.size() != dimension_) { + return false; + } + features_.emplace_back(key, std::move(vec)); + + total_sparse_count_ += vec.sparse_count(); + + return true; + } + + //! Request a change in capacity + void reserve(size_t size) { + features_.reserve(size); + } + + private: + //! Disable them + MultiPassIndexHybridHolderBase(void) = delete; + + //! Members + size_t dimension_{0}; + std::vector>> features_; + size_t total_sparse_count_{0}; +}; + +/*! One-Pass Index Sparse Holder + */ +template +class OnePassIndexSparseHolderBase : public IndexSparseHolder { + public: + /*! One-Pass Index Holder Iterator + */ + class Iterator : public IndexSparseHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(OnePassIndexSparseHolderBase *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + virtual ~Iterator(void) {} + + //! 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_++); + } + + //! Retrieve primary key + uint32_t sparse_count() const override { + return features_iter_->second.sparse_count(); + } + + //! Retrieve primary key + const uint32_t *sparse_indices() const override { + return features_iter_->second.sparse_indices(); + } + + //! Retrieve primary key + const void *sparse_data() const override { + return features_iter_->second.sparse_data(); + } + + private: + OnePassIndexSparseHolderBase *holder_{nullptr}; + typename std::list>>::iterator + features_iter_{}; + }; + + //! Constructor + OnePassIndexSparseHolderBase() {} + + //! Retrieve count of elements in holder (-1 indicates unknown) + size_t count(void) const override { + return features_.size(); + } + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return false; + } + + //! Create a new iterator + IndexSparseHolder::Iterator::Pointer create_iterator(void) override { + return IndexSparseHolder::Iterator::Pointer( + new OnePassIndexSparseHolderBase::Iterator(this)); + } + + //! Retrieve sparse count summing up over all the docs + size_t total_sparse_count(void) const override { + return total_sparse_count_; + ; + } + + //! Append an element into holder + bool emplace(uint64_t key, const ailego::SparseVector &vec) { + features_.emplace_back(key, vec); + + total_sparse_count_ += vec.sparse_count(); + + return true; + } + + //! Append an element into holder + bool emplace(uint64_t key, ailego::SparseVector &&vec) { + features_.emplace_back(key, std::move(vec)); + + total_sparse_count_ += vec.sparse_count(); + + return true; + } + + private: + //! Members + std::list>> features_; + size_t total_sparse_count_{0}; +}; + +/*! Multi-Pass Index Sparse Holder Base + */ +template +class MultiPassIndexSparseHolderBase : public IndexSparseHolder { + public: + /*! Multi-Pass Index Holder Iterator + */ + class Iterator : public IndexSparseHolder::Iterator { + public: + //! Index Holder Iterator Pointer + typedef std::unique_ptr Pointer; + + //! Constructor + Iterator(MultiPassIndexSparseHolderBase *owner) : holder_(owner) { + features_iter_ = holder_->features_.begin(); + } + + //! Destructor + virtual ~Iterator(void) {} + + //! 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 { + ++features_iter_; + } + + //! Retrieve primary key + uint32_t sparse_count() const override { + return features_iter_->second.sparse_count(); + } + + //! Retrieve primary key + const uint32_t *sparse_indices() const override { + return features_iter_->second.sparse_indices(); + } + + //! Retrieve primary key + const void *sparse_data() const override { + return features_iter_->second.sparse_data(); + } + + private: + MultiPassIndexSparseHolderBase *holder_{nullptr}; + typename std::vector>>::iterator + features_iter_{}; + }; + + //! Constructor + MultiPassIndexSparseHolderBase() {} + + //! Retrieve count of elements in holder (-1 indicates unknown) + size_t count(void) const override { + return features_.size(); + } + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve if it can multi-pass + bool multipass(void) const override { + return true; + } + + //! Create a new iterator + IndexSparseHolder::Iterator::Pointer create_iterator(void) override { + return IndexSparseHolder::Iterator::Pointer( + new MultiPassIndexSparseHolderBase::Iterator(this)); + } + + //! Retrieve sparse count summing up over all the docs + size_t total_sparse_count(void) const override { + return 0; + } + + //! Append an element into holder + bool emplace(uint64_t key, const ailego::SparseVector &vec) { + features_.emplace_back(key, vec); + + total_sparse_count_ += vec.sparse_count(); + + return true; + } + + //! Append an element into holder + bool emplace(uint64_t key, ailego::SparseVector &&vec) { + features_.emplace_back(key, std::move(vec)); + + total_sparse_count_ += vec.sparse_count(); + + return true; + } + + //! Request a change in capacity + void reserve(size_t size) { + features_.reserve(size); + } + + private: + //! Members + std::vector>> features_; + size_t total_sparse_count_{0}; +}; + +/*! One-Pass Index Holder + */ +template +struct OnePassIndexHolder; + +/*! One-Pass Index Holder (BINARY32) + */ +template <> +struct OnePassIndexHolder + : public OnePassBinaryIndexHolder { + //! Constructor + using OnePassBinaryIndexHolder::OnePassBinaryIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_BINARY32; + } +}; + +/*! One-Pass Index Holder (BINARY64) + */ +template <> +struct OnePassIndexHolder + : public OnePassBinaryIndexHolder { + //! Constructor + using OnePassBinaryIndexHolder::OnePassBinaryIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_BINARY64; + } +}; + +/*! One-Pass Index Holder (FP16) + */ +template <> +struct OnePassIndexHolder + : public OnePassNumericalIndexHolder { + //! Constructor + using OnePassNumericalIndexHolder::OnePassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP16; + } +}; + +/*! One-Pass Index Holder (FP32) + */ +template <> +struct OnePassIndexHolder + : public OnePassNumericalIndexHolder { + //! Constructor + using OnePassNumericalIndexHolder::OnePassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP32; + } +}; + +/*! One-Pass Index Holder (FP64) + */ +template <> +struct OnePassIndexHolder + : public OnePassNumericalIndexHolder { + //! Constructor + using OnePassNumericalIndexHolder::OnePassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP64; + } +}; + +/*! One-Pass Index Holder (INT8) + */ +template <> +struct OnePassIndexHolder + : public OnePassNumericalIndexHolder { + //! Constructor + using OnePassNumericalIndexHolder::OnePassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT8; + } +}; + +/*! One-Pass Index Holder (INT16) + */ +template <> +struct OnePassIndexHolder + : public OnePassNumericalIndexHolder { + //! Constructor + using OnePassNumericalIndexHolder::OnePassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT16; + } +}; + +/*! Multi-Pass Index Holder + */ +template +struct MultiPassIndexHolder; + +/*! Multi-Pass Index Holder (BINARY32) + */ +template <> +struct MultiPassIndexHolder + : public MultiPassBinaryIndexHolder { + //! Constructor + using MultiPassBinaryIndexHolder::MultiPassBinaryIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_BINARY32; + } +}; + +/*! Multi-Pass Index Holder (BINARY64) + */ +template <> +struct MultiPassIndexHolder + : public MultiPassBinaryIndexHolder { + //! Constructor + using MultiPassBinaryIndexHolder::MultiPassBinaryIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_BINARY64; + } +}; + +/*! Multi-Pass Index Holder (FP16) + */ +template <> +struct MultiPassIndexHolder + : public MultiPassNumericalIndexHolder { + //! Constructor + using MultiPassNumericalIndexHolder::MultiPassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP16; + } +}; + +/*! Multi-Pass Index Holder (FP32) + */ +template <> +struct MultiPassIndexHolder + : public MultiPassNumericalIndexHolder { + //! Constructor + using MultiPassNumericalIndexHolder::MultiPassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP32; + } +}; + +/*! Multi-Pass Index Holder (FP64) + */ +template <> +struct MultiPassIndexHolder + : public MultiPassNumericalIndexHolder { + //! Constructor + using MultiPassNumericalIndexHolder::MultiPassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP64; + } +}; + +/*! Multi-Pass Index Holder (INT8) + */ +template <> +struct MultiPassIndexHolder + : public MultiPassNumericalIndexHolder { + //! Constructor + using MultiPassNumericalIndexHolder::MultiPassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT8; + } +}; + +/*! Multi-Pass Index Holder (INT16) + */ +template <> +struct MultiPassIndexHolder + : public MultiPassNumericalIndexHolder { + //! Constructor + using MultiPassNumericalIndexHolder::MultiPassNumericalIndexHolder; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT16; + } +}; + +/*! One-Pass Index Hybrid Holder + */ +template +struct OnePassIndexHybridHolder; + +/*! One-Pass Index Hybrid Holder (FP16) + */ +template <> +struct OnePassIndexHybridHolder + : public OnePassIndexHybridHolderBase { + //! Constructor + using OnePassIndexHybridHolderBase::OnePassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP16; + } +}; + +/*! One-Pass Index Hybrid Holder (FP32) + */ +template <> +struct OnePassIndexHybridHolder + : public OnePassIndexHybridHolderBase { + //! Constructor + using OnePassIndexHybridHolderBase::OnePassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP32; + } +}; + +/*! One-Pass Index Hybrid Holder (FP64) + */ +template <> +struct OnePassIndexHybridHolder + : public OnePassIndexHybridHolderBase { + //! Constructor + using OnePassIndexHybridHolderBase::OnePassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP64; + } +}; + +/*! One-Pass Index Hybrid Holder (INT8) + */ +template <> +struct OnePassIndexHybridHolder + : public OnePassIndexHybridHolderBase { + //! Constructor + using OnePassIndexHybridHolderBase::OnePassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT8; + } +}; + +/*! One-Pass Index Hybrid Holder (INT16) + */ +template <> +struct OnePassIndexHybridHolder + : public OnePassIndexHybridHolderBase { + //! Constructor + using OnePassIndexHybridHolderBase::OnePassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT16; + } +}; + +/*! Multi-Pass Index Hybrid Holder + */ +template +struct MultiPassIndexHybridHolder; + +/*! Multi-Pass Index Hybrid Holder (FP16) + */ +template <> +struct MultiPassIndexHybridHolder + : public MultiPassIndexHybridHolderBase { + //! Constructor + using MultiPassIndexHybridHolderBase::MultiPassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP16; + } +}; + +/*! Multi-Pass Index Hybrid Holder (FP32) + */ +template <> +struct MultiPassIndexHybridHolder + : public MultiPassIndexHybridHolderBase { + //! Constructor + using MultiPassIndexHybridHolderBase::MultiPassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP32; + } +}; + +/*! Multi-Pass Index Hybrid Holder (FP64) + */ +template <> +struct MultiPassIndexHybridHolder + : public MultiPassIndexHybridHolderBase { + //! Constructor + using MultiPassIndexHybridHolderBase::MultiPassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP64; + } +}; + +/*! Multi-Pass Index Hybrid Holder (INT8) + */ +template <> +struct MultiPassIndexHybridHolder + : public MultiPassIndexHybridHolderBase { + //! Constructor + using MultiPassIndexHybridHolderBase::MultiPassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT8; + } +}; + +/*! Multi-Pass Index Hybrid Holder (INT16) + */ +template <> +struct MultiPassIndexHybridHolder + : public MultiPassIndexHybridHolderBase { + //! Constructor + using MultiPassIndexHybridHolderBase::MultiPassIndexHybridHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT16; + } +}; + +/*! One-Pass Index Sparse Holder + */ +template +struct OnePassIndexSparseHolder; + +/*! One-Pass Index Sparse Holder (FP16) + */ +template <> +struct OnePassIndexSparseHolder + : public OnePassIndexSparseHolderBase { + //! Constructor + using OnePassIndexSparseHolderBase::OnePassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP16; + } +}; + +/*! One-Pass Index Sparse Holder (FP32) + */ +template <> +struct OnePassIndexSparseHolder + : public OnePassIndexSparseHolderBase { + //! Constructor + using OnePassIndexSparseHolderBase::OnePassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP32; + } +}; + +/*! One-Pass Index Sparse Holder (FP64) + */ +template <> +struct OnePassIndexSparseHolder + : public OnePassIndexSparseHolderBase { + //! Constructor + using OnePassIndexSparseHolderBase::OnePassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP64; + } +}; + +/*! One-Pass Index Sparse Holder (INT8) + */ +template <> +struct OnePassIndexSparseHolder + : public OnePassIndexSparseHolderBase { + //! Constructor + using OnePassIndexSparseHolderBase::OnePassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT8; + } +}; + +/*! One-Pass Index Sparse Holder (INT16) + */ +template <> +struct OnePassIndexSparseHolder + : public OnePassIndexSparseHolderBase { + //! Constructor + using OnePassIndexSparseHolderBase::OnePassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT16; + } +}; + +/*! Multi-Pass Index Sparse Holder + */ +template +struct MultiPassIndexSparseHolder; + +/*! Multi-Pass Index Sparse Holder (FP16) + */ +template <> +struct MultiPassIndexSparseHolder + : public MultiPassIndexSparseHolderBase { + //! Constructor + using MultiPassIndexSparseHolderBase::MultiPassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP16; + } +}; + +/*! Multi-Pass Index Sparse Holder (FP32) + */ +template <> +struct MultiPassIndexSparseHolder + : public MultiPassIndexSparseHolderBase { + //! Constructor + using MultiPassIndexSparseHolderBase::MultiPassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP32; + } +}; + +/*! Multi-Pass Index Sparse Holder (FP64) + */ +template <> +struct MultiPassIndexSparseHolder + : public MultiPassIndexSparseHolderBase { + //! Constructor + using MultiPassIndexSparseHolderBase::MultiPassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_FP64; + } +}; + +/*! Multi-Pass Index Sparse Holder (INT8) + */ +template <> +struct MultiPassIndexSparseHolder + : public MultiPassIndexSparseHolderBase { + //! Constructor + using MultiPassIndexSparseHolderBase::MultiPassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT8; + } +}; + +/*! Multi-Pass Index Sparse Holder (INT16) + */ +template <> +struct MultiPassIndexSparseHolder + : public MultiPassIndexSparseHolderBase { + //! Constructor + using MultiPassIndexSparseHolderBase::MultiPassIndexSparseHolderBase; + + //! Retrieve type information + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT16; + } +}; + +/*! 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 + virtual ~Iterator(void) {} + + //! Retrieve pointer of data + virtual const void *data(void) const override { + return holder_->element(id_); + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return id_ < holder_->count(); + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + return holder_->key(id_); + } + + //! Next iterator + virtual 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) + virtual size_t count(void) const override { + return features_->count(); + } + + //! Retrieve dimension + virtual size_t dimension(void) const override { + return features_->dimension(); + } + + //! Retrieve type information + virtual IndexMeta::DataType data_type(void) const override { + return features_->data_type(); + } + + //! Retrieve element size in bytes + virtual size_t element_size(void) const override { + return features_->element_size(); + } + + //! Retrieve if it can multi-pass + virtual bool multipass(void) const override { + return true; + } + + //! Create a new iterator + virtual 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_{}; +}; + +} // 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..b54f318 --- /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 "index_logger.h" +#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_logger.h b/src/core/framework/index_logger.h new file mode 100644 index 0000000..58d6db9 --- /dev/null +++ b/src/core/framework/index_logger.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 +#include +#include +#include "ailego/container/params.h" +#include "ailego/logger/logger.h" +#include "index_module.h" + +#define ZVEC_LOG_IMPL(level, format, ...) \ + do { \ + if (zvec::core::IndexLoggerBroker::IsLevelEnabled(level)) { \ + zvec::core::IndexLoggerBroker::Log(level, __FILE__, __LINE__, format, \ + ##__VA_ARGS__); \ + } \ + } while (0) + +//! Log Debug Message +#ifndef LOG_DEBUG +#define LOG_DEBUG(format, ...) \ + ZVEC_LOG_IMPL(zvec::core::IndexLogger::LEVEL_DEBUG, format, ##__VA_ARGS__) +#endif + +//! Log Information Message +#ifndef LOG_INFO +#define LOG_INFO(format, ...) \ + ZVEC_LOG_IMPL(zvec::core::IndexLogger::LEVEL_INFO, format, ##__VA_ARGS__) +#endif + +//! Log Warn Message +#ifndef LOG_WARN +#define LOG_WARN(format, ...) \ + ZVEC_LOG_IMPL(zvec::core::IndexLogger::LEVEL_WARN, format, ##__VA_ARGS__) +#endif + +//! Log Error Message +#ifndef LOG_ERROR +#define LOG_ERROR(format, ...) \ + ZVEC_LOG_IMPL(zvec::core::IndexLogger::LEVEL_ERROR, format, ##__VA_ARGS__) +#endif + +//! Log Fatal Message +#ifndef LOG_FATAL +#define LOG_FATAL(format, ...) \ + ZVEC_LOG_IMPL(zvec::core::IndexLogger::LEVEL_FATAL, format, ##__VA_ARGS__) +#endif + +namespace zvec { +namespace core { + +/*! Index Logger + */ +struct IndexLogger : public IndexModule { + //! Index Logger Pointer + typedef std::shared_ptr Pointer; + + static const int LEVEL_DEBUG; + static const int LEVEL_INFO; + static const int LEVEL_WARN; + static const int LEVEL_ERROR; + static const int LEVEL_FATAL; + + //! Retrieve string of level + static const char *LevelString(int level) { + static const char *info[] = {"DEBUG", " INFO", " WARN", "ERROR", "FATAL"}; + if (level < (int)(sizeof(info) / sizeof(info[0]))) { + return info[level]; + } + return ""; + } + + //! Retrieve symbol of level + static char LevelSymbol(int level) { + static const char info[5] = {'D', 'I', 'W', 'E', 'F'}; + if (level < (int)(sizeof(info) / sizeof(info[0]))) { + return info[level]; + } + return ' '; + } + + //! Destructor + virtual ~IndexLogger(void) {} + + //! Initialize Logger + virtual int init(const zvec::ailego::Params ¶ms) = 0; + + //! Cleanup Logger + virtual int cleanup(void) = 0; + + //! Log Message + virtual void log(int level, const char *file, int line, const char *format, + va_list args) = 0; +}; + +/*! Index Logger Broker + */ +class IndexLoggerBroker { + public: + //! Register Logger + static IndexLogger::Pointer Register(IndexLogger::Pointer logger) { + IndexLogger::Pointer ret = std::move(logger_); + logger_ = std::move(logger); + return ret; + } + + //! Register Logger with init params + static int Register(IndexLogger::Pointer logger, + const ailego::Params ¶ms) { + //! Cleanup the previous, before initizlizing the new one + if (logger_) { + logger_->cleanup(); + } + logger_ = std::move(logger); + return logger_->init(params); + } + + //! Unregister Logger + static void Unregister(void) { + logger_ = nullptr; + } + + //! Set Level of Logger + static void SetLevel(int level) { + logger_level_ = level; + } + + //! Check if log level is enabled + static bool IsLevelEnabled(int level) { + return logger_level_ <= level && logger_; + } + + //! Log Message + __attribute__((format(printf, 4, 5))) static void Log( + int level, const char *file, int line, const char *format, ...) { + if (IsLevelEnabled(level)) { + va_list args; + va_start(args, format); + logger_->log(level, file, line, format, args); + va_end(args); + } + } + + private: + //! Disable them + IndexLoggerBroker(void) = delete; + IndexLoggerBroker(const IndexLoggerBroker &) = delete; + IndexLoggerBroker(IndexLoggerBroker &&) = delete; + + //! Members + static int logger_level_; + static IndexLogger::Pointer logger_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_mapping.cc b/src/core/framework/index_mapping.cc new file mode 100644 index 0000000..470e026 --- /dev/null +++ b/src/core/framework/index_mapping.cc @@ -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. + +#include "index_mapping.h" +#include +#include "index_error.h" +#include "index_logger.h" + +#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, errno %d, %s", errno, std::strerror(errno)); + 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) { + 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, errno %d, %s", path.c_str(), errno, + std::strerror(errno)); + 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, errno %d, %s", path.c_str(), errno, + std::strerror(errno)); + return IndexError_SeekFile; + } + return this->init_index_mapping(mapping_size); +} + +int IndexMapping::create(const std::string &path, size_t segs_size) { + if (!file_.create(path.c_str(), 0)) { + LOG_ERROR("Failed to create file %s, errno %d, %s", path.c_str(), errno, + std::strerror(errno)); + return IndexError_CreateFile; + } + huge_page_ = Ishugetlbfs(path); + if (huge_page_) { + return create_hugepage(segs_size); + } + size_t len = CalcPageAlignedSize(segs_size + 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_.write(&meta_header, sizeof(meta_header)) != sizeof(meta_header)) { + LOG_ERROR("Failed to write file: %s, errno %d, %s", path.c_str(), errno, + std::strerror(errno)); + 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, errno %d, %s", path.c_str(), errno, + std::strerror(errno)); + 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, errno %d, %s", path.c_str(), errno, + std::strerror(errno)); + return IndexError_WriteData; + } + return this->init_index_mapping(len); +} + +int IndexMapping::create_hugepage(size_t segs_size) { + size_t len = CalcPageAlignedSize(segs_size + sizeof(IndexFormat::MetaHeader) + + sizeof(IndexFormat::MetaFooter), + true); + int file_offset = 0; + 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_) { + ailego::File::MemoryUnmap(header_, header_->content_offset); + } + // Reset members + segment_ids_offset_ = 0; + segment_start_ = nullptr; + header_ = nullptr; + 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) { + 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_) { + return IndexError_NoBuffer; + } + + if (!copy_on_write_ && !file_.truncate(index_size_ + size)) { + LOG_ERROR("Failed to truncate file, errno %d, %s", errno, + std::strerror(errno)); + 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; + 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, segment); + 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; + } + Segment *item = &iter->second; + if (!item->data()) { + auto meta = item->meta(); + size_t mapping_size = meta->data_size + meta->padding_size; + size_t offset = meta->data_index + header_->content_offset; + + 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()) { + Segment *item = &iter->second; + + 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) { + Segment *item = &iter->second; + + 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, errno %d, %s", index_size_, + errno, std::strerror(errno)); + return IndexError_TruncateFile; + } + + for (auto iter = segments_.begin(); iter != segments_.end(); ++iter) { + const Segment *item = &iter->second; + 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 = 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, errno %d, %s", + segment_size, errno, std::strerror(errno)); + 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_) { + if (file_.write(0, header_, header_->content_offset) != + header_->content_offset) { + LOG_ERROR("Failed to write segment, size %u, errno %d, %s", + header_->content_offset, errno, std::strerror(errno)); + return IndexError_WriteData; + } + } else { + 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(), 0, len, opts)); + if (!start) { + LOG_ERROR("Failed to map file, errno %d, %s", errno, std::strerror(errno)); + return IndexError_MMapFile; + } + + // Unpack header + header_ = reinterpret_cast(start); + 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; + } + + // 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), + iter); + } + if (sizeof(IndexFormat::SegmentMeta) * footer_->segment_count > + segment_ids_offset_) { + return IndexError_InvalidLength; + } + 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 + return false; +#endif +} + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_mapping.h b/src/core/framework/index_mapping.h new file mode 100644 index 0000000..96bf347 --- /dev/null +++ b/src/core/framework/index_mapping.h @@ -0,0 +1,206 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_format.h" + +namespace zvec { +namespace core { + +/*! Index Mapping + */ +class IndexMapping { + public: + /*! Index Mapping Segment + */ + class Segment { + public: + //! Constructor + Segment(void) {} + + //! Constructor + Segment(IndexFormat::SegmentMeta *segmeta) : meta_(segmeta) {} + + //! Flush the segment + bool flush(void) const { + ailego_false_if_false(this->meta_ && this->data_); + return ailego::File::MemoryFlush( + this->data_, this->meta_->data_size + this->meta_->padding_size); + } + + //! Retrieve mapping address of the segment + void *data(void) const { + return data_; + } + + //! Retrieve pointer of SegmentMeta + IndexFormat::SegmentMeta *meta(void) const { + return meta_; + } + + //! Retrieve dirty flag of the segment + bool dirty(void) const { + return dirty_; + } + + //! Set the segment as dirty + void set_dirty(void) const { + dirty_ = true; + } + + private: + friend class IndexMapping; + + //! Set the mapping address of the segment + void set_data(void *addr) { + data_ = addr; + } + + //! Clear the dirty flag + void reset_dirty(void) const { + dirty_ = false; + } + + private: + //! Members + IndexFormat::SegmentMeta *meta_{nullptr}; + void *data_{nullptr}; + mutable bool dirty_{false}; + }; + + //! Constructor + IndexMapping(void) {} + + //! Constructor + IndexMapping(IndexMapping &&rhs) + : segment_ids_offset_(rhs.segment_ids_offset_), + segment_start_(rhs.segment_start_), + header_(rhs.header_), + footer_(rhs.footer_), + segments_(std::move(rhs.segments_)), + file_(std::move(rhs.file_)) { + rhs.segment_ids_offset_ = 0; + rhs.segment_start_ = nullptr; + rhs.header_ = nullptr; + rhs.footer_ = nullptr; + } + + //! Assignment + IndexMapping &operator=(IndexMapping &&rhs) { + segment_ids_offset_ = rhs.segment_ids_offset_; + segment_start_ = rhs.segment_start_; + header_ = rhs.header_; + footer_ = rhs.footer_; + segments_ = std::move(rhs.segments_); + file_ = std::move(rhs.file_); + rhs.segment_ids_offset_ = 0; + rhs.segment_start_ = nullptr; + rhs.header_ = nullptr; + rhs.footer_ = nullptr; + return *this; + } + + //! Open a index file + int open(const std::string &path, bool cow, bool full_mode); + + //! Create a index file + int create(const std::string &path, size_t segs_size); + + int create_hugepage(size_t segs_size); + + //! Close the index + void close(void); + + //! Refresh meta information (checksum, update time, etc.) + void refresh(uint64_t check_point); + + //! Append a segment into index + int append(const std::string &id, size_t size); + + //! Map a segment by id + Segment *map(const std::string &id, bool warmup, bool lock); + + //! Unmap a segment by id + void unmap(const std::string &id); + + //! Unmap all segments + void unmap_all(void); + + //! Flush the index mapping + int flush(void); + + //! Test if the segment is exist + bool has(const std::string &id) const { + return (segments_.find(id) != segments_.end()); + } + + //! Retrieve count of segments + size_t segment_count(void) const { + return segments_.size(); + } + + //! Retrieve size of index mapping + size_t index_size(void) const { + return index_size_; + } + + //! Retrieve magic number of index + uint32_t magic(void) const { + return (header_ ? header_->magic : 0); + } + + //! Retrieve header information + const IndexFormat::MetaHeader &header(void) const { + return *header_; + } + + //! Retrieve footer information + const IndexFormat::MetaFooter &footer(void) const { + return *footer_; + } + + bool huge_page() const { + return huge_page_; + } + + protected: + //! Initialize index file mapping + int init_index_mapping(size_t len); + + bool Ishugetlbfs(const std::string &path) const; + + private: + //! Disable them + IndexMapping(const IndexMapping &) = delete; + IndexMapping &operator=(const IndexMapping &) = delete; + + //! Members + uint32_t segment_ids_offset_{0}; + IndexFormat::SegmentMeta *segment_start_{nullptr}; + IndexFormat::MetaHeader *header_{nullptr}; + IndexFormat::MetaFooter *footer_{nullptr}; + std::map segments_{}; + size_t index_size_{0u}; + ailego::File file_{}; + bool copy_on_write_{false}; + bool full_mode_{false}; + bool header_dirty_{false}; + bool huge_page_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_memory.h b/src/core/framework/index_memory.h new file mode 100644 index 0000000..5b7c312 --- /dev/null +++ b/src/core/framework/index_memory.h @@ -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. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +/*! Index Memory + */ +class IndexMemory { + public: + /*! Index Memory Block + */ + class Block { + public: + //! Constructor + Block(size_t sz) : buffer_(sz) {} + + //! Constructor + Block(const Block &rhs) : buffer_(rhs.buffer_) {} + + //! Constructor + Block(Block &&rhs) noexcept : buffer_(std::move(rhs.buffer_)) {} + + //! Assignment + Block &operator=(const Block &rhs) { + buffer_ = rhs.buffer_; + return *this; + } + + //! Assignment + Block &operator=(Block &&rhs) { + buffer_ = std::move(rhs.buffer_); + return *this; + } + + //! Retrieve size of buffer + size_t size(void) const { + return buffer_.size(); + } + + //! Append data into the block + size_t append(const void *data, size_t len) { + std::copy(reinterpret_cast(data), + reinterpret_cast(data) + len, + std::back_inserter(buffer_)); + return len; + } + + //! Write data into the block + size_t write(size_t off, const void *data, size_t len) { + size_t region_size = buffer_.size(); + if (off + len > region_size) { + if (off > region_size) { + off = region_size; + } + len = region_size - off; + } + std::copy(reinterpret_cast(data), + reinterpret_cast(data) + len, + buffer_.data() + off); + return len; + } + + //! Fetch data from the storage (with own buffer) + size_t fetch(size_t off, void *buf, size_t len) const { + size_t region_size = buffer_.size(); + if (off + len > region_size) { + if (off > region_size) { + off = region_size; + } + len = region_size - off; + } + std::copy(buffer_.data(), buffer_.data() + len, + reinterpret_cast(buf)); + return len; + } + + //! Read data from the storage (Zero-copy) + size_t read(size_t off, const void **data, size_t len) { + size_t region_size = buffer_.size(); + if (off + len > region_size) { + if (off > region_size) { + off = region_size; + } + len = region_size - off; + } + *data = buffer_.data() + off; + return len; + } + + private: + //! Members + std::vector buffer_{}; + }; + + /*! Index Memory Rope + */ + class Rope { + public: + //! Index Memory Rope Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + Rope(void) {} + + //! Constructor + Rope(const Rope &rhs) : blocks_(rhs.blocks_) {} + + //! Constructor + Rope(Rope &&rhs) : blocks_(std::move(rhs.blocks_)) {} + + //! Assignment + Rope &operator=(const Rope &rhs) { + blocks_ = rhs.blocks_; + return *this; + } + + //! Assignment + Rope &operator=(Rope &&rhs) { + blocks_ = std::move(rhs.blocks_); + return *this; + } + + //! Retrieve the block at index n + Block &operator[](size_t n) { + return blocks_[n]; + } + + //! Retrieve the block at index n + const Block &operator[](size_t n) const { + return blocks_[n]; + } + + //! Retrieve count of blocks + size_t count(void) const { + return blocks_.size(); + } + + //! Retrieve memory size of rope + size_t size(void) const { + size_t sum = 0u; + for (const auto &it : blocks_) { + sum += it.size(); + } + return sum; + } + + //! Test if the rope is empty + bool empty(void) const { + return blocks_.empty(); + } + + //! Append a new memory block + Block &append(size_t init_size) { + return *blocks_.emplace(blocks_.end(), init_size); + } + + private: + //! Members + std::vector blocks_{}; + }; + + //! Constructor + IndexMemory(void) {} + + //! Constructor + IndexMemory(IndexMemory &&rhs) { + std::lock_guard latch(rhs.mutex_); + pool_ = std::move(rhs.pool_); + } + + //! Assignment + IndexMemory &operator=(IndexMemory &&rhs) { + std::lock_guard latch1(mutex_); + { + std::lock_guard latch2(rhs.mutex_); + pool_ = std::move(rhs.pool_); + } + return *this; + } + + //! Retrieve the singleton memory + static IndexMemory *Instance(void) { + static IndexMemory mem; + return (&mem); + } + + //! Clear the memory + void clear(void) { + std::lock_guard latch(mutex_); + pool_.clear(); + } + + //! Test if the element is exist + bool has(const std::string &key) const { + std::lock_guard latch(mutex_); + return (pool_.find(key) != pool_.end()); + } + + //! Create or overwrite a new memory rope + Rope::Pointer create(const std::string &key) { + std::lock_guard latch(mutex_); + auto it = pool_.emplace(key, nullptr).first; + it->second = std::make_shared(); + return it->second; + } + + //! Create or overwrite a new memory rope + Rope::Pointer create(std::string &&key) { + std::lock_guard latch(mutex_); + auto it = pool_.emplace(std::move(key), nullptr).first; + it->second = std::make_shared(); + return it->second; + } + + //! Open a memory rope (read only) + Rope::Pointer open(const std::string &key) const { + std::lock_guard latch(mutex_); + auto it = pool_.find(key); + if (it == pool_.end()) { + return nullptr; + } + return it->second; + } + + //! Remove a memory rope + void remove(const std::string &key) { + std::lock_guard latch(mutex_); + auto it = pool_.find(key); + if (it != pool_.end()) { + pool_.erase(it); + } + } + + private: + //! Disable them + IndexMemory(const IndexMemory &) = delete; + IndexMemory &operator=(const IndexMemory &) = delete; + + //! Members + std::map pool_{}; + mutable std::mutex mutex_{}; +}; + +} // 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..920a134 --- /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 "index_meta.h" +#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_meta.h b/src/core/framework/index_meta.h new file mode 100644 index 0000000..bdbecc7 --- /dev/null +++ b/src/core/framework/index_meta.h @@ -0,0 +1,723 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "ailego/container/params.h" + +namespace zvec { +namespace core { + +/*! Index Meta + */ +class IndexMeta { + public: + /*! Meta Types + */ + enum MetaType { MT_UNDEFINED = 0, MT_DENSE = 1, MT_SPARSE = 2 }; + + /*! Data Types + */ + enum DataType { + DT_UNDEFINED = 0, + DT_FP16 = 1, + DT_FP32 = 2, + DT_FP64 = 3, + DT_INT8 = 4, + DT_INT16 = 5, + DT_INT4 = 6, + DT_BINARY32 = 7, + DT_BINARY64 = 8, + }; + + /*! Major Orders + */ + enum MajorOrder { + MO_UNDEFINED = 0, + MO_ROW = 1, + MO_COLUMN = 2, + }; + + /*! Metric Type + */ + enum MetricType { + METRIC_EUCLIDEAN = 1, + METRIC_SQUARED_EUCLIDEAN = 2, + METRIC_INNER_PRODUCT = 3, + METRIC_MIPS_SQUARED_EUCLIDEAN = 4 + }; + + //! Constructor + IndexMeta(void) { + this->set_meta(DataType::DT_FP32, 128u); + this->set_metric("SquaredEuclidean", 0, ailego::Params()); + } + + //! Constructor + IndexMeta(DataType data_type, uint32_t dim) { + meta_type_ = MT_DENSE; + this->set_meta(data_type, dim); + this->set_metric("SquaredEuclidean", 0, ailego::Params()); + } + + //! Constructor + IndexMeta(MetaType meta_type, DataType data_type) { + meta_type_ = meta_type; + + this->set_meta(data_type, 0); + this->set_metric("SquaredEuclidean", 0, ailego::Params()); + } + + //! Constructor + IndexMeta(const IndexMeta &rhs) + : meta_type_{rhs.meta_type_}, + major_order_(rhs.major_order_), + data_type_(rhs.data_type_), + dimension_(rhs.dimension_), + unit_size_(rhs.unit_size_), + element_size_(rhs.element_size_), + space_id_(rhs.space_id_), + metric_revision_(rhs.metric_revision_), + converter_revision_(rhs.converter_revision_), + reformer_revision_(rhs.reformer_revision_), + trainer_revision_(rhs.trainer_revision_), + builder_revision_(rhs.builder_revision_), + reducer_revision_(rhs.reducer_revision_), + searcher_revision_(rhs.searcher_revision_), + streamer_revision_(rhs.streamer_revision_), + metric_name_(rhs.metric_name_), + converter_name_(rhs.converter_name_), + reformer_name_(rhs.reformer_name_), + trainer_name_(rhs.trainer_name_), + builder_name_(rhs.builder_name_), + reducer_name_(rhs.reducer_name_), + searcher_name_(rhs.searcher_name_), + streamer_name_(rhs.streamer_name_), + metric_params_(rhs.metric_params_), + converter_params_(rhs.converter_params_), + reformer_params_(rhs.reformer_params_), + trainer_params_(rhs.trainer_params_), + builder_params_(rhs.builder_params_), + reducer_params_(rhs.reducer_params_), + searcher_params_(rhs.searcher_params_), + streamer_params_(rhs.streamer_params_), + attributes_(rhs.attributes_) {} + + //! Constructor + IndexMeta(IndexMeta &&rhs) + : meta_type_{rhs.meta_type_}, + major_order_(rhs.major_order_), + data_type_(rhs.data_type_), + dimension_(rhs.dimension_), + unit_size_(rhs.unit_size_), + element_size_(rhs.element_size_), + space_id_(rhs.space_id_), + metric_revision_(rhs.metric_revision_), + converter_revision_(rhs.converter_revision_), + reformer_revision_(rhs.reformer_revision_), + trainer_revision_(rhs.trainer_revision_), + builder_revision_(rhs.builder_revision_), + reducer_revision_(rhs.reducer_revision_), + searcher_revision_(rhs.searcher_revision_), + streamer_revision_(rhs.streamer_revision_), + metric_name_(std::move(rhs.metric_name_)), + converter_name_(std::move(rhs.converter_name_)), + reformer_name_(std::move(rhs.reformer_name_)), + trainer_name_(std::move(rhs.trainer_name_)), + builder_name_(std::move(rhs.builder_name_)), + reducer_name_(std::move(rhs.reducer_name_)), + searcher_name_(std::move(rhs.searcher_name_)), + streamer_name_(std::move(rhs.streamer_name_)), + metric_params_(std::move(rhs.metric_params_)), + converter_params_(std::move(rhs.converter_params_)), + reformer_params_(std::move(rhs.reformer_params_)), + trainer_params_(std::move(rhs.trainer_params_)), + builder_params_(std::move(rhs.builder_params_)), + reducer_params_(std::move(rhs.reducer_params_)), + searcher_params_(std::move(rhs.searcher_params_)), + streamer_params_(std::move(rhs.streamer_params_)), + attributes_(std::move(rhs.attributes_)) {} + + //! Assignment + IndexMeta &operator=(const IndexMeta &rhs) { + meta_type_ = rhs.meta_type_; + major_order_ = rhs.major_order_; + data_type_ = rhs.data_type_; + dimension_ = rhs.dimension_; + unit_size_ = rhs.unit_size_; + element_size_ = rhs.element_size_; + space_id_ = rhs.space_id_; + metric_revision_ = rhs.metric_revision_; + converter_revision_ = rhs.converter_revision_; + reformer_revision_ = rhs.reformer_revision_; + trainer_revision_ = rhs.trainer_revision_; + builder_revision_ = rhs.builder_revision_; + reducer_revision_ = rhs.reducer_revision_; + searcher_revision_ = rhs.searcher_revision_; + streamer_revision_ = rhs.streamer_revision_; + metric_name_ = std::move(rhs.metric_name_); + converter_name_ = std::move(rhs.converter_name_); + reformer_name_ = std::move(rhs.reformer_name_); + trainer_name_ = std::move(rhs.trainer_name_); + builder_name_ = std::move(rhs.builder_name_); + reducer_name_ = std::move(rhs.reducer_name_); + searcher_name_ = std::move(rhs.searcher_name_); + streamer_name_ = std::move(rhs.streamer_name_); + metric_params_ = std::move(rhs.metric_params_); + converter_params_ = std::move(rhs.converter_params_); + reformer_params_ = std::move(rhs.reformer_params_); + trainer_params_ = std::move(rhs.trainer_params_); + builder_params_ = std::move(rhs.builder_params_); + reducer_params_ = std::move(rhs.reducer_params_); + searcher_params_ = std::move(rhs.searcher_params_); + streamer_params_ = std::move(rhs.streamer_params_); + attributes_ = std::move(rhs.attributes_); + + return *this; + } + + //! Assignment + IndexMeta &operator=(IndexMeta &&rhs) { + meta_type_ = rhs.meta_type_; + major_order_ = rhs.major_order_; + data_type_ = rhs.data_type_; + dimension_ = rhs.dimension_; + unit_size_ = rhs.unit_size_; + element_size_ = rhs.element_size_; + space_id_ = rhs.space_id_; + metric_revision_ = rhs.metric_revision_; + converter_revision_ = rhs.converter_revision_; + reformer_revision_ = rhs.reformer_revision_; + trainer_revision_ = rhs.trainer_revision_; + builder_revision_ = rhs.builder_revision_; + reducer_revision_ = rhs.reducer_revision_; + searcher_revision_ = rhs.searcher_revision_; + streamer_revision_ = rhs.streamer_revision_; + metric_name_ = std::move(rhs.metric_name_); + converter_name_ = std::move(rhs.converter_name_); + reformer_name_ = std::move(rhs.reformer_name_); + trainer_name_ = std::move(rhs.trainer_name_); + builder_name_ = std::move(rhs.builder_name_); + reducer_name_ = std::move(rhs.reducer_name_); + searcher_name_ = std::move(rhs.searcher_name_); + streamer_name_ = std::move(rhs.streamer_name_); + metric_params_ = std::move(rhs.metric_params_); + converter_params_ = std::move(rhs.converter_params_); + reformer_params_ = std::move(rhs.reformer_params_); + trainer_params_ = std::move(rhs.trainer_params_); + builder_params_ = std::move(rhs.builder_params_); + reducer_params_ = std::move(rhs.reducer_params_); + searcher_params_ = std::move(rhs.searcher_params_); + streamer_params_ = std::move(rhs.streamer_params_); + attributes_ = std::move(rhs.attributes_); + + return *this; + } + + //! Reset the meta + void clear(void) { + meta_type_ = MetaType::MT_DENSE; + major_order_ = MajorOrder::MO_UNDEFINED; + data_type_ = DataType::DT_UNDEFINED; + dimension_ = 0; + unit_size_ = 0; + element_size_ = 0; + space_id_ = 0; + metric_revision_ = 0; + converter_revision_ = 0; + reformer_revision_ = 0; + trainer_revision_ = 0; + builder_revision_ = 0; + reducer_revision_ = 0; + searcher_revision_ = 0; + streamer_revision_ = 0; + metric_name_.clear(); + converter_name_.clear(); + reformer_name_.clear(); + trainer_name_.clear(); + builder_name_.clear(); + reducer_name_.clear(); + searcher_name_.clear(); + streamer_name_.clear(); + metric_params_.clear(); + converter_params_.clear(); + reformer_params_.clear(); + trainer_params_.clear(); + builder_params_.clear(); + reducer_params_.clear(); + searcher_params_.clear(); + streamer_params_.clear(); + attributes_.clear(); + } + + //! Retrieve major order information + MetaType meta_type(void) const { + return meta_type_; + } + + //! Retrieve major order information + MajorOrder major_order(void) const { + return major_order_; + } + + //! Retrieve type information + DataType data_type(void) const { + return data_type_; + } + + //! Retrieve dimension + uint32_t dimension(void) const { + return dimension_; + } + + //! Retrieve unit size in bytes + uint32_t unit_size(void) const { + return unit_size_; + } + + //! Retrieve element size in bytes + uint32_t element_size(void) const { + return element_size_; + } + + //! Retrieve space id + uint64_t space_id(void) const { + return space_id_; + } + + //! Retrieve revision of metric + uint32_t metric_revision(void) const { + return metric_revision_; + } + + //! Retrieve revision of converter + uint32_t converter_revision(void) const { + return converter_revision_; + } + + //! Retrieve revision of reformer + uint32_t reformer_revision(void) const { + return reformer_revision_; + } + + //! Retrieve revision of trainer + uint32_t trainer_revision(void) const { + return trainer_revision_; + } + + //! Retrieve revision of builder + uint32_t builder_revision(void) const { + return builder_revision_; + } + + //! Retrieve revision of searcher + uint32_t searcher_revision(void) const { + return searcher_revision_; + } + + //! Retrieve revision of reducer + uint32_t reducer_revision(void) const { + return reducer_revision_; + } + + //! Retrieve revision of streamer + uint32_t streamer_revision(void) const { + return streamer_revision_; + } + + //! Retrieve name of metric + const std::string &metric_name(void) const { + return metric_name_; + } + + //! Retrieve name of converter + const std::string &converter_name(void) const { + return converter_name_; + } + + //! Retrieve name of reformer + const std::string &reformer_name(void) const { + return reformer_name_; + } + + //! Retrieve name of trainer + const std::string &trainer_name(void) const { + return trainer_name_; + } + + //! Retrieve name of builder + const std::string &builder_name(void) const { + return builder_name_; + } + + //! Retrieve name of reducer + const std::string &reducer_name(void) const { + return reducer_name_; + } + + //! Retrieve name of searcher + const std::string &searcher_name(void) const { + return searcher_name_; + } + + //! Retrieve name of streamer + const std::string &streamer_name(void) const { + return streamer_name_; + } + + //! Retrieve metric params + const ailego::Params &metric_params(void) const { + return metric_params_; + } + + //! Retrieve converter params + const ailego::Params &converter_params(void) const { + return converter_params_; + } + + //! Retrieve reformer params + const ailego::Params &reformer_params(void) const { + return reformer_params_; + } + + //! Retrieve trainer params + const ailego::Params &trainer_params(void) const { + return trainer_params_; + } + + //! Retrieve builder params + const ailego::Params &builder_params(void) const { + return builder_params_; + } + + //! Retrieve reducer params + const ailego::Params &reducer_params(void) const { + return reducer_params_; + } + + //! Retrieve searcher params + const ailego::Params &searcher_params(void) const { + return searcher_params_; + } + + //! Retrieve streamer params + const ailego::Params &streamer_params(void) const { + return streamer_params_; + } + + //! Retrieve attributes + const ailego::Params &attributes(void) const { + return attributes_; + } + + //! Retrieve mutable attributes + ailego::Params *mutable_attributes(void) { + return &attributes_; + } + + //! Set meta type + void set_meta_type(MetaType meta_type) { + meta_type_ = meta_type; + } + + //! Set major order of features + void set_major_order(MajorOrder major_order) { + major_order_ = major_order; + } + + //! Set dimension of feature + void set_dimension(uint32_t dim) { + dimension_ = dim; + element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dim); + } + + //! Set meta information of feature + void set_data_type(DataType data_type) { + data_type_ = data_type; + unit_size_ = UnitSizeof(data_type); + } + + //! Set meta information of feature + void set_meta(DataType data_type, uint32_t unit, uint32_t dim) { + data_type_ = data_type; + dimension_ = dim; + unit_size_ = unit; + element_size_ = ElementSizeof(data_type, unit, dim); + } + + //! Set meta information of feature + void set_meta(DataType data_type, uint32_t dim) { + this->set_meta(data_type, UnitSizeof(data_type), dim); + } + + //! Set information of metric + template + void set_metric(TName &&name, uint32_t rev, TParams &¶ms) { + metric_name_ = std::forward(name); + metric_revision_ = rev; + metric_params_ = std::forward(params); + } + + //! Set information of converter + template + void set_converter(TName &&name, uint32_t rev, TParams &¶ms) { + converter_name_ = std::forward(name); + converter_revision_ = rev; + converter_params_ = std::forward(params); + } + + //! Set information of reformer + template + void set_reformer(TName &&name, uint32_t rev, TParams &¶ms) { + reformer_name_ = std::forward(name); + reformer_revision_ = rev; + reformer_params_ = std::forward(params); + } + + //! Set information of trainer + template + void set_trainer(TName &&name, uint32_t rev, TParams &¶ms) { + trainer_name_ = std::forward(name); + trainer_revision_ = rev; + trainer_params_ = std::forward(params); + } + + //! Set information of builder + template + void set_builder(TName &&name, uint32_t rev, TParams &¶ms) { + builder_name_ = std::forward(name); + builder_revision_ = rev; + builder_params_ = std::forward(params); + } + + //! Set information of reducer + template + void set_reducer(TName &&name, uint32_t rev, TParams &¶ms) { + reducer_name_ = std::forward(name); + reducer_revision_ = rev; + reducer_params_ = std::forward(params); + } + + //! Set information of searcher + template + void set_searcher(TName &&name, uint32_t rev, TParams &¶ms) { + searcher_name_ = std::forward(name); + searcher_revision_ = rev; + searcher_params_ = std::forward(params); + } + + //! Set information of streamer + template + void set_streamer(TName &&name, uint32_t rev, TParams &¶ms) { + streamer_name_ = std::forward(name); + streamer_revision_ = rev; + streamer_params_ = std::forward(params); + } + + //! Serialize meta information into buffer + void serialize(std::string *out) const; + + //! Derialize meta information from buffer + bool deserialize(const void *data, size_t len); + + //! Calculate unit size of feature + static uint32_t UnitSizeof(DataType data_type) { + static const uint32_t unit_size_table[] = { + 0u, // DT_UNDEFINED + sizeof(uint16_t), // DT_FP16 + sizeof(float), // DT_FP32 + sizeof(double), // DT_FP64 + sizeof(int8_t), // DT_INT8 + sizeof(int16_t), // DT_INT16 + sizeof(uint8_t), // DT_INT4 + sizeof(uint32_t), // DT_BINARY32 + sizeof(uint64_t) // DT_BINARY64 + }; + return unit_size_table[data_type]; + } + + //! Calculate align size of feature + static uint32_t AlignSizeof(DataType ft) { + static const uint32_t align_size_table[] = { + 0u, // DT_UNDEFINED + sizeof(uint16_t), // DT_FP16 + sizeof(float), // DT_FP32 + sizeof(double), // DT_FP64 + sizeof(int8_t) * 4, // DT_INT8 + sizeof(int16_t), // DT_INT16 + sizeof(uint8_t) * 4, // DT_INT4 + sizeof(uint32_t), // DT_BINARY32 + sizeof(uint64_t) // DT_BINARY64 + }; + return align_size_table[ft]; + } + + //! Calculate element size of feature + static uint32_t ElementSizeof(DataType data_type, uint32_t unit, + uint32_t dim) { + switch (data_type) { + case DataType::DT_UNDEFINED: + return 0; + case DataType::DT_FP16: + case DataType::DT_FP32: + case DataType::DT_FP64: + case DataType::DT_INT8: + case DataType::DT_INT16: + return (dim * unit); + case DataType::DT_INT4: + return (dim + unit * 2 - 1) / (unit * 2) * unit; + case DataType::DT_BINARY32: + case DataType::DT_BINARY64: + return (dim + unit * 8 - 1) / (unit * 8) * unit; + } + return 0; + } + + //! Calculate element size of vector + static uint32_t ElementSizeof(DataType data_type, uint32_t dim) { + return ElementSizeof(data_type, UnitSizeof(data_type), dim); + } + + private: + MetaType meta_type_{MetaType::MT_DENSE}; + MajorOrder major_order_{MajorOrder::MO_UNDEFINED}; + DataType data_type_{DataType::DT_UNDEFINED}; + uint32_t dimension_{0}; + uint32_t unit_size_{0}; + uint32_t element_size_{0}; + uint64_t space_id_{0}; + uint32_t metric_revision_{0}; + uint32_t converter_revision_{0}; + uint32_t reformer_revision_{0}; + uint32_t trainer_revision_{0}; + uint32_t builder_revision_{0}; + uint32_t reducer_revision_{0}; + uint32_t searcher_revision_{0}; + uint32_t streamer_revision_{0}; + + std::string metric_name_{}; + std::string converter_name_{}; + std::string reformer_name_{}; + std::string trainer_name_{}; + std::string builder_name_{}; + std::string reducer_name_{}; + std::string searcher_name_{}; + std::string streamer_name_{}; + + ailego::Params metric_params_{}; + ailego::Params converter_params_{}; + ailego::Params reformer_params_{}; + ailego::Params trainer_params_{}; + ailego::Params builder_params_{}; + ailego::Params reducer_params_{}; + ailego::Params searcher_params_{}; + ailego::Params streamer_params_{}; + ailego::Params attributes_{}; +}; + +/*! Index Query Meta + */ +class IndexQueryMeta { + public: + //! Constructor + IndexQueryMeta(void) {} + + //! Constructor + IndexQueryMeta(IndexMeta::MetaType meta_type, IndexMeta::DataType data_type, + uint32_t unit, uint32_t dim) + : meta_type_(meta_type), + data_type_(data_type), + dimension_(dim), + unit_size_(unit), + element_size_(IndexMeta::ElementSizeof(data_type, unit, dim)) {} + + //! Constructor + IndexQueryMeta(IndexMeta::DataType data_type, uint32_t dim) + : IndexQueryMeta{IndexMeta::MetaType::MT_DENSE, data_type, + IndexMeta::UnitSizeof(data_type), dim} {} + + //! Constructor + IndexQueryMeta(IndexMeta::DataType data_type) + : IndexQueryMeta{IndexMeta::MetaType::MT_SPARSE, data_type, + IndexMeta::UnitSizeof(data_type), 0} {} + + //! Constructor + IndexQueryMeta(IndexMeta::MetaType meta_type, IndexMeta::DataType data_type, + uint32_t dim = 0) + : IndexQueryMeta{meta_type, data_type, IndexMeta::UnitSizeof(data_type), + dim} {} + + //! Retrieve meta type + IndexMeta::MetaType meta_type(void) const { + return meta_type_; + } + + //! Retrieve data + IndexMeta::DataType data_type(void) const { + return data_type_; + } + + //! Retrieve dimension of features + uint32_t dimension(void) const { + return dimension_; + } + + //! Retrieve unit size of feature + uint32_t unit_size(void) const { + return unit_size_; + } + + //! Retrieve element size of feature + uint32_t element_size(void) const { + return element_size_; + } + + //! Set dimension of feature + void set_dimension(uint32_t dim) { + dimension_ = dim; + element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dim); + } + + //! Set meta type + void set_meta_type(IndexMeta::MetaType meta_type) { + meta_type_ = meta_type; + } + + //! Set data type + void set_data_type(IndexMeta::DataType data_type) { + data_type_ = data_type; + } + + //! Set meta information of feature + void set_meta(IndexMeta::DataType data_type, uint32_t unit, uint32_t dim) { + data_type_ = data_type; + dimension_ = dim; + unit_size_ = unit; + element_size_ = IndexMeta::ElementSizeof(data_type, unit, dim); + } + + //! Set meta information of feature + void set_meta(IndexMeta::DataType data_type, uint32_t dim) { + this->set_meta(data_type, IndexMeta::UnitSizeof(data_type), dim); + } + + private: + IndexMeta::MetaType meta_type_{IndexMeta::MetaType::MT_DENSE}; + IndexMeta::DataType data_type_{IndexMeta::DataType::DT_UNDEFINED}; + uint32_t dimension_{0}; + uint32_t unit_size_{0}; + uint32_t element_size_{0}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_metric.h b/src/core/framework/index_metric.h new file mode 100644 index 0000000..ef1a992 --- /dev/null +++ b/src/core/framework/index_metric.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 "index_error.h" +#include "index_meta.h" +#include "index_module.h" + +namespace zvec { +namespace core { + +/*! Index Metric + */ +struct IndexMetric : public IndexModule { + //! Index Metric Pointer + typedef std::shared_ptr Pointer; + + //! Matrix Distance Function + typedef void (*MatrixDistanceHandle)(const void *m, const void *q, size_t dim, + float *out); + + //! Matrix Distance Function Object + using MatrixDistance = + std::function; + + //! Matrix Sparse Distance Function + typedef void (*MatrixSparseDistanceHandle)(const void *m_sparse_data, + const void *q_sparse_data, + float *out); + + //! Matrix Sparse Distance Function Object + using MatrixSparseDistance = std::function; + + + //! Matrix Batch Distance Function + typedef void (*MatrixBatchDistanceHandle)(const void **m, const void *q, + size_t num, size_t dim, float *out); + + //! Matrix Batch Distance Function Object + using MatrixBatchDistance = std::function; + + //! Destructor + virtual ~IndexMetric(void) {} + + //! Initialize Metric + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) = 0; + + //! Cleanup Metric + virtual int cleanup(void) = 0; + + //! Retrieve if it matched + virtual bool is_matched(const IndexMeta &meta) const = 0; + + //! Retrieve if it matched + virtual bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const = 0; + + //! Retrieve distance function for query + virtual MatrixDistance distance(void) const { + return nullptr; + } + + //! Retrieve hybrid distance function for query + virtual MatrixSparseDistance sparse_distance(void) const { + return nullptr; + }; + + //! Retrieve distance function for query + virtual MatrixBatchDistance batch_distance(void) const { + return nullptr; + } + + //! Retrieve distance function for index features + virtual MatrixDistance distance_matrix(size_t /*m*/, size_t /*n*/) const { + return nullptr; + } + + //! Retrieve params of Metric + virtual const ailego::Params ¶ms(void) const = 0; + + //! Retrieve query metric object of this index metric + virtual Pointer query_metric(void) const = 0; + + //! Normalize result + virtual void normalize(float *score) const { + (void)score; + } + + //! Denormalize result + virtual void denormalize(float *score) const { + (void)score; + } + + //! Retrieve if it supports normalization + virtual bool support_normalize(void) const { + return false; + } + + //! Train the metric + virtual int train(const void *vec, size_t dim) { + (void)vec; + (void)dim; + return 0; + } + + //! Retrieve if it supports training + virtual bool support_train(void) const { + return false; + } + + //! Compute the distance between feature and query + float distance(const void *m, const void *q, size_t dim) const { + float dist; + (this->distance())(m, q, dim, &dist); + return dist; + } + + using DistanceBatchQueryPreprocessFunc = + ailego::DistanceBatch::DistanceBatchQueryPreprocessFunc; + + virtual DistanceBatchQueryPreprocessFunc get_query_preprocess_func() const { + return nullptr; + } +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_module.h b/src/core/framework/index_module.h new file mode 100644 index 0000000..d82fdbb --- /dev/null +++ b/src/core/framework/index_module.h @@ -0,0 +1,67 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 { + +/*! Index Module + */ +class IndexModule { + public: + //! Index Module Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexModule(void) {} + + //! Retrieve debug information + virtual std::string debug_string(void) const { + return std::string(); + } + + //! Retrieve name of module + const std::string &name(void) const { + return name_; + } + + //! Retrieve revision of module + uint32_t revision(void) const { + return revision_; + } + + protected: + friend struct IndexFactory; + + //! Set name of module + void set_name(const std::string &str) { + name_ = str; + } + + //! Set revision of module + void set_revision(uint32_t val) { + revision_ = val; + } + + private: + //! Members + uint32_t revision_{0u}; + std::string name_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_packer.h b/src/core/framework/index_packer.h new file mode 100644 index 0000000..7181e0e --- /dev/null +++ b/src/core/framework/index_packer.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 "index_error.h" +#include "index_format.h" +#include "index_version.h" + +namespace zvec { +namespace core { + +/*! Index Packer + */ +class IndexPacker { + public: + /*! Index Packer Segment Meta + */ + class SegmentMeta { + public: + //! Constructor + SegmentMeta(const std::string &str, size_t dsz, size_t psz, uint32_t crc) + : data_size_(dsz), padding_size_(psz), data_crc_(crc), id_(str) {} + + //! Constructor + SegmentMeta(std::string &&str, size_t dsz, size_t psz, uint32_t crc) + : data_size_(dsz), + padding_size_(psz), + data_crc_(crc), + id_(std::forward(str)) {} + + //! Constructor + SegmentMeta(const SegmentMeta &rhs) + : data_size_(rhs.data_size_), + padding_size_(rhs.padding_size_), + data_crc_(rhs.data_crc_), + id_(rhs.id_) {} + + //! Constructor + SegmentMeta(SegmentMeta &&rhs) + : data_size_(rhs.data_size_), + padding_size_(rhs.padding_size_), + data_crc_(rhs.data_crc_), + id_(std::move(rhs.id_)) {} + + //! Retrieve id of segment + const std::string &id(void) const { + return id_; + } + + //! Retrieve size of data + size_t data_size(void) const { + return data_size_; + } + + //! Retrieve crc of data + uint32_t data_crc(void) const { + return data_crc_; + } + + //! Retrieve size of padding + size_t padding_size(void) const { + return padding_size_; + } + + private: + size_t data_size_{0}; + size_t padding_size_{0}; + uint32_t data_crc_{0}; + std::string id_{}; + }; + + //! Retrieve magic number of index + uint32_t magic(void) const { + return magic_; + } + + //! Reset the packer + void reset(void) { + magic_ = 0; + data_crc_ = 0u; + data_size_ = 0u; + } + + //! Setup header of index pacakge + template + bool setup(TFunc write_data) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + IndexFormat::MetaHeader header; + IndexFormat::SetupMetaHeader( + &header, (uint32_t)(0 - sizeof(IndexFormat::MetaFooter)), + sizeof(IndexFormat::MetaHeader)); + + if (write_data(&header, sizeof(header)) != sizeof(header)) { + return false; + } + magic_ = header.magic; + return true; + } + + //! Pack index data + template + size_t pack(TFunc write_data, const void *data, size_t len) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + size_t wrlen = write_data(data, len); + if (wrlen > 0u) { + data_crc_ = ailego::Crc32c::Hash(data, wrlen, data_crc_); + data_size_ += wrlen; + } + return wrlen; + } + + //! Finish packing data + template + bool finish(TFunc write_data, std::vector &stab) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + + size_t content_size = 0u; + for (const auto &it : stab) { + content_size += it.data_size() + it.padding_size(); + } + + if (content_size != data_size_) { + return false; + } + + if (!this->pack_version(write_data, stab)) { + return false; + } + + // Write the padding if need + size_t content_padding_size = ailego_align(data_size_, 32) - data_size_; + if (content_padding_size) { + std::string padding(content_padding_size, '\0'); + + if (write_data(padding.data(), padding.size()) != padding.size()) { + return false; + } + } + + // Prepare segment meta buffer + IndexFormat::SegmentMetaBuffer buffer(stab.size()); + for (const auto &it : stab) { + buffer.append(it.id(), it.data_size(), it.padding_size(), it.data_crc()); + } + buffer.resize(ailego_align(buffer.size(), 32)); + + // Write segment table into file + if (write_data(buffer.data(), buffer.size()) != buffer.size()) { + return false; + } + + // Update footer + IndexFormat::MetaFooter footer; + IndexFormat::SetupMetaFooter(&footer); + footer.segments_meta_crc = buffer.crc(); + footer.content_crc = data_crc_; + footer.segment_count = stab.size(); + footer.segments_meta_size = buffer.size(); + footer.content_size = data_size_; + footer.content_padding_size = content_padding_size; + footer.total_size = footer.content_size + footer.content_padding_size + + footer.segments_meta_size + + sizeof(IndexFormat::MetaHeader) + + sizeof(IndexFormat::MetaFooter); + IndexFormat::UpdateMetaFooter(&footer, 0); + + // Write footer into file + if (write_data(&footer, sizeof(footer)) != sizeof(footer)) { + return false; + } + return true; + } + + //! Pack index version + template + bool pack_version(TFunc write_data, std::vector &stab) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + std::string buffer(IndexVersion::Details()); + + 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 (write_data(buffer.data(), buffer.size()) != buffer.size()) { + return false; + } + data_crc_ = ailego::Crc32c::Hash(buffer.data(), buffer.size(), data_crc_); + data_size_ += buffer.size(); + stab.emplace_back(std::string("IndexVersion"), data_size, + buffer.size() - data_size, data_crc); + return true; + } + + private: + uint32_t magic_{0u}; + uint32_t data_crc_{0u}; + size_t data_size_{0u}; +}; + +} // 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..5c0e5b8 --- /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 "index_plugin.h" +#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_plugin.h b/src/core/framework/index_plugin.h new file mode 100644 index 0000000..954ed98 --- /dev/null +++ b/src/core/framework/index_plugin.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 + +namespace zvec { +namespace core { + +/*! Index Plugin + */ +class IndexPlugin { + public: + //! Constructor + IndexPlugin(void) : handle_(nullptr) {} + + //! Constructor + IndexPlugin(IndexPlugin &&plugin) : handle_(plugin.handle_) { + plugin.handle_ = nullptr; + } + + //! Constructor + explicit IndexPlugin(const std::string &path) : handle_(nullptr) { + this->load(path); + } + + //! Destructor + ~IndexPlugin(void) {} + + //! Test if the plugin is valid + bool is_valid(void) const { + return (!!handle_); + } + + //! Retrieve the handle + void *handle(void) const { + return handle_; + } + + //! Load the library path + bool load(const std::string &path); + + //! Load the library path + bool load(const std::string &path, std::string *err); + + //! Unload plugin + void unload(void); + + private: + //! Disable them + IndexPlugin(const IndexPlugin &) = delete; + IndexPlugin &operator=(const IndexPlugin &) = delete; + + //! Members + void *handle_; +}; + +/*! Index Plugin Broker + */ +class IndexPluginBroker { + public: + //! Constructor + IndexPluginBroker(void) : plugins_() {} + + //! Constructor + IndexPluginBroker(IndexPluginBroker &&broker) + : plugins_(std::move(broker.plugins_)) {} + + //! Destructor + ~IndexPluginBroker(void) {} + + //! Emplace a plugin + bool emplace(IndexPlugin &&plugin); + + //! Emplace a plugin via library path + bool emplace(const std::string &path) { + return this->emplace(IndexPlugin(path)); + } + + //! Emplace a plugin via library path + bool emplace(const std::string &path, std::string *err) { + IndexPlugin plugin; + if (!plugin.load(path, err)) { + return false; + } + return this->emplace(std::move(plugin)); + } + + //! Retrieve count of plugins in broker + size_t count(void) const { + return plugins_.size(); + } + + private: + //! Disable them + IndexPluginBroker(const IndexPluginBroker &) = delete; + IndexPluginBroker &operator=(const IndexPluginBroker &) = delete; + + //! Members + std::vector plugins_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_provider.h b/src/core/framework/index_provider.h new file mode 100644 index 0000000..93d5671 --- /dev/null +++ b/src/core/framework/index_provider.h @@ -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. +#pragma once + +#include "index_error.h" +#include "index_holder.h" +#include "index_storage.h" + +namespace zvec { +namespace core { + +/*! Index Provider + */ +struct IndexProvider : public IndexHolder { + //! Index Provider Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexProvider(void) {} + + bool multipass() const override { + return true; + } + + public: // Provider's unique method + //! Retrieve a vector using a primary key + virtual const void *get_vector(const uint64_t key) const = 0; + + //! Retrieve a vector using a primary key + virtual int get_vector(const uint64_t /*key*/, + IndexStorage::MemoryBlock & /*block*/) const { + return IndexError_NotImplemented; + } + + //! Retrieve the owner class + virtual const std::string &owner_class(void) const = 0; +}; + +/*! Index SparseProvider + */ +struct IndexSparseProvider : IndexSparseHolder { + //! Index Provider Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexSparseProvider(void) {} + + bool multipass() const override { + return true; + } + + public: // Provider's unique method + //! Retrieve a vector using a primary key + virtual int get_sparse_vector(uint64_t key, uint32_t *sparse_count, + std::string *sparse_indices_buffer, + std::string *sparse_values_buffer) const = 0; + + //! Retrieve the owner class + virtual const std::string &owner_class(void) const = 0; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_reducer.h b/src/core/framework/index_reducer.h new file mode 100644 index 0000000..5c85bc5 --- /dev/null +++ b/src/core/framework/index_reducer.h @@ -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. + +#pragma once + +#include "ailego/parallel/thread_pool.h" +#include "index_builder.h" +#include "index_converter.h" +#include "index_dumper.h" +#include "index_error.h" +#include "index_filter.h" +#include "index_reformer.h" +#include "index_stats.h" +#include "index_streamer.h" + +namespace zvec { +namespace core { + +/*! Index Reducer Base + */ +class IndexReducerBase : public IndexModule { + public: + //! Index Reducer Pointer + typedef std::shared_ptr Pointer; + + /*! Index Reducer Stats + */ + class Stats : public IndexStats { + public: + Stats() {} + Stats(const Stats &stats) { + *this = stats; + } + Stats &operator=(const Stats &stats) { + this->loaded_count_.store(stats.loaded_count_.load()); + this->dumped_count_.store(stats.dumped_count_.load()); + this->filtered_count_.store(stats.filtered_count_.load()); + this->duplicated_count_.store(stats.duplicated_count_.load()); + this->reduced_costtime_.store(stats.reduced_costtime_.load()); + this->dumped_costtime_.store(stats.dumped_costtime_.load()); + return *this; + } + //! Set count of documents loaded + void set_loaded_count(size_t count) { + loaded_count_ = count; + } + + //! Set count of documents dumped + void set_dumped_count(size_t count) { + dumped_count_ = count; + } + + //! Set count of documents filtered + void set_filtered_count(size_t count) { + filtered_count_ = count; + } + + //! Set count of documents duplicated + void set_duplicated_count(size_t count) { + duplicated_count_ = count; + } + + //! Set time cost of documents reduced + void set_reduced_costtime(uint64_t cost) { + reduced_costtime_ = cost; + } + + //! Set time cost of documents dumped + void set_dumped_costtime(uint64_t cost) { + dumped_costtime_ = cost; + } + + //! Retrieve count of documents loaded + size_t loaded_count(void) const { + return loaded_count_; + } + + //! Retrieve count of documents dumped + size_t dumped_count(void) const { + return dumped_count_; + } + + //! Retrieve count of documents filtered + size_t filtered_count(void) const { + return filtered_count_; + } + + //! Retrieve count of documents duplicated + size_t duplicated_count(void) const { + return duplicated_count_; + } + + //! Retrieve time cost of documents reduced + uint64_t reduced_costtime(void) const { + return reduced_costtime_; + } + + //! Retrieve time cost of documents dumped + uint64_t dumped_costtime(void) const { + return dumped_costtime_; + } + + //! Retrieve count of documents loaded (mutable) + std::atomic *mutable_loaded_count(void) { + return &loaded_count_; + } + + //! Retrieve count of documents dumped (mutable) + std::atomic *mutable_dumped_count(void) { + return &dumped_count_; + } + + //! Retrieve count of documents filtered (mutable) + std::atomic *mutable_filtered_count(void) { + return &filtered_count_; + } + + //! Retrieve count of documents duplicated (mutable) + std::atomic *mutable_duplicated_count(void) { + return &duplicated_count_; + } + + //! Retrieve time cost of documents reduced (mutable) + std::atomic *mutable_reduced_costtime(void) { + return &reduced_costtime_; + } + + //! Retrieve time cost of documents dumped (mutable) + std::atomic *mutable_dumped_costtime(void) { + return &dumped_costtime_; + } + + private: + //! Members + std::atomic loaded_count_{0u}; + std::atomic dumped_count_{0u}; + std::atomic filtered_count_{0u}; + std::atomic duplicated_count_{0u}; + std::atomic reduced_costtime_{0u}; + std::atomic dumped_costtime_{0u}; + }; + + //! Destructor + virtual ~IndexReducerBase(void) = default; + + //! Initialize Reducer + virtual int init(const ailego::Params ¶ms) = 0; + + //! Cleanup Reducer + virtual int cleanup(void) = 0; + + //! Reduce operator (with filter) + virtual int reduce(const IndexFilter &filter) = 0; + + //! Dump index into storage + virtual int dump(const IndexDumper::Pointer &dumper) = 0; + + //! Retrieve statistics + virtual const Stats &stats(void) const = 0; + + //! Set thread pool + void set_thread_pool(ailego::ThreadPool *pool) { + thread_pool_ = pool; + } + + //! Set stop flag + void set_stop_flag(std::atomic *stop_flag) { + stop_flag_ = stop_flag; + } + + protected: + ailego::ThreadPool *thread_pool_{nullptr}; + std::atomic *stop_flag_{nullptr}; +}; + +/*! Index Reducer + */ +class IndexReducer : public IndexReducerBase { + public: + //! Index Reducer Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexReducer(void) = default; +}; + +/*! Index Sparse Reducer + */ +class IndexSparseReducer : public IndexReducerBase { + public: + //! Index Reducer Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexSparseReducer(void) = default; +}; + +/*! Index Streamer Reducer + */ +class IndexStreamerReducer : public IndexReducerBase { + public: + //! Index Reducer Pointer + typedef std::shared_ptr Pointer; + + virtual int set_target_streamer_wiht_info( + const IndexBuilder::Pointer builder, + const IndexStreamer::Pointer streamer, + const IndexConverter::Pointer converter, + const IndexReformer::Pointer reformer = nullptr, + const IndexQueryMeta &original_query_meta = IndexQueryMeta()) = 0; + virtual int feed_streamer_with_reformer( + IndexStreamer::Pointer streamer, + const IndexReformer::Pointer reformer) = 0; + + virtual ~IndexStreamerReducer(void) = default; +}; +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_refiner.h b/src/core/framework/index_refiner.h new file mode 100644 index 0000000..3b73e50 --- /dev/null +++ b/src/core/framework/index_refiner.h @@ -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. + +#pragma once + +#include +#include +#include "index_context.h" +#include "index_helper.h" +#include "index_searcher.h" +#include "index_streamer.h" + +namespace zvec { +namespace core { + +/*! Index Refiner + */ +class IndexRefiner : public IndexModule { + public: + //! Index Refiner Pointer + typedef std::shared_ptr Pointer; + + /*! Index Searcher Context + */ + struct Context : public IndexContext { + public: + Context() = default; + ~Context() = default; + + virtual int set_contexts(IndexRunner::Context::Pointer base_ctx, + IndexRunner::Context::Pointer refine_ctx) = 0; + }; + + //! Initialize refiner with streamer + virtual int init(IndexRunner::Pointer base_runner, + IndexRunner::Pointer refine_runner, + const ailego::Params ¶ms) = 0; + + //! Cleanup + virtual int cleanup() = 0; + + //! Create a context + virtual Context::Pointer create_context(void) const = 0; + + //! Add a vector into index + virtual int add_impl(uint64_t key, const void *base_query, + const IndexQueryMeta &base_qmeta, + const void *refine_query, + const IndexQueryMeta &refine_qmeta, + Context::Pointer &context) = 0; + + //! Similarity search + virtual int search_impl(const void *base_query, + const IndexQueryMeta &base_qmeta, + const void *refine_query, + const IndexQueryMeta &refine_qmeta, + Context::Pointer &context) const = 0; + //! Similarity search + virtual int search_impl(const void *base_query, + const IndexQueryMeta &base_qmeta, + const void *refine_query, + const IndexQueryMeta &refine_qmeta, uint32_t count, + Context::Pointer &context) const = 0; + + //! Similarity brute force search + virtual int search_bf_impl(const void *base_query, + const IndexQueryMeta &base_qmeta, + const void *refine_query, + const IndexQueryMeta &refine_qmeta, + Context::Pointer &context) const = 0; + + //! Similarity brute force search + virtual int search_bf_impl(const void *base_query, + const IndexQueryMeta &base_qmeta, + const void *refine_query, + const IndexQueryMeta &refine_qmeta, uint32_t count, + Context::Pointer &context) const = 0; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_reformer.h b/src/core/framework/index_reformer.h new file mode 100644 index 0000000..290dbb3 --- /dev/null +++ b/src/core/framework/index_reformer.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 "index_document.h" +#include "index_meta.h" + +namespace zvec { +namespace core { + +/*! Index Reformer + */ +class IndexReformer : public IndexModule { + public: + //! Index Reformer Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexReformer(void) {} + + //! Initialize Reformer + virtual int init(const ailego::Params ¶ms) = 0; + + //! Cleanup Reformer + virtual int cleanup(void) = 0; + + //! Load index from container + virtual int load(IndexStorage::Pointer cntr) = 0; + + //! Unload index + virtual int unload(void) = 0; + + //! Transform a query + virtual int transform(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const { + return IndexError_NotImplemented; + } + + //! Transform queries + virtual int transform(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, uint32_t /*count*/, + std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const { + return IndexError_NotImplemented; + } + + //! Convert a record + virtual int convert(const void *record, const IndexQueryMeta &rmeta, + std::string *out, IndexQueryMeta *ometa) const { + return this->transform(record, rmeta, out, ometa); + } + + //! Convert records + virtual int convert(const void *records, const IndexQueryMeta &rmeta, + uint32_t count, std::string *out, + IndexQueryMeta *ometa) const { + return this->transform(records, rmeta, count, out, ometa); + } + + //! Normalize results + virtual int normalize(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, + IndexDocumentList & /*result*/) const { + return IndexError_NotImplemented; + } + + virtual bool need_revert() const { + return false; + } + + virtual int revert(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + //! Transform a query + virtual int transform(uint32_t /*sparse_count*/, + const uint32_t * /*sparse_indices*/, + const void * /*sparse_query*/, + const IndexQueryMeta & /*qmeta*/, std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const { + return IndexError_NotImplemented; + } + + //! Transform queries + virtual int transform(const uint32_t * /*sparse_count*/, + const uint32_t * /*sparse_indices*/, + const void * /*sparse_query*/, + const IndexQueryMeta & /*qmeta*/, uint32_t /*count*/, + std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const { + return IndexError_NotImplemented; + } + + //! Convert a record + virtual int convert(uint32_t sparse_count, const uint32_t *sparse_indices, + const void *sparse_query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + return this->transform(sparse_count, sparse_indices, sparse_query, qmeta, + out, ometa); + } + + //! Convert records + virtual int convert(const uint32_t *sparse_count, + const uint32_t *sparse_indices, const void *sparse_query, + const IndexQueryMeta &qmeta, uint32_t count, + std::string *out, IndexQueryMeta *ometa) const { + return this->transform(sparse_count, sparse_indices, sparse_query, qmeta, + count, out, ometa); + } + + virtual int revert(const uint32_t /*sparse_count*/, + const uint32_t * /*sparse_indices*/, + const void * /*sparse_query*/, + const IndexQueryMeta & /*qmeta*/, + std::string * /*sparse_query_out*/) const { + return 0; + } +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_runner.h b/src/core/framework/index_runner.h new file mode 100644 index 0000000..3b032ca --- /dev/null +++ b/src/core/framework/index_runner.h @@ -0,0 +1,743 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_context.h" +#include "index_dumper.h" +#include "index_meta.h" +#include "index_metric.h" +#include "index_module.h" +#include "index_provider.h" +#include "index_stats.h" +#include "index_threads.h" +#include "index_trainer.h" + +namespace zvec { +namespace core { + +/*! Index Runner + */ +class IndexRunner : public IndexModule { + public: + //! Index Searcher Pointer + typedef std::shared_ptr Pointer; + + /*! Index Searcher Context + */ + using Context = IndexContext; + + /*! Index Searcher Provider + */ + using Provider = IndexProvider; + + /*! Index Sparse Searcher Provider + */ + using SparseProvider = IndexSparseProvider; + + /*! Index Streamer Stats + */ + class Stats : public IndexStats { + public: + Stats() {} + Stats(const Stats &stats) { + *this = stats; + } + Stats &operator=(const Stats &stats) { + this->revision_id_ = stats.revision_id_; + + this->trained_count_.store(stats.trained_count_.load()); + this->built_count_.store(stats.built_count_.load()); + this->dumped_count_.store(stats.dumped_count_.load()); + this->loaded_count_.store(stats.loaded_count_.load()); + this->added_count_.store(stats.added_count_.load()); + this->discarded_count_.store(stats.discarded_count_.load()); + this->updated_count_.store(stats.updated_count_.load()); + this->deleted_count_.store(stats.deleted_count_.load()); + + this->index_size_.store(stats.index_size_.load()); + this->dumped_size_.store(stats.dumped_size_.load()); + + this->check_point_.store(stats.check_point_.load()); + + this->create_time_.store(stats.create_time_.load()); + this->update_time_.store(stats.update_time_.load()); + this->loaded_costtime_.store(stats.loaded_costtime_.load()); + this->trained_costtime_.store(stats.trained_costtime_.load()); + this->built_costtime_.store(stats.built_costtime_.load()); + this->dumped_costtime_.store(stats.dumped_costtime_.load()); + + return *this; + } + //! Set revision id + void set_revision_id(size_t rev) { + revision_id_ = rev; + } + + //! Set count of documents trained + void set_trained_count(size_t count) { + trained_count_ = count; + } + + //! Set count of documents built + void set_built_count(size_t count) { + built_count_ = count; + } + + //! Set count of documents dumped + void set_dumped_count(size_t count) { + dumped_count_ = count; + } + + //! Set count of documents loaded + void set_loaded_count(size_t count) { + loaded_count_ = count; + } + + //! Set count of documents added + void set_added_count(size_t count) { + added_count_ = count; + } + + //! Set count of documents discarded + void set_discarded_count(size_t count) { + discarded_count_ = count; + } + + //! Set count of documents updated + void set_updated_count(size_t count) { + updated_count_ = count; + } + + //! Set count of documents deleted + void set_deleted_count(size_t count) { + deleted_count_ = count; + } + + //! Set size of index + void set_index_size(size_t count) { + index_size_ = count; + } + + //! Set size of index dumped + void set_dumped_size(size_t count) { + dumped_size_ = count; + } + + //! Set size of index dumped + void set_check_point(uint64_t val) { + check_point_ = val; + } + + //! Retrieve create time + void set_create_time(uint64_t val) { + create_time_ = val; + } + + //! Retrieve update time + void set_update_time(uint64_t val) { + update_time_ = val; + } + + //! Retrieve loaded costtime + void set_loaded_costtime(uint64_t val) { + loaded_costtime_ = val; + } + + //! Retrieve train costtime + void set_trained_costtime(uint64_t val) { + trained_costtime_ = val; + } + + //! Retrieve built costtime + void set_built_costtime(uint64_t val) { + built_costtime_ = val; + } + + //! Retrieve update time + void set_dumped_costtime(uint64_t val) { + dumped_costtime_ = val; + } + + //! Retrieve revision id + size_t revision_id(void) const { + return revision_id_; + } + + //! Retrieve count of documents trained + size_t trained_count(void) const { + return trained_count_; + } + + //! Retrieve count of documents built + size_t built_count(void) const { + return built_count_; + } + + //! Retrieve count of documents dumped + size_t dumped_count(void) const { + return dumped_count_; + } + + //! Retrieve count of documents loaded + size_t loaded_count(void) const { + return loaded_count_; + } + + //! Retrieve count of documents added + size_t added_count(void) const { + return added_count_; + } + + //! Retrieve count of documents discarded + size_t discarded_count(void) const { + return discarded_count_; + } + + //! Retrieve count of documents updated + size_t updated_count(void) const { + return updated_count_; + } + + //! Retrieve count of documents deleted + size_t deleted_count(void) const { + return deleted_count_; + } + + //! Retrieve size of index + size_t index_size(void) const { + return index_size_; + } + + //! Retrieve size of index dumped + size_t dumped_size(void) const { + return dumped_size_; + } + + //! Retrieve check point of index + uint64_t check_point(void) const { + return check_point_; + } + + //! Retrieve create time of index + uint64_t create_time(void) const { + return create_time_; + } + + //! Retrieve update time of index + uint64_t update_time(void) const { + return update_time_; + } + + //! Retrieve loaded cost time of index + uint64_t loaded_costtime(void) const { + return loaded_costtime_; + } + + //! Retrieve trained cost time of index + uint64_t trained_costtime(void) const { + return trained_costtime_; + } + + //! Retrieve built cost time of index + uint64_t built_costtime(void) const { + return built_costtime_; + } + + //! Retrieve dumped cost time of index + uint64_t dumped_costtime(void) const { + return dumped_costtime_; + } + + //! Retrieve count of documents loaded (mutable) + std::atomic *mutable_trained_count(void) { + return &loaded_count_; + } + + //! Retrieve count of documents built (mutable) + std::atomic *mutable_built_count(void) { + return &built_count_; + } + + //! Retrieve count of documents dumped (mutable) + std::atomic *mutable_dumped_count(void) { + return &dumped_count_; + } + + //! Retrieve count of documents loaded (mutable) + std::atomic *mutable_loaded_count(void) { + return &loaded_count_; + } + + //! Retrieve count of documents added (mutable) + std::atomic *mutable_added_count(void) { + return &added_count_; + } + + //! Retrieve count of documents discarded (mutable) + std::atomic *mutable_discarded_count(void) { + return &discarded_count_; + } + + //! Retrieve count of documents updated (mutable) + std::atomic *mutable_updated_count(void) { + return &updated_count_; + } + + //! Retrieve count of documents deleted (mutable) + std::atomic *mutable_deleted_count(void) { + return &deleted_count_; + } + + //! Retrieve size of index (mutable) + std::atomic *mutable_index_size(void) { + return &index_size_; + } + + //! Retrieve size of index dumped (mutable) + std::atomic *mutable_dumped_size(void) { + return &dumped_size_; + } + + //! Retrieve check point of index (mutable) + std::atomic *mutable_check_point(void) { + return &check_point_; + } + + //! Retrieve create time of index (mutable) + std::atomic *mutable_create_time(void) { + return &create_time_; + } + + //! Retrieve update time of index (mutable) + std::atomic *mutable_update_time(void) { + return &update_time_; + } + + //! Retrieve loaded time of index (mutable) + std::atomic *mutable_loaded_costtime(void) { + return &loaded_costtime_; + } + + //! Retrieve trained costtime of index (mutable) + std::atomic *mutable_trained_costtime(void) { + return &trained_costtime_; + } + + //! Retrieve built costtime of index (mutable) + std::atomic *mutable_built_costtime(void) { + return &built_costtime_; + } + + //! Retrieve dump costtime of index (mutable) + std::atomic *mutable_dumped_costtime(void) { + return &dumped_costtime_; + } + + void clear() { + this->clear_attributes(); + + revision_id_ = 0u; + + trained_count_ = 0u; + built_count_ = 0u; + dumped_count_ = 0u; + loaded_count_ = 0u; + added_count_ = 0u; + discarded_count_ = 0u; + updated_count_ = 0u; + deleted_count_ = 0u; + + index_size_ = 0u; + dumped_size_ = 0u; + check_point_ = 0u; + + create_time_ = 0u; + update_time_ = 0u; + loaded_costtime_ = 0u; + trained_costtime_ = 0u; + built_costtime_ = 0u; + dumped_costtime_ = 0u; + } + + private: + //! Members + size_t revision_id_{0u}; + + std::atomic trained_count_{0u}; + std::atomic built_count_{0u}; + std::atomic dumped_count_{0u}; + std::atomic loaded_count_{0u}; + std::atomic added_count_{0u}; + std::atomic discarded_count_{0u}; + std::atomic updated_count_{0u}; + std::atomic deleted_count_{0u}; + + std::atomic index_size_{0u}; + std::atomic dumped_size_{0u}; + std::atomic check_point_{0u}; + + std::atomic create_time_{0u}; + std::atomic update_time_{0u}; + std::atomic loaded_costtime_{0u}; + std::atomic trained_costtime_{0u}; + std::atomic built_costtime_{0u}; + std::atomic dumped_costtime_{0u}; + }; + + //! Constructor + IndexRunner() = default; + + //! Destructor + virtual ~IndexRunner() = default; + + //! Retrieve statistics + virtual const Stats &stats(void) const = 0; + + //! Cleanup Searcher + virtual int cleanup() = 0; + + //! Unload Searcher + virtual int unload() { + return IndexError_NotImplemented; + } + + //! Print debug info + virtual void print_debug_info() {}; + + //! Create a searcher context + virtual Context::Pointer create_context(void) const { + return Context::Pointer(); + } + + //! Create a searcher provider + virtual Provider::Pointer create_provider(void) const { + return Provider::Pointer(); + } + + //! Create a searcher sparse provider + virtual SparseProvider::Pointer create_sparse_provider(void) const { + return SparseProvider::Pointer(); + } + + //! Get vector by key + virtual const void *get_vector(uint64_t /*key*/) const { + return nullptr; + } + + virtual int get_vector(const uint64_t /*key*/, + IndexStorage::MemoryBlock & /*block*/) const { + return IndexError_NotImplemented; + } + + //! Fetch vector by id + virtual const void *get_vector_by_id(uint32_t /*id*/) const { + return nullptr; + } + + virtual int get_vector_by_id(const uint32_t /*id*/, + IndexStorage::MemoryBlock & /*block*/) const { + return IndexError_NotImplemented; + } + + virtual int get_vector_by_key(const uint64_t /*key*/, + IndexStorage::MemoryBlock & /*block*/) const { + return IndexError_NotImplemented; + } + + //! Get vector by key + virtual int get_sparse_vector(uint64_t /*key*/, uint32_t * /*sparse_count*/, + std::string * /*sparse_indices_buffer*/, + std::string * /*sparse_values_buffer*/) const { + return IndexError_NotImplemented; + } + + //! Fetch vector by id + virtual int get_sparse_vector_by_id( + uint32_t /*id*/, uint32_t * /*sparse_count*/, + std::string * /*sparse_indices_buffer*/, + std::string * /*sparse_values_buffer*/) const { + return IndexError_NotImplemented; + } + + //! Add a vector into index + virtual int add_impl(uint64_t /*key*/, const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, + Context::Pointer & /*context*/) { + return IndexError_NotImplemented; + } + + //! Add a vector with id into index + virtual int add_with_id_impl(uint32_t /*id*/, const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, + Context::Pointer & /*context*/) { + return IndexError_NotImplemented; + } + + //! Similarity search + virtual int search_impl(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, + Context::Pointer & /*context*/) const { + return IndexError_NotImplemented; + } + //! Similarity search + virtual int search_impl(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, uint32_t /*count*/, + Context::Pointer & /*context*/) const { + return IndexError_NotImplemented; + } + + virtual 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 { + return IndexError_NotImplemented; + } + + //! Similarity search with sparse inputs + virtual int search_impl(const uint32_t /*sparse_count*/, + const uint32_t * /*sparse_indices*/, + const void * /*sparse_query*/, + const IndexQueryMeta & /*qmeta*/, + Context::Pointer & /*context*/) const { + return IndexError_NotImplemented; + } + + //! Similarity brute force search + virtual int search_bf_impl(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, + Context::Pointer & /*context*/) const { + return IndexError_NotImplemented; + } + //! Similarity brute force search + virtual int search_bf_impl(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, + uint32_t /*count*/, + Context::Pointer & /*context*/) const { + return IndexError_NotImplemented; + } + + //! Add a vector into index with dense and sparse inputs + virtual 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 */) { + return IndexError_NotImplemented; + } + + //! Add a vector with id into index + virtual 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 */) { + return IndexError_NotImplemented; + } + + //! Bruteforce search with sparse inputs + virtual 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 { + return IndexError_NotImplemented; + } + + //! Bruteforce search with sparse inputs + virtual 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 { + return IndexError_NotImplemented; + } + + //! Similarity brute force search by primary keys + virtual int search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, Context::Pointer &context) const { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + //! Similarity brute force search by primary keys + virtual 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 { + return IndexError_NotImplemented; + } + + //! Linear search by primary keys with dense and sparse inputs + virtual 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 */, + Context::Pointer & /* context */) const { + return IndexError_NotImplemented; + } + + //! Linear search by primary keys with dense and sparse inputs + virtual 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 */, + Context::Pointer & /* context */) const { + return IndexError_NotImplemented; + } + + //! Linear search by primary keys with dense and sparse inputs + virtual int search_bf_by_p_keys_impl( + const void * /*dense_query*/, const uint32_t /*sparse_count*/, + const uint32_t * /*sparse_indices*/, const void * /*sparse_query*/, + const std::vector> & /*p_keys*/, + const IndexQueryMeta & /*qmeta*/, Context::Pointer & /*context*/) const { + return IndexError_NotImplemented; + } + + //! Linear search by primary keys with dense and sparse inputs + virtual int search_bf_by_p_keys_impl( + const void * /*dense_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*/, + Context::Pointer & /*context*/) const { + return IndexError_NotImplemented; + } + + //! Update the vector in index + virtual int update_impl(uint64_t /*key*/, const void * /*query*/, + const IndexQueryMeta & /*qmeta*/, + Context::Pointer & /*context*/) { + return IndexError_NotImplemented; + } + + //! Delete the vector in index + virtual int remove_impl(uint64_t /*key*/, Context::Pointer & /*context*/) { + return IndexError_NotImplemented; + } + + //! Optimize the index + virtual int optimize_impl(IndexThreads::Pointer) { + return IndexError_NotImplemented; + } + + //! Delete the vector in index + int remove(uint64_t key, Context::Pointer &context) { + return this->remove_impl(key, context); + } + + //! Optimize the index + int optimize(IndexThreads::Pointer threads) { + return this->optimize_impl(threads); + } + + //! Train the data + virtual int train(IndexHolder::Pointer holder) { + return this->train(nullptr, std::move(holder)); + } + + //! Train the data + virtual int train(IndexThreads::Pointer /*threads*/, + IndexHolder::Pointer /*holder*/) { + return IndexError_NotImplemented; + } + + //! Train the data + virtual int train(const IndexTrainer::Pointer & /*trainer*/) { + return IndexError_NotImplemented; + } + + //! Train the data + virtual int train(IndexSparseHolder::Pointer holder) { + return this->train(nullptr, std::move(holder)); + } + + //! Train the data + virtual int train(IndexThreads::Pointer /*threads*/, + IndexSparseHolder::Pointer /*holder*/) { + return IndexError_NotImplemented; + } + + //! Build the index + virtual int build(IndexHolder::Pointer holder) { + return this->build(nullptr, std::move(holder)); + } + + //! Build the index + virtual int build(IndexThreads::Pointer /*threads*/, + IndexHolder::Pointer /*holder*/) { + return IndexError_NotImplemented; + } + + //! Build the index + virtual int build(IndexSparseHolder::Pointer holder) { + return this->build(nullptr, std::move(holder)); + } + + //! Build the index + virtual int build(IndexThreads::Pointer /*threads*/, + IndexSparseHolder::Pointer /*holder*/) { + return IndexError_NotImplemented; + } + + //! Build the index with indptr format + virtual int build(size_t count, const uint64_t *keys, + const uint64_t *sparse_indptr, + const uint32_t *sparse_indices, const void *sparse_data) { + return this->build(nullptr, count, keys, sparse_indptr, sparse_indices, + sparse_data); + } + + virtual int build(const IndexQueryMeta &qmeta, size_t count, + const uint64_t *keys, const uint64_t *sparse_indptr, + const uint32_t *sparse_indices, const void *sparse_data) { + return this->build(nullptr, qmeta, count, keys, sparse_indptr, + sparse_indices, sparse_data); + } + + //! Build the index with indptr format + virtual 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*/) { + return IndexError_NotImplemented; + } + + //! Build the index with indptr format + virtual 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*/) { + return IndexError_NotImplemented; + } + + //! Dump index into storage + virtual int dump(const IndexDumper::Pointer & /*dumper*/) { + return IndexError_NotImplemented; + } +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_searcher.h b/src/core/framework/index_searcher.h new file mode 100644 index 0000000..0349bd8 --- /dev/null +++ b/src/core/framework/index_searcher.h @@ -0,0 +1,57 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_context.h" +#include "index_meta.h" +#include "index_metric.h" +#include "index_module.h" +#include "index_provider.h" +#include "index_runner.h" +#include "index_stats.h" + +namespace zvec { +namespace core { + +/*! Index Searcher + */ +class IndexSearcher : public IndexRunner { + public: + //! Index Searcher Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + IndexSearcher() = default; + + //! Destructor + virtual ~IndexSearcher() = default; + + //! Initialize Searcher + virtual int init(const ailego::Params & /*params*/) = 0; + + //! Retrieve meta of index + virtual const IndexMeta &meta(void) const = 0; + + //! Retrieve params of index + virtual const ailego::Params ¶ms(void) const = 0; + + virtual int load(IndexStorage::Pointer /*container*/, + IndexMetric::Pointer /*metric*/) { + return IndexError_NotImplemented; + } +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_segment_storage.h b/src/core/framework/index_segment_storage.h new file mode 100644 index 0000000..0a8e121 --- /dev/null +++ b/src/core/framework/index_segment_storage.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 "ailego/container/params.h" +#include "index_module.h" +#include "index_storage.h" +#include "index_unpacker.h" + +namespace zvec { +namespace core { + +/*! Index Segment Container + */ +class IndexSegmentStorage : public IndexStorage { + public: + //! Index Segment Container Pointer + typedef std::shared_ptr Pointer; + + /*! Index Container Segment + */ + class Segment : public IndexStorage::Segment, + public std::enable_shared_from_this { + public: + //! Index Container Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + Segment(const Segment &rhs) = delete; + + //! Constructor + Segment(const IndexStorage::Segment::Pointer &parent, + const IndexUnpacker::SegmentMeta &segment) + : data_offset_(segment.data_offset()), + data_size_(segment.data_size()), + padding_size_(segment.padding_size()), + region_size_(segment.data_size() + segment.padding_size()), + data_crc_(segment.data_crc()), + parent_(parent->clone()) {} + + //! Destructor + virtual ~Segment(void) {} + + //! Retrieve size of data + size_t data_size(void) const override { + return data_size_; + } + + //! Retrieve crc of data + uint32_t data_crc(void) const override { + return data_crc_; + } + + //! Retrieve size of padding + size_t padding_size(void) const override { + return padding_size_; + } + + size_t capacity(void) const override { + return region_size_; + } + + //! Fetch data from segment (with own buffer) + size_t fetch(size_t offset, void *buf, size_t len) const override { + return parent_->fetch(data_offset_ + offset, buf, len); + } + + //! Read data from segment + size_t read(size_t offset, const void **data, size_t len) override { + return parent_->read(data_offset_ + offset, data, len); + } + + size_t read(size_t offset, MemoryBlock &data, size_t len) override { + const void **data_ptr = nullptr; + size_t ret = parent_->read(data_offset_ + offset, data_ptr, len); + data.reset((void *)*data_ptr); + return ret; + } + + //! Read data from segment + bool read(SegmentData *iovec, size_t count) override { + for (SegmentData *it = iovec, *end = iovec + count; it != end; ++it) { + it->offset += data_offset_; + } + bool success = parent_->read(iovec, count); + for (SegmentData *it = iovec, *end = iovec + count; it != end; ++it) { + it->offset -= data_offset_; + } + return success; + } + + size_t write(size_t, const void *, size_t) override { + return IndexError_NotImplemented; + } + + size_t resize(size_t) override { + return IndexError_NotImplemented; + } + + void update_data_crc(uint32_t) override { + return; + } + + //! Clone the segment + IndexStorage::Segment::Pointer clone(void) override { + return shared_from_this(); + } + + private: + size_t data_offset_{0u}; + size_t data_size_{0u}; + size_t padding_size_{0u}; + size_t region_size_{0u}; + uint32_t data_crc_{0u}; + IndexStorage::Segment::Pointer parent_{nullptr}; + }; + + //! Constructor + IndexSegmentStorage(IndexStorage::Segment::Pointer &&seg) + : parent_(std::move(seg)) {} + + //! Constructor + IndexSegmentStorage(const IndexStorage::Segment::Pointer &seg) + : parent_(seg) {} + + //! Destructor + virtual ~IndexSegmentStorage(void) {} + + //! Initialize container + int init(const ailego::Params &) override { + return 0; + } + + //! Cleanup container + int cleanup(void) override { + return 0; + } + + //! Load the current segment, ignore path + int open(const std::string &, bool) override { + if (!parent_) { + LOG_ERROR("Failed to load an empty segment"); + return IndexError_NoReady; + } + + auto read_data = [this](size_t offset, const void **data, size_t len) { + return this->parent_->read(offset, data, len); + }; + + IndexUnpacker unpacker; + if (!unpacker.unpack(read_data, parent_->data_size(), false)) { + LOG_ERROR("Failed to unpack segment data"); + return IndexError_UnpackIndex; + } + segments_ = std::move(*unpacker.mutable_segments()); + magic_ = unpacker.magic(); + return 0; + } + + //! Retrieve a segment by id + IndexStorage::Segment::Pointer get(const std::string &id, int) override { + if (!parent_) { + return IndexStorage::Segment::Pointer(); + } + auto it = segments_.find(id); + if (it == segments_.end()) { + return IndexStorage::Segment::Pointer(); + } + return std::make_shared(parent_, it->second); + } + + //! Test if it a segment exists + bool has(const std::string &id) const override { + return (segments_.find(id) != segments_.end()); + } + + //! Retrieve all segments + std::map get_all( + void) const override { + std::map result; + if (parent_) { + for (const auto &it : segments_) { + result.emplace(it.first, std::make_shared( + parent_, it.second)); + } + } + return result; + } + + //! Unload all indexes + int close(void) override { + parent_ = nullptr; + segments_.clear(); + return 0; + } + + //! Retrieve magic number of index + uint32_t magic(void) const override { + return magic_; + } + + int flush(void) override { + return IndexError_NotImplemented; + } + + int append(const std::string & /*id*/, size_t /*size*/) override { + return IndexError_NotImplemented; + } + + void refresh(uint64_t) override { + return; + } + + uint64_t check_point(void) const override { + return 0; + } + + private: + uint32_t magic_{0}; + std::map segments_{}; + IndexStorage::Segment::Pointer parent_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_stats.h b/src/core/framework/index_stats.h new file mode 100644 index 0000000..868ec51 --- /dev/null +++ b/src/core/framework/index_stats.h @@ -0,0 +1,68 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "ailego/container/params.h" + +namespace zvec { +namespace core { + +/*! Index Stats + */ +class IndexStats { + public: + //! Test if the element is exist + bool has_attribute(const std::string &key) const { + return attributes_.has(key); + } + + //! Set the value of key in T + template + bool set_attribute(const std::string &key, T &&val) { + return attributes_.set(key, std::forward(val)); + } + + //! Retrieve attribute with key + template + bool get_attribute(const std::string &key, T *out) const { + return attributes_.get(key, out); + } + + //! Erase the pair via a key + bool erase_attribute(const std::string &key) { + return attributes_.erase(key); + } + + //! Clear the attributes + void clear_attributes(void) { + attributes_.clear(); + } + + //! Retrieve attributes + const ailego::Params &attributes(void) const { + return attributes_; + } + + //! Retrieve mutable attributes + ailego::Params *mutable_attributes(void) { + return &attributes_; + } + + private: + //! Members + ailego::Params attributes_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_storage.h b/src/core/framework/index_storage.h new file mode 100644 index 0000000..ca3e98e --- /dev/null +++ b/src/core/framework/index_storage.h @@ -0,0 +1,271 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_error.h" +#include "index_module.h" + +namespace zvec { +namespace core { + +/*! Index Storage + */ +class IndexStorage : public IndexModule { + public: + //! Index Storage Pointer + typedef std::shared_ptr Pointer; + + struct MemoryBlock { + enum MemoryBlockType { + MBT_UNKNOWN = 0, + MBT_MMAP = 1, + MBT_BUFFERPOOL = 2, + }; + + MemoryBlock() {} + MemoryBlock(ailego::BufferHandle::Pointer &&buffer_handle) + : type_(MemoryBlockType::MBT_BUFFERPOOL), + buffer_handle_(std::move(buffer_handle)) { + data_ = buffer_handle_->pin_vector_data(); + } + MemoryBlock(void *data) : type_(MemoryBlockType::MBT_MMAP), data_(data) {} + + MemoryBlock(MemoryBlock &rhs) { + switch (rhs.type_) { + case MemoryBlockType::MBT_MMAP: + this->reset(rhs.data_); + break; + case MemoryBlockType::MBT_BUFFERPOOL: + this->reset(rhs.buffer_handle_); + break; + default: + break; + } + } + + MemoryBlock(MemoryBlock &&rhs) { + switch (rhs.type_) { + case MemoryBlockType::MBT_MMAP: + this->reset(std::move(rhs.data_)); + break; + case MemoryBlockType::MBT_BUFFERPOOL: + this->reset(std::move(rhs.buffer_handle_)); + break; + default: + break; + } + } + + MemoryBlock &operator=(MemoryBlock &rhs) { + if (this != &rhs) { + switch (rhs.type_) { + case MemoryBlockType::MBT_MMAP: + this->reset(rhs.data_); + break; + case MemoryBlockType::MBT_BUFFERPOOL: + this->reset(rhs.buffer_handle_); + break; + default: + break; + } + } + return *this; + } + + MemoryBlock &operator=(MemoryBlock &&rhs) { + if (this != &rhs) { + switch (rhs.type_) { + case MemoryBlockType::MBT_MMAP: + this->reset(std::move(rhs.data_)); + break; + case MemoryBlockType::MBT_BUFFERPOOL: + this->reset(std::move(rhs.buffer_handle_)); + break; + default: + break; + } + } + return *this; + } + + ~MemoryBlock() { + switch (type_) { + case MemoryBlockType::MBT_MMAP: + break; + case MemoryBlockType::MBT_BUFFERPOOL: + if (buffer_handle_) { + buffer_handle_->unpin_vector_data(); + // buffer_handle_.reset(); + } + break; + default: + break; + } + data_ = nullptr; + } + + const void *data() const { + return data_; + } + + void reset(ailego::BufferHandle::Pointer &buffer_handle) { + if (type_ == MemoryBlockType::MBT_BUFFERPOOL) { + buffer_handle_->unpin_vector_data(); + buffer_handle_.reset(); + } + type_ = MemoryBlockType::MBT_BUFFERPOOL; + if (buffer_handle) { + buffer_handle_.reset(buffer_handle.release()); + } + data_ = buffer_handle_->pin_vector_data(); + } + + void reset(ailego::BufferHandle::Pointer &&buffer_handle) { + if (type_ == MemoryBlockType::MBT_BUFFERPOOL) { + buffer_handle_->unpin_vector_data(); + buffer_handle_.reset(); + } + type_ = MemoryBlockType::MBT_BUFFERPOOL; + if (buffer_handle) { + buffer_handle_ = std::move(buffer_handle); + } + data_ = buffer_handle_->pin_vector_data(); + } + + void reset(void *data) { + if (type_ == MemoryBlockType::MBT_BUFFERPOOL) { + buffer_handle_->unpin_vector_data(); + buffer_handle_.reset(); + } + type_ = MemoryBlockType::MBT_MMAP; + data_ = data; + } + + MemoryBlockType type_{MBT_UNKNOWN}; + void *data_{nullptr}; + mutable ailego::BufferHandle::Pointer buffer_handle_{nullptr}; + }; + + struct SegmentData { + //! Constructor + SegmentData(void) : offset(0u), length(0u), data(nullptr) {} + + //! Constructor + SegmentData(size_t off, size_t len) + : offset(off), length(len), data(nullptr) {} + + //! Members + size_t offset; + size_t length; + const void *data; + }; + + /*! Index Storage Segment + */ + struct Segment { + //! Index Storage Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~Segment(void) {} + + //! Retrieve size of data + virtual size_t data_size(void) const = 0; + + //! Retrieve crc of data + virtual uint32_t data_crc(void) const = 0; + + //! Retrieve size of padding + virtual size_t padding_size(void) const = 0; + + //! Retrieve capacity of segment + virtual size_t capacity(void) const = 0; + + //! Fetch data from segment (with own buffer) + virtual size_t fetch(size_t offset, void *buf, size_t len) const = 0; + + //! Read data from segment + virtual size_t read(size_t offset, const void **data, size_t len) = 0; + + virtual size_t read(size_t offset, MemoryBlock &data, size_t len) = 0; + + virtual bool read(SegmentData *, size_t) { + return false; + } + + //! Write data into the storage with offset + virtual size_t write(size_t offset, const void *data, size_t len) = 0; + + //! Resize size of data + virtual size_t resize(size_t size) = 0; + + //! Update crc of data + virtual void update_data_crc(uint32_t crc) = 0; + + //! Clone the segment + virtual Pointer clone(void) = 0; + }; + + //! Destructor + virtual ~IndexStorage(void) {} + + //! Initialize storage + virtual int init(const ailego::Params ¶ms) = 0; + + //! Cleanup storage + virtual int cleanup(void) = 0; + + //! Open storage + virtual int open(const std::string &path, bool create) = 0; + + //! Flush storage + virtual int flush(void) = 0; + + //! Close storage + virtual int close(void) = 0; + + //! Append a segment into storage + virtual int append(const std::string &id, size_t size) = 0; + + //! Refresh meta information (checksum, update time, etc.) + virtual void refresh(uint64_t check_point) = 0; + + //! Retrieve check point of storage + virtual uint64_t check_point(void) const = 0; + + //! Retrieve a segment by id + virtual Segment::Pointer get(const std::string &id, int level = -1) = 0; + + virtual std::map get_all(void) const { + // LOG_ERROR("get_all() Not Implemented"); + std::map result; + return result; + } + + //! Test if it a segment exists + virtual bool has(const std::string &id) const = 0; + + //! Retrieve magic number of index + virtual uint32_t magic(void) const = 0; + + //! huge page + virtual bool isHugePage(void) const { + return false; + } +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_streamer.h b/src/core/framework/index_streamer.h new file mode 100644 index 0000000..7335154 --- /dev/null +++ b/src/core/framework/index_streamer.h @@ -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. +#pragma once + +#include +#include "index_context.h" +#include "index_helper.h" +#include "index_provider.h" +#include "index_runner.h" +#include "index_stats.h" + +namespace zvec { +namespace core { + +/*! Index Streamer + */ +class IndexStreamer : public IndexRunner { + public: + //! Index Streamer Pointer + typedef std::shared_ptr Pointer; + + //! Destructor + virtual ~IndexStreamer(void) = default; + + //! Initialize the builder + virtual int init(const IndexMeta & /*meta*/, + const ailego::Params & /*params*/) { + return IndexError_NotImplemented; + } + + //! Open a index from storage + virtual int open(IndexStorage::Pointer stg) = 0; + + //! Flush index + virtual int flush(uint64_t check_point) = 0; + + //! Close index + virtual int close(void) = 0; + + //! Retrieve meta of index + virtual const IndexMeta &meta(void) const = 0; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_threads.h b/src/core/framework/index_threads.h new file mode 100644 index 0000000..5b9fa0d --- /dev/null +++ b/src/core/framework/index_threads.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 +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +/*! Index Threads + * Index ThreadPool maintains multiple threads resources to execute the tasks + * concurrently + */ +class IndexThreads { + public: + using Pointer = std::shared_ptr; + + /*! Threads Task Group + * Manage of a group of sub-tasks which can be seen as a big task, + * so we can wait all sub-tasks finished, or get the status of them + */ + class TaskGroup { + public: + using Pointer = std::shared_ptr; + + //! Destructor + virtual ~TaskGroup(void) {} + + //! Submit a task to be executed asynchronous + virtual void submit(ailego::ClosureHandler &&task) = 0; + + //! Check if the group is finished + virtual bool is_finished(void) const = 0; + + //! Wait until all tasks in group finished + virtual void wait_finish(void) = 0; + }; + + //! Destructor + virtual ~IndexThreads(void) {} + + //! Retrieve thread count in pool + virtual size_t count(void) const = 0; + + //! Stop all threads + virtual void stop(void) = 0; + + //! Submit a task to be executed asynchronous + virtual void submit(ailego::ClosureHandler &&task) = 0; + + //! Make a task group + virtual TaskGroup::Pointer make_group(void) = 0; + + //! Get the current work thread index + virtual int indexof_this(void) const = 0; +}; + +/*! Single Queue Index Threads + */ +class SingleQueueIndexThreads : public IndexThreads { + public: + /*! Single Queue Index Threads Task Group + */ + class SingleQueueTaskGroup : public TaskGroup { + public: + using Pointer = std::shared_ptr; + + //! Constructor + explicit SingleQueueTaskGroup( + ailego::ThreadPool::TaskGroup::Pointer task_group) + : task_group_(std::move(task_group)) {} + + //! Submit a task to be executed asynchronous + void submit(ailego::ClosureHandler &&task) override { + while (task_group_->pending_count() >= kMaxQueueSize) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + task_group_->enqueue_and_wake(std::move(task)); + } + + //! Check if the group is finished + bool is_finished(void) const override { + return task_group_->is_finished(); + } + + //! Wait until all tasks in group finished + void wait_finish(void) override { + return task_group_->wait_finish(); + } + + private: + //! Members + ailego::ThreadPool::TaskGroup::Pointer task_group_{}; + }; + + //! Constructor + SingleQueueIndexThreads(uint32_t size, bool binding) + : pool_( + size > 0 ? size : std::max(std::thread::hardware_concurrency(), 1u), + binding) {} + + //! Constructor + explicit SingleQueueIndexThreads(bool binding) + : SingleQueueIndexThreads(0, binding) {} + + //! Constructor + SingleQueueIndexThreads(void) : SingleQueueIndexThreads{false} {} + + //! Destructor + virtual ~SingleQueueIndexThreads(void) {} + + //! Retrieve thread count in pool + size_t count(void) const override { + return pool_.count(); + } + + //! Stop all threads + void stop(void) override { + pool_.stop(); + } + + //! Submit a task to be executed asynchronous + void submit(ailego::ClosureHandler &&task) override { + while (pool_.pending_count() >= kMaxQueueSize) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + pool_.enqueue_and_wake(std::move(task)); + } + + //! Make a task group + TaskGroup::Pointer make_group(void) override { + return std::make_shared(pool_.make_group()); + } + + //! Get the current work thread index + int indexof_this(void) const override { + return pool_.indexof_this(); + } + + private: + static constexpr size_t kMaxQueueSize = 4096u; + + //! Disable them + SingleQueueIndexThreads(const SingleQueueIndexThreads &) = delete; + SingleQueueIndexThreads(SingleQueueIndexThreads &&) = delete; + SingleQueueIndexThreads &operator=(const SingleQueueIndexThreads &) = delete; + + //! Members + ailego::ThreadPool pool_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_trainer.h b/src/core/framework/index_trainer.h new file mode 100644 index 0000000..6ad6fdc --- /dev/null +++ b/src/core/framework/index_trainer.h @@ -0,0 +1,111 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_bundle.h" +#include "index_dumper.h" +#include "index_error.h" +#include "index_holder.h" +#include "index_meta.h" +#include "index_stats.h" +#include "index_storage.h" +#include "index_threads.h" + +namespace zvec { +namespace core { + +/*! Index Trainer + */ +class IndexTrainer : public IndexModule { + public: + //! Index Trainer Pointer + typedef std::shared_ptr Pointer; + + /*! Index Trainer Stats + */ + class Stats : public IndexStats { + public: + //! Set count of documents trained + void set_trained_count(size_t count) { + trained_count_ = count; + } + + //! Set count of documents discarded + void set_discarded_count(size_t count) { + discarded_count_ = count; + } + + //! Set time cost of documents trained + void set_trained_costtime(uint64_t cost) { + trained_costtime_ = cost; + } + + //! Retrieve count of documents trained + size_t trained_count(void) const { + return trained_count_; + } + + //! Retrieve count of documents discarded + size_t discarded_count(void) const { + return discarded_count_; + } + + //! Retrieve time cost of documents trained + uint64_t trained_costtime(void) const { + return trained_costtime_; + } + + private: + //! Members + size_t trained_count_{0u}; + size_t discarded_count_{0u}; + uint64_t trained_costtime_{0u}; + }; + + //! Destructor + virtual ~IndexTrainer(void) {} + + //! Initialize Trainer + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) = 0; + + //! Cleanup Trainer + virtual int cleanup(void) = 0; + + //! Train the data + virtual int train(IndexHolder::Pointer holder) { + return this->train(nullptr, holder); + } + + //! Train the data + virtual int train(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) = 0; + + //! Load index from container + virtual int load(IndexStorage::Pointer cntr) = 0; + + //! Dump index into storage + virtual int dump(const IndexDumper::Pointer &dumper) = 0; + + //! Retrieve Index Meta + virtual const IndexMeta &meta(void) const = 0; + + //! Retrieve statistics + virtual const Stats &stats(void) const = 0; + + //! Retrieve the output indexes + virtual IndexBundle::Pointer indexes(void) const = 0; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/framework/index_unpacker.h b/src/core/framework/index_unpacker.h new file mode 100644 index 0000000..8d607a4 --- /dev/null +++ b/src/core/framework/index_unpacker.h @@ -0,0 +1,304 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index_error.h" +#include "index_format.h" +#include "index_logger.h" + +namespace zvec { +namespace core { + +/*! Index Unpacker + */ +class IndexUnpacker { + public: + /*! Index Unpacker Segment Meta + */ + class SegmentMeta { + public: + //! Constructor + SegmentMeta(size_t offset, size_t dsz, size_t psz, uint32_t crc) + : data_offset_(offset), + data_size_(dsz), + padding_size_(psz), + data_crc_(crc) {} + + //! Retrieve offset of data + size_t data_offset(void) const { + return data_offset_; + } + + //! Retrieve size of data + size_t data_size(void) const { + return data_size_; + } + + //! Retrieve crc of data + uint32_t data_crc(void) const { + return data_crc_; + } + + //! Retrieve size of padding + size_t padding_size(void) const { + return padding_size_; + } + + private: + size_t data_offset_{0}; + size_t data_size_{0}; + size_t padding_size_{0}; + uint32_t data_crc_{0}; + }; + + //! Reset the unpacker + void reset(void) { + segments_.clear(); + } + + //! Retrieve segments of index package + const std::map &segments(void) const { + return segments_; + } + + //! Retrieve magic number of index + uint32_t magic(void) const { + return header_.magic; + } + + //! Retrieve header of index package + const IndexFormat::MetaHeader &header(void) const { + return header_; + } + + //! Retrieve footer of index package + const IndexFormat::MetaFooter &footer(void) const { + return footer_; + } + + //! Retrieve version information + const std::string &version(void) const { + return version_; + } + + //! Retrieve mutable segments of index package + std::map *mutable_segments(void) { + return &segments_; + } + + //! Unpack index data + template + bool unpack(TFunc read_data, size_t total, bool checksum) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + + if (!this->unpack_header(read_data)) { + LOG_ERROR("Failed to unpack index header"); + return false; + } + if (!this->unpack_footer(read_data, total)) { + LOG_ERROR("Failed to unpack index footer"); + return false; + } + if (!this->unpack_segments(read_data, total)) { + LOG_ERROR("Failed to unpack index segments' meta"); + return false; + } + if (checksum && !this->validate_checksum(read_data)) { + LOG_ERROR("Failed to validate checksum of index content"); + return false; + } + if (!this->unpack_version(read_data)) { + LOG_ERROR("Failed to unpack index version"); + return false; + } + return true; + } + + //! Unpack index header + template + bool unpack_header(TFunc read_data) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + const void *data = nullptr; + if (read_data(0u, &data, sizeof(header_)) != sizeof(header_)) { + return false; + } + + memcpy(&header_, data, sizeof(header_)); + if (header_.meta_header_size != sizeof(header_)) { + return false; + } + if (ailego::Crc32c::Hash(&header_, sizeof(header_), header_.header_crc) != + header_.header_crc) { + return false; + } + return true; + } + + //! Unpack index footer + template + bool unpack_footer(TFunc read_data, size_t total) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + if (header_.meta_footer_size != sizeof(footer_)) { + return false; + } + + size_t footer_offset = ((int32_t)header_.meta_footer_offset < 0) + ? total + (int32_t)header_.meta_footer_offset + : header_.meta_footer_offset; + if (footer_offset + sizeof(footer_) > total) { + return false; + } + + const void *data = nullptr; + if (read_data(footer_offset, &data, sizeof(footer_)) != sizeof(footer_)) { + return false; + } + + memcpy(&footer_, data, sizeof(footer_)); + if ((footer_.total_size != total) || + (footer_.content_size + footer_.content_padding_size + + header_.content_offset > + total)) { + return false; + } + if (ailego::Crc32c::Hash(&footer_, sizeof(footer_), footer_.footer_crc) != + footer_.footer_crc) { + return false; + } + return true; + } + + //! Unpack segments' meta + template + bool unpack_segments(TFunc read_data, size_t total) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + if (sizeof(IndexFormat::SegmentMeta) * footer_.segment_count > + footer_.segments_meta_size) { + return false; + } + size_t offset = ((int32_t)header_.meta_footer_offset < 0) + ? total + (int32_t)header_.meta_footer_offset + : header_.meta_footer_offset; + if (offset < footer_.segments_meta_size || offset > total) { + return false; + } + offset -= footer_.segments_meta_size; + + const void *data = nullptr; + if (read_data(offset, &data, footer_.segments_meta_size) != + footer_.segments_meta_size) { + return false; + } + if (ailego::Crc32c::Hash(data, footer_.segments_meta_size, 0u) != + footer_.segments_meta_crc) { + return false; + } + + IndexFormat::SegmentMeta *seg = (IndexFormat::SegmentMeta *)data; + for (size_t i = 0; i < footer_.segment_count; ++i, ++seg) { + if (seg->segment_id_offset > footer_.segments_meta_size) { + return false; + } + if (seg->data_index > footer_.content_size) { + return false; + } + if (seg->data_index + seg->data_size > footer_.content_size) { + return false; + } + segments_.emplace( + std::string(reinterpret_cast(data) + + seg->segment_id_offset), + SegmentMeta(seg->data_index + header_.content_offset, seg->data_size, + seg->padding_size, seg->data_crc)); + } + return true; + } + + //! Unpack index version + template + bool unpack_version(TFunc read_data) { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + + auto it = segments_.find("IndexVersion"); + if (it == segments_.end()) { + return false; + } + + const SegmentMeta &segment = it->second; + const void *data = nullptr; + + if (read_data(segment.data_offset(), &data, segment.data_size()) != + segment.data_size()) { + return false; + } + if (segment.data_crc() != 0u && + ailego::Crc32c::Hash(data, segment.data_size(), 0u) != + segment.data_crc()) { + return false; + } + version_.assign(reinterpret_cast(data), segment.data_size()); + return true; + } + + //! Validate checksum of content + template + bool validate_checksum(TFunc read_data) const { + static_assert(ailego::IsInvocableWithResult::value, + "Invocable function type"); + if (footer_.content_crc == 0) { + return true; + } + const size_t block_size = 4096u; + const void *data = nullptr; + uint32_t checksum = 0u; + size_t total = footer_.content_size; + size_t offset = sizeof(header_); + + while (total >= block_size) { + if (read_data(offset, &data, block_size) != block_size) { + return false; + } + checksum = ailego::Crc32c::Hash(data, block_size, checksum); + total -= block_size; + offset += block_size; + } + if (read_data(offset, &data, total) != total) { + return false; + } + checksum = ailego::Crc32c::Hash(data, total, checksum); + return (checksum == footer_.content_crc); + } + + private: + IndexFormat::MetaHeader header_; + IndexFormat::MetaFooter footer_; + std::string version_{}; + std::map segments_{}; +}; + +} // 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..cd206b0 --- /dev/null +++ b/src/core/framework/index_version.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 "index_version.h" +#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/framework/index_version.h b/src/core/framework/index_version.h new file mode 100644 index 0000000..0037667 --- /dev/null +++ b/src/core/framework/index_version.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 + +namespace zvec { +namespace core { + +/*! Index Version + */ +struct IndexVersion { + //! Retrieve the version number in string + static const char *String(void); + + //! Retrieve the detailed version information + static const char *Details(void); +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/interface/CMakeLists.txt b/src/core/interface/CMakeLists.txt new file mode 100644 index 0000000..819b89d --- /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_static sparsehash magic_enum + VERSION "${PROXIMA_ZVEC_VERSION}" +) diff --git a/src/core/interface/constants.h b/src/core/interface/constants.h new file mode 100644 index 0000000..aeb941a --- /dev/null +++ b/src/core/interface/constants.h @@ -0,0 +1,13 @@ +#pragma once + +#include + +namespace zvec::core_interface { + +constexpr static uint32_t kDefaultHnswEfConstruction = 500; +constexpr static uint32_t kDefaultHnswNeighborCnt = 100; + +constexpr static uint32_t kDefaultHnswEfSearch = 300; + + +} // namespace zvec::core_interface \ No newline at end of file diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc new file mode 100644 index 0000000..3b50d92 --- /dev/null +++ b/src/core/interface/index.cc @@ -0,0 +1,773 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "index.h" +#include "framework/index_error.h" +#include "framework/index_storage.h" + +namespace zvec::core_interface { + +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 (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::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; + default: + LOG_ERROR("Unsupported quantizer 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(); + 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); + 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; + } + + // converter/reformer/metric are created in IndexFactory::CreateIndex + // TODO: init + + // 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; +} + +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; + } + + if (is_sparse_) { + return _sparse_add(vector_data, doc_id, context); + } + return _dense_add(vector_data, doc_id, context); +} + + +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; + } + + 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"); + return core::IndexError_Runtime; + } + + if (is_sparse_) { + return _sparse_search(vector_data, search_param, result, context); + } + + // dense support refiner, but sparse doesn't + if (search_param->refiner_param == nullptr) { + return _dense_search(vector_data, search_param, result, context); + } else { + auto &reference_index = search_param->refiner_param->reference_index; + if (reference_index == nullptr) { + LOG_ERROR("Reference index is not set"); + 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"); + 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"); + 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); + + return reference_index->Search(vector_data, flat_search_param, result); + } +} + + +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(); + } + // TODO: group by + if (search_param->bf_pks != nullptr) { + // should we elimate 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; + } + } + result->doc_list_ = std::move(context->result()); + + if (metric_->support_normalize()) { + for (uint32_t i = 0; i < result->doc_list_.size(); ++i) { + metric_->normalize(result->doc_list_[i].mutable_score()); + } + } + if (reformer_) { + 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()) { + // TODO: use std::pmr to optimize memory allocation + result->reverted_vector_list_.resize(context->result().size()); + for (uint32_t i = 0; i < context->result().size(); ++i) { + std::string &reverted_vector = result->reverted_vector_list_[i]; + reverted_vector.resize(input_vector_meta_.dimension() * + input_vector_meta_.unit_size()); + if (reformer_->revert(context->result()[i].vector(), new_meta, + &reverted_vector) != 0) { + LOG_ERROR("Failed to revert vector"); + return core::IndexError_Runtime; + } + } + } + } + + 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; + } + } + result->doc_list_ = std::move(context->result()); + + if (metric_->support_normalize()) { + for (uint32_t i = 0; i < result->doc_list_.size(); ++i) { + metric_->normalize(result->doc_list_[i].mutable_score()); + } + } + if (reformer_) { + // TODO: no need to call reformer_->normalize() when sparse? + if (context->fetch_vector() && reformer_->need_revert()) { + // TODO: use std::pmr to optimize memory allocation + auto &result_doc_list = context->result(); + result->reverted_sparse_values_list_.resize(result_doc_list.size()); + for (uint32_t i = 0; i < result_doc_list.size(); ++i) { + auto &result_doc = result_doc_list[i].sparse_doc(); + std::string &reverted_sparse_values = + result->reverted_sparse_values_list_[i]; + 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"); + return core::IndexError_Runtime; + } + } + } + } + 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); +} + +} // namespace zvec::core_interface diff --git a/src/core/interface/index.h b/src/core/interface/index.h new file mode 100644 index 0000000..4b5c88c --- /dev/null +++ b/src/core/interface/index.h @@ -0,0 +1,321 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "core/framework/index_context.h" +#include "core/framework/index_converter.h" +#include "core/framework/index_factory.h" +#include "core/framework/index_filter.h" +#include "core/framework/index_holder.h" +#include "core/framework/index_meta.h" +#include "core/framework/index_metric.h" +#include "core/framework/index_reducer.h" +#include "core/framework/index_reformer.h" +#include "core/framework/index_searcher.h" +#include "core/framework/index_storage.h" +#include "core/mixed_reducer/mixed_reducer_params.h" +#include "index_param.h" +#include "index_param_builders.h" + +namespace zvec::core_interface { + +class IndexFactory; + +struct DenseVector { + const void *data; + // core::IndexQueryMeta meta; + // DenseVector(void *data) : data(data) { + // meta.set_meta_type(core::IndexMeta::MetaType::MT_DENSE); + // }; +}; + +struct SparseVector { + uint32_t count; + const void *indices; + const void *values; + + const uint32_t *get_indices() const { + return reinterpret_cast(indices); + } + + template + const T *get_values() const { + return reinterpret_cast(values); + } +}; + +struct VectorData { + std::variant vector; + + // DenseVector dense_vector; + // SparseVector sparse_vector; +}; + +// Used to pass mutable vectors +struct DenseVectorBuffer { + std::string data; // use string to manage memory +}; + +struct SparseVectorBuffer { + uint32_t count; + std::string indices; + std::string values; + + uint32_t *get_indices() { + return reinterpret_cast(indices.data()); + } + + template + T *get_values() { + return reinterpret_cast(values.data()); + } +}; + +struct VectorDataBuffer { + std::variant vector_buffer; +}; + + +struct SearchResult { + core::IndexDocumentList doc_list_; + // use string to manage memory + std::vector reverted_vector_list_{}; + std::vector reverted_sparse_values_list_{}; +}; + +// eliminate the pre-alloc of the context pool +thread_local static std::array() - 1) * 2> + _context_list; + +class Index { + public: + typedef std::shared_ptr Pointer; + virtual ~Index() = default; + + // static Index::Pointer Create(const BaseIndexParam ¶m); //IndexFactory + virtual int Open(const std::string &file_path, + StorageOptions storage_options); + int Close(); + int Flush(); + // virtual int Serialize(const std::string &file_path); + // virtual int Deserialize(const std::string &file_path); + + // // TODO: use holder + // virtual int Build() = 0; + virtual int Train() { + is_trained_ = true; + return 0; + } + + // virtual int Dump(const std::string &file_path) = 0; + virtual int Merge(const std::vector &indexes, + const IndexFilter &filter, + const MergeOptions &options = {}); + // TODO: static reduce + + virtual int Add(const VectorData &vector, uint32_t doc_id); + virtual int Fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer); + virtual int Search(const VectorData &query, + const BaseIndexQueryParam::Pointer &search_param, + SearchResult *result); + + virtual BaseIndexParam::Pointer GetParam() const { + return std::make_shared(param_); + } + + virtual bool IsTrained() const { + return is_trained_; + } + + uint32_t GetDocCount() const { + if (streamer_ == nullptr) { + return -1; + } + if (is_sparse_) { + return streamer_->create_sparse_provider()->count(); + } else { + return streamer_->create_provider()->count(); + } + } + + core::IndexStreamer::Pointer index_searcher() { + return streamer_; + } + + protected: + int _sparse_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer); + virtual int _dense_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer); + + int _sparse_add(const VectorData &vector, const uint32_t doc_id, + core::IndexContext::Pointer &context); + int _dense_add(const VectorData &vector, const uint32_t doc_id, + core::IndexContext::Pointer &context); + int _sparse_search(const VectorData &query, + const BaseIndexQueryParam::Pointer &search_param, + SearchResult *result, + core::IndexContext::Pointer &context); + int _dense_search(const VectorData &query, + const BaseIndexQueryParam::Pointer &search_param, + SearchResult *result, core::IndexContext::Pointer &context); + virtual int _prepare_for_search( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) = 0; + virtual int _get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param); + + protected: + friend class IndexFactory; + Index() = default; + int Init(const BaseIndexParam ¶m); + + + protected: + int ParseMetricName(const BaseIndexParam ¶m); + int CreateAndInitMetric(const BaseIndexParam ¶m); + int CreateAndInitConverterReformer(const QuantizerParam ¶m, + const BaseIndexParam &index_param); + virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) = 0; + + protected: + bool 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 &acquire_context() { + init_context(); + return _context_list[context_index_]; + } + + void release_context() { + // context_list_[get_context_index()]->reset(); + } + + protected: + bool is_trained_{false}; + + BaseIndexParam param_; + ailego::Params proxima_index_params_{}; + core::IndexMeta proxima_index_meta_{}; // IndexQueryMeta + other index config + core::IndexQueryMeta input_vector_meta_; // input + core::IndexQueryMeta streamer_vector_meta_; // after reformer.convert() + + core::IndexBuilder::Pointer builder_{}; + core::IndexStreamer::Pointer streamer_{}; + core::IndexReformer::Pointer reformer_{}; + core::IndexConverter::Pointer converter_{}; // for build() + core::IndexMetric::Pointer metric_{}; // to do normalization + + size_t context_index_; + core::IndexStorage::Pointer storage_{}; + + bool is_open_{false}; + bool is_sparse_{false}; + bool is_huge_page_{false}; + bool is_read_only_{false}; +}; + + +class FlatIndex : public Index { + public: + FlatIndex() = default; + // FlatIndex(const FlatIndexParam ¶m) : param_(param) {} + // FlatIndex(FlatIndexParam &¶m) : param(std::move(param)) {} + + + protected: + virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) override; + + virtual int _prepare_for_search( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) override; + + private: + FlatIndexParam param_{}; +}; + +class IVFIndex : public Index { + public: + IVFIndex() = default; + + protected: + virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) override; + + virtual int _prepare_for_search( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) override; + + virtual int Add(const VectorData &vector, uint32_t doc_id) override; + + virtual int Train() override; + + virtual int Open(const std::string &file_path, + StorageOptions storage_options) override; + + virtual int _dense_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer) override; + virtual int Merge(const std::vector &indexes, + const IndexFilter &filter, + const MergeOptions &options) override; + int GenerateHolder(); + + private: + IVFIndexParam param_{}; + std::mutex mutex_{}; + std::vector> doc_cache_; + core::IndexHolder::Pointer holder_{}; + std::string file_path_; +}; + + +class HNSWIndex : public Index { + public: + HNSWIndex() = default; + + protected: + virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) override; + + virtual int _prepare_for_search( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) override; + virtual int _get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param); + + + private: + HNSWIndexParam param_{}; +}; + + +} // 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..f924632 --- /dev/null +++ b/src/core/interface/index_factory.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. +// +// Created by wangjianning.wjn on 8/28/25. +// + +#include "core/framework/index_factory.h" +#include "ailego/container/params.h" +#include "core/framework/index_meta.h" +#include "utils/utils.h" +#include "index_factory.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 { + 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; + } + default: + LOG_ERROR("Unsupported index type: %s", + magic_enum::enum_name(index_type).data()); + return nullptr; + } +} + +} // namespace zvec::core_interface diff --git a/src/core/interface/index_factory.h b/src/core/interface/index_factory.h new file mode 100644 index 0000000..a5414a2 --- /dev/null +++ b/src/core/interface/index_factory.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. +// +// Created by wangjianning.wjn on 8/28/25. +// + +#ifndef ZVEC_INDEX_FACTORY_H +#define ZVEC_INDEX_FACTORY_H + +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include "index.h" +#include "index_param.h" + +namespace zvec::core_interface { + +// 索引的工厂类 +class IndexFactory { + public: + static Index::Pointer CreateAndInitIndex(const BaseIndexParam ¶m); + + static BaseIndexParam::Pointer DeserializeIndexParamFromJson( + const std::string &json_str); + + + static std::string QueryParamSerializeToJson( + const BaseIndexQueryParam ¶m); + + + template < + typename QueryParamType, + std::enable_if_t, + bool> = true> + static std::string QueryParamSerializeToJson(const QueryParamType ¶m, + bool omit_empty_value = false); + + template < + typename QueryParamType, + std::enable_if_t, + bool> = true> + static typename QueryParamType::Pointer QueryParamDeserializeFromJson( + const std::string &json_str); + + // register() -- Index class should have a `create` interface +}; + + +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; + 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)); + } + 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))); + } + + 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 , 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; + } + 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 { + 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; + } + } 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 { + LOG_ERROR("Unsupported index type: %s", + magic_enum::enum_name(index_type).data()); + return nullptr; + } + return param; + } +} + + +} // namespace zvec::core_interface + + +#endif // ZVEC_INDEX_FACTORY_H diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc new file mode 100644 index 0000000..394788b --- /dev/null +++ b/src/core/interface/index_param.cc @@ -0,0 +1,206 @@ +#include "index_param.h" +#include +#include + +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 (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)); + 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); + + 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); + + 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/index_param.h b/src/core/interface/index_param.h new file mode 100644 index 0000000..626e36b --- /dev/null +++ b/src/core/interface/index_param.h @@ -0,0 +1,340 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Created by wangjianning.wjn on 8/28/25. +// + +#ifndef ZVEC_INDEX_PARAM_H +#define ZVEC_INDEX_PARAM_H +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "core/framework/index_filter.h" +#include "core/framework/index_meta.h" +#include "utils/utils.h" +#include "constants.h" + +namespace zvec::core_interface { +#define MAX_DIMENSION 65536 +// #define MAX_EF_CONSTRUCTION 65536 +// #define MAX_EF_SEARCH 100 + +class IndexFactory; +class Index; +class BaseIndexParam; +class BaseIndexQueryParam; + +struct StorageOptions { + enum class StorageType { kNone, kMMAP, kMemory, kBufferPool }; + + StorageType type = StorageType::kNone; + bool create_new = false; + bool read_only = false; +}; + +struct MergeOptions { + uint32_t write_concurrency = 1; + ailego::ThreadPool *pool = nullptr; +}; + +using IndexMeta = core::IndexMeta; +using IndexQueryMeta = core::IndexQueryMeta; +using DataType = core::IndexMeta::DataType; +using IndexFilter = core::IndexFilter; + + +// 定义支持的索引类型 +enum class IndexType { + // to do: support factory's register, may change to + // `static constexpr std::string_view`, which may incur str comp overhead + kNone, + kFlat, + kIVF, // it's actual a two-layer index + kHNSW, +}; + +enum class IVFSearchMethod { kBF, kHNSW }; + +enum class MetricType { + kNone, + kL2sq, // Euclidean + kInnerProduct, + kCosine, + kMIPSL2sq // spherical? +}; + +enum class QuantizerType { + kNone, + kPQ, // Product Quantization + kQuickADC, // TODO: +refiner ? // should be a type of index? + kAQ, + kFP16, + kInt8, + kInt4 +}; + +struct SerializableBase { + std::string SerializeToJson(bool omit_empty_value = false) const { + return zvec::ailego::JsonValue(SerializeToJsonObject(omit_empty_value)) + .as_json_string() + .as_stl_string(); + } + + bool DeserializeFromJson(const std::string &json_str) { + ailego::JsonValue json_value; + if (!json_value.parse(json_str)) { + return false; + } + return DeserializeFromJsonObject(json_value.as_object()); + } + + protected: + virtual ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const = 0; + virtual bool DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) = 0; +}; + +// TODO: maybe a base class for quantizer? +struct QuantizerParam : public SerializableBase { + QuantizerType type = QuantizerType::kNone; + int num_subquantizers = 8; // M + int num_bits = 8; // bits per subquantizer + + // Constructors + // QuantizerParam() = default; + QuantizerParam(QuantizerType t = QuantizerType::kNone, int subquantizers = 8, + int bits = 8) + : type(t), num_subquantizers(subquantizers), num_bits(bits) {} + + + protected: + friend class BaseIndexParam; + virtual ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override { + ailego::JsonObject json_obj; + if (!omit_empty_value || type != QuantizerType::kNone) { + json_obj.set("type", + zvec::ailego::JsonValue(magic_enum::enum_name(type).data())); + } + return json_obj; + } + + virtual bool DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) override { + DESERIALIZE_ENUM_FIELD(json_obj, type, QuantizerType); + return true; + } +}; + +// preprocessor +enum class PreprocessorType { + kNone, + kPCA, + kOPQ, +}; + +struct PreprocessorParam { + PreprocessorType type = PreprocessorType::kNone; + + // Constructors + // PreprocessorParam() = default; + explicit PreprocessorParam(PreprocessorType t = PreprocessorType::kNone) + : type(t) {} +}; + +struct RefinerParam { + using Pointer = std::shared_ptr; + + float scale_factor_{0}; + std::shared_ptr reference_index = nullptr; +}; + +// --- Query Parameters (can be passed to search methods) --- +class BaseIndexQueryParam { + public: + using Pointer = std::shared_ptr; + + virtual ~BaseIndexQueryParam() = default; + + uint32_t topk = 10; + bool fetch_vector = false; + std::shared_ptr filter = nullptr; + std::shared_ptr> bf_pks = nullptr; + float radius = 0.0f; + bool is_linear = false; + RefinerParam::Pointer refiner_param = nullptr; + + virtual Pointer Clone() const = 0; +}; + +struct FlatQueryParam : public BaseIndexQueryParam { + using Pointer = std::shared_ptr; + + BaseIndexQueryParam::Pointer Clone() const override { + return std::make_shared(*this); + } +}; + +struct HNSWQueryParam : public BaseIndexQueryParam { + using Pointer = std::shared_ptr; + + uint32_t ef_search = kDefaultHnswEfSearch; + + BaseIndexQueryParam::Pointer Clone() const override { + return std::make_shared(*this); + } +}; + +struct IVFQueryParam : public BaseIndexQueryParam { + int nprobe = 10; + std::shared_ptr l1QueryParam = nullptr; + std::shared_ptr l2QueryParam = nullptr; + + using Pointer = std::shared_ptr; + + BaseIndexQueryParam::Pointer Clone() const override { + auto cloned_this = std::make_shared(*this); + cloned_this->l1QueryParam = l1QueryParam ? l1QueryParam->Clone() : nullptr; + cloned_this->l2QueryParam = l2QueryParam ? l2QueryParam->Clone() : nullptr; + return cloned_this; + } +}; + +// --- Construction Parameters --- +// template +class BaseIndexParam : public SerializableBase { + public: + using Pointer = std::shared_ptr; + + explicit BaseIndexParam(IndexType type = IndexType::kNone, + MetricType metric = MetricType::kL2sq, int dim = 0, + int ver = 0) + : index_type(type), metric_type(metric), dimension(dim), version(ver) {} + + virtual ~BaseIndexParam() = default; + + IndexType index_type = IndexType::kNone; + MetricType metric_type = MetricType::kL2sq; + int dimension = 0; // [1, MAX_DIMENSION] + int version = 0; // for compatibility + bool is_sparse = false; + bool is_huge_page = false; + DataType data_type = DataType::DT_UNDEFINED; + bool use_id_map = true; + + // IndexMeta meta; + ailego::Params params; + + // pipeline + PreprocessorParam preprocess_param; + QuantizerParam quantizer_param; + + BaseIndexQueryParam::Pointer default_query_param = nullptr; + // virtual std::shared_ptr GetDefaultQueryParam() const + // { + // return std::make_shared(); + // } + // + + protected: + virtual bool DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) override; + virtual ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override; +}; + +struct FlatIndexParam : public BaseIndexParam { + using Pointer = std::shared_ptr; + FlatIndexParam() : BaseIndexParam(IndexType::kFlat) {} + + IndexMeta::MajorOrder major_order = IndexMeta::MajorOrder::MO_ROW; + + protected: + bool DeserializeFromJsonObject(const ailego::JsonObject &json_obj) override; + ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override; +}; + +struct IVFIndexParam : public BaseIndexParam { + using Pointer = std::shared_ptr; + int nlist = 1024; + int niters = 10; + std::shared_ptr l1Index = nullptr; + std::shared_ptr l2Index = nullptr; + bool use_soar = false; + + // Constructors with delegation + IVFIndexParam() : BaseIndexParam(IndexType::kIVF) {} + + IVFIndexParam(int nlist, int niters, std::shared_ptr l1Index, + std::shared_ptr l2Index) + : BaseIndexParam(IndexType::kIVF), + nlist(nlist), + niters(niters), + l1Index(std::move(l1Index)), + l2Index(std::move(l2Index)) {} + + IVFIndexParam(MetricType metric, int dim, int nlist, int niters, + std::shared_ptr l1Index, + std::shared_ptr l2Index) + : BaseIndexParam(IndexType::kIVF, metric, dim), + nlist(nlist), + niters(niters), + l1Index(std::move(l1Index)), + l2Index(std::move(l2Index)) {} + + // query param: + // topk of l1Index's param ==== IVFIndexQueryParam.nprobe + // topk of l2Index's param ==== IVFIndexQueryParam.topK + + // IVFIndexParam.metric_type === l2Index's metric_type + // IVFIndexParam.quantization === l2Index's quantization +}; + +struct HNSWIndexParam : public BaseIndexParam { + using Pointer = std::shared_ptr; + int m = kDefaultHnswNeighborCnt; // 图中每个点的邻居数 + int ef_construction = kDefaultHnswEfConstruction; // 构建时的邻居候选大小 + + // Constructors with delegation + HNSWIndexParam() : BaseIndexParam(IndexType::kHNSW) {} + + HNSWIndexParam(int m, int ef_construction) + : BaseIndexParam(IndexType::kHNSW), + m(m), + ef_construction(ef_construction) {} + + HNSWIndexParam(MetricType metric, int dim, int m, int ef_construction) + : BaseIndexParam(IndexType::kHNSW, metric, dim), + m(m), + ef_construction(ef_construction) {} + + protected: + bool DeserializeFromJsonObject(const ailego::JsonObject &json_obj) override; + ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override; +}; + +} // namespace zvec::core_interface +#endif // ZVEC_INDEX_PARAM_H diff --git a/src/core/interface/index_param_builders.h b/src/core/interface/index_param_builders.h new file mode 100644 index 0000000..2f8afb7 --- /dev/null +++ b/src/core/interface/index_param_builders.h @@ -0,0 +1,341 @@ +#ifndef ZVEC_INDEX_PARAM_BUILDERS_H +#define ZVEC_INDEX_PARAM_BUILDERS_H +#pragma once + +#include +#include "index_param.h" + +namespace zvec::core_interface { + + +// struct ConditionalIndexParam { +// // predicate / rule / threshold +// // candidate +// }; + + +// chaining calls builder +template +class BaseIndexParamBuilder { // : public + // std::enable_shared_from_this + public: + BaseIndexParamBuilder() : param(std::make_shared()) {}; + virtual ~BaseIndexParamBuilder() = default; + + ActualIndexParamBuilderType &WithVersion(int version) { + param.version = version; + return static_cast(*this); + } + ActualIndexParamBuilderType &WithIndexType(IndexType index_type) { + param->index_type = index_type; + return static_cast(*this); + } + ActualIndexParamBuilderType &WithMetricType(MetricType metric_type) { + param->metric_type = metric_type; + return static_cast(*this); + } + ActualIndexParamBuilderType &WithDimension(int dimension) { + param->dimension = dimension; + return static_cast(*this); + } + ActualIndexParamBuilderType &WithPreprocessParam( + const PreprocessorParam &preprocess_param) { + param->preprocess_param = + std::make_shared(preprocess_param); + return static_cast(*this); + } + ActualIndexParamBuilderType &WithQuantizerParam( + const QuantizerParam &quantizer_param) { + param->quantizer_param = quantizer_param; + return static_cast(*this); + } + // ActualIndexParamBuilderType &WithRefinerParam( + // const RefinerParam &refiner_param) { + // param->refiner_param = refiner_param; + // return static_cast(*this); + // } + // ActualIndexParamBuilderType &WithDefaultQueryParam( + // const BaseIndexQueryParam &default_query_param) { + // param->default_query_param = default_query_param; + // return static_cast(*this); + // } + + ActualIndexParamBuilderType &WithIsSparse(bool is_sparse) { + param->is_sparse = is_sparse; + return static_cast(*this); + } + ActualIndexParamBuilderType &WithDataType(DataType data_type) { + param->data_type = data_type; + return static_cast(*this); + } + + ActualIndexParamBuilderType &WithUseIDMap(bool use_id_map) { + param->use_id_map = use_id_map; + return static_cast(*this); + } + + virtual std::shared_ptr Build() = 0; + + protected: + std::shared_ptr param; +}; + +class FlatIndexParamBuilder + : public BaseIndexParamBuilder { + public: + FlatIndexParamBuilder() = default; + std::shared_ptr Build() override { + return param; + } +}; + +class IVFIndexParamBuilder + : public BaseIndexParamBuilder { + public: + IVFIndexParamBuilder() = default; + IVFIndexParamBuilder &WithNList(int nlist) { + param->nlist = nlist; + return *this; + } + IVFIndexParamBuilder &WithNiters(int niters) { + param->niters = niters; + return *this; + } + IVFIndexParamBuilder &WithL1Index(const BaseIndexParam &l1Index) { + param->l1Index = std::make_shared(l1Index); + return *this; + } + IVFIndexParamBuilder &WithL2Index(const BaseIndexParam &l2Index) { + param->l2Index = std::make_shared(l2Index); + return *this; + } + IVFIndexParamBuilder &WithUseSoar(bool use_soar) { + param->use_soar = use_soar; + return *this; + } + + std::shared_ptr Build() override { + return param; + } +}; + +class HNSWIndexParamBuilder + : public BaseIndexParamBuilder { + public: + HNSWIndexParamBuilder() = default; + HNSWIndexParamBuilder &WithM(int m) { + param->m = m; + return *this; + } + HNSWIndexParamBuilder &WithEFConstruction(int ef_construction) { + param->ef_construction = ef_construction; + return *this; + } + + std::shared_ptr Build() override { + return param; + } +}; + +// class CompositeIndexParamBuilder : public +// BaseIndexParamBuilder +// { public: +// CompositeIndexParamBuilder() = default; +// CompositeIndexParamBuilder &WithLayers(const +// std::vector> &layers) { +// param.layers = layers; +// return *this; +// } +// // with layer +// CompositeIndexParamBuilder &WithLayer(const BaseIndexParam &layer) +// { +// param.layers.push_back(std::make_shared(layer)); +// return *this; +// } + +// CompositeIndexParamBuilder &WithLayer(const BaseIndexParam &layer, +// const BaseIndexQueryParam +// &default_query_param) { +// param.layers.push_back(std::make_shared(layer)); +// param.layers.back()->default_query_param = +// std::make_shared(default_query_param); +// return *this; +// } +// std::shared_ptr Build() { return +// std::make_shared(param); } + +// private: +// CompositeIndexParam param; +// }; + + +#include +#include + +template +class BaseIndexQueryParamBuilder { + public: + // This allows derived builders to access the protected member + T m_param; + + // Fluent setters for BaseIndexQueryParam fields + Derived &with_topk(int topk) { + m_param.topk = topk; + return static_cast(*this); + } + + Derived &with_fetch_vector(bool fetch_vector) { + m_param.fetch_vector = fetch_vector; + return static_cast(*this); + } + + Derived &with_filter(std::shared_ptr filter) { + m_param.filter = std::move(filter); + return static_cast(*this); + } + + // Using a vector of uint64_t for the next one + Derived &with_bf_pks(std::shared_ptr> bf_pks) { + m_param.bf_pks = std::move(bf_pks); + return static_cast(*this); + } + + Derived &with_radius(float radius) { + m_param.radius = radius; + return static_cast(*this); + } + + Derived &with_is_linear(bool is_linear) { + m_param.is_linear = is_linear; + return static_cast(*this); + } + + Derived &with_refiner_param(RefinerParam::Pointer refiner_param) { + m_param.refiner_param = std::move(refiner_param); + return static_cast(*this); + } +}; + +// FLAT builder (no extra fields, just inherits base functionality) +class FlatQueryParamBuilder + : public BaseIndexQueryParamBuilder { + public: + FlatQueryParam::Pointer build() { + return std::make_shared(std::move(m_param)); + } +}; + +// Example Usage: +// FlatQueryParam::Pointer flat_config = FlatQueryParamBuilder() +// .with_topk(20) +// .with_fetch_vector(true) +// .build(); + +// HNSW builder (adds one specific field: ef_search) +class HNSWQueryParamBuilder + : public BaseIndexQueryParamBuilder { + public: + HNSWQueryParamBuilder &with_ef_search(int ef_search) { + m_param.ef_search = ef_search; + return *this; + } + + HNSWQueryParam::Pointer build() { + return std::make_shared(std::move(m_param)); + } +}; + +// Example Usage: +// HNSWQueryParam::Pointer hnsw_config = HNSWQueryParamBuilder() +// .with_topk(5) +// .with_ef_search(128) // HNSW specific +// .with_is_linear(false) +// .build(); + +// IVF builder (adds specific fields: nprobe, l1QueryParam, l2QueryParam) +class IVFQueryParamBuilder + : public BaseIndexQueryParamBuilder { + public: + IVFQueryParamBuilder &with_nprobe(int nprobe) { + m_param.nprobe = nprobe; + return *this; + } + + // Since l1QueryParam and l2QueryParam are shared_ptr to BaseIndexQueryParam, + // they can accept ANY derived configuration object. + IVFQueryParamBuilder &with_l1_query_param( + BaseIndexQueryParam::Pointer l1QueryParam) { + m_param.l1QueryParam = std::move(l1QueryParam); + return *this; + } + + IVFQueryParamBuilder &with_l2_query_param( + BaseIndexQueryParam::Pointer l2QueryParam) { + m_param.l2QueryParam = std::move(l2QueryParam); + return *this; + } + + IVFQueryParam::Pointer build() { + return std::make_shared(std::move(m_param)); + } +}; + +// Example Usage: +// // First, build the required nested params +// auto nested_hnsw = HNSWQueryParamBuilder().with_ef_search(64).build(); +// +// // Then, build the IVF param +// IVFQueryParam::Pointer ivf_config = IVFQueryParamBuilder() +// .with_topk(10) +// .with_nprobe(50) // IVF specific +// .with_l1_query_param(nested_hnsw) // Set a nested config object +// .build(); + + +namespace predefined { +// some predefined index param builders, e.g., SCANN +class SCANNIndexParamBuilder { + public: + // alias SCANNIIndexParam = xxxxx + std::shared_ptr Build() { + // SCANN + auto param_ptr = + IVFIndexParamBuilder() + .WithNList(40000) // 10000000 -> 40000 + .WithUseSoar( + true) // 由于1个数据点可能对应2个partition，因此140个点中可能有重复，需要去重（保留一个取均值） + .WithQuantizerParam(QuantizerParam(QuantizerType::kQuickADC)) + // .WithDefaultQueryParam( + // IVFQueryParamBuilder().with_topk(140).with_nprobe(68).build()) + // .WithRefinerParam(RefinerParam{ + // 10, // 140 -> 10 + // nullptr, + // std::make_shared( + // QuantizerParam{QuantizerType::kFP16}), + // }) + .WithL1Index(*( + IVFIndexParamBuilder() + .WithMetricType( + MetricType::kInnerProduct) // Layer2 flat index + .WithNList(700) // 40000 -> 700 + .WithQuantizerParam( + QuantizerParam{QuantizerType::kQuickADC}) + // .WithDefaultQueryParam(IVFQueryParamBuilder() + // .with_topk(68) + // .with_nprobe(20) + // .build()) + .WithL1Index(*( + FlatIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + // implicit : + // .WithDefaultQueryParam(FlatQueryParamBuilder().with_topk(20).build()) + .Build())) + .Build())) + .Build(); + + return param_ptr; + } +}; +} // namespace predefined +} // namespace zvec::core_interface +#endif // ZVEC_INDEX_PARAM_BUILDERS_H \ 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..ec196bf --- /dev/null +++ b/src/core/interface/indexes/flat_index.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. +// +// Created by wangjianning.wjn on 8/28/25. +// + +#include +#include +#include "algorithm/flat/flat_utility.h" +#include "interface/index.h" + +namespace zvec::core_interface { + +int FlatIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { + param_ = dynamic_cast(param); + + proxima_index_params_.set(core::FLAT_BUILDER_COLUMN_MAJOR_ORDER, + param_.major_order == IndexMeta::MO_COLUMN); + proxima_index_params_.set(core::FLAT_SEARCHER_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); + } + + 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..7d42691 --- /dev/null +++ b/src/core/interface/indexes/hnsw_index.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 "algorithm/hnsw/hnsw_params.h" +#include "algorithm/hnsw_sparse/hnsw_sparse_params.h" +#include "interface/index.h" + +namespace zvec::core_interface { + +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); + 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); + context->update(params); + 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/ivf_index.cc b/src/core/interface/indexes/ivf_index.cc new file mode 100644 index 0000000..debd80e --- /dev/null +++ b/src/core/interface/indexes/ivf_index.cc @@ -0,0 +1,292 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "algorithm/ivf/ivf_params.h" +#include "interface/index.h" + +namespace zvec::core_interface { + +static constexpr uint64_t kInvalidKey = std::numeric_limits::max(); + +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: { + 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; + } + } + + 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; + } + is_trained_ = true; + } + is_open_ = true; + return 0; +} + +int IVFIndex::GenerateHolder() { + if (param_.data_type == DataType::DT_FP16) { + auto holder = + std::make_shared>( + param_.dimension); + for (auto doc : doc_cache_) { + ailego::NumericalVector vec(doc.second); + if (doc.first == kInvalidKey) { + continue; + } + if (!holder->emplace(doc.first, vec)) { + LOG_ERROR("Failed to add vector"); + return core::IndexError_Runtime; + } + } + holder_ = holder; + } else if (param_.data_type == DataType::DT_FP32) { + auto holder = + std::make_shared>( + param_.dimension); + for (auto doc : doc_cache_) { + ailego::NumericalVector vec(doc.second); + if (doc.first == kInvalidKey) { + continue; + } + if (!holder->emplace(doc.first, vec)) { + LOG_ERROR("Failed to add vector"); + return core::IndexError_Runtime; + } + } + holder_ = holder; + } else if (param_.data_type == DataType::DT_INT8) { + auto holder = + std::make_shared>( + param_.dimension); + for (auto doc : doc_cache_) { + ailego::NumericalVector vec(doc.second); + if (doc.first == kInvalidKey) { + continue; + } + if (!holder->emplace(doc.first, vec)) { + LOG_ERROR("Failed to add vector"); + return core::IndexError_Runtime; + } + } + holder_ = holder; + } else { + LOG_ERROR("data_type is not support"); + return core::IndexError_Runtime; + } + if (converter_) { + core::IndexConverter::TrainAndTransform(converter_, holder_); + holder_ = converter_->result(); + } + return 0; +} + +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); + 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; + } + 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); + + 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) { + // TODO: 1. sparse; 2. default ef + ailego::Params params; + // need fix + params.set(core::PARAM_IVF_BUILDER_CENTROID_COUNT, + 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); + 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; + } + is_trained_ = true; + return 0; +} +} // namespace zvec::core_interface \ No newline at end of file diff --git a/src/core/interface/utils/utils.h b/src/core/interface/utils/utils.h new file mode 100644 index 0000000..74c542f --- /dev/null +++ b/src/core/interface/utils/utils.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 + +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..cbc1049 --- /dev/null +++ b/src/core/metric/CMakeLists.txt @@ -0,0 +1,11 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +cc_library( + NAME core_metric + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS zvec_ailego core_framework + INCS . ${PROJECT_ROOT_DIR}/src/core + 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..525d3fb --- /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 "framework/index_error.h" +#include "framework/index_factory.h" + +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..8c74b7d --- /dev/null +++ b/src/core/metric/euclidean_metric.cc @@ -0,0 +1,1098 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_error.h" +#include "framework/index_factory.h" +#include "framework/index_metric.h" + +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)]; +} + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle HammingDistanceMatrix32( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::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)]; +} + +#if defined(AILEGO_M64) +static inline IndexMetric::MatrixDistanceHandle HammingDistanceMatrix64( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::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)]; +} +#endif // AILEGO_M64 + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle +HammingSquareRootDistanceMatrix32(size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::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)]; +} + +#if defined(AILEGO_M64) +static inline IndexMetric::MatrixDistanceHandle +HammingSquareRootDistanceMatrix64(size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::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)]; +} +#endif // AILEGO_M64 + +/*! 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 && + dt != IndexMeta::DataType::DT_BINARY32 && + dt != IndexMeta::DataType::DT_BINARY64) { + 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_BINARY32: + return reinterpret_cast( + ailego::HammingDistanceMatrix::Compute); + +#if defined(AILEGO_M64) + case IndexMeta::DataType::DT_BINARY64: + return reinterpret_cast( + ailego::HammingDistanceMatrix::Compute); +#endif // AILEGO_M64 + + 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_BINARY32: + return HammingDistanceMatrix32(m, n); + +#if defined(AILEGO_M64) + case IndexMeta::DataType::DT_BINARY64: + return HammingDistanceMatrix64(m, n); +#endif // AILEGO_M64 + + 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_BINARY32: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + +#if defined(AILEGO_M64) + case IndexMeta::DataType::DT_BINARY64: + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); +#endif // AILEGO_M64 + + 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 && + dt != IndexMeta::DataType::DT_BINARY32 && + dt != IndexMeta::DataType::DT_BINARY64) { + 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_BINARY32: + return reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute); + +#if defined(AILEGO_M64) + case IndexMeta::DataType::DT_BINARY64: + return reinterpret_cast( + ailego::HammingSquareRootDistanceMatrix::Compute); +#endif // AILEGO_M64 + + 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_BINARY32: + return HammingSquareRootDistanceMatrix32(m, n); + +#if defined(AILEGO_M64) + case IndexMeta::DataType::DT_BINARY64: + return HammingSquareRootDistanceMatrix64(m, n); +#endif // AILEGO_M64 + + 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 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/hamming_metric.cc b/src/core/metric/hamming_metric.cc new file mode 100644 index 0000000..b3125dc --- /dev/null +++ b/src/core/metric/hamming_metric.cc @@ -0,0 +1,237 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT 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 "framework/index_error.h" +#include "framework/index_factory.h" + +namespace zvec { +namespace core { + +//! Retrieve distance function for index features +static inline IndexMetric::MatrixDistanceHandle HammingDistanceMatrix32( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::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)]; +} + +#if defined(AILEGO_M64) +static inline IndexMetric::MatrixDistanceHandle HammingDistanceMatrix64( + size_t m, size_t n) { + static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = { + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr, nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + nullptr}, + {reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::Compute), + reinterpret_cast( + ailego::HammingDistanceMatrix::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)]; +} +#endif // AILEOG_M64 + +/*! Hamming Metric + */ +class HammingMetric : public IndexMetric { + public: + //! Initialize Metric + int init(const IndexMeta &meta, const ailego::Params &index_params) override { + if (meta.data_type() != IndexMeta::DataType::DT_BINARY32 && + meta.data_type() != IndexMeta::DataType::DT_BINARY64) { + return IndexError_Unsupported; + } + if (IndexMeta::UnitSizeof(meta.data_type()) != meta.unit_size()) { + return IndexError_Unsupported; + } + feature_type_ = meta.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() == feature_type_ && + meta.unit_size() == IndexMeta::UnitSizeof(feature_type_)); + } + + //! Retrieve if it matched + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return (qmeta.data_type() == feature_type_ && + qmeta.unit_size() == IndexMeta::UnitSizeof(feature_type_) && + qmeta.dimension() == meta.dimension()); + } + + //! Retrieve distance function for query + MatrixDistance distance(void) const override { +#if defined(AILEGO_M64) + if (feature_type_ == IndexMeta::DataType::DT_BINARY64) { + return reinterpret_cast( + ailego::HammingDistanceMatrix::Compute); + } +#endif + if (feature_type_ == IndexMeta::DataType::DT_BINARY32) { + return reinterpret_cast( + ailego::HammingDistanceMatrix::Compute); + } + return nullptr; + } + + MatrixBatchDistance batch_distance(void) const override { +#if defined(AILEGO_M64) + if (feature_type_ == IndexMeta::DataType::DT_BINARY64) { + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + } +#endif + if (feature_type_ == IndexMeta::DataType::DT_BINARY32) { + return reinterpret_cast( + ailego::BaseDistance::ComputeBatch); + } + return nullptr; + } + + //! Retrieve distance function for index features + MatrixDistance distance_matrix(size_t m, size_t n) const override { +#if defined(AILEGO_M64) + if (feature_type_ == IndexMeta::DataType::DT_BINARY64) { + return HammingDistanceMatrix64(m, n); + } +#endif + if (feature_type_ == IndexMeta::DataType::DT_BINARY32) { + return HammingDistanceMatrix32(m, n); + } + 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 feature_type_{IndexMeta::DataType::DT_BINARY32}; + ailego::Params params_{}; +}; + +INDEX_FACTORY_REGISTER_METRIC_ALIAS(Hamming, HammingMetric); + +} // 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..9f6a746 --- /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 "framework/index_error.h" +#include "framework/index_factory.h" +#include "framework/index_metric.h" + +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..b5719e8 --- /dev/null +++ b/src/core/metric/metric_params.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 + +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"; + +} // 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..b8140c7 --- /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 "framework/index_error.h" +#include "framework/index_factory.h" +#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